later 0.0.49

#![cfg(feature = "sqlite")]

#[tokio::test]
async fn graceful_shutdown_commits_the_running_partition_head() -> anyhow::Result<()> {
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let gate = Arc::new(tokio::sync::Notify::new());
    let server = start_server_with_flaky(
        "partition-graceful-shutdown",
        pool.clone(),
        Arc::new(Mutex::new(Vec::new())),
        Arc::new(Mutex::new(None)),
        Arc::new(Mutex::new(std::collections::HashMap::new())),
        gate.clone(),
        1,
    )
    .await;
    server.enqueue_to_partition(GateJob).await?;

    let deadline = std::time::Instant::now() + Duration::from_secs(5);
    loop {
        let leases: i64 = sqlx::query_scalar(
            "SELECT COUNT(*) FROM later_partition_lease \
             WHERE namespace = 'partition-graceful-shutdown'",
        )
        .fetch_one(&pool)
        .await?;
        if leases == 1 {
            break;
        }
        if std::time::Instant::now() >= deadline {
            return Err(anyhow::anyhow!(
                "partition handler did not acquire its head"
            ));
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }
    gate.notify_one();
    server.shutdown(Duration::from_secs(2)).await?;
    let remaining: i64 = sqlx::query_scalar(
        "SELECT COUNT(*) FROM later_partition_job \
         WHERE namespace = 'partition-graceful-shutdown'",
    )
    .fetch_one(&pool)
    .await?;
    assert_eq!(remaining, 0);
    Ok(())
}

use later::topic::{partition_for_key, JobPartition, PartitionId, PartitionKey, TopicConfig};
use later::{BackgroundJobServer, Config};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;

const TOPIC_PARTITION_COUNT: u32 = 4;

/// Returns `n` key values that are guaranteed to hash to distinct partitions
/// under `TOPIC_PARTITION_COUNT`, so tests that need independent partitions
/// don't rely on two arbitrary keys happening not to collide.
fn distinct_keys(n: usize) -> Vec<u64> {
    let mut seen_partitions = std::collections::HashSet::new();
    let mut chosen = Vec::new();
    let mut candidate = 0u64;
    while chosen.len() < n {
        let partition = partition_for_key(&PartitionKey::from(candidate), TOPIC_PARTITION_COUNT);
        if seen_partitions.insert(partition) {
            chosen.push(candidate);
        }
        candidate += 1;
    }
    chosen
}

#[derive(Serialize, Deserialize, Debug, Clone)]
struct OrderedJob {
    key: u64,
    seq: u32,
}

impl JobPartition for OrderedJob {
    fn partition_key(&self) -> PartitionKey {
        PartitionKey::from(self.key)
    }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
struct UntopicedJob {
    key: u64,
}

impl JobPartition for UntopicedJob {
    fn partition_key(&self) -> PartitionKey {
        PartitionKey::from(self.key)
    }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
struct FlakyJob {
    key: u64,
    seq: u32,
}

impl JobPartition for FlakyJob {
    fn partition_key(&self) -> PartitionKey {
        PartitionKey::from(self.key)
    }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
struct GateJob;

impl JobPartition for GateJob {
    fn partition_key(&self) -> PartitionKey {
        PartitionKey::from(0u64)
    }
}

#[derive(Serialize, Deserialize, Debug, Clone)]
struct OtherTopicJob {
    key: u64,
    seq: u32,
}

impl JobPartition for OtherTopicJob {
    fn partition_key(&self) -> PartitionKey {
        PartitionKey::from(self.key)
    }
}

/// Declared with a topic this server never registers, to prove an unknown
/// topic is rejected without writing anything.
#[derive(Serialize, Deserialize, Debug, Clone)]
struct UnregisteredTopicJob;

impl JobPartition for UnregisteredTopicJob {
    fn partition_key(&self) -> PartitionKey {
        PartitionKey::from("unused")
    }
}

pub struct AppContext {
    recorded: Arc<Mutex<Vec<(u64, u32)>>>,
    slow_key: Arc<Mutex<Option<u64>>>,
    fail_until_attempt: Arc<Mutex<std::collections::HashMap<u64, usize>>>,
    gate: Arc<tokio::sync::Notify>,
}

later::background_job! {
    struct Jobs {
        #[topic("orders")]
        ordered_job: OrderedJob,
        #[topic("orders")]
        untopiced_job: UntopicedJob,
        #[topic("orders")]
        flaky_job: FlakyJob,
        #[topic("payments")]
        other_topic_job: OtherTopicJob,
        #[topic("orders")]
        gate_job: GateJob,
        #[topic("unregistered")]
        unregistered_topic_job: UnregisteredTopicJob,
    }
}

async fn handle_gate(ctx: JobsContext<AppContext>, _payload: GateJob) -> anyhow::Result<()> {
    ctx.app.gate.notified().await;
    Ok(())
}

async fn handle_other_topic(
    ctx: JobsContext<AppContext>,
    payload: OtherTopicJob,
) -> anyhow::Result<()> {
    ctx.app
        .recorded
        .lock()
        .await
        .push((payload.key, payload.seq));
    Ok(())
}

async fn handle_unregistered_topic(
    _ctx: JobsContext<AppContext>,
    _payload: UnregisteredTopicJob,
) -> anyhow::Result<()> {
    Ok(())
}

async fn handle_ordered(ctx: JobsContext<AppContext>, payload: OrderedJob) -> anyhow::Result<()> {
    if let Some(slow) = *ctx.app.slow_key.lock().await {
        if slow == payload.key {
            tokio::time::sleep(Duration::from_millis(150)).await;
        }
    }
    ctx.app
        .recorded
        .lock()
        .await
        .push((payload.key, payload.seq));
    Ok(())
}

async fn handle_untopiced(
    _ctx: JobsContext<AppContext>,
    _payload: UntopicedJob,
) -> anyhow::Result<()> {
    Ok(())
}

async fn handle_flaky(ctx: JobsContext<AppContext>, payload: FlakyJob) -> anyhow::Result<()> {
    let should_fail = {
        let mut attempts = ctx.app.fail_until_attempt.lock().await;
        let remaining = attempts.entry(payload.key).or_insert(0);
        if *remaining > 0 {
            *remaining -= 1;
            true
        } else {
            false
        }
    };
    if should_fail {
        return Err(anyhow::anyhow!("flaky failure"));
    }
    ctx.app
        .recorded
        .lock()
        .await
        .push((payload.key, payload.seq));
    Ok(())
}

async fn start_server(
    namespace: &str,
    pool: sqlx::SqlitePool,
    recorded: Arc<Mutex<Vec<(u64, u32)>>>,
    slow_key: Arc<Mutex<Option<u64>>>,
    worker_count: u8,
) -> BackgroundJobServer<AppContext, Jobs<AppContext>> {
    start_server_with_flaky(
        namespace,
        pool,
        recorded,
        slow_key,
        Arc::new(Mutex::new(std::collections::HashMap::new())),
        Arc::new(tokio::sync::Notify::new()),
        worker_count,
    )
    .await
}

async fn start_server_with_flaky(
    namespace: &str,
    pool: sqlx::SqlitePool,
    recorded: Arc<Mutex<Vec<(u64, u32)>>>,
    slow_key: Arc<Mutex<Option<u64>>>,
    fail_until_attempt: Arc<Mutex<std::collections::HashMap<u64, usize>>>,
    gate: Arc<tokio::sync::Notify>,
    worker_count: u8,
) -> BackgroundJobServer<AppContext, Jobs<AppContext>> {
    let backend = later::backend::SqliteBackend::from_pool(namespace, pool)
        .await
        .expect("create sqlite backend");
    start_server_with_backend(
        Box::new(backend),
        recorded,
        slow_key,
        fail_until_attempt,
        gate,
        worker_count,
    )
    .await
}

async fn start_server_with_backend(
    backend: Box<dyn later::backend::Backend>,
    recorded: Arc<Mutex<Vec<(u64, u32)>>>,
    slow_key: Arc<Mutex<Option<u64>>>,
    fail_until_attempt: Arc<Mutex<std::collections::HashMap<u64, usize>>>,
    gate: Arc<tokio::sync::Notify>,
    worker_count: u8,
) -> BackgroundJobServer<AppContext, Jobs<AppContext>> {
    let ctx = AppContext {
        recorded,
        slow_key,
        fail_until_attempt,
        gate,
    };
    JobsBuilder::new(
        Config::builder()
            .context(ctx)
            .backend(backend)
            .worker_count(worker_count)
            .default_retry_policy(later::retry::RetryPolicy::immediate(6))
            .topics(vec![
                TopicConfig::new("orders", TOPIC_PARTITION_COUNT).unwrap(),
                TopicConfig::new("payments", TOPIC_PARTITION_COUNT).unwrap(),
            ])
            .build(),
    )
    .with_ordered_job_handler(handle_ordered)
    .with_untopiced_job_handler(handle_untopiced)
    .with_gate_job_handler(handle_gate)
    .with_flaky_job_handler(handle_flaky)
    .with_other_topic_job_handler(handle_other_topic)
    .with_unregistered_topic_job_handler(handle_unregistered_topic)
    .build()
    .await
    .expect("start job server")
}

async fn wait_until(condition: impl Fn() -> bool, timeout: Duration) {
    let start = std::time::Instant::now();
    while start.elapsed() < timeout {
        if condition() {
            return;
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
}

#[tokio::test]
async fn one_partition_completes_in_exact_enqueue_order_with_many_workers() -> anyhow::Result<()> {
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let slow_key = Arc::new(Mutex::new(None));
    let server = start_server("order-test", pool, recorded.clone(), slow_key, 8).await;

    for seq in 0..20u32 {
        server
            .enqueue_to_partition(OrderedJob { key: 0, seq })
            .await?;
    }

    wait_until(
        || recorded.try_lock().map(|r| r.len() >= 20).unwrap_or(false),
        Duration::from_secs(10),
    )
    .await;

    let result = recorded.lock().await.clone();
    let expected: Vec<(u64, u32)> = (0..20).map(|seq| (0, seq)).collect();
    assert_eq!(result, expected);
    Ok(())
}

#[tokio::test]
async fn idle_partition_assignments_release_the_poller_for_other_partitions() -> anyhow::Result<()>
{
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let keys = distinct_keys(usize::try_from(TOPIC_PARTITION_COUNT)?);
    let server = start_server(
        "idle-assignment-release",
        pool,
        recorded.clone(),
        Arc::new(Mutex::new(None)),
        1,
    )
    .await;

    for key in &keys {
        server
            .enqueue_to_partition(OrderedJob { key: *key, seq: 0 })
            .await?;
    }

    let completed = tokio::time::timeout(Duration::from_secs(3), async {
        loop {
            let current = recorded.lock().await.clone();
            if current.len() == keys.len() {
                return current;
            }
            tokio::time::sleep(Duration::from_millis(25)).await;
        }
    })
    .await
    .map_err(|_| anyhow::anyhow!("idle partition lease held the only poller"))?;
    let mut completed_keys: Vec<u64> = completed.into_iter().map(|(key, _)| key).collect();
    let mut expected_keys = keys;
    completed_keys.sort_unstable();
    expected_keys.sort_unstable();
    assert_eq!(completed_keys, expected_keys);
    Ok(())
}

#[tokio::test]
async fn final_failure_releases_the_next_job_in_its_partition() -> anyhow::Result<()> {
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let failures = Arc::new(Mutex::new(std::collections::HashMap::from([(0, 7)])));
    let server = start_server_with_flaky(
        "final-failure-release",
        pool,
        recorded.clone(),
        Arc::new(Mutex::new(None)),
        failures,
        Arc::new(tokio::sync::Notify::new()),
        4,
    )
    .await;

    server
        .enqueue_to_partition(FlakyJob { key: 0, seq: 0 })
        .await?;
    server
        .enqueue_to_partition(OrderedJob { key: 0, seq: 1 })
        .await?;

    wait_until(
        || {
            recorded
                .try_lock()
                .map(|items| items.len() == 1)
                .unwrap_or(false)
        },
        Duration::from_secs(10),
    )
    .await;
    assert_eq!(*recorded.lock().await, vec![(0, 1)]);
    Ok(())
}

#[tokio::test]
async fn sqlite_equal_partition_ids_in_different_topics_do_not_block_each_other(
) -> anyhow::Result<()> {
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let gate = Arc::new(tokio::sync::Notify::new());
    let server = start_server_with_flaky(
        "sqlite-topic-isolation",
        pool,
        recorded.clone(),
        Arc::new(Mutex::new(None)),
        Arc::new(Mutex::new(std::collections::HashMap::new())),
        gate.clone(),
        4,
    )
    .await;

    server.enqueue_to_partition(GateJob).await?;
    server
        .enqueue_to_partition(OtherTopicJob { key: 0, seq: 1 })
        .await?;

    wait_until(
        || {
            recorded
                .try_lock()
                .map(|items| items.len() == 1)
                .unwrap_or(false)
        },
        Duration::from_secs(10),
    )
    .await;
    assert_eq!(*recorded.lock().await, vec![(0, 1)]);
    gate.notify_waiters();
    Ok(())
}

#[tokio::test]
async fn dropping_a_worker_rebalances_its_partition_without_heartbeat_expiry() -> anyhow::Result<()>
{
    let dir = tempfile::tempdir()?;
    let db_path = dir.path().join("graceful-rebalance.db");
    let pool = later::storage::Sqlite::new(&format!("sqlite://{}?mode=rwc", db_path.display()))
        .await?
        .pool()
        .clone();
    let namespace = "graceful-rebalance";
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let first = start_server(
        namespace,
        pool.clone(),
        recorded.clone(),
        Arc::new(Mutex::new(None)),
        1,
    )
    .await;

    drop(first);
    let deadline = std::time::Instant::now() + Duration::from_secs(3);
    loop {
        let remaining: i64 =
            sqlx::query_scalar("SELECT COUNT(*) FROM later_worker WHERE namespace = ?1")
                .bind(namespace)
                .fetch_one(&pool)
                .await?;
        if remaining == 0 {
            break;
        }
        if std::time::Instant::now() >= deadline {
            return Err(anyhow::anyhow!("dropped worker was not deregistered"));
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    let second = start_server(
        namespace,
        pool,
        recorded.clone(),
        Arc::new(Mutex::new(None)),
        1,
    )
    .await;
    second
        .enqueue_to_partition(OrderedJob { key: 0, seq: 1 })
        .await?;
    wait_until(
        || {
            recorded
                .try_lock()
                .map(|items| items.len() == 1)
                .unwrap_or(false)
        },
        Duration::from_secs(5),
    )
    .await;
    assert_eq!(*recorded.lock().await, vec![(0, 1)]);
    Ok(())
}

#[tokio::test]
async fn different_partitions_in_one_topic_run_concurrently() -> anyhow::Result<()> {
    let dir = tempfile::tempdir()?;
    let db_path = dir.path().join("concurrent.db");
    let url = format!("sqlite://{}?mode=rwc", db_path.display());
    let pool = later::storage::Sqlite::new(&url).await?.pool().clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let keys = distinct_keys(2);
    let (slow, fast) = (keys[0], keys[1]);
    let slow_key = Arc::new(Mutex::new(Some(slow)));
    let server = start_server("concurrent-test", pool, recorded.clone(), slow_key, 8).await;

    // `slow`'s partition takes 150ms/job; `fast`'s partition does not.
    for seq in 0..5u32 {
        server
            .enqueue_to_partition(OrderedJob { key: slow, seq })
            .await?;
    }
    for seq in 0..5u32 {
        server
            .enqueue_to_partition(OrderedJob { key: fast, seq })
            .await?;
    }

    wait_until(
        || recorded.try_lock().map(|r| r.len() >= 10).unwrap_or(false),
        Duration::from_secs(10),
    )
    .await;
    let all = recorded.lock().await.clone();
    assert_eq!(all.len(), 10, "both partitions must finish");

    let slow_results: Vec<u32> = all
        .iter()
        .filter(|(k, _)| *k == slow)
        .map(|(_, s)| *s)
        .collect();
    let fast_results: Vec<u32> = all
        .iter()
        .filter(|(k, _)| *k == fast)
        .map(|(_, s)| *s)
        .collect();
    assert_eq!(slow_results, vec![0, 1, 2, 3, 4]);
    assert_eq!(fast_results, vec![0, 1, 2, 3, 4]);

    // Proves the partitions ran concurrently rather than one queued behind
    // the other: the whole fast partition finishes interleaved with (not
    // strictly after) the slow one, i.e. most of it lands before the slow
    // partition's last job completes.
    let last_slow_index = all
        .iter()
        .position(|entry| *entry == (slow, 4))
        .expect("slow partition's last job must complete");
    let fast_before_slow_finished = all[..last_slow_index]
        .iter()
        .filter(|(k, _)| *k == fast)
        .count();
    assert!(
        fast_before_slow_finished >= 2,
        "expected the fast partition to interleave with the slow one, \
         but only {fast_before_slow_finished} of its jobs finished before \
         the slow partition's last job: {all:?}"
    );
    Ok(())
}

#[tokio::test]
async fn partition_ordering_cannot_be_requested_without_a_known_topic() -> anyhow::Result<()> {
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let server = start_server(
        "unknown-topic-test",
        pool.clone(),
        recorded,
        Arc::new(Mutex::new(None)),
        2,
    )
    .await;

    // `UnregisteredTopicJob` is declared with a topic this server's Config
    // never registers.
    let result = server.enqueue_to_partition(UnregisteredTopicJob).await;
    assert!(result.is_err());

    let job_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM later_partition_job")
        .fetch_one(&pool)
        .await?;
    assert_eq!(job_count, 0, "rejected enqueue must not store a job");
    Ok(())
}

#[tokio::test]
async fn force_fail_partition_head_unblocks_a_poison_job() -> anyhow::Result<()> {
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let slow_key = Arc::new(Mutex::new(None));
    let server = start_server("force-fail-test", pool, recorded.clone(), slow_key, 4).await;

    let key = distinct_keys(1)[0];
    // Stands in for a "poison" job: delayed far into the future so it never
    // becomes ready on its own, blocking its partition until an operator
    // intervenes.
    let poison_id = server
        .enqueue_to_partition_delayed(OrderedJob { key, seq: 0 }, Duration::from_secs(3600))
        .await?;
    server
        .enqueue_to_partition(OrderedJob { key, seq: 1 })
        .await?;

    tokio::time::sleep(Duration::from_millis(300)).await;
    assert!(
        recorded.lock().await.is_empty(),
        "the second job must stay blocked behind the poison head"
    );

    let partition = partition_for_key(&PartitionKey::from(key), TOPIC_PARTITION_COUNT);
    let failed_id = server
        .force_fail_partition_head("orders", partition, "operator override in test".to_owned())
        .await?
        .ok_or_else(|| anyhow::anyhow!("expected a head to be force-failed"))?;
    assert_eq!(failed_id, poison_id);

    wait_until(
        || recorded.try_lock().map(|r| !r.is_empty()).unwrap_or(false),
        Duration::from_secs(10),
    )
    .await;
    assert_eq!(recorded.lock().await.clone(), vec![(key, 1)]);
    Ok(())
}

#[tokio::test]
async fn two_sqlite_pools_sharing_one_file_never_overlap_one_topic_partition() -> anyhow::Result<()>
{
    let dir = tempfile::tempdir()?;
    let db_path = dir.path().join("shared.db");
    let url = format!("sqlite://{}?mode=rwc", db_path.display());

    let pool_a = later::storage::Sqlite::new(&url).await?.pool().clone();
    let pool_b = later::storage::Sqlite::new(&url).await?.pool().clone();

    let recorded = Arc::new(Mutex::new(Vec::new()));
    let namespace = "shared-file-test";
    let server_a = start_server(
        namespace,
        pool_a,
        recorded.clone(),
        Arc::new(Mutex::new(None)),
        4,
    )
    .await;
    let server_b = start_server(
        namespace,
        pool_b,
        recorded.clone(),
        Arc::new(Mutex::new(None)),
        4,
    )
    .await;

    for seq in 0..30u32 {
        server_a
            .enqueue_to_partition(OrderedJob { key: 2, seq })
            .await?;
    }

    wait_until(
        || recorded.try_lock().map(|r| r.len() >= 30).unwrap_or(false),
        Duration::from_secs(10),
    )
    .await;

    let result = recorded.lock().await.clone();
    let expected: Vec<(u64, u32)> = (0..30).map(|seq| (2, seq)).collect();
    assert_eq!(result, expected);

    drop(server_b);
    Ok(())
}

#[tokio::test]
async fn retry_blocks_the_partition_then_releases_it_in_order() -> anyhow::Result<()> {
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let fail_until_attempt = Arc::new(Mutex::new(std::collections::HashMap::from([(0, 2)])));
    let server = start_server_with_flaky(
        "retry-test",
        pool,
        recorded.clone(),
        Arc::new(Mutex::new(None)),
        fail_until_attempt,
        Arc::new(tokio::sync::Notify::new()),
        4,
    )
    .await;

    server
        .enqueue_to_partition(FlakyJob { key: 0, seq: 0 })
        .await?;
    server
        .enqueue_to_partition(FlakyJob { key: 0, seq: 1 })
        .await?;

    wait_until(
        || recorded.try_lock().map(|r| r.len() >= 2).unwrap_or(false),
        Duration::from_secs(10),
    )
    .await;

    let result = recorded.lock().await.clone();
    assert_eq!(result, vec![(0, 0), (0, 1)]);
    Ok(())
}

#[tokio::test]
async fn continuation_blocks_its_partition_until_the_parent_succeeds() -> anyhow::Result<()> {
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let gate = Arc::new(tokio::sync::Notify::new());
    let server = start_server_with_flaky(
        "continuation-test",
        pool,
        recorded.clone(),
        Arc::new(Mutex::new(None)),
        Arc::new(Mutex::new(std::collections::HashMap::new())),
        gate.clone(),
        4,
    )
    .await;

    // The parent (`GateJob`) is not itself topic-partitioned, proving a
    // continuation can be blocked by a job outside its own partition.
    let parent_id = server.enqueue(GateJob).await?;
    server
        .enqueue_continue_to_partition(parent_id, OrderedJob { key: 3, seq: 100 })
        .await?;
    server
        .enqueue_to_partition(OrderedJob { key: 3, seq: 200 })
        .await?;

    // Neither partitioned job may run while the parent is still running.
    tokio::time::sleep(Duration::from_millis(300)).await;
    assert!(
        recorded.lock().await.is_empty(),
        "continuation and the job behind it must not run before the parent succeeds"
    );

    gate.notify_one();

    wait_until(
        || recorded.try_lock().map(|r| r.len() >= 2).unwrap_or(false),
        Duration::from_secs(10),
    )
    .await;

    let result = recorded.lock().await.clone();
    assert_eq!(result, vec![(3, 100), (3, 200)]);
    Ok(())
}

#[tokio::test]
async fn continuation_to_partition_advances_when_the_parent_already_succeeded() -> anyhow::Result<()>
{
    // Reproduces a crash-recovery gap: if a worker dies after committing the
    // parent's success but before `enqueue_continuations` runs, nothing would
    // otherwise ever wake a topic-partitioned continuation up, permanently
    // stalling its whole partition. Calling `enqueue_continue_to_partition`
    // against an already-succeeded parent hits the exact same code path, so
    // it doubles as a deterministic, non-crash reproduction of that gap.
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let slow_key = Arc::new(Mutex::new(None));
    let server = start_server(
        "late-continuation-test",
        pool,
        recorded.clone(),
        slow_key,
        4,
    )
    .await;

    let parent_key = distinct_keys(1)[0];
    let parent_id = server
        .enqueue(OrderedJob {
            key: parent_key,
            seq: 0,
        })
        .await?;
    wait_until(
        || {
            recorded
                .try_lock()
                .map(|r| r.contains(&(parent_key, 0)))
                .unwrap_or(false)
        },
        Duration::from_secs(5),
    )
    .await;
    // Give the success transition a moment to be saved after the handler
    // pushed to `recorded`, so the parent is genuinely terminal in storage.
    tokio::time::sleep(Duration::from_millis(200)).await;
    recorded.lock().await.clear();

    let child_key = distinct_keys(1)[0];
    server
        .enqueue_continue_to_partition(
            parent_id,
            OrderedJob {
                key: child_key,
                seq: 1,
            },
        )
        .await?;

    wait_until(
        || recorded.try_lock().map(|r| !r.is_empty()).unwrap_or(false),
        Duration::from_secs(10),
    )
    .await;

    assert_eq!(recorded.lock().await.clone(), vec![(child_key, 1)]);
    Ok(())
}

#[test]
fn same_key_always_resolves_to_the_same_partition() {
    let key = PartitionKey::from(42u64);
    let first = partition_for_key(&key, TOPIC_PARTITION_COUNT);
    let second = partition_for_key(&key, TOPIC_PARTITION_COUNT);
    assert_eq!(first, second);
    assert!(first.0 < TOPIC_PARTITION_COUNT);
}

#[test]
fn distinct_keys_helper_returns_keys_in_different_partitions() {
    let keys = distinct_keys(2);
    let partitions: Vec<PartitionId> = keys
        .iter()
        .map(|key| partition_for_key(&PartitionKey::from(*key), TOPIC_PARTITION_COUNT))
        .collect();
    assert_ne!(partitions[0], partitions[1]);
}

#[tokio::test]
async fn adding_a_worker_picks_up_idle_partitions_without_waiting_for_lease_expiry(
) -> anyhow::Result<()> {
    // A dedicated topic/job type so this test's two servers don't share
    // rendezvous assignments with the shared `Jobs` server used above.
    const REBALANCE_PARTITIONS: u32 = 8;
    type RebalanceLog = Arc<Mutex<Vec<(&'static str, u64, u32)>>>;

    #[derive(Serialize, Deserialize, Debug, Clone)]
    struct RebalanceJob {
        key: u64,
        seq: u32,
    }
    impl JobPartition for RebalanceJob {
        fn partition_key(&self) -> PartitionKey {
            PartitionKey::from(self.key)
        }
    }

    pub struct RebalanceContext {
        recorded: RebalanceLog,
        server_label: &'static str,
    }

    later::background_job! {
        struct RebalanceJobs {
            #[topic("rebalance")]
            rebalance_job: RebalanceJob,
        }
    }

    async fn handle_rebalance(
        ctx: RebalanceJobsContext<RebalanceContext>,
        payload: RebalanceJob,
    ) -> anyhow::Result<()> {
        ctx.app
            .recorded
            .lock()
            .await
            .push((ctx.app.server_label, payload.key, payload.seq));
        Ok(())
    }

    async fn start_rebalance_server(
        pool: sqlx::SqlitePool,
        recorded: RebalanceLog,
        server_label: &'static str,
    ) -> BackgroundJobServer<RebalanceContext, RebalanceJobs<RebalanceContext>> {
        let backend = later::backend::SqliteBackend::from_pool("rebalance-test", pool)
            .await
            .expect("create sqlite backend");
        RebalanceJobsBuilder::new(
            Config::builder()
                .context(RebalanceContext {
                    recorded,
                    server_label,
                })
                .backend(Box::new(backend))
                .worker_count(4)
                .topics(vec![
                    TopicConfig::new("rebalance", REBALANCE_PARTITIONS).unwrap()
                ])
                .build(),
        )
        .with_rebalance_job_handler(handle_rebalance)
        .build()
        .await
        .expect("start rebalance job server")
    }

    let dir = tempfile::tempdir()?;
    let db_path = dir.path().join("rebalance.db");
    let url = format!("sqlite://{}?mode=rwc", db_path.display());
    let pool_a = later::storage::Sqlite::new(&url).await?.pool().clone();

    let recorded = Arc::new(Mutex::new(Vec::new()));
    let keys = distinct_keys_for(REBALANCE_PARTITIONS as usize, REBALANCE_PARTITIONS);
    let server_a = start_rebalance_server(pool_a, recorded.clone(), "server-a").await;

    // Alone, server-a claims every partition; process one round.
    for &key in &keys {
        server_a
            .enqueue_to_partition(RebalanceJob { key, seq: 0 })
            .await?;
    }
    wait_until(
        || {
            recorded
                .try_lock()
                .map(|r| r.len() >= keys.len())
                .unwrap_or(false)
        },
        Duration::from_secs(10),
    )
    .await;
    assert_eq!(recorded.lock().await.len(), keys.len());

    // server-b joins the same database and namespace. Give membership a
    // couple of refresh cycles (2s each) to converge on both sides before
    // enqueuing more work, so this test observes steady-state rebalancing
    // rather than racing server-a's still-stale view of who else is live.
    let pool_b = later::storage::Sqlite::new(&url).await?.pool().clone();
    let server_b = start_rebalance_server(pool_b, recorded.clone(), "server-b").await;
    tokio::time::sleep(Duration::from_secs(5)).await;

    // Enqueue a second round into the same partitions. Some must move to
    // server-b well before the ~30s partition lease would ever expire.
    for &key in &keys {
        server_a
            .enqueue_to_partition(RebalanceJob { key, seq: 1 })
            .await?;
    }
    wait_until(
        || {
            recorded
                .try_lock()
                .map(|r| r.len() >= keys.len() * 2)
                .unwrap_or(false)
        },
        Duration::from_secs(20),
    )
    .await;

    let all = recorded.lock().await.clone();
    assert_eq!(
        all.len(),
        keys.len() * 2,
        "both rounds must complete: {all:?}"
    );
    let second_round_servers: std::collections::HashSet<&'static str> = all
        .iter()
        .filter(|(_, _, seq)| *seq == 1)
        .map(|(server, _, _)| *server)
        .collect();
    assert!(
        second_round_servers.contains("server-b"),
        "expected server-b to take over at least one partition within a few \
         poll cycles, well short of the ~30s partition lease duration; \
         second round was handled by: {second_round_servers:?}"
    );

    drop(server_b);
    Ok(())
}

/// Like `distinct_keys`, but for an arbitrary partition count.
fn distinct_keys_for(n: usize, partition_count: u32) -> Vec<u64> {
    let mut seen_partitions = std::collections::HashSet::new();
    let mut chosen = Vec::new();
    let mut candidate = 0u64;
    while chosen.len() < n && seen_partitions.len() < partition_count as usize {
        let partition = partition_for_key(&PartitionKey::from(candidate), partition_count);
        if seen_partitions.insert(partition) {
            chosen.push(candidate);
        }
        candidate += 1;
    }
    chosen
}

#[tokio::test]
async fn dropping_a_server_deregisters_its_partition_worker() -> anyhow::Result<()> {
    let pool = later::storage::Sqlite::new("sqlite::memory:")
        .await?
        .pool()
        .clone();
    let recorded = Arc::new(Mutex::new(Vec::new()));
    let server = start_server(
        "deregister-test",
        pool.clone(),
        recorded,
        Arc::new(Mutex::new(None)),
        2,
    )
    .await;

    // The server heartbeats itself synchronously on startup.
    let live_before: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM later_worker")
        .fetch_one(&pool)
        .await?;
    assert_eq!(live_before, 1, "server did not register a heartbeat");

    drop(server);
    // Deregistration is a best-effort task spawned from Drop; give it a
    // moment to run.
    let deadline = std::time::Instant::now() + Duration::from_secs(5);
    let mut live_after;
    loop {
        live_after = sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM later_worker")
            .fetch_one(&pool)
            .await?;
        if live_after == 0 || std::time::Instant::now() >= deadline {
            break;
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
    assert_eq!(
        live_after, 0,
        "dropping the server should deregister its partition worker promptly"
    );
    Ok(())
}