apalis_core/worker/

context.rs

//! Worker context and task tracking.
//!
//! [`WorkerContext`] is responsible for managing
//! the execution lifecycle of a worker, tracking tasks, handling shutdown, and emitting
//! lifecycle events. It also provides [`Tracked`] for wrapping and monitoring asynchronous
//! tasks within the worker domain.
//!
//! ## Lifecycle
//! A `WorkerContext` goes through distinct phases of operation:
//!
//! - **Pending**: Created via [`WorkerContext::new`] and must be explicitly started.
//! - **Running**: Activated by calling [`WorkerContext::start`]. The worker becomes ready to accept and track tasks.
//! - **Paused**: Temporarily halted via [`WorkerContext::pause`]. New tasks are blocked from execution.
//! - **Resumed**: Brought back to `Running` using [`WorkerContext::resume`].
//! - **Stopped**: Finalized via [`WorkerContext::stop`]. The worker shuts down gracefully, allowing tracked tasks to complete.
//!
//! The `WorkerContext` itself implements [`Future`], and can be `.await`ed — it resolves
//! once the worker is shut down and all tasks have completed.
//!
//! ## Task Management
//! Asynchronous tasks can be tracked using [`WorkerContext::track`], which wraps a future
//! in a [`Tracked`] type. This ensures:
//! - Task count is incremented before execution and decremented on completion
//! - Shutdown is automatically triggered once all tasks are done
//!
//! Use [`task_count`](WorkerContext::task_count) and [`has_pending_tasks`](WorkerContext::has_pending_tasks) to inspect
//! ongoing task state.
//!
//! ## Shutdown Semantics
//! The worker is considered shutting down if:
//! - `stop()` has been called
//! - A shutdown signal (if configured) has been triggered
//!
//! Once shutdown begins, no new tasks should be accepted. Internally, a stored [`Waker`] is
//! used to drive progress toward shutdown completion.
//!
//! ## Event Handling
//! Worker lifecycle events (e.g., `Start`, `Stop`) can be emitted using [`WorkerContext::emit`].
//! Custom handlers can be registered via [`WorkerContext::wrap_listener`] to hook into these transitions.
//!
//! ## Request Integration
//! `WorkerContext` implements [`FromRequest`] so it can be extracted automatically in request
//! handlers when using a compatible framework or service layer.
//!
//! ## Types
//! - [`WorkerContext`] — shared state container for a worker
//! - [`Tracked`] — future wrapper for task lifecycle tracking
use std::{
    any::type_name,
    fmt,
    future::Future,
    pin::Pin,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicUsize, Ordering},
    },
    task::{Context, Poll, Waker},
};

use crate::{
    error::{WorkerError, WorkerStateError},
    monitor::shutdown::Shutdown,
    task::{Task, data::MissingDataError},
    task_fn::FromRequest,
    worker::{
        event::{Event, EventListener, RawEventListener},
        state::{InnerWorkerState, WorkerState},
    },
};

/// Utility for managing a worker's context
///
/// A worker context is created for each worker thread and is responsible for managing
/// the worker's state, task tracking, and event handling.
///
///  **Tip**: All fields are wrapped inside [`Arc`] so it should be cheap to clone
#[derive(Clone)]
pub struct WorkerContext {
    pub(super) name: Arc<String>,
    task_count: Arc<AtomicUsize>,
    waker: Arc<Mutex<Option<Waker>>>,
    state: Arc<WorkerState>,
    pub(crate) shutdown: Option<Shutdown>,
    event_handler: EventListener,
    pub(super) is_ready: Arc<AtomicBool>,
    pub(super) service: &'static str,
}

impl fmt::Debug for WorkerContext {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("WorkerContext")
            .field("shutdown", &["Shutdown handle"])
            .field("task_count", &self.task_count)
            .field("state", &self.state.load(Ordering::SeqCst))
            .field("service", &self.service)
            .field("is_ready", &self.is_ready)
            .finish()
    }
}

/// A future tracked by the worker
#[pin_project::pin_project(PinnedDrop)]
#[derive(Debug)]
pub struct Tracked<F> {
    ctx: WorkerContext,
    #[pin]
    task: F,
}

impl<F: Future> Future for Tracked<F> {
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
        let this = self.project();

        match this.task.poll(cx) {
            res @ Poll::Ready(_) => res,
            Poll::Pending => Poll::Pending,
        }
    }
}

#[pin_project::pinned_drop]
impl<F> PinnedDrop for Tracked<F> {
    fn drop(self: Pin<&mut Self>) {
        self.ctx.end_task();
    }
}

impl WorkerContext {
    /// Create a new worker context
    #[must_use]
    pub fn new<S>(name: &str) -> Self {
        Self {
            name: Arc::new(name.to_owned()),
            service: type_name::<S>(),
            task_count: Default::default(),
            waker: Default::default(),
            state: Default::default(),
            shutdown: Default::default(),
            event_handler: Arc::new(Box::new(|_, _| {
                // noop
            })),
            is_ready: Default::default(),
        }
    }

    /// Get the worker id
    #[must_use]
    pub fn name(&self) -> &String {
        &self.name
    }

    /// Start running the worker
    pub fn start(&mut self) -> Result<(), WorkerError> {
        let current_state = self.state.load(Ordering::SeqCst);
        if current_state != InnerWorkerState::Pending {
            return Err(WorkerError::StateError(WorkerStateError::AlreadyStarted));
        }
        self.state
            .store(InnerWorkerState::Running, Ordering::SeqCst);
        self.is_ready.store(false, Ordering::SeqCst);
        self.emit(&Event::Start);
        info!("Worker {} started", self.name());
        Ok(())
    }

    /// Restart running the worker
    pub fn restart(&mut self) -> Result<(), WorkerError> {
        self.state
            .store(InnerWorkerState::Pending, Ordering::SeqCst);
        self.is_ready.store(false, Ordering::SeqCst);
        info!("Worker {} restarted", self.name());
        Ok(())
    }

    /// Start a task that is tracked by the worker
    pub fn track<F: Future>(&self, task: F) -> Tracked<F> {
        self.start_task();
        Tracked {
            ctx: self.clone(),
            task,
        }
    }
    /// Pauses a worker, preventing any new jobs from being polled
    pub fn pause(&self) -> Result<(), WorkerError> {
        if !self.is_running() {
            return Err(WorkerError::StateError(WorkerStateError::NotRunning));
        }
        self.state.store(InnerWorkerState::Paused, Ordering::SeqCst);
        info!("Worker {} paused", self.name());
        Ok(())
    }

    /// Resume a worker that is paused
    pub fn resume(&self) -> Result<(), WorkerError> {
        if !self.is_paused() {
            return Err(WorkerError::StateError(WorkerStateError::NotPaused));
        }
        if self.is_shutting_down() {
            return Err(WorkerError::StateError(WorkerStateError::ShuttingDown));
        }
        self.state
            .store(InnerWorkerState::Running, Ordering::SeqCst);
        self.wake();
        info!("Worker {} resumed", self.name());
        Ok(())
    }

    /// Calling this function triggers shutting down the worker while waiting for any tasks to complete
    pub fn stop(&self) -> Result<(), WorkerError> {
        let current_state = self.state.load(Ordering::SeqCst);
        if current_state == InnerWorkerState::Pending {
            return Err(WorkerError::StateError(WorkerStateError::NotStarted));
        }
        self.state
            .store(InnerWorkerState::Stopped, Ordering::SeqCst);
        self.wake();
        self.emit_ref(&Event::Stop);
        info!("Worker {} stopped", self.name());
        Ok(())
    }

    /// Checks if the worker is ready to consume new tasks
    #[must_use]
    pub fn is_ready(&self) -> bool {
        self.is_running() && !self.is_shutting_down() && self.is_ready.load(Ordering::SeqCst)
    }

    /// Get the type of service
    #[must_use]
    pub fn get_service(&self) -> &str {
        self.service
    }

    /// Checks whether the worker is running
    #[must_use]
    pub fn is_running(&self) -> bool {
        self.state.load(Ordering::SeqCst) == InnerWorkerState::Running
    }

    /// Checks whether the worker is pending
    #[must_use]
    pub fn is_pending(&self) -> bool {
        self.state.load(Ordering::SeqCst) == InnerWorkerState::Pending
    }

    /// Checks whether the worker is paused
    #[must_use]
    pub fn is_paused(&self) -> bool {
        self.state.load(Ordering::SeqCst) == InnerWorkerState::Paused
    }

    /// Checks whether the worker has been stopped
    #[must_use]
    pub fn is_stopped(&self) -> bool {
        self.state.load(Ordering::SeqCst) == InnerWorkerState::Stopped
    }

    /// Checks the current futures in the worker domain
    /// This include futures spawned via `worker.track`
    #[must_use]
    pub fn task_count(&self) -> usize {
        self.task_count.load(Ordering::Relaxed)
    }

    /// Checks whether the worker has pending tasks
    #[must_use]
    pub fn has_pending_tasks(&self) -> bool {
        self.task_count.load(Ordering::Relaxed) > 0
    }

    /// Is the shutdown token called
    #[must_use]
    pub fn is_shutting_down(&self) -> bool {
        self.is_stopped() || self.shutdown.as_ref().is_some_and(|s| s.is_shutting_down())
    }

    /// Allows workers to emit events
    pub fn emit(&mut self, event: &Event) {
        self.emit_ref(event);
    }

    fn emit_ref(&self, event: &Event) {
        let handler = self.event_handler.as_ref();
        handler(self, event);
    }

    /// Wraps the event listener with a new function
    pub fn wrap_listener<F: Fn(&Self, &Event) + Send + Sync + 'static>(&mut self, f: F) {
        let cur = self.event_handler.clone();
        let new: RawEventListener = Box::new(move |ctx, ev| {
            f(ctx, ev);
            cur(ctx, ev);
        });
        self.event_handler = Arc::new(new);
    }

    pub(crate) fn add_waker(&self, cx: &Context<'_>) {
        if let Ok(mut waker_guard) = self.waker.lock() {
            if waker_guard
                .as_ref()
                .is_none_or(|stored_waker| !stored_waker.will_wake(cx.waker()))
            {
                *waker_guard = Some(cx.waker().clone());
            }
        }
    }

    /// Checks if the stored waker matches the current one.
    fn has_recent_waker(&self, cx: &Context<'_>) -> bool {
        if let Ok(waker_guard) = self.waker.lock() {
            if let Some(stored_waker) = &*waker_guard {
                return stored_waker.will_wake(cx.waker());
            }
        }
        false
    }

    fn start_task(&self) {
        self.task_count.fetch_add(1, Ordering::Relaxed);
    }

    fn end_task(&self) {
        if self.task_count.fetch_sub(1, Ordering::Relaxed) == 1 {
            self.wake();
        }
    }

    pub(crate) fn wake(&self) {
        if let Ok(waker) = self.waker.lock() {
            if let Some(waker) = &*waker {
                waker.wake_by_ref();
            }
        }
    }
}

impl Future for WorkerContext {
    type Output = Result<(), WorkerError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let task_count = self.task_count.load(Ordering::Relaxed);
        let state = self.state.load(Ordering::SeqCst);
        if state == InnerWorkerState::Pending {
            return Poll::Ready(Err(WorkerError::StateError(WorkerStateError::NotStarted)));
        }
        if self.is_shutting_down() && task_count == 0 {
            Poll::Ready(Ok(()))
        } else {
            if !self.has_recent_waker(cx) {
                self.add_waker(cx);
            }
            Poll::Pending
        }
    }
}

impl<Args: Sync, Ctx: Sync, IdType: Sync + Send> FromRequest<Task<Args, Ctx, IdType>>
    for WorkerContext
{
    type Error = MissingDataError;
    async fn from_request(task: &Task<Args, Ctx, IdType>) -> Result<Self, Self::Error> {
        task.parts.data.get_checked().cloned()
    }
}

impl Drop for WorkerContext {
    fn drop(&mut self) {
        if Arc::strong_count(&self.state) > 1 {
            // There are still other references to this context, so we shouldn't log a warning.
            return;
        }
        if self.is_running() && self.has_pending_tasks() {
            error!(
                "Worker '{}' is being dropped while running with `{}` tasks. Consider calling stop() before dropping.",
                self.name(),
                self.task_count()
            );
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        backend::memory::MemoryStorage, error::BoxDynError, worker::builder::WorkerBuilder,
    };
    use std::time::Duration;

    use super::*;

    #[tokio::test]
    async fn test_worker_state_transitions() {
        let backend = MemoryStorage::<u32>::new();

        let worker = WorkerBuilder::new("test-worker")
            .backend(backend)
            .build(|_task: u32| async { Ok::<_, BoxDynError>(()) });

        let mut ctx = WorkerContext::new::<()>("test-worker");
        let ctx_handle = ctx.clone();

        let worker_handle = tokio::spawn(async move { worker.run_with_ctx(&mut ctx).await });
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Initial state: worker should be running
        assert!(ctx_handle.is_running());
        assert!(!ctx_handle.is_shutting_down());
        assert!(!ctx_handle.is_stopped());

        // Pause the worker
        ctx_handle.pause().unwrap();
        assert!(ctx_handle.is_paused());
        assert!(
            !ctx_handle.is_shutting_down(),
            "Paused worker should NOT be considered shutting down"
        );

        // Resume the worker
        ctx_handle.resume().unwrap();
        assert!(ctx_handle.is_running());
        assert!(!ctx_handle.is_paused());

        // Stop the worker
        ctx_handle.stop().unwrap();
        assert!(ctx_handle.is_stopped());
        assert!(ctx_handle.is_shutting_down());

        // Try to resume a stopped worker (should fail with NotPaused error since state is Stopped)
        assert!(
            matches!(
                ctx_handle.resume(),
                Err(WorkerError::StateError(WorkerStateError::NotPaused))
            ),
            "Resuming a stopped worker should fail with NotPaused error"
        );

        worker_handle.await.unwrap().unwrap();
    }
}