taskvisor 0.8.0

In-process Tokio task supervisor with retries, graceful shutdown, reliable final outcomes, and per-key admission control
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
//! Commits final registry state after the actor join produces a result.
//!
//! Every removal source reaches this module with one [`RemovalReport`]. The join result is first
//! mapped to its public outcome. Terminal commit then removes the identity and label indexes under
//! the state lock. After the lock is released, it publishes terminal events, resolves a watched outcome,
//! and transfers user values with their capacity reservation to physical and deferred cleanup.
//!
//! Logical completion, pending-join release, and empty-registry notification form a guarded tail.
//! They still run if reporting unwinds. If a future waiting for the state lock is cancelled,
//! its report moves to a detached continuation when a Tokio runtime is available.
//! No user-owned terminal value is dropped on the listener or actor task by this module.

use std::sync::Arc;

use tokio::sync::{Notify, RwLock};

use super::{JoinCompletion, PendingJoins, RemovalReport};
use crate::{
    core::{
        actor::ActorExitReason,
        deferred_drop::DropBundle,
        outcome::TaskOutcome,
        registry::{
            Registry,
            completion::{OutcomeTx, RemovalCompletion},
            scheduler::{ActorJoinError, AttemptReaper},
            state::{EntryState, Inner},
        },
    },
    events::{Bus, Event, EventKind},
    identity::TaskId,
};

/// Keeps an undelivered outcome on the isolated destruction path.
///
/// The guard retains ownership if event publication or outcome delivery unwinds.
struct OutcomeDropGuard<'a> {
    /// Outcome retained until its watched channel accepts ownership.
    outcome: Option<TaskOutcome>,
    /// Bundle used when the outcome cannot be delivered.
    cleanup: &'a mut DropBundle,
}

impl<'a> OutcomeDropGuard<'a> {
    /// Creates a guard for one classified terminal outcome.
    fn new(outcome: TaskOutcome, cleanup: &'a mut DropBundle) -> Self {
        Self {
            outcome: Some(outcome),
            cleanup,
        }
    }

    /// Borrows the outcome while the guard retains ownership.
    fn get(&self) -> &TaskOutcome {
        self.outcome
            .as_ref()
            .expect("the terminal outcome remains guarded until delivery")
    }

    /// Transfers the outcome to its watched channel at most once.
    fn take(&mut self) -> TaskOutcome {
        self.outcome
            .take()
            .expect("the terminal outcome is delivered at most once")
    }
}

impl Drop for OutcomeDropGuard<'_> {
    /// Moves an undelivered outcome into the cleanup bundle.
    fn drop(&mut self) {
        if let Some(outcome) = self.outcome.take() {
            self.cleanup.attach_outcome(outcome);
        }
    }
}

/// Owns one terminal report until registry membership is removed.
///
/// Cancellation while waiting for the state lock can spawn one detached commit.
/// Physical ownership remains retained if no runtime can run that continuation.
struct PendingTerminalReport {
    /// Authoritative membership state.
    state: Arc<RwLock<Inner>>,
    /// Notification used when terminal removal empties the registry.
    empty_notify: Arc<Notify>,
    /// Barrier for every claimed removal owner.
    pending_joins: Arc<PendingJoins>,
    /// Event destination for terminal diagnostics.
    bus: Bus,
    /// Identity whose removing entry must be committed.
    id: TaskId,
    /// Classified terminal outcome awaiting delivery.
    outcome: Option<TaskOutcome>,
    /// Optional watched outcome sender.
    done: Option<OutcomeTx>,
    /// User values and reserved capacity awaiting isolated destruction.
    cleanup: Option<DropBundle>,
    /// Physical ownership sink for terminal values.
    reaper: AttemptReaper,
    /// Two-phase completion for this removal.
    completion: RemovalCompletion,
    /// Whether cancellation should spawn one commit continuation.
    detach_on_drop: bool,
}

impl PendingTerminalReport {
    /// Takes the outcome, sender, and cleanup bundle for terminal reporting.
    fn take(&mut self) -> (TaskOutcome, Option<OutcomeTx>, DropBundle) {
        (
            self.outcome
                .take()
                .expect("one terminal outcome is classified"),
            self.done.take(),
            self.cleanup
                .take()
                .expect("one charged terminal bundle is retained"),
        )
    }

    /// Removes membership and commits terminal side effects without another await.
    ///
    /// Completion waiters wake only after both indexes are updated, reporting is attempted, and the state lock is released.
    async fn commit(&mut self) {
        let removed = {
            let mut st = self.state.write().await;
            let is_removing = st
                .tasks
                .get(&self.id)
                .is_some_and(|entry| matches!(&entry.state, EntryState::Removing { .. }));
            if !is_removing {
                None
            } else {
                let entry = st
                    .tasks
                    .remove(&self.id)
                    .expect("the removing entry was checked above");
                let EntryState::Removing {
                    completion: state_completion,
                } = entry.state
                else {
                    unreachable!("the removing entry was checked above")
                };
                if st.by_label.get(entry.label.as_ref()) == Some(&self.id) {
                    st.by_label.remove(entry.label.as_ref());
                }
                let is_empty = st.tasks.is_empty();
                Some((entry.label, state_completion, is_empty))
            }
        };

        let Some((label, state_completion, is_empty)) = removed else {
            self.finish_without_membership();
            return;
        };

        let (terminal_outcome, outcome, cleanup) = self.take();
        let mut finalizer = TerminalFinalizer {
            id: self.id,
            empty_notify: &self.empty_notify,
            pending_joins: &self.pending_joins,
            state_completion: Some(state_completion),
            report_completion: self.completion.clone(),
            is_empty,
            terminal: Some((self.reaper.clone(), cleanup)),
        };
        let cleanup = &mut finalizer
            .terminal
            .as_mut()
            .expect("terminal ownership is installed")
            .1;

        Registry::report_outcome(
            &self.bus,
            self.id,
            &label,
            terminal_outcome,
            outcome,
            cleanup,
        );
        drop(finalizer);
    }

    /// Completes a stale report after the state lock confirms missing membership.
    fn finish_without_membership(&mut self) {
        if self.retain_terminal() {
            self.pending_joins.dec(self.id);
            self.completion.complete_logical();
        }
    }

    /// Transfers terminal ownership without completing registry barriers.
    fn retain_terminal(&mut self) -> bool {
        let Some(mut cleanup) = self.cleanup.take() else {
            return false;
        };
        if let Some(outcome) = self.outcome.take() {
            match self.done.take() {
                Some(done) => {
                    if let Err(undelivered) = done.send(outcome) {
                        cleanup.attach_outcome(undelivered);
                    }
                }
                None => cleanup.attach_outcome(outcome),
            }
        }
        self.reaper
            .attach_terminal(self.id, cleanup, None, self.completion.clone());
        true
    }

    /// Moves an interrupted commit into a non-detaching continuation guard.
    fn take_continuation(&mut self) -> Self {
        Self {
            state: Arc::clone(&self.state),
            empty_notify: Arc::clone(&self.empty_notify),
            pending_joins: Arc::clone(&self.pending_joins),
            bus: self.bus.clone(),
            id: self.id,
            outcome: self.outcome.take(),
            done: self.done.take(),
            cleanup: self.cleanup.take(),
            reaper: self.reaper.clone(),
            completion: self.completion.clone(),
            detach_on_drop: false,
        }
    }

    /// Retains physical ownership when a detached commit cannot finish.
    ///
    /// Membership and logical completion remain pending because no state commit occurred.
    fn retain_without_logical_completion(&mut self) {
        let _ = self.retain_terminal();
    }
}

impl Drop for PendingTerminalReport {
    /// Continues an interrupted commit or retains its physical ownership.
    fn drop(&mut self) {
        if self.cleanup.is_none() {
            return;
        }

        if self.detach_on_drop
            && let Ok(runtime) = tokio::runtime::Handle::try_current()
        {
            let mut continuation = self.take_continuation();
            drop(runtime.spawn(async move {
                continuation.commit().await;
            }));
            return;
        }

        self.retain_without_logical_completion();
    }
}

/// Completes the terminal tail even if outcome reporting unwinds.
pub(in crate::core::registry) struct TerminalFinalizer<'a> {
    /// Identity whose terminal ownership is released.
    pub(in crate::core::registry) id: TaskId,
    /// Notification used when removal empties the registry.
    pub(in crate::core::registry) empty_notify: &'a Notify,
    /// Join-owner barrier decremented by this finalizer.
    pub(in crate::core::registry) pending_joins: &'a PendingJoins,
    /// Completion stored in the authoritative removing entry.
    pub(in crate::core::registry) state_completion: Option<RemovalCompletion>,
    /// Completion carried by the terminal report.
    pub(in crate::core::registry) report_completion: RemovalCompletion,
    /// Whether membership removal left the registry empty.
    pub(in crate::core::registry) is_empty: bool,
    /// Reaper and cleanup bundle that retain physical ownership.
    pub(in crate::core::registry) terminal: Option<(AttemptReaper, DropBundle)>,
}

impl Drop for TerminalFinalizer<'_> {
    /// Transfers terminal ownership and completes every logical barrier.
    fn drop(&mut self) {
        if let Some((reaper, bundle)) = self.terminal.take() {
            reaper.attach_terminal(
                self.id,
                bundle,
                self.state_completion.clone(),
                self.report_completion.clone(),
            );
        }
        self.pending_joins.dec(self.id);
        if let Some(completion) = &self.state_completion {
            completion.complete_logical();
        }
        self.report_completion.complete_logical();
        if self.is_empty {
            self.empty_notify.notify_waiters();
        }
    }
}

impl Registry {
    /// Commits terminal cleanup for one removing entry.
    ///
    /// Membership removal, reporting, and pending-join accounting finish before an empty-registry waiter can continue.
    pub(in crate::core::registry) async fn finish_removal(
        state: &Arc<RwLock<Inner>>,
        empty_notify: &Arc<Notify>,
        pending_joins: &Arc<PendingJoins>,
        bus: &Bus,
        reaper: &AttemptReaper,
        report: RemovalReport,
    ) {
        let RemovalReport {
            id,
            outcome,
            join,
            completion: removal_completion,
            mut cleanup,
        } = report;
        let terminal_outcome = match join {
            JoinCompletion::Joined(result) => Self::outcome_of(result, &mut cleanup),
            JoinCompletion::ForceAborted => TaskOutcome::ForceAborted,
        };
        let mut pending_report = PendingTerminalReport {
            state: Arc::clone(state),
            empty_notify: Arc::clone(empty_notify),
            pending_joins: Arc::clone(pending_joins),
            bus: bus.clone(),
            id,
            outcome: Some(terminal_outcome),
            done: outcome,
            cleanup: Some(cleanup),
            reaper: reaper.clone(),
            completion: removal_completion,
            detach_on_drop: true,
        };
        pending_report.commit().await;
    }

    /// Publishes terminal events and delivers the watched outcome.
    ///
    /// The event and waiter use the same [`TaskOutcome`] classification.
    fn report_outcome(
        bus: &Bus,
        id: TaskId,
        label: &str,
        outcome: TaskOutcome,
        done: Option<OutcomeTx>,
        cleanup: &mut DropBundle,
    ) {
        let mut outcome = OutcomeDropGuard::new(outcome, cleanup);
        bus.publish_lazy(|| {
            let mut finished = Event::new(EventKind::TaskFinished)
                .with_task(label)
                .with_id(id)
                .with_outcome_kind(outcome.get().kind());
            match outcome.get() {
                TaskOutcome::Failed {
                    reason, exit_code, ..
                }
                | TaskOutcome::Fatal {
                    reason, exit_code, ..
                } => {
                    finished = finished.with_reason(Arc::clone(reason));
                    if let Some(code) = exit_code {
                        finished = finished.with_exit_code(*code);
                    }
                }
                TaskOutcome::ForceAborted => {
                    finished =
                        finished.with_reason("task did not stop within grace; force-aborted");
                }
                TaskOutcome::Panicked => {
                    finished = finished.with_reason("internal task runner panicked");
                }
                TaskOutcome::Completed | TaskOutcome::Canceled | TaskOutcome::Rejected { .. } => {}
            }
            finished
        });
        if let Some(done) = done
            && let Err(undelivered) = done.send(outcome.take())
        {
            outcome.cleanup.attach_outcome(undelivered);
        }
        bus.publish_lazy(|| {
            Event::new(EventKind::TaskRemoved)
                .with_task(label)
                .with_id(id)
        });
    }

    /// Maps an actor join result to its public terminal outcome.
    fn outcome_of(
        res: Result<ActorExitReason, ActorJoinError>,
        cleanup: &mut DropBundle,
    ) -> TaskOutcome {
        match res {
            Ok(ActorExitReason::Completed) => TaskOutcome::Completed,
            Ok(ActorExitReason::Canceled) => TaskOutcome::Canceled,
            Ok(ActorExitReason::Panicked { cleanup_poisoned }) => {
                if cleanup_poisoned {
                    cleanup.poison();
                }
                TaskOutcome::Panicked
            }
            Ok(ActorExitReason::Exhausted {
                reason,
                exit_code,
                source,
            }) => TaskOutcome::Failed {
                reason,
                exit_code,
                source,
            },
            Ok(ActorExitReason::Fatal {
                reason,
                exit_code,
                source,
            }) => TaskOutcome::Fatal {
                reason,
                exit_code,
                source,
            },
            Err(e) if e.is_panic() => {
                if e.cleanup_poisoned() {
                    cleanup.poison();
                }
                TaskOutcome::Panicked
            }
            Err(_aborted) => TaskOutcome::ForceAborted,
        }
    }
}