fix: fix unregister on quit

src/app.rs

@@ -1,3 +1,4 @@
+use crate::media::{cache::set_cache_dir, engine::StreamEngine};
 use crate::{
     call::{ActiveCallRef, sip::Invitation},
     callrecord::{
@@ -14,10 +15,9 @@ use crate::{
         registration::{RegistrationHandle, UserCredential},
     },
 };
-
-use crate::media::{cache::set_cache_dir, engine::StreamEngine};
 use anyhow::Result;
 use chrono::{DateTime, Local};
+use futures::{FutureExt, future};
 use humantime::parse_duration;
 use rsip::prelude::HeadersExt;
 use rsipstack::transaction::{
@@ -39,6 +39,7 @@ use tracing::{info, warn};
 pub struct AppStateInner {
     pub config: Arc<Config>,
     pub token: CancellationToken,
+    pub endpoint_token: CancellationToken,
     pub stream_engine: Arc<StreamEngine>,
     pub callrecord_sender: Option<CallRecordSender>,
     pub endpoint: Endpoint,
@@ -139,37 +140,52 @@ impl AppStateInner {
             }
         }
 
-        tokio::select! {
-            _ = token.cancelled() => {
-                info!("cancelled");
-            }
-            result = endpoint_inner.serve() => {
-                if let Err(e) = result {
-                    info!("endpoint serve error: {:?}", e);
+        let serving = endpoint_inner.serve();
+        let processing =
+            app_state_clone.process_incoming_request(dialog_layer.clone(), incoming_txs);
+        let mut cancelled = false;
+        let stop_registration = future::pending().boxed();
+        tokio::pin!(serving);
+        tokio::pin!(processing);
+        tokio::pin!(stop_registration);
+
+        loop {
+            tokio::select! {
+                _ = token.cancelled(), if !cancelled=> {
+                    cancelled = true;
+                    let timeout = self
+                        .config
+                        .graceful_shutdown
+                        .map(|_| Duration::from_secs(5));
+                    *stop_registration = self.stop_registration(timeout).boxed();
                 }
-            }
-            result = app_state_clone.process_incoming_request(dialog_layer.clone(), incoming_txs) => {
-                if let Err(e) = result {
-                    info!("process incoming request error: {:?}", e);
+                result = &mut serving => {
+                    if let Err(e) = result {
+                        info!("endpoint serve error: {:?}", e);
+                    }
+                    break;
                 }
-            },
-        }
-
-        // Wait for registration to stop, if not stopped within 50 seconds,
-        // force stop it.
-        let timeout = self
-            .config
-            .graceful_shutdown
-            .map(|_| Duration::from_secs(50));
-
-        match self.stop_registration(timeout).await {
-            Ok(_) => {
-                info!("registration stopped, waiting for clear");
-            }
-            Err(e) => {
-                warn!("failed to stop registration: {:?}", e);
+                result = &mut processing => {
+                    if let Err(e) = result {
+                        info!("process incoming request error: {:?}", e);
+                    }
+                    break;
+                },
+                result = &mut stop_registration => {
+                    match result {
+                        Ok(_) => {
+                            info!("registration stopped, waiting for clear");
+                        }
+                        Err(e) => {
+                            warn!("failed to stop registration: {:?}", e);
+                        }
+                    }
+                    break;
+                },
             }
         }
+
+        self.endpoint_token.cancel();
         info!("stopping");
         Ok(())
     }
@@ -673,13 +689,16 @@ impl AppStateBuilder {
         } else {
             crate::net_tool::get_first_non_loopback_interface()?
         };
-        let transport_layer = rsipstack::transport::TransportLayer::new(token.clone());
+
+        let endpoint_cancel_token = CancellationToken::new();
+        let transport_layer =
+            rsipstack::transport::TransportLayer::new(endpoint_cancel_token.clone());
         let local_addr: SocketAddr = format!("{}:{}", local_ip, config.udp_port).parse()?;
 
         let udp_conn = rsipstack::transport::udp::UdpConnection::create_connection(
             local_addr,
             None,
-            Some(token.child_token()),
+            Some(endpoint_cancel_token.child_token()),
         )
         .await
         .map_err(|e| anyhow::anyhow!("Create useragent UDP connection: {} {}", local_addr, e))?;
@@ -694,7 +713,7 @@ impl AppStateBuilder {
         }
 
         let mut endpoint_builder = endpoint_builder
-            .with_cancel_token(token.child_token())
+            .with_cancel_token(endpoint_cancel_token.clone())
             .with_transport_layer(transport_layer)
             .with_option(endpoint_option);
 
@@ -765,6 +784,7 @@ impl AppStateBuilder {
             total_calls: AtomicU64::new(0),
             total_failed_calls: AtomicU64::new(0),
             uptime: Local::now(),
+            endpoint_token: endpoint_cancel_token,
         });
 
         Ok(app_state)

src/config.rs

@@ -90,6 +90,10 @@ fn default_config_rtp_latching() -> Option<bool> {
     Some(true)
 }
 
+fn default_graceful_shutdown() -> Option<bool> {
+    Some(true)
+}
+
 fn default_config_useragent() -> Option<String> {
     Some(format!(
         "active-call({} miuda.ai)",
@@ -182,6 +186,7 @@ pub struct Config {
     #[serde(default = "default_config_useragent")]
     pub useragent: Option<String>,
     pub register_users: Option<Vec<RegisterOption>>,
+    #[serde(default = "default_graceful_shutdown")]
     pub graceful_shutdown: Option<bool>,
     pub handler: Option<InviteHandlerConfig>,
     pub accept_timeout: Option<String>,

src/main.rs

@@ -4,6 +4,7 @@ use axum::response::IntoResponse;
 use axum::routing::get;
 use clap::Parser;
 use dotenvy::dotenv;
+use futures::{FutureExt, future};
 use reqwest::StatusCode;
 use std::sync::Arc;
 use tokio::signal;
@@ -273,21 +274,47 @@ async fn main() -> Result<()> {
         .nest_service("/static", ServeDir::new("static"))
         .with_state(app_state.clone());
 
-    tokio::select! {
-        result = axum::serve(listener, app) => {
-            if let Err(e) = result {
-                warn!("axum serve error: {:?}", e);
+    let app_state_clone = app_state.clone();
+    let graceful_shutdown = config.graceful_shutdown.unwrap_or_default();
+
+    let axum_serving = axum::serve(listener, app).into_future();
+    let app_state_serving = app_state_clone.serve();
+    let mut canceled = false;
+    let cancel_timeout = future::pending().boxed();
+
+    tokio::pin!(axum_serving);
+    tokio::pin!(app_state_serving);
+    tokio::pin!(cancel_timeout);
+
+    loop {
+        tokio::select! {
+            result = &mut axum_serving => {
+                if let Err(e) = result {
+                    warn!("axum serve error: {:?}", e);
+                }
+                break;
             }
-        }
-        res = app_state.serve() => {
-            if let Err(e) = res {
-                warn!("AppState server error: {}", e);
+            res = &mut app_state_serving => {
+                if let Err(e) = res {
+                    warn!("AppState server error: {}", e);
+                }
+                break;
+            }
+            _ = signal::ctrl_c(), if !canceled => {
+                info!("Shutdown signal received");
+                if graceful_shutdown {
+                    app_state.stop();
+                    *cancel_timeout = tokio::time::sleep(tokio::time::Duration::from_secs(5)).boxed();
+                    canceled = true;
+                } else {
+                    break;
+                }
+            }
+            _ = &mut cancel_timeout => {
+                warn!("Shutdown timeout");
+                break;
             }
-        }
-        _ = signal::ctrl_c() => {
-            info!("Shutdown signal received");
         }
     }
-    info!("Shutting down...");
     Ok(())
 }