Skip to main content

apalis_workflow/dag/
node.rs

1use apalis_core::backend::BackendExt;
2use apalis_core::backend::codec::Codec;
3use apalis_core::error::BoxDynError;
4use apalis_core::task::Task;
5use std::future::Future;
6use std::pin::Pin;
7use std::task::{Context, Poll};
8use tower::Service;
9
10use crate::dag::decode::DagCodec;
11
12/// A service that wraps another service to handle encoding and decoding
13/// of task inputs and outputs using the backend's codec.
14pub struct NodeService<S, B, Input>
15where
16    S: Service<Task<Input, B::Context, B::IdType>>,
17    B: BackendExt,
18{
19    inner: S,
20    _phantom: std::marker::PhantomData<(B, Input)>,
21}
22
23impl<S, B, Input> std::fmt::Debug for NodeService<S, B, Input>
24where
25    S: Service<Task<Input, B::Context, B::IdType>>,
26    B: BackendExt,
27{
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        f.debug_struct("NodeService")
30            .field("inner", &"<service>")
31            .field("_phantom", &std::any::type_name::<(B, Input)>())
32            .finish()
33    }
34}
35
36impl<S, B, Input> Clone for NodeService<S, B, Input>
37where
38    S: Service<Task<Input, B::Context, B::IdType>> + Clone,
39    B: BackendExt,
40{
41    fn clone(&self) -> Self {
42        Self {
43            inner: self.inner.clone(),
44            _phantom: std::marker::PhantomData,
45        }
46    }
47}
48
49impl<S, B, Input> NodeService<S, B, Input>
50where
51    S: Service<Task<Input, B::Context, B::IdType>>,
52    B: BackendExt,
53{
54    /// Creates a new `NodeService` wrapping the provided service.
55    pub fn new(inner: S) -> Self {
56        Self {
57            inner,
58            _phantom: std::marker::PhantomData,
59        }
60    }
61}
62
63impl<S, B, Input, CdcErr> Service<Task<B::Compact, B::Context, B::IdType>>
64    for NodeService<S, B, Input>
65where
66    S: Service<Task<Input, B::Context, B::IdType>>,
67    S::Error: Into<BoxDynError>,
68    B: BackendExt,
69    B::Codec: Codec<Input, Compact = B::Compact, Error = CdcErr>
70        + Codec<S::Response, Compact = B::Compact, Error = CdcErr>,
71    Input: DagCodec<B, Error = CdcErr>,
72    CdcErr: Into<BoxDynError> + Send + 'static,
73    S::Future: Send + 'static,
74{
75    type Response = B::Compact;
76    type Error = BoxDynError;
77    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
78
79    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
80        self.inner.poll_ready(cx).map_err(|e| e.into())
81    }
82
83    fn call(&mut self, req: Task<B::Compact, B::Context, B::IdType>) -> Self::Future {
84        let decoded_req = match Input::decode(&req.args) {
85            Ok(decoded) => req.map(|_| decoded),
86            Err(e) => {
87                return Box::pin(async move { Err(CdcErr::into(e)) });
88            }
89        };
90
91        let fut = self.inner.call(decoded_req);
92
93        Box::pin(async move {
94            let response = fut.await.map_err(|e| e.into())?;
95            B::Codec::encode(&response).map_err(|e| e.into())
96        })
97    }
98}