Add server shutdown logic

also change Dockerfile server start from `CMD` to `ENTRYPOINT`

Author: OreQr

Cargo.lock

@@ -13,6 +13,7 @@ tokio = { version = "1.45.0", features = [
   "rt-multi-thread",
   "net",
   "io-util",
+  "signal",
 ] }
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.140"

Cargo.toml

@@ -20,4 +20,4 @@ FROM gcr.io/distroless/cc:latest
 COPY --from=builder /code/target/release/server /usr/local/bin/server
 
 # start server
-CMD [ "/usr/local/bin/server" ]
+ENTRYPOINT [ "/usr/local/bin/server" ]

Dockerfile

@@ -2,10 +2,12 @@ mod config;
 mod models;
 mod protocol;
 mod server;
+mod shutdown;
 
 use config::Config;
+use shutdown::Shutdown;
 use std::{env, io};
-use tokio::net::TcpListener;
+use tokio::{net::TcpListener, select};
 use tracing::{debug, info, info_span, Instrument};
 use tracing_subscriber::EnvFilter;
 
@@ -39,24 +41,37 @@ async fn main() -> io::Result<()> {
     let listener = TcpListener::bind(format!("{}:{}", config.host, config.port)).await?;
     info!("Server listening on port {}", config.port);
 
+    let mut shutdown = Shutdown::new()?;
+
     loop {
-        match listener.accept().await {
-            Ok((stream, addr)) => {
-                debug!("New connection from: {}", addr);
-
-                let config = config.clone();
-                tokio::spawn(
-                    async move {
-                        if let Err(e) = server::handle_client(stream, config).await {
-                            debug!("Error handling client: {}", e);
-                        }
+        select! {
+            conn = listener.accept() => {
+                match conn {
+                    Ok((stream, addr)) => {
+                        debug!("New connection from: {}", addr);
+
+                        let config = config.clone();
+                        tokio::spawn(
+                            async move {
+                                if let Err(e) = server::handle_client(stream, config).await {
+                                    debug!("Error handling client: {}", e);
+                                }
+                            }
+                            .instrument(info_span!("client", "{}", addr)),
+                        );
+                    }
+                    Err(e) => {
+                        debug!("Error accepting connection: {}", e);
                     }
-                    .instrument(info_span!("client", "{}", addr)),
-                );
+                }
             }
-            Err(e) => {
-                debug!("Error accepting connection: {}", e);
+
+            _ = shutdown.wait_for_shutdown() => {
+                break;
             }
         }
     }
+
+    info!("Shutting down server...");
+    Ok(())
 }

src/main.rs

@@ -0,0 +1,30 @@
+use std::io;
+use tokio::{select, signal};
+
+#[cfg(unix)]
+use tokio::signal::unix::{signal, SignalKind};
+
+pub struct Shutdown {
+    #[cfg(unix)]
+    sigterm: tokio::signal::unix::Signal,
+}
+impl Shutdown {
+    pub fn new() -> io::Result<Self> {
+        Ok(Self {
+            #[cfg(unix)]
+            sigterm: signal(SignalKind::terminate())?,
+        })
+    }
+
+    pub async fn wait_for_shutdown(&mut self) {
+        #[cfg(unix)]
+        let unix_signal = self.sigterm.recv();
+        #[cfg(not(unix))]
+        let unix_signal = std::future::pending::<()>();
+
+        select! {
+            _ = signal::ctrl_c() => {},
+            _ = unix_signal => {},
+        }
+    }
+}