awa_worker/

dispatcher.rs

use crate::executor::{DispatchedJob, JobExecutor};
use crate::runtime::InFlightMap;
use crate::storage::RuntimeStorage;
use awa_model::JobRow;
use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex as StdMutex, OnceLock};
use std::time::{Duration, Instant};
use tokio::sync::{Mutex, Notify, Semaphore};
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};
use uuid::Uuid;

const DEFAULT_CLAIM_BATCH_LIMIT: usize = 512;
const MAX_CLAIMERS_PER_QUEUE: i16 = 4;
const CLAIMER_LEASE_TTL: Duration = Duration::from_secs(3);
const CLAIMER_IDLE_THRESHOLD: Duration = Duration::from_millis(500);

fn max_claimers_per_queue() -> i16 {
    static MAX_CLAIMERS: OnceLock<i16> = OnceLock::new();
    *MAX_CLAIMERS.get_or_init(|| {
        let Ok(raw) = std::env::var("AWA_MAX_CLAIMERS_PER_QUEUE") else {
            return MAX_CLAIMERS_PER_QUEUE;
        };

        match raw.parse::<i16>() {
            Ok(value) if value > 0 => value,
            Ok(value) => {
                warn!(
                    value,
                    default = MAX_CLAIMERS_PER_QUEUE,
                    "AWA_MAX_CLAIMERS_PER_QUEUE must be positive; using default"
                );
                MAX_CLAIMERS_PER_QUEUE
            }
            Err(error) => {
                warn!(
                    raw = %raw,
                    %error,
                    default = MAX_CLAIMERS_PER_QUEUE,
                    "Failed to parse AWA_MAX_CLAIMERS_PER_QUEUE; using default"
                );
                MAX_CLAIMERS_PER_QUEUE
            }
        }
    })
}

#[derive(Debug, Clone, Copy)]
enum WakeReason {
    Notify,
    Capacity,
    Poll,
}

impl WakeReason {
    fn as_str(self) -> &'static str {
        match self {
            WakeReason::Notify => "notify",
            WakeReason::Capacity => "capacity",
            WakeReason::Poll => "poll",
        }
    }
}

/// Rate limit configuration for a queue.
#[derive(Debug, Clone)]
pub struct RateLimit {
    /// Maximum sustained dispatch rate (jobs per second).
    pub max_rate: f64,
    /// Maximum burst size. Defaults to ceil(max_rate) if 0.
    pub burst: u32,
}

/// Internal token bucket state for rate limiting.
pub(crate) struct TokenBucket {
    tokens: f64,
    max_tokens: f64,
    refill_rate: f64,
    last_refill: Instant,
}

impl TokenBucket {
    fn new(rate_limit: &RateLimit) -> Self {
        let burst = if rate_limit.burst == 0 {
            (rate_limit.max_rate.ceil() as u32).max(1)
        } else {
            rate_limit.burst
        };
        Self {
            tokens: burst as f64,
            max_tokens: burst as f64,
            refill_rate: rate_limit.max_rate,
            last_refill: Instant::now(),
        }
    }

    /// Return how many whole tokens are available after refilling.
    fn available(&mut self) -> u32 {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
        self.tokens = (self.tokens + elapsed * self.refill_rate).min(self.max_tokens);
        self.last_refill = now;
        self.tokens.floor() as u32
    }

    /// Consume `n` tokens (caller must ensure n <= available()).
    fn consume(&mut self, n: u32) {
        self.tokens -= n as f64;
    }

    fn refund(&mut self, n: u32) {
        self.tokens = (self.tokens + n as f64).min(self.max_tokens);
    }
}

pub(crate) fn shared_rate_limiter(config: &QueueConfig) -> Option<Arc<StdMutex<TokenBucket>>> {
    config
        .rate_limit
        .as_ref()
        .map(TokenBucket::new)
        .map(|bucket| Arc::new(StdMutex::new(bucket)))
}

/// Configuration for a single queue.
#[derive(Debug, Clone)]
pub struct QueueConfig {
    pub max_workers: u32,
    pub poll_interval: Duration,
    pub deadline_duration: Duration,
    pub priority_aging_interval: Duration,
    /// Optional rate limit for this queue. None means unlimited.
    pub rate_limit: Option<RateLimit>,
    /// Minimum guaranteed workers in weighted mode (default: 0).
    pub min_workers: u32,
    /// Weight for overflow allocation in weighted mode (default: 1).
    pub weight: u32,
    /// Number of dispatcher/claimer loops for this queue.
    ///
    /// `max_workers` / `min_workers` still define total queue capacity;
    /// claimers share those permits. Raising this above 1 lets one runtime
    /// use multiple queue-storage claimer leases for hot queues.
    pub claimers: u16,
    /// Maximum jobs a dispatcher attempts to claim in one DB round-trip.
    pub claim_batch_size: usize,
}

impl Default for QueueConfig {
    fn default() -> Self {
        Self {
            max_workers: 50,
            poll_interval: Duration::from_millis(200),
            deadline_duration: Duration::from_secs(300), // 5 minutes
            priority_aging_interval: Duration::from_secs(60),
            rate_limit: None,
            min_workers: 0,
            weight: 1,
            claimers: 1,
            claim_batch_size: DEFAULT_CLAIM_BATCH_LIMIT,
        }
    }
}

/// Wraps permits so the correct resource is released on drop.
/// The OwnedSemaphorePermit fields are held purely for their Drop behavior.
#[allow(dead_code)]
pub(crate) enum DispatchPermit {
    /// Hard-reserved semaphore permit (current default behavior).
    Hard(tokio::sync::OwnedSemaphorePermit),
    /// Local (guaranteed minimum) semaphore permit in weighted mode.
    Local(tokio::sync::OwnedSemaphorePermit),
    /// Overflow permit from the shared OverflowPool.
    Overflow {
        pool: Arc<OverflowPool>,
        queue: String,
    },
}

impl Drop for DispatchPermit {
    fn drop(&mut self) {
        if let DispatchPermit::Overflow { pool, queue } = self {
            pool.release(queue, 1);
        }
        // OwnedSemaphorePermit auto-releases on drop for Hard/Local
    }
}

/// Concurrency mode for a dispatcher.
pub(crate) enum ConcurrencyMode {
    /// Each queue has its own semaphore. No sharing. Default behavior.
    HardReserved { semaphore: Arc<Semaphore> },
    /// Queues share a global overflow pool with per-queue minimum guarantees.
    Weighted {
        local_semaphore: Arc<Semaphore>,
        overflow_pool: Arc<OverflowPool>,
        queue_name: String,
    },
}

/// Centralized overflow capacity allocator for weighted mode.
/// Thread-safe: called from multiple dispatcher poll loops via Mutex.
pub(crate) struct OverflowPool {
    total: u32,
    state: std::sync::Mutex<OverflowState>,
}

struct OverflowState {
    /// Per-queue: currently held overflow permits (decremented on release).
    held: HashMap<String, u32>,
    /// Per-queue: last-declared demand (updated every try_acquire call).
    demand: HashMap<String, u32>,
    /// Per-queue: configured weight (immutable after construction).
    weights: HashMap<String, u32>,
}

impl OverflowPool {
    pub fn new(total: u32, weights: HashMap<String, u32>) -> Self {
        Self {
            total,
            state: std::sync::Mutex::new(OverflowState {
                held: HashMap::new(),
                demand: HashMap::new(),
                weights,
            }),
        }
    }

    /// Try to acquire up to `wanted` overflow permits for `queue`.
    /// Returns the number actually granted (0..=wanted).
    ///
    /// Calling with wanted=0 is valid — it clears this queue's demand signal.
    pub fn try_acquire(&self, queue: &str, wanted: u32) -> u32 {
        let mut state = self.state.lock().unwrap();

        // Always update demand — this is the key signal for fairness
        state.demand.insert(queue.to_string(), wanted);

        if wanted == 0 {
            return 0;
        }

        let currently_used: u32 = state.held.values().sum();
        let available = self.total.saturating_sub(currently_used);
        if available == 0 {
            return 0;
        }

        let my_weight = state.weights.get(queue).copied().unwrap_or(1);

        // Contending = queues with demand > 0 OR held > 0
        let contending_weight: u32 = state
            .weights
            .iter()
            .filter(|(q, _)| {
                state.demand.get(q.as_str()).copied().unwrap_or(0) > 0
                    || state.held.get(q.as_str()).copied().unwrap_or(0) > 0
            })
            .map(|(_, w)| *w)
            .sum();

        if contending_weight == 0 {
            return 0;
        }

        // My fair share of the TOTAL pool (not just available)
        let my_fair_share =
            ((self.total as f64) * (my_weight as f64 / contending_weight as f64)).ceil() as u32;
        let my_held = state.held.get(queue).copied().unwrap_or(0);
        let room = my_fair_share.saturating_sub(my_held);

        let granted = wanted.min(available).min(room);
        if granted > 0 {
            *state.held.entry(queue.to_string()).or_insert(0) += granted;
        }
        granted
    }

    /// Release `n` overflow permits back to the pool.
    pub fn release(&self, queue: &str, n: u32) {
        let mut state = self.state.lock().unwrap();
        if let Some(held) = state.held.get_mut(queue) {
            *held = held.saturating_sub(n);
        }
    }

    /// Get the number of overflow permits currently held by a queue.
    pub fn held(&self, queue: &str) -> u32 {
        let state = self.state.lock().unwrap();
        state.held.get(queue).copied().unwrap_or(0)
    }
}

/// Dispatcher polls a single queue for available jobs and dispatches them.
pub struct Dispatcher {
    queue: String,
    runtime_instance_id: Uuid,
    config: QueueConfig,
    pool: PgPool,
    executor: Arc<JobExecutor>,
    metrics: crate::metrics::AwaMetrics,
    _in_flight: InFlightMap,
    concurrency: ConcurrencyMode,
    alive: Arc<AtomicBool>,
    cancel: CancellationToken,
    job_set: Arc<Mutex<JoinSet<()>>>,
    rate_limiter: Option<Arc<StdMutex<TokenBucket>>>,
    storage: RuntimeStorage,
    capacity_wake: Arc<Notify>,
    claimer_owner_id: Uuid,
}

impl Dispatcher {
    /// Create a dispatcher with a specific concurrency mode (used for weighted mode).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn with_concurrency(
        queue: String,
        runtime_instance_id: Uuid,
        config: QueueConfig,
        pool: PgPool,
        executor: Arc<JobExecutor>,
        metrics: crate::metrics::AwaMetrics,
        in_flight: InFlightMap,
        alive: Arc<AtomicBool>,
        cancel: CancellationToken,
        job_set: Arc<Mutex<JoinSet<()>>>,
        concurrency: ConcurrencyMode,
        rate_limiter: Option<Arc<StdMutex<TokenBucket>>>,
        capacity_wake: Arc<Notify>,
        claimer_owner_id: Uuid,
        storage: RuntimeStorage,
    ) -> Self {
        Self {
            queue,
            runtime_instance_id,
            config,
            pool,
            executor,
            metrics,
            _in_flight: in_flight,
            concurrency,
            alive,
            cancel,
            job_set,
            rate_limiter,
            storage,
            capacity_wake,
            claimer_owner_id,
        }
    }

    /// Run the poll loop. Returns when cancelled.
    pub async fn run(mut self) {
        self.alive.store(true, Ordering::SeqCst);
        info!(
            queue = %self.queue,
            runtime_instance_id = %self.runtime_instance_id,
            claimer_owner_id = %self.claimer_owner_id,
            poll_interval_ms = self.config.poll_interval.as_millis(),
            "Dispatcher started"
        );

        // Set up LISTEN/NOTIFY for this queue
        let notify_channel = format!("awa:{}", self.queue);
        let mut listener = match sqlx::postgres::PgListener::connect_with(&self.pool).await {
            Ok(listener) => listener,
            Err(err) => {
                warn!(
                    queue = %self.queue,
                    error = %err,
                    "Failed to create PG listener, falling back to polling only"
                );
                // Fall back to poll-only mode
                self.poll_loop_only().await;
                self.alive.store(false, Ordering::SeqCst);
                return;
            }
        };

        if let Err(err) = listener.listen(&notify_channel).await {
            warn!(
                queue = %self.queue,
                error = %err,
                channel = %notify_channel,
                "Failed to LISTEN, falling back to polling"
            );
            self.poll_loop_only().await;
            self.alive.store(false, Ordering::SeqCst);
            return;
        }

        debug!(
            queue = %self.queue,
            channel = %notify_channel,
            "Listening for job notifications"
        );

        loop {
            tokio::select! {
                _ = self.cancel.cancelled() => {
                    debug!(queue = %self.queue, "Dispatcher shutting down");
                    break;
                }
                // Wait for either a notification or the poll interval
                notification = listener.recv() => {
                    match notification {
                        Ok(_) => {
                            debug!(queue = %self.queue, "Woken by NOTIFY");
                            self.drain_ready(WakeReason::Notify, Instant::now()).await;
                        }
                        Err(err) => {
                            warn!(
                                queue = %self.queue,
                                error = %err,
                                "PG listener error, will retry"
                            );
                            tokio::time::sleep(Duration::from_secs(1)).await;
                        }
                    }
                }
                _ = self.capacity_wake.notified() => {
                    self.drain_ready(WakeReason::Capacity, Instant::now()).await;
                }
                _ = tokio::time::sleep(self.config.poll_interval) => {
                    self.drain_ready(WakeReason::Poll, Instant::now()).await;
                }
            }
        }

        self.alive.store(false, Ordering::SeqCst);
    }

    /// Poll-only fallback (no LISTEN/NOTIFY).
    async fn poll_loop_only(&mut self) {
        loop {
            tokio::select! {
                _ = self.cancel.cancelled() => {
                    debug!(queue = %self.queue, "Dispatcher (poll-only) shutting down");
                    break;
                }
                _ = self.capacity_wake.notified() => {
                    self.drain_ready(WakeReason::Capacity, Instant::now()).await;
                }
                _ = tokio::time::sleep(self.config.poll_interval) => {
                    self.drain_ready(WakeReason::Poll, Instant::now()).await;
                }
            }
        }
    }

    /// Pre-acquire permits (non-blocking). Returns a vec of permits.
    fn acquire_permits(&mut self) -> Vec<DispatchPermit> {
        let mut permits = Vec::new();
        match &self.concurrency {
            ConcurrencyMode::HardReserved { semaphore } => {
                for _ in 0..self.config.claim_batch_size {
                    match semaphore.clone().try_acquire_owned() {
                        Ok(p) => permits.push(DispatchPermit::Hard(p)),
                        Err(_) => break,
                    }
                }
            }
            ConcurrencyMode::Weighted {
                local_semaphore,
                overflow_pool,
                queue_name,
            } => {
                // First: local (guaranteed) permits
                for _ in 0..self.config.claim_batch_size {
                    match local_semaphore.clone().try_acquire_owned() {
                        Ok(p) => permits.push(DispatchPermit::Local(p)),
                        Err(_) => break,
                    }
                }
                // Then: overflow permits up to the claim batch limit.
                let overflow_wanted =
                    (self.config.claim_batch_size.saturating_sub(permits.len())) as u32;
                let granted = overflow_pool.try_acquire(queue_name, overflow_wanted);
                for _ in 0..granted {
                    permits.push(DispatchPermit::Overflow {
                        pool: overflow_pool.clone(),
                        queue: queue_name.clone(),
                    });
                }
            }
        }
        permits
    }

    fn refund_rate_limit(&self, n: usize) {
        if n == 0 {
            return;
        }
        if let Some(rate_limiter) = &self.rate_limiter {
            let mut rate_limiter = rate_limiter.lock().expect("rate limiter lock poisoned");
            rate_limiter.refund(n as u32);
        }
    }

    /// Drain immediately available work after a wake-up until the queue is empty,
    /// capacity is exhausted, rate limiting stops us, or shutdown is requested.
    async fn drain_ready(&mut self, wake_reason: WakeReason, woke_at: Instant) {
        self.metrics
            .record_dispatch_wake(&self.queue, wake_reason.as_str());
        let mut first_iteration = true;
        while !self.cancel.is_cancelled() {
            let wake_context = first_iteration.then_some((wake_reason, woke_at));
            if !self.poll_once(wake_context).await {
                break;
            }
            first_iteration = false;
        }
    }

    /// Single poll iteration: pre-acquire permits, claim jobs, dispatch.
    ///
    /// Returns whether the caller should keep polling.
    // Debug-level root span: see #449, #455.
    #[tracing::instrument(
        level = "debug",
        parent = None,
        skip(self),
        fields(queue = %self.queue)
    )]
    async fn poll_once(&mut self, wake_context: Option<(WakeReason, Instant)>) -> bool {
        // Phase 1: Pre-acquire permits (non-blocking)
        let mut permits = self.acquire_permits();
        if permits.is_empty() {
            return false;
        }
        if let Some((reason, woke_at)) = wake_context {
            self.metrics.record_dispatch_wake_to_claim(
                &self.queue,
                reason.as_str(),
                woke_at.elapsed(),
            );
            self.metrics.record_dispatch_capacity_available(
                &self.queue,
                reason.as_str(),
                permits.len() as u64,
            );
        }

        // Phase 2: Apply rate limit
        let batch_size = if let Some(rate_limiter) = &self.rate_limiter {
            let mut rate_limiter = rate_limiter.lock().expect("rate limiter lock poisoned");
            let rate_available = rate_limiter.available() as usize;
            let batch_size = permits
                .len()
                .min(rate_available)
                .min(self.config.claim_batch_size);
            if batch_size > 0 {
                rate_limiter.consume(batch_size as u32);
            }
            batch_size
        } else {
            permits.len().min(self.config.claim_batch_size)
        };
        if batch_size == 0 {
            // Drop all permits — rate limited
            if let Some((reason, _)) = wake_context {
                self.metrics
                    .record_dispatch_rate_limited(&self.queue, reason.as_str());
            }
            return false;
        }
        // Release excess permits beyond what rate limit allows
        while permits.len() > batch_size {
            permits.pop(); // Drop releases the permit
        }

        // Phase 3: Claim jobs from DB.
        //
        // Uses a CTE (not a FROM-subquery) so the LIMIT is enforced as a
        // materialization barrier. PostgreSQL's planner can merge a
        // FROM-subquery with the UPDATE target when both reference the same
        // table, which under concurrent load causes the LIMIT to be ignored.
        //
        // Single index scan on idx_awa_jobs_hot_dequeue with FOR UPDATE SKIP LOCKED
        // acquires row locks during the scan. Priority ordering is strict
        // (priority ASC, run_at ASC, id ASC); cross-priority fairness is handled
        // by the maintenance leader's age_waiting_priorities task (ADR-005).
        let deadline_secs = self.config.deadline_duration.as_secs_f64();
        let claim_start = Instant::now();

        let jobs: Vec<DispatchedJob> = match &self.storage {
            RuntimeStorage::Canonical => match sqlx::query_as::<_, JobRow>(
                r#"
                WITH claimed AS (
                    SELECT id
                    FROM awa.jobs_hot
                    WHERE state = 'available'
                      AND queue = $1
                      AND run_at <= now()
                      AND NOT EXISTS (
                          SELECT 1 FROM awa.queue_meta
                          WHERE queue = $1 AND paused = TRUE
                      )
                    ORDER BY priority ASC, run_at ASC, id ASC
                    LIMIT $2
                    FOR UPDATE SKIP LOCKED
                )
                UPDATE awa.jobs_hot
                SET state = 'running',
                    attempt = attempt + 1,
                    run_lease = run_lease + 1,
                    attempted_at = now(),
                    heartbeat_at = now(),
                    deadline_at = now() + make_interval(secs => $3)
                FROM claimed
                WHERE awa.jobs_hot.id = claimed.id
                  AND awa.jobs_hot.state = 'available'
                RETURNING awa.jobs_hot.*
                "#,
            )
            .bind(&self.queue)
            .bind(batch_size as i32)
            .bind(deadline_secs)
            .fetch_all(&self.pool)
            .await
            {
                Ok(jobs) => jobs
                    .into_iter()
                    .map(|job| DispatchedJob {
                        job,
                        queue_storage_claim: None,
                        queue_storage_unique_states: None,
                    })
                    .collect(),
                Err(err) => {
                    warn!(queue = %self.queue, error = %err, "Failed to claim jobs");
                    self.refund_rate_limit(batch_size);
                    return false;
                }
            },
            RuntimeStorage::QueueStorage(runtime) => match runtime
                .store
                .claim_runtime_batch_with_aging_for_instance(
                    &self.pool,
                    &self.queue,
                    batch_size as i64,
                    self.config.deadline_duration,
                    self.config.priority_aging_interval,
                    self.claimer_owner_id,
                    max_claimers_per_queue().max(self.config.claimers as i16),
                    CLAIMER_LEASE_TTL,
                    CLAIMER_IDLE_THRESHOLD,
                )
                .await
            {
                Ok(jobs) => jobs
                    .into_iter()
                    .map(|claimed| DispatchedJob {
                        job: claimed.job,
                        queue_storage_claim: Some(claimed.claim),
                        queue_storage_unique_states: claimed.unique_states,
                    })
                    .collect(),
                Err(err) => {
                    warn!(
                        queue = %self.queue,
                        error = %err,
                        "Failed to claim queue storage jobs"
                    );
                    self.refund_rate_limit(batch_size);
                    return false;
                }
            },
        };
        self.metrics
            .record_claim_batch(&self.queue, jobs.len() as u64, claim_start.elapsed());
        if !jobs.is_empty() {
            // Queue-storage carries an `enqueue_shard` on every claim;
            // bucket the batch so dashboards can sum
            // `awa.job.claimed` by `awa.enqueue.shard` and confirm the
            // claim ordering rotates across shards. Canonical claims
            // have no shard and report through the un-decorated form.
            let mut by_shard: std::collections::BTreeMap<i16, u64> =
                std::collections::BTreeMap::new();
            let mut canonical_count: u64 = 0;
            for dispatched in &jobs {
                match &dispatched.queue_storage_claim {
                    Some(claim) => *by_shard.entry(claim.enqueue_shard).or_default() += 1,
                    None => canonical_count += 1,
                }
            }
            for (shard, count) in by_shard {
                self.metrics
                    .record_job_claimed_by_shard(&self.queue, shard, count);
            }
            if canonical_count > 0 {
                self.metrics
                    .record_job_claimed(&self.queue, canonical_count);
            }
            // Wait duration = created_at → now() (claim moment).
            //
            // Earlier this used `attempted_at - created_at`, but the
            // queue_storage claim path only populates `attempted_at`
            // when deadline_duration > 0. Receipt-plane jobs (zero
            // deadline, the 0.6 default) get NULL attempted_at on the
            // ClaimedRuntimeJob, so the previous version silently
            // skipped the metric — it never appeared on the
            // dashboard. Falling back to `Utc::now()` measures the
            // same operator-visible quantity ("time from enqueue to
            // claim") regardless of whether the receipt-plane
            // optimisation skipped the attempted_at write.
            let now = chrono::Utc::now();
            for job in &jobs {
                let claim_at = job.job.attempted_at.unwrap_or(now);
                let wait_secs = (claim_at - job.job.created_at).num_milliseconds() as f64 / 1000.0;
                if wait_secs >= 0.0 {
                    self.metrics.record_wait_duration(&self.queue, wait_secs);
                }
            }
        }
        if jobs.len() < batch_size {
            self.refund_rate_limit(batch_size - jobs.len());
        }

        // Phase 4: Release excess permits if DB had fewer jobs
        let unused_permits = permits.len().saturating_sub(jobs.len());
        while permits.len() > jobs.len() {
            permits.pop();
        }
        if unused_permits > 0 {
            self.metrics
                .record_dispatch_unused_permits(&self.queue, unused_permits as u64);
        }

        // Phase 5: Clear overflow demand if no jobs found
        if jobs.is_empty() {
            if let Some((reason, _)) = wake_context {
                self.metrics
                    .record_dispatch_empty_claim(&self.queue, reason.as_str());
            }
            if let ConcurrencyMode::Weighted {
                overflow_pool,
                queue_name,
                ..
            } = &self.concurrency
            {
                overflow_pool.try_acquire(queue_name, 0);
            }
            return false;
        }

        debug!(queue = %self.queue, count = jobs.len(), "Claimed jobs");

        // Phase 6: Dispatch (each job takes one pre-acquired permit)
        let mut set = self.job_set.lock().await;
        // Reap completed task handles. JoinSet retains JoinHandles (and the
        // task Cell they keep alive) until join_next() consumes them; under
        // steady-state load that pins the entire execute_task closure
        // captures and leaks ~3 GB/h/replica. Drain here so the set only
        // holds in-flight tasks.
        while set.try_join_next().is_some() {}
        for (job, permit) in jobs.into_iter().zip(permits) {
            let cancel_flag = Arc::new(AtomicBool::new(false));
            let task = self.executor.execute_task(job, cancel_flag);
            let capacity_wake = self.capacity_wake.clone();
            set.spawn(async move {
                task.await;
                drop(permit);
                capacity_wake.notify_one();
            });
        }

        true
    }
}