apalis_workflow/sequential/fold/
mod.rs

1use std::{marker::PhantomData, task::Context};
2
3use apalis_core::{
4    backend::{Backend, BackendConfig, TaskSinkError, WireFormatBackend, codec::Codec},
5    error::BoxDynError,
6    task::task_fn::{TaskFn, task_fn},
7    task::{
8        Task,
9        builder::TaskBuilder,
10        metadata::{Metadata, MetadataError, MetadataStore},
11        task_id::GenerateId,
12    },
13};
14use futures_util::{FutureExt, Sink, SinkExt, future::BoxFuture};
15use serde::{Deserialize, Serialize};
16use serde_json::to_value;
17use tower::Service;
18
19use crate::{
20    SteppedService,
21    sequential::{
22        context::{StepContext, WorkflowContext},
23        router::{GoTo, StepResponse, WorkflowRouter},
24        step::{Layer, Stack, Step},
25        workflow::SteppedFlow,
26    },
27};
28
29/// The fold layer that folds over a collection of items.
30#[derive(Clone, Debug)]
31pub struct Fold<F, Init> {
32    fold: F,
33    _marker: std::marker::PhantomData<Init>,
34}
35
36impl<F, Init, S> Layer<S> for Fold<F, Init>
37where
38    F: Clone,
39    Init: Clone,
40{
41    type Step = FoldStep<S, F, Init>;
42
43    fn layer(&self, step: S) -> Self::Step {
44        FoldStep {
45            inner: step,
46            fold: self.fold.clone(),
47            _marker: std::marker::PhantomData,
48        }
49    }
50}
51impl<Start, C, L, I: IntoIterator<Item = C>, B: Backend> SteppedFlow<Start, I, B, L> {
52    /// Folds over a collection of items in the workflow.
53    #[allow(clippy::type_complexity)]
54    pub fn fold<F, Output, FnArgs, Init>(
55        self,
56        fold: F,
57    ) -> SteppedFlow<Start, Output, B, Stack<Fold<TaskFn<F, (Init, C), FnArgs>, Init>, L>>
58    where
59        TaskFn<F, (Init, C), FnArgs>: Service<Task<(Init, C)>, Response = Output>,
60    {
61        self.add_step(Fold {
62            fold: task_fn(fold),
63            _marker: PhantomData,
64        })
65    }
66}
67
68/// The fold step that folds over a collection of items.
69#[derive(Clone, Debug)]
70pub struct FoldStep<S, F, Init> {
71    inner: S,
72    fold: F,
73    _marker: std::marker::PhantomData<Init>,
74}
75
76impl<S, F, Input, I: IntoIterator<Item = Input>, Init, B, Err, CodecError> Step<I, B>
77    for FoldStep<S, F, Init>
78where
79    F: Service<Task<(Init, Input)>, Response = Init> + Send + Sync + 'static + Clone,
80    S: Step<Init, B>,
81    B: Backend<Error = Err>
82        + WireFormatBackend
83        + BackendConfig
84        + Send
85        + Sync
86        + Clone
87        + Sink<Task<B::Compact>, Error = Err>
88        + Unpin
89        + 'static,
90    I: IntoIterator<Item = Input> + Send + Sync + 'static,
91    B::Codec: Codec<(Init, Vec<Input>), Error = CodecError, Compact = B::Compact>
92        + Codec<Init, Error = CodecError, Compact = B::Compact>
93        + Codec<I, Error = CodecError, Compact = B::Compact>
94        + Codec<(Init, Input), Error = CodecError, Compact = B::Compact>
95        + Send
96        + Sync
97        + Clone
98        + 'static,
99    B::Id: GenerateId + Sync + Send + 'static + Clone,
100    Init: Default + Serialize + Send + Sync + 'static,
101    Err: std::error::Error + Send + Sync + 'static,
102    CodecError: std::error::Error + Send + Sync + 'static,
103    F::Error: Into<BoxDynError> + Send + 'static,
104    F::Future: Send + 'static,
105    B::Compact: Send + 'static,
106    Input: Send + 'static,
107{
108    type Response = Init;
109    type Error = F::Error;
110    fn register(&mut self, ctx: &mut WorkflowRouter<B>) -> Result<(), BoxDynError> {
111        let svc = SteppedService::new(FoldService {
112            fold: self.fold.clone(),
113            _marker: PhantomData::<(Init, I, B)>,
114        });
115        let count = ctx.steps.len();
116        ctx.steps.insert(count, svc);
117        self.inner.register(ctx)
118    }
119}
120
121/// The fold service that handles folding over a collection of items.
122#[derive(Debug)]
123pub struct FoldService<F, Init, I, B> {
124    fold: F,
125    _marker: std::marker::PhantomData<(Init, I, B)>,
126}
127
128impl<F: Clone, Init, I, B> Clone for FoldService<F, Init, I, B> {
129    fn clone(&self) -> Self {
130        Self {
131            fold: self.fold.clone(),
132            _marker: std::marker::PhantomData,
133        }
134    }
135}
136
137impl<F, Init, I, B> FoldService<F, Init, I, B> {
138    /// Creates a new `FoldService` with the given fold function.
139    pub fn new(fold: F) -> Self {
140        Self {
141            fold,
142            _marker: std::marker::PhantomData,
143        }
144    }
145}
146
147impl<F, Init, I, B, Input, CodecError, Err> Service<Task<B::Compact>> for FoldService<F, Init, I, B>
148where
149    F: Service<Task<(Init, Input)>, Response = Init> + Send + 'static + Clone,
150    B: Backend<Error = Err>
151        + WireFormatBackend
152        + BackendConfig
153        + Clone
154        + Sink<Task<B::Compact>, Error = Err>
155        + Send
156        + Sync
157        + Unpin
158        + 'static,
159    I: IntoIterator<Item = Input> + Send + 'static,
160    B::Codec: Codec<(Init, Vec<Input>), Error = CodecError, Compact = B::Compact>
161        + Codec<Init, Error = CodecError, Compact = B::Compact>
162        + Codec<I, Error = CodecError, Compact = B::Compact>
163        + Codec<(Init, Input), Error = CodecError, Compact = B::Compact>
164        + Send
165        + Sync
166        + Clone
167        + 'static,
168    B::Id: GenerateId + Sync + Send + 'static,
169    Init: Default + Serialize + Send + 'static,
170    Err: std::error::Error + Send + Sync + 'static,
171    CodecError: std::error::Error + Send + Sync + 'static,
172    F::Error: Into<BoxDynError> + Send + 'static,
173    F::Future: Send + 'static,
174    B::Compact: Send + 'static,
175    Input: Send + 'static,
176{
177    type Response = GoTo<StepResponse>;
178    type Error = BoxDynError;
179    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
180
181    fn poll_ready(&mut self, cx: &mut Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
182        self.fold.poll_ready(cx).map_err(|e| e.into())
183    }
184
185    fn call(&mut self, task: Task<B::Compact>) -> Self::Future {
186        let state = FoldState::extract(task.metadata()).unwrap_or(FoldState::Init);
187        let mut ctx = task.data().get::<StepContext<B>>().cloned().unwrap();
188        let codec = ctx.backend.codec().clone();
189        let mut fold = self.fold.clone();
190
191        match state {
192            FoldState::Init => async move {
193                let task_id = B::Id::generate();
194                let steps: Task<I> = task.try_map_args(|arg| B::Codec::decode(&codec, &arg))?;
195                let steps = steps.args.into_iter().collect::<Vec<_>>();
196                let task = TaskBuilder::new(B::Codec::encode(&codec, &(Init::default(), steps))?)
197                    .metadata(&WorkflowContext {
198                        step_index: ctx.current_step,
199                    })
200                    .task_id(task_id.clone())
201                    .metadata(&FoldState::Collection)
202                    .build();
203                ctx.backend
204                    .send(task)
205                    .await
206                    .map_err(TaskSinkError::PushError)?;
207                Ok(GoTo::Next(StepResponse {
208                    result: to_value(Init::default())?,
209                    next_task_id: Some(task_id),
210                }))
211            }
212            .boxed(),
213            FoldState::Collection => async move {
214                let args: (Init, Vec<Input>) = B::Codec::decode(&codec, &task.args)?;
215                let (acc, items) = args;
216
217                let mut items = items.into_iter();
218                let next = items.next().unwrap();
219                let rest = items.collect::<Vec<_>>();
220                let fold_task = task.map_args(|_| (acc, next));
221                let response = fold.call(fold_task).await.map_err(|e| e.into())?;
222
223                match rest.len() {
224                    0 if ctx.has_next => {
225                        let task_id = B::Id::generate();
226                        let result = B::Codec::encode(&codec, &response)?;
227                        let next_step = TaskBuilder::new(result)
228                            .task_id(task_id.clone())
229                            .metadata(&WorkflowContext {
230                                step_index: ctx.current_step + 1,
231                            })
232                            .build();
233                        ctx.backend
234                            .send(next_step)
235                            .await
236                            .map_err(TaskSinkError::PushError)?;
237                        Ok(GoTo::Break(StepResponse {
238                            result: to_value(&response)?,
239                            next_task_id: Some(task_id),
240                        }))
241                    }
242                    0 => Ok(GoTo::Break(StepResponse {
243                        result: to_value(&response)?,
244                        next_task_id: None,
245                    })),
246                    1.. => {
247                        // Shouldn't this be limited?
248                        let task_id = B::Id::generate();
249                        let result = to_value(&response)?;
250                        let steps = TaskBuilder::new(B::Codec::encode(&codec, &(response, rest))?)
251                            .task_id(task_id.clone())
252                            .metadata(&WorkflowContext {
253                                step_index: ctx.current_step,
254                            })
255                            .metadata(&FoldState::Collection)
256                            .build();
257                        ctx.backend
258                            .send(steps)
259                            .await
260                            .map_err(TaskSinkError::PushError)?;
261                        Ok(GoTo::Next(StepResponse {
262                            result,
263                            next_task_id: Some(task_id),
264                        }))
265                    }
266                }
267            }
268            .boxed(),
269        }
270    }
271}
272
273/// The state of the fold operation
274#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
275#[non_exhaustive]
276pub enum FoldState {
277    /// Initializing state
278    Init,
279    /// Collection has started
280    Collection,
281}
282
283const FOLD_STATE_KEY: &str = "apalis_workflow.fold.state";
284
285/// An error representing an invalid [`FoldState`]
286#[derive(Debug, thiserror::Error)]
287#[non_exhaustive]
288pub enum FoldStateError {
289    /// The fold state key is missing
290    #[error("the data for key {FOLD_STATE_KEY} is missing")]
291    MissingKey,
292
293    /// Duplicate entry
294    #[error("Duplicate entry: {0}")]
295    DuplicateEntry(#[from] MetadataError),
296}
297
298impl Metadata for FoldState {
299    type Error = FoldStateError;
300
301    fn extract(map: &MetadataStore) -> Result<Self, Self::Error> {
302        let value = map.get(FOLD_STATE_KEY).ok_or(FoldStateError::MissingKey)?;
303
304        match value.as_str() {
305            "Collection" => Ok(Self::Collection),
306            _ => Ok(Self::Init),
307        }
308    }
309
310    fn inject(&self, map: &mut MetadataStore) -> Result<(), FoldStateError> {
311        let value = match self {
312            Self::Init => "Init",
313            Self::Collection => "Collection",
314        };
315        map.insert(FOLD_STATE_KEY, value)?;
316        Ok(())
317    }
318}
apalis_workflow/sequential/fold/mod.rs

apalis_workflow/sequential/fold/
mod.rs
