systemprompt_traits/
background_tasks.rs

1//! Owned background work: one tracker per process, drained at shutdown.
2//!
3//! [`BackgroundTasks`] is the single owner for work that outlives the call
4//! that started it — audit completion after a response, analytics writes,
5//! periodic loops. It is built once at the composition root, cloned into
6//! whatever needs to spawn, and drained by the server's graceful shutdown, so
7//! no task is detached: shutdown and tests observe the moment nothing is left
8//! in flight.
9//!
10//! Two kinds of work run on it:
11//!
12//! - [`spawn`](BackgroundTasks::spawn) — one-shot work that runs to completion
13//!   (a write, a settlement). Cancellation does not interrupt it; a drain waits
14//!   for it within its timeout.
15//! - [`spawn_cancellable`](BackgroundTasks::spawn_cancellable) — a loop or
16//!   watcher handed a [`CancellationToken`] it must observe;
17//!   [`shutdown`](BackgroundTasks::shutdown) cancels the token before it
18//!   drains.
19//!
20//! Every task carries the tracing span and subscriber of the code that spawned
21//! it, so its events correlate with the originating request and reach the same
22//! collector. A task that panics is logged by name instead of vanishing.
23//!
24//! Work whose handle a single owner must join or abort (a listener, a relay,
25//! a request-scoped worker) uses [`OwnedTask`], which stores its
26//! [`JoinHandle`] and aborts it on drop.
27//!
28//! Copyright (c) systemprompt.io — Business Source License 1.1.
29//! See <https://systemprompt.io> for licensing details.
30
31use std::future::Future;
32use std::panic::AssertUnwindSafe;
33use std::time::Duration;
34
35use futures::FutureExt;
36use tokio::task::{JoinError, JoinHandle};
37use tokio_util::sync::CancellationToken;
38use tokio_util::task::TaskTracker;
39use tracing::Instrument;
40use tracing::instrument::WithSubscriber;
41
42/// Process-wide owner of detached work; cloning shares the same tracker.
43#[derive(Debug, Clone, Default)]
44pub struct BackgroundTasks {
45    tracker: TaskTracker,
46    cancel: CancellationToken,
47}
48
49/// Result of waiting for in-flight background work.
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51#[must_use]
52pub enum DrainOutcome {
53    Drained,
54    TimedOut { in_flight: usize },
55}
56
57impl DrainOutcome {
58    #[must_use]
59    pub const fn is_drained(self) -> bool {
60        matches!(self, Self::Drained)
61    }
62}
63
64impl BackgroundTasks {
65    #[must_use]
66    pub fn new() -> Self {
67        Self::default()
68    }
69
70    pub fn spawn<F>(&self, name: &'static str, task: F)
71    where
72        F: Future<Output = ()> + Send + 'static,
73    {
74        self.tracker.spawn(
75            supervised(name, task)
76                .in_current_span()
77                .with_current_subscriber(),
78        );
79    }
80
81    pub fn spawn_cancellable<F, Fut>(&self, name: &'static str, task: F)
82    where
83        F: FnOnce(CancellationToken) -> Fut,
84        Fut: Future<Output = ()> + Send + 'static,
85    {
86        let token = self.cancel.child_token();
87        self.spawn(name, task(token));
88    }
89
90    #[must_use]
91    pub fn cancellation_token(&self) -> CancellationToken {
92        self.cancel.child_token()
93    }
94
95    #[must_use]
96    pub fn in_flight(&self) -> usize {
97        self.tracker.len()
98    }
99
100    #[must_use]
101    pub fn is_shutting_down(&self) -> bool {
102        self.cancel.is_cancelled()
103    }
104
105    pub async fn drain(&self, timeout: Duration) -> DrainOutcome {
106        self.tracker.close();
107        let waited = tokio::time::timeout(timeout, self.tracker.wait()).await;
108        self.tracker.reopen();
109        match waited {
110            Ok(()) => DrainOutcome::Drained,
111            Err(_elapsed) => DrainOutcome::TimedOut {
112                in_flight: self.tracker.len(),
113            },
114        }
115    }
116
117    pub async fn shutdown(&self, timeout: Duration) -> DrainOutcome {
118        self.cancel.cancel();
119        self.drain(timeout).await
120    }
121}
122
123async fn supervised<F>(name: &'static str, task: F)
124where
125    F: Future<Output = ()> + Send + 'static,
126{
127    if AssertUnwindSafe(task).catch_unwind().await.is_err() {
128        tracing::error!(task = name, "Background task panicked");
129    }
130}
131
132/// A single task whose owner joins or aborts it; dropping it aborts the task.
133#[derive(Debug)]
134#[must_use = "dropping an OwnedTask aborts it"]
135pub struct OwnedTask<T> {
136    name: &'static str,
137    handle: JoinHandle<T>,
138}
139
140impl<T: Send + 'static> OwnedTask<T> {
141    pub fn spawn<F>(name: &'static str, task: F) -> Self
142    where
143        F: Future<Output = T> + Send + 'static,
144    {
145        Self {
146            name,
147            handle: tokio::spawn(task.in_current_span().with_current_subscriber()),
148        }
149    }
150
151    #[must_use]
152    pub const fn name(&self) -> &'static str {
153        self.name
154    }
155
156    #[must_use]
157    pub fn is_finished(&self) -> bool {
158        self.handle.is_finished()
159    }
160
161    pub fn abort(&self) {
162        self.handle.abort();
163    }
164
165    pub async fn join(mut self) -> Result<T, JoinError> {
166        (&mut self.handle).await
167    }
168
169    pub async fn abort_and_join(mut self) -> Option<T> {
170        self.handle.abort();
171        match (&mut self.handle).await {
172            Ok(value) => Some(value),
173            Err(e) if e.is_cancelled() => None,
174            Err(e) => {
175                tracing::error!(task = self.name, error = %e, "Owned task panicked");
176                None
177            },
178        }
179    }
180}
181
182impl<T> Drop for OwnedTask<T> {
183    fn drop(&mut self) {
184        self.handle.abort();
185    }
186}
systemprompt_traits/background_tasks.rs

systemprompt_traits/
background_tasks.rs
