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    state_override: Option<Arc<StateOverride>>,
23}
24
25struct StateOverride {
26    key: String,
27    value: Mutex<Box<dyn Any + Send>>,
28    parent: Option<Arc<StateOverride>>,
29}
30
31impl DagContext {
32    /// Create a new empty context.
33    pub fn new() -> Self {
34        Self {
35            outputs: Arc::new(Mutex::new(HashMap::new())),
36            user_state: Arc::new(Mutex::new(HashMap::new())),
37            state_override: None,
38        }
39    }
40
41    /// Get a task's output, downcasting to the expected type.
42    pub fn get_output<T: Clone + 'static>(&self, id: &TaskId) -> Option<T> {
43        let map = self.outputs.lock().unwrap();
44        map.get(id)?.downcast_ref::<T>().cloned()
45    }
46
47    /// Store a task's output.
48    pub fn set_output<T: Send + 'static>(&self, id: TaskId, value: T) {
49        self.outputs.lock().unwrap().insert(id, Box::new(value));
50    }
51
52    /// Check if a task has produced output.
53    pub fn has_output(&self, id: &TaskId) -> bool {
54        self.outputs.lock().unwrap().contains_key(id)
55    }
56
57    /// Get user-provided state by key.
58    pub fn get_state<T: Clone + 'static>(&self, key: &str) -> Option<T> {
59        let mut current = self.state_override.as_ref();
60        while let Some(state) = current {
61            if state.key == key {
62                return state.value.lock().unwrap().downcast_ref::<T>().cloned();
63            }
64            current = state.parent.as_ref();
65        }
66        let map = self.user_state.lock().unwrap();
67        map.get(key)?.downcast_ref::<T>().cloned()
68    }
69
70    /// Override one read in a task-local view; outputs and ordinary state writes
71    /// remain shared. Existing overrides are inherited without mutating siblings.
72    pub fn with_state_override<T: Send + 'static>(&self, key: &str, value: T) -> Self {
73        Self {
74            outputs: Arc::clone(&self.outputs),
75            user_state: Arc::clone(&self.user_state),
76            state_override: Some(Arc::new(StateOverride {
77                key: key.into(),
78                value: Mutex::new(Box::new(value)),
79                parent: self.state_override.clone(),
80            })),
81        }
82    }
83
84    /// Set user-provided state.
85    pub fn set_state<T: Send + 'static>(&self, key: &str, value: T) {
86        self.user_state
87            .lock()
88            .unwrap()
89            .insert(key.to_string(), Box::new(value));
90    }
91}
92
93impl Default for DagContext {
94    fn default() -> Self {
95        Self::new()
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn test_output_roundtrip() {
105        let ctx = DagContext::new();
106        let id = TaskId::from("build:hp01");
107        ctx.set_output(id.clone(), "/nix/store/abc".to_string());
108        let val: String = ctx.get_output(&id).unwrap();
109        assert_eq!(val, "/nix/store/abc");
110    }
111
112    #[test]
113    fn test_state_roundtrip() {
114        let ctx = DagContext::new();
115        ctx.set_state("fanout", 64usize);
116        let val: usize = ctx.get_state("fanout").unwrap();
117        assert_eq!(val, 64);
118    }
119
120    #[test]
121    fn test_missing_output() {
122        let ctx = DagContext::new();
123        let val: Option<String> = ctx.get_output(&TaskId::from("nonexistent"));
124        assert!(val.is_none());
125    }
126
127    #[test]
128    fn test_context_clone_shares_state() {
129        let ctx1 = DagContext::new();
130        let ctx2 = ctx1.clone();
131        ctx1.set_state("shared", true);
132        let val: bool = ctx2.get_state("shared").unwrap();
133        assert!(val);
134    }
135
136    #[test]
137    fn task_overrides_isolate_reads_and_share_pipeline_outputs() {
138        let base = DagContext::new();
139        base.set_state("executor", "base".to_owned());
140        let first = base.with_state_override("executor", "first".to_owned());
141        let second = base.with_state_override("executor", "second".to_owned());
142        let nested = first.with_state_override("other", 5_u8);
143        assert_eq!(nested.get_state::<String>("executor").unwrap(), "first");
144        assert_eq!(second.get_state::<String>("executor").unwrap(), "second");
145        assert_eq!(base.get_state::<String>("executor").unwrap(), "base");
146        assert_eq!(nested.get_state::<u8>("other"), Some(5));
147        assert_eq!(first.get_state::<u8>("other"), None);
148        std::thread::scope(|scope| {
149            for (index, view) in [first, second].into_iter().enumerate() {
150                scope.spawn(move || {
151                    view.set_output(TaskId::from(format!("out:{index}")), index);
152                    view.set_state(&format!("state:{index}"), index);
153                });
154            }
155        });
156        for index in 0..2 {
157            assert_eq!(
158                base.get_output::<usize>(&TaskId::from(format!("out:{index}"))),
159                Some(index)
160            );
161            assert_eq!(
162                base.get_state::<usize>(&format!("state:{index}")),
163                Some(index)
164            );
165        }
166    }
167}
consortium/dag/context.rs

consortium/dag/
context.rs
