runledger-runtime 0.12.0

Async worker, scheduler, and reaper runtime for the Runledger job system
Documentation
use std::future::pending;
use std::sync::Arc;
use std::time::Duration;

use sqlx::postgres::PgPoolOptions;
use tokio::sync::{Notify, watch};
use tokio::time::timeout;

use super::super::WorkerLoop;
use crate::RuntimeLoopExit;
use crate::config::JobsConfig;
use crate::observer::JobLifecycleObservers;
use crate::registry::JobRegistry;

const UNUSED_LAZY_POOL_URL: &str =
    "postgres://postgres:postgres@127.0.0.1:65535/runledger_worker_capacity_test";

fn worker_with_capacity(max_global_concurrency: usize) -> WorkerLoop {
    let pool = PgPoolOptions::new()
        .connect_lazy(UNUSED_LAZY_POOL_URL)
        .expect("construct lazy capacity-test pool");
    let (_, shutdown) = watch::channel(false);
    WorkerLoop::new(
        pool,
        JobRegistry::new(),
        JobsConfig {
            worker_id: "worker-capacity-test".to_owned(),
            poll_interval: Duration::from_millis(10),
            claim_batch_size: 4,
            lease_ttl_seconds: 30,
            max_global_concurrency,
            reaper_interval: Duration::from_secs(30),
            schedule_poll_interval: Duration::from_secs(30),
            reaper_retry_delay_ms: 1_000,
        },
        shutdown,
        JobLifecycleObservers::empty(),
    )
}

#[tokio::test]
async fn join_set_occupancy_is_worker_capacity() {
    let mut worker = worker_with_capacity(2);
    assert_eq!(worker.available_capacity(), 2);

    worker.join_set.spawn(pending());
    assert_eq!(worker.available_capacity(), 1);
    worker.join_set.spawn(pending());
    assert_eq!(worker.available_capacity(), 0);

    worker.join_set.abort_all();
    while worker.join_set.join_next().await.is_some() {}
    assert_eq!(worker.available_capacity(), 2);
}

#[tokio::test]
async fn crashed_job_task_is_drained_and_restores_capacity() {
    let mut worker = worker_with_capacity(1);
    worker
        .join_set
        .spawn(async { panic!("capacity test panic") });
    assert_eq!(worker.available_capacity(), 0);

    for _ in 0..100 {
        worker.drain_finished_tasks().await;
        if worker.join_set.is_empty() {
            break;
        }
        tokio::task::yield_now().await;
    }

    assert!(worker.join_set.is_empty());
    assert_eq!(worker.available_capacity(), 1);
}

#[tokio::test]
async fn shutdown_drain_waits_for_join_set_tasks() {
    let mut worker = worker_with_capacity(1);
    let started = Arc::new(Notify::new());
    let release = Arc::new(Notify::new());
    let task_started = Arc::clone(&started);
    let task_release = Arc::clone(&release);
    worker.join_set.spawn(async move {
        task_started.notify_one();
        task_release.notified().await;
    });
    started.notified().await;

    let mut drain = tokio::spawn(worker.drain(RuntimeLoopExit::Shutdown));
    assert!(
        timeout(Duration::from_millis(25), &mut drain)
            .await
            .is_err(),
        "worker drain returned before its JoinSet task completed"
    );

    release.notify_waiters();
    assert_eq!(
        timeout(Duration::from_secs(1), drain)
            .await
            .expect("worker drain should finish after release")
            .expect("worker drain task should not panic"),
        RuntimeLoopExit::Shutdown
    );
}