laminar_connectors/connector/
task_tracking.rs

1//! Detached connector-task admission and terminal lifetime tracking.
2
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::sync::{Arc, Weak};
5
6use tokio::sync::Notify;
7
8/// What the runtime must do when a started connector operation is cancelled.
9///
10/// This is an internal connector/driver capability, not a deployment option.
11/// Cancellation always respects the runtime-owned deadline. A connector may be
12/// reused only when dropping the exact future is known to preserve its state;
13/// otherwise the runtime retires the complete connector generation.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum ConnectorCancellationPolicy {
16    /// Dropping an in-flight future leaves the connector valid for recovery or reuse.
17    CancelSafe,
18    /// Dropping an in-flight future may leave its external outcome unknown, so
19    /// the connector instance must not process another operation.
20    RetireConnector,
21}
22
23const CONNECTOR_TASK_OWNER_DROPPED: usize = 1usize << (usize::BITS - 1);
24
25pub(super) struct ConnectorTaskState {
26    state: AtomicUsize,
27    terminated: Notify,
28}
29
30/// Sole admission authority for detached tasks owned by one connector generation.
31///
32/// The owner must live inside the connector. Dropping it seals the generation so
33/// terminal completion can be observed after every admitted task guard is gone.
34pub struct ConnectorTaskOwner {
35    inner: Arc<ConnectorTaskState>,
36}
37
38/// Cloneable, non-owning admission handle for dynamically spawned connector tasks.
39///
40/// The handle does not keep the task generation open. Admission fails after
41/// the sole [`ConnectorTaskOwner`] is dropped, including when existing task
42/// guards still keep the generation observable.
43#[derive(Clone)]
44pub struct ConnectorTaskAdmission {
45    pub(super) inner: Weak<ConnectorTaskState>,
46}
47
48/// Cloneable observer for terminal completion of one connector generation.
49#[derive(Clone)]
50pub struct ConnectorTaskTracker {
51    inner: Arc<ConnectorTaskState>,
52}
53
54/// RAII proof that one connector-owned task is still active.
55///
56/// Move the guard into the task before spawning it and retain it for the task's
57/// full lifetime.
58#[must_use = "dropping the guard marks its connector task complete"]
59pub struct ConnectorTaskGuard {
60    inner: Arc<ConnectorTaskState>,
61}
62
63impl ConnectorTaskOwner {
64    /// Create the sole task owner and its cloneable terminal observer.
65    #[must_use]
66    pub fn new() -> (Self, ConnectorTaskTracker) {
67        let inner = Arc::new(ConnectorTaskState {
68            state: AtomicUsize::new(0),
69            terminated: Notify::new(),
70        });
71        (
72            Self {
73                inner: Arc::clone(&inner),
74            },
75            ConnectorTaskTracker { inner },
76        )
77    }
78
79    /// Create a non-owning admission handle for tasks discovered by owned work.
80    ///
81    /// This is intended for accept loops and similar dynamic task producers.
82    /// Cloning the handle never extends the connector generation's admission
83    /// lifetime.
84    #[must_use]
85    pub fn admission(&self) -> ConnectorTaskAdmission {
86        ConnectorTaskAdmission {
87            inner: Arc::downgrade(&self.inner),
88        }
89    }
90
91    /// Admit one task into this live connector generation.
92    ///
93    /// The returned guard must be created before the task is spawned and moved
94    /// into that task. `None` means the generation can no longer admit work.
95    #[must_use]
96    pub fn track(&self) -> Option<ConnectorTaskGuard> {
97        track_connector_task(&self.inner)
98    }
99}
100
101impl ConnectorTaskAdmission {
102    /// Admit one dynamic task while the connector generation remains open.
103    ///
104    /// Returns `None` once the sole owner has been dropped. A successful
105    /// admission remains tracked until its returned guard is dropped.
106    #[must_use]
107    pub fn track(&self) -> Option<ConnectorTaskGuard> {
108        let inner = self.inner.upgrade()?;
109        track_connector_task(&inner)
110    }
111}
112
113fn track_connector_task(inner: &Arc<ConnectorTaskState>) -> Option<ConnectorTaskGuard> {
114    let mut observed = inner.state.load(Ordering::Acquire);
115    loop {
116        if observed & CONNECTOR_TASK_OWNER_DROPPED != 0 {
117            return None;
118        }
119        let next = observed.checked_add(1)?;
120        if next & CONNECTOR_TASK_OWNER_DROPPED != 0 {
121            return None;
122        }
123        match inner
124            .state
125            .compare_exchange_weak(observed, next, Ordering::AcqRel, Ordering::Acquire)
126        {
127            Ok(_) => {
128                return Some(ConnectorTaskGuard {
129                    inner: Arc::clone(inner),
130                });
131            }
132            Err(actual) => observed = actual,
133        }
134    }
135}
136
137impl Drop for ConnectorTaskOwner {
138    fn drop(&mut self) {
139        let previous = self
140            .inner
141            .state
142            .fetch_or(CONNECTOR_TASK_OWNER_DROPPED, Ordering::AcqRel);
143        debug_assert_eq!(previous & CONNECTOR_TASK_OWNER_DROPPED, 0);
144        if previous == 0 {
145            self.inner.terminated.notify_waiters();
146        }
147    }
148}
149
150impl ConnectorTaskTracker {
151    /// Whether the owner and all task guards have been dropped.
152    #[must_use]
153    pub fn is_terminated(&self) -> bool {
154        self.inner.state.load(Ordering::Acquire) == CONNECTOR_TASK_OWNER_DROPPED
155    }
156
157    /// Wait until the owner and all task guards have been dropped.
158    pub async fn wait_terminated(&self) {
159        loop {
160            let notified = self.inner.terminated.notified();
161            tokio::pin!(notified);
162            notified.as_mut().enable();
163            if self.is_terminated() {
164                return;
165            }
166            notified.await;
167        }
168    }
169}
170
171impl Drop for ConnectorTaskGuard {
172    fn drop(&mut self) {
173        let previous = self.inner.state.fetch_sub(1, Ordering::AcqRel);
174        debug_assert_ne!(previous & !CONNECTOR_TASK_OWNER_DROPPED, 0);
175        if previous == CONNECTOR_TASK_OWNER_DROPPED | 1 {
176            self.inner.terminated.notify_waiters();
177        }
178    }
179}
laminar_connectors/connector/task_tracking.rs

laminar_connectors/connector/
task_tracking.rs
