diff --git a/src/app.rs b/src/app.rs index f22b6d3..e326117 100644 --- a/src/app.rs +++ b/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, pub token: CancellationToken, + pub endpoint_token: CancellationToken, pub stream_engine: Arc, pub callrecord_sender: Option, 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) diff --git a/src/config.rs b/src/config.rs index 1d813d0..5dad7a5 100644 --- a/src/config.rs +++ b/src/config.rs @@ -90,6 +90,10 @@ fn default_config_rtp_latching() -> Option { Some(true) } +fn default_graceful_shutdown() -> Option { + Some(true) +} + fn default_config_useragent() -> Option { Some(format!( "active-call({} miuda.ai)", @@ -182,6 +186,7 @@ pub struct Config { #[serde(default = "default_config_useragent")] pub useragent: Option, pub register_users: Option>, + #[serde(default = "default_graceful_shutdown")] pub graceful_shutdown: Option, pub handler: Option, pub accept_timeout: Option, diff --git a/src/main.rs b/src/main.rs index 770c990..d64d7cb 100644 --- a/src/main.rs +++ b/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(()) }