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#[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#[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#[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#[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 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 #[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}