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

use std::time::Duration;

use remora_admission::{
    AdmissionConfig, AdmissionDecision, AdmissionService, migrate as admission_migrate,
};
use remora_attempt::{
    ActivationOutcome, AttemptConfig, AttemptMutationOutcome, AttemptService, AttemptState,
    CheckpointProjection, CheckpointRecord, CommittedResultRecord, DurableDescriptor,
    FailureRecord, MutationRejection, ProgressRecord, migrate as attempt_migrate,
};
use remora_contracts::{HttpSourceSubmissionDto, JobSubmissionDto, SourceSubmissionDto};
use remora_coordination::{CoordinationStore, StoreConfig};
use remora_credentials::migrate as credentials_migrate;
use remora_domain::{
    BucketName, FailureDisposition, JobPhase, OpaqueId, RuntimeFailureClass, StorageScope,
    TargetPrefix, TargetStoreId,
};
use remora_durable::{MutationDisposition, PoolConfig, connect, migrate as durable_migrate};
use remora_identity::{DefaultTarget, Principal};
use remora_result::{PublicationInput, migrate as result_migrate};
use remora_target::{ContentSha256, ObjectIdentity, StoredObject, attempt_manifest_key};
use sqlx::PgPool;

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

impl Harness {
    async fn new(name: &str) -> 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: 16,
                acquire_timeout: Duration::from_secs(10),
                idle_timeout: Some(Duration::from_secs(30)),
            },
        )
        .await
        .expect("connect postgres");
        durable_migrate(&pool).await.expect("durable migrate");
        credentials_migrate(&pool)
            .await
            .expect("credentials migrate");
        admission_migrate(&pool).await.expect("admission migrate");
        attempt_migrate(&pool).await.expect("attempt migrate");
        result_migrate(&pool).await.expect("result migrate");
        let coordination = CoordinationStore::connect(
            &coordination_url,
            format!("{}wp16:{name}:{}:", coordination_namespace(), nonce()),
            StoreConfig {
                global_limit: 16,
                lease_ttl: Duration::from_secs(10),
                reserve_ttl: Duration::from_secs(10),
                progress_ttl: Duration::from_secs(10),
                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,
                local_limit: 8,
            },
        )
        .expect("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: "mutation-tests-v1".into(),
            },
        )
        .unwrap();
        let outcome = admission
            .submit(
                &principal,
                unique_id("submit").as_str(),
                JobSubmissionDto {
                    source: SourceSubmissionDto::Http(HttpSourceSubmissionDto {
                        url: "https://example.invalid/artifact".into(),
                        result_path: "artifact.bin".into(),
                        credential_id: None,
                    }),
                    target: None,
                    label: Some("mutation-test".into()),
                },
            )
            .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
    }

    async fn executing(
        &self,
        service: &AttemptService,
        job: &OpaqueId,
    ) -> remora_attempt::AttemptAuthority {
        let worker = unique_id("worker");
        let claim = service
            .claim(job, &worker, &unique_id("claimop"))
            .await
            .expect("claim");
        let remora_attempt::ClaimOutcome::Claimed { authority, .. } = claim else {
            panic!("claim rejected")
        };
        assert_eq!(authority.state, AttemptState::Authorized);
        assert!(matches!(
            service
                .activate(&authority, &unique_id("activateop"))
                .await
                .expect("activate"),
            ActivationOutcome::Activated(_)
        ));
        service
            .authority(&authority.attempt_id)
            .await
            .unwrap()
            .unwrap()
    }

    async fn sealed_result(
        &self,
        attempt_id: &OpaqueId,
        member_count: u64,
        total_bytes: u64,
    ) -> PublicationInput {
        let manifest_id = unique_id("manifest");
        let sha256 = format!("sha256:{}", "11".repeat(32));
        let key = attempt_manifest_key(attempt_id, &manifest_id);
        let input = PublicationInput {
            attempt_id: attempt_id.clone(),
            manifest_id: manifest_id.clone(),
            schema_version: 1,
            manifest: StoredObject {
                identity: ObjectIdentity {
                    target_store_id: "primary".into(),
                    bucket: "remora-results".into(),
                    key: key.clone(),
                    version_id: Some("version-1".into()),
                },
                content_sha256: ContentSha256::parse(&sha256).unwrap(),
                size: 128,
            },
            member_count,
            total_bytes,
        };
        sqlx::query(
            "INSERT INTO result_inventories (
                attempt_id, state, manifest_schema_version, manifest_id, manifest_store,
                manifest_bucket, manifest_key, manifest_version_id, manifest_sha256,
                manifest_size, member_count, total_bytes
             ) VALUES ($1, 'sealed', 1, $2, 'primary', 'remora-results', $3,
                       'version-1', $4, 128, $5, $6)",
        )
        .bind(attempt_id.as_str())
        .bind(manifest_id.as_str())
        .bind(key)
        .bind(sha256)
        .bind(i64::try_from(member_count).unwrap())
        .bind(i64::try_from(total_bytes).unwrap())
        .execute(&self.pool)
        .await
        .unwrap();
        input
    }
}

#[tokio::test]
async fn resolved_source_checkpoint_and_progress_are_fenced_and_replayable() {
    let Some(harness) = Harness::new("mutations").await else {
        return;
    };
    let service = harness.service(3);
    let job = harness.job().await;
    let worker = unique_id("worker");
    let claim = service
        .claim(&job, &worker, &unique_id("claimop"))
        .await
        .unwrap();
    let remora_attempt::ClaimOutcome::Claimed { authority, .. } = claim else {
        panic!("claim rejected")
    };
    let resolved = descriptor("resolved", 4096, Some(3), 'a');
    assert_rejected(
        service
            .pin_resolved_source(&authority, &unique_id("resolveop"), resolved.clone())
            .await
            .unwrap(),
        MutationRejection::AuthorityLost,
    );
    service
        .activate(&authority, &unique_id("activateop"))
        .await
        .unwrap();
    let authority = service
        .authority(&authority.attempt_id)
        .await
        .unwrap()
        .unwrap();
    let operation = unique_id("resolveop");
    assert_eq!(
        service
            .pin_resolved_source(&authority, &operation, resolved.clone())
            .await
            .unwrap(),
        AttemptMutationOutcome::Accepted(MutationDisposition::First)
    );
    assert_eq!(
        service
            .pin_resolved_source(&authority, &operation, resolved.clone())
            .await
            .unwrap(),
        AttemptMutationOutcome::Accepted(MutationDisposition::Replay)
    );
    assert_eq!(
        service
            .pin_resolved_source(&authority, &unique_id("resolveop"), resolved.clone())
            .await
            .unwrap(),
        AttemptMutationOutcome::Accepted(MutationDisposition::First)
    );
    assert_rejected(
        service
            .pin_resolved_source(
                &authority,
                &unique_id("resolveop"),
                descriptor("resolved", 4096, Some(3), 'b'),
            )
            .await
            .unwrap(),
        MutationRejection::ResolvedSourceConflict,
    );

    assert!(
        service
            .persist_checkpoint(
                &authority,
                &unique_id("checkpointop"),
                CheckpointRecord {
                    descriptor: descriptor("checkpoint", 1, Some(0), '0'),
                    resolved_source_id: resolved.id.clone(),
                    validated_bytes: 0,
                    validated_members: 0,
                },
            )
            .await
            .is_err(),
        "an all-zero checkpoint does not refine PersistCheckpoint"
    );

    let checkpoint = CheckpointRecord {
        descriptor: descriptor("checkpoint", 512, Some(1), 'c'),
        resolved_source_id: resolved.id.clone(),
        validated_bytes: 1024,
        validated_members: 1,
    };
    let mut wrong_source = checkpoint.clone();
    wrong_source.resolved_source_id = unique_id("resolved");
    assert_rejected(
        service
            .persist_checkpoint(&authority, &unique_id("checkpointop"), wrong_source)
            .await
            .unwrap(),
        MutationRejection::CheckpointSourceMismatch,
    );
    assert!(matches!(
        service
            .persist_checkpoint(&authority, &unique_id("checkpointop"), checkpoint.clone())
            .await
            .unwrap(),
        AttemptMutationOutcome::Accepted(_)
    ));
    let mut regression = checkpoint.clone();
    regression.validated_bytes = 1000;
    assert_rejected(
        service
            .persist_checkpoint(&authority, &unique_id("checkpointop"), regression)
            .await
            .unwrap(),
        MutationRejection::CheckpointRegression,
    );

    let progress = ProgressRecord {
        completed_bytes: 2048,
        total_bytes: Some(4096),
        completed_members: 1,
        total_members: Some(3),
    };
    let progress_operation = unique_id("progressop");
    assert_eq!(
        service
            .report_progress(&authority, &progress_operation, progress)
            .await
            .unwrap(),
        AttemptMutationOutcome::Accepted(MutationDisposition::First)
    );
    assert_eq!(
        service
            .report_progress(&authority, &progress_operation, progress)
            .await
            .unwrap(),
        AttemptMutationOutcome::Accepted(MutationDisposition::Replay)
    );
    let hint = harness
        .coordination
        .get_progress(job.as_str())
        .await
        .unwrap()
        .unwrap();
    assert_eq!(hint.bytes, 2048);
    assert_eq!(hint.total_bytes, Some(4096));
    assert_eq!(hint.completed_members, 1);
    assert_eq!(hint.total_members, Some(3));
    assert_rejected(
        service
            .report_progress(
                &authority,
                &unique_id("progressop"),
                ProgressRecord {
                    completed_bytes: 1024,
                    ..progress
                },
            )
            .await
            .unwrap(),
        MutationRejection::ProgressRegression,
    );
    assert!(
        service
            .persist_checkpoint(&authority, &progress_operation, checkpoint)
            .await
            .is_err(),
        "an operation identity cannot change mutation kind"
    );
}

#[tokio::test]
async fn success_is_an_exact_fenced_commit_cas_and_terminal_state_is_immutable() {
    let Some(harness) = Harness::new("success").await else {
        return;
    };
    let service = harness.service(3);
    let job = harness.job().await;
    let authority = harness.executing(&service, &job).await;
    let result = harness.sealed_result(&authority.attempt_id, 2, 8192).await;
    let operation = unique_id("successop");
    assert_eq!(
        service
            .publish_success(&authority, &operation, result.clone())
            .await
            .unwrap(),
        AttemptMutationOutcome::Accepted(MutationDisposition::First)
    );
    assert_eq!(
        service
            .publish_success(&authority, &operation, result.clone())
            .await
            .unwrap(),
        AttemptMutationOutcome::Accepted(MutationDisposition::Replay)
    );
    let mut changed = result.clone();
    changed.total_bytes += 1;
    assert!(
        service
            .publish_success(&authority, &operation, changed)
            .await
            .is_err()
    );
    assert_rejected(
        service
            .publish_failure(
                &authority,
                &unique_id("failureop"),
                RuntimeFailureClass::Internal,
            )
            .await
            .unwrap(),
        MutationRejection::AuthorityLost,
    );
    assert_rejected(
        service
            .pin_resolved_source(
                &authority,
                &unique_id("resolveop"),
                descriptor("resolved", 1, None, 'd'),
            )
            .await
            .unwrap(),
        MutationRejection::AuthorityLost,
    );
    let projection = service.mutation_projection(&job).await.unwrap().unwrap();
    assert_eq!(projection.phase, JobPhase::Succeeded);
    assert_eq!(projection.authority_attempt_id, None);
    assert_eq!(
        projection.checkpoint,
        Some(CheckpointProjection::Complete {
            validated_bytes: 8192,
            validated_members: 2,
        })
    );
    assert_eq!(
        projection.committed_result,
        Some(CommittedResultRecord::from(&result))
    );
    assert_eq!(projection.failure, None);
    let state: String = sqlx::query_scalar("SELECT state FROM attempts WHERE attempt_id = $1")
        .bind(authority.attempt_id.as_str())
        .fetch_one(&harness.pool)
        .await
        .unwrap();
    assert_eq!(state, "committed");
}

#[tokio::test]
async fn invalid_result_stale_authority_and_cancellation_cannot_publish() {
    let Some(harness) = Harness::new("rejections").await else {
        return;
    };
    let service = harness.service(3);
    let job = harness.job().await;
    let authority = harness.executing(&service, &job).await;
    let mut invalid = harness.sealed_result(&authority.attempt_id, 1, 64).await;
    sqlx::query("UPDATE result_inventories SET manifest_store = 'other' WHERE attempt_id = $1")
        .bind(authority.attempt_id.as_str())
        .execute(&harness.pool)
        .await
        .unwrap();
    invalid.manifest.identity.target_store_id = "other".into();
    assert_rejected(
        service
            .publish_success(&authority, &unique_id("successop"), invalid.clone())
            .await
            .unwrap(),
        MutationRejection::InvalidResult,
    );
    sqlx::query("UPDATE result_inventories SET manifest_store = 'primary' WHERE attempt_id = $1")
        .bind(authority.attempt_id.as_str())
        .execute(&harness.pool)
        .await
        .unwrap();
    invalid.manifest.identity.target_store_id = "primary".into();
    invalid.manifest.size += 1;
    assert_rejected(
        service
            .publish_success(&authority, &unique_id("successop"), invalid)
            .await
            .unwrap(),
        MutationRejection::InvalidResult,
    );
    sqlx::query("UPDATE jobs SET cancel_requested = TRUE WHERE job_id = $1")
        .bind(job.as_str())
        .execute(&harness.pool)
        .await
        .unwrap();
    assert_rejected(
        service
            .publish_failure(
                &authority,
                &unique_id("failureop"),
                RuntimeFailureClass::TransientSource,
            )
            .await
            .unwrap(),
        MutationRejection::CancellationRequested,
    );
    let valid = publication_from_database(&harness.pool, &authority.attempt_id).await;
    assert_rejected(
        service
            .publish_success(&authority, &unique_id("successop"), valid)
            .await
            .unwrap(),
        MutationRejection::CancellationRequested,
    );
    assert_eq!(
        service
            .mutation_projection(&job)
            .await
            .unwrap()
            .unwrap()
            .phase,
        JobPhase::Running
    );
}

#[tokio::test]
async fn failure_disposition_retries_with_budget_and_becomes_terminal_when_exhausted() {
    let Some(harness) = Harness::new("retry").await else {
        return;
    };
    let service = harness.service(2);
    let job = harness.job().await;
    let first = harness.executing(&service, &job).await;
    assert_eq!(first.fence, 1);
    let retry_source = descriptor("resolved", 2048, Some(1), 'f');
    let retry_checkpoint = CheckpointRecord {
        descriptor: descriptor("checkpoint", 256, Some(1), 'e'),
        resolved_source_id: retry_source.id.clone(),
        validated_bytes: 512,
        validated_members: 1,
    };
    assert_rejected(
        service
            .persist_checkpoint(&first, &unique_id("checkpointop"), retry_checkpoint.clone())
            .await
            .unwrap(),
        MutationRejection::ResolvedSourceRequired,
    );
    service
        .pin_resolved_source(&first, &unique_id("resolveop"), retry_source)
        .await
        .unwrap();
    service
        .persist_checkpoint(&first, &unique_id("checkpointop"), retry_checkpoint)
        .await
        .unwrap();
    assert!(matches!(
        service
            .publish_failure(
                &first,
                &unique_id("failureop"),
                RuntimeFailureClass::TransientTarget,
            )
            .await
            .unwrap(),
        AttemptMutationOutcome::Accepted(_)
    ));
    let retry = service.mutation_projection(&job).await.unwrap().unwrap();
    assert_eq!(retry.phase, JobPhase::RetryWait);
    assert_eq!(retry.failure, None);
    assert_attempt_failure(
        &harness.pool,
        &first.attempt_id,
        FailureRecord {
            class: RuntimeFailureClass::TransientTarget,
            disposition: FailureDisposition::Retryable,
        },
    )
    .await;
    sqlx::query(
        "UPDATE jobs SET phase = 'queued', dispatch_needed = TRUE,
            revision = revision + 1 WHERE job_id = $1",
    )
    .bind(job.as_str())
    .execute(&harness.pool)
    .await
    .unwrap();
    harness
        .coordination
        .release(
            job.as_str(),
            &remora_coordination::LeaseOwner::new(
                first.attempt_id.as_str(),
                first.fence,
                first.worker_incarnation.as_str(),
            )
            .unwrap(),
        )
        .await
        .unwrap();
    service.release_local(&job);
    harness.coordination.enqueue(job.as_str()).await.unwrap();
    let second = harness.executing(&service, &job).await;
    assert_eq!(second.fence, 2);
    let initial_progress: i64 =
        sqlx::query_scalar("SELECT initial_progress FROM attempts WHERE attempt_id = $1")
            .bind(second.attempt_id.as_str())
            .fetch_one(&harness.pool)
            .await
            .unwrap();
    assert_eq!(initial_progress, 512);
    service
        .publish_failure(
            &second,
            &unique_id("failureop"),
            RuntimeFailureClass::Internal,
        )
        .await
        .unwrap();
    let exhausted = service.mutation_projection(&job).await.unwrap().unwrap();
    assert_eq!(exhausted.phase, JobPhase::Failed);
    assert_eq!(
        exhausted.failure,
        Some(FailureRecord {
            class: RuntimeFailureClass::Internal,
            disposition: FailureDisposition::Retryable,
        })
    );

    let permanent_job = harness.job().await;
    let permanent = harness.executing(&service, &permanent_job).await;
    service
        .publish_failure(
            &permanent,
            &unique_id("failureop"),
            RuntimeFailureClass::Authorization,
        )
        .await
        .unwrap();
    let projection = service
        .mutation_projection(&permanent_job)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(projection.phase, JobPhase::Failed);
    assert_eq!(
        projection.failure.unwrap().disposition,
        FailureDisposition::Permanent
    );
}

#[tokio::test]
async fn concurrent_success_and_failure_have_one_terminal_winner() {
    let Some(harness) = Harness::new("race").await else {
        return;
    };
    let service = harness.service(3);
    let job = harness.job().await;
    let authority = harness.executing(&service, &job).await;
    let result = harness.sealed_result(&authority.attempt_id, 1, 64).await;
    let success_service = service.clone();
    let success_authority = authority.clone();
    let success = tokio::spawn(async move {
        success_service
            .publish_success(&success_authority, &unique_id("successop"), result)
            .await
            .unwrap()
    });
    let failure_service = service.clone();
    let failure_authority = authority.clone();
    let failure = tokio::spawn(async move {
        failure_service
            .publish_failure(
                &failure_authority,
                &unique_id("failureop"),
                RuntimeFailureClass::Authorization,
            )
            .await
            .unwrap()
    });
    let (success, failure) = tokio::join!(success, failure);
    let outcomes = [success.unwrap(), failure.unwrap()];
    assert_eq!(
        outcomes
            .iter()
            .filter(|outcome| matches!(outcome, AttemptMutationOutcome::Accepted(_)))
            .count(),
        1
    );
    let projection = service.mutation_projection(&job).await.unwrap().unwrap();
    assert!(matches!(
        projection.phase,
        JobPhase::Succeeded | JobPhase::Failed
    ));
    assert_ne!(
        projection.committed_result.is_some(),
        projection.failure.is_some()
    );
}

fn descriptor(
    prefix: &str,
    byte_length: u64,
    item_count: Option<u64>,
    digest: char,
) -> DurableDescriptor {
    DurableDescriptor {
        id: unique_id(prefix),
        schema_version: 1,
        byte_length,
        item_count,
        checksum: format!("sha256:{}", digest.to_string().repeat(64)),
    }
}

fn assert_rejected(outcome: AttemptMutationOutcome, expected: MutationRejection) {
    assert!(
        matches!(outcome, AttemptMutationOutcome::Rejected { reason, .. } if reason == expected),
        "unexpected outcome: {outcome:?}"
    );
}

async fn assert_attempt_failure(pool: &PgPool, attempt_id: &OpaqueId, expected: FailureRecord) {
    let (state, failure): (String, serde_json::Value) =
        sqlx::query_as("SELECT state, failure FROM attempts WHERE attempt_id = $1")
            .bind(attempt_id.as_str())
            .fetch_one(pool)
            .await
            .unwrap();
    assert_eq!(state, "stopped");
    assert_eq!(failure, serde_json::to_value(expected).unwrap());
}

async fn publication_from_database(pool: &PgPool, attempt_id: &OpaqueId) -> PublicationInput {
    let row: (
        String,
        String,
        String,
        Option<String>,
        String,
        i64,
        i64,
        i64,
    ) = sqlx::query_as(
        "SELECT manifest_id, manifest_store, manifest_bucket, manifest_version_id,
                manifest_sha256, manifest_size, member_count, total_bytes
         FROM result_inventories WHERE attempt_id = $1",
    )
    .bind(attempt_id.as_str())
    .fetch_one(pool)
    .await
    .unwrap();
    let manifest_id = OpaqueId::parse(row.0).unwrap();
    PublicationInput {
        attempt_id: attempt_id.clone(),
        manifest_id: manifest_id.clone(),
        schema_version: 1,
        manifest: StoredObject {
            identity: ObjectIdentity {
                target_store_id: row.1,
                bucket: row.2,
                key: attempt_manifest_key(attempt_id, &manifest_id),
                version_id: row.3,
            },
            content_sha256: ContentSha256::parse(row.4).unwrap(),
            size: u64::try_from(row.5).unwrap(),
        },
        member_count: u64::try_from(row.6).unwrap(),
        total_bytes: u64::try_from(row.7).unwrap(),
    }
}

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()
}