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

apalis_core/worker/
context.rs
