Fix shutdown immediately happening when time not provided

src/core/server/distant/mod.rs

@@ -11,7 +11,6 @@ use crate::core::{
         PortRange,
     },
 };
-use futures::future::OptionFuture;
 use log::*;
 use std::{net::IpAddr, sync::Arc};
 use tokio::{
@@ -91,7 +90,7 @@ async fn connection_loop(
     state: Arc<Mutex<State>>,
     auth_key: Arc<SecretKey>,
     tracker: Option<Arc<Mutex<ConnTracker>>>,
-    shutdown: OptionFuture<ShutdownTask>,
+    shutdown: Option<ShutdownTask>,
     max_msg_capacity: usize,
 ) {
     let inner = async move {
@@ -121,11 +120,14 @@ async fn connection_loop(
         }
     };
 
-    tokio::select! {
-        _ = inner => {}
-        _ = shutdown => {
-            warn!("Reached shutdown timeout, so terminating");
-        }
+    match shutdown {
+        Some(shutdown) => tokio::select! {
+            _ = inner => {}
+            _ = shutdown => {
+                warn!("Reached shutdown timeout, so terminating");
+            }
+        },
+        None => inner.await,
     }
 }
 

src/core/server/relay.rs

@@ -110,11 +110,14 @@ impl RelayServer {
                 }
             };
 
-            tokio::select! {
-                _ = inner => {}
-                _ = shutdown => {
-                    warn!("Reached shutdown timeout, so terminating");
-                }
+            match shutdown {
+                Some(shutdown) => tokio::select! {
+                    _ = inner => {}
+                    _ = shutdown => {
+                        warn!("Reached shutdown timeout, so terminating");
+                    }
+                },
+                None => inner.await,
             }
         });
 

src/core/server/utils.rs

@@ -1,4 +1,3 @@
-use futures::future::OptionFuture;
 use log::*;
 use std::{
     future::Future,
@@ -29,19 +28,17 @@ impl Future for ShutdownTask {
 
 impl ShutdownTask {
     /// Given an optional timeout, will either create the shutdown task or not,
-    /// returning an optional future for the completion of the shutdown task
-    /// alongside an optional connection tracker
+    /// returning an optional shutdown task alongside an optional connection tracker
     pub fn maybe_initialize(
         duration: Option<Duration>,
-    ) -> (OptionFuture<ShutdownTask>, Option<Arc<Mutex<ConnTracker>>>) {
+    ) -> (Option<ShutdownTask>, Option<Arc<Mutex<ConnTracker>>>) {
         match duration {
             Some(duration) => {
                 let task = Self::initialize(duration);
                 let tracker = task.tracker();
-                let task: OptionFuture<_> = Some(task).into();
-                (task, Some(tracker))
+                (Some(task), Some(tracker))
             }
-            None => (None.into(), None),
+            None => (None, None),
         }
     }