read expiration from Set command + tests

src/server.rs

@@ -1,6 +1,8 @@
 use std::sync::Arc;
 
 use tokio::net::ToSocketAddrs;
+use tokio::sync::oneshot;
+use tokio::task::JoinHandle;
 use tokio::{net::TcpListener, sync::Semaphore};
 
 use crate::Result;
@@ -14,6 +16,7 @@ pub struct Server {
     db: Database,
     listener: TcpListener,
     connection_limit: Arc<Semaphore>,
+    expiration_manager_handle: JoinHandle<()>,
 }
 
 impl Server {
@@ -27,23 +30,39 @@ impl Server {
 
         let db = Database::new();
 
-        tokio::spawn(key_expiration_manager(db.clone()));
+        let expiration_manager_handle = tokio::spawn(key_expiration_manager(db.clone()));
 
         Ok(Self {
             db,
             connection_limit: Arc::new(Semaphore::const_new(max_connections)),
             listener,
+            expiration_manager_handle,
         })
     }
 
-    pub async fn run(&mut self) -> Result<()> {
+    pub async fn run(&mut self, mut shutdown: oneshot::Receiver<()>) -> Result<()> {
+        let shutdown = &mut shutdown;
+
         loop {
             let permit = Arc::clone(&self.connection_limit)
                 .acquire_owned()
                 .await
                 .unwrap();
 
-            let socket = self.listener.accept().await?.0;
+            let Some(socket) = ({
+                tokio::select! {
+                    socket = self.listener.accept() => Some(socket?.0),
+                    _ = &mut *shutdown => None,
+                }
+            }) else {
+                log::info!("Shutting down");
+
+                self.db.shutdown().await;
+                let _ = (&mut self.expiration_manager_handle).await;
+
+                return Ok(());
+            };
+
             let addr = socket.peer_addr()?;
 
             let connection = Connection::new(socket);