remora/crates/remora-attempt/tests/authority.rs

use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use remora_admission::{AdmissionConfig, AdmissionDecision, AdmissionService};
use remora_attempt::{
    ActivationOutcome, AttemptConfig, AttemptFaultInjector, AttemptFaultPoint, AttemptService,
    AttemptState, ClaimOutcome, RecoveryOutcome, RenewalOutcome,
};
use remora_contracts::{HttpSourceSubmissionDto, JobSubmissionDto, SourceSubmissionDto};
use remora_coordination::{ActivateOptions, CoordinationStore, LeaseOwner, StoreConfig};
use remora_domain::{BucketName, OpaqueId, StorageScope, TargetPrefix, TargetStoreId};
use remora_durable::{MutationDisposition, PoolConfig, connect};
use remora_identity::{DefaultTarget, Principal};
use remora_schema::migrate as migrate_schema;
use sqlx::PgPool;
use sqlx::postgres::PgPoolOptions;
use tokio::time::sleep;

struct Harness {
    pool: PgPool,
    coordination: CoordinationStore,
}

impl Harness {
    async fn new(name: &str, ttl: Duration, global_limit: u64) -> Option<Self> {
        let database_url = required_env("DATABASE_URL")?;
        let coordination_url = required_env("REMORA_COORDINATION_URL")?;
        let pool = connect(
            &database_url,
            PoolConfig {
                max_connections: 12,
                acquire_timeout: Duration::from_secs(10),
                idle_timeout: Some(Duration::from_secs(30)),
            },
        )
        .await
        .expect("connect postgres");
        migrate_schema(&pool).await.expect("schema migrate");
        let namespace = format!("{}wp15:{name}:{}:", coordination_namespace(), nonce());
        let coordination = CoordinationStore::connect(
            &coordination_url,
            namespace,
            StoreConfig {
                global_limit,
                lease_ttl: ttl,
                reserve_ttl: ttl,
                progress_ttl: ttl,
                cancel_ttl: None,
            },
        )
        .await
        .expect("connect coordination store");
        Some(Self { pool, coordination })
    }

    fn service(&self, max_attempts: u64) -> AttemptService {
        AttemptService::new(
            self.pool.clone(),
            self.coordination.clone(),
            AttemptConfig { max_attempts },
        )
        .expect("attempt service")
    }

    fn service_with_fault(&self, point: AttemptFaultPoint, max_attempts: u64) -> AttemptService {
        AttemptService::new_with_fault_injector(
            self.pool.clone(),
            self.coordination.clone(),
            AttemptConfig { max_attempts },
            Arc::new(PanicOnce::new(point)),
        )
        .expect("fault-injected attempt service")
    }

    async fn job(&self) -> OpaqueId {
        let owner = unique_id("owner");
        let store = TargetStoreId::parse("primary").unwrap();
        let bucket = BucketName::parse("remora-results").unwrap();
        let prefix = TargetPrefix::parse("users/alice", ["_remora"]).unwrap();
        let principal = Principal {
            id: owner,
            issuer: "https://issuer.example".into(),
            subject: "subject".into(),
            administrator: false,
            scopes: [StorageScope {
                store: store.clone(),
                bucket: bucket.clone(),
                prefix: prefix.clone(),
            }]
            .into_iter()
            .collect(),
        };
        let admission = AdmissionService::new(
            self.pool.clone(),
            AdmissionConfig {
                default_target: DefaultTarget {
                    store,
                    bucket,
                    prefix,
                },
                reserved_target_segments: vec!["_remora".into()],
                policy_version: "attempt-tests-v1".into(),
            },
        )
        .unwrap();
        let request = JobSubmissionDto {
            source: SourceSubmissionDto::Http(HttpSourceSubmissionDto {
                url: "https://example.invalid/artifact".into(),
                result_path: "artifact.bin".into(),
                credential_id: None,
            }),
            target: None,
            label: Some("attempt-test".into()),
        };
        let outcome = admission
            .submit(&principal, unique_id("submit").as_str(), request)
            .await
            .expect("submit job");
        let AdmissionDecision::Accepted(job) = outcome.decision else {
            panic!("job rejected");
        };
        self.coordination
            .enqueue(job.job_id.as_str())
            .await
            .expect("project queue");
        job.job_id
    }
}

#[tokio::test]
async fn claim_replay_conflict_and_fence_allocation_are_durable() {
    let Some(harness) = Harness::new("claim", Duration::from_secs(2), 8).await else {
        return;
    };
    let job = harness.job().await;
    let worker = unique_id("worker");
    let operation = unique_id("claimop");
    let service = harness.service(2);

    let first = claimed(
        service
            .claim(&job, &worker, &operation)
            .await
            .expect("claim"),
    );
    assert_eq!(first.0, 1);
    assert_eq!(first.1, MutationDisposition::First);
    let replay = claimed(
        service
            .claim(&job, &worker, &operation)
            .await
            .expect("claim replay"),
    );
    assert_eq!(replay.0, 1);
    assert_eq!(replay.1, MutationDisposition::Replay);
    assert_eq!(first.2, replay.2);

    let other_job = harness.job().await;
    assert!(
        service
            .claim(&other_job, &worker, &operation)
            .await
            .is_err(),
        "one operation identity cannot be rebound to another job"
    );

    sleep(Duration::from_millis(2200)).await;
    let authority = service.authority(&first.2).await.unwrap().unwrap();
    assert_eq!(
        service.recover_inactive(&authority).await.unwrap(),
        RecoveryOutcome::Reaped
    );
    harness.coordination.enqueue(job.as_str()).await.unwrap();
    let second = claimed(
        service
            .claim(&job, &worker, &unique_id("claimop"))
            .await
            .unwrap(),
    );
    assert_eq!(second.0, 2);
    sleep(Duration::from_millis(2200)).await;
    let second_authority = service.authority(&second.2).await.unwrap().unwrap();
    service.recover_inactive(&second_authority).await.unwrap();
    harness.coordination.enqueue(job.as_str()).await.unwrap();
    let exhausted = service
        .claim(&job, &worker, &unique_id("claimop"))
        .await
        .unwrap();
    assert!(matches!(exhausted, ClaimOutcome::Rejected { .. }));
}

#[tokio::test]
async fn activation_replay_requires_the_complete_live_authority_tuple() {
    let Some(harness) = Harness::new("activation", Duration::from_secs(3), 8).await else {
        return;
    };
    let job = harness.job().await;
    let worker = unique_id("worker");
    let service = harness.service(3);
    let claim_operation = unique_id("claimop");
    let (_, _, attempt) = claimed(
        service
            .claim(&job, &worker, &claim_operation)
            .await
            .unwrap(),
    );
    let authority = service.authority(&attempt).await.unwrap().unwrap();
    assert!(matches!(
        service
            .claim(&job, &worker, &unique_id("claimop"))
            .await
            .unwrap(),
        ClaimOutcome::Rejected { .. }
    ));
    let unavailable_pool = PgPoolOptions::new()
        .acquire_timeout(Duration::from_millis(100))
        .connect_lazy("postgres://remora:[email protected]:1/remora")
        .unwrap();
    let unavailable = AttemptService::new(
        unavailable_pool,
        harness.coordination.clone(),
        AttemptConfig { max_attempts: 3 },
    )
    .unwrap();
    assert!(
        unavailable
            .claim(&job, &worker, &claim_operation)
            .await
            .is_err()
    );
    assert_eq!(
        harness
            .coordination
            .get_reservation(job.as_str())
            .await
            .unwrap()
            .unwrap()
            .worker_id,
        worker.as_str(),
        "durable outage before mutation observation must preserve the claim reservation"
    );
    assert!(
        service
            .activate(&authority, &claim_operation)
            .await
            .is_err(),
        "an operation identity cannot change mutation kind"
    );
    let operation = unique_id("activateop");
    assert_eq!(
        service.activate(&authority, &operation).await.unwrap(),
        ActivationOutcome::Activated(MutationDisposition::First)
    );
    let executing = service.authority(&attempt).await.unwrap().unwrap();
    assert_eq!(executing.state, AttemptState::Executing);
    assert_eq!(
        service.activate(&executing, &operation).await.unwrap(),
        ActivationOutcome::Activated(MutationDisposition::Replay)
    );
    assert_eq!(
        service
            .activate(&executing, &unique_id("activateop"))
            .await
            .unwrap(),
        ActivationOutcome::Activated(MutationDisposition::First)
    );
    assert_eq!(
        service.renew(&authority).await.unwrap(),
        RenewalOutcome::Renewed
    );

    let mut forged = authority.clone();
    forged.worker_incarnation = unique_id("worker");
    assert!(service.activate(&forged, &operation).await.is_err());
    assert_eq!(
        service.renew(&authority).await.unwrap(),
        RenewalOutcome::Renewed,
        "conflicting operation must preserve the established lease"
    );
    assert_eq!(
        service.recover_inactive(&authority).await.unwrap(),
        RecoveryOutcome::StillOwned
    );
    assert_eq!(
        service.renew(&authority).await.unwrap(),
        RenewalOutcome::Renewed,
        "inactive recovery must preserve an exact executing authority"
    );

    harness
        .coordination
        .release(job.as_str(), &owner(&authority))
        .await
        .unwrap();
    assert_eq!(
        service.activate(&authority, &operation).await.unwrap(),
        ActivationOutcome::PreconditionFailed
    );
    assert_eq!(
        service.renew(&authority).await.unwrap(),
        RenewalOutcome::AuthorityLost
    );
    assert_eq!(
        service.recover_inactive(&authority).await.unwrap(),
        RecoveryOutcome::StillOwned,
        "missing coordination never reverses executing durable authority"
    );
}

#[tokio::test]
async fn activation_and_inactive_recovery_serialize_on_durable_authority() {
    let Some(harness) = Harness::new("serialize", Duration::from_millis(700), 8).await else {
        return;
    };
    let job = harness.job().await;
    let worker = unique_id("worker");
    let service = harness.service(3);
    let (_, _, attempt) = claimed(
        service
            .claim(&job, &worker, &unique_id("claimop"))
            .await
            .unwrap(),
    );
    let authority = service.authority(&attempt).await.unwrap().unwrap();
    sleep(Duration::from_millis(800)).await;

    let activation_service = service.clone();
    let activation_authority = authority.clone();
    let activation = tokio::spawn(async move {
        activation_service
            .activate(&activation_authority, &unique_id("activateop"))
            .await
            .unwrap()
    });
    let recovery_service = service.clone();
    let recovery_authority = authority.clone();
    let recovery = tokio::spawn(async move {
        recovery_service
            .recover_inactive(&recovery_authority)
            .await
            .unwrap()
    });
    let activation = activation.await.unwrap();
    let recovery = recovery.await.unwrap();
    assert!(
        matches!(
            (activation, recovery),
            (
                ActivationOutcome::PreconditionFailed,
                RecoveryOutcome::Reaped
            )
        ),
        "expired reservation cannot activate while recovery reaps: {activation:?} {recovery:?}"
    );
    assert!(service.authority(&attempt).await.unwrap().is_none());
}

#[tokio::test]
async fn global_permits_bound_claims_and_survive_lease_expiry() {
    let Some(harness) = Harness::new("permits", Duration::from_millis(500), 1).await else {
        return;
    };
    let service = harness.service(3);
    let worker = unique_id("worker");
    let first_job = harness.job().await;
    let (_, _, attempt) = claimed(
        service
            .claim(&first_job, &worker, &unique_id("claimop"))
            .await
            .unwrap(),
    );
    let authority = service.authority(&attempt).await.unwrap().unwrap();
    service
        .activate(&authority, &unique_id("activateop"))
        .await
        .unwrap();
    sleep(Duration::from_millis(600)).await;
    let second_job = harness.job().await;
    assert!(matches!(
        service
            .claim(&second_job, &worker, &unique_id("claimop"))
            .await
            .unwrap(),
        ClaimOutcome::Rejected { .. }
    ));
    harness
        .coordination
        .expire_lease(first_job.as_str())
        .await
        .unwrap();
}

#[tokio::test]
async fn precommit_lease_is_compensated_when_durable_authority_was_reaped() {
    let Some(harness) = Harness::new("compensate", Duration::from_millis(700), 8).await else {
        return;
    };
    let job = harness.job().await;
    let worker = unique_id("worker");
    let service = harness.service(3);
    let (_, _, attempt) = claimed(
        service
            .claim(&job, &worker, &unique_id("claimop"))
            .await
            .unwrap(),
    );
    let authority = service.authority(&attempt).await.unwrap().unwrap();
    let lease = harness
        .coordination
        .activate(
            job.as_str(),
            &owner(&authority),
            ActivateOptions {
                seed_progress: true,
                require_reserve: true,
            },
        )
        .await
        .unwrap();
    assert!(lease.ok());
    harness
        .coordination
        .release(job.as_str(), &owner(&authority))
        .await
        .unwrap();
    sleep(Duration::from_millis(800)).await;
    assert_eq!(
        service.recover_inactive(&authority).await.unwrap(),
        RecoveryOutcome::Reaped
    );
    assert_eq!(
        service
            .activate(&authority, &unique_id("activateop"))
            .await
            .unwrap(),
        ActivationOutcome::PreconditionFailed
    );
    assert!(
        harness
            .coordination
            .get_live_lease(job.as_str())
            .await
            .unwrap()
            .is_none()
    );
}

#[tokio::test]
async fn crash_injection_covers_every_cross_store_and_lost_response_boundary() {
    let Some(harness) = Harness::new("crash-matrix", Duration::from_millis(500), 32).await else {
        return;
    };
    let worker = unique_id("worker");
    let normal = harness.service(32);

    let before_reserve_job = harness.job().await;
    crash_claim(
        harness.service_with_fault(AttemptFaultPoint::ClaimBeforeReserve, 32),
        before_reserve_job.clone(),
        worker.clone(),
        unique_id("claimop"),
    )
    .await;
    assert!(
        harness
            .coordination
            .get_reservation(before_reserve_job.as_str())
            .await
            .unwrap()
            .is_none()
    );

    let after_reserve_job = harness.job().await;
    crash_claim(
        harness.service_with_fault(AttemptFaultPoint::ClaimAfterReserve, 32),
        after_reserve_job.clone(),
        worker.clone(),
        unique_id("claimop"),
    )
    .await;
    assert!(
        harness
            .coordination
            .get_reservation(after_reserve_job.as_str())
            .await
            .unwrap()
            .is_some()
    );
    harness
        .coordination
        .free_reserve(after_reserve_job.as_str(), Some(worker.as_str()))
        .await
        .unwrap();

    let lost_claim_job = harness.job().await;
    let lost_claim_operation = unique_id("claimop");
    crash_claim(
        harness.service_with_fault(AttemptFaultPoint::ClaimAfterDurable, 32),
        lost_claim_job.clone(),
        worker.clone(),
        lost_claim_operation.clone(),
    )
    .await;
    let (_, disposition, lost_claim_attempt) = claimed(
        normal
            .claim(&lost_claim_job, &worker, &lost_claim_operation)
            .await
            .unwrap(),
    );
    assert_eq!(disposition, MutationDisposition::Replay);
    let lost_claim_authority = normal
        .authority(&lost_claim_attempt)
        .await
        .unwrap()
        .unwrap();

    let before_lease_operation = unique_id("activateop");
    crash_activate(
        harness.service_with_fault(AttemptFaultPoint::ActivationBeforeLease, 32),
        lost_claim_authority.clone(),
        before_lease_operation.clone(),
    )
    .await;
    assert!(
        harness
            .coordination
            .get_live_lease(lost_claim_job.as_str())
            .await
            .unwrap()
            .is_none()
    );
    assert_eq!(
        normal
            .activate(&lost_claim_authority, &before_lease_operation)
            .await
            .unwrap(),
        ActivationOutcome::Activated(MutationDisposition::First)
    );

    let after_lease_job = harness.job().await;
    let (_, _, after_lease_attempt) = claimed(
        normal
            .claim(&after_lease_job, &worker, &unique_id("claimop"))
            .await
            .unwrap(),
    );
    let after_lease_authority = normal
        .authority(&after_lease_attempt)
        .await
        .unwrap()
        .unwrap();
    let after_lease_operation = unique_id("activateop");
    crash_activate(
        harness.service_with_fault(AttemptFaultPoint::ActivationAfterLease, 32),
        after_lease_authority.clone(),
        after_lease_operation.clone(),
    )
    .await;
    assert!(
        harness
            .coordination
            .get_live_lease(after_lease_job.as_str())
            .await
            .unwrap()
            .is_some()
    );
    assert_eq!(
        normal
            .activate(&after_lease_authority, &after_lease_operation)
            .await
            .unwrap(),
        ActivationOutcome::Activated(MutationDisposition::First)
    );

    let lost_activation_job = harness.job().await;
    let (_, _, lost_activation_attempt) = claimed(
        normal
            .claim(&lost_activation_job, &worker, &unique_id("claimop"))
            .await
            .unwrap(),
    );
    let lost_activation_authority = normal
        .authority(&lost_activation_attempt)
        .await
        .unwrap()
        .unwrap();
    let lost_activation_operation = unique_id("activateop");
    crash_activate(
        harness.service_with_fault(AttemptFaultPoint::ActivationAfterDurable, 32),
        lost_activation_authority.clone(),
        lost_activation_operation.clone(),
    )
    .await;
    let executing = normal
        .authority(&lost_activation_attempt)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(
        normal
            .activate(&executing, &lost_activation_operation)
            .await
            .unwrap(),
        ActivationOutcome::Activated(MutationDisposition::Replay)
    );

    crash_renew(
        harness.service_with_fault(AttemptFaultPoint::RenewalBeforeLease, 32),
        executing.clone(),
    )
    .await;
    crash_renew(
        harness.service_with_fault(AttemptFaultPoint::RenewalAfterLease, 32),
        executing.clone(),
    )
    .await;
    assert_eq!(
        normal.renew(&executing).await.unwrap(),
        RenewalOutcome::Renewed
    );

    let mut recovery_cases = Vec::new();
    for point in [
        AttemptFaultPoint::RecoveryBeforeRevalidation,
        AttemptFaultPoint::RecoveryBeforeDurable,
        AttemptFaultPoint::RecoveryAfterDurable,
    ] {
        let job = harness.job().await;
        let (_, _, attempt) = claimed(
            normal
                .claim(&job, &worker, &unique_id("claimop"))
                .await
                .unwrap(),
        );
        let authority = normal.authority(&attempt).await.unwrap().unwrap();
        recovery_cases.push((point, authority));
    }
    sleep(Duration::from_millis(600)).await;
    for (point, authority) in recovery_cases {
        let used_before = harness.coordination.global_used().await.unwrap();
        crash_recovery(harness.service_with_fault(point, 32), authority.clone()).await;
        let expected = if point == AttemptFaultPoint::RecoveryAfterDurable {
            RecoveryOutcome::AlreadyInactive
        } else {
            RecoveryOutcome::Reaped
        };
        assert_eq!(normal.recover_inactive(&authority).await.unwrap(), expected);
        assert_eq!(
            harness.coordination.global_used().await.unwrap(),
            used_before - 1,
            "recovery retry must release the expired attempt's global slot"
        );
    }
}

fn required_env(name: &str) -> Option<String> {
    match std::env::var(name).ok().filter(|value| !value.is_empty()) {
        Some(value) => Some(value),
        None if std::env::var("REMORA_ATTEMPT_REQUIRE").is_ok() => {
            panic!("{name} is required when REMORA_ATTEMPT_REQUIRE is set")
        }
        None => None,
    }
}

fn coordination_namespace() -> String {
    std::env::var("REMORA_COORDINATION_NAMESPACE")
        .ok()
        .filter(|value| !value.is_empty())
        .unwrap_or_else(|| "remora:test:".into())
}

fn unique_id(prefix: &str) -> OpaqueId {
    OpaqueId::parse(format!("{prefix}{:032x}", nonce())).unwrap()
}

fn nonce() -> u128 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_nanos()
}

fn claimed(outcome: ClaimOutcome) -> (u64, MutationDisposition, OpaqueId) {
    let ClaimOutcome::Claimed {
        authority,
        disposition: Some(disposition),
    } = outcome
    else {
        panic!("claim rejected")
    };
    assert_eq!(authority.state, AttemptState::Authorized);
    (authority.fence, disposition, authority.attempt_id)
}

fn owner(authority: &remora_attempt::AttemptAuthority) -> LeaseOwner {
    LeaseOwner::new(
        authority.attempt_id.as_str(),
        authority.fence,
        authority.worker_incarnation.as_str(),
    )
    .unwrap()
}

struct PanicOnce {
    armed: Mutex<HashSet<AttemptFaultPoint>>,
}

impl PanicOnce {
    fn new(point: AttemptFaultPoint) -> Self {
        Self {
            armed: Mutex::new(HashSet::from([point])),
        }
    }
}

impl AttemptFaultInjector for PanicOnce {
    fn hit(&self, point: AttemptFaultPoint) {
        if self.armed.lock().unwrap().remove(&point) {
            panic!("injected crash at {point:?}");
        }
    }
}

async fn crash_claim(
    service: AttemptService,
    job: OpaqueId,
    worker: OpaqueId,
    operation: OpaqueId,
) {
    assert!(
        tokio::spawn(async move { service.claim(&job, &worker, &operation).await })
            .await
            .expect_err("claim must crash")
            .is_panic()
    );
}

async fn crash_activate(
    service: AttemptService,
    authority: remora_attempt::AttemptAuthority,
    operation: OpaqueId,
) {
    assert!(
        tokio::spawn(async move { service.activate(&authority, &operation).await })
            .await
            .expect_err("activation must crash")
            .is_panic()
    );
}

async fn crash_renew(service: AttemptService, authority: remora_attempt::AttemptAuthority) {
    assert!(
        tokio::spawn(async move { service.renew(&authority).await })
            .await
            .expect_err("renewal must crash")
            .is_panic()
    );
}

async fn crash_recovery(service: AttemptService, authority: remora_attempt::AttemptAuthority) {
    assert!(
        tokio::spawn(async move { service.recover_inactive(&authority).await })
            .await
            .expect_err("recovery must crash")
            .is_panic()
    );
}