later 0.0.49

#[cfg(feature = "prometheus")]
use crate::metrics;
use crate::{
    commands,
    core::BgJobHandler,
    encoder,
    models::{AmqpCommand, ChannelCommand},
    mq::MqClient,
    BackgroundJobServer, BackgroundJobServerPublisher, ServerConfig, UtcDateTime,
};
use async_std::channel::{Receiver, Sender};
use std::{marker::PhantomData, sync::Arc, time::Duration};

pub(crate) struct WorkerTask {
    stop: Option<tokio::sync::oneshot::Sender<()>>,
    task: tokio::task::JoinHandle<anyhow::Result<()>>,
    /// Epoch milliseconds, updated by the worker itself once per delivery
    /// loop iteration (see `start_distributed_job_worker`) - a purely
    /// in-memory write, independent of whatever database/network call the
    /// worker might currently be blocked on, so it stays a reliable
    /// liveness signal even when the worker is stuck rather than crashed.
    /// `supervise_workers` uses staleness here (not just
    /// `task.is_finished()`) to catch a hang - a worker whose delivery
    /// consumer never returns (for example, blocked forever on an
    /// exhausted database connection pool) never panics and never exits,
    /// so nothing about the task itself signals trouble.
    last_active: Arc<std::sync::atomic::AtomicI64>,
    worker_id: i32,
}

fn now_millis() -> i64 {
    chrono::Utc::now().timestamp_millis()
}

impl<C, H> BackgroundJobServer<C, H>
where
    C: Sync + Send + 'static,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    /// Starts worker tasks using an already-built handler and delivery client.
    ///
    /// Generated builders call this after constructing the handler. The method
    /// waits for every worker consumer to start and returns an error when
    /// `worker_count` is zero or delivery setup fails.
    pub async fn start(
        handler: H,
        mq_client: Arc<Box<dyn MqClient>>,
        config: ServerConfig,
    ) -> anyhow::Result<Self> {
        let mut worker_tasks = Vec::new();
        let mut maintenance_tasks = Vec::new();
        let mut partition_membership_task = None;
        let mut partition_tasks = Vec::new();
        let mut partition_stop = None;
        let handler = Arc::new(handler);
        let publisher = handler.get_publisher();
        let num_bg_workers = config.worker_count;
        if num_bg_workers == 0 {
            return Err(anyhow::anyhow!("worker_count must be at least 1"));
        }
        // Only the *starting* regular-worker count - never the partition
        // poller count below, which stays `num_bg_workers` regardless of
        // autoscaling (see `Config::worker_autoscale`'s own doc comment on
        // why autoscaling is deliberately scoped to the regular pool only).
        let regular_worker_count = match config.worker_autoscale {
            Some(autoscale) => {
                let min = autoscale.min.max(1);
                let max = autoscale.max.max(min);
                num_bg_workers.clamp(min, max)
            }
            None => num_bg_workers,
        };
        let (tx, rx) = async_std::channel::unbounded::<ChannelCommand>();
        let (worker_ready_tx, worker_ready_rx) =
            async_std::channel::bounded::<()>(regular_worker_count.into());
        #[cfg(feature = "dashboard")]
        let server_instance_id = crate::generate_id();

        // workers to process jobs (distributed)
        for id in 0..regular_worker_count {
            worker_tasks.push(spawn_worker(
                handler.clone(),
                mq_client.clone(),
                publisher.routing_key.clone(),
                tx.clone(),
                id.into(),
                worker_ready_tx.clone(),
                #[cfg(feature = "dashboard")]
                server_instance_id.clone(),
            ));
        }
        drop(worker_ready_tx);

        for _ in 0..regular_worker_count {
            if worker_ready_rx.recv().await.is_err() {
                for worker in &worker_tasks {
                    worker.task.abort();
                }
                return Err(anyhow::anyhow!("job worker stopped during startup"));
            }
        }

        let worker_tasks = Arc::new(std::sync::Mutex::new(worker_tasks));
        let next_worker_id = Arc::new(std::sync::atomic::AtomicU32::new(
            regular_worker_count.into(),
        ));
        {
            let worker_tasks = worker_tasks.clone();
            let next_worker_id = next_worker_id.clone();
            let handler_for_supervisor = handler.clone();
            let mq_client_for_supervisor = mq_client.clone();
            let routing_key_for_supervisor = publisher.routing_key.clone();
            let tx_for_supervisor = tx.clone();
            #[cfg(feature = "dashboard")]
            let server_instance_id_for_supervisor = server_instance_id.clone();
            maintenance_tasks.push(tokio::spawn(async move {
                supervise_workers(
                    worker_tasks,
                    next_worker_id,
                    handler_for_supervisor,
                    mq_client_for_supervisor,
                    routing_key_for_supervisor,
                    tx_for_supervisor,
                    #[cfg(feature = "dashboard")]
                    server_instance_id_for_supervisor,
                )
                .await;
                Ok(())
            }));
        }

        if let Some(autoscale) = config.worker_autoscale {
            let worker_tasks = worker_tasks.clone();
            let next_worker_id = next_worker_id.clone();
            let handler_for_autoscale = handler.clone();
            let mq_client_for_autoscale = mq_client.clone();
            let routing_key_for_autoscale = publisher.routing_key.clone();
            let tx_for_autoscale = tx.clone();
            #[cfg(feature = "dashboard")]
            let server_instance_id_for_autoscale = server_instance_id.clone();
            maintenance_tasks.push(tokio::spawn(async move {
                autoscale_workers(
                    autoscale,
                    worker_tasks,
                    next_worker_id,
                    handler_for_autoscale,
                    mq_client_for_autoscale,
                    routing_key_for_autoscale,
                    tx_for_autoscale,
                    #[cfg(feature = "dashboard")]
                    server_instance_id_for_autoscale,
                )
                .await;
                Ok(())
            }));
        }

        let handler_for_ensure_ops = handler.clone();
        let ensure_ops_tx = tx.clone();
        maintenance_tasks.push(tokio::spawn(async move {
            start_bg_worker_system_ops_ensure_ops_run_on_certain_interval(
                handler_for_ensure_ops,
                ensure_ops_tx,
            )
            .await
        }));

        #[cfg(feature = "dashboard")]
        {
            let stats_handler = handler.clone();
            maintenance_tasks.push(tokio::spawn(async move {
                // Dashboard throughput/wait samples are summed in memory and
                // written together - see `Stats::flush_metrics`.
                let mut interval = tokio::time::interval(Duration::from_secs(1));
                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
                // Seeds the live-stage counters from jobs that predate their
                // triggers; from then on the counters follow the index.
                stats_handler
                    .get_publisher()
                    .reconcile_dashboard_counts()
                    .await;
                // Then clears rows earlier dashboard designs left behind, in
                // small batches alongside the flush below.
                let purge_handler = stats_handler.clone();
                tokio::spawn(async move {
                    purge_handler
                        .get_publisher()
                        .purge_unused_dashboard_rows()
                        .await;
                });
                let mut ticks: u64 = 0;
                loop {
                    interval.tick().await;
                    let publisher = stats_handler.get_publisher();
                    publisher.flush_dashboard_metrics().await;
                    // A queue-length reading once a minute, for the graph.
                    if ticks % 60 == 0 {
                        publisher.record_queue_sample().await;
                    }
                    ticks = ticks.wrapping_add(1);
                }
            }));
        }

        #[cfg(feature = "prometheus")]
        {
            let metrics_handler = handler.clone();
            maintenance_tasks.push(tokio::spawn(async move {
                publish_queue_depths(metrics_handler).await
            }));
        }

        let mut partition_owner = None;
        if handler.get_publisher().has_topics() {
            // One owner ID per process: rendezvous hashing assigns
            // partitions to a process, and every poller in that process
            // acts as the same owner when claiming and renewing leases.
            let owner = crate::generate_id();
            handler
                .get_publisher()
                .heartbeat_partition_worker(&owner)
                .await?;
            let live_workers = std::sync::Arc::new(tokio::sync::Mutex::new(vec![owner.clone()]));
            let (stop_tx, _) = tokio::sync::watch::channel(false);
            {
                let membership_handler = handler.clone();
                let membership_owner = owner.clone();
                let membership_live_workers = live_workers.clone();
                partition_membership_task = Some(tokio::spawn(async move {
                    run_partition_membership(
                        membership_handler,
                        membership_owner,
                        membership_live_workers,
                    )
                    .await
                }));
            }
            for _ in 0..num_bg_workers {
                let partition_handler = handler.clone();
                let owner = owner.clone();
                let live_workers = live_workers.clone();
                let stop = stop_tx.subscribe();
                partition_tasks.push(tokio::spawn(async move {
                    run_partition_poller(partition_handler, owner, live_workers, stop).await
                }));
            }
            partition_owner = Some(owner);
            partition_stop = Some(stop_tx);
        }

        for _ in 1..2 {
            let rx_clone = rx.clone();
            let tx_clone = tx.clone();
            let handler_for_check_ops = handler.clone();

            maintenance_tasks.push(tokio::spawn(async move {
                start_bg_worker_system_ops_inproc_cmd_to_amqp_cmd(
                    handler_for_check_ops,
                    tx_clone,
                    rx_clone,
                )
                .await
            }));
        }

        Ok(Self {
            ctx: PhantomData,
            handler,
            mq_client,
            inproc_cmd_tx: tx,
            worker_tasks,
            worker_scaling: tokio::sync::Mutex::new(()),
            next_worker_id,
            #[cfg(feature = "dashboard")]
            server_instance_id,
            maintenance_tasks,
            partition_owner,
            partition_membership_task,
            partition_tasks: std::sync::Mutex::new(partition_tasks),
            partition_stop,
        })
    }

    // `enqueue`, `enqueue_continue` etc. available as
    // self impl Deref to BackgroundJobServer

    /// Returns the number of worker tasks currently owned by this server.
    pub fn worker_count(&self) -> usize {
        self.with_worker_tasks(|workers| {
            workers
                .iter()
                .filter(|worker| !worker.task.is_finished())
                .count()
        })
    }

    /// Starts one worker and waits until its delivery consumer is ready.
    ///
    /// The new worker uses the same backend, namespace, handlers, and
    /// application context as the existing workers. A server can own at most
    /// 255 workers.
    pub async fn add_worker(&self) -> anyhow::Result<usize> {
        let _scaling = self.worker_scaling.lock().await;
        let current_count = self.with_worker_tasks(|workers| {
            workers.retain(|worker| !worker.task.is_finished());
            workers.len()
        });
        if current_count >= usize::from(u8::MAX) {
            return Err(anyhow::anyhow!("worker count cannot exceed 255"));
        }

        let worker_id = self
            .next_worker_id
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let worker_id = i32::try_from(worker_id)
            .map_err(|_| anyhow::anyhow!("worker identifier is out of range"))?;
        let (worker_ready_tx, worker_ready_rx) = async_std::channel::bounded(1);
        let worker = spawn_worker(
            self.handler.clone(),
            self.mq_client.clone(),
            self.handler.get_publisher().routing_key.clone(),
            self.inproc_cmd_tx.clone(),
            worker_id,
            worker_ready_tx,
            #[cfg(feature = "dashboard")]
            self.server_instance_id.clone(),
        );
        if worker_ready_rx.recv().await.is_err() {
            worker.task.abort();
            return Err(anyhow::anyhow!("worker stopped during startup"));
        }
        Ok(self.with_worker_tasks(|workers| {
            workers.push(worker);
            workers.len()
        }))
    }

    /// Stops one worker after its current job finishes.
    ///
    /// At least one worker remains so the server can keep processing system
    /// commands and queued jobs.
    pub async fn remove_worker(&self) -> anyhow::Result<usize> {
        let _scaling = self.worker_scaling.lock().await;
        let (worker, remaining) = self
            .with_worker_tasks(|workers| {
                workers.retain(|worker| !worker.task.is_finished());
                if workers.len() <= 1 {
                    return None;
                }
                let worker = workers.pop()?;
                Some((worker, workers.len()))
            })
            .ok_or_else(|| anyhow::anyhow!("at least one worker must remain"))?;

        let WorkerTask { stop, task, .. } = worker;
        if let Some(stop) = stop {
            let _ = stop.send(());
        }
        task.await
            .map_err(|error| anyhow::anyhow!("worker task stopped unexpectedly: {error}"))??;
        Ok(remaining)
    }

    /// Stops claiming new work, then waits for active handlers to commit.
    ///
    /// Call this from the application's SIGTERM or equivalent shutdown
    /// handler. It stops normal delivery workers after their current command,
    /// and stops partition pollers after their current head is completed or
    /// released. The partition membership heartbeat stays alive while heads
    /// drain so their leases remain valid. If `timeout` expires, remaining
    /// tasks are aborted and durable jobs recover through their normal lease.
    pub async fn shutdown(mut self, timeout: Duration) -> anyhow::Result<()> {
        for task in &self.maintenance_tasks {
            task.abort();
        }
        self.maintenance_tasks.clear();
        if let Some(stop) = self.partition_stop.take() {
            let _ = stop.send(true);
        }

        let mut workers = self.with_worker_tasks(std::mem::take);
        for worker in &mut workers {
            if let Some(stop) = worker.stop.take() {
                let _ = stop.send(());
            }
        }
        let mut partition_tasks = match self.partition_tasks.lock() {
            Ok(mut tasks) => std::mem::take(&mut *tasks),
            Err(poisoned) => std::mem::take(&mut *poisoned.into_inner()),
        };
        let drained = async {
            for worker in &mut workers {
                (&mut worker.task).await.map_err(|error| {
                    anyhow::anyhow!("worker task stopped unexpectedly: {error}")
                })??;
            }
            for task in &mut partition_tasks {
                task.await.map_err(|error| {
                    anyhow::anyhow!("partition poller stopped unexpectedly: {error}")
                })??;
            }
            Ok::<(), anyhow::Error>(())
        };
        let result = match tokio::time::timeout(timeout, drained).await {
            Ok(result) => result,
            Err(_) => {
                for worker in &workers {
                    worker.task.abort();
                }
                for task in &partition_tasks {
                    task.abort();
                }
                Err(anyhow::anyhow!(
                    "graceful shutdown timed out; durable jobs will recover through their leases"
                ))
            }
        };
        if let Some(task) = self.partition_membership_task.take() {
            task.abort();
        }
        // Whatever the last second of work recorded and the flush loop hasn't
        // written yet.
        #[cfg(feature = "dashboard")]
        self.handler.get_publisher().flush_dashboard_metrics().await;
        if let Some(owner) = self.partition_owner.take() {
            self.handler
                .get_publisher()
                .deregister_partition_worker(&owner)
                .await?;
        }
        result
    }

    fn with_worker_tasks<T>(&self, operation: impl FnOnce(&mut Vec<WorkerTask>) -> T) -> T {
        match self.worker_tasks.lock() {
            Ok(mut workers) => operation(&mut workers),
            Err(poisoned) => {
                let mut workers = poisoned.into_inner();
                operation(&mut workers)
            }
        }
    }
}

#[cfg(feature = "dashboard")]
async fn publish_worker_heartbeats<C, H>(handler: Arc<H>, worker_id: String) -> anyhow::Result<()>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    let mut interval = tokio::time::interval(Duration::from_secs(5));
    loop {
        interval.tick().await;
        if let Err(error) = handler
            .get_publisher()
            .publish_worker_heartbeat(worker_id.clone())
            .await
        {
            tracing::warn!(%worker_id, %error, "Could not publish worker heartbeat");
        }
    }
}

#[cfg(feature = "prometheus")]
async fn publish_queue_depths<C, H>(handler: Arc<H>) -> anyhow::Result<()>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    let mut interval = tokio::time::interval(Duration::from_secs(5));
    loop {
        interval.tick().await;
        match handler.get_publisher().storage.queue_depths().await {
            Ok(depths) => {
                for depth in depths {
                    metrics::set_queue_depth(
                        handler.get_publisher().metrics_queue(),
                        depth.state.metric_name(),
                        depth.count,
                    );
                }
            }
            Err(error) => tracing::warn!(%error, "Could not collect queue depth metrics"),
        }
    }
}

/// Heartbeats this process's partition worker owner and refreshes the shared
/// live-worker list every 5 seconds.
///
/// Membership views across processes can briefly differ; the partition
/// lease remains the final authority on who may run a head, so a stale or
/// slightly-ahead view here only affects which partitions a poller attempts
/// to claim, never correctness.
async fn run_partition_membership<C, H>(
    handler: Arc<H>,
    owner: String,
    live_workers: Arc<tokio::sync::Mutex<Vec<String>>>,
) -> anyhow::Result<()>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    loop {
        if let Err(error) = handler
            .get_publisher()
            .heartbeat_partition_worker(&owner)
            .await
        {
            tracing::warn!(%error, "Failed to send partition worker heartbeat");
        }
        match handler.get_publisher().list_live_partition_workers().await {
            Ok(mut workers) => {
                if !workers.iter().any(|live| live == &owner) {
                    // Our own heartbeat may not be visible yet right after
                    // the first write on some backends; always include
                    // ourselves so we never starve for our own partitions.
                    workers.push(owner.clone());
                }
                #[cfg(feature = "prometheus")]
                metrics::set_partition_workers_live(
                    handler.get_publisher().metrics_queue(),
                    workers.len(),
                );
                *live_workers.lock().await = workers;
            }
            Err(error) => tracing::warn!(%error, "Failed to list live partition workers"),
        }
        #[cfg(feature = "prometheus")]
        match handler.get_publisher().partition_backlog_summary().await {
            Ok(summary) => {
                // The query only returns topics with at least one ready
                // head, so every *configured* topic must still be visited
                // here — otherwise a topic that just drained to zero would
                // leave its gauge stuck at its last nonzero value forever.
                let mut by_topic: std::collections::HashMap<String, (usize, f64)> = summary
                    .into_iter()
                    .map(|(topic, ready_count, age)| (topic, (ready_count as usize, age)))
                    .collect();
                for topic in handler.get_publisher().topic_names() {
                    let (ready_count, oldest_head_age_seconds) =
                        by_topic.remove(topic).unwrap_or((0, 0.0));
                    metrics::set_partition_backlog(
                        handler.get_publisher().metrics_queue(),
                        topic,
                        ready_count,
                        oldest_head_age_seconds,
                    );
                }
            }
            Err(error) => tracing::warn!(%error, "Failed to summarize partition backlog"),
        }
        #[cfg(feature = "prometheus")]
        match handler
            .get_publisher()
            .partition_queue_depth_summary()
            .await
        {
            Ok(summary) => {
                // Same reasoning as the backlog block above: only
                // partitions with at least one queued job are returned, so
                // every configured partition must still be visited to zero
                // out a gauge for one that just fully drained.
                let mut by_partition: std::collections::HashMap<(String, u32), u64> = summary
                    .into_iter()
                    .map(|(topic, partition, depth)| ((topic, partition.0), depth))
                    .collect();
                for topic_config in handler.get_publisher().topics() {
                    for partition in 0..topic_config.partition_count() {
                        let depth = by_partition
                            .remove(&(topic_config.name().to_string(), partition))
                            .unwrap_or(0);
                        metrics::set_partition_queue_depth(
                            handler.get_publisher().metrics_queue(),
                            topic_config.name(),
                            &partition.to_string(),
                            i64::try_from(depth).unwrap_or(i64::MAX),
                        );
                    }
                }
            }
            Err(error) => tracing::warn!(%error, "Failed to summarize partition queue depth"),
        }
        #[cfg(all(feature = "prometheus", feature = "retained-log"))]
        if let Some(source) = handler.get_publisher().retained_log_lag().source() {
            for topic in handler.get_publisher().retained_log_lag().topics() {
                match source.consumer_lag(topic).await {
                    Ok(lags) => {
                        for lag in lags {
                            metrics::set_retained_log_consumer_lag(
                                handler.get_publisher().metrics_queue(),
                                topic,
                                &lag.group,
                                lag.lag,
                            );
                        }
                    }
                    Err(error) => {
                        tracing::warn!(%error, topic, "Failed to compute consumer lag")
                    }
                }
            }
        }
        sleep_ms(2_000).await;
    }
}

/// Repeatedly tries to claim and run one topic partition head.
///
/// Runs at full speed while work is found, and backs off briefly when no
/// partition has a ready, unowned head, so idle servers with topics
/// configured do not busy-loop the shared state backend.
async fn run_partition_poller<C, H>(
    handler: Arc<H>,
    owner: String,
    live_workers: Arc<tokio::sync::Mutex<Vec<String>>>,
    mut stop: tokio::sync::watch::Receiver<bool>,
) -> anyhow::Result<()>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    #[cfg(feature = "prometheus")]
    let worker_id_label = owner.clone();
    let wake = handler.get_publisher().partition_wake().clone();
    // Held across iterations so consecutive ready jobs in the same partition
    // don't each pay for a fresh claim; see `partition::poll_partition_head`.
    let mut assignment: Option<crate::partition::PartitionAssignment> = None;
    loop {
        if *stop.borrow() {
            return Ok(());
        }
        let workers = live_workers.lock().await.clone();
        match crate::partition::poll_partition_head(
            &handler,
            &owner,
            &workers,
            &mut assignment,
            #[cfg(feature = "prometheus")]
            &worker_id_label,
        )
        .await
        {
            Ok(true) => continue,
            // Idle: wait for either the backoff or an explicit wake-up
            // signal from a local enqueue or continuation, whichever comes
            // first. A missed
            // or spurious wake is harmless, it only changes how soon this
            // loop tries again.
            Ok(false) => {
                tokio::select! {
                    _ = sleep_ms(250) => {}
                    _ = wake.notified() => {}
                    changed = stop.changed() => {
                        if changed.is_ok() && *stop.borrow() {
                            return Ok(());
                        }
                    }
                }
            }
            Err(error) => {
                tracing::warn!(%error, "Partition head poll failed");
                sleep_ms(1_000).await;
            }
        }
    }
}

/// Periodically checks every regular delivery worker for two independent
/// failure modes and replaces either kind with a fresh worker:
///
/// - **Exited** - `task.is_finished()`, i.e. the worker's future already
///   returned (crashed, returned an error, or its delivery stream ended).
///   `spawn_worker` already logs the reason when this happens; this just
///   notices and replaces it.
/// - **Hung** - the task is still technically running, but
///   [`WorkerTask::last_active`] hasn't advanced in `STALL_THRESHOLD_MS`.
///   This is the failure mode `task.is_finished()` alone can never catch: a
///   worker blocked forever inside `consumer.next()` (for example, a
///   database connection pool starved by other load) never panics and
///   never returns, so nothing about the task itself signals trouble - only
///   the absence of its own progress marker does. This exact gap once let
///   every regular worker in production sit idle for hours with a
///   79,000+ job backlog and not one line logged about why.
#[allow(clippy::too_many_arguments)]
async fn supervise_workers<C, H>(
    worker_tasks: Arc<std::sync::Mutex<Vec<WorkerTask>>>,
    next_worker_id: Arc<std::sync::atomic::AtomicU32>,
    handler: Arc<H>,
    mq_client: Arc<Box<dyn MqClient>>,
    routing_key: String,
    inproc_cmd_tx: Sender<ChannelCommand>,
    #[cfg(feature = "dashboard")] server_instance_id: String,
) where
    C: Sync + Send + 'static,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    // Shortened under `cfg(test)` so tests exercising this exercise real
    // wall-clock waits (there's no injectable clock here) without taking
    // minutes - the ratio between the two stays the same.
    #[cfg(not(test))]
    const CHECK_INTERVAL_MS: u64 = 30_000;
    #[cfg(test)]
    const CHECK_INTERVAL_MS: u64 = 200;
    #[cfg(not(test))]
    const STALL_THRESHOLD_MS: i64 = 5 * 60 * 1_000;
    #[cfg(test)]
    const STALL_THRESHOLD_MS: i64 = 500;

    loop {
        sleep_ms(CHECK_INTERVAL_MS).await;

        let stale_or_dead: Vec<usize> = {
            let tasks = match worker_tasks.lock() {
                Ok(tasks) => tasks,
                Err(poisoned) => poisoned.into_inner(),
            };
            let now = now_millis();
            tasks
                .iter()
                .enumerate()
                .filter(|(_, worker)| {
                    worker.task.is_finished()
                        || now
                            - worker
                                .last_active
                                .load(std::sync::atomic::Ordering::Relaxed)
                            > STALL_THRESHOLD_MS
                })
                .map(|(index, _)| index)
                .collect()
        };
        if stale_or_dead.is_empty() {
            continue;
        }

        // Highest index first so earlier ones stay valid as each is removed.
        for &index in stale_or_dead.iter().rev() {
            let worker = {
                let mut tasks = match worker_tasks.lock() {
                    Ok(tasks) => tasks,
                    Err(poisoned) => poisoned.into_inner(),
                };
                if index >= tasks.len() {
                    // Already handled by a previous pass, or removed
                    // concurrently by `remove_worker` - nothing to do.
                    continue;
                }
                tasks.remove(index)
            };
            let worker_id = worker.worker_id;
            if worker.task.is_finished() {
                tracing::error!(
                    worker_id,
                    "[Worker#{worker_id}] Found dead (already exited) during a routine health \
                     check - replacing it"
                );
            } else {
                let stalled_for_ms = now_millis()
                    - worker
                        .last_active
                        .load(std::sync::atomic::Ordering::Relaxed);
                tracing::error!(
                    worker_id,
                    stalled_for_ms,
                    "[Worker#{worker_id}] Appears hung (no progress in {stalled_for_ms}ms) - \
                     forcibly restarting it"
                );
                worker.task.abort();
            }

            match spawn_and_register_worker(
                &worker_tasks,
                &next_worker_id,
                &handler,
                &mq_client,
                &routing_key,
                &inproc_cmd_tx,
                #[cfg(feature = "dashboard")]
                &server_instance_id,
            )
            .await
            {
                Ok(new_worker_id) => tracing::warn!(
                    worker_id = new_worker_id,
                    replaces = worker_id,
                    "[Worker#{new_worker_id}] Started as a replacement for Worker#{worker_id}"
                ),
                Err(error) => tracing::error!(
                    %error,
                    "Could not start a replacement for dead Worker#{worker_id}"
                ),
            }
        }
    }
}

/// Assigns a fresh worker ID, starts a worker with it, waits for its
/// delivery consumer to be ready, and registers it in `worker_tasks` -
/// shared by [`supervise_workers`] (replacing a dead/hung worker) and
/// [`autoscale_workers`] (adding one under load). Returns the new worker's
/// ID on success; the worker is already running and registered by the time
/// this returns.
#[allow(clippy::too_many_arguments)]
async fn spawn_and_register_worker<C, H>(
    worker_tasks: &Arc<std::sync::Mutex<Vec<WorkerTask>>>,
    next_worker_id: &Arc<std::sync::atomic::AtomicU32>,
    handler: &Arc<H>,
    mq_client: &Arc<Box<dyn MqClient>>,
    routing_key: &str,
    inproc_cmd_tx: &Sender<ChannelCommand>,
    #[cfg(feature = "dashboard")] server_instance_id: &str,
) -> anyhow::Result<i32>
where
    C: Sync + Send + 'static,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    let new_worker_id = next_worker_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    let new_worker_id = i32::try_from(new_worker_id)
        .map_err(|_| anyhow::anyhow!("worker identifier is out of range"))?;
    let (worker_ready_tx, worker_ready_rx) = async_std::channel::bounded(1);
    let worker = spawn_worker(
        handler.clone(),
        mq_client.clone(),
        routing_key.to_string(),
        inproc_cmd_tx.clone(),
        new_worker_id,
        worker_ready_tx,
        #[cfg(feature = "dashboard")]
        server_instance_id.to_string(),
    );
    if worker_ready_rx.recv().await.is_err() {
        worker.task.abort();
        return Err(anyhow::anyhow!("[Worker#{new_worker_id}] Failed to start"));
    }
    match worker_tasks.lock() {
        Ok(mut tasks) => tasks.push(worker),
        Err(poisoned) => poisoned.into_inner().push(worker),
    }
    Ok(new_worker_id)
}

/// Adds or removes regular workers between `bounds.min` and `bounds.max`
/// based on the regular delivery queue's own backlog - see
/// [`crate::WorkerAutoscale`]'s own doc comment for the exact thresholds
/// and why they're deliberately conservative (never more than one worker
/// changed per tick).
#[allow(clippy::too_many_arguments)]
async fn autoscale_workers<C, H>(
    bounds: crate::WorkerAutoscale,
    worker_tasks: Arc<std::sync::Mutex<Vec<WorkerTask>>>,
    next_worker_id: Arc<std::sync::atomic::AtomicU32>,
    handler: Arc<H>,
    mq_client: Arc<Box<dyn MqClient>>,
    routing_key: String,
    inproc_cmd_tx: Sender<ChannelCommand>,
    #[cfg(feature = "dashboard")] server_instance_id: String,
) where
    C: Sync + Send + 'static,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    #[cfg(not(test))]
    const CHECK_INTERVAL_MS: u64 = 15_000;
    #[cfg(test)]
    const CHECK_INTERVAL_MS: u64 = 200;
    const BACKLOG_PER_WORKER_SCALE_UP_THRESHOLD: usize = 20;
    const CONSECUTIVE_IDLE_TICKS_BEFORE_SCALE_DOWN: u32 = 3;

    let min = bounds.min.max(1);
    let max = bounds.max.max(min);
    let mut consecutive_idle_ticks = 0u32;

    loop {
        sleep_ms(CHECK_INTERVAL_MS).await;

        let current_count = {
            let tasks = match worker_tasks.lock() {
                Ok(tasks) => tasks,
                Err(poisoned) => poisoned.into_inner(),
            };
            tasks.iter().filter(|w| !w.task.is_finished()).count()
        };

        let depth = match mq_client.queue_depth(&routing_key).await {
            Ok(depth) => depth,
            Err(error) => {
                tracing::warn!(%error, "Could not read delivery queue depth for autoscaling");
                continue;
            }
        };
        if depth == 0 {
            consecutive_idle_ticks += 1;
        } else {
            consecutive_idle_ticks = 0;
        }

        let current_count_u8 = u8::try_from(current_count).unwrap_or(u8::MAX);
        if depth > current_count * BACKLOG_PER_WORKER_SCALE_UP_THRESHOLD && current_count_u8 < max {
            match spawn_and_register_worker(
                &worker_tasks,
                &next_worker_id,
                &handler,
                &mq_client,
                &routing_key,
                &inproc_cmd_tx,
                #[cfg(feature = "dashboard")]
                &server_instance_id,
            )
            .await
            {
                Ok(new_worker_id) => tracing::warn!(
                    worker_id = new_worker_id,
                    depth,
                    current_count,
                    max,
                    "[Worker#{new_worker_id}] Started - scaling up, backlog is {depth} against \
                     {current_count} workers"
                ),
                Err(error) => tracing::error!(%error, "Could not scale up a worker"),
            }
        } else if consecutive_idle_ticks >= CONSECUTIVE_IDLE_TICKS_BEFORE_SCALE_DOWN
            && current_count_u8 > min
        {
            consecutive_idle_ticks = 0;
            let worker = {
                let mut tasks = match worker_tasks.lock() {
                    Ok(tasks) => tasks,
                    Err(poisoned) => poisoned.into_inner(),
                };
                tasks.retain(|worker| !worker.task.is_finished());
                if tasks.len() <= usize::from(min) {
                    None
                } else {
                    tasks.pop()
                }
            };
            let Some(mut worker) = worker else {
                continue;
            };
            let worker_id = worker.worker_id;
            tracing::warn!(
                worker_id,
                current_count,
                min,
                "[Worker#{worker_id}] Stopping - backlog has been empty, scaling down"
            );
            if let Some(stop) = worker.stop.take() {
                let _ = stop.send(());
            }
            if let Err(error) = worker.task.await {
                tracing::warn!(worker_id, %error, "[Worker#{worker_id}] Error while stopping for scale-down");
            }
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn spawn_worker<C, H>(
    handler: Arc<H>,
    mq_client: Arc<Box<dyn MqClient>>,
    routing_key: String,
    inproc_cmd_tx: Sender<ChannelCommand>,
    worker_id: i32,
    worker_ready_tx: Sender<()>,
    #[cfg(feature = "dashboard")] server_instance_id: String,
) -> WorkerTask
where
    C: Sync + Send + 'static,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    let (stop, stop_rx) = tokio::sync::oneshot::channel();
    let last_active = Arc::new(std::sync::atomic::AtomicI64::new(now_millis()));
    let task = {
        let last_active = last_active.clone();
        tokio::spawn(async move {
            #[cfg(feature = "dashboard")]
            let heartbeat = {
                let heartbeat_handler = handler.clone();
                let heartbeat_worker_id = format!("{server_instance_id}:{worker_id}");
                tokio::spawn(async move {
                    publish_worker_heartbeats(heartbeat_handler, heartbeat_worker_id).await
                })
            };
            let result = start_distributed_job_worker(
                handler,
                worker_id,
                mq_client,
                &routing_key,
                inproc_cmd_tx,
                worker_ready_tx,
                stop_rx,
                last_active,
            )
            .await;
            #[cfg(feature = "dashboard")]
            heartbeat.abort();
            match &result {
                Ok(()) => tracing::info!(
                    worker_id,
                    "[Worker#{worker_id}] Exited (stop requested or delivery stream ended)"
                ),
                Err(error) => {
                    tracing::error!(worker_id, %error, "[Worker#{worker_id}] Exited with an error")
                }
            }
            result
        })
    };
    WorkerTask {
        stop: Some(stop),
        task,
        last_active,
        worker_id,
    }
}

pub(crate) async fn sleep_ms(ms: u64) {
    tokio::time::sleep(Duration::from_millis(ms)).await;
}

impl<C, H> std::ops::Deref for BackgroundJobServer<C, H>
where
    C: Sync + Send + 'static,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    type Target = BackgroundJobServerPublisher;

    fn deref(&self) -> &Self::Target {
        self.handler.get_publisher()
    }
}

/// Listens for commands on an in-process queue, then waits a few moments
/// and runs the corresponding system operation (polling for certain jobs,
/// sweeping expired storage, etc.) directly in this task.
///
/// Deliberately *not* routed through [`AmqpCommand`] and the shared
/// delivery queue the way [`AmqpCommand::ExecuteJob`] is - these
/// maintenance operations are cheap, idempotent, and safe to run
/// concurrently from every server instance sharing a database (each one
/// atomically claims/deletes only what it actually finds ready), so
/// nothing is gained by making one instance's worker dequeue a message for
/// another instance's poll trigger. Going through the shared queue used
/// to mean these could get stuck behind a deep `ExecuteJob` backlog under
/// sustained load - FIFO delivery order gives every regular job ahead of
/// a poll command priority over it, so a backlog large enough to keep
/// workers busy for minutes could starve `PollDelayedJobs`/
/// `PollRequeuedJobs`/`PollRecurringJobs` (delaying scheduled work) and
/// `PollExpiredStorage`/`PollStuckJobs` (the two problems they exist to
/// prevent - unbounded storage growth and permanently stuck jobs -
/// getting worse for exactly as long as the backlog that caused them
/// persists). Running locally removes that dependency entirely.
async fn start_bg_worker_system_ops_inproc_cmd_to_amqp_cmd<C, H>(
    handler: Arc<H>,
    tx: Sender<ChannelCommand>,
    rx: Receiver<ChannelCommand>,
) -> anyhow::Result<()>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    loop {
        let channel_command = match rx.recv().await {
            Ok(command) => command,
            Err(_) => {
                tracing::info!("System operation channel closed");
                return Ok(());
            }
        };
        sleep_ms(2000).await; // run system ops after 2s delay

        let config = handler.get_publisher().storage.config();
        // Re-queues itself once done, the same self-perpetuating cadence
        // `handle_amqp_command`'s poll arms used to drive (roughly every
        // 2s, bounded below by however long the operation itself takes) -
        // `start_bg_worker_system_ops_ensure_ops_run_on_certain_interval`
        // is only the backstop that restarts this chain if it ever breaks,
        // not the normal steady-state cadence.
        //
        // This is the *only* task running any of these five polls (see
        // `bg_job_server.rs`'s `start`, `for _ in 1..2`) and nothing
        // supervises or restarts it - `tokio::spawn`'s `JoinHandle` is only
        // ever awaited from `shutdown()`. A transient error from any single
        // poll (a busy/locked SQLite connection under load is the one this
        // was actually seen with in production) used to propagate out via
        // `?`, ending this loop for good: every poll it drives - including
        // `PollExpiredStorage` - then silently stopped forever, letting
        // `later_storage_range` grow unbounded (millions of un-swept rows)
        // until the dashboard itself timed out/errored under the resulting
        // lock contention. Logging and re-queuing instead means one bad
        // tick is retried next cycle rather than ending the poll for the
        // life of the process.
        let result = match &channel_command {
            ChannelCommand::PollDelayedJobs => {
                commands::handle_poll_delayed_job_command(handler.clone())
                    .await
                    .and(config.poll_delayed_jobs_last_run_set().await)
            }
            ChannelCommand::PollRequeuedJobs => {
                commands::handle_poll_requeued_job_command(handler.clone())
                    .await
                    .and(config.poll_requeued_jobs_last_run_set().await)
            }
            ChannelCommand::PollRecurringJobs => {
                commands::handle_poll_recurring_job_command(handler.clone())
                    .await
                    .and(config.poll_recurring_jobs_last_run_set().await)
            }
            ChannelCommand::PollExpiredStorage => {
                commands::handle_poll_expired_storage_command(handler.clone())
                    .await
                    .and(config.poll_expired_storage_last_run_set().await)
            }
            ChannelCommand::PollStuckJobs => {
                commands::handle_poll_stuck_jobs_command(handler.clone())
                    .await
                    .and(config.poll_stuck_jobs_last_run_set().await)
            }
            ChannelCommand::VacuumDatabase => {
                commands::handle_vacuum_database_command(handler.clone())
                    .await
                    .and(config.poll_vacuum_last_run_set().await)
            }
        };
        if let Err(error) = result {
            tracing::warn!(
                command = ?channel_command,
                %error,
                "System operation failed; will retry next cycle"
            );
        }
        // `VacuumDatabase` is interval-gated (see `Config::vacuum_interval`)
        // and far too expensive to run on this loop's normal ~10s
        // self-perpetuating cadence - only the staleness backstop
        // (`start_bg_worker_system_ops_ensure_ops_run_on_certain_interval`)
        // re-enqueues it, once its own much longer interval has elapsed.
        if !matches!(channel_command, ChannelCommand::VacuumDatabase) {
            tx.send(channel_command).await?;
        }
    }
}

/// Every 10 secs, ensure some system operations are run - otherwise enqueue them
async fn start_bg_worker_system_ops_ensure_ops_run_on_certain_interval<C, H>(
    handler: Arc<H>,
    tx: Sender<ChannelCommand>,
) -> anyhow::Result<()>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    sleep_ms(3_000).await; // wait for existing commands to be processed

    loop {
        let config = handler.get_publisher().storage.config();

        enqueue_if(
            &tx,
            ChannelCommand::PollDelayedJobs,
            config.poll_delayed_jobs_last_run().await?,
            10,
        )
        .await?;

        enqueue_if(
            &tx,
            ChannelCommand::PollRequeuedJobs,
            config.poll_requeued_jobs_last_run().await?,
            10,
        )
        .await?;

        enqueue_if(
            &tx,
            ChannelCommand::PollRecurringJobs,
            config.poll_recurring_jobs_last_run().await?,
            10,
        )
        .await?;

        enqueue_if(
            &tx,
            ChannelCommand::PollExpiredStorage,
            config.poll_expired_storage_last_run().await?,
            15,
        )
        .await?;

        enqueue_if(
            &tx,
            ChannelCommand::PollStuckJobs,
            config.poll_stuck_jobs_last_run().await?,
            30,
        )
        .await?;

        // Unlike every poll above, `VacuumDatabase` only ever runs through
        // this staleness check - see `ChannelCommand::VacuumDatabase` and
        // `start_bg_worker_system_ops_inproc_cmd_to_amqp_cmd`'s own note on
        // why it deliberately doesn't re-queue itself. Disabled entirely
        // (the pre-existing, surprise-free default) unless
        // `Config::vacuum_interval` is set.
        if let Some(vacuum_interval) = handler.get_publisher().vacuum_interval() {
            enqueue_if(
                &tx,
                ChannelCommand::VacuumDatabase,
                config.poll_vacuum_last_run().await?,
                i64::try_from(vacuum_interval.as_secs()).unwrap_or(i64::MAX),
            )
            .await?;
        }

        sleep_ms(10_000).await; // check again 10s later
    }

    // unreachable

    async fn enqueue_if(
        tx: &Sender<ChannelCommand>,
        cmd: ChannelCommand,
        last_run: UtcDateTime,
        if_not_sec: i64,
    ) -> anyhow::Result<()> {
        let last_run_since = chrono::Utc::now() - last_run;
        if last_run_since.num_seconds() > if_not_sec {
            tracing::warn!("Ops {} did not run for a while: Enqueuing", cmd);

            tx.send(cmd).await?;
        }
        Ok(())
    }
}

async fn start_distributed_job_worker<C, H>(
    handler: Arc<H>,
    worker_id: i32,
    mq_client: Arc<Box<dyn MqClient>>,
    routing_key: &str,
    inproc_cmd_tx: Sender<ChannelCommand>,
    worker_ready_tx: Sender<()>,
    mut stop: tokio::sync::oneshot::Receiver<()>,
    last_active: Arc<std::sync::atomic::AtomicI64>,
) -> anyhow::Result<()>
where
    C: Sync + Send,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    tracing::info!("[Worker#{}] Starting", worker_id);
    #[cfg(feature = "prometheus")]
    let worker_id_label = worker_id.to_string();
    #[cfg(feature = "prometheus")]
    let _worker_guard = metrics::WorkerGuard::new(routing_key, &worker_id_label);

    let mut consumer = mq_client.new_consumer(routing_key, worker_id).await?;
    worker_ready_tx
        .send(())
        .await
        .map_err(|_| anyhow::anyhow!("server stopped during worker startup"))?;

    loop {
        let message = tokio::select! {
            biased;
            _ = &mut stop => return Ok(()),
            message = consumer.next() => message,
        };
        // A purely in-memory write, deliberately placed right after
        // `consumer.next()` returns rather than before - see
        // `WorkerTask::last_active`'s own doc comment on why this is the
        // signal `supervise_workers` watches for a hung (not crashed)
        // worker: it only advances once this call has actually returned,
        // so it stops moving the moment that call blocks forever instead
        // of just returning slowly.
        last_active.store(now_millis(), std::sync::atomic::Ordering::Relaxed);
        let Some(message) = message else {
            tracing::warn!(
                worker_id,
                "[Worker#{worker_id}] Delivery consumer's message stream ended - exiting so the \
                 supervisor can replace this worker with one holding a fresh consumer"
            );
            return Ok(());
        };
        match message {
            Ok(delivery) => match encoder::decode::<AmqpCommand>(delivery.data()) {
                Ok(command) => {
                    let headers = delivery.get_headers();
                    #[cfg(feature = "prometheus")]
                    let command_name = metrics::command_name(&command);

                    match commands::handle_amqp_command(
                        command,
                        worker_id,
                        #[cfg(feature = "prometheus")]
                        &worker_id_label,
                        &handler,
                        &inproc_cmd_tx,
                        headers,
                    )
                    .await
                    {
                        Ok(()) => {
                            #[cfg(feature = "prometheus")]
                            metrics::record_command(
                                routing_key,
                                &worker_id_label,
                                command_name,
                                "success",
                            );
                            delivery.ack().await?
                        }
                        Err(error) => {
                            #[cfg(feature = "prometheus")]
                            metrics::record_command(
                                routing_key,
                                &worker_id_label,
                                command_name,
                                "error",
                            );
                            tracing::warn!(
                                worker_id,
                                %error,
                                "Job command failed and will be retried"
                            );
                            delivery.nack_requeue().await?;
                        }
                    }
                }
                Err(err) => {
                    #[cfg(feature = "prometheus")]
                    metrics::record_command(routing_key, &worker_id_label, "invalid", "discarded");
                    tracing::warn!(
                        "[Worker#{}] Unknown message received [{} bytes]: {}",
                        worker_id,
                        delivery.data().len(),
                        err
                    );

                    delivery.ack().await?;
                }
            },
            Err(e) => {
                // Previously looped straight back to `consumer.next()` with
                // no delay - a persistent (not transient) claim failure
                // turned into a silent, unbounded busy-loop that never
                // logged more than once (this same line, endlessly) and
                // never gave `supervise_workers` anything to catch, since
                // `last_active` above kept advancing every iteration. A
                // brief backoff bounds the spin; `last_active` still
                // advances during it, since a merely-erroring worker isn't
                // the "stuck with a DB pool starved" case the watchdog
                // exists for.
                tracing::error!(worker_id, error = ?e, "[Worker#{worker_id}] Consumer error - backing off before retrying");
                sleep_ms(1_000).await;
            }
        }
    }
}

impl<C, H> Drop for BackgroundJobServer<C, H>
where
    C: Sync + Send + 'static,
    H: BgJobHandler<C> + Sync + Send + 'static,
{
    fn drop(&mut self) {
        match self.worker_tasks.lock() {
            Ok(workers) => {
                for worker in workers.iter() {
                    worker.task.abort();
                }
            }
            Err(poisoned) => {
                for worker in poisoned.into_inner().iter() {
                    worker.task.abort();
                }
            }
        }
        for task in &self.maintenance_tasks {
            task.abort();
        }

        if let Some(task) = &self.partition_membership_task {
            task.abort();
        }
        match self.partition_tasks.lock() {
            Ok(tasks) => {
                for task in tasks.iter() {
                    task.abort();
                }
            }
            Err(poisoned) => {
                for task in poisoned.into_inner().iter() {
                    task.abort();
                }
            }
        }

        // Best-effort: let other workers stop considering this process for
        // rendezvous assignment immediately, instead of waiting out its
        // heartbeat TTL. Only attempted when a Tokio runtime is still
        // available (for example, not after the runtime itself has already
        // started shutting down).
        if let Some(owner) = self.partition_owner.take() {
            if let Ok(handle) = tokio::runtime::Handle::try_current() {
                let handler = self.handler.clone();
                handle.spawn(async move {
                    if let Err(error) = handler
                        .get_publisher()
                        .deregister_partition_worker(&owner)
                        .await
                    {
                        tracing::warn!(%error, "Failed to deregister partition worker on shutdown");
                    }
                });
            }
        }
    }
}