From 6309acc57f351433cf52a5c4fdc05cc342fae0b3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 28 Aug 2025 13:26:36 +0000 Subject: [PATCH 1/4] Implement graceful shutdown with signal handling and task draining Co-authored-by: miles.frankel --- TODO.md | 2 +- src/bin/queueber/main.rs | 33 +++++++- src/server/coalescer.rs | 24 +++++- src/server/mod.rs | 32 ++++++-- tests/server_graceful_shutdown.rs | 127 ++++++++++++++++++++++++++++++ 5 files changed, 206 insertions(+), 12 deletions(-) create mode 100644 tests/server_graceful_shutdown.rs diff --git a/TODO.md b/TODO.md index e20ba65..45dfa37 100644 --- a/TODO.md +++ b/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 diff --git a/src/bin/queueber/main.rs b/src/bin/queueber/main.rs index 5fe624b..2f3fae7 100644 --- a/src/bin/queueber/main.rs +++ b/src/bin/queueber/main.rs @@ -13,6 +13,8 @@ use queueber::{ }; use socket2::{Domain, Protocol, Socket, Type}; use std::sync::Arc; +#[cfg(unix)] +use tokio::signal::unix::{SignalKind, signal}; use tokio::sync::Notify; use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; @@ -116,8 +118,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> = Vec::new(); loop { tokio::select! { _ = async { if *shutdown_rx_worker.borrow() { } else { let _ = shutdown_rx_worker.changed().await; } } => { @@ -129,7 +131,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 +154,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 +172,27 @@ async fn main() -> Result<()> { worker_handles.push(handle); } + // Signal handling tasks: SIGINT and SIGTERM (unix) + { + let shutdown_tx_for_signals = shutdown_tx.clone(); + tokio::spawn(async move { + let _ = tokio::signal::ctrl_c().await; + tracing::info!("received Ctrl-C; initiating shutdown"); + let _ = shutdown_tx_for_signals.send(true); + }); + } + #[cfg(unix)] + { + let shutdown_tx_for_signals = shutdown_tx.clone(); + tokio::spawn(async move { + if let Ok(mut sigterm) = signal(SignalKind::terminate()) { + let _ = sigterm.recv().await; + tracing::info!("received SIGTERM; initiating shutdown"); + let _ = shutdown_tx_for_signals.send(true); + } + }); + } + // Wait for shutdown signal, then join workers let _ = async { if *shutdown_rx.borrow() { diff --git a/src/server/coalescer.rs b/src/server/coalescer.rs index 772719f..f9da89c 100644 --- a/src/server/coalescer.rs +++ b/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>>, // pending batch wake: Arc, + shutdown: watch::Receiver, } impl PollCoalescer { - pub fn new(storage: Arc>) -> Self { + pub fn new(storage: Arc>, shutdown: watch::Receiver) -> 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 = 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"); } diff --git a/src/server/mod.rs b/src/server/mod.rs index 063f741..0de2c40 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -29,7 +29,10 @@ impl Server { notify: Arc, shutdown_tx: watch::Sender, ) -> 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,6 +51,8 @@ pub fn spawn_background_tasks( let bg_notify = Arc::clone(¬ify); 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() @@ -55,6 +60,9 @@ pub fn spawn_background_tasks( .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(()); }, + } } } } diff --git a/tests/server_graceful_shutdown.rs b/tests/server_graceful_shutdown.rs new file mode 100644 index 0000000..0f009e1 --- /dev/null +++ b/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(¬ify), + 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"); +} From c03c8ffb56ce1fbc0dbdb77c1d1aa05041222f08 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 28 Aug 2025 15:55:20 +0000 Subject: [PATCH 2/4] Improve signal handling with graceful shutdown and force exit Co-authored-by: miles.frankel --- src/bin/queueber/main.rs | 58 ++++++++++++++++++++++++++++++---------- 1 file changed, 44 insertions(+), 14 deletions(-) diff --git a/src/bin/queueber/main.rs b/src/bin/queueber/main.rs index 2f3fae7..9a71812 100644 --- a/src/bin/queueber/main.rs +++ b/src/bin/queueber/main.rs @@ -172,23 +172,53 @@ async fn main() -> Result<()> { worker_handles.push(handle); } - // Signal handling tasks: SIGINT and SIGTERM (unix) + // Signals: first SIGINT/SIGTERM triggers graceful shutdown; second exits immediately { let shutdown_tx_for_signals = shutdown_tx.clone(); tokio::spawn(async move { - let _ = tokio::signal::ctrl_c().await; - tracing::info!("received Ctrl-C; initiating shutdown"); - let _ = shutdown_tx_for_signals.send(true); - }); - } - #[cfg(unix)] - { - let shutdown_tx_for_signals = shutdown_tx.clone(); - tokio::spawn(async move { - if let Ok(mut sigterm) = signal(SignalKind::terminate()) { - let _ = sigterm.recv().await; - tracing::info!("received SIGTERM; initiating shutdown"); - let _ = shutdown_tx_for_signals.send(true); + let mut first_seen = false; + #[cfg(unix)] + let mut sigterm_opt = signal(SignalKind::terminate()).ok(); + loop { + #[cfg(unix)] + { + tokio::select! { + _ = tokio::signal::ctrl_c() => { + if !first_seen { + first_seen = true; + tracing::info!("received Ctrl-C; initiating graceful shutdown"); + let _ = shutdown_tx_for_signals.send(true); + } else { + tracing::warn!("received second Ctrl-C; exiting now"); + std::process::exit(130); + } + } + _ = async { + if let Some(ref mut s) = sigterm_opt { let _ = s.recv().await; } else { futures::future::pending::>().await; } + } => { + if !first_seen { + first_seen = true; + tracing::info!("received SIGTERM; initiating graceful shutdown"); + let _ = shutdown_tx_for_signals.send(true); + } else { + tracing::warn!("received second SIGTERM; exiting now"); + std::process::exit(143); + } + } + } + } + #[cfg(not(unix))] + { + let _ = tokio::signal::ctrl_c().await; + if !first_seen { + first_seen = true; + tracing::info!("received Ctrl-C; initiating graceful shutdown"); + let _ = shutdown_tx_for_signals.send(true); + } else { + tracing::warn!("received second Ctrl-C; exiting now"); + std::process::exit(130); + } + } } }); } From 6d75a1ce9033f9f88eb32cb2932a3c1c9438ce39 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 28 Aug 2025 16:07:16 +0000 Subject: [PATCH 3/4] Simplify signal handling for graceful shutdown on Unix and non-Unix systems Co-authored-by: miles.frankel --- src/bin/queueber/main.rs | 58 ++++++++++------------------------------ 1 file changed, 14 insertions(+), 44 deletions(-) diff --git a/src/bin/queueber/main.rs b/src/bin/queueber/main.rs index 9a71812..420bf82 100644 --- a/src/bin/queueber/main.rs +++ b/src/bin/queueber/main.rs @@ -172,53 +172,23 @@ async fn main() -> Result<()> { worker_handles.push(handle); } - // Signals: first SIGINT/SIGTERM triggers graceful shutdown; second exits immediately + // Signals: single-shot handlers for SIGINT and SIGTERM (unix) { let shutdown_tx_for_signals = shutdown_tx.clone(); tokio::spawn(async move { - let mut first_seen = false; - #[cfg(unix)] - let mut sigterm_opt = signal(SignalKind::terminate()).ok(); - loop { - #[cfg(unix)] - { - tokio::select! { - _ = tokio::signal::ctrl_c() => { - if !first_seen { - first_seen = true; - tracing::info!("received Ctrl-C; initiating graceful shutdown"); - let _ = shutdown_tx_for_signals.send(true); - } else { - tracing::warn!("received second Ctrl-C; exiting now"); - std::process::exit(130); - } - } - _ = async { - if let Some(ref mut s) = sigterm_opt { let _ = s.recv().await; } else { futures::future::pending::>().await; } - } => { - if !first_seen { - first_seen = true; - tracing::info!("received SIGTERM; initiating graceful shutdown"); - let _ = shutdown_tx_for_signals.send(true); - } else { - tracing::warn!("received second SIGTERM; exiting now"); - std::process::exit(143); - } - } - } - } - #[cfg(not(unix))] - { - let _ = tokio::signal::ctrl_c().await; - if !first_seen { - first_seen = true; - tracing::info!("received Ctrl-C; initiating graceful shutdown"); - let _ = shutdown_tx_for_signals.send(true); - } else { - tracing::warn!("received second Ctrl-C; exiting now"); - std::process::exit(130); - } - } + let _ = tokio::signal::ctrl_c().await; + tracing::info!("received Ctrl-C; initiating graceful shutdown"); + let _ = shutdown_tx_for_signals.send(true); + }); + } + #[cfg(unix)] + { + let shutdown_tx_for_signals = shutdown_tx.clone(); + tokio::spawn(async move { + if let Ok(mut sigterm) = signal(SignalKind::terminate()) { + let _ = sigterm.recv().await; + tracing::info!("received SIGTERM; initiating graceful shutdown"); + let _ = shutdown_tx_for_signals.send(true); } }); } From 33e3906cde0b6409453a1b17d711171fca9b5084 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 28 Aug 2025 16:59:51 +0000 Subject: [PATCH 4/4] Simplify signal handling with tokio::select! for SIGINT and SIGTERM Co-authored-by: miles.frankel --- src/bin/queueber/main.rs | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/src/bin/queueber/main.rs b/src/bin/queueber/main.rs index 420bf82..88cb908 100644 --- a/src/bin/queueber/main.rs +++ b/src/bin/queueber/main.rs @@ -13,7 +13,6 @@ use queueber::{ }; use socket2::{Domain, Protocol, Socket, Type}; use std::sync::Arc; -#[cfg(unix)] use tokio::signal::unix::{SignalKind, signal}; use tokio::sync::Notify; use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; @@ -172,24 +171,21 @@ async fn main() -> Result<()> { worker_handles.push(handle); } - // Signals: single-shot handlers for SIGINT and SIGTERM (unix) + // 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 _ = tokio::signal::ctrl_c().await; - tracing::info!("received Ctrl-C; initiating graceful shutdown"); - let _ = shutdown_tx_for_signals.send(true); - }); - } - #[cfg(unix)] - { - let shutdown_tx_for_signals = shutdown_tx.clone(); - tokio::spawn(async move { - if let Ok(mut sigterm) = signal(SignalKind::terminate()) { - let _ = sigterm.recv().await; - tracing::info!("received SIGTERM; initiating graceful shutdown"); - let _ = shutdown_tx_for_signals.send(true); + 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); }); }