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

//! Applies authoritative runtime results to slot state.
//!
//! The controller driver passes direct registry Add decisions and reliable physical completion signals to this module.
//! A successful Add moves the matching admission to `Running`, unless an earlier replacement requires immediate removal.
//! A rejected Add clears the admission and advances queued work.
//!
//! Physical completion is the normal signal that releases a `Running` or `Terminating` owner.
//! Slot name and [`TaskId`] correlation make stale results harmless.

use std::sync::Arc;

use tokio::time::Instant;

use crate::{
    RuntimeError,
    events::{Event, EventKind},
    identity::TaskId,
};

use super::super::{
    AdmissionResult, CompletionResult, Controller, TrackedOperations, state::AdmissionTransition,
};

impl Controller {
    /// Routes one tracked registry Add decision to its slot.
    pub(in crate::controller::engine) async fn handle_admission_result(
        &self,
        result: AdmissionResult,
        operations: &mut TrackedOperations,
    ) {
        let AdmissionResult {
            id,
            slot_name,
            decision,
        } = result;
        self.handle_registry_decision(id, slot_name, decision, operations)
            .await;
    }

    /// Updates a slot from one registry Add decision.
    async fn handle_registry_decision(
        &self,
        id: TaskId,
        slot_name: Arc<str>,
        decision: Result<crate::core::RemovalCompletion, RuntimeError>,
        operations: &mut TrackedOperations,
    ) {
        let Some(slot_arc) = self.slot(&slot_name) else {
            return;
        };
        let mut slot = slot_arc.lock().await;

        match decision {
            Ok(completion) => match slot.confirm_admission(id, Instant::now()) {
                AdmissionTransition::Running => {
                    Self::track_completion(
                        &operations.completions,
                        id,
                        Arc::clone(&slot_name),
                        completion,
                    );
                    self.bus.publish_lazy(|| {
                        Event::new(EventKind::ControllerSlotTransition)
                            .with_task(Arc::clone(&slot_name))
                            .with_reason("admitting→running")
                    });
                }
                AdmissionTransition::RemoveNow(owner) => {
                    Self::track_completion(
                        &operations.completions,
                        id,
                        Arc::clone(&slot_name),
                        completion,
                    );
                    let Some(sup) = self.supervisor.upgrade() else {
                        return;
                    };
                    Self::track_removal(&operations.removals, sup, owner, slot_name);
                }
                AdmissionTransition::Stale => {}
            },
            Err(_) => {
                if !slot.reject_admission(id) {
                    return;
                }

                let deferred_drops = if !self.is_shutting_down()
                    && let Some(sup) = self.supervisor.upgrade()
                {
                    self.start_next_from_queue(&sup, &mut slot, &slot_name, operations)
                } else {
                    Vec::new()
                };

                self.gc_if_idle(&slot_name, slot);
                self.drop_pending_submissions(deferred_drops);
            }
        }
    }

    /// Releases the matching owner after reliable physical completion.
    pub(in crate::controller::engine) async fn handle_completion_result(
        &self,
        result: CompletionResult,
        operations: &mut TrackedOperations,
    ) {
        let Some(sup) = self.supervisor.upgrade() else {
            return;
        };
        let Some(slot_arc) = self.slot(&result.slot_name) else {
            return;
        };
        let mut slot = slot_arc.lock().await;
        if !slot.complete_owner(result.id) {
            return;
        }

        let deferred_drops = if !self.is_shutting_down() {
            self.start_next_from_queue(&sup, &mut slot, &result.slot_name, operations)
        } else {
            Vec::new()
        };

        self.gc_if_idle(&result.slot_name, slot);
        self.drop_pending_submissions(deferred_drops);
    }
}