feat(router): added a simple wrapper around background tasks

bin/router/Cargo.toml

@@ -42,3 +42,4 @@ mimalloc = { version = "0.1.47", features = ["override"] }
 moka = { version = "0.12.10", features = ["future"] }
 ulid = "1.2.1"
 ntex = { version = "2", features = ["tokio"] }
+tokio-util = "0.7.16"

bin/router/src/background_tasks/mod.rs

@@ -0,0 +1,51 @@
+use async_trait::async_trait;
+use futures::future::join_all;
+use std::sync::Arc;
+use tokio::task::JoinHandle;
+use tokio_util::sync::CancellationToken;
+use tracing::{debug, info};
+
+#[async_trait]
+pub trait BackgroundTask: Send + Sync {
+    fn id(&self) -> &str;
+    async fn run(&self, token: CancellationToken);
+}
+
+pub struct BackgroundTasksManager {
+    cancellation_token: CancellationToken,
+    task_handles: Vec<JoinHandle<()>>,
+}
+
+impl BackgroundTasksManager {
+    pub fn new() -> Self {
+        Self {
+            cancellation_token: CancellationToken::new(),
+            task_handles: Vec::new(),
+        }
+    }
+
+    pub fn register_task<T>(&mut self, task: T)
+    where
+        T: BackgroundTask + 'static,
+    {
+        info!("registering background task: {}", task.id());
+        let child_token = self.cancellation_token.clone();
+        let task_arc = Arc::new(task);
+
+        let handle = tokio::spawn(async move {
+            task_arc.run(child_token).await;
+        });
+
+        self.task_handles.push(handle);
+    }
+
+    pub async fn shutdown(self) {
+        info!("shutdown triggered, stopping all background tasks...");
+        self.cancellation_token.cancel();
+
+        debug!("waiting for background tasks to finish...");
+        join_all(self.task_handles).await;
+
+        println!("all background tasks have been shut down gracefully.");
+    }
+}

bin/router/src/lib.rs

@@ -1,3 +1,4 @@
+mod background_tasks;
 mod http_utils;
 mod logger;
 mod pipeline;
@@ -6,6 +7,7 @@
 use std::sync::Arc;
 
 use crate::{
+    background_tasks::BackgroundTasksManager,
     http_utils::{health::health_check_handler, landing_page::landing_page_handler},
     logger::configure_logging,
     pipeline::graphql_request_handler,
@@ -19,6 +21,7 @@
 };
 
 use hive_router_query_planner::utils::parsing::parse_schema;
+use tracing::info;
 
 async fn graphql_endpoint_handler(
     mut request: HttpRequest,
@@ -37,8 +40,9 @@
     let parsed_schema = parse_schema(&supergraph_sdl);
     let addr = router_config.http.address();
     let shared_state = RouterSharedState::new(parsed_schema, router_config);
+    let mut bg_tasks_manager = BackgroundTasksManager::new();
 
-    web::HttpServer::new(move || {
+    let maybe_error = web::HttpServer::new(move || {
         web::App::new()
             .state(shared_state.clone())
             .route("/graphql", web::to(graphql_endpoint_handler))
@@ -48,5 +52,10 @@
     .bind(addr)?
     .run()
     .await
-    .map_err(|err| err.into())
+    .map_err(|err| err.into());
+
+    info!("server stopped, clearning background tasks");
+    bg_tasks_manager.shutdown().await;
+
+    maybe_error
 }