Skip to main content

apalis_workflow/sequential/delay/
mod.rs

1use std::time::Duration;
2
3use apalis_core::{
4    backend::{Backend, BackendConfig, WireFormatBackend, codec::Codec},
5    error::BoxDynError,
6    task::{Task, builder::TaskBuilder, task_id::GenerateId},
7};
8use futures_util::SinkExt;
9use futures_util::{FutureExt, Sink, future::BoxFuture};
10use serde_json::to_value;
11use tower::Service;
12
13use crate::{
14    SteppedFlow, SteppedService,
15    sequential::{
16        context::{StepContext, WorkflowContext},
17        router::{GoTo, StepResponse, WorkflowRouter},
18        step::{Layer, Stack, Step},
19    },
20};
21
22/// Layer that delays execution by a specified duration
23#[derive(Clone, Debug)]
24pub struct DelayFor {
25    duration: Duration,
26}
27
28impl<S> Layer<S> for DelayFor
29where
30    S: Clone,
31{
32    type Step = DelayForStep<S>;
33
34    fn layer(&self, step: S) -> Self::Step {
35        DelayForStep {
36            inner: step,
37            duration: self.duration,
38        }
39    }
40}
41
42/// Step that delays execution by a specified duration
43#[derive(Clone, Debug)]
44pub struct DelayForStep<S> {
45    inner: S,
46    duration: Duration,
47}
48
49impl<Input, B, S, Err> Step<Input, B> for DelayForStep<S>
50where
51    B::Id: GenerateId + Send + Sync + 'static,
52    B::Compact: Send + 'static,
53    B: Sink<Task<B::Compact>, Error = Err>
54        + WireFormatBackend
55        + BackendConfig
56        + Unpin
57        + Send
58        + Sync
59        + Clone
60        + 'static,
61    Err: std::error::Error + Send + Sync + 'static,
62    S: Clone + Send + Sync + 'static,
63    S::Response: Send + 'static,
64    B::Codec: Codec<Duration, Compact = B::Compact>
65        + Codec<Input, Compact = B::Compact>
66        + Send
67        + Clone
68        + 'static,
69    <B::Codec as Codec<Duration>>::Error: Into<BoxDynError>,
70    Input: Send + Sync + 'static,
71    <B::Codec as Codec<Input>>::Error: Into<BoxDynError>,
72    B: Backend,
73    S: Step<Input, B>,
74{
75    type Response = Input;
76    type Error = BoxDynError;
77    fn register(&mut self, ctx: &mut WorkflowRouter<B>) -> Result<(), BoxDynError> {
78        let duration = self.duration;
79        let svc = SteppedService::new(DelayWithStep {
80            f: Box::new(move |_| duration),
81            inner: self.inner.clone(),
82            _marker: std::marker::PhantomData,
83        });
84        let count = ctx.steps.len();
85        ctx.steps.insert(count, svc);
86        self.inner.register(ctx)
87    }
88}
89
90/// Step that delays execution by a specified duration
91#[derive(Clone, Debug)]
92pub struct DelayWith<F, B, Input> {
93    f: F,
94    _marker: std::marker::PhantomData<(B, Input)>,
95}
96
97impl<S, F: Clone, B, I> Layer<S> for DelayWith<F, B, I> {
98    type Step = DelayWithStep<S, F, B, I>;
99
100    fn layer(&self, step: S) -> Self::Step {
101        DelayWithStep {
102            f: self.f.clone(),
103            inner: step,
104            _marker: std::marker::PhantomData,
105        }
106    }
107}
108
109/// Step that delays execution by a specified duration
110#[derive(Debug)]
111pub struct DelayWithStep<S, F, B, Input> {
112    f: F,
113    inner: S,
114    _marker: std::marker::PhantomData<(B, Input)>,
115}
116
117impl<S: Clone, F: Clone, B, Input> Clone for DelayWithStep<S, F, B, Input> {
118    fn clone(&self) -> Self {
119        Self {
120            f: self.f.clone(),
121            inner: self.inner.clone(),
122            _marker: std::marker::PhantomData,
123        }
124    }
125}
126
127impl<Input, F, B, S, Err> Step<Input, B> for DelayWithStep<S, F, B, Input>
128where
129    F: FnMut(Task<Input>) -> Duration + Send + Sync + 'static + Clone,
130    B::Id: GenerateId + Sync + Send + 'static,
131    B::Compact: Send + 'static,
132    B: Sink<Task<B::Compact>, Error = Err>
133        + BackendConfig
134        + WireFormatBackend
135        + Unpin
136        + Send
137        + Sync
138        + Clone
139        + 'static,
140    Err: std::error::Error + Send + Sync + 'static,
141    S: Step<Input, B> + Clone + Send + Sync + 'static,
142    S::Response: Send + 'static,
143    B::Codec: Codec<Duration, Compact = B::Compact>
144        + Codec<Input, Compact = B::Compact>
145        + Send
146        + Clone
147        + 'static,
148    <B::Codec as Codec<Duration>>::Error: Into<BoxDynError>,
149    Input: Send + Sync + 'static,
150    <B::Codec as Codec<Input>>::Error: Into<BoxDynError>,
151    B: Backend,
152{
153    type Response = Input;
154    type Error = BoxDynError;
155    fn register(&mut self, ctx: &mut WorkflowRouter<B>) -> Result<(), BoxDynError> {
156        let svc = SteppedService::new(Self {
157            f: self.f.clone(),
158            inner: self.inner.clone(),
159            _marker: std::marker::PhantomData,
160        });
161        let count = ctx.steps.len();
162        ctx.steps.insert(count, svc);
163        self.inner.register(ctx)
164    }
165}
166
167impl<S, F, B: Backend + Send + Sync + 'static + Clone, Input, Err> Service<Task<B::Compact>>
168    for DelayWithStep<S, F, B, Input>
169where
170    F: FnMut(Task<Input>) -> Duration + Send + 'static + Clone,
171    S: Step<Input, B> + Send + 'static,
172    S::Response: Send + 'static,
173    B::Id: GenerateId + Sync + Send + 'static,
174    B::Compact: Send + 'static,
175    B: Sink<Task<B::Compact>, Error = Err>
176        + WireFormatBackend
177        + BackendConfig
178        + Unpin
179        + Send
180        + Sync,
181    Err: std::error::Error + Send + Sync + 'static,
182    B::Codec: Codec<Duration, Compact = B::Compact>
183        + Codec<Input, Compact = B::Compact>
184        + Send
185        + Clone
186        + 'static,
187    <B::Codec as Codec<Duration>>::Error: Into<BoxDynError>,
188    <B::Codec as Codec<Input>>::Error: Into<BoxDynError>,
189{
190    type Response = GoTo<StepResponse>;
191    type Error = BoxDynError;
192    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
193
194    fn poll_ready(
195        &mut self,
196        _cx: &mut std::task::Context<'_>,
197    ) -> std::task::Poll<Result<(), Self::Error>> {
198        std::task::Poll::Ready(Ok(()))
199    }
200
201    fn call(&mut self, req: Task<B::Compact>) -> Self::Future {
202        let mut step_context: StepContext<B> = req.data().get().cloned().unwrap();
203        let mut f = self.f.clone();
204        let codec = step_context.backend.codec().clone();
205        let task_id = B::Id::generate();
206        async move {
207            let decoded: Input = B::Codec::decode(&codec, &req.args)
208                .map_err(|e: <B::Codec as Codec<Input>>::Error| e.into())?;
209            let (args, ctx) = req.take();
210            let delay_duration = f(Task::new_with_ctx(decoded, ctx));
211
212            let task = TaskBuilder::new(args)
213                .task_id(task_id.clone())
214                .metadata(&WorkflowContext {
215                    step_index: step_context.current_step + 1,
216                })
217                .run_after(delay_duration)
218                .build();
219            step_context
220                .backend
221                .send(task)
222                .await
223                .map_err(|e| BoxDynError::from(e))?;
224            Ok(GoTo::DelayFor(
225                delay_duration,
226                StepResponse {
227                    result: to_value(delay_duration)?,
228                    next_task_id: Some(task_id),
229                },
230            ))
231        }
232        .boxed()
233    }
234}
235
236impl<Start, Cur, B, L> SteppedFlow<Start, Cur, B, L> {
237    /// Delay the workflow by a fixed duration
238    pub fn delay_for(self, delay: Duration) -> SteppedFlow<Start, Cur, B, Stack<DelayFor, L>> {
239        self.add_step(DelayFor { duration: delay })
240    }
241}
242impl<Start, Cur, B, L> SteppedFlow<Start, Cur, B, L> {
243    /// Delay the workflow by a duration determined by a function
244    #[allow(clippy::type_complexity)]
245    pub fn delay_with<F>(self, f: F) -> SteppedFlow<Start, Cur, B, Stack<DelayWith<F, B, Cur>, L>>
246    where
247        F: FnMut(Task<Cur>) -> Duration + Send + 'static,
248    {
249        self.add_step(DelayWith {
250            f,
251            _marker: std::marker::PhantomData,
252        })
253    }
254}