Refactor --shutdown-after into --shutdown (#49)

CHANGELOG.md

@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.18.0] - ???
+### Changed
+
+- `shutdown-after` replaced with `shutdown` that supports three options:
+    1. `never` - server will never shutdown automatically
+    2. `after=N` - server will shutdown after N seconds
+    3. `lonely=N` - server will shutdown N seconds after no connections
+
 ## [0.17.6] - 2022-08-18
 ### Fixed
 

distant-net/src/server/config.rs

@@ -1,10 +1,88 @@
+use derive_more::{Display, Error};
 use serde::{Deserialize, Serialize};
-use std::time::Duration;
+use std::{num::ParseFloatError, str::FromStr, time::Duration};
 
 /// Represents a general-purpose set of properties tied with a server instance
 #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
 pub struct ServerConfig {
-    /// If provided, will cause server to shut down if duration is exceeded with no active
-    /// connections
-    pub shutdown_after: Option<Duration>,
+    /// Rules for how a server will shutdown automatically
+    pub shutdown: Shutdown,
+}
+
+/// Rules for how a server will shut itself down automatically
+#[derive(Copy, Clone, Debug, Display, PartialEq, Eq, Serialize, Deserialize)]
+pub enum Shutdown {
+    /// Server should shutdown immediately after duration exceeded
+    #[display(fmt = "after={}", "_0.as_secs_f32()")]
+    After(Duration),
+
+    /// Server should shutdown after no connections for over duration time
+    #[display(fmt = "lonely={}", "_0.as_secs_f32()")]
+    Lonely(Duration),
+
+    /// No shutdown logic will be applied to the server
+    #[display(fmt = "never")]
+    Never,
+}
+
+impl Shutdown {
+    /// Return duration associated with shutdown if it has one
+    pub fn duration(&self) -> Option<Duration> {
+        match self {
+            Self::Never => None,
+            Self::After(x) | Self::Lonely(x) => Some(*x),
+        }
+    }
+}
+
+impl Default for Shutdown {
+    /// By default, shutdown is never
+    fn default() -> Self {
+        Self::Never
+    }
+}
+
+/// Parsing errors that can occur for [`Shutdown`]
+#[derive(Clone, Debug, Display, Error, PartialEq, Eq)]
+pub enum ShutdownParseError {
+    #[display(fmt = "Bad value for after: {}", _0)]
+    BadValueForAfter(ParseFloatError),
+
+    #[display(fmt = "Bad value for lonely: {}", _0)]
+    BadValueForLonely(ParseFloatError),
+
+    #[display(fmt = "Missing key")]
+    MissingKey,
+
+    #[display(fmt = "Unknown key")]
+    UnknownKey,
+}
+
+impl FromStr for Shutdown {
+    type Err = ShutdownParseError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if s.eq_ignore_ascii_case("never") {
+            Ok(Self::Never)
+        } else {
+            let (left, right) = s.split_once('=').ok_or(ShutdownParseError::MissingKey)?;
+            let left = left.trim();
+            let right = right.trim();
+            if left.eq_ignore_ascii_case("after") {
+                Ok(Self::After(Duration::from_secs_f32(
+                    right
+                        .parse()
+                        .map_err(ShutdownParseError::BadValueForAfter)?,
+                )))
+            } else if left.eq_ignore_ascii_case("lonely") {
+                Ok(Self::Lonely(Duration::from_secs_f32(
+                    right
+                        .parse()
+                        .map_err(ShutdownParseError::BadValueForLonely)?,
+                )))
+            } else {
+                Err(ShutdownParseError::UnknownKey)
+            }
+        }
+    }
 }

distant-net/src/server/ext.rs

@@ -1,6 +1,6 @@
 use crate::{
     utils::Timer, GenericServerRef, Listener, Request, Response, Server, ServerConnection,
-    ServerCtx, ServerRef, ServerReply, ServerState, TypedAsyncRead, TypedAsyncWrite,
+    ServerCtx, ServerRef, ServerReply, ServerState, Shutdown, TypedAsyncRead, TypedAsyncWrite,
 };
 use log::*;
 use serde::{de::DeserializeOwned, Serialize};
@@ -86,11 +86,19 @@ where
     //       shutdown receiver to last forever in the event that there is no shutdown configured,
     //       not return immediately, which is what would happen if the sender was dropped.
     #[allow(clippy::manual_map)]
-    let mut shutdown_timer = match config.shutdown_after {
-        Some(duration) => Some(Timer::new(duration, async move {
+    let mut shutdown_timer = match config.shutdown {
+        // Create a timer, start it, and drop it so it will always happen
+        Shutdown::After(duration) => {
+            Timer::new(duration, async move {
+                let _ = shutdown_tx.send(()).await;
+            })
+            .start();
+            None
+        }
+        Shutdown::Lonely(duration) => Some(Timer::new(duration, async move {
             let _ = shutdown_tx.send(()).await;
         })),
-        None => None,
+        Shutdown::Never => None,
     };
 
     if let Some(timer) = shutdown_timer.as_mut() {
@@ -124,7 +132,7 @@ where
             _ = shutdown_rx.recv() => {
                 info!(
                     "Server shutdown triggered after {}s",
-                    config.shutdown_after.unwrap_or_default().as_secs_f32(),
+                    config.shutdown.duration().unwrap_or_default().as_secs_f32(),
                 );
                 break;
             }
@@ -289,12 +297,12 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn should_shutdown_if_no_connections_received_after_n_secs_when_config_set() {
+    async fn should_lonely_shutdown_if_no_connections_received_after_n_secs_when_config_set() {
         let (_tx, listener) = make_listener(100);
 
         let server = ServerExt::start(
             TestServer(ServerConfig {
-                shutdown_after: Some(Duration::from_millis(100)),
+                shutdown: Shutdown::Lonely(Duration::from_millis(100)),
             }),
             listener,
         )
@@ -307,7 +315,8 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn should_shutdown_if_last_connection_terminated_and_then_no_connections_after_n_secs() {
+    async fn should_lonely_shutdown_if_last_connection_terminated_and_then_no_connections_after_n_secs(
+    ) {
         // Create a test listener where we will forward a connection
         let (tx, listener) = make_listener(100);
 
@@ -319,7 +328,7 @@ mod tests {
 
         let server = ServerExt::start(
             TestServer(ServerConfig {
-                shutdown_after: Some(Duration::from_millis(100)),
+                shutdown: Shutdown::Lonely(Duration::from_millis(100)),
             }),
             listener,
         )
@@ -335,7 +344,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn should_not_shutdown_as_long_as_a_connection_exists() {
+    async fn should_not_lonely_shutdown_as_long_as_a_connection_exists() {
         // Create a test listener where we will forward a connection
         let (tx, listener) = make_listener(100);
 
@@ -347,7 +356,7 @@ mod tests {
 
         let server = ServerExt::start(
             TestServer(ServerConfig {
-                shutdown_after: Some(Duration::from_millis(100)),
+                shutdown: Shutdown::Lonely(Duration::from_millis(100)),
             }),
             listener,
         )
@@ -360,12 +369,54 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn should_never_shutdown_if_config_not_set() {
+    async fn should_shutdown_after_n_seconds_even_with_connections_if_config_set_to_after() {
+        let (tx, listener) = make_listener(100);
+
+        // Make bounded transport pair and send off one of them to act as our connection
+        let (_transport, connection) = MpscTransport::<Request<u16>, Response<String>>::pair(100);
+        tx.send(connection.into_split())
+            .await
+            .expect("Failed to feed listener a connection");
+
+        let server = ServerExt::start(
+            TestServer(ServerConfig {
+                shutdown: Shutdown::After(Duration::from_millis(100)),
+            }),
+            listener,
+        )
+        .expect("Failed to start server");
+
+        // Wait for some time
+        tokio::time::sleep(Duration::from_millis(300)).await;
+
+        assert!(server.is_finished(), "Server shutdown not triggered!");
+    }
+
+    #[tokio::test]
+    async fn should_shutdown_after_n_seconds_if_config_set_to_after() {
+        let (_tx, listener) = make_listener(100);
+
+        let server = ServerExt::start(
+            TestServer(ServerConfig {
+                shutdown: Shutdown::After(Duration::from_millis(100)),
+            }),
+            listener,
+        )
+        .expect("Failed to start server");
+
+        // Wait for some time
+        tokio::time::sleep(Duration::from_millis(300)).await;
+
+        assert!(server.is_finished(), "Server shutdown not triggered!");
+    }
+
+    #[tokio::test]
+    async fn should_never_shutdown_if_config_set_to_never() {
         let (_tx, listener) = make_listener(100);
 
         let server = ServerExt::start(
             TestServer(ServerConfig {
-                shutdown_after: None,
+                shutdown: Shutdown::Never,
             }),
             listener,
         )

distant-net/src/utils.rs

@@ -161,5 +161,20 @@ mod tests {
             assert!(timer.callback.is_finished(), "Callback not finished");
             assert!(rx.try_recv().is_ok(), "Callback not triggered");
         }
+
+        #[tokio::test]
+        async fn should_trigger_callback_even_if_timer_dropped() {
+            let (tx, mut rx) = mpsc::channel(1);
+
+            let mut timer = Timer::new(Duration::default(), async move {
+                let _ = tx.send(()).await;
+            });
+            timer.start();
+            drop(timer);
+
+            tokio::time::sleep(Duration::from_millis(300)).await;
+
+            assert!(rx.try_recv().is_ok(), "Callback not triggered");
+        }
     }
 }

distant-ssh2/tests/sshd/mod.rs

@@ -529,7 +529,7 @@ pub async fn launched_client(
     let client = ssh_client
         .launch_and_connect(DistantLaunchOpts {
             binary,
-            args: "--shutdown-after 10".to_string(),
+            args: "--shutdown after=10".to_string(),
             ..Default::default()
         })
         .await

src/cli/commands/server.rs

@@ -12,10 +12,7 @@ use distant_core::{
     DistantApiServer, DistantSingleKeyCredentials, Host,
 };
 use log::*;
-use std::{
-    io::{self, Read, Write},
-    time::Duration,
-};
+use std::io::{self, Read, Write};
 
 #[derive(Debug, Subcommand)]
 pub enum ServerSubcommand {
@@ -177,7 +174,7 @@ impl ServerSubcommand {
                     }
                 );
                 let server = DistantApiServer::local(NetServerConfig {
-                    shutdown_after: get!(shutdown_after).map(Duration::from_secs_f32),
+                    shutdown: get!(shutdown).unwrap_or_default(),
                 })
                 .context("Failed to create local distant api")?
                 .start(addr, get!(port).unwrap_or_else(|| 0.into()), codec)

src/config/server/listen.rs

@@ -1,7 +1,10 @@
 use anyhow::Context;
 use clap::Args;
 use derive_more::Display;
-use distant_core::{net::PortRange, Map};
+use distant_core::{
+    net::{PortRange, Shutdown},
+    Map,
+};
 use serde::{Deserialize, Serialize};
 use std::{
     env,
@@ -41,12 +44,15 @@ pub struct ServerListenConfig {
     #[clap(short = '6', long)]
     pub use_ipv6: bool,
 
-    /// The time in seconds before shutting down the server if there are no active
-    /// connections. The countdown begins once all connections have closed and
-    /// stops when a new connection is made. In not specified, the server will not
-    /// shutdown at any point when there are no active connections.
+    /// Logic to apply to server when determining when to shutdown automatically
+    ///
+    /// 1. "never" means the server will never automatically shut down
+    /// 2. "after=<N>" means the server will shut down after N seconds
+    /// 3. "lonely=<N>" means the server will shut down after N seconds with no connections
+    ///
+    /// Default is to never shut down
     #[clap(long)]
-    pub shutdown_after: Option<f32>,
+    pub shutdown: Option<Shutdown>,
 
     /// Changes the current working directory (cwd) to the specified directory
     #[clap(long)]
@@ -64,9 +70,9 @@ impl From<Map> for ServerListenConfig {
                 .remove("use_ipv6")
                 .and_then(|x| x.parse::<bool>().ok())
                 .unwrap_or_default(),
-            shutdown_after: map
-                .remove("shutdown_after")
-                .and_then(|x| x.parse::<f32>().ok()),
+            shutdown: map
+                .remove("shutdown")
+                .and_then(|x| x.parse::<Shutdown>().ok()),
             current_dir: map
                 .remove("current_dir")
                 .and_then(|x| x.parse::<PathBuf>().ok()),
@@ -88,8 +94,8 @@ impl From<ServerListenConfig> for Map {
 
         this.insert("use_ipv6".to_string(), config.use_ipv6.to_string());
 
-        if let Some(x) = config.shutdown_after {
-            this.insert("shutdown_after".to_string(), x.to_string());
+        if let Some(x) = config.shutdown {
+            this.insert("shutdown".to_string(), x.to_string());
         }
 
         if let Some(x) = config.current_dir {