consortium/dag/
context.rs

1//! Shared context for DAG task execution.
2//!
3//! `DagContext` provides type-safe output passing between dependent tasks
4//! and arbitrary user state accessible to all tasks.
5
6use std::any::Any;
7use std::collections::HashMap;
8use std::sync::{Arc, Mutex};
9
10use crate::dag::types::TaskId;
11
12/// Shared state accessible to all tasks during DAG execution.
13///
14/// Tasks use `DagContext` to:
15/// - Read outputs from predecessor tasks (`get_output`)
16/// - Write their own outputs for dependent tasks (`set_output`)
17/// - Access user-provided state like configuration (`get_state`/`set_state`)
18#[derive(Clone)]
19pub struct DagContext {
20    outputs: Arc<Mutex<HashMap<TaskId, Box<dyn Any + Send>>>>,
21    user_state: Arc<Mutex<HashMap<String, Box<dyn Any + Send>>>>,
22}
23
24impl DagContext {
25    /// Create a new empty context.
26    pub fn new() -> Self {
27        Self {
28            outputs: Arc::new(Mutex::new(HashMap::new())),
29            user_state: Arc::new(Mutex::new(HashMap::new())),
30        }
31    }
32
33    /// Get a task's output, downcasting to the expected type.
34    pub fn get_output<T: Clone + 'static>(&self, id: &TaskId) -> Option<T> {
35        let map = self.outputs.lock().unwrap();
36        map.get(id)?.downcast_ref::<T>().cloned()
37    }
38
39    /// Store a task's output.
40    pub fn set_output<T: Send + 'static>(&self, id: TaskId, value: T) {
41        self.outputs.lock().unwrap().insert(id, Box::new(value));
42    }
43
44    /// Check if a task has produced output.
45    pub fn has_output(&self, id: &TaskId) -> bool {
46        self.outputs.lock().unwrap().contains_key(id)
47    }
48
49    /// Get user-provided state by key.
50    pub fn get_state<T: Clone + 'static>(&self, key: &str) -> Option<T> {
51        let map = self.user_state.lock().unwrap();
52        map.get(key)?.downcast_ref::<T>().cloned()
53    }
54
55    /// Set user-provided state.
56    pub fn set_state<T: Send + 'static>(&self, key: &str, value: T) {
57        self.user_state
58            .lock()
59            .unwrap()
60            .insert(key.to_string(), Box::new(value));
61    }
62}
63
64impl Default for DagContext {
65    fn default() -> Self {
66        Self::new()
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[test]
75    fn test_output_roundtrip() {
76        let ctx = DagContext::new();
77        let id = TaskId::from("build:hp01");
78        ctx.set_output(id.clone(), "/nix/store/abc".to_string());
79        let val: String = ctx.get_output(&id).unwrap();
80        assert_eq!(val, "/nix/store/abc");
81    }
82
83    #[test]
84    fn test_state_roundtrip() {
85        let ctx = DagContext::new();
86        ctx.set_state("fanout", 64usize);
87        let val: usize = ctx.get_state("fanout").unwrap();
88        assert_eq!(val, 64);
89    }
90
91    #[test]
92    fn test_missing_output() {
93        let ctx = DagContext::new();
94        let val: Option<String> = ctx.get_output(&TaskId::from("nonexistent"));
95        assert!(val.is_none());
96    }
97
98    #[test]
99    fn test_context_clone_shares_state() {
100        let ctx1 = DagContext::new();
101        let ctx2 = ctx1.clone();
102        ctx1.set_state("shared", true);
103        let val: bool = ctx2.get_state("shared").unwrap();
104        assert!(val);
105    }
106}
consortium/dag/context.rs

consortium/dag/
context.rs
