crank/crates/crank-registry/tests/integration/artifact_reconciliation.rs

use std::time::Duration as StdDuration;

use crank_artifacts::ArtifactStore;
use crank_registry::{
    ArtifactClaimFinalization, ArtifactClaimFinalizeOutcome, ArtifactClaimOutcome,
    ArtifactClaimToken, ClaimArtifactReconciliationRequest,
    ClaimExpiredArtifactReconciliationRequest, RegistryError,
};

use super::{
    artifact_sources::{TestRoot, timestamp},
    common::TestDatabase,
};

#[tokio::test]
async fn locked_database_time_and_expired_recovery_are_fenced() {
    let database = TestDatabase::new().await;
    let registry = database.registry().await;
    let raw_pool = database.raw_pool().await;
    let root = TestRoot::new("database-time");
    let store = ArtifactStore::open(&root.0).unwrap();
    let caller_time = timestamp("2020-01-01T00:00:00Z");

    let blocked = store.put_registered(b"row lock expiry\n").unwrap();
    assert!(matches!(
        registry
            .claim_artifact_reconciliation(ClaimArtifactReconciliationRequest {
                artifact: &blocked,
                token: ArtifactClaimToken::generate(),
                claimed_at: caller_time,
                lease_expires_at: caller_time + time::Duration::seconds(1),
                detached_grace: time::Duration::hours(24),
            })
            .await
            .unwrap(),
        ArtifactClaimOutcome::Claimed(_)
    ));
    let mut blocker = raw_pool.begin().await.unwrap();
    sqlx::query("select digest from artifact_blobs where digest = $1 for update")
        .bind(blocked.artifact_ref().digest_hex())
        .fetch_one(&mut *blocker)
        .await
        .unwrap();
    let waiting_registry = registry.clone();
    let waiting_artifact = blocked.clone();
    let waiting = tokio::spawn(async move {
        waiting_registry
            .claim_artifact_reconciliation(ClaimArtifactReconciliationRequest {
                artifact: &waiting_artifact,
                token: ArtifactClaimToken::generate(),
                claimed_at: caller_time,
                lease_expires_at: caller_time + time::Duration::minutes(5),
                detached_grace: time::Duration::hours(24),
            })
            .await
    });
    tokio::time::sleep(StdDuration::from_millis(1_100)).await;
    blocker.commit().await.unwrap();
    assert!(matches!(
        waiting.await.unwrap().unwrap(),
        ArtifactClaimOutcome::Claimed(_)
    ));

    let recoverable = store.put_registered(b"single winner\n").unwrap();
    let prior_token = ArtifactClaimToken::generate();
    let prior_claim = match registry
        .claim_artifact_reconciliation(ClaimArtifactReconciliationRequest {
            artifact: &recoverable,
            token: prior_token.clone(),
            claimed_at: caller_time,
            lease_expires_at: caller_time + time::Duration::minutes(5),
            detached_grace: time::Duration::hours(24),
        })
        .await
        .unwrap()
    {
        ArtifactClaimOutcome::Claimed(claim) => claim,
        outcome => panic!("unexpected claim outcome: {outcome:?}"),
    };
    sqlx::query(
        "update artifact_blobs
            set claim_expires_at = clock_timestamp() - interval '1 second'
          where digest = $1",
    )
    .bind(recoverable.artifact_ref().digest_hex())
    .execute(&raw_pool)
    .await
    .unwrap();
    assert!(matches!(
        registry
            .claim_artifact_reconciliation(ClaimArtifactReconciliationRequest {
                artifact: &recoverable,
                token: prior_token.clone(),
                claimed_at: caller_time,
                lease_expires_at: caller_time + time::Duration::minutes(5),
                detached_grace: time::Duration::hours(24),
            })
            .await
            .unwrap(),
        ArtifactClaimOutcome::HeldByOther
    ));
    let recovery_now = registry.artifact_reconciliation_now().await.unwrap();
    let probe = registry
        .list_expired_artifact_reconciliation_probes(recovery_now, 1)
        .await
        .unwrap()
        .pop()
        .unwrap();
    assert!(matches!(
        registry
            .claim_expired_artifact_reconciliation(
                &probe,
                ClaimExpiredArtifactReconciliationRequest {
                    token: prior_token,
                    claimed_at: caller_time,
                    lease_expires_at: caller_time + time::Duration::minutes(5),
                },
            )
            .await,
        Err(RegistryError::InvalidArtifactSource {
            field: "claim_token"
        })
    ));

    let first_registry = registry.clone();
    let second_registry = registry.clone();
    let first_probe = probe.clone();
    let (first, second) = tokio::join!(
        async move {
            first_registry
                .claim_expired_artifact_reconciliation(
                    &first_probe,
                    ClaimExpiredArtifactReconciliationRequest {
                        token: ArtifactClaimToken::generate(),
                        claimed_at: caller_time,
                        lease_expires_at: caller_time + time::Duration::minutes(5),
                    },
                )
                .await
        },
        async move {
            second_registry
                .claim_expired_artifact_reconciliation(
                    &probe,
                    ClaimExpiredArtifactReconciliationRequest {
                        token: ArtifactClaimToken::generate(),
                        claimed_at: caller_time,
                        lease_expires_at: caller_time + time::Duration::minutes(5),
                    },
                )
                .await
        }
    );
    let first = first.unwrap();
    let second = second.unwrap();
    assert_eq!(
        usize::from(first.is_some()) + usize::from(second.is_some()),
        1
    );
    let winner = first.or(second).unwrap();

    sqlx::query("update artifact_blobs set claim_expires_at = clock_timestamp() where digest = $1")
        .bind(recoverable.artifact_ref().digest_hex())
        .execute(&raw_pool)
        .await
        .unwrap();
    for claim in [&winner, &prior_claim] {
        assert_eq!(
            registry
                .finalize_artifact_reconciliation_claim(
                    claim,
                    ArtifactClaimFinalization::AlreadyAbsent,
                    caller_time,
                )
                .await
                .unwrap(),
            ArtifactClaimFinalizeOutcome::Stale
        );
    }

    let invalid = store.put_registered(b"invalid bounds\n").unwrap();
    assert!(matches!(
        registry
            .claim_artifact_reconciliation(ClaimArtifactReconciliationRequest {
                artifact: &invalid,
                token: ArtifactClaimToken::generate(),
                claimed_at: caller_time,
                lease_expires_at: caller_time + time::Duration::minutes(5),
                detached_grace: time::Duration::MAX,
            })
            .await,
        Err(RegistryError::InvalidArtifactSource {
            field: "detached_grace"
        })
    ));
    assert!(matches!(
        registry
            .claim_artifact_reconciliation(ClaimArtifactReconciliationRequest {
                artifact: &invalid,
                token: ArtifactClaimToken::generate(),
                claimed_at: caller_time,
                lease_expires_at: caller_time
                    + time::Duration::minutes(5)
                    + time::Duration::seconds(1),
                detached_grace: time::Duration::ZERO,
            })
            .await,
        Err(RegistryError::InvalidArtifactSource {
            field: "claim_lease"
        })
    ));

    database.cleanup().await;
}