awa-worker 0.6.2

use crate::runtime::RunLease;
use crate::storage::RuntimeStorage;
use awa_model::{AwaError, ClaimedEntry, ClaimedRuntimeJob};
use sqlx::PgPool;
use std::collections::HashSet;
use std::env;
use std::time::Duration;
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use tracing::{debug, warn};

// Default completion batch size and flush cadence. Queue storage short-job
// completion uses a fused receipt-close + terminal-insert statement, so the
// default can keep finalization latency low while still batching enough work
// to amortise the durable completion path. Queue storage starts with one
// completion shard for ordinary runtimes, then uses four shards for very
// large single-runtime worker pools where one serialized flusher cannot keep
// permits full. Extra flushers can multiply fleet-wide contention in
// multi-process deployments. Tunable per-deployment via
// `AWA_COMPLETION_BATCH_SIZE`, `AWA_COMPLETION_FLUSH_MS`, and
// `AWA_COMPLETION_SHARDS`.
const COMPLETION_BATCH_SIZE: usize = 512;
const COMPLETION_FLUSH_INTERVAL: Duration = Duration::from_millis(1);
const COMPLETION_CHANNEL_CAPACITY: usize = 4096;
const QUEUE_STORAGE_COMPLETION_SHARD_WORKER_THRESHOLD: u32 = 512;

fn default_completion_shards(storage: &RuntimeStorage, runtime_worker_capacity: u32) -> usize {
    match storage {
        RuntimeStorage::Canonical => 8,
        RuntimeStorage::QueueStorage(_) => {
            if runtime_worker_capacity >= QUEUE_STORAGE_COMPLETION_SHARD_WORKER_THRESHOLD {
                4
            } else {
                1
            }
        }
    }
}

fn completion_batch_size() -> usize {
    env::var("AWA_COMPLETION_BATCH_SIZE")
        .ok()
        .and_then(|value| value.parse::<usize>().ok())
        .filter(|value| *value > 0)
        .unwrap_or(COMPLETION_BATCH_SIZE)
}

fn completion_flush_interval() -> Duration {
    env::var("AWA_COMPLETION_FLUSH_MS")
        .ok()
        .and_then(|value| value.parse::<u64>().ok())
        .filter(|value| *value > 0)
        .map(Duration::from_millis)
        .unwrap_or(COMPLETION_FLUSH_INTERVAL)
}

fn completion_shards(storage: &RuntimeStorage, runtime_worker_capacity: u32) -> usize {
    env::var("AWA_COMPLETION_SHARDS")
        .ok()
        .and_then(|value| value.parse::<usize>().ok())
        .filter(|value| *value > 0)
        .unwrap_or_else(|| default_completion_shards(storage, runtime_worker_capacity))
}
const COMPLETE_BATCH_SQL: &str = r#"
    WITH completed (id, run_lease) AS (
        SELECT * FROM unnest($1::bigint[], $2::bigint[])
    ),
    locked AS (
        SELECT jobs.ctid, jobs.id, jobs.run_lease
        FROM awa.jobs_hot AS jobs
        JOIN completed
          ON jobs.id = completed.id
         AND jobs.run_lease = completed.run_lease
        WHERE jobs.state = 'running'
        ORDER BY jobs.id
        FOR UPDATE OF jobs
    )
    UPDATE awa.jobs_hot AS jobs
    SET state = 'completed',
        finalized_at = now(),
        progress = NULL
    FROM locked
    WHERE jobs.ctid = locked.ctid
    RETURNING locked.id, locked.run_lease
"#;

struct CompletionRequest {
    job_id: i64,
    run_lease: RunLease,
    claim: Option<ClaimedEntry>,
    runtime_job: Option<ClaimedRuntimeJob>,
    response: oneshot::Sender<Result<bool, AwaError>>,
}

fn completion_sort_key(request: &CompletionRequest) -> (i32, i64, RunLease) {
    let claim_slot = request
        .runtime_job
        .as_ref()
        .map(|runtime_job| runtime_job.claim.claim_slot)
        .or_else(|| request.claim.as_ref().map(|claim| claim.claim_slot))
        .unwrap_or(i32::MAX);
    (claim_slot, request.job_id, request.run_lease)
}

#[derive(Clone)]
pub(crate) struct CompletionBatcherHandle {
    shards: Vec<mpsc::Sender<CompletionRequest>>,
}

impl CompletionBatcherHandle {
    pub async fn complete(&self, job_id: i64, run_lease: RunLease) -> Result<bool, AwaError> {
        self.complete_inner(job_id, run_lease, None, None).await
    }

    pub async fn complete_runtime_job(
        &self,
        runtime_job: ClaimedRuntimeJob,
    ) -> Result<bool, AwaError> {
        self.complete_inner(
            runtime_job.job.id,
            runtime_job.job.run_lease,
            Some(runtime_job.claim.clone()),
            Some(runtime_job),
        )
        .await
    }

    async fn complete_inner(
        &self,
        job_id: i64,
        run_lease: RunLease,
        claim: Option<ClaimedEntry>,
        runtime_job: Option<ClaimedRuntimeJob>,
    ) -> Result<bool, AwaError> {
        let shard = (job_id.rem_euclid(self.shards.len() as i64)) as usize;
        let (response_tx, response_rx) = oneshot::channel();
        self.shards[shard]
            .send(CompletionRequest {
                job_id,
                run_lease,
                claim,
                runtime_job,
                response: response_tx,
            })
            .await
            .map_err(|_| AwaError::Validation("completion batcher stopped".into()))?;

        response_rx
            .await
            .map_err(|_| AwaError::Validation("completion batcher dropped response".into()))?
    }
}

pub(crate) struct CompletionBatcher {
    workers: Vec<CompletionWorker>,
}

impl CompletionBatcher {
    pub fn new(
        pool: PgPool,
        cancel: CancellationToken,
        metrics: crate::metrics::AwaMetrics,
        storage: RuntimeStorage,
        runtime_worker_capacity: u32,
    ) -> (Self, CompletionBatcherHandle) {
        let shard_count = completion_shards(&storage, runtime_worker_capacity);
        let batch_size = completion_batch_size();
        let flush_interval = completion_flush_interval();
        let mut shards = Vec::with_capacity(shard_count);
        let mut workers = Vec::with_capacity(shard_count);

        for shard_id in 0..shard_count {
            let (tx, rx) = mpsc::channel(COMPLETION_CHANNEL_CAPACITY);
            shards.push(tx);
            workers.push(CompletionWorker {
                shard_id,
                shard_count,
                pool: pool.clone(),
                rx,
                cancel: cancel.clone(),
                metrics: metrics.clone(),
                storage: storage.clone(),
                batch_size,
                flush_interval,
            });
        }

        (Self { workers }, CompletionBatcherHandle { shards })
    }

    pub fn spawn(self) -> Vec<tokio::task::JoinHandle<()>> {
        self.workers
            .into_iter()
            .map(|worker| tokio::spawn(async move { worker.run().await }))
            .collect()
    }
}

struct CompletionWorker {
    shard_id: usize,
    shard_count: usize,
    pool: PgPool,
    rx: mpsc::Receiver<CompletionRequest>,
    cancel: CancellationToken,
    metrics: crate::metrics::AwaMetrics,
    storage: RuntimeStorage,
    batch_size: usize,
    flush_interval: Duration,
}

impl CompletionWorker {
    async fn run(mut self) {
        let mut pending = Vec::with_capacity(self.batch_size);

        loop {
            if pending.is_empty() {
                tokio::select! {
                    _ = self.cancel.cancelled() => break,
                    request = self.rx.recv() => {
                        match request {
                            Some(request) => pending.push(request),
                            None => break,
                        }
                    }
                }
            } else {
                let timer = tokio::time::sleep(self.flush_interval);
                tokio::pin!(timer);

                tokio::select! {
                    _ = self.cancel.cancelled() => break,
                    _ = &mut timer => {
                        self.flush(&mut pending).await;
                    }
                    request = self.rx.recv() => {
                        match request {
                            Some(request) => {
                                pending.push(request);
                                if pending.len() >= self.batch_size {
                                    self.flush(&mut pending).await;
                                }
                            }
                            None => break,
                        }
                    }
                }
            }
        }

        while let Ok(request) = self.rx.try_recv() {
            pending.push(request);
            if pending.len() >= self.batch_size {
                self.flush(&mut pending).await;
            }
        }
        if !pending.is_empty() {
            self.flush(&mut pending).await;
        }
        debug!(shard = self.shard_id, "Completion batcher shard stopped");
    }

    #[tracing::instrument(
        skip(self, pending),
        fields(shard = self.shard_id, batch_size = pending.len())
    )]
    async fn flush(&self, pending: &mut Vec<CompletionRequest>) {
        if pending.is_empty() {
            return;
        }

        let mut batch: Vec<_> = std::mem::take(pending);
        batch.sort_unstable_by_key(completion_sort_key);
        let job_ids: Vec<i64> = batch.iter().map(|request| request.job_id).collect();
        // Completion is routed to a shard by job_id, and each shard completes
        // its batch independently, so a misrouted job would let two shards race
        // the same job's terminal write. The routing is a pure function of
        // job_id, so this can only fire on a routing-logic regression.
        debug_assert!(
            job_ids
                .iter()
                .all(|id| id.rem_euclid(self.shard_count as i64) as usize == self.shard_id),
            "completion shard {} received a job routed to another shard (shard_count {})",
            self.shard_id,
            self.shard_count,
        );
        let run_leases: Vec<i64> = batch.iter().map(|request| request.run_lease).collect();
        let flush_start = std::time::Instant::now();

        let updated = match &self.storage {
            RuntimeStorage::Canonical => sqlx::query_as::<_, (i64, i64)>(COMPLETE_BATCH_SQL)
                .bind(&job_ids)
                .bind(&run_leases)
                .fetch_all(&self.pool)
                .await
                .map_err(AwaError::Database),
            RuntimeStorage::QueueStorage(runtime) => {
                if let Some(runtime_jobs) = batch
                    .iter()
                    .map(|request| request.runtime_job.clone())
                    .collect::<Option<Vec<ClaimedRuntimeJob>>>()
                {
                    runtime
                        .store
                        .complete_runtime_batch(&self.pool, &runtime_jobs)
                        .await
                } else if let Some(claimed) = batch
                    .iter()
                    .map(|request| request.claim.clone())
                    .collect::<Option<Vec<ClaimedEntry>>>()
                {
                    runtime
                        .store
                        .complete_claimed_batch(&self.pool, &claimed)
                        .await
                } else {
                    runtime
                        .store
                        .complete_job_batch_by_id(
                            &self.pool,
                            &job_ids
                                .iter()
                                .copied()
                                .zip(run_leases.iter().copied())
                                .collect::<Vec<_>>(),
                        )
                        .await
                }
            }
        };

        match updated {
            Ok(updated_rows) => {
                let updated: HashSet<(i64, i64)> = updated_rows.into_iter().collect();
                let updated_count = updated.len();
                self.metrics.record_completion_flush(
                    self.shard_id,
                    job_ids.len() as u64,
                    flush_start.elapsed(),
                );
                for request in batch {
                    let _ = request
                        .response
                        .send(Ok(updated.contains(&(request.job_id, request.run_lease))));
                }
                debug!(
                    shard = self.shard_id,
                    batch_size = job_ids.len(),
                    updated = updated_count,
                    "Flushed completed job batch"
                );
            }
            Err(err) => {
                self.metrics.record_completion_flush(
                    self.shard_id,
                    job_ids.len() as u64,
                    flush_start.elapsed(),
                );
                warn!(
                    error = %err,
                    shard = self.shard_id,
                    batch_size = job_ids.len(),
                    "Failed to flush completed job batch"
                );
                let message = format!("completion batch flush failed: {err}");
                for request in batch {
                    let _ = request
                        .response
                        .send(Err(AwaError::Validation(message.clone())));
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use awa_model::migrations;
    use sqlx::postgres::PgPoolOptions;
    use std::sync::Arc;
    use std::time::{Duration, Instant};

    fn base_database_url() -> String {
        std::env::var("DATABASE_URL")
            .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string())
    }

    fn replace_database_name(url: &str, database_name: &str) -> String {
        let (without_query, query_suffix) = match url.split_once('?') {
            Some((prefix, query)) => (prefix, Some(query)),
            None => (url, None),
        };
        let (base, _) = without_query
            .rsplit_once('/')
            .expect("database URL should include a database name");
        let mut out = format!("{base}/{database_name}");
        if let Some(query) = query_suffix {
            out.push('?');
            out.push_str(query);
        }
        out
    }

    fn database_url() -> String {
        std::env::var("DATABASE_URL_WORKER_COMPLETION").unwrap_or_else(|_| {
            replace_database_name(&base_database_url(), "awa_test_worker_completion")
        })
    }

    async fn ensure_database_exists(url: &str) {
        let database_name = url
            .split_once('?')
            .map(|(prefix, _)| prefix)
            .unwrap_or(url)
            .rsplit_once('/')
            .map(|(_, database_name)| database_name.to_string())
            .expect("database URL should include a database name");
        let admin_url = replace_database_name(url, "postgres");
        let admin_pool = PgPoolOptions::new()
            .max_connections(1)
            .connect(&admin_url)
            .await
            .expect("Failed to connect to admin database for completion tests");
        let create_sql = format!("CREATE DATABASE {database_name}");
        match sqlx::query(&create_sql).execute(&admin_pool).await {
            Ok(_) => {}
            Err(sqlx::Error::Database(db_err)) if db_err.code().as_deref() == Some("42P04") => {}
            Err(sqlx::Error::Database(db_err)) if db_err.code().as_deref() == Some("23505") => {}
            Err(err) => {
                panic!("Failed to create completion test database {database_name}: {err}")
            }
        }
    }

    async fn setup(max_conns: u32) -> PgPool {
        let url = database_url();
        ensure_database_exists(&url).await;
        let pool = PgPoolOptions::new()
            .max_connections(max_conns)
            .acquire_timeout(Duration::from_secs(5))
            .connect(&url)
            .await
            .expect("Failed to connect to database");
        migrations::run(&pool).await.expect("Failed to migrate");
        pool
    }

    fn completion_request_with_claim_slot(
        job_id: i64,
        run_lease: i64,
        claim_slot: i32,
    ) -> CompletionRequest {
        let (response, _rx) = oneshot::channel();
        CompletionRequest {
            job_id,
            run_lease,
            claim: Some(ClaimedEntry {
                queue: "default".to_string(),
                priority: 0,
                lane_seq: job_id,
                ready_slot: 0,
                ready_generation: 0,
                lease_slot: 0,
                lease_generation: 0,
                claim_slot,
                receipt_id: Some(job_id),
                claim_batch_id: None,
                claim_batch_index: None,
                lease_claim_receipt: true,
                enqueue_shard: 0,
            }),
            runtime_job: None,
            response,
        }
    }

    #[test]
    fn completion_sort_groups_receipt_slots_before_job_id() {
        let mut requests = [
            completion_request_with_claim_slot(10, 1, 3),
            completion_request_with_claim_slot(2, 1, 1),
            completion_request_with_claim_slot(1, 1, 3),
        ];

        requests.sort_unstable_by_key(completion_sort_key);

        let ordered: Vec<_> = requests
            .iter()
            .map(|request| {
                (
                    request.claim.as_ref().unwrap().claim_slot,
                    request.job_id,
                    request.run_lease,
                )
            })
            .collect();
        assert_eq!(ordered, vec![(1, 2, 1), (3, 1, 1), (3, 10, 1)]);
    }

    #[test]
    fn queue_storage_completion_shards_scale_with_runtime_capacity() {
        let runtime = crate::storage::QueueStorageRuntime::new(
            awa_model::QueueStorageConfig::default(),
            Duration::from_millis(1_000),
            Duration::from_millis(50),
        )
        .expect("queue storage runtime config should be valid");
        let high_capacity_runtime = runtime.clone();

        assert_eq!(
            default_completion_shards(&crate::storage::RuntimeStorage::Canonical, 50),
            8
        );
        assert_eq!(
            default_completion_shards(&crate::storage::RuntimeStorage::QueueStorage(runtime), 50),
            1
        );
        assert_eq!(
            default_completion_shards(
                &crate::storage::RuntimeStorage::QueueStorage(high_capacity_runtime),
                QUEUE_STORAGE_COMPLETION_SHARD_WORKER_THRESHOLD,
            ),
            4
        );
    }

    #[test]
    fn completion_defaults_use_throughput_oriented_batch_size() {
        assert_eq!(COMPLETION_BATCH_SIZE, 512);
        assert_eq!(COMPLETION_FLUSH_INTERVAL, Duration::from_millis(1));
    }

    async fn clean_queue(pool: &PgPool, queue: &str) {
        sqlx::query("DELETE FROM awa.jobs_hot WHERE queue = $1")
            .bind(queue)
            .execute(pool)
            .await
            .expect("Failed to clean canonical queue jobs");
        sqlx::query("DELETE FROM awa.jobs WHERE queue = $1")
            .bind(queue)
            .execute(pool)
            .await
            .expect("Failed to clean queue jobs");
        sqlx::query("DELETE FROM awa.queue_meta WHERE queue = $1")
            .bind(queue)
            .execute(pool)
            .await
            .expect("Failed to clean queue meta");
    }

    async fn seed_running_jobs(pool: &PgPool, queue: &str, total_jobs: i64) -> Vec<(i64, i64)> {
        sqlx::query(
            r#"
            INSERT INTO awa.jobs_hot (
                kind, queue, args, state, priority, attempt, run_lease,
                max_attempts, run_at, attempted_at, heartbeat_at, deadline_at,
                metadata, tags
            )
            SELECT
                'bench_job',
                $1,
                jsonb_build_object('seq', g),
                'running'::awa.job_state,
                2,
                1,
                1,
                25,
                now(),
                now(),
                now(),
                now() + interval '5 minutes',
                '{}'::jsonb,
                '{}'::text[]
            FROM generate_series(1, $2) AS g
            "#,
        )
        .bind(queue)
        .bind(total_jobs)
        .execute(pool)
        .await
        .expect("Failed to seed running jobs");

        sqlx::query_as::<_, (i64, i64)>(
            "SELECT id, run_lease FROM awa.jobs_hot WHERE queue = $1 ORDER BY id ASC",
        )
        .bind(queue)
        .fetch_all(pool)
        .await
        .expect("Failed to load seeded rows")
    }

    async fn reset_running_jobs(pool: &PgPool, queue: &str) {
        sqlx::query(
            r#"
            UPDATE awa.jobs_hot
            SET state = 'running',
                finalized_at = NULL,
                heartbeat_at = now(),
                progress = NULL
            WHERE queue = $1
            "#,
        )
        .bind(queue)
        .execute(pool)
        .await
        .expect("Failed to reset running jobs");
    }

    async fn set_run_lease(pool: &PgPool, job_id: i64, run_lease: i64) {
        sqlx::query(
            r#"
            UPDATE awa.jobs_hot
            SET run_lease = $2,
                attempt = GREATEST(attempt, $2::smallint),
                heartbeat_at = now(),
                finalized_at = NULL,
                progress = NULL
            WHERE id = $1
            "#,
        )
        .bind(job_id)
        .bind(run_lease)
        .execute(pool)
        .await
        .expect("Failed to set run lease");
    }

    async fn complete_jobs_in_lock_order(
        pool: &PgPool,
        first_id: i64,
        first_lease: i64,
        second_id: i64,
        second_lease: i64,
    ) -> Vec<(i64, i64)> {
        sqlx::query_as::<_, (i64, i64)>(COMPLETE_BATCH_SQL)
            .bind(vec![second_id, first_id])
            .bind(vec![second_lease, first_lease])
            .fetch_all(pool)
            .await
            .expect("completion sweep failed")
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn completion_and_heartbeat_batches_do_not_deadlock_with_reversed_input_order() {
        let pool = setup(10).await;
        let queue = "test_completion_heartbeat_lock_order";
        clean_queue(&pool, queue).await;

        let jobs = seed_running_jobs(&pool, queue, 2).await;
        let (first_id, first_lease) = jobs[0];
        let (second_id, second_lease) = jobs[1];

        for iteration in 0..20 {
            if iteration > 0 {
                reset_running_jobs(&pool, queue).await;
            }

            let barrier = Arc::new(tokio::sync::Barrier::new(3));

            let completion_pool = pool.clone();
            let completion_barrier = barrier.clone();
            let completion_task = tokio::spawn(async move {
                completion_barrier.wait().await;
                complete_jobs_in_lock_order(
                    &completion_pool,
                    first_id,
                    first_lease,
                    second_id,
                    second_lease,
                )
                .await
            });

            let heartbeat_pool = pool.clone();
            let heartbeat_barrier = barrier.clone();
            let heartbeat_task = tokio::spawn(async move {
                heartbeat_barrier.wait().await;
                sqlx::query(
                    r#"
                    WITH inflight AS (
                        SELECT * FROM unnest($1::bigint[], $2::bigint[], $3::jsonb[]) AS v(id, run_lease, progress)
                    ),
                    locked AS (
                        SELECT jobs.ctid, inflight.progress
                        FROM awa.jobs_hot AS jobs
                        JOIN inflight
                          ON jobs.id = inflight.id
                         AND jobs.run_lease = inflight.run_lease
                        WHERE jobs.state = 'running'
                        ORDER BY jobs.id
                        FOR UPDATE OF jobs
                    )
                    UPDATE awa.jobs_hot AS jobs
                    SET heartbeat_at = now(),
                        progress = locked.progress
                    FROM locked
                    WHERE jobs.ctid = locked.ctid
                    "#,
                )
                .bind(vec![first_id, second_id])
                .bind(vec![first_lease, second_lease])
                .bind(vec![
                    serde_json::json!({ "iteration": iteration, "job": 1 }),
                    serde_json::json!({ "iteration": iteration, "job": 2 }),
                ])
                .execute(&heartbeat_pool)
                .await
            });

            barrier.wait().await;

            let (completed_ids, heartbeat_result) =
                tokio::time::timeout(Duration::from_secs(5), async move {
                    let completed_ids = completion_task.await.expect("completion task panicked");
                    let heartbeat_result = heartbeat_task.await.expect("heartbeat task panicked");
                    (completed_ids, heartbeat_result)
                })
                .await
                .expect("concurrent completion/heartbeat batches timed out");

            heartbeat_result.expect("heartbeat batch query failed");
            assert!(
                completed_ids.len() <= 2,
                "completion batch returned too many rows in iteration {iteration}"
            );

            complete_jobs_in_lock_order(&pool, first_id, first_lease, second_id, second_lease)
                .await;

            let completed: i64 = sqlx::query_scalar(
                "SELECT count(*) FROM awa.jobs_hot WHERE queue = $1 AND state = 'completed'",
            )
            .bind(queue)
            .fetch_one(&pool)
            .await
            .expect("Failed to count completed rows");
            assert_eq!(
                completed, 2,
                "expected both rows completed after concurrent batch iteration {iteration}"
            );
        }
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn stale_completion_in_same_batch_returns_false_for_old_run_lease() {
        let pool = setup(10).await;
        let queue = "test_completion_stale_lease_batch";
        clean_queue(&pool, queue).await;

        let jobs = seed_running_jobs(&pool, queue, 1).await;
        let (job_id, stale_lease) = jobs[0];
        let current_lease = stale_lease + 1;
        set_run_lease(&pool, job_id, current_lease).await;

        let (batcher, handle) = CompletionBatcher::new(
            pool.clone(),
            CancellationToken::new(),
            crate::metrics::AwaMetrics::from_global(),
            crate::storage::RuntimeStorage::Canonical,
            2,
        );
        let workers = batcher.spawn();

        let barrier = Arc::new(tokio::sync::Barrier::new(3));

        let stale_handle = handle.clone();
        let stale_barrier = barrier.clone();
        let stale_task = tokio::spawn(async move {
            stale_barrier.wait().await;
            stale_handle.complete(job_id, stale_lease).await
        });

        let current_handle = handle.clone();
        let current_barrier = barrier.clone();
        let current_task = tokio::spawn(async move {
            current_barrier.wait().await;
            current_handle.complete(job_id, current_lease).await
        });

        barrier.wait().await;

        let stale_result = stale_task
            .await
            .expect("stale completion task panicked")
            .expect("stale completion request failed");
        let current_result = current_task
            .await
            .expect("current completion task panicked")
            .expect("current completion request failed");

        assert!(
            !stale_result,
            "stale completion should be reported as ignored"
        );
        assert!(
            current_result,
            "current completion should be reported as applied"
        );

        let row: (String, i64) =
            sqlx::query_as("SELECT state::text, run_lease FROM awa.jobs_hot WHERE id = $1")
                .bind(job_id)
                .fetch_one(&pool)
                .await
                .expect("Failed to load completed row");
        assert_eq!(row.0, "completed");
        assert_eq!(row.1, current_lease);

        for worker in workers {
            worker.abort();
        }
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn completion_and_heartbeat_only_batches_do_not_deadlock_with_reversed_input_order() {
        let pool = setup(10).await;
        let queue = "test_completion_heartbeat_only_lock_order";
        clean_queue(&pool, queue).await;

        let jobs = seed_running_jobs(&pool, queue, 2).await;
        let (first_id, first_lease) = jobs[0];
        let (second_id, second_lease) = jobs[1];

        for iteration in 0..20 {
            if iteration > 0 {
                reset_running_jobs(&pool, queue).await;
            }

            let barrier = Arc::new(tokio::sync::Barrier::new(3));

            let completion_pool = pool.clone();
            let completion_barrier = barrier.clone();
            let completion_task = tokio::spawn(async move {
                completion_barrier.wait().await;
                complete_jobs_in_lock_order(
                    &completion_pool,
                    first_id,
                    first_lease,
                    second_id,
                    second_lease,
                )
                .await
            });

            let heartbeat_pool = pool.clone();
            let heartbeat_barrier = barrier.clone();
            let heartbeat_task = tokio::spawn(async move {
                heartbeat_barrier.wait().await;
                sqlx::query(
                    r#"
                    WITH inflight AS (
                        SELECT * FROM unnest($1::bigint[], $2::bigint[]) AS v(id, run_lease)
                    ),
                    locked AS (
                        SELECT jobs.ctid
                        FROM awa.jobs_hot AS jobs
                        JOIN inflight
                          ON jobs.id = inflight.id
                         AND jobs.run_lease = inflight.run_lease
                        WHERE jobs.state = 'running'
                        ORDER BY jobs.id
                        FOR UPDATE OF jobs
                    )
                    UPDATE awa.jobs_hot AS jobs
                    SET heartbeat_at = now()
                    FROM locked
                    WHERE jobs.ctid = locked.ctid
                    "#,
                )
                .bind(vec![first_id, second_id])
                .bind(vec![first_lease, second_lease])
                .execute(&heartbeat_pool)
                .await
            });

            barrier.wait().await;

            let (completed_ids, heartbeat_result) =
                tokio::time::timeout(Duration::from_secs(5), async move {
                    let completed_ids = completion_task.await.expect("completion task panicked");
                    let heartbeat_result = heartbeat_task.await.expect("heartbeat task panicked");
                    (completed_ids, heartbeat_result)
                })
                .await
                .expect("concurrent completion/heartbeat-only batches timed out");

            heartbeat_result.expect("heartbeat-only batch query failed");
            assert!(
                completed_ids.len() <= 2,
                "completion batch returned too many rows in iteration {iteration}"
            );

            complete_jobs_in_lock_order(&pool, first_id, first_lease, second_id, second_lease)
                .await;

            let completed: i64 = sqlx::query_scalar(
                "SELECT count(*) FROM awa.jobs_hot WHERE queue = $1 AND state = 'completed'",
            )
            .bind(queue)
            .fetch_one(&pool)
            .await
            .expect("Failed to count completed rows");
            assert_eq!(
                completed, 2,
                "expected both rows completed after concurrent heartbeat-only iteration {iteration}"
            );
        }
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    #[ignore]
    async fn benchmark_completion_batcher_ack_throughput() {
        let pool = setup(20).await;
        let queue = "bench_completion_batcher";
        clean_queue(&pool, queue).await;

        let jobs = seed_running_jobs(&pool, queue, 10_000).await;
        let (batcher, handle) = CompletionBatcher::new(
            pool.clone(),
            CancellationToken::new(),
            crate::metrics::AwaMetrics::from_global(),
            crate::storage::RuntimeStorage::Canonical,
            2,
        );
        let workers = batcher.spawn();

        let start = Instant::now();
        let mut set = tokio::task::JoinSet::new();
        for (job_id, run_lease) in jobs {
            let handle = handle.clone();
            set.spawn(async move { handle.complete(job_id, run_lease).await });
        }

        let mut completed = 0usize;
        while let Some(result) = set.join_next().await {
            let updated = result
                .expect("completion task panicked")
                .expect("completion failed");
            assert!(updated, "Expected completion batcher to update row");
            completed += 1;
        }
        let elapsed = start.elapsed();

        let completed_rows: i64 = sqlx::query_scalar(
            "SELECT count(*) FROM awa.jobs WHERE queue = $1 AND state = 'completed'",
        )
        .bind(queue)
        .fetch_one(&pool)
        .await
        .unwrap();

        println!(
            "[completion] batcher_ack completed={} in {:.3}s ({:.0}/s)",
            completed,
            elapsed.as_secs_f64(),
            completed as f64 / elapsed.as_secs_f64()
        );
        assert_eq!(completed_rows, completed as i64);

        for worker in workers {
            worker.abort();
        }
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    #[ignore]
    async fn benchmark_completion_direct_batched_query_throughput() {
        let pool = setup(20).await;
        let queue = "bench_completion_direct";
        clean_queue(&pool, queue).await;

        let jobs = seed_running_jobs(&pool, queue, 10_000).await;
        let start = Instant::now();
        for chunk in jobs.chunks(COMPLETION_BATCH_SIZE) {
            let job_ids: Vec<i64> = chunk.iter().map(|(job_id, _)| *job_id).collect();
            let run_leases: Vec<i64> = chunk.iter().map(|(_, run_lease)| *run_lease).collect();
            let updated: Vec<i64> = sqlx::query_scalar(
                r#"
                WITH completed (id, run_lease) AS (
                    SELECT * FROM unnest($1::bigint[], $2::bigint[])
                )
                UPDATE awa.jobs_hot AS jobs
                SET state = 'completed',
                    finalized_at = now()
                FROM completed
                WHERE jobs.id = completed.id
                  AND jobs.run_lease = completed.run_lease
                  AND jobs.state = 'running'
                RETURNING jobs.id
                "#,
            )
            .bind(&job_ids)
            .bind(&run_leases)
            .fetch_all(&pool)
            .await
            .expect("Direct completion batch failed");
            assert_eq!(updated.len(), chunk.len());
        }
        let elapsed = start.elapsed();

        println!(
            "[completion] direct_batched_sql completed={} in {:.3}s ({:.0}/s)",
            jobs.len(),
            elapsed.as_secs_f64(),
            jobs.len() as f64 / elapsed.as_secs_f64()
        );
    }
}