apalis_workflow/graph/
node.rs

1use apalis_core::backend::codec::Codec;
2use apalis_core::backend::{Backend, BackendConfig, WireFormatBackend};
3use apalis_core::error::BoxDynError;
4use apalis_core::task::Task;
5use serde::Serialize;
6use serde::de::DeserializeOwned;
7use std::future::Future;
8use std::pin::Pin;
9use std::task::{Context, Poll};
10use tower::Service;
11
12use crate::graph::NodeInput;
13use crate::graph::decode::GraphCodec;
14
15/// A service that wraps another service to handle encoding and decoding
16/// of task inputs and outputs using the backend's codec.
17pub struct GraphNodeService<S, B, Input>
18where
19    S: Service<Task<Input>>,
20    B: Backend,
21{
22    inner: S,
23    _phantom: std::marker::PhantomData<(B, Input)>,
24}
25
26impl<S, B, Input> std::fmt::Debug for GraphNodeService<S, B, Input>
27where
28    S: Service<Task<Input>>,
29    B: Backend,
30{
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        f.debug_struct("NodeService")
33            .field("inner", &"<service>")
34            .field("_phantom", &std::any::type_name::<(B, Input)>())
35            .finish()
36    }
37}
38
39impl<S, B, Input> Clone for GraphNodeService<S, B, Input>
40where
41    S: Service<Task<Input>> + Clone,
42    B: Backend,
43{
44    fn clone(&self) -> Self {
45        Self {
46            inner: self.inner.clone(),
47            _phantom: std::marker::PhantomData,
48        }
49    }
50}
51
52impl<S, B, Input> GraphNodeService<S, B, Input>
53where
54    S: Service<Task<Input>>,
55    B: Backend,
56{
57    /// Creates a new `NodeService` wrapping the provided service.
58    pub fn new(inner: S) -> Self {
59        Self {
60            inner,
61            _phantom: std::marker::PhantomData,
62        }
63    }
64}
65
66impl<S, B, Input, CdcErr> Service<Task<NodeInput<B::Compact>>> for GraphNodeService<S, B, Input>
67where
68    S: Service<Task<Input>>,
69    S::Error: Into<BoxDynError>,
70    B: Backend + WireFormatBackend + BackendConfig + Send + Sync + 'static,
71    B::Codec: Codec<Input, Compact = B::Compact, Error = CdcErr>
72        + Codec<S::Response, Compact = B::Compact, Error = CdcErr>
73        + Send
74        + Sync
75        + Clone,
76    Input: GraphCodec<B, Error = CdcErr> + DeserializeOwned,
77    CdcErr: Into<BoxDynError> + Send + 'static,
78    S::Future: Send + 'static,
79    S::Response: Serialize,
80{
81    // Here we return both the encoded version and the json version
82    // We push next nodes with the compact version
83    // We store results in json - hence why we need the json version
84    type Response = (B::Compact, serde_json::Value);
85    type Error = BoxDynError;
86    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
87
88    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
89        self.inner.poll_ready(cx).map_err(|e| e.into())
90    }
91
92    fn call(&mut self, req: Task<NodeInput<B::Compact>>) -> Self::Future {
93        let codec = req
94            .data()
95            .get::<B::Codec>()
96            .cloned()
97            .expect("GraphExecutor should be injected");
98
99        let req = req.try_map_args(|args| match args {
100            NodeInput::Single(args) => match Input::decode(&args, &codec) {
101                Ok(decoded) => Ok(decoded),
102                Err(e) => Err(CdcErr::into(e)),
103            },
104            NodeInput::FanIn(fan_in) => {
105                let value = serde_json::Value::Array(fan_in);
106                let result: Input = serde_json::from_value(value)?;
107                Ok(result)
108            }
109        });
110
111        let decoded_req = match req {
112            Ok(req) => req,
113            Err(e) => {
114                return Box::pin(async move { Err(e) });
115            }
116        };
117
118        let fut = self.inner.call(decoded_req);
119
120        Box::pin(async move {
121            let response = fut.await.map_err(|e| e.into())?;
122            let compact = B::Codec::encode(&codec, &response).map_err(|e| e.into())?;
123            let res = serde_json::to_value(&response)?;
124            Ok((compact, res))
125        })
126    }
127}
apalis_workflow/graph/node.rs

apalis_workflow/graph/
node.rs
