From 2c45c85eb2929bea352742893e3244a74b1d1ebc Mon Sep 17 00:00:00 2001
From: ITSYABOINEUTRAL <okenwa.nwagu@mcis.sch.ng>
Date: Fri, 25 Sep 2026 18:01:58 +0100
Subject: [PATCH] fix: gracefully shut down queue worker

---
 backend/src/main.rs                         | 36 +++++++++++++++------
 backend/src/services/disbursement_worker.rs | 17 ++++++++--
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/backend/src/main.rs b/backend/src/main.rs
index 24ef3d3..26cf37a 100644
--- a/backend/src/main.rs
+++ b/backend/src/main.rs
@@ -393,8 +393,7 @@ async fn main() {
         )
         .layer(cors_layer(&config.cors_allowed_origins));
 
-    // #726 — Track background worker tasks so they can be allowed to finish
-    // their current batch item before the process exits on shutdown.
+    // #726 — Keep background worker tasks alive for the lifetime of the server.
     let mut worker_handles: Vec<tokio::task::JoinHandle<()>> = Vec::new();
     worker_handles.push(event_streamer_handle);
 
@@ -445,9 +444,11 @@ async fn main() {
     // BE-554 / Issue #936: Background queue worker processing bulk disbursement requests asynchronously
     let disbursement_pool = pool.clone();
     let disbursement_config = services::disbursement_worker::DisbursementWorkerConfig::from_env();
-    worker_handles.push(tokio::spawn(async move {
-        services::disbursement_worker::run(disbursement_pool, disbursement_config).await;
-    }));
+    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
+    let disbursement_handle = tokio::spawn(async move {
+        services::disbursement_worker::run(disbursement_pool, disbursement_config, shutdown_rx)
+            .await;
+    });
 
     let addr = SocketAddr::from(([0, 0, 0, 0], 8080));
     tracing::info!("Listening on {}", addr);
@@ -457,12 +458,29 @@ async fn main() {
     // on every request so IP-based rate limiting (auth::auth_rate_limit,
     // rate_limiter_middleware) has a real peer address to fall back on when
     // there's no `X-Forwarded-For` header (i.e. no reverse proxy in front).
-    axum::serve(
+    let server = axum::serve(
         listener,
         app.into_make_service_with_connect_info::<SocketAddr>(),
-    )
-    .await
-    .unwrap();
+    );
+
+    tokio::select! {
+        result = server => {
+            if let Err(error) = result {
+                tracing::error!("HTTP server stopped: {error}");
+            }
+        }
+        result = tokio::signal::ctrl_c() => {
+            match result {
+                Ok(()) => tracing::info!("Shutdown signal received"),
+                Err(error) => tracing::error!("Failed to listen for shutdown signal: {error}"),
+            }
+        }
+    }
+
+    let _ = shutdown_tx.send(());
+    if let Err(error) = disbursement_handle.await {
+        tracing::error!("Disbursement worker task failed during shutdown: {error}");
+    }
 }
 
 // ── /api/v1/config — mobile minimum-version gate (#805) ───────────────────────
diff --git a/backend/src/services/disbursement_worker.rs b/backend/src/services/disbursement_worker.rs
index a9c9ff1..7b15d79 100644
--- a/backend/src/services/disbursement_worker.rs
+++ b/backend/src/services/disbursement_worker.rs
@@ -111,8 +111,12 @@ impl DisbursementWorkerConfig {
 /// Outcome of one dispatch attempt (alias of SdpOutcome for backward compat).
 type DispatchOutcome = SdpOutcome;
 
-/// Entry point. Runs until the process exits.
-pub async fn run(pool: PgPool, config: DisbursementWorkerConfig) {
+/// Entry point. Runs until shutdown is requested.
+pub async fn run(
+    pool: PgPool,
+    config: DisbursementWorkerConfig,
+    mut shutdown: tokio::sync::oneshot::Receiver<()>,
+) {
     tracing::info!(
         worker_id = %config.worker_id,
         "Starting disbursement worker (interval={:?}, claim_size={}, max_attempts={})",
@@ -132,7 +136,14 @@ pub async fn run(pool: PgPool, config: DisbursementWorkerConfig) {
     let mut interval = tokio::time::interval(config.poll_interval);
 
     loop {
-        interval.tick().await;
+        tokio::select! {
+            biased;
+            _ = &mut shutdown => {
+                tracing::info!("Disbursement worker shutdown requested");
+                break;
+            }
+            _ = interval.tick() => {}
+        }
 
         if let Err(err) = reclaim_stale_leases(&pool, config.lease_timeout_secs).await {
             tracing::error!("Failed to reclaim stale leases: {err:?}");