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
11pub trait DagNode: Send + Sync + 'static {
13 fn id(&self) -> &str;
15
16 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
26pub 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 #[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}