later 0.0.27

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

use async_std::sync::{Arc, Mutex};
use later::{backend::SqliteBackend, storage::Sqlite, BackgroundJobServer, Config};
use serde::{Deserialize, Serialize};
use std::{
    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
    time::{Duration, Instant},
};

#[derive(Debug, Clone)]
struct Event {
    label: &'static str,
    start: Instant,
    end: Instant,
}

#[derive(Default)]
struct Concurrency {
    current: AtomicUsize,
    peak: AtomicUsize,
}

impl Concurrency {
    fn enter(&self) {
        let now = self.current.fetch_add(1, Ordering::SeqCst) + 1;
        self.peak.fetch_max(now, Ordering::SeqCst);
    }

    fn exit(&self) {
        self.current.fetch_sub(1, Ordering::SeqCst);
    }
}

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

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

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

struct AppContext {
    events: Arc<Mutex<Vec<Event>>>,
    tick_sleep: Duration,
    chain_continuation_once: Arc<AtomicBool>,
    concurrency: Arc<Concurrency>,
}

later::background_job! {
    struct Jobs {
        sequential_tick: SequentialTick,
        sequential_continuation: SequentialContinuation,
        normal_tick: NormalTick,
    }
}

async fn handle_sequential_tick(
    ctx: JobsContext<AppContext>,
    _payload: SequentialTick,
) -> anyhow::Result<()> {
    if ctx
        .app
        .chain_continuation_once
        .swap(false, Ordering::SeqCst)
    {
        ctx.enqueue_recurring_continue(ctx.job_id().clone(), SequentialContinuation)
            .await?;
    }

    let start = Instant::now();
    tokio::time::sleep(ctx.app.tick_sleep).await;
    let end = Instant::now();
    ctx.app.events.lock().await.push(Event {
        label: "tick",
        start,
        end,
    });
    Ok(())
}

async fn handle_sequential_continuation(
    ctx: JobsContext<AppContext>,
    _payload: SequentialContinuation,
) -> anyhow::Result<()> {
    let now = Instant::now();
    ctx.app.events.lock().await.push(Event {
        label: "continuation",
        start: now,
        end: now,
    });
    Ok(())
}

async fn handle_normal_tick(
    ctx: JobsContext<AppContext>,
    _payload: NormalTick,
) -> anyhow::Result<()> {
    ctx.app.concurrency.enter();
    tokio::time::sleep(ctx.app.tick_sleep).await;
    ctx.app.concurrency.exit();
    Ok(())
}

async fn start_server(
    namespace: &str,
    pool: sqlx::SqlitePool,
    events: Arc<Mutex<Vec<Event>>>,
    concurrency: Arc<Concurrency>,
    tick_sleep: Duration,
) -> BackgroundJobServer<AppContext, Jobs<AppContext>> {
    let backend = SqliteBackend::from_pool(namespace, pool)
        .await
        .expect("create sqlite backend");
    let ctx = AppContext {
        events,
        tick_sleep,
        chain_continuation_once: Arc::new(AtomicBool::new(true)),
        concurrency,
    };
    JobsBuilder::new(
        Config::builder()
            .context(ctx)
            .backend(Box::new(backend))
            .worker_count(4)
            .default_retry_policy(later::retry::RetryPolicy::immediate(6))
            .recurring_sequential_partitions(4)
            .build(),
    )
    .with_sequential_tick_handler(handle_sequential_tick)
    .with_sequential_continuation_handler(handle_sequential_continuation)
    .with_normal_tick_handler(handle_normal_tick)
    .build()
    .await
    .expect("start job server")
}

async fn wait_until_events(events: &Arc<Mutex<Vec<Event>>>, count: usize, timeout: Duration) {
    let start = Instant::now();
    while start.elapsed() < timeout {
        if events.lock().await.len() >= count {
            return;
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
}

async fn wait_until_peak(concurrency: &Arc<Concurrency>, peak: usize, timeout: Duration) {
    let start = Instant::now();
    while start.elapsed() < timeout {
        if concurrency.peak.load(Ordering::SeqCst) >= peak {
            return;
        }
        tokio::time::sleep(Duration::from_millis(50)).await;
    }
}

/// Every second's occurrence sleeps well past the next tick. The recurring
/// scheduler must never let two occurrences run concurrently, and must not
/// drop the ticks that pile up behind a slow one.
#[tokio::test]
async fn sequential_recurring_jobs_never_overlap_and_queue_a_backlog() -> anyhow::Result<()> {
    let pool = Sqlite::new("sqlite::memory:").await?.pool().clone();
    let events = Arc::new(Mutex::new(Vec::new()));
    let concurrency = Arc::new(Concurrency::default());
    let server = start_server(
        "recurring-sequential-overlap",
        pool,
        events.clone(),
        concurrency,
        Duration::from_millis(2500),
    )
    .await;

    server
        .enqueue_recurring_sequential(
            "overlap-test".to_string(),
            SequentialTick,
            "* * * * * *".to_string(),
        )
        .await?;

    wait_until_events(&events, 3, Duration::from_secs(30)).await;

    let recorded = events.lock().await.clone();
    assert!(
        recorded.len() >= 3,
        "expected at least 3 occurrences to run despite a 2.5s handler on a 1s schedule, got {}",
        recorded.len()
    );
    for pair in recorded.windows(2) {
        let [a, b] = pair else { unreachable!() };
        assert!(
            a.end <= b.start,
            "occurrences overlapped: {:?} finished after {:?} started",
            a,
            b
        );
    }

    let _ = server.shutdown(Duration::from_secs(5)).await;
    Ok(())
}

/// A continuation chained off a sequential occurrence's own job id (from
/// inside its handler) must run as part of that occurrence's chain: after
/// its parent, and before the next scheduled occurrence.
#[tokio::test]
async fn sequential_recurring_continuation_runs_as_part_of_the_chain() -> anyhow::Result<()> {
    let pool = Sqlite::new("sqlite::memory:").await?.pool().clone();
    let events = Arc::new(Mutex::new(Vec::new()));
    let concurrency = Arc::new(Concurrency::default());
    let server = start_server(
        "recurring-sequential-continuation",
        pool,
        events.clone(),
        concurrency,
        Duration::from_millis(50),
    )
    .await;

    server
        .enqueue_recurring_sequential(
            "continuation-test".to_string(),
            SequentialTick,
            "*/3 * * * * *".to_string(),
        )
        .await?;

    wait_until_events(&events, 3, Duration::from_secs(20)).await;

    let recorded = events.lock().await.clone();
    assert!(
        recorded.len() >= 3,
        "expected at least [tick, continuation, tick], got {:?}",
        recorded
    );
    assert_eq!(recorded[0].label, "tick");
    assert_eq!(recorded[1].label, "continuation");
    assert_eq!(
        recorded[2].label, "tick",
        "the continuation must run before the next scheduled occurrence"
    );

    let _ = server.shutdown(Duration::from_secs(2)).await;
    Ok(())
}

/// Normal-mode recurring jobs run on the plain unordered path, unlike
/// sequential ones: nothing about being a recurring occurrence should make a
/// job serialize against other concurrent work. A far-future cron keeps the
/// proactive poller from creating a second occurrence during the test, so
/// the overlap is forced deterministically with a sibling job due at
/// roughly the same time as the recurring job's first occurrence, rather
/// than by racing the (comparatively slow and jittery) recurring-job poller
/// cadence.
#[tokio::test]
async fn normal_recurring_jobs_may_overlap() -> anyhow::Result<()> {
    let pool = Sqlite::new("sqlite::memory:").await?.pool().clone();
    let events = Arc::new(Mutex::new(Vec::new()));
    let concurrency = Arc::new(Concurrency::default());
    let server = start_server(
        "recurring-normal-overlap",
        pool,
        events,
        concurrency.clone(),
        Duration::from_millis(3000),
    )
    .await;

    server
        .enqueue_recurring(
            "normal-overlap-test".to_string(),
            NormalTick,
            "* * * * * *".to_string(),
        )
        .await?;
    server
        .enqueue_delayed(NormalTick, Duration::from_millis(500))
        .await?;

    wait_until_peak(&concurrency, 2, Duration::from_secs(20)).await;
    assert!(
        concurrency.peak.load(Ordering::SeqCst) >= 2,
        "expected the recurring occurrence to overlap with a concurrently-due sibling job"
    );

    let _ = server.shutdown(Duration::from_secs(5)).await;
    Ok(())
}

/// Re-registering the same identifier must update the definition in place,
/// not add a second entry to the index the recurring-job poller scans, and
/// - the actual production bug this regression-tests - must not enqueue a
/// second occurrence either. Every process restart calls
/// `enqueue_recurring_sequential` again to recreate its schedule; if that
/// unconditionally enqueued a fresh occurrence, a long-running deployment
/// that restarts often would permanently queue one more job per restart,
/// forever, each blocking every other job sharing its reserved-topic
/// partition until its turn came - exactly what happened in production
/// (a daily job's restart-accumulated occurrences blocked a once-a-minute
/// job behind them for hours).
#[tokio::test]
async fn reregistering_a_sequential_recurring_job_does_not_duplicate_its_schedule_entry(
) -> anyhow::Result<()> {
    let pool = Sqlite::new("sqlite::memory:").await?.pool().clone();
    let events = Arc::new(Mutex::new(Vec::new()));
    let concurrency = Arc::new(Concurrency::default());
    let server = start_server(
        "recurring-upsert",
        pool.clone(),
        events,
        concurrency,
        Duration::from_millis(50),
    )
    .await;

    let first = server
        .enqueue_recurring_sequential(
            "upsert-test".to_string(),
            SequentialTick,
            "*/30 * * * * *".to_string(),
        )
        .await?;
    assert!(
        first.is_some(),
        "a brand-new identifier must enqueue its first occurrence"
    );

    let second = server
        .enqueue_recurring_sequential(
            "upsert-test".to_string(),
            SequentialTick,
            "*/30 * * * * *".to_string(),
        )
        .await?;
    assert!(
        second.is_none(),
        "re-registering the same identifier on the same schedule must not enqueue another occurrence"
    );

    let count: i64 = sqlx::query_scalar(
        "SELECT COUNT(*) FROM later_storage_range WHERE range_key = 'later-recurring-upsert-all-recurring-jobs'",
    )
    .fetch_one(&pool)
    .await?;
    assert_eq!(
        count, 1,
        "upserting the same identifier must not add a duplicate schedule entry"
    );

    // At most 1, not exactly 1: the single legitimate occurrence may have
    // already been claimed and run (this cron is due within 30s and the
    // handler is fast) by the time this check runs - that race is fine.
    // What must never happen is a *second* occurrence queued behind it,
    // which is what `queued == 2` would mean.
    let queued: i64 = sqlx::query_scalar(
        "SELECT COUNT(*) FROM later_partition_job WHERE namespace = 'later-recurring-upsert' AND topic = '__later_recurring__'",
    )
    .fetch_one(&pool)
    .await?;
    assert!(
        queued <= 1,
        "re-registering the same identifier must not leave a second occurrence queued behind the first, got {queued}"
    );

    let _ = server.shutdown(Duration::from_secs(2)).await;
    Ok(())
}

/// Re-registering with a *different* cron schedule is treated as a real
/// change, not a no-op upsert - it must still enqueue a fresh occurrence
/// under the new schedule, unlike re-registering with the same schedule.
#[tokio::test]
async fn reregistering_with_a_changed_cron_enqueues_a_fresh_occurrence() -> anyhow::Result<()> {
    let pool = Sqlite::new("sqlite::memory:").await?.pool().clone();
    let events = Arc::new(Mutex::new(Vec::new()));
    let concurrency = Arc::new(Concurrency::default());
    let server = start_server(
        "recurring-cron-change",
        pool,
        events,
        concurrency,
        Duration::from_millis(50),
    )
    .await;

    let first = server
        .enqueue_recurring_sequential(
            "cron-change-test".to_string(),
            SequentialTick,
            "*/30 * * * * *".to_string(),
        )
        .await?;
    assert!(first.is_some());

    let second = server
        .enqueue_recurring_sequential(
            "cron-change-test".to_string(),
            SequentialTick,
            "0 * * * * *".to_string(),
        )
        .await?;
    assert!(
        second.is_some(),
        "a changed cron schedule must enqueue a fresh occurrence, not be treated as a no-op"
    );

    let _ = server.shutdown(Duration::from_secs(2)).await;
    Ok(())
}

/// Two servers sharing one database and namespace - the normal shape of a
/// horizontally-scaled deployment - registering the exact same brand-new
/// recurring identifier at the same moment must not both decide "I'm
/// first": exactly one enqueues the first occurrence, and the other
/// defers to it.
#[tokio::test]
async fn simultaneous_first_registration_from_two_servers_enqueues_only_once() -> anyhow::Result<()>
{
    let pool = Sqlite::new("sqlite::memory:").await?.pool().clone();
    let namespace = "recurring-simultaneous-registration";
    let first_server = start_server(
        namespace,
        pool.clone(),
        Arc::new(Mutex::new(Vec::new())),
        Arc::new(Concurrency::default()),
        Duration::from_millis(50),
    )
    .await;
    let second_server = start_server(
        namespace,
        pool,
        Arc::new(Mutex::new(Vec::new())),
        Arc::new(Concurrency::default()),
        Duration::from_millis(50),
    )
    .await;

    let (first, second) = tokio::join!(
        first_server.enqueue_recurring_sequential(
            "cold-start-race".to_string(),
            SequentialTick,
            "*/30 * * * * *".to_string(),
        ),
        second_server.enqueue_recurring_sequential(
            "cold-start-race".to_string(),
            SequentialTick,
            "*/30 * * * * *".to_string(),
        ),
    );
    let outcomes = [first?.is_some(), second?.is_some()];
    assert_eq!(
        outcomes.iter().filter(|enqueued| **enqueued).count(),
        1,
        "exactly one of the two simultaneous first registrations must enqueue the first \
         occurrence, got {outcomes:?}"
    );

    let _ = first_server.shutdown(Duration::from_secs(2)).await;
    let _ = second_server.shutdown(Duration::from_secs(2)).await;
    Ok(())
}