Add CancellationToken as argument to server methods

This is the advice tokio way to stop a task
locka99 · Nov 30, 2024 · 1cfeccf · 1cfeccf
1 parent 0c98f9b
commit 1cfeccf
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/lib/Cargo.toml b/lib/Cargo.toml
@@ -48,7 +48,7 @@ chrono = { version = "0.4", features = ["serde"] }
 parking_lot = { version = "0.12", features = ["send_guard"] }
 futures = "0.3"
 tokio = { version = "1", features = ["full"] }
-tokio-util = { version = "0.6", features = ["codec"] }
+tokio-util = { version = "0.7", features = ["codec"] }
 lazy_static = "1.4.0"
 regex = "1.7"
 serde = "1.0"

diff --git a/lib/src/server/server.rs b/lib/src/server/server.rs
@@ -12,6 +12,7 @@ use tokio::{
     sync::oneshot::{self, Sender},
     time::{interval_at, Duration, Instant},
 };
+use tokio_util::sync::CancellationToken;
 
 use crate::core::{config::Config, prelude::*};
 use crate::crypto::*;
@@ -220,19 +221,21 @@ impl Server {
     /// Calling this function consumes the server.
     pub fn run(self) {
         let server = Arc::new(RwLock::new(self));
-        Self::run_server(server);
+        //FIXME; What API do we we want? for now just creating a dummy token
+        let cancel_token = CancellationToken::new();
+        Self::run_server(server, cancel_token);
     }
 
     /// Runs the supplied server and blocks until it completes either by aborting or
     /// by error.
-    pub fn run_server(server: Arc<RwLock<Server>>) {
+    pub fn run_server(server: Arc<RwLock<Server>>, cancel_token: CancellationToken) {
         let single_threaded_executor = {
             let server = trace_read_lock!(server);
             let server_state = trace_read_lock!(server.server_state);
             let config = trace_read_lock!(server_state.config);
             config.performance.single_threaded_executor
         };
-        let server_task = Self::new_server_task(server);
+        let server_task = Self::new_server_task(server, cancel_token);
         // Launch
         let mut builder = if !single_threaded_executor {
             tokio::runtime::Builder::new_multi_thread()
@@ -266,7 +269,7 @@ impl Server {
     }
 
     /// Returns the main server task - the loop that waits for connections and processes them.
-    pub async fn new_server_task(server: Arc<RwLock<Server>>) {
+    pub async fn new_server_task(server: Arc<RwLock<Server>>, cancel_token: CancellationToken) {
         // Get the address and discovery url
         let (sock_addr, discovery_server_url) = {
             let server = trace_read_lock!(server);
@@ -296,14 +299,17 @@ impl Server {
             None => {
                 error!("Cannot resolve server address, check configuration of server");
             }
-            Some(sock_addr) => Self::server_task(server, sock_addr, discovery_server_url).await,
+            Some(sock_addr) => {
+                Self::server_task(server, sock_addr, discovery_server_url, cancel_token).await
+            }
         }
     }
 
     async fn server_task<A: ToSocketAddrs>(
         server: Arc<RwLock<Server>>,
         sock_addr: A,
         discovery_server_url: Option<String>,
+        cancel_token: CancellationToken,
     ) {
         // This is returned as the main server task
         info!("Waiting for Connection");
@@ -375,6 +381,9 @@ impl Server {
             _ = rx_abort => {
                 info!("abort received");
             }
+            _ = cancel_token.cancelled() => {
+                info!("Cancellation token triggered, shutting down server");
+            }
         }
         info!("main server task is finished");
     }