diff --git a/src/router/run.rs b/src/router/run.rs index 57f1c048..50d0c2f2 100644 --- a/src/router/run.rs +++ b/src/router/run.rs @@ -3,8 +3,11 @@ use std::{ time::Duration, }; -use futures::FutureExt; -use tuwunel_core::{Error, Result, Server, debug, debug_error, debug_info, error, info}; +use futures::{FutureExt, pin_mut}; +use tuwunel_core::{ + Error, Result, Server, debug, debug_error, debug_info, error, info, + utils::{BoolExt, future::OptionFutureExt}, +}; use tuwunel_service::Services; use crate::{handle::ServerHandle, serve}; @@ -24,27 +27,23 @@ pub(crate) async fn run(services: Arc) -> Result { .runtime() .spawn(signal(server.clone(), handle.clone())); - let mut listener = if services.config.listening { - let future = serve::serve(services.clone(), handle); - server - .runtime() - .spawn(future) - .map(|res| res.map_err(Error::from).unwrap_or_else(Err)) - .boxed() - } else { - let server = server.clone(); - async move { - server.until_shutdown().await; - Ok(()) - } - .boxed() - }; + let listener = services + .config + .listening + .then_async(|| { + server + .runtime() + .spawn(serve::serve(services.clone(), handle)) + .map(|res| res.map_err(Error::from).unwrap_or_else(Err)) + }) + .unwrap_or_else_async(|| server.until_shutdown().map(Ok)); // Focal point debug!("Running"); + pin_mut!(listener); let res = tokio::select! { - res = &mut listener => res, - res = services.poll() => handle_services_poll(server, res, listener).await, + res = &mut listener => res.unwrap_or(Ok(())), + res = services.poll() => handle_services_finish(server, res, listener.await), }; // Join the signal handler before we leave. @@ -127,10 +126,10 @@ fn handle_shutdown(server: &Arc, handle: &ServerHandle) { handle.graceful_shutdown(Some(timeout)); } -async fn handle_services_poll( +fn handle_services_finish( server: &Arc, result: Result, - listener: impl Future, + listener: Option, ) -> Result { debug!("Service manager finished: {result:?}"); @@ -140,7 +139,7 @@ async fn handle_services_poll( error!("Failed to send shutdown signal: {e}"); } - if let Err(e) = listener.await { + if let Some(Err(e)) = listener { error!("Client listener task finished with error: {e}"); }