Skip to main content

max / makenotwork

6.5 KB · 162 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. `db::page_views::PageViewTx` is the same shape for page
6 //! views; the webhook hot path
7 //! (`routes/stripe/webhook/checkout_helpers.rs`) is the other caller.
8 //!
9 //! Two bounds applied:
10 //! - **Queue capacity**: `try_send` returns Err once `CAPACITY` tasks are
11 //! queued; the task is dropped and a warning is logged. Caller is never
12 //! blocked and never panics.
13 //! - **Concurrent execution**: a semaphore caps the number of background
14 //! tasks running at once, so a burst of email sends can't exhaust the
15 //! 25-slot DB pool out from under request handlers.
16 //!
17 //! Use `state.bg.spawn(name, fut)` instead of `tokio::spawn(fut)` for any
18 //! work that (a) doesn't need its result and (b) acquires DB pool conns or
19 //! makes outbound network calls.
20
21 use std::future::Future;
22 use std::pin::Pin;
23 use std::sync::Arc;
24 use tokio::sync::{Semaphore, mpsc, watch};
25 use tokio::task::JoinHandle;
26
27 /// Max queued tasks before `spawn` starts dropping with a warning.
28 const CAPACITY: usize = 1024;
29
30 /// Max background tasks running concurrently. Set well below
31 /// `DB_POOL_MAX_CONNECTIONS` (25) so background work can't starve requests.
32 const PARALLELISM: usize = 8;
33
34 type BoxFuture = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
35
36 #[derive(Clone)]
37 pub struct BackgroundTx {
38 tx: mpsc::Sender<(&'static str, BoxFuture)>,
39 }
40
41 impl BackgroundTx {
42 /// Submit a fire-and-forget task. On queue overflow the task is dropped
43 /// and a warning is logged with `name` as the task identifier. Never
44 /// blocks; never panics.
45 pub fn spawn<F>(&self, name: &'static str, fut: F)
46 where
47 F: Future<Output = ()> + Send + 'static,
48 {
49 if self.tx.try_send((name, Box::pin(fut))).is_err() {
50 // A dropped task is silent data loss (a missed email / cache purge /
51 // MT thread), so surface it as a metric, not just a log line, so a
52 // saturated queue is visible on the dashboard (audit Run 17 Resilience).
53 metrics::counter!("background_tasks_dropped_total", "task" => name).increment(1);
54 tracing::warn!(task = name, "background queue full, task dropped");
55 }
56 }
57 }
58
59 /// Spawn the drainer + return the sender to install on `AppState`, plus the
60 /// drainer's `JoinHandle` so shutdown can await an orderly drain. The drainer
61 /// pulls tasks off the channel and runs each under a semaphore permit so
62 /// concurrent execution is bounded.
63 ///
64 /// On shutdown (the `shutdown` watch fires, a value change or all senders
65 /// dropped) the drainer stops taking new work, runs every already-queued task,
66 /// then waits for in-flight tasks to finish before exiting. Without this the
67 /// `bg` pool was the one undrained primitive: in-flight emails / cache purges /
68 /// MT thread creation were abandoned at the 10s hard exit on every restart.
69 /// The wait is best-effort and bounded by the
70 /// caller's shutdown timeout.
71 pub fn spawn_pool(mut shutdown: watch::Receiver<()>) -> (BackgroundTx, JoinHandle<()>) {
72 let (tx, mut rx) = mpsc::channel::<(&'static str, BoxFuture)>(CAPACITY);
73 let sem = Arc::new(Semaphore::new(PARALLELISM));
74 let handle = tokio::spawn(async move {
75 loop {
76 tokio::select! {
77 maybe = rx.recv() => match maybe {
78 Some((_, task)) => run_task(&sem, task).await,
79 None => break, // all senders dropped → nothing left to drain
80 },
81 // `changed()` resolves with Err once every sender is dropped,
82 // which is exactly how main.rs signals shutdown (drop(shutdown_tx)).
83 _ = shutdown.changed() => break,
84 }
85 }
86
87 // Drain everything already enqueued, then wait for in-flight tasks by
88 // reacquiring all permits (each running task holds one; once all
89 // PARALLELISM are free, the pool is idle).
90 while let Ok((_, task)) = rx.try_recv() {
91 run_task(&sem, task).await;
92 }
93 let _ = sem.acquire_many(PARALLELISM as u32).await;
94 tracing::info!("background-task drainer exited");
95 });
96 (BackgroundTx { tx }, handle)
97 }
98
99 /// Spawn a background pool with no external shutdown wiring, for test harnesses
100 /// and embeddings that manage lifecycle purely via `AppState` drop. The shutdown
101 /// sender is retained internally (its `changed()` never fires), so the pool
102 /// drains only when every `BackgroundTx` clone is dropped and the channel closes.
103 pub fn spawn_pool_detached() -> BackgroundTx {
104 let (tx, rx) = watch::channel(());
105 // Deliberately keep the sender alive so the drainer's shutdown arm never
106 // fires; the pool then behaves exactly as it did before shutdown draining
107 // existed (runs until the queue closes). Only reachable off the request path.
108 std::mem::forget(tx);
109 spawn_pool(rx).0
110 }
111
112 /// Acquire a permit and spawn `task`, releasing the permit when it finishes.
113 async fn run_task(sem: &Arc<Semaphore>, task: BoxFuture) {
114 let Ok(permit) = sem.clone().acquire_owned().await else {
115 return; // semaphore closed → shutting down
116 };
117 tokio::spawn(async move {
118 task.await;
119 drop(permit);
120 });
121 }
122
123 #[cfg(test)]
124 mod tests {
125 use super::*;
126 use std::sync::atomic::{AtomicUsize, Ordering};
127
128 #[tokio::test]
129 async fn drains_queued_tasks_on_shutdown() {
130 let (shutdown_tx, shutdown_rx) = watch::channel(());
131 let (bg, handle) = spawn_pool(shutdown_rx);
132
133 let counter = Arc::new(AtomicUsize::new(0));
134 for _ in 0..50 {
135 let c = counter.clone();
136 bg.spawn("test", async move {
137 c.fetch_add(1, Ordering::SeqCst);
138 });
139 }
140
141 // Signal shutdown (drop the only sender) and await the orderly drain.
142 drop(shutdown_tx);
143 handle.await.expect("drainer joins cleanly");
144
145 assert_eq!(
146 counter.load(Ordering::SeqCst),
147 50,
148 "every queued task must run before the drainer exits"
149 );
150 }
151
152 #[tokio::test]
153 async fn drainer_exits_when_all_senders_drop() {
154 let (_shutdown_tx, shutdown_rx) = watch::channel(());
155 let (bg, handle) = spawn_pool(shutdown_rx);
156 // Dropping every sender closes the channel; the drainer must exit on its
157 // own (recv() -> None) without needing the shutdown signal.
158 drop(bg);
159 handle.await.expect("drainer exits when the queue closes");
160 }
161 }
162