use std::time::Duration; use remora_admission::{AdmissionConfig, AdmissionDecision, AdmissionService}; use remora_attempt::{ ActivationOutcome, AttemptConfig, AttemptMutationOutcome, AttemptService, AttemptState, CheckpointProjection, CheckpointRecord, CommittedResultRecord, DurableDescriptor, FailureRecord, MutationRejection, ProgressRecord, }; use remora_contracts::{HttpSourceSubmissionDto, JobSubmissionDto, SourceSubmissionDto}; use remora_coordination::{CoordinationStore, StoreConfig}; use remora_domain::{ BucketName, FailureDisposition, JobPhase, OpaqueId, RuntimeFailureClass, StorageScope, TargetPrefix, TargetStoreId, }; use remora_durable::{MutationDisposition, PoolConfig, connect}; use remora_identity::{DefaultTarget, Principal}; use remora_result::PublicationInput; use remora_schema::migrate as migrate_schema; use remora_target::{ContentSha256, ObjectIdentity, StoredObject, attempt_manifest_key_in}; use sqlx::PgPool; struct Harness { pool: PgPool, coordination: CoordinationStore, } impl Harness { async fn new(name: &str) -> Option { 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"); migrate_schema(&pool).await.expect("schema migrate"); let coordination = CoordinationStore::connect( &coordination_url, format!( "{}attempt-mutations:{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 }, ) .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_in("users/alice", 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, staging_origin_attempt_id, target_prefix ) VALUES ($1, 'sealed', 1, $2, 'primary', 'remora-results', $3, 'version-1', $4, 128, $5, $6, $1, 'users/alice')", ) .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, staging_multipart: None, }, ) .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, staging_multipart: None, }; 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, staging_multipart: None, }; 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(); sqlx::query( "INSERT INTO result_inventories ( attempt_id, staging_origin_attempt_id, target_prefix ) VALUES ($1, $1, 'users/alice')", ) .bind(first.attempt_id.as_str()) .execute(&harness.pool) .await .unwrap(); sqlx::query( "INSERT INTO result_inventory_members ( attempt_id, member_id, path, size, object_store, object_bucket, object_key, object_sha256 ) VALUES ($1, $2, 'retry.bin', 512, 'target', 'bucket', $3, $4)", ) .bind(first.attempt_id.as_str()) .bind(unique_id("member").as_str()) .bind(format!( "users/alice/_remora/a/{}/o/retry/retry.bin", first.attempt_id.as_str() )) .bind(format!("sha256:{}", "0".repeat(64))) .execute(&harness.pool) .await .unwrap(); sqlx::query( "UPDATE result_inventories SET member_count = 1, total_bytes = 512 WHERE attempt_id = $1", ) .bind(first.attempt_id.as_str()) .execute(&harness.pool) .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.into(), 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(); 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); let staging_origin: String = sqlx::query_scalar("SELECT staging_origin_attempt_id FROM attempts WHERE attempt_id = $1") .bind(second.attempt_id.as_str()) .fetch_one(&harness.pool) .await .unwrap(); assert_eq!(staging_origin, first.attempt_id.as_str()); let inventory_owner: String = sqlx::query_scalar("SELECT attempt_id FROM result_inventories WHERE attempt_id = $1") .bind(second.attempt_id.as_str()) .fetch_one(&harness.pool) .await .unwrap(); assert_eq!(inventory_owner, second.attempt_id.as_str()); 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.into(), 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, 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, 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_in("users/alice", 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 { 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() }