apalis_workflow/graph/
context.rs

1use std::{
2    collections::{HashMap, HashSet},
3    num::ParseIntError,
4};
5
6use apalis_core::{
7    error::BoxDynError,
8    task::{
9        metadata::{Metadata, MetadataError, MetadataStore},
10        task_id::TaskId,
11    },
12};
13use petgraph::graph::NodeIndex;
14use serde::{Deserialize, Serialize};
15
16/// Metadata stored in each node's task for workflow processing
17#[derive(Debug, Deserialize, Serialize, Default)]
18pub struct GraphFlowContext {
19    /// Previous node executed in the Graph
20    /// This is the source node that led to the current node's execution
21    pub prev_node: Option<NodeIndex>,
22    /// The current node being executed in the Graph
23    pub current_node: NodeIndex,
24
25    /// All nodes that have been completed in this execution
26    pub completed_nodes: HashSet<NodeIndex>,
27
28    /// Map of node indices to their task IDs for result lookup
29    pub node_task_ids: HashMap<NodeIndex, TaskId>,
30
31    /// Current position in the topological order
32    pub current_position: usize,
33
34    /// Whether this is the initial execution
35    pub is_initial: bool,
36
37    /// The original task ID that started this Graph execution
38    pub root_task_id: Option<TaskId>,
39}
40
41impl Clone for GraphFlowContext {
42    fn clone(&self) -> Self {
43        Self {
44            prev_node: self.prev_node,
45            current_node: self.current_node,
46            completed_nodes: self.completed_nodes.clone(),
47            node_task_ids: self.node_task_ids.clone(),
48            current_position: self.current_position,
49            is_initial: self.is_initial,
50            root_task_id: self.root_task_id.clone(),
51        }
52    }
53}
54
55impl GraphFlowContext {
56    /// Create initial context for Graph execution
57    #[must_use]
58    pub fn new(root_task_id: Option<TaskId>) -> Self {
59        Self {
60            prev_node: None,
61            current_node: NodeIndex::new(0),
62            completed_nodes: HashSet::new(),
63            node_task_ids: HashMap::new(),
64            current_position: 0,
65            is_initial: true,
66            root_task_id,
67        }
68    }
69    /// Get task IDs for dependencies of a given node
70    #[must_use]
71    pub fn get_dependency_task_ids(
72        &self,
73        dependencies: &[NodeIndex],
74    ) -> HashMap<NodeIndex, TaskId> {
75        dependencies
76            .iter()
77            .filter_map(|dep| {
78                self.node_task_ids
79                    .get(dep)
80                    .cloned()
81                    .map(|task_id| (*dep, task_id))
82            })
83            .collect()
84    }
85}
86
87const DAG_FLOW_PREV_NODE_KEY: &str = "apalis_workflow.graph.prev_node";
88
89const DAG_FLOW_CURRENT_NODE_KEY: &str = "apalis_workflow.graph.current_node";
90
91const DAG_FLOW_COMPLETED_NODES_KEY: &str = "apalis_workflow.graph.completed_nodes";
92
93const DAG_FLOW_NODE_TASK_IDS_KEY: &str = "apalis_workflow.graph.node_task_ids";
94
95const DAG_FLOW_CURRENT_POSITION_KEY: &str = "apalis_workflow.graph.current_position";
96
97const DAG_FLOW_IS_INITIAL_KEY: &str = "apalis_workflow.graph.is_initial";
98
99const DAG_FLOW_ROOT_TASK_ID_KEY: &str = "apalis_workflow.graph.root_task_id";
100
101/// An error representing an invalid [`GraphFlowContext`]
102#[derive(Debug, thiserror::Error)]
103#[non_exhaustive]
104pub enum GraphFlowContextError {
105    /// Missing current node key
106    #[error("missing key {DAG_FLOW_CURRENT_NODE_KEY}")]
107    MissingCurrentNode,
108
109    /// Missing current position key
110    #[error("missing key {DAG_FLOW_CURRENT_POSITION_KEY}")]
111    MissingCurrentPosition,
112
113    /// Could not parse a node index
114    #[error("could not parse node index")]
115    ParseNodeIndex(#[from] ParseIntError),
116
117    /// Could not parse a task_id
118    #[error("could not parse task id: {0}")]
119    ParseTaskId(BoxDynError),
120
121    /// Duplicate entry
122    #[error("Duplicate entry: {0}")]
123    DuplicateEntry(#[from] MetadataError),
124}
125
126impl Metadata for GraphFlowContext {
127    type Error = GraphFlowContextError;
128
129    fn extract(map: &MetadataStore) -> Result<Self, Self::Error> {
130        let prev_node = map
131            .get(DAG_FLOW_PREV_NODE_KEY)
132            .map(|v| v.parse::<usize>())
133            .transpose()?
134            .map(NodeIndex::new);
135
136        let current_node = map
137            .get(DAG_FLOW_CURRENT_NODE_KEY)
138            .ok_or(GraphFlowContextError::MissingCurrentNode)?
139            .parse::<usize>()?;
140
141        let completed_nodes = map
142            .get(DAG_FLOW_COMPLETED_NODES_KEY)
143            .map(|v| {
144                v.split(',')
145                    .filter(|s| !s.is_empty())
146                    .map(|s| s.parse::<usize>().map(NodeIndex::new))
147                    .collect::<Result<HashSet<_>, _>>()
148            })
149            .transpose()?
150            .unwrap_or_default();
151
152        let node_task_ids = map
153            .get(DAG_FLOW_NODE_TASK_IDS_KEY)
154            .map(|v| {
155                v.split(',')
156                    .filter(|s| !s.is_empty())
157                    .map(|s| {
158                        s.split_once('=')
159                            .ok_or(GraphFlowContextError::ParseTaskId(
160                                "Invalid delimiter".into(),
161                            ))
162                            .and_then(|(k, v)| {
163                                let node = k
164                                    .parse::<usize>()
165                                    .map(NodeIndex::new)
166                                    .map_err(GraphFlowContextError::ParseNodeIndex)?;
167                                let task_id = v
168                                    .parse::<TaskId>()
169                                    .map_err(|e| GraphFlowContextError::ParseTaskId(e.into()))?;
170                                Ok((node, task_id))
171                            })
172                    })
173                    .collect::<Result<HashMap<_, _>, _>>()
174            })
175            .transpose()?
176            .unwrap_or_default();
177
178        let current_position = map
179            .get(DAG_FLOW_CURRENT_POSITION_KEY)
180            .ok_or(GraphFlowContextError::MissingCurrentPosition)?
181            .parse::<usize>()?;
182
183        let is_initial = map
184            .get(DAG_FLOW_IS_INITIAL_KEY)
185            .map(|v| v.parse::<bool>())
186            .transpose()
187            .unwrap_or(None)
188            .unwrap_or(true);
189
190        let root_task_id = map
191            .get(DAG_FLOW_ROOT_TASK_ID_KEY)
192            .map(|v| v.parse::<TaskId>())
193            .transpose()
194            .map_err(|e| GraphFlowContextError::ParseTaskId(e.into()))?;
195
196        Ok(Self {
197            prev_node,
198            current_node: NodeIndex::new(current_node),
199            completed_nodes,
200            node_task_ids,
201            current_position,
202            is_initial,
203            root_task_id,
204        })
205    }
206
207    fn inject(&self, map: &mut MetadataStore) -> Result<(), GraphFlowContextError> {
208        if let Some(prev_node) = self.prev_node {
209            map.insert(DAG_FLOW_PREV_NODE_KEY, prev_node.index().to_string())?;
210        }
211
212        map.insert(
213            DAG_FLOW_CURRENT_NODE_KEY,
214            self.current_node.index().to_string(),
215        )?;
216
217        let completed_nodes = self
218            .completed_nodes
219            .iter()
220            .map(|n| n.index())
221            .collect::<Vec<_>>();
222
223        map.insert(
224            DAG_FLOW_COMPLETED_NODES_KEY,
225            completed_nodes
226                .iter()
227                .map(ToString::to_string)
228                .collect::<Vec<_>>()
229                .join(","),
230        )?;
231
232        let node_task_ids = self
233            .node_task_ids
234            .iter()
235            .map(|(k, v)| format!("{}={v}", k.index()))
236            .collect::<Vec<_>>()
237            .join(",");
238
239        map.insert(DAG_FLOW_NODE_TASK_IDS_KEY, node_task_ids)?;
240
241        map.insert(
242            DAG_FLOW_CURRENT_POSITION_KEY,
243            self.current_position.to_string(),
244        )?;
245
246        map.insert(DAG_FLOW_IS_INITIAL_KEY, self.is_initial.to_string())
247            .expect("A value already exists");
248
249        if let Some(root_task_id) = &self.root_task_id {
250            map.insert(DAG_FLOW_ROOT_TASK_ID_KEY, root_task_id.to_string())?;
251        }
252        Ok(())
253    }
254}
apalis_workflow/graph/context.rs

apalis_workflow/graph/
context.rs
