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..88cb908 100644 --- a/src/bin/queueber/main.rs +++ b/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> = 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() { 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"); +}