palladium_actor/

pool.rs

use crate::errors::{AskError, SendError};
use crate::message::Message;
use crate::path::AddrHash;
use crate::policy::{RestartPolicy, ShutdownPolicy};
use crate::stable::StableAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;

// ── PoolConfig ─────────────────────────────────────────────────────────────────

/// Configuration for a [`WorkerPool`].
///
/// Pass to [`EngineHandle::spawn_worker_pool`] to create and supervise N
/// identical worker actors under the `/user` supervisor.
#[derive(Debug, Clone)]
pub struct PoolConfig {
    /// Number of worker actors to spawn.  Must be ≥ 1.
    pub size: usize,
    /// Prefix used to derive per-worker actor names.
    ///
    /// Workers are spawned with paths `{parent}/{name_prefix}-0`,
    /// `{parent}/{name_prefix}-1`, … under whichever supervisor the pool
    /// builder targets.
    pub name_prefix: String,
    /// Per-worker mailbox capacity.  Default: `1024`.
    pub mailbox_capacity: usize,
    /// Restart policy applied to each worker.  Default: [`RestartPolicy::Permanent`].
    pub restart: RestartPolicy,
    /// Shutdown policy applied to each worker.  Default: 5-second graceful stop.
    pub shutdown: ShutdownPolicy,
    /// Timeout for `ask()` calls routed through this pool.  Default: 30 seconds.
    pub ask_timeout: Duration,
}

impl Default for PoolConfig {
    fn default() -> Self {
        Self {
            size: 4,
            name_prefix: "worker".to_string(),
            mailbox_capacity: 1024,
            restart: RestartPolicy::Permanent,
            shutdown: ShutdownPolicy::Timeout(Duration::from_secs(5)),
            ask_timeout: Duration::from_secs(30),
        }
    }
}

// ── WorkerPool ─────────────────────────────────────────────────────────────────

/// A supervision-aware, round-robin routing handle for a pool of identical
/// worker actors.
///
/// Created by [`EngineHandle::spawn_worker_pool`].  Workers are spawned and
/// supervised by the `/user` supervisor; this struct is a thin routing facade
/// over N [`StableAddr<M>`]s that remain valid across individual worker
/// restarts.
///
/// ## Routing semantics
///
/// - **Round-robin**: each call to [`send`](WorkerPool::send) or
///   [`ask`](WorkerPool::ask) selects the next worker in rotation.
/// - **Shared counter across clones**: all clones share the rotation counter
///   (`Arc<AtomicUsize>`), so concurrent senders interleave rather than
///   duplicate.
/// - **Backpressure**: [`SendError::MailboxFull`] is returned immediately when
///   the selected worker's mailbox is at capacity.
/// - **Restart window**: [`SendError::Unroutable`] is returned during the brief
///   interval between a worker crash and the supervisor refreshing its
///   [`StableAddr`].  The window is transient; callers should yield and retry.
///
/// ## In-flight message loss
///
/// Messages that were already in a crashed worker's mailbox at crash time are
/// **lost** (best-effort semantics).  Full at-least-once delivery requires
/// application-level acknowledgement — outside the scope of this primitive.
pub struct WorkerPool<M: Message> {
    workers: Vec<StableAddr<M>>,
    next: Arc<AtomicUsize>,
}

impl<M: Message> Clone for WorkerPool<M> {
    fn clone(&self) -> Self {
        Self {
            workers: self.workers.clone(),
            next: Arc::clone(&self.next),
        }
    }
}

impl<M: Message> WorkerPool<M> {
    /// Create a pool handle from a pre-built list of [`StableAddr<M>`].
    ///
    /// Called by the runtime builder; not intended for direct use in
    /// application code.
    ///
    /// # Panics
    ///
    /// Panics if `workers` is empty.
    pub fn new(workers: Vec<StableAddr<M>>) -> Self {
        assert!(
            !workers.is_empty(),
            "WorkerPool requires at least one worker"
        );
        Self {
            workers,
            next: Arc::new(AtomicUsize::new(0)),
        }
    }

    /// Round-robin fire-and-forget send to one worker.
    ///
    /// Advances the rotation counter and calls [`StableAddr::send`] on the
    /// selected worker.  Does **not** retry on `Unroutable`; callers that
    /// need retry-on-crash semantics should yield and call `send` again.
    ///
    /// # Errors
    ///
    /// - [`SendError::MailboxFull`] — the selected worker's mailbox is full.
    /// - [`SendError::Unroutable`] — the worker is between crash and restart.
    pub fn send(&self, msg: M) -> Result<(), SendError> {
        let idx = self.next.fetch_add(1, Ordering::Relaxed) % self.workers.len();
        self.workers[idx].send(msg)
    }

    /// Round-robin request/response.
    ///
    /// Selects one worker and awaits its reply.
    /// `AskError::Send(Unroutable)` may be returned during a crash window;
    /// callers decide retry.
    pub async fn ask(&self, msg: M) -> Result<M::Response, AskError> {
        let idx = self.next.fetch_add(1, Ordering::Relaxed) % self.workers.len();
        self.workers[idx].ask(msg).await
    }

    /// Number of worker slots in the pool.
    pub fn size(&self) -> usize {
        self.workers.len()
    }

    /// Number of workers whose [`StableAddr`] has been populated by the
    /// supervisor (i.e., the worker has been spawned at least once).
    ///
    /// Useful in tests and health checks to wait until all workers are live
    /// before sending messages.
    pub fn workers_ready_count(&self) -> usize {
        self.workers
            .iter()
            .filter(|w| w.current().is_some())
            .count()
    }

    /// Current [`AddrHash`] of the worker at `index`, or `None` if the worker
    /// has not yet been spawned or is in the crash/restart window.
    ///
    /// Index is stable (slot 0 always refers to `{prefix}-0`); the hash
    /// changes on each restart as a new generation is assigned.
    pub fn worker_addr(&self, index: usize) -> Option<AddrHash> {
        self.workers.get(index)?.current().map(|a| a.target())
    }
}