khora_core/agent/
completion.rs

1// Copyright 2025 eraflo
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8
9//! Per-frame agent completion tracking — the synchronization primitive used
10//! by the scheduler to coordinate agents that declare `Hard` dependencies.
11//!
12//! At frame start the scheduler builds an [`AgentCompletionMap`] populated
13//! with one [`StageHandle<AgentDone>`] per known agent. After each agent
14//! `execute()` (or skip), the scheduler calls [`AgentCompletionMap::mark`]
15//! with a [`CompletionOutcome`].  Dependents either inspect the outcome
16//! synchronously via [`AgentCompletionMap::outcome`] or — once the parallel
17//! scheduler is enabled — `await` it via [`AgentCompletionMap::wait`].
18
19use std::collections::HashMap;
20use std::sync::OnceLock;
21
22use crate::control::gorna::AgentId;
23use crate::renderer::api::core::StageHandle;
24
25/// Marker type for the "an agent finished its frame work" stage.
26pub struct AgentDone;
27
28/// What happened to an agent this frame, from the scheduler's point of view.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum CompletionOutcome {
31    /// `execute()` ran to completion.
32    Completed,
33    /// The agent was skipped (budget pressure, missing dep, etc.).
34    Skipped,
35}
36
37/// One entry in the completion map: a level-triggered handle plus a
38/// once-set outcome flag.
39struct AgentCompletion {
40    handle: StageHandle<AgentDone>,
41    outcome: OnceLock<CompletionOutcome>,
42}
43
44impl AgentCompletion {
45    fn new() -> Self {
46        Self {
47            handle: StageHandle::<AgentDone>::default(),
48            outcome: OnceLock::new(),
49        }
50    }
51}
52
53/// Frame-scoped map of agent completion handles.
54///
55/// Created by the scheduler at the top of each frame and inserted into
56/// the per-frame [`ServiceRegistry`](crate::ServiceRegistry) overlay.
57/// Agents must NOT retain a reference past the current frame.
58pub struct AgentCompletionMap {
59    entries: HashMap<AgentId, AgentCompletion>,
60}
61
62impl AgentCompletionMap {
63    /// Builds a completion map with one empty handle per agent ID.
64    pub fn new(agent_ids: &[AgentId]) -> Self {
65        let mut entries = HashMap::with_capacity(agent_ids.len());
66        for id in agent_ids {
67            entries.insert(*id, AgentCompletion::new());
68        }
69        Self { entries }
70    }
71
72    /// Marks an agent as having finished this frame with the given outcome.
73    ///
74    /// Idempotent: subsequent calls for the same agent are no-ops.
75    /// Returns `false` if the agent ID is unknown to this map.
76    pub fn mark(&self, id: AgentId, outcome: CompletionOutcome) -> bool {
77        match self.entries.get(&id) {
78            Some(entry) => {
79                let _ = entry.outcome.set(outcome);
80                entry.handle.mark_done();
81                true
82            }
83            None => false,
84        }
85    }
86
87    /// Returns the recorded outcome for `id`, or `None` if the agent has
88    /// not yet been marked (or the ID is unknown).
89    pub fn outcome(&self, id: AgentId) -> Option<CompletionOutcome> {
90        self.entries.get(&id)?.outcome.get().copied()
91    }
92
93    /// Returns `true` if a [`mark`](Self::mark) has been recorded for `id`.
94    pub fn is_done(&self, id: AgentId) -> bool {
95        self.entries
96            .get(&id)
97            .is_some_and(|entry| entry.handle.is_done())
98    }
99
100    /// Awaits an agent's completion. Returns the outcome, or `None` if the
101    /// agent ID is unknown to this map (signals a configuration bug rather
102    /// than blocking forever).
103    pub async fn wait(&self, id: AgentId) -> Option<CompletionOutcome> {
104        let entry = self.entries.get(&id)?;
105        entry.handle.wait().await;
106        entry.outcome.get().copied()
107    }
108
109    /// Returns all agent IDs tracked by this map (for tests / introspection).
110    pub fn known_ids(&self) -> impl Iterator<Item = AgentId> + '_ {
111        self.entries.keys().copied()
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    fn rt() -> tokio::runtime::Runtime {
120        tokio::runtime::Builder::new_current_thread()
121            .enable_all()
122            .build()
123            .unwrap()
124    }
125
126    #[test]
127    fn wait_returns_immediately_after_mark() {
128        let map = AgentCompletionMap::new(&[AgentId::Renderer]);
129        map.mark(AgentId::Renderer, CompletionOutcome::Completed);
130
131        let result = rt().block_on(async { map.wait(AgentId::Renderer).await });
132        assert_eq!(result, Some(CompletionOutcome::Completed));
133    }
134
135    #[test]
136    fn wait_unblocks_when_mark_arrives_later() {
137        use std::sync::Arc;
138
139        let map = Arc::new(AgentCompletionMap::new(&[AgentId::ShadowRenderer]));
140        let map_clone = Arc::clone(&map);
141
142        let result = rt().block_on(async move {
143            let waiter = tokio::spawn(async move { map_clone.wait(AgentId::ShadowRenderer).await });
144            // Yield so the waiter has a chance to register.
145            tokio::task::yield_now().await;
146            map.mark(AgentId::ShadowRenderer, CompletionOutcome::Skipped);
147            waiter.await.unwrap()
148        });
149        assert_eq!(result, Some(CompletionOutcome::Skipped));
150    }
151
152    #[test]
153    fn wait_on_unknown_agent_returns_none() {
154        let map = AgentCompletionMap::new(&[AgentId::Renderer]);
155        let result = rt().block_on(async { map.wait(AgentId::Audio).await });
156        assert_eq!(result, None);
157    }
158
159    #[test]
160    fn outcome_distinguishes_completed_from_skipped() {
161        let map = AgentCompletionMap::new(&[AgentId::Renderer, AgentId::ShadowRenderer]);
162        map.mark(AgentId::Renderer, CompletionOutcome::Completed);
163        map.mark(AgentId::ShadowRenderer, CompletionOutcome::Skipped);
164
165        assert_eq!(
166            map.outcome(AgentId::Renderer),
167            Some(CompletionOutcome::Completed)
168        );
169        assert_eq!(
170            map.outcome(AgentId::ShadowRenderer),
171            Some(CompletionOutcome::Skipped)
172        );
173        assert_eq!(map.outcome(AgentId::Physics), None);
174    }
175
176    #[test]
177    fn mark_is_idempotent() {
178        let map = AgentCompletionMap::new(&[AgentId::Renderer]);
179        assert!(map.mark(AgentId::Renderer, CompletionOutcome::Completed));
180        // Second mark is accepted (returns true: agent is known) but the
181        // outcome stays at the first value.
182        assert!(map.mark(AgentId::Renderer, CompletionOutcome::Skipped));
183        assert_eq!(
184            map.outcome(AgentId::Renderer),
185            Some(CompletionOutcome::Completed)
186        );
187    }
188
189    #[test]
190    fn mark_unknown_agent_returns_false() {
191        let map = AgentCompletionMap::new(&[AgentId::Renderer]);
192        assert!(!map.mark(AgentId::Audio, CompletionOutcome::Completed));
193    }
194}
khora_core/agent/completion.rs

khora_core/agent/
completion.rs
