apalis_workflow/sequential/
service.rs

1use apalis_core::{
2    backend::{Backend, BackendConfig, TaskSinkError, WireFormatBackend, codec::Codec},
3    error::BoxDynError,
4    task::{Task, builder::TaskBuilder, metadata::Metadata, task_id::GenerateId},
5};
6use futures_util::{FutureExt, Sink, SinkExt, future::BoxFuture};
7use serde::Serialize;
8use serde_json::to_value;
9use std::{
10    collections::{HashMap, VecDeque},
11    marker::PhantomData,
12    task::{Context, Poll},
13};
14use tower::Service;
15
16use crate::{
17    SteppedService,
18    sequential::{
19        context::{StepContext, WorkflowContext},
20        router::{GoTo, StepResponse},
21    },
22};
23
24/// The main workflow service that orchestrates the execution of workflow steps.
25#[derive(Debug, Clone)]
26pub struct WorkflowService<B, Input, Output>
27where
28    B: Backend + WireFormatBackend,
29{
30    services: HashMap<usize, SteppedService<B::Compact>>,
31    not_ready: VecDeque<usize>,
32    backend: B,
33    _marker: PhantomData<(Input, Output)>,
34}
35impl<B, Input, Output> WorkflowService<B, Input, Output>
36where
37    B: Backend + WireFormatBackend,
38{
39    /// Creates a new `WorkflowService` with the given services and backend.
40    pub fn new(services: HashMap<usize, SteppedService<B::Compact>>, backend: B) -> Self {
41        Self {
42            services,
43            not_ready: VecDeque::new(),
44            backend,
45            _marker: PhantomData,
46        }
47    }
48}
49
50impl<B, Err, Input, Output> Service<Task<B::Compact>> for WorkflowService<B, Input, Output>
51where
52    B: Sink<Task<B::Compact>, Error = Err>
53        + Unpin
54        + WireFormatBackend
55        + BackendConfig<Args = Input>
56        + Clone
57        + Send
58        + Sync
59        + 'static
60        + Backend<Error = Err>,
61    B::Compact: Send + 'static,
62    Err: std::error::Error + Send + Sync + 'static,
63    B::Id: GenerateId + Send + 'static,
64{
65    type Response = GoTo<StepResponse>;
66    type Error = BoxDynError;
67    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
68
69    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
70        loop {
71            // must wait for *all* services to be ready.
72            // this will cause head-of-line blocking unless the underlying services are always ready.
73            if self.not_ready.is_empty() {
74                return Poll::Ready(Ok(()));
75            } else {
76                if self
77                    .services
78                    .get_mut(&self.not_ready[0])
79                    .unwrap()
80                    .poll_ready(cx)?
81                    .is_pending()
82                {
83                    return Poll::Pending;
84                }
85
86                self.not_ready.pop_front();
87            }
88        }
89    }
90
91    fn call(&mut self, mut req: Task<B::Compact>) -> Self::Future {
92        assert!(
93            self.not_ready.is_empty(),
94            "Workflow must wait for all services to be ready. Did you forget to call poll_ready()?"
95        );
96        let meta = WorkflowContext::extract(req.metadata()).unwrap_or_default();
97        let idx = meta.step_index;
98
99        let has_next = self.services.contains_key(&(idx + 1));
100        let step_ctx: StepContext<B> = StepContext::new(self.backend.clone(), idx, has_next);
101
102        let svc = self
103            .services
104            .get_mut(&idx)
105            .expect("Attempted to run a step that doesn't exist");
106
107        req.inject_data(step_ctx);
108
109        self.not_ready.push_back(idx);
110        svc.call(req).boxed()
111    }
112}
113
114/// Handle the result of a workflow step, scheduling the next step if necessary
115pub async fn handle_step_result<N, Compact, B, Err>(
116    ctx: &mut StepContext<B>,
117    result: GoTo<N>,
118) -> Result<GoTo<StepResponse>, TaskSinkError<Err>>
119where
120    B: Sink<Task<Compact>, Error = Err>
121        + Backend<Error = Err>
122        + WireFormatBackend<Compact = Compact>
123        + BackendConfig
124        + Send
125        + Unpin,
126    Err: Into<BoxDynError>,
127    B::Codec: Codec<N, Compact = Compact> + Clone,
128    <B::Codec as Codec<N>>::Error: Into<BoxDynError>,
129    N: Serialize,
130    Compact: 'static,
131    N: 'static,
132    B::Id: GenerateId + Send + 'static,
133{
134    let codec = ctx.backend.codec().clone();
135    match result {
136        GoTo::Next(next) if ctx.has_next => {
137            let task_id = B::Id::generate();
138            let task = TaskBuilder::new(
139                B::Codec::encode(&codec, &next).map_err(|e| TaskSinkError::CodecError(e.into()))?,
140            )
141            .task_id(task_id.clone())
142            .metadata(&WorkflowContext {
143                step_index: ctx.current_step + 1,
144            })
145            .build();
146            ctx.backend.send(task).await?;
147            Ok(GoTo::Next(StepResponse {
148                result: to_value(&next).map_err(|e| TaskSinkError::CodecError(e.into()))?,
149                next_task_id: Some(task_id),
150            }))
151        }
152        GoTo::DelayFor(delay, next) if ctx.has_next => {
153            let task_id = B::Id::generate();
154
155            let task = TaskBuilder::new(
156                B::Codec::encode(&codec, &next).map_err(|e| TaskSinkError::CodecError(e.into()))?,
157            )
158            .run_after(delay)
159            .task_id(task_id.clone())
160            .metadata(&WorkflowContext {
161                step_index: ctx.current_step + 1,
162            })
163            .build();
164            ctx.backend.send(task).await?;
165            Ok(GoTo::DelayFor(
166                delay,
167                StepResponse {
168                    result: to_value(&next).map_err(|e| TaskSinkError::CodecError(e.into()))?,
169                    next_task_id: Some(task_id),
170                },
171            ))
172        }
173        #[allow(clippy::match_same_arms)]
174        GoTo::Done => Ok(GoTo::Done),
175        GoTo::Break(res) => Ok(GoTo::Break(StepResponse {
176            result: to_value(&res).map_err(|e| TaskSinkError::CodecError(e.into()))?,
177            next_task_id: None,
178        })),
179        _ => Ok(GoTo::Done),
180    }
181}
apalis_workflow/sequential/service.rs

apalis_workflow/sequential/
service.rs
