apalis_core/worker/
context.rs

1//! Worker context and task tracking.
2//!
3//! [`WorkerContext`] is responsible for managing
4//! the execution lifecycle of a worker, tracking tasks, handling shutdown, and emitting
5//! lifecycle events.
6//!
7//! ## Lifecycle
8//! A `WorkerContext` goes through distinct phases of operation:
9//!
10//! - **Pending**: Created via [`WorkerContext::new`] and must be explicitly started.
11//! - **Running**: Activated by calling [`WorkerContext::start`]. The worker becomes ready to accept and track tasks.
12//! - **Paused**: Temporarily halted via [`WorkerContext::pause`]. New tasks are blocked from execution.
13//! - **Resumed**: Brought back to `Running` using [`WorkerContext::resume`].
14//! - **Stopped**: Finalized via [`WorkerContext::stop`]. The worker shuts down gracefully, allowing tracked tasks to complete.
15//!
16//! The `WorkerContext` itself implements [`Future`], and can be `.await`ed — it resolves
17//! once the worker is shut down and all tasks have completed.
18//!
19//! ## Task Management
20//! Asynchronous tasks can be tracked which ensures:
21//! - Task count is incremented before execution and decremented on completion
22//! - Shutdown is automatically triggered once all tasks are done
23//!
24//! Use [`task_count`](WorkerContext::task_count) and [`has_pending_tasks`](WorkerContext::has_pending_tasks) to inspect
25//! ongoing task state.
26//!
27//! ## Shutdown Semantics
28//! The worker is considered shutting down if:
29//! - `stop()` has been called
30//! - A shutdown signal (if configured) has been triggered
31//!
32//! Once shutdown begins, no new tasks should be accepted. Internally, a stored [`Waker`] is
33//! used to drive progress toward shutdown completion.
34//!
35//! ## Event Handling
36//! Worker lifecycle events (e.g., `Start`, `Stop`) are emitted automatically
37//! and custom ones can be emitted using [`WorkerContext::emit`].
38//!
39//! ## Request Integration
40//! `WorkerContext` implements [`FromRequest`] so it can be extracted automatically in request
41//! handlers when using a compatible framework or service layer.
42//!
43//! ## Types
44//! - [`WorkerContext`] — shared state container for a worker
45use std::{
46    fmt::{self},
47    sync::{
48        Arc, Mutex,
49        atomic::{AtomicBool, AtomicUsize, Ordering},
50    },
51    task::{Context, Waker},
52    time::{Duration, Instant},
53};
54
55use dashmap::DashSet;
56
57use crate::{
58    error::{WorkerError, WorkerStateError},
59    monitor::shutdown::Shutdown,
60    task::from_request::FromRequest,
61    task::{
62        ExecutionContext, Task,
63        context::{TaskContext, TaskStateError},
64        data::MissingDataError,
65    },
66    worker::{
67        event::{Event, EventListener, RawEventListener},
68        lifecycle::TaskLifecycleError,
69        state::{InnerWorkerState, WorkerState},
70    },
71};
72
73/// Utility for managing a worker's context
74///
75/// A worker context is created for each worker thread and is responsible for managing
76/// the worker's state, task tracking, and event handling.
77///
78///  **Tip**: All fields are wrapped inside [`Arc`] so it should be cheap to clone
79#[derive(Clone)]
80pub struct WorkerContext {
81    pub(crate) name: Arc<String>,
82    /// The waker used to wake the worker when tasks complete or shutdown is triggered.
83    waker: Arc<Mutex<Option<Waker>>>,
84    state: Arc<WorkerState>,
85    pub(crate) shutdown: Option<Shutdown>,
86    event_handler: EventListener,
87    pub(super) is_ready: Arc<AtomicBool>,
88    service: &'static str,
89    tasks: Arc<DashSet<TaskContext>>,
90    instant: Instant,
91    restarts: Arc<AtomicUsize>,
92}
93
94impl fmt::Debug for WorkerContext {
95    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
96        f.debug_struct("WorkerContext")
97            .field("shutdown", &["Shutdown handle"])
98            .field("task_count", &self.task_count())
99            .field("state", &self.state.load(Ordering::SeqCst))
100            .field("service", &self.service)
101            .field("is_ready", &self.is_ready)
102            .field("tasks", &"[..]")
103            .field("elapsed", &self.instant.elapsed())
104            .field("restarts", &self.restarts)
105            .finish()
106    }
107}
108
109impl WorkerContext {
110    /// Create a new worker context
111    #[must_use]
112    pub fn new(name: &str) -> Self {
113        Self {
114            name: Arc::new(name.to_owned()),
115            service: "Unspecified",
116            waker: Default::default(),
117            state: Default::default(),
118            shutdown: Default::default(),
119            event_handler: Arc::new(Box::new(|_, _| {
120                // noop
121            })),
122            is_ready: Default::default(),
123            tasks: Default::default(),
124            instant: Instant::now(),
125            restarts: Default::default(),
126        }
127    }
128
129    /// Get the worker id
130    #[must_use]
131    pub fn name(&self) -> &str {
132        &self.name
133    }
134
135    /// Start running the worker
136    pub fn start(&mut self) -> Result<(), WorkerError> {
137        let current_state = self.state.load(Ordering::SeqCst);
138        if current_state != InnerWorkerState::Pending {
139            return Err(WorkerError::StateError(WorkerStateError::AlreadyStarted));
140        }
141        self.state
142            .store(InnerWorkerState::Running, Ordering::SeqCst);
143        self.is_ready.store(false, Ordering::SeqCst);
144        info!("Worker {} started", self.name());
145        self.wake();
146        Ok(())
147    }
148
149    /// Restart running the worker
150    pub(crate) fn restart(&self) -> Result<(), WorkerError> {
151        self.state
152            .store(InnerWorkerState::Pending, Ordering::SeqCst);
153        self.is_ready.store(false, Ordering::SeqCst);
154        self.cleanup();
155        self.restarts.fetch_add(1, Ordering::SeqCst);
156        info!("Worker {} restarted", self.name());
157        self.wake();
158        Ok(())
159    }
160
161    /// Pauses a worker, preventing any new jobs from being polled
162    pub fn pause(&self) -> Result<(), WorkerError> {
163        if !self.is_running() {
164            return Err(WorkerError::StateError(WorkerStateError::NotRunning));
165        }
166        self.state.store(InnerWorkerState::Paused, Ordering::SeqCst);
167        info!("Worker {} paused", self.name());
168        Ok(())
169    }
170
171    /// Resume a worker that is paused
172    pub fn resume(&self) -> Result<(), WorkerError> {
173        if !self.is_paused() {
174            return Err(WorkerError::StateError(WorkerStateError::NotPaused));
175        }
176        if self.is_shutting_down() {
177            return Err(WorkerError::StateError(WorkerStateError::ShuttingDown));
178        }
179        self.state
180            .store(InnerWorkerState::Running, Ordering::SeqCst);
181        self.wake();
182        info!("Worker {} resumed", self.name());
183        Ok(())
184    }
185
186    /// Calling this function triggers shutting down the worker while waiting for any tasks to complete
187    pub fn stop(&self) -> Result<(), WorkerError> {
188        let current_state = self.state.load(Ordering::SeqCst);
189        if current_state == InnerWorkerState::Pending {
190            return Err(WorkerError::StateError(WorkerStateError::NotStarted));
191        }
192        self.state
193            .store(InnerWorkerState::Stopped, Ordering::SeqCst);
194        self.wake();
195        self.emit_ref(&Event::Stop);
196        info!("Worker {} stopped", self.name());
197        Ok(())
198    }
199
200    /// Checks if the worker is ready to consume new tasks
201    #[must_use]
202    pub fn is_ready(&self) -> bool {
203        self.is_running() && !self.is_shutting_down() && self.is_ready.load(Ordering::SeqCst)
204    }
205
206    /// Get the `type_name` of the service used
207    ///
208    /// ## Example
209    /// ```ignore
210    /// async fn send_email(email: Email) {}
211    ///
212    /// // Might be something like:
213    /// TaskFn<send_email, Email, ()>>
214    /// ```
215    #[must_use]
216    pub fn get_service(&self) -> &str {
217        self.service
218    }
219
220    pub(super) fn bind_service<T>(&mut self) {
221        let service = std::any::type_name::<T>();
222
223        const RULES: &[(&str, &str, &str)] = &[
224            (
225                "Retry<",
226                "Trace<",
227                "`retries()` must be before `enable_tracing()`; traces produced will provide invalid attempt information",
228            ),
229            (
230                "Retry<",
231                "PrometheusService<",
232                "`retries()` must be before `prometheus()`; metrics inside retries will be invalid",
233            ),
234            (
235                "Retry<",
236                "Timeout<",
237                "`retries()` must be before `timeout()`; timeouts will be applied to the total retry process",
238            ),
239            (
240                "Timeout<",
241                "Trace<",
242                "`timeout()` should be before `enable_tracing()`; otherwise timeout failures may not be reflected correctly in traces",
243            ),
244            (
245                "Timeout<",
246                "PrometheusService<",
247                "`timeout()` should be before `prometheus()`; otherwise timeout failures may not be reflected correctly in metrics",
248            ),
249            (
250                "ConcurrencyLimit<",
251                "Retry<",
252                "`concurrency()` should generally be before `retries()`; otherwise each retry may consume a separate concurrency slot",
253            ),
254            (
255                "ConcurrencyLimit<",
256                "Timeout<",
257                "`concurrency()` should generally be before `timeout()`; otherwise queued requests may consume timeout duration",
258            ),
259            (
260                "RateLimit<",
261                "Retry<",
262                "`rate_limit()` should generally be before `retries()`; otherwise retries may consume rate-limit capacity",
263            ),
264            (
265                "LoadShed<",
266                "Retry<",
267                "`load_shed()` should generally be before `retries()`; otherwise retries may repeatedly encounter load-shed failures",
268            ),
269            (
270                "Retry<",
271                "CatchPanicService<",
272                "`catch_panic()` should generally be after `retries()` if panics are intended to participate in retry handling",
273            ),
274        ];
275
276        for &(outer, inner, message) in RULES {
277            if let (Some(a), Some(b)) = (service.find(outer), service.find(inner)) {
278                if a > b {
279                    warn!("{message}");
280                }
281            }
282        }
283
284        self.service = service;
285    }
286
287    /// Checks whether the worker is running
288    #[must_use]
289    pub fn is_running(&self) -> bool {
290        self.state.load(Ordering::SeqCst) == InnerWorkerState::Running
291    }
292
293    /// Checks whether the worker is pending
294    #[must_use]
295    pub fn is_pending(&self) -> bool {
296        self.state.load(Ordering::SeqCst) == InnerWorkerState::Pending
297    }
298
299    /// Checks whether the worker is paused
300    #[must_use]
301    pub fn is_paused(&self) -> bool {
302        self.state.load(Ordering::SeqCst) == InnerWorkerState::Paused
303    }
304
305    /// Checks whether the worker has been stopped
306    #[must_use]
307    pub fn is_stopped(&self) -> bool {
308        self.state.load(Ordering::SeqCst) == InnerWorkerState::Stopped || self.is_terminated()
309    }
310
311    /// Checks whether the worker is terminated
312    ///
313    #[must_use]
314    pub fn is_terminated(&self) -> bool {
315        self.state.load(Ordering::SeqCst) == InnerWorkerState::Terminated
316    }
317
318    /// Checks the current futures in the worker domain
319    /// This include futures spawned via `worker.track`
320    #[must_use]
321    pub fn task_count(&self) -> usize {
322        self.tasks.len()
323    }
324
325    /// Checks whether the worker has pending tasks
326    #[must_use]
327    pub fn has_pending_tasks(&self) -> bool {
328        self.task_count() > 0
329    }
330
331    /// Is the shutdown token called
332    #[must_use]
333    pub fn is_shutting_down(&self) -> bool {
334        self.is_stopped() || self.shutdown.as_ref().is_some_and(|s| s.is_shutting_down())
335    }
336
337    /// Get the current worker state
338    #[must_use]
339    pub fn state(&self) -> &str {
340        self.state.as_str()
341    }
342
343    /// Emits an event to the worker's event handler
344    pub(crate) fn emit_event(&self, event: &Event) {
345        self.emit_ref(event);
346    }
347
348    /// Emits a [`Event::Custom`] to the worker's event handler
349    pub fn emit<T: Send + Sync + 'static>(&self, data: T) {
350        self.emit_ref(&Event::custom(data));
351    }
352
353    fn emit_ref(&self, event: &Event) {
354        let handler = self.event_handler.as_ref();
355        handler(self, event);
356    }
357
358    /// Calls a method to signify a heartbeat with the worker
359    pub fn heartbeat(&self, cx: &mut Context<'_>) {
360        self.register_waker(cx);
361        self.emit_ref(&Event::HeartBeat);
362        // Mark the worker as ready/alive.
363        self.is_ready.store(true, Ordering::SeqCst);
364    }
365
366    /// Wraps the event listener with a new function
367    pub(crate) fn add_listener<F: Fn(&Self, &Event) + Send + Sync + 'static>(&mut self, f: F) {
368        let cur = self.event_handler.clone();
369        let new: RawEventListener = Box::new(move |ctx, ev| {
370            f(ctx, ev);
371            cur(ctx, ev);
372        });
373        self.event_handler = Arc::new(new);
374    }
375
376    /// Register the current waker for the worker
377    ///
378    /// This is used to wake the worker when tasks complete or shutdown is triggered.
379    pub(crate) fn register_waker(&self, cx: &Context<'_>) {
380        if let Ok(mut guard) = self.waker.lock() {
381            if guard
382                .as_ref()
383                .is_none_or(|stored| !stored.will_wake(cx.waker()))
384            {
385                *guard = Some(cx.waker().clone());
386            }
387        }
388    }
389
390    pub(crate) fn wake(&self) {
391        if let Ok(waker) = self.waker.lock() {
392            if let Some(waker) = &*waker {
393                waker.wake_by_ref();
394            }
395        }
396    }
397
398    /// Register the [`ExecutionContext`] to get the [`TaskContext`]
399    pub(super) fn register_task(
400        &self,
401        ctx: &Arc<ExecutionContext>,
402    ) -> Result<TaskContext, TaskLifecycleError> {
403        let task_id = ctx
404            .task_id()
405            .ok_or(TaskLifecycleError::MissingTaskId)?
406            .to_string();
407
408        let tasks = &self.tasks;
409
410        if tasks.contains(task_id.as_str()) {
411            return Err(TaskLifecycleError::Duplicate);
412        }
413
414        let token = TaskContext::new(ctx);
415        tasks.insert(token.clone());
416        Ok(token)
417    }
418
419    /// Cancel a specific task, if it's tracked.
420    pub fn cancel_task(&self, context: &TaskContext) -> Result<(), TaskStateError> {
421        if let Some(token) = self.tasks.get(context) {
422            token.cancel()
423        } else {
424            Err(TaskStateError::TaskNotFound)
425        }
426    }
427
428    /// Remove the token once the task completes, to avoid unbounded growth.
429    pub(super) fn remove_task(&self, ctx: &TaskContext) -> bool {
430        let task_id = ctx.task_id();
431        self.tasks.remove(task_id).is_some()
432    }
433
434    /// Extracts the [`TaskContext`] from the [`WorkerContext`]
435    pub(crate) fn get_task_context(
436        &self,
437        ctx: &Arc<ExecutionContext>,
438    ) -> Result<TaskContext, MissingDataError> {
439        self.get_task(ctx.task_id().unwrap().to_string().as_str())
440    }
441
442    /// Get the task context for a task attached to a worker
443    pub fn get_task(&self, task_id: &str) -> Result<TaskContext, MissingDataError> {
444        let tasks = &self.tasks;
445        Ok(tasks
446            .get(task_id)
447            .ok_or(MissingDataError::NotFound("TaskContext".to_owned()))?
448            .clone())
449    }
450
451    /// Remove any completed tasks
452    pub fn cleanup(&self) {
453        let tasks = &self.tasks;
454        tasks.retain(|token| !(token.is_completed() && token.is_empty()));
455    }
456
457    /// Get the context of each running task
458    #[must_use]
459    pub fn tasks(&self) -> Vec<TaskContext> {
460        self.tasks.iter().map(|s| s.clone()).collect()
461    }
462
463    /// Returns the amount of time elapsed since this worker started.
464    #[must_use]
465    pub fn elapsed(&self) -> Duration {
466        self.instant.elapsed()
467    }
468
469    /// Returns the number of times the worker has been restated:
470    ///
471    /// See also [`Monitor::should_restart`]
472    ///
473    /// [`Monitor::should_restart`]: crate::monitor::Monitor::should_restart
474    #[must_use]
475    pub fn restarts(&self) -> usize {
476        self.restarts.load(Ordering::SeqCst)
477    }
478
479    /// This forces a shutting down worker to exit.
480    pub fn kill(&mut self) -> Result<(), WorkerError> {
481        if !self.is_shutting_down() {
482            return Err(WorkerError::StateError(WorkerStateError::InvalidState(
483                "Worker is not shutting down".to_owned(),
484            )));
485        }
486        if self.task_count() != 0 {
487            self.tasks()
488                .into_iter()
489                .map(|a| a.cancel())
490                .collect::<Result<Vec<_>, _>>()
491                .map_err(|e| {
492                    WorkerError::StateError(WorkerStateError::InvalidState(e.to_string()))
493                })?;
494        }
495        self.state
496            .store(InnerWorkerState::Terminated, Ordering::SeqCst);
497        self.wake();
498        Ok(())
499    }
500}
501
502impl From<&str> for WorkerContext {
503    fn from(name: &str) -> Self {
504        Self::new(name)
505    }
506}
507
508impl From<String> for WorkerContext {
509    fn from(name: String) -> Self {
510        Self::new(&name)
511    }
512}
513
514impl From<&Self> for WorkerContext {
515    fn from(context: &Self) -> Self {
516        context.clone()
517    }
518}
519
520impl<Args: Sync> FromRequest<Task<Args>> for WorkerContext {
521    type Error = MissingDataError;
522    async fn from_request(task: &Task<Args>) -> Result<Self, Self::Error> {
523        task.data().get_checked().cloned()
524    }
525}
526
527impl Drop for WorkerContext {
528    fn drop(&mut self) {
529        if Arc::strong_count(&self.state) > 1 {
530            // There are still other references to this context, so we shouldn't log a warning.
531            return;
532        }
533        if self.is_running() && self.has_pending_tasks() {
534            error!(
535                "Worker '{}' is being dropped while running with `{}` tasks. Consider calling stop() before dropping.",
536                self.name(),
537                self.task_count()
538            );
539        }
540    }
541}
542
543#[cfg(test)]
544mod tests {
545
546    use futures_util::FutureExt;
547
548    use crate::{
549        backend::memory::MemoryStorage, error::BoxDynError, worker::builder::WorkerBuilder,
550    };
551    use std::time::Duration;
552
553    use super::*;
554
555    #[tokio::test]
556    async fn test_worker_state_transitions() {
557        let backend = MemoryStorage::<u32>::new();
558
559        let ctx = WorkerContext::new("test-worker");
560
561        let worker = WorkerBuilder::new(&ctx)
562            .backend(backend)
563            .build(|_task: u32| async { Ok::<_, BoxDynError>(()) });
564
565        let worker_handle = tokio::spawn(async move { worker.run().boxed().await });
566        tokio::time::sleep(Duration::from_millis(50)).await;
567
568        // Initial state: worker should be running
569        assert!(ctx.is_running());
570        assert!(!ctx.is_shutting_down());
571        assert!(!ctx.is_stopped());
572
573        // Pause the worker
574        ctx.pause().unwrap();
575        assert!(ctx.is_paused());
576        assert!(
577            !ctx.is_shutting_down(),
578            "Paused worker should NOT be considered shutting down"
579        );
580
581        // Resume the worker
582        ctx.resume().unwrap();
583        assert!(ctx.is_running());
584        assert!(!ctx.is_paused());
585
586        // Stop the worker
587        ctx.stop().unwrap();
588        assert!(ctx.is_stopped());
589        assert!(ctx.is_shutting_down());
590
591        // Try to resume a stopped worker (should fail with NotPaused error since state is Stopped)
592        assert!(
593            matches!(
594                ctx.resume(),
595                Err(WorkerError::StateError(WorkerStateError::NotPaused))
596            ),
597            "Resuming a stopped worker should fail with NotPaused error"
598        );
599
600        worker_handle.await.unwrap().unwrap();
601    }
602}
apalis_core/worker/context.rs

apalis_core/worker/
context.rs
