Skip to main content

rskit_dag/
node.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::marker::PhantomData;
4use std::pin::Pin;
5use std::sync::Arc;
6
7use rskit_errors::{AppError, AppResult, ErrorCode};
8use serde::Serialize;
9use tokio_util::sync::CancellationToken;
10
11/// A single executable node in a DAG pipeline.
12pub trait DagNode: Send + Sync + 'static {
13    /// Unique identifier for this node within the DAG.
14    fn id(&self) -> &str;
15
16    /// Execute the node with the collected outputs from upstream dependencies.
17    fn execute(
18        &self,
19        inputs: HashMap<String, serde_json::Value>,
20        cancel: CancellationToken,
21    ) -> Pin<Box<dyn Future<Output = AppResult<serde_json::Value>> + Send + '_>>;
22}
23
24type InputMapper<I> = Arc<dyn Fn(HashMap<String, serde_json::Value>) -> AppResult<I> + Send + Sync>;
25
26/// Typed adapter for a DAG node with explicit input and output contracts.
27pub struct TypedDagNode<I, O, F> {
28    id: String,
29    map_inputs: InputMapper<I>,
30    execute: F,
31    _output: PhantomData<fn() -> O>,
32}
33
34impl<I, O, F> TypedDagNode<I, O, F> {
35    /// Create a typed node from an input mapper and async operation.
36    #[must_use]
37    pub fn new(
38        id: impl Into<String>,
39        map_inputs: impl Fn(HashMap<String, serde_json::Value>) -> AppResult<I> + Send + Sync + 'static,
40        execute: F,
41    ) -> Self {
42        Self {
43            id: id.into(),
44            map_inputs: Arc::new(map_inputs),
45            execute,
46            _output: PhantomData,
47        }
48    }
49}
50
51impl<I, O, F, Fut> DagNode for TypedDagNode<I, O, F>
52where
53    I: Send + 'static,
54    O: Serialize + Send + 'static,
55    F: Fn(I, CancellationToken) -> Fut + Send + Sync + 'static,
56    Fut: Future<Output = AppResult<O>> + Send + 'static,
57{
58    fn id(&self) -> &str {
59        &self.id
60    }
61
62    fn execute(
63        &self,
64        inputs: HashMap<String, serde_json::Value>,
65        cancel: CancellationToken,
66    ) -> Pin<Box<dyn Future<Output = AppResult<serde_json::Value>> + Send + '_>> {
67        Box::pin(async move {
68            let input = (self.map_inputs)(inputs)?;
69            let output = (self.execute)(input, cancel).await?;
70            serde_json::to_value(output).map_err(|error| {
71                AppError::new(
72                    ErrorCode::Internal,
73                    format!("failed to encode typed DAG node output: {error}"),
74                )
75            })
76        })
77    }
78}