Skip to main content

max / makenotwork

6.6 KB · 163 lines History Blame Raw
1 //! Bounded background-task queue for fire-and-forget work.
2 //!
3 //! Replaces per-request `tokio::spawn(...)` for low-priority work that
4 //! competes with request handlers for the DB pool, email sends, mailing-list
5 //! subscriptions, etc. Run #4 fixed the same shape for page views via
6 //! `db::page_views::PageViewTx`; Run #8 surfaced it again on the webhook
7 //! hot path (`routes/stripe/webhook/checkout_helpers.rs`), so this module
8 //! generalizes the pattern.
9 //!
10 //! Two bounds applied:
11 //! - **Queue capacity**: `try_send` returns Err once `CAPACITY` tasks are
12 //! queued; the task is dropped and a warning is logged. Caller is never
13 //! blocked and never panics.
14 //! - **Concurrent execution**: a semaphore caps the number of background
15 //! tasks running at once, so a burst of email sends can't exhaust the
16 //! 25-slot DB pool out from under request handlers.
17 //!
18 //! Use `state.bg.spawn(name, fut)` instead of `tokio::spawn(fut)` for any
19 //! work that (a) doesn't need its result and (b) acquires DB pool conns or
20 //! makes outbound network calls.
21
22 use std::future::Future;
23 use std::pin::Pin;
24 use std::sync::Arc;
25 use tokio::sync::{Semaphore, mpsc, watch};
26 use tokio::task::JoinHandle;
27
28 /// Max queued tasks before `spawn` starts dropping with a warning.
29 const CAPACITY: usize = 1024;
30
31 /// Max background tasks running concurrently. Set well below
32 /// `DB_POOL_MAX_CONNECTIONS` (25) so background work can't starve requests.
33 const PARALLELISM: usize = 8;
34
35 type BoxFuture = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
36
37 #[derive(Clone)]
38 pub struct BackgroundTx {
39 tx: mpsc::Sender<(&'static str, BoxFuture)>,
40 }
41
42 impl BackgroundTx {
43 /// Submit a fire-and-forget task. On queue overflow the task is dropped
44 /// and a warning is logged with `name` as the task identifier. Never
45 /// blocks; never panics.
46 pub fn spawn<F>(&self, name: &'static str, fut: F)
47 where
48 F: Future<Output = ()> + Send + 'static,
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);
55 tracing::warn!(task = name, "background queue full, task dropped");
56 }
57 }
58 }
59
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<()>) {
73 let (tx, mut rx) = mpsc::channel::<(&'static str, BoxFuture)>(CAPACITY);
74 let sem = Arc::new(Semaphore::new(PARALLELISM));
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 }
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;
95 tracing::info!("background-task drainer exited");
96 });
97 (BackgroundTx { tx }, handle)
98 }
99
100 /// Spawn a background pool with no external shutdown wiring, for test harnesses
101 /// and embeddings that manage lifecycle purely via `AppState` drop. The shutdown
102 /// sender is retained internally (its `changed()` never fires), so the pool
103 /// drains only when every `BackgroundTx` clone is dropped and the channel closes.
104 pub fn spawn_pool_detached() -> BackgroundTx {
105 let (tx, rx) = watch::channel(());
106 // Deliberately keep the sender alive so the drainer's shutdown arm never
107 // fires; the pool then behaves exactly as it did before shutdown draining
108 // existed (runs until the queue closes). Only reachable off the request path.
109 std::mem::forget(tx);
110 spawn_pool(rx).0
111 }
112
113 /// Acquire a permit and spawn `task`, releasing the permit when it finishes.
114 async fn run_task(sem: &Arc<Semaphore>, task: BoxFuture) {
115 let Ok(permit) = sem.clone().acquire_owned().await else {
116 return; // semaphore closed → shutting down
117 };
118 tokio::spawn(async move {
119 task.await;
120 drop(permit);
121 });
122 }
123
124 #[cfg(test)]
125 mod tests {
126 use super::*;
127 use std::sync::atomic::{AtomicUsize, Ordering};
128
129 #[tokio::test]
130 async fn drains_queued_tasks_on_shutdown() {
131 let (shutdown_tx, shutdown_rx) = watch::channel(());
132 let (bg, handle) = spawn_pool(shutdown_rx);
133
134 let counter = Arc::new(AtomicUsize::new(0));
135 for _ in 0..50 {
136 let c = counter.clone();
137 bg.spawn("test", async move {
138 c.fetch_add(1, Ordering::SeqCst);
139 });
140 }
141
142 // Signal shutdown (drop the only sender) and await the orderly drain.
143 drop(shutdown_tx);
144 handle.await.expect("drainer joins cleanly");
145
146 assert_eq!(
147 counter.load(Ordering::SeqCst),
148 50,
149 "every queued task must run before the drainer exits"
150 );
151 }
152
153 #[tokio::test]
154 async fn drainer_exits_when_all_senders_drop() {
155 let (_shutdown_tx, shutdown_rx) = watch::channel(());
156 let (bg, handle) = spawn_pool(shutdown_rx);
157 // Dropping every sender closes the channel; the drainer must exit on its
158 // own (recv() -> None) without needing the shutdown signal.
159 drop(bg);
160 handle.await.expect("drainer exits when the queue closes");
161 }
162 }
163