Implement graceful server shutdown

TODO.md

@@ -55,7 +55,7 @@
 - [ ] (testing) test startup recovery / durability
 - [ ] (testing) extend property/fuzz tests to RPC surface with failure injection
 - [ ] (feat) multiple-queues/namespaces with per-queue isolation and limits
-- [ ] (productionize) implement graceful shutdown: handle SIGTERM/SIGINT, stop accepting, drain in-flight RPCs, signal background tasks to exit
+- [X] (productionize) implement graceful shutdown: handle SIGTERM/SIGINT, stop accepting, drain in-flight RPCs, signal background tasks to exit
 - [ ] (productionize) add configurablerate limiting to RPCs with backpressure
 - [ ] (productionize) RocksDB durability/tuning surfaced via config: WAL sync policy, background jobs, compaction, block cache, rate limiter, max open files
 - [ ] (productionize) data directory hardening: avoid /tmp default, permissions/ownership checks, separate WAL dir, disk space checks

src/bin/queueber/main.rs

@@ -13,6 +13,7 @@ use queueber::{
 };
 use socket2::{Domain, Protocol, Socket, Type};
 use std::sync::Arc;
+use tokio::signal::unix::{SignalKind, signal};
 use tokio::sync::Notify;
 use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};
 
@@ -116,8 +117,8 @@ async fn main() -> Result<()> {
                                 .expect("set nonblocking");
                             let listener = tokio::net::TcpListener::from_std(std_listener)
                                 .expect("tokio listener from std");
-
                             let mut conn_id: u64 = 0;
+                            let mut conn_handles: Vec<tokio::task::JoinHandle<()>> = Vec::new();
                             loop {
                                 tokio::select! {
                                     _ = async { if *shutdown_rx_worker.borrow() { } else { let _ = shutdown_rx_worker.changed().await; } } => {
@@ -129,7 +130,7 @@ async fn main() -> Result<()> {
                                         let client = queue_client.clone();
                                         let this_conn = conn_id;
                                         conn_id = conn_id.wrapping_add(1);
-                                        let _jh = tokio::task::Builder::new()
+                                        let jh = tokio::task::Builder::new()
                                             .name("rpc_server")
                                             .spawn_local(async move {
                                                 let (reader, writer) =
@@ -152,9 +153,15 @@ async fn main() -> Result<()> {
                                                 let _ = this_conn; // reserved for future naming/metrics
                                             })
                                             .unwrap();
+                                        conn_handles.push(jh);
                                     }
                                 }
                             }
+
+                            // Drain in-flight RPC tasks
+                            for jh in conn_handles {
+                                let _ = jh.await;
+                            }
                         })
                         .await;
                 });
@@ -164,6 +171,24 @@ async fn main() -> Result<()> {
         worker_handles.push(handle);
     }
 
+    // Signals: listen for SIGINT and SIGTERM once; trigger graceful shutdown on first
+    {
+        let shutdown_tx_for_signals = shutdown_tx.clone();
+        tokio::spawn(async move {
+            let mut sigint = signal(SignalKind::interrupt()).expect("install SIGINT handler");
+            let mut sigterm = signal(SignalKind::terminate()).expect("install SIGTERM handler");
+            tokio::select! {
+                _ = sigint.recv() => {
+                    tracing::info!("received SIGINT; initiating graceful shutdown");
+                }
+                _ = sigterm.recv() => {
+                    tracing::info!("received SIGTERM; initiating graceful shutdown");
+                }
+            }
+            let _ = shutdown_tx_for_signals.send(true);
+        });
+    }
+
     // Wait for shutdown signal, then join workers
     let _ = async {
         if *shutdown_rx.borrow() {

src/server/coalescer.rs

@@ -1,5 +1,5 @@
 use std::sync::Arc;
-use tokio::sync::{Notify, oneshot};
+use tokio::sync::{Notify, oneshot, watch};
 use tokio::time::Duration;
 
 use capnp::message::{Builder, HeapAllocator, TypedReader};
@@ -40,15 +40,17 @@ pub struct PollCoalescer {
     cfg: PollCoalescingConfig,
     state: Arc<tokio::sync::Mutex<VecDeque<PendingPoll>>>, // pending batch
     wake: Arc<Notify>,
+    shutdown: watch::Receiver<bool>,
 }
 
 impl PollCoalescer {
-    pub fn new(storage: Arc<RetriedStorage<Storage>>) -> Self {
+    pub fn new(storage: Arc<RetriedStorage<Storage>>, shutdown: watch::Receiver<bool>) -> Self {
         let this = Self {
             storage,
             cfg: PollCoalescingConfig::default(),
             state: Arc::new(tokio::sync::Mutex::new(VecDeque::new())),
             wake: Arc::new(Notify::new()),
+            shutdown,
         };
         this.spawn_batcher();
         this
@@ -59,15 +61,23 @@ impl PollCoalescer {
         let cfg = self.cfg;
         let state = Arc::clone(&self.state);
         let wake = Arc::clone(&self.wake);
+        let mut shutdown = self.shutdown.clone();
         tokio::task::Builder::new()
             .name("poll_coalescer")
             .spawn(async move {
                 loop {
                     if state.lock().await.is_empty() {
-                        wake.notified().await;
+                        tokio::select! {
+                            _ = wake.notified() => {}
+                            _ = async { if *shutdown.borrow() { } else { let _ = shutdown.changed().await; } } => { break; }
+                        }
                     }
 
-                    tokio::time::sleep(Duration::from_millis(cfg.batch_window_ms)).await;
+                    // Wait the batch window or abort on shutdown
+                    tokio::select! {
+                        _ = tokio::time::sleep(Duration::from_millis(cfg.batch_window_ms)) => {}
+                        _ = async { if *shutdown.borrow() { } else { let _ = shutdown.changed().await; } } => { break; }
+                    }
 
                     let mut batch: Vec<PendingPoll> = Vec::new();
                     {
@@ -144,6 +154,12 @@ impl PollCoalescer {
                         };
                     }
                 }
+
+                // Shutdown: drain any pending polls with an empty result to unblock waiters
+                let mut guard = state.lock().await;
+                while let Some(p) = guard.pop_front() {
+                    let _ = p.tx.send(Ok(None));
+                }
             })
             .expect("spawn poll coalescer");
     }

src/server/mod.rs

@@ -29,7 +29,10 @@ impl Server {
         notify: Arc<Notify>,
         shutdown_tx: watch::Sender<bool>,
     ) -> Self {
-        let coalescer = Arc::new(PollCoalescer::new(Arc::clone(&storage)));
+        let coalescer = Arc::new(PollCoalescer::new(
+            Arc::clone(&storage),
+            shutdown_tx.subscribe(),
+        ));
         Self {
             storage,
             notify,
@@ -48,13 +51,18 @@ pub fn spawn_background_tasks(
     let bg_notify = Arc::clone(&notify);
     let lease_expiry_shutdown = shutdown_tx.clone();
     let visibility_wakeup_shutdown = shutdown_tx.clone();
+    let shutdown_rx_for_lease = shutdown_tx.subscribe();
+    let mut shutdown_rx_for_visibility = shutdown_tx.subscribe();
 
     // Background task to expire leases periodically
     tokio::task::Builder::new()
         .name("lease_expiry")
         .spawn(async move {
             let st = Arc::clone(&bg_storage);
             loop {
+                if *shutdown_rx_for_lease.borrow() {
+                    break;
+                }
                 match st.expire_due_leases().await {
                     Ok(n) => {
                         if n > 0 {
@@ -91,16 +99,25 @@ pub fn spawn_background_tasks(
                         if ts_secs <= now_secs {
                             vis_notify.notify_one();
                             // Avoid busy loop; small sleep before checking again.
-                            tokio::time::sleep(Duration::from_millis(50)).await;
+                            tokio::select! {
+                                _ = tokio::time::sleep(Duration::from_millis(50)) => {}
+                                _ = async { if *shutdown_rx_for_visibility.borrow() { } else { let _ = shutdown_rx_for_visibility.changed().await; } } => { break; }
+                            }
                         } else {
                             let sleep_dur = std::time::Duration::from_secs(ts_secs - now_secs);
-                            tokio::time::sleep(sleep_dur).await;
+                            tokio::select! {
+                                _ = tokio::time::sleep(sleep_dur) => {}
+                                _ = async { if *shutdown_rx_for_visibility.borrow() { } else { let _ = shutdown_rx_for_visibility.changed().await; } } => { break; }
+                            }
                             vis_notify.notify_one();
                         }
                     }
                     Ok(None) => {
                         // No items; back off
-                        tokio::time::sleep(Duration::from_millis(200)).await;
+                        tokio::select! {
+                            _ = tokio::time::sleep(Duration::from_millis(200)) => {}
+                            _ = async { if *shutdown_rx_for_visibility.borrow() { } else { let _ = shutdown_rx_for_visibility.changed().await; } } => { break; }
+                        }
                     }
                     Err(e) => {
                         tracing::error!("peek_next_visibility_ts_secs: {}", e);
@@ -196,6 +213,7 @@ impl crate::protocol::queue::Server for Server {
         let _storage = Arc::clone(&self.storage);
         let notify = Arc::clone(&self.notify);
         let coalescer = Arc::clone(&self.coalescer);
+        let mut shutdown_rx = self.shutdown_tx.subscribe();
 
         Promise::from_future(async move {
             let req = params.get()?.get_req()?;
@@ -237,10 +255,14 @@ impl crate::protocol::queue::Server for Server {
                         tokio::select! {
                             _ = notify.notified() => {},
                             _ = tokio::time::sleep(remaining) => { return Ok(()); },
+                            _ = async { if *shutdown_rx.borrow() { } else { let _ = shutdown_rx.changed().await; } } => { return Ok(()); },
                         }
                     }
                     None => {
-                        notify.notified().await;
+                        tokio::select! {
+                            _ = notify.notified() => {},
+                            _ = async { if *shutdown_rx.borrow() { } else { let _ = shutdown_rx.changed().await; } } => { return Ok(()); },
+                        }
                     }
                 }
             }

tests/server_graceful_shutdown.rs

@@ -0,0 +1,127 @@
+use std::net::TcpListener;
+use std::time::Duration;
+
+use capnp_rpc::{RpcSystem, rpc_twoparty_capnp, twoparty};
+use futures::AsyncReadExt;
+use tokio::sync::watch;
+
+#[tokio::test(flavor = "current_thread")]
+async fn server_drains_inflight_requests_on_shutdown() {
+    // Bind ephemeral port
+    let listener = TcpListener::bind("127.0.0.1:0").expect("bind");
+    let addr = listener.local_addr().unwrap();
+    drop(listener);
+
+    // Start minimal server like in other tests and signal readiness
+    let data_dir = tempfile::TempDir::new().unwrap();
+    let data_path = data_dir.path().to_path_buf();
+    let (shutdown_tx, _rx_unused) = watch::channel(false);
+    let mut shutdown_rx_for_thread = shutdown_tx.subscribe();
+    let shutdown_tx_for_later = shutdown_tx.clone();
+    let (ready_tx, ready_rx) = std::sync::mpsc::sync_channel::<()>(1);
+
+    let server_thread = std::thread::spawn(move || {
+        let rt = tokio::runtime::Builder::new_current_thread()
+            .enable_io()
+            .enable_time()
+            .build()
+            .unwrap();
+        rt.block_on(async move {
+            use queueber::protocol::queue::Client as QueueClient;
+            use queueber::server::Server;
+            use queueber::storage::{RetriedStorage, Storage};
+            use std::sync::Arc;
+            use tokio::sync::Notify;
+            let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
+            let storage = Arc::new(RetriedStorage::new(Storage::new(&data_path).unwrap()));
+            let notify = Arc::new(Notify::new());
+            queueber::server::spawn_background_tasks(
+                Arc::clone(&storage),
+                Arc::clone(&notify),
+                shutdown_tx.clone(),
+            );
+            let server = Server::new(storage, notify, shutdown_tx.clone());
+            let queue_client: QueueClient = capnp_rpc::new_client(server);
+
+            tokio::task::LocalSet::new()
+                .run_until(async move {
+                    let _ = ready_tx.send(());
+                    loop {
+                        tokio::select! {
+                            _ = async { if *shutdown_rx_for_thread.borrow() { } else { let _ = shutdown_rx_for_thread.changed().await; } } => { break; }
+                            accept_res = listener.accept() => {
+                                let (stream, _) = accept_res.unwrap();
+                                stream.set_nodelay(true).unwrap();
+                                let (reader, writer) = tokio_util::compat::TokioAsyncReadCompatExt::compat(stream).split();
+                                let network = twoparty::VatNetwork::new(
+                                    futures::io::BufReader::new(reader),
+                                    futures::io::BufWriter::new(writer),
+                                    rpc_twoparty_capnp::Side::Server,
+                                    Default::default(),
+                                );
+                                let rpc_system = RpcSystem::new(Box::new(network), Some(queue_client.clone().client));
+                                let _jh = tokio::task::Builder::new().name("rpc_system").spawn_local(rpc_system).unwrap();
+                            }
+                        }
+                    }
+                })
+                .await;
+        });
+    });
+
+    // Wait until server is ready to accept
+    let _ = ready_rx.recv();
+
+    // Create a client and start a long poll that should be in-flight during shutdown
+    let stream = tokio::net::TcpStream::connect(addr).await.unwrap();
+    stream.set_nodelay(true).unwrap();
+    let (reader, writer) = tokio_util::compat::TokioAsyncReadCompatExt::compat(stream).split();
+    let rpc_network = Box::new(twoparty::VatNetwork::new(
+        futures::io::BufReader::new(reader),
+        futures::io::BufWriter::new(writer),
+        rpc_twoparty_capnp::Side::Client,
+        Default::default(),
+    ));
+    let mut rpc_system = RpcSystem::new(rpc_network, None);
+    let queue_client: queueber::protocol::queue::Client =
+        rpc_system.bootstrap(rpc_twoparty_capnp::Side::Server);
+
+    let local = tokio::task::LocalSet::new();
+    local
+        .run_until(async move {
+            let (done_tx, done_rx) = tokio::sync::oneshot::channel::<()>();
+            tokio::task::spawn_local(rpc_system);
+            tokio::task::spawn_local(async move {
+                let mut poll = queue_client.poll_request();
+                {
+                    let mut req = poll.get().init_req();
+                    req.set_lease_validity_secs(30);
+                    req.set_num_items(1);
+                    req.set_timeout_secs(2); // longish request
+                }
+                let _ = poll.send().promise.await; // ignore result
+                let _ = done_tx.send(());
+            });
+
+            // Let the request run for a short time then initiate shutdown and ensure it drains
+            tokio::time::sleep(Duration::from_millis(200)).await;
+            let _ = shutdown_tx_for_later.send(true);
+
+            tokio::time::timeout(Duration::from_secs(5), async move {
+                let _ = done_rx.await;
+            })
+            .await
+            .expect("poll finished");
+        })
+        .await;
+
+    // Join the server thread within a reasonable timeout
+    let (tx_join, rx_join) = std::sync::mpsc::channel();
+    std::thread::spawn(move || {
+        let _ = server_thread.join();
+        let _ = tx_join.send(());
+    });
+    rx_join
+        .recv_timeout(Duration::from_secs(5))
+        .expect("server exited");
+}