core/test/rust/runtime/job/lease.rs

use super::*;
use std::{sync::Arc, time::Duration};

#[test]
fn expired_ack_is_rejected_after_waiting_for_the_write_lock() {
    let directory = Directory::new();
    let runtime = Runtime::new(Limits::default()).unwrap();
    runtime
        .run(async |scope| {
            let database = Database::open(&directory.0, options(), scope)
                .await
                .unwrap();
            let mut owner = database.executor().await.unwrap();
            owner
                .schema(scope, job::schema::TABLE, job::schema::SQL)
                .await
                .unwrap();
            job::enqueue(&mut owner, scope, submission("locked"), 1000)
                .await
                .unwrap();
            let claim = Arc::new(
                job::claim(&mut owner, scope, 1000, 100)
                    .await
                    .unwrap()
                    .unwrap(),
            );
            let mut waiting = database.executor().await.unwrap();
            let transaction = owner.begin(scope).await.unwrap();
            let (ready, started) = tokio::sync::oneshot::channel();
            let claim_copy = claim.clone();
            let task = scope
                .spawn(move |child| async move {
                    ready.send(()).unwrap();
                    job::complete(&mut waiting, &child, &claim_copy, 1001).await
                })
                .unwrap();
            started.await.unwrap();
            // 租约必须包含 SQL 锁外的等待时间，不能复用等待前的时间快照。
            tokio::time::sleep(Duration::from_millis(200)).await;
            owner.finish(scope, transaction, true).await.unwrap();
            assert!(
                task.wait().await.unwrap().is_err(),
                "expired lease acknowledged after lock wait"
            );
            let next = job::claim(&mut owner, scope, 1300, 100)
                .await
                .unwrap()
                .unwrap();
            assert_eq!(next.id, claim.id);
            assert_ne!(next.token, claim.token);
            assert_eq!(next.attempt, 2);
        })
        .unwrap();
}

#[test]
fn claim_lease_starts_after_its_write_lock_wait() {
    let directory = Directory::new();
    let runtime = Runtime::new(Limits::default()).unwrap();
    runtime
        .run(async |scope| {
            let database = Database::open(&directory.0, options(), scope)
                .await
                .unwrap();
            let mut owner = database.executor().await.unwrap();
            owner
                .schema(scope, job::schema::TABLE, job::schema::SQL)
                .await
                .unwrap();
            job::enqueue(&mut owner, scope, submission("claim"), 1000)
                .await
                .unwrap();
            let mut waiting = database.executor().await.unwrap();
            let transaction = owner.begin(scope).await.unwrap();
            let (ready, started) = tokio::sync::oneshot::channel();
            let task = scope
                .spawn(move |child| async move {
                    ready.send(()).unwrap();
                    job::claim(&mut waiting, &child, 1000, 100)
                        .await
                        .unwrap()
                        .unwrap()
                })
                .unwrap();
            started.await.unwrap();
            tokio::time::sleep(Duration::from_millis(200)).await;
            owner.finish(scope, transaction, true).await.unwrap();
            let claim = task.wait().await.unwrap();
            assert!(
                claim.lease_until >= 1300,
                "lease lost its duration while waiting for the write lock: {}",
                claim.lease_until
            );
        })
        .unwrap();
}

#[test]
fn cancelling_a_waiting_ack_leaves_the_claim_recoverable() {
    let directory = Directory::new();
    let runtime = Runtime::new(Limits::default()).unwrap();
    runtime
        .run(async |scope| {
            let database = Database::open(&directory.0, options(), scope)
                .await
                .unwrap();
            let mut owner = database.executor().await.unwrap();
            owner
                .schema(scope, job::schema::TABLE, job::schema::SQL)
                .await
                .unwrap();
            job::enqueue(&mut owner, scope, submission("cancelled"), 1000)
                .await
                .unwrap();
            let claim = Arc::new(
                job::claim(&mut owner, scope, 1000, 1000)
                    .await
                    .unwrap()
                    .unwrap(),
            );
            let mut waiting = database.executor().await.unwrap();
            let transaction = owner.begin(scope).await.unwrap();
            let (ready, started) = tokio::sync::oneshot::channel();
            let claim_copy = claim.clone();
            let task = scope
                .spawn(move |child| async move {
                    ready.send(()).unwrap();
                    job::complete(&mut waiting, &child, &claim_copy, 1001).await
                })
                .unwrap();
            started.await.unwrap();
            tokio::time::sleep(Duration::from_millis(50)).await;
            task.cancel();
            tokio::time::sleep(Duration::from_millis(10)).await;
            owner.finish(scope, transaction, true).await.unwrap();
            assert!(matches!(
                task.wait().await,
                Err(dever_runtime::task::Error::Cancelled)
            ));
            let next = job::claim(&mut owner, scope, 3000, 100)
                .await
                .unwrap()
                .unwrap();
            assert_eq!(next.id, claim.id);
            assert_ne!(next.token, claim.token);
            assert_eq!(next.attempt, 2);
        })
        .unwrap();
}

#[test]
fn expired_ack_rolls_back_the_callers_business_transaction() {
    let directory = Directory::new();
    let runtime = Runtime::new(Limits::default()).unwrap();
    runtime
        .run(async |scope| {
            let database = Database::open(&directory.0, options(), scope)
                .await
                .unwrap();
            let mut executor = database.executor().await.unwrap();
            executor
                .schema(scope, job::schema::TABLE, job::schema::SQL)
                .await
                .unwrap();
            executor
                .schema(
                    scope,
                    "business",
                    "CREATE TABLE business(id INTEGER PRIMARY KEY);",
                )
                .await
                .unwrap();
            job::enqueue(&mut executor, scope, submission("atomic"), 1000)
                .await
                .unwrap();
            let claim = job::claim(&mut executor, scope, 1000, 100)
                .await
                .unwrap()
                .unwrap();
            let owner = executor.begin(scope).await.unwrap();
            executor
                .execute(scope, "INSERT INTO business(id) VALUES (1)", vec![])
                .await
                .unwrap();
            assert!(
                job::complete(&mut executor, scope, &claim, claim.lease_until)
                    .await
                    .is_err()
            );
            assert!(executor.finish(scope, owner, true).await.is_err());
            let rows = executor
                .query(scope, "SELECT count(*) FROM business", vec![], 1, |row| {
                    row.get::<_, i64>(0)
                })
                .await
                .unwrap();
            assert_eq!(rows, vec![0]);
            let next = job::claim(&mut executor, scope, claim.lease_until, 100)
                .await
                .unwrap()
                .unwrap();
            assert_eq!(next.id, claim.id);
            assert_eq!(next.attempt, 2);
        })
        .unwrap();
}