a3s_code_core/capability/
supervisor.rs

1use std::collections::BTreeMap;
2use std::future::Future;
3use std::num::NonZeroU64;
4use std::sync::{Arc, Mutex, MutexGuard, Weak};
5use std::time::Duration;
6
7use async_trait::async_trait;
8use tokio::sync::Notify;
9use tokio::task::{JoinError, JoinHandle, JoinSet};
10use tokio::time::Instant;
11use tokio_util::sync::CancellationToken;
12
13use super::{CapabilityEffect, CapabilityEffectError, CapabilityScopeError, RetainedUseGeneration};
14
15pub const MAX_SCOPE_EFFECTS: usize = 1_024;
16pub const MAX_SCOPE_TASKS: usize = 1_024;
17pub const MAX_SCOPE_CHILDREN: usize = 1_024;
18pub const DEFAULT_SCOPE_CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
19pub const MAX_SCOPE_CLOSE_TIMEOUT: Duration = Duration::from_secs(60);
20
21const MAX_LIFECYCLE_NAME_BYTES: usize = 128;
22const ABORT_SETTLE_GRACE: Duration = Duration::from_millis(100);
23
24/// Bounded close policy inherited by every child scope.
25#[derive(Clone, Copy, Debug, Eq, PartialEq)]
26pub struct ScopeClosePolicy {
27    timeout: Duration,
28}
29
30impl ScopeClosePolicy {
31    pub fn new(timeout: Duration) -> Result<Self, CapabilityScopeError> {
32        if timeout.is_zero() {
33            return Err(CapabilityScopeError::InvalidExecutionLimit {
34                field: "scope_close_timeout",
35            });
36        }
37        if timeout > MAX_SCOPE_CLOSE_TIMEOUT {
38            return Err(CapabilityScopeError::BoundExceeded {
39                field: "scope_close_timeout_ms",
40                max: MAX_SCOPE_CLOSE_TIMEOUT.as_millis() as usize,
41            });
42        }
43        Ok(Self { timeout })
44    }
45
46    pub const fn timeout(self) -> Duration {
47        self.timeout
48    }
49}
50
51impl Default for ScopeClosePolicy {
52    fn default() -> Self {
53        Self {
54            timeout: DEFAULT_SCOPE_CLOSE_TIMEOUT,
55        }
56    }
57}
58
59/// Opaque identity of one task owned by a capability scope supervisor.
60#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
61pub struct SupervisedTaskId(NonZeroU64);
62
63impl SupervisedTaskId {
64    pub const fn get(self) -> u64 {
65        self.0.get()
66    }
67}
68
69/// Deterministic, bounded outcome of one idempotent scope close.
70#[derive(Clone, Debug, Default, Eq, PartialEq)]
71pub struct ScopeCloseReport {
72    pub tasks_completed: usize,
73    pub tasks_failed: usize,
74    pub tasks_cancelled: usize,
75    pub tasks_timed_out: usize,
76    pub child_scopes_closed: usize,
77    pub child_scopes_failed: usize,
78    pub child_scopes_timed_out: usize,
79    pub effects_closed: usize,
80    pub effects_failed: usize,
81    pub effects_timed_out: usize,
82    pub generation_leases_released: usize,
83}
84
85impl ScopeCloseReport {
86    pub const fn is_clean(&self) -> bool {
87        self.tasks_failed == 0
88            && self.tasks_timed_out == 0
89            && self.child_scopes_failed == 0
90            && self.child_scopes_timed_out == 0
91            && self.effects_failed == 0
92            && self.effects_timed_out == 0
93    }
94}
95
96struct SupervisedTaskOutcome {
97    name: Box<str>,
98    result: Result<(), CapabilityEffectError>,
99}
100
101struct RegisteredEffect {
102    name: Box<str>,
103    effect: Box<dyn CapabilityEffect>,
104}
105
106#[async_trait]
107pub(super) trait SupervisedChild: Send {
108    fn name(&self) -> &str;
109
110    fn cancel(&self);
111
112    async fn close(self: Box<Self>) -> Result<ScopeCloseReport, CapabilityScopeError>;
113}
114
115struct OpenSupervisor {
116    tasks: JoinSet<SupervisedTaskOutcome>,
117    children: BTreeMap<u64, Box<dyn SupervisedChild>>,
118    effects: Vec<RegisteredEffect>,
119    generation_leases: Vec<Box<dyn RetainedUseGeneration>>,
120    next_task_id: u64,
121    next_child_id: u64,
122}
123
124impl Default for OpenSupervisor {
125    fn default() -> Self {
126        Self {
127            tasks: JoinSet::new(),
128            children: BTreeMap::new(),
129            effects: Vec::new(),
130            generation_leases: Vec::new(),
131            next_task_id: 1,
132            next_child_id: 1,
133        }
134    }
135}
136
137enum SupervisorState {
138    Open(OpenSupervisor),
139    Closing { driver: Option<JoinHandle<()>> },
140    Closed(ScopeCloseReport),
141}
142
143pub(super) struct SupervisorInner {
144    scope_id: Box<str>,
145    cancellation: CancellationToken,
146    policy: ScopeClosePolicy,
147    state: Mutex<SupervisorState>,
148    closed: Notify,
149}
150
151/// Cloneable weak handle that can add work to an open capability scope.
152///
153/// The handle cannot keep the scope or its A3S Use generation lease alive by
154/// itself. Successfully registered tasks transfer ownership to the scope's
155/// structured-concurrency supervisor.
156#[derive(Clone)]
157pub(crate) struct SupervisedTaskSpawner {
158    inner: Weak<SupervisorInner>,
159    scope_id: Arc<str>,
160}
161
162impl SupervisedTaskSpawner {
163    pub(crate) fn spawn_task<F>(
164        &self,
165        name: impl Into<String>,
166        task: F,
167    ) -> Result<SupervisedTaskId, CapabilityScopeError>
168    where
169        F: Future<Output = Result<(), CapabilityEffectError>> + Send + 'static,
170    {
171        let Some(inner) = self.inner.upgrade() else {
172            return Err(CapabilityScopeError::SupervisorClosed {
173                scope_id: self.scope_id.to_string(),
174            });
175        };
176        spawn_task(&inner, name.into(), task)
177    }
178}
179
180impl SupervisorInner {
181    fn lock_state(&self) -> MutexGuard<'_, SupervisorState> {
182        self.state
183            .lock()
184            .unwrap_or_else(std::sync::PoisonError::into_inner)
185    }
186}
187
188impl Drop for SupervisorInner {
189    fn drop(&mut self) {
190        self.cancellation.cancel();
191        let state = self
192            .state
193            .get_mut()
194            .unwrap_or_else(std::sync::PoisonError::into_inner);
195        if let SupervisorState::Open(open) = state {
196            open.tasks.abort_all();
197        }
198    }
199}
200
201/// Owns every asynchronous task, child scope, reversible effect, and exact
202/// upstream generation lease for one capability scope.
203pub(super) struct EffectSupervisor {
204    inner: Arc<SupervisorInner>,
205}
206
207impl EffectSupervisor {
208    pub(super) fn new(
209        scope_id: impl Into<String>,
210        cancellation: CancellationToken,
211        policy: ScopeClosePolicy,
212    ) -> Self {
213        Self {
214            inner: Arc::new(SupervisorInner {
215                scope_id: scope_id.into().into_boxed_str(),
216                cancellation,
217                policy,
218                state: Mutex::new(SupervisorState::Open(OpenSupervisor::default())),
219                closed: Notify::new(),
220            }),
221        }
222    }
223
224    pub(super) fn policy(&self) -> ScopeClosePolicy {
225        self.inner.policy
226    }
227
228    pub(super) fn cancellation(&self) -> CancellationToken {
229        self.inner.cancellation.clone()
230    }
231
232    pub(super) fn is_open(&self) -> bool {
233        matches!(*self.inner.lock_state(), SupervisorState::Open(_))
234            && !self.inner.cancellation.is_cancelled()
235    }
236
237    pub(super) fn register_effect(
238        &self,
239        effect: Box<dyn CapabilityEffect>,
240    ) -> Result<(), CapabilityScopeError> {
241        let name = effect.name().to_owned();
242        validate_lifecycle_name(&name)?;
243        let mut state = self.inner.lock_state();
244        let SupervisorState::Open(open) = &mut *state else {
245            return Err(self.closed_error());
246        };
247        if self.inner.cancellation.is_cancelled() {
248            return Err(self.closed_error());
249        }
250        if open.effects.len() >= MAX_SCOPE_EFFECTS {
251            return Err(CapabilityScopeError::BoundExceeded {
252                field: "scope_effects",
253                max: MAX_SCOPE_EFFECTS,
254            });
255        }
256        open.effects.push(RegisteredEffect {
257            name: name.into_boxed_str(),
258            effect,
259        });
260        Ok(())
261    }
262
263    pub(super) fn register_generation_lease(
264        &self,
265        lease: Box<dyn RetainedUseGeneration>,
266    ) -> Result<(), CapabilityScopeError> {
267        let mut state = self.inner.lock_state();
268        let SupervisorState::Open(open) = &mut *state else {
269            return Err(self.closed_error());
270        };
271        if self.inner.cancellation.is_cancelled() {
272            return Err(self.closed_error());
273        }
274        if !open.generation_leases.is_empty() {
275            return Err(CapabilityScopeError::BoundExceeded {
276                field: "use_generation_leases",
277                max: 1,
278            });
279        }
280        open.generation_leases.push(lease);
281        Ok(())
282    }
283
284    pub(super) fn spawn_task<F>(
285        &self,
286        name: impl Into<String>,
287        task: F,
288    ) -> Result<SupervisedTaskId, CapabilityScopeError>
289    where
290        F: Future<Output = Result<(), CapabilityEffectError>> + Send + 'static,
291    {
292        spawn_task(&self.inner, name.into(), task)
293    }
294
295    pub(super) fn task_spawner(&self) -> SupervisedTaskSpawner {
296        SupervisedTaskSpawner {
297            inner: Arc::downgrade(&self.inner),
298            scope_id: Arc::from(self.inner.scope_id.as_ref()),
299        }
300    }
301
302    pub(super) fn register_child(
303        &self,
304        child: Box<dyn SupervisedChild>,
305    ) -> Result<u64, CapabilityScopeError> {
306        let mut state = self.inner.lock_state();
307        let SupervisorState::Open(open) = &mut *state else {
308            return Err(self.closed_error());
309        };
310        if self.inner.cancellation.is_cancelled() {
311            return Err(self.closed_error());
312        }
313        if open
314            .children
315            .values()
316            .any(|existing| existing.name() == child.name())
317        {
318            return Err(CapabilityScopeError::DuplicateChildScope {
319                scope_id: child.name().to_owned(),
320            });
321        }
322        if open.children.len() >= MAX_SCOPE_CHILDREN {
323            return Err(CapabilityScopeError::BoundExceeded {
324                field: "scope_children",
325                max: MAX_SCOPE_CHILDREN,
326            });
327        }
328        let id = open.next_child_id;
329        open.next_child_id = id
330            .checked_add(1)
331            .ok_or(CapabilityScopeError::ChildIdentityExhausted)?;
332        open.children.insert(id, child);
333        Ok(id)
334    }
335
336    pub(super) fn downgrade(&self) -> Weak<SupervisorInner> {
337        Arc::downgrade(&self.inner)
338    }
339
340    pub(super) fn cancel(&self) {
341        self.inner.cancellation.cancel();
342        let mut state = self.inner.lock_state();
343        match &mut *state {
344            SupervisorState::Open(open) => {
345                open.tasks.abort_all();
346                for child in open.children.values() {
347                    child.cancel();
348                }
349            }
350            SupervisorState::Closing { driver } => {
351                let _close_driver_started = driver.is_some();
352            }
353            SupervisorState::Closed(_) => {}
354        }
355    }
356
357    pub(super) async fn close(&self) -> Result<ScopeCloseReport, CapabilityScopeError> {
358        let runtime = tokio::runtime::Handle::try_current()
359            .map_err(|_| CapabilityScopeError::TokioRuntimeUnavailable)?;
360
361        loop {
362            let notified = self.inner.closed.notified();
363            tokio::pin!(notified);
364            notified.as_mut().enable();
365
366            let start = {
367                let mut state = self.inner.lock_state();
368                match &*state {
369                    SupervisorState::Closed(report) => return Ok(report.clone()),
370                    SupervisorState::Closing { driver } => {
371                        let _close_driver_started = driver.is_some();
372                        None
373                    }
374                    SupervisorState::Open(_) => {
375                        let previous = std::mem::replace(
376                            &mut *state,
377                            SupervisorState::Closing { driver: None },
378                        );
379                        match previous {
380                            SupervisorState::Open(open) => Some(open),
381                            SupervisorState::Closing { .. } | SupervisorState::Closed(_) => None,
382                        }
383                    }
384                }
385            };
386
387            if let Some(open) = start {
388                let inner = Arc::clone(&self.inner);
389                let driver = runtime.spawn(async move {
390                    drive_close(inner, open).await;
391                });
392                let mut state = self.inner.lock_state();
393                if let SupervisorState::Closing { driver: slot } = &mut *state {
394                    *slot = Some(driver);
395                }
396            }
397
398            notified.await;
399        }
400    }
401
402    fn closed_error(&self) -> CapabilityScopeError {
403        CapabilityScopeError::SupervisorClosed {
404            scope_id: self.inner.scope_id.to_string(),
405        }
406    }
407}
408
409fn spawn_task<F>(
410    inner: &Arc<SupervisorInner>,
411    name: String,
412    task: F,
413) -> Result<SupervisedTaskId, CapabilityScopeError>
414where
415    F: Future<Output = Result<(), CapabilityEffectError>> + Send + 'static,
416{
417    validate_lifecycle_name(&name)?;
418    tokio::runtime::Handle::try_current()
419        .map_err(|_| CapabilityScopeError::TokioRuntimeUnavailable)?;
420
421    let mut state = inner.lock_state();
422    let SupervisorState::Open(open) = &mut *state else {
423        return Err(CapabilityScopeError::SupervisorClosed {
424            scope_id: inner.scope_id.to_string(),
425        });
426    };
427    if inner.cancellation.is_cancelled() {
428        return Err(CapabilityScopeError::SupervisorClosed {
429            scope_id: inner.scope_id.to_string(),
430        });
431    }
432    if open.tasks.len() >= MAX_SCOPE_TASKS {
433        return Err(CapabilityScopeError::BoundExceeded {
434            field: "scope_tasks",
435            max: MAX_SCOPE_TASKS,
436        });
437    }
438    let raw_id = open.next_task_id;
439    let id = NonZeroU64::new(raw_id).ok_or(CapabilityScopeError::TaskIdentityExhausted)?;
440    open.next_task_id = raw_id
441        .checked_add(1)
442        .ok_or(CapabilityScopeError::TaskIdentityExhausted)?;
443    let task_name = name.into_boxed_str();
444    let _abort_handle = open.tasks.spawn(async move {
445        SupervisedTaskOutcome {
446            name: task_name,
447            result: task.await,
448        }
449    });
450    Ok(SupervisedTaskId(id))
451}
452
453impl Drop for EffectSupervisor {
454    fn drop(&mut self) {
455        self.cancel();
456    }
457}
458
459pub(super) fn remove_registered_child(parent: &Weak<SupervisorInner>, id: u64) {
460    let Some(parent) = parent.upgrade() else {
461        return;
462    };
463    let mut state = parent.lock_state();
464    if let SupervisorState::Open(open) = &mut *state {
465        open.children.remove(&id);
466    }
467}
468
469async fn drive_close(inner: Arc<SupervisorInner>, mut open: OpenSupervisor) {
470    inner.cancellation.cancel();
471    let deadline = Instant::now() + inner.policy.timeout();
472    let mut report = ScopeCloseReport::default();
473
474    settle_tasks(&mut open.tasks, deadline, &mut report).await;
475    close_children(open.children, deadline, &mut report).await;
476    close_effects(open.effects, deadline, &mut report).await;
477    while let Some(lease) = open.generation_leases.pop() {
478        drop(lease);
479        report.generation_leases_released += 1;
480    }
481
482    {
483        let mut state = inner.lock_state();
484        *state = SupervisorState::Closed(report);
485    }
486    inner.closed.notify_waiters();
487}
488
489async fn settle_tasks(
490    tasks: &mut JoinSet<SupervisedTaskOutcome>,
491    deadline: Instant,
492    report: &mut ScopeCloseReport,
493) {
494    loop {
495        if tasks.is_empty() {
496            return;
497        }
498        match tokio::time::timeout_at(deadline, tasks.join_next()).await {
499            Ok(Some(result)) => record_task_result(result, report),
500            Ok(None) => return,
501            Err(_) => {
502                let remaining = tasks.len();
503                report.tasks_timed_out += remaining;
504                tasks.abort_all();
505                let _ = tokio::time::timeout(ABORT_SETTLE_GRACE, tasks.shutdown()).await;
506                return;
507            }
508        }
509    }
510}
511
512fn record_task_result(
513    result: Result<SupervisedTaskOutcome, JoinError>,
514    report: &mut ScopeCloseReport,
515) {
516    match result {
517        Ok(outcome) => match outcome.result {
518            Ok(()) => report.tasks_completed += 1,
519            Err(error) => {
520                report.tasks_failed += 1;
521                tracing::warn!(task = %outcome.name, error = %error, "Capability scope task failed");
522            }
523        },
524        Err(error) if error.is_cancelled() => report.tasks_cancelled += 1,
525        Err(error) => {
526            report.tasks_failed += 1;
527            tracing::warn!(error = %error, "Capability scope task panicked");
528        }
529    }
530}
531
532async fn close_children(
533    children: BTreeMap<u64, Box<dyn SupervisedChild>>,
534    deadline: Instant,
535    report: &mut ScopeCloseReport,
536) {
537    for (_, child) in children.into_iter().rev() {
538        let name = child.name().to_owned();
539        if Instant::now() >= deadline {
540            drop(child);
541            report.child_scopes_timed_out += 1;
542            continue;
543        }
544        let mut close = tokio::spawn(async move { child.close().await });
545        match tokio::time::timeout_at(deadline, &mut close).await {
546            Ok(Ok(Ok(child_report))) => {
547                report.child_scopes_closed += 1;
548                if !child_report.is_clean() {
549                    report.child_scopes_failed += 1;
550                }
551            }
552            Ok(Ok(Err(error))) => {
553                report.child_scopes_failed += 1;
554                tracing::warn!(scope = %name, error = %error, "Child capability scope close failed");
555            }
556            Ok(Err(error)) => {
557                report.child_scopes_failed += 1;
558                tracing::warn!(scope = %name, error = %error, "Child capability scope close panicked");
559            }
560            Err(_) => {
561                close.abort();
562                let _ = tokio::time::timeout(ABORT_SETTLE_GRACE, &mut close).await;
563                report.child_scopes_timed_out += 1;
564            }
565        }
566    }
567}
568
569async fn close_effects(
570    effects: Vec<RegisteredEffect>,
571    deadline: Instant,
572    report: &mut ScopeCloseReport,
573) {
574    for registered in effects.into_iter().rev() {
575        let name = registered.name;
576        let effect = registered.effect;
577        if Instant::now() >= deadline {
578            drop(effect);
579            report.effects_timed_out += 1;
580            continue;
581        }
582        let mut close = tokio::spawn(async move { effect.close().await });
583        match tokio::time::timeout_at(deadline, &mut close).await {
584            Ok(Ok(Ok(()))) => report.effects_closed += 1,
585            Ok(Ok(Err(error))) => {
586                report.effects_failed += 1;
587                tracing::warn!(effect = %name, error = %error, "Capability effect close failed");
588            }
589            Ok(Err(error)) => {
590                report.effects_failed += 1;
591                tracing::warn!(effect = %name, error = %error, "Capability effect close panicked");
592            }
593            Err(_) => {
594                close.abort();
595                let _ = tokio::time::timeout(ABORT_SETTLE_GRACE, &mut close).await;
596                report.effects_timed_out += 1;
597            }
598        }
599    }
600}
601
602fn validate_lifecycle_name(value: &str) -> Result<(), CapabilityScopeError> {
603    if value.is_empty() {
604        return Err(CapabilityScopeError::InvalidLifecycleName {
605            reason: "it is empty",
606        });
607    }
608    if value.len() > MAX_LIFECYCLE_NAME_BYTES {
609        return Err(CapabilityScopeError::BoundExceeded {
610            field: "lifecycle_name",
611            max: MAX_LIFECYCLE_NAME_BYTES,
612        });
613    }
614    if !value.bytes().all(|byte| {
615        byte.is_ascii_lowercase() || byte.is_ascii_digit() || matches!(byte, b'.' | b'-' | b'_')
616    }) || !value
617        .as_bytes()
618        .first()
619        .is_some_and(u8::is_ascii_alphanumeric)
620        || !value
621            .as_bytes()
622            .last()
623            .is_some_and(u8::is_ascii_alphanumeric)
624    {
625        return Err(CapabilityScopeError::InvalidLifecycleName {
626            reason: "it contains non-canonical characters or boundaries",
627        });
628    }
629    Ok(())
630}
a3s_code_core/capability/supervisor.rs

a3s_code_core/capability/
supervisor.rs
