Skip to main content

max / makenotwork

server: drain the background pool on graceful shutdown The fire-and-forget bg pool was the one undrained primitive: main only awaited the scheduler, so in-flight state.bg.spawn work (transactional/announcement emails, MT thread creation, Cloudflare cache purges of quarantined objects, platform-credit announcements) was abandoned at the 10s hard exit on every restart. spawn_pool now takes the shutdown watch receiver and returns its JoinHandle; on signal it stops taking new work, runs everything already queued, then waits for in-flight tasks (reacquiring all permits) before exiting. main awaits that drain (bounded 5s, like the scheduler) before the hard-exit deadline. Also count dropped-on-full tasks via a background_tasks_dropped_total metric so a saturated queue is visible, not just logged. Adds drain regression tests. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Author: Max Johnson <me@maxj.phd> · 2026-07-02 16:56 UTC
Commit: b208bdfa3b928b5f0995587e7809ae581121b054
Parent: 3bcd9e3
2 files changed, +106 insertions, -18 deletions
M server/src/background.rs +90 -16
@@ -22,7 +22,8 @@
22 22 use std::future::Future;
23 23 use std::pin::Pin;
24 24 use std::sync::Arc;
25 - use tokio::sync::{mpsc, Semaphore};
25 + use tokio::sync::{mpsc, watch, Semaphore};
26 + use tokio::task::JoinHandle;
26 27
27 28 /// Max queued tasks before `spawn` starts dropping with a warning.
28 29 const CAPACITY: usize = 1024;
@@ -47,29 +48,102 @@ impl BackgroundTx {
47 48 F: Future<Output = ()> + Send + 'static,
48 49 {
49 50 if self.tx.try_send((name, Box::pin(fut))).is_err() {
51 + // A dropped task is silent data loss (a missed email / cache purge /
52 + // MT thread), so surface it as a metric, not just a log line, so a
53 + // saturated queue is visible on the dashboard (audit Run 17 Resilience).
54 + metrics::counter!("background_tasks_dropped_total", "task" => name).increment(1);
50 55 tracing::warn!(task = name, "background queue full, task dropped");
51 56 }
52 57 }
53 58 }
54 59
55 - /// Spawn the drainer + return the sender to install on `AppState`. The
56 - /// drainer pulls tasks off the channel and runs each under a semaphore
57 - /// permit so concurrent execution is bounded.
58 - pub fn spawn_pool() -> BackgroundTx {
60 + /// Spawn the drainer + return the sender to install on `AppState`, plus the
61 + /// drainer's `JoinHandle` so shutdown can await an orderly drain. The drainer
62 + /// pulls tasks off the channel and runs each under a semaphore permit so
63 + /// concurrent execution is bounded.
64 + ///
65 + /// On shutdown (the `shutdown` watch fires — a value change or all senders
66 + /// dropped) the drainer stops taking new work, runs every already-queued task,
67 + /// then waits for in-flight tasks to finish before exiting. Without this the
68 + /// `bg` pool was the one undrained primitive: in-flight emails / cache purges /
69 + /// MT thread creation were abandoned at the 10s hard exit on every restart
70 + /// (audit Run 17 Resilience). The wait is best-effort and bounded by the
71 + /// caller's shutdown timeout.
72 + pub fn spawn_pool(mut shutdown: watch::Receiver<()>) -> (BackgroundTx, JoinHandle<()>) {
59 73 let (tx, mut rx) = mpsc::channel::<(&'static str, BoxFuture)>(CAPACITY);
60 74 let sem = Arc::new(Semaphore::new(PARALLELISM));
61 - tokio::spawn(async move {
62 - while let Some((_, task)) = rx.recv().await {
63 - let permit = match sem.clone().acquire_owned().await {
64 - Ok(p) => p,
65 - Err(_) => break, // semaphore closed → shutting down
66 - };
67 - tokio::spawn(async move {
68 - task.await;
69 - drop(permit);
70 - });
75 + let handle = tokio::spawn(async move {
76 + loop {
77 + tokio::select! {
78 + maybe = rx.recv() => match maybe {
79 + Some((_, task)) => run_task(&sem, task).await,
80 + None => break, // all senders dropped → nothing left to drain
81 + },
82 + // `changed()` resolves with Err once every sender is dropped,
83 + // which is exactly how main.rs signals shutdown (drop(shutdown_tx)).
84 + _ = shutdown.changed() => break,
85 + }
71 86 }
87 +
88 + // Drain everything already enqueued, then wait for in-flight tasks by
89 + // reacquiring all permits (each running task holds one; once all
90 + // PARALLELISM are free, the pool is idle).
91 + while let Ok((_, task)) = rx.try_recv() {
92 + run_task(&sem, task).await;
93 + }
94 + let _ = sem.acquire_many(PARALLELISM as u32).await;
72 95 tracing::info!("background-task drainer exited");
73 96 });
74 - BackgroundTx { tx }
97 + (BackgroundTx { tx }, handle)
98 + }
99 +
100 + /// Acquire a permit and spawn `task`, releasing the permit when it finishes.
101 + async fn run_task(sem: &Arc<Semaphore>, task: BoxFuture) {
102 + let Ok(permit) = sem.clone().acquire_owned().await else {
103 + return; // semaphore closed → shutting down
104 + };
105 + tokio::spawn(async move {
106 + task.await;
107 + drop(permit);
108 + });
109 + }
110 +
111 + #[cfg(test)]
112 + mod tests {
113 + use super::*;
114 + use std::sync::atomic::{AtomicUsize, Ordering};
115 +
116 + #[tokio::test]
117 + async fn drains_queued_tasks_on_shutdown() {
118 + let (shutdown_tx, shutdown_rx) = watch::channel(());
119 + let (bg, handle) = spawn_pool(shutdown_rx);
120 +
121 + let counter = Arc::new(AtomicUsize::new(0));
122 + for _ in 0..50 {
123 + let c = counter.clone();
124 + bg.spawn("test", async move {
125 + c.fetch_add(1, Ordering::SeqCst);
126 + });
127 + }
128 +
129 + // Signal shutdown (drop the only sender) and await the orderly drain.
130 + drop(shutdown_tx);
131 + handle.await.expect("drainer joins cleanly");
132 +
133 + assert_eq!(
134 + counter.load(Ordering::SeqCst),
135 + 50,
136 + "every queued task must run before the drainer exits"
137 + );
138 + }
139 +
140 + #[tokio::test]
141 + async fn drainer_exits_when_all_senders_drop() {
142 + let (_shutdown_tx, shutdown_rx) = watch::channel(());
143 + let (bg, handle) = spawn_pool(shutdown_rx);
144 + // Dropping every sender closes the channel; the drainer must exit on its
145 + // own (recv() -> None) without needing the shutdown signal.
146 + drop(bg);
147 + handle.await.expect("drainer exits when the queue closes");
148 + }
75 149 }
M server/src/main.rs +16 -2
@@ -344,8 +344,11 @@ async fn main() {
344 344
345 345 let started_at = chrono::Utc::now();
346 346 let start_instant = Instant::now();
347 + // Shutdown broadcast: dropped at graceful-shutdown time to signal the
348 + // background pool, monitor, scheduler, and scan workers to drain and stop.
349 + let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(());
347 350 let page_view_tx = makenotwork::db::page_views::spawn_batcher(db.clone());
348 - let bg = makenotwork::background::spawn_pool();
351 + let (bg, bg_handle) = makenotwork::background::spawn_pool(shutdown_tx.subscribe());
349 352 let state = AppState {
350 353 db,
351 354 config: config.clone(),
@@ -393,7 +396,6 @@ async fn main() {
393 396 );
394 397
395 398 // Start background health monitor and scheduler
396 - let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(());
397 399 let _monitor_handle = makenotwork::monitor::spawn_monitor(state.clone(), shutdown_rx);
398 400 let scheduler_shutdown_rx = shutdown_tx.subscribe();
399 401 let scheduler_handle = makenotwork::scheduler::spawn_scheduler(state.clone(), scheduler_shutdown_rx);
@@ -498,6 +500,18 @@ async fn main() {
498 500 ),
499 501 }
500 502
503 + // Drain the fire-and-forget background pool: it runs the queued emails /
504 + // cache purges / MT thread creations before the armed hard-exit deadline,
505 + // instead of abandoning them mid-flight (audit Run 17 Resilience). Bounded
506 + // like the scheduler; if it overruns we proceed to the hard exit.
507 + match tokio::time::timeout(Duration::from_secs(5), bg_handle).await {
508 + Ok(Ok(())) => tracing::info!("background pool drained"),
509 + Ok(Err(e)) => tracing::warn!(error = ?e, "background pool task ended abnormally on shutdown"),
510 + Err(_) => {
511 + tracing::warn!("background pool did not drain within 5s of shutdown; proceeding")
512 + }
513 + }
514 +
501 515 tracing::info!("Server shut down gracefully");
502 516 }
503 517