apalis_core/worker/ext/parallelize/
mod.rs

1//! Worker extension for parallel execution and spawning child tasks.
2//!
3//! This module provides functionality to parallelize task execution within a worker
4//! and spawn child tasks. It includes the [`ParallelizeLayer`] middleware for parallel execution
5//! and the [`ParallelizeService`] service.
6//!
7//! # Example
8//! ```rust,no_run
9//! # use apalis_core::worker::builder::WorkerBuilder;
10//! # use apalis_core::backend::memory::MemoryStorage;
11//! # use apalis_core::worker::ext::parallelize::ParallelizeExt;
12//! # use apalis_core::worker::ext::event_listener::EventListenerExt;
13//! #[tokio::main]
14//! async fn main() {
15//!     async fn task(task: u32) {
16//!         println!("Processing task: {task}");
17//!     }
18//!     let in_memory = MemoryStorage::new();
19//!     let worker = WorkerBuilder::new("rango-tango")
20//!         .backend(in_memory)
21//!         .parallelize(tokio::spawn)
22//!         .on_event(|wrk, ev| {
23//!             println!("CTX {:?}, On Event = {:?}", wrk.name(), ev);
24//!         })
25//!         .build(task);
26//!     worker.run().await.unwrap();
27//! }
28//! ```
29use std::future::ready;
30
31use futures_core::future::BoxFuture;
32use futures_util::FutureExt;
33use futures_util::TryFutureExt;
34use tower_layer::{Layer, Stack};
35use tower_service::Service;
36
37use crate::{backend::Backend, error::BoxDynError, task::Task, worker::builder::WorkerBuilder};
38
39/// Worker extension for parallel execution
40pub trait ParallelizeExt<Args, Source, Middleware, Executor>: Sized {
41    /// Register the executor for parallel task execution.
42    fn parallelize(
43        self,
44        f: Executor,
45    ) -> WorkerBuilder<Args, Source, Stack<ParallelizeLayer<Executor>, Middleware>>;
46}
47
48/// Middleware for emitting events
49#[derive(Debug, Clone, Default)]
50pub struct ParallelizeLayer<Executor> {
51    executor: Executor,
52}
53
54impl<Executor> ParallelizeLayer<Executor> {
55    /// Create a new event listener layer
56    pub fn new(executor: Executor) -> Self {
57        Self { executor }
58    }
59}
60
61impl<S, Executor: Clone> Layer<S> for ParallelizeLayer<Executor> {
62    type Service = ParallelizeService<S, Executor>;
63
64    fn layer(&self, service: S) -> Self::Service {
65        ParallelizeService {
66            service,
67            executor: self.executor.clone(),
68        }
69    }
70}
71
72/// Service to allow spawning tasks with utils like `tokio::spawn`
73#[derive(Debug, Clone)]
74pub struct ParallelizeService<S, Executor> {
75    service: S,
76    executor: Executor,
77}
78
79impl<S, Args, Fut, T, Executor, ExecErr> Service<Task<Args>> for ParallelizeService<S, Executor>
80where
81    S: Service<Task<Args>, Future = Fut>,
82    Executor: Fn(Fut) -> T + Send + 'static,
83    Fut: Future<Output = Result<S::Response, S::Error>> + Send + 'static,
84    T: Future<Output = Result<Result<S::Response, S::Error>, ExecErr>> + Send + 'static,
85    S::Error: Into<BoxDynError> + Send + 'static,
86    ExecErr: Into<BoxDynError>,
87    S::Response: Send + 'static,
88{
89    type Response = S::Response;
90    type Error = BoxDynError;
91    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
92
93    fn poll_ready(
94        &mut self,
95        cx: &mut std::task::Context<'_>,
96    ) -> std::task::Poll<Result<(), Self::Error>> {
97        self.service.poll_ready(cx).map_err(|e| e.into())
98    }
99
100    fn call(&mut self, request: Task<Args>) -> Self::Future {
101        (self.executor)(self.service.call(request))
102            .map_err(|e| e.into())
103            .and_then(|s| ready(s.map_err(|e| e.into())))
104            .boxed()
105    }
106}
107
108impl<Args, P, M, Executor> ParallelizeExt<Args, P, M, Executor> for WorkerBuilder<Args, P, M>
109where
110    P: Backend,
111    M: Layer<ParallelizeLayer<Executor>>,
112{
113    fn parallelize(
114        self,
115        f: Executor,
116    ) -> WorkerBuilder<Args, P, Stack<ParallelizeLayer<Executor>, M>> {
117        self.layer(ParallelizeLayer::new(f))
118    }
119}
apalis_core/worker/ext/parallelize/mod.rs

apalis_core/worker/ext/parallelize/
mod.rs
