remora/crates/remora-worker/tests/runtime.rs

use std::{
    collections::BTreeMap,
    future::Future,
    net::TcpListener,
    pin::Pin,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::Duration,
};

use actix_web::{App, HttpRequest, HttpResponse, HttpServer, http::header, web};
use futures_util::stream;
use rcgen::{
    BasicConstraints, Certificate, CertificateParams, ExtendedKeyUsagePurpose, IsCa, Issuer,
    KeyPair, KeyUsagePurpose,
};
use remora_admission::{
    AdmissionConfig, AdmissionDecision, AdmissionService, migrate as admission_migrate,
};
use remora_api::{WorkerApiState, authenticate_rustls_worker, configure_worker_api};
use remora_attempt::{
    ActivationOutcome, AttemptConfig, AttemptService, CheckpointRecord, DurableDescriptor,
    RecoveryConfig, StagingObject, migrate as attempt_migrate,
};
use remora_contracts::{
    CheckpointDescriptorRefDto, HttpSourceSubmissionDto, HuggingFaceRepositoryTypeDto,
    HuggingFaceSourceSubmissionDto, JobSubmissionDto, ObjectIdentityDto,
    ObjectStoreSourceSubmissionDto, PositiveIntegerDto, ProgressValueDto,
    ResolvedSourceDescriptorRefDto, ResultDescriptorDto, ResultKindDto, SafeIntegerDto,
    SourceProjectionDto, SourceSubmissionDto,
};
use remora_coordination::{CoordinationStore, LeaseOwner, StoreConfig};
use remora_credentials::{
    CreateCredential, CredentialRegistry, GrantBoundResolver, SecretMaterial, VaultClient,
    VaultConfig, migrate as credentials_migrate,
};
use remora_domain::{BucketName, JobPhase, OpaqueId, StorageScope, TargetPrefix, TargetStoreId};
use remora_durable::{PoolConfig, connect, migrate as durable_migrate};
use remora_identity::{DefaultTarget, Principal};
use remora_outbound::OutboundPolicyBuilder;
use remora_result::{ResultInventory, ResultMember, migrate as result_migrate};
use remora_stream::{CancelToken, CopyConfig, RetryTiming};
use remora_target::{CreateObjectRequest, TargetStore, TargetStoreConfig, attempt_manifest_key_in};
use remora_worker::{
    AdapterExecutor, AttemptContext, AttemptExecution, AttemptExecutor, DescriptorStore,
    ObjectSourceEndpoint, WorkerApiClient, WorkerConfig, WorkerRunOutcome, WorkerRuntime,
    migrate as worker_migrate,
};
use reqwest::{Certificate as ClientCertificate, Client, Identity};
use rustls::{
    RootCertStore, ServerConfig,
    pki_types::{PrivateKeyDer, PrivatePkcs8KeyDer},
    server::WebPkiClientVerifier,
};
use sha2::{Digest, Sha256};
use sqlx::PgPool;

struct Harness {
    pool: PgPool,
    coordination: CoordinationStore,
    attempts: AttemptService,
    api: WorkerApiClient,
    owner: OpaqueId,
    server: actix_web::dev::ServerHandle,
}

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: 24,
                acquire_timeout: Duration::from_secs(10),
                idle_timeout: Some(Duration::from_secs(30)),
            },
        )
        .await
        .unwrap();
        durable_migrate(&pool).await.unwrap();
        credentials_migrate(&pool).await.unwrap();
        admission_migrate(&pool).await.unwrap();
        attempt_migrate(&pool).await.unwrap();
        result_migrate(&pool).await.unwrap();
        worker_migrate(&pool).await.unwrap();
        sqlx::query("TRUNCATE TABLE jobs CASCADE")
            .execute(&pool)
            .await
            .unwrap();
        let coordination = CoordinationStore::connect(
            &coordination_url,
            format!("{}wp19:{name}:{}:", coordination_namespace(), nonce()),
            StoreConfig {
                global_limit: 16,
                lease_ttl: Duration::from_secs(2),
                reserve_ttl: Duration::from_secs(2),
                progress_ttl: Duration::from_secs(2),
                cancel_ttl: None,
            },
        )
        .await
        .unwrap();
        let attempts = AttemptService::new(
            pool.clone(),
            coordination.clone(),
            AttemptConfig {
                max_attempts: 3,
                local_limit: 16,
            },
        )
        .unwrap();
        let (base_url, server, client) = start_server(WorkerApiState::new(attempts.clone())).await;
        let api = WorkerApiClient::new(base_url, client).unwrap();
        let owner = unique_id("owner");
        Some(Self {
            pool,
            coordination,
            attempts,
            api,
            owner,
            server,
        })
    }

    fn runtime(&self, executor: Arc<dyn AttemptExecutor>) -> WorkerRuntime {
        WorkerRuntime::new(
            self.coordination.clone(),
            self.api.clone(),
            executor,
            WorkerConfig {
                local_limit: 2,
                idle_delay: Duration::from_millis(10),
                renew_interval: Duration::from_millis(200),
                cancel_poll_interval: Duration::from_millis(20),
            },
        )
        .unwrap()
    }

    async fn submit(&self, source: SourceSubmissionDto) -> OpaqueId {
        let store = TargetStoreId::parse(
            std::env::var("REMORA_TARGET_PROVIDER")
                .ok()
                .filter(|value| !value.is_empty() && value != "-")
                .as_deref()
                .unwrap_or("primary"),
        )
        .unwrap();
        let bucket = BucketName::parse(
            std::env::var("REMORA_TARGET_BUCKET")
                .ok()
                .filter(|value| !value.is_empty())
                .as_deref()
                .unwrap_or("remora-results"),
        )
        .unwrap();
        let prefix = TargetPrefix::parse("users/alice", ["_remora"]).unwrap();
        let principal = Principal {
            id: self.owner.clone(),
            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 service = AdmissionService::new(
            self.pool.clone(),
            AdmissionConfig {
                default_target: DefaultTarget {
                    store,
                    bucket,
                    prefix,
                },
                reserved_target_segments: vec!["_remora".into()],
                policy_version: "worker-runtime-v1".into(),
            },
        )
        .unwrap();
        let outcome = service
            .submit(
                &principal,
                unique_id("submit").as_str(),
                JobSubmissionDto {
                    source,
                    target: None,
                    label: Some("worker-runtime-test".into()),
                },
            )
            .await
            .unwrap();
        let job = match outcome.decision {
            AdmissionDecision::Accepted(job) => job,
            AdmissionDecision::Rejected { reason } => panic!("job rejected: {reason:?}"),
        };
        self.coordination
            .enqueue(job.job_id.as_str())
            .await
            .unwrap();
        job.job_id
    }
}

impl Drop for Harness {
    fn drop(&mut self) {
        let server = self.server.clone();
        tokio::spawn(async move { server.stop(true).await });
    }
}

#[tokio::test]
async fn all_source_kinds_run_through_activation_progress_checkpoint_and_outcome() {
    let Some(harness) = Harness::new("sources").await else {
        return;
    };
    let executor = Arc::new(ProtocolExecutor::success(harness.pool.clone()));
    let runtime = harness.runtime(executor.clone());
    let first_incarnation = runtime.incarnation().clone();
    let restarted = harness.runtime(executor);
    assert_ne!(first_incarnation, *restarted.incarnation());

    let sources = [
        SourceSubmissionDto::Http(HttpSourceSubmissionDto {
            url: "https://example.invalid/artifact?private=value".into(),
            result_path: "http.bin".into(),
            credential_id: None,
        }),
        SourceSubmissionDto::ObjectStore(ObjectStoreSourceSubmissionDto {
            store: "source-store".into(),
            bucket: "source-bucket".into(),
            key: "objects/item.bin".into(),
            version_id: Some("version-1".into()),
            result_path: "object.bin".into(),
            credential_id: None,
        }),
        SourceSubmissionDto::HuggingFace(HuggingFaceSourceSubmissionDto {
            repository_type: HuggingFaceRepositoryTypeDto::Model,
            repository_id: "owner/model".into(),
            revision: "0123456789abcdef0123456789abcdef01234567".into(),
            selection: None,
            credential_id: None,
        }),
    ];
    for source in sources {
        let job_id = harness.submit(source).await;
        assert_eq!(
            runtime.run_once().await.unwrap(),
            WorkerRunOutcome::Succeeded
        );
        let phase: String = sqlx::query_scalar("SELECT phase FROM jobs WHERE job_id = $1")
            .bind(job_id.as_str())
            .fetch_one(&harness.pool)
            .await
            .unwrap();
        assert_eq!(phase, "succeeded");
        assert!(
            harness
                .coordination
                .get_live_lease(job_id.as_str())
                .await
                .unwrap()
                .is_none()
        );
    }
}

#[tokio::test]
async fn production_adapters_complete_one_job_per_source_kind() {
    let Some(harness) = Harness::new("production-sources").await else {
        return;
    };
    let fixture = SourceFixture::start().await;
    let target_config = TargetStoreConfig::from_env().unwrap();
    let target = TargetStore::connect(target_config.clone()).unwrap();
    let source_key = CreateObjectRequest::for_attempt_object(
        &unique_id("sourceattempt"),
        &unique_id("sourceobject"),
        Some("source.bin"),
    );
    let source_object = target
        .put_create_only(&source_key, SOURCE_BODY)
        .await
        .unwrap();
    let (credential_id, resolver) = object_credential(&harness, false).await;
    let executor = production_executor(&harness, &fixture, resolver, target_config, "source-store");
    let runtime = harness.runtime(executor);
    let sources = [
        SourceSubmissionDto::Http(HttpSourceSubmissionDto {
            url: format!("{}/http.bin", fixture.base_url),
            result_path: "http.bin".into(),
            credential_id: None,
        }),
        SourceSubmissionDto::ObjectStore(ObjectStoreSourceSubmissionDto {
            store: "source-store".into(),
            bucket: std::env::var("REMORA_TARGET_BUCKET").unwrap(),
            key: source_object.identity.key,
            version_id: None,
            result_path: "object.bin".into(),
            credential_id: Some(credential_id),
        }),
        SourceSubmissionDto::HuggingFace(HuggingFaceSourceSubmissionDto {
            repository_type: HuggingFaceRepositoryTypeDto::Model,
            repository_id: "owner/model".into(),
            revision: HF_COMMIT.into(),
            selection: None,
            credential_id: None,
        }),
    ];
    for source in sources {
        let kind = match &source {
            SourceSubmissionDto::Http(_) => "http",
            SourceSubmissionDto::ObjectStore(_) => "object_store",
            SourceSubmissionDto::HuggingFace(_) => "hugging_face",
        };
        let job_id = harness.submit(source).await;
        let outcome = runtime.run_once().await.unwrap();
        if outcome != WorkerRunOutcome::Succeeded {
            let failure: Option<serde_json::Value> = sqlx::query_scalar(
                "SELECT failure FROM attempts
                 WHERE job_id = $1 ORDER BY fence DESC LIMIT 1",
            )
            .bind(job_id.as_str())
            .fetch_one(&harness.pool)
            .await
            .unwrap();
            let inventory: Option<(String, i64, i64)> = sqlx::query_as(
                "SELECT i.state, i.member_count, i.total_bytes
                 FROM result_inventories i
                 JOIN attempts a ON a.attempt_id = i.attempt_id
                 WHERE a.job_id = $1 ORDER BY a.fence DESC LIMIT 1",
            )
            .bind(job_id.as_str())
            .fetch_optional(&harness.pool)
            .await
            .unwrap();
            panic!("{kind} failed with {outcome:?}: {failure:?}, inventory={inventory:?}");
        }
    }
}

#[tokio::test]
async fn production_adapter_resumes_a_transferred_checkpoint_after_restart() {
    let Some(harness) = Harness::new("production-resume").await else {
        return;
    };
    let fixture = SourceFixture::start().await;
    let target_config = TargetStoreConfig::from_env().unwrap();
    let target = TargetStore::connect(target_config.clone()).unwrap();
    let source_url = format!("{}/http.bin", fixture.base_url);
    let job_id = harness
        .submit(SourceSubmissionDto::Http(HttpSourceSubmissionDto {
            url: source_url.clone(),
            result_path: "http.bin".into(),
            credential_id: None,
        }))
        .await;
    assert_eq!(
        harness.coordination.dequeue().await.unwrap().as_deref(),
        Some(job_id.as_str())
    );
    let first_worker = unique_id("worker");
    let first = match harness
        .attempts
        .claim(&job_id, &first_worker, &unique_id("claimop"))
        .await
        .unwrap()
    {
        remora_attempt::ClaimOutcome::Claimed { authority, .. } => authority,
        remora_attempt::ClaimOutcome::Rejected { reason, .. } => {
            panic!("initial claim rejected: {reason:?}")
        }
    };
    assert!(matches!(
        harness
            .attempts
            .activate(&first, &unique_id("activateop"))
            .await
            .unwrap(),
        ActivationOutcome::Activated(_)
    ));
    let first = harness
        .attempts
        .authority(&first.attempt_id)
        .await
        .unwrap()
        .unwrap();
    let descriptors = DescriptorStore::new(harness.pool.clone());
    let identity = serde_json::json!({
        "kind": "http",
        "requested_url_sha256": digest_text(&source_url),
        "final_url_sha256": digest_text(&source_url),
        "etag_sha256": digest_text("\"source-v1\""),
        "representation_length": SOURCE_BODY.len(),
        "content_encoding": "identity"
    });
    let resolved_ref = descriptors
        .store_resolved(&first.attempt_id, &identity, Some(1))
        .await
        .unwrap();
    harness
        .attempts
        .pin_resolved_source(
            &first,
            &unique_id("resolveop"),
            resolved_descriptor(&resolved_ref),
        )
        .await
        .unwrap();
    let object_request = CreateObjectRequest::for_target_attempt_object(
        "users/alice",
        &first.attempt_id,
        &unique_id("member"),
        Some("http.bin"),
    );
    let stored = target
        .put_create_only(&object_request, SOURCE_BODY)
        .await
        .unwrap();
    let inventory = ResultInventory::new(harness.pool.clone());
    inventory
        .create_bound(&first.attempt_id, &first.attempt_id, "users/alice")
        .await
        .unwrap();
    inventory
        .add_member(
            &first.attempt_id,
            &ResultMember::new(
                unique_id("member"),
                "http.bin",
                stored.size,
                stored.identity.clone(),
                stored.content_sha256.clone(),
            )
            .unwrap(),
        )
        .await
        .unwrap();
    harness
        .attempts
        .register_staging_object(
            &first.attempt_id,
            &StagingObject {
                target_store: stored.identity.target_store_id.clone(),
                bucket: stored.identity.bucket.clone(),
                key: stored.identity.key.clone(),
                version_id: stored.identity.version_id.clone(),
            },
        )
        .await
        .unwrap();
    let checkpoint_ref = descriptors
        .store_checkpoint(
            &first.attempt_id,
            &serde_json::json!({
                "staging_origin_attempt_id": first.attempt_id,
                "resolved_source_id": resolved_ref.id,
                "validated_bytes": SOURCE_BODY.len(),
                "validated_members": 1
            }),
            Some(1),
        )
        .await
        .unwrap();
    harness
        .attempts
        .persist_checkpoint(
            &first,
            &unique_id("checkpointop"),
            CheckpointRecord {
                descriptor: checkpoint_descriptor(&checkpoint_ref),
                resolved_source_id: resolved_ref.id.clone(),
                validated_bytes: SOURCE_BODY.len() as u64,
                validated_members: 1,
                staging_multipart: None,
            },
        )
        .await
        .unwrap();
    harness
        .attempts
        .publish_failure(
            &first,
            &unique_id("failureop"),
            remora_domain::RuntimeFailureClass::TransientTarget,
        )
        .await
        .unwrap();
    harness
        .coordination
        .release(
            job_id.as_str(),
            &LeaseOwner::new(
                first.attempt_id.as_str(),
                first.fence,
                first.worker_incarnation.as_str(),
            )
            .unwrap(),
        )
        .await
        .unwrap();
    harness.attempts.release_local(&job_id);
    let recovery = RecoveryConfig {
        retry_delay: Duration::ZERO,
        ..RecoveryConfig::default()
    };
    let report = harness.attempts.reconcile_once(&recovery).await.unwrap();
    assert_eq!(report.retries_scheduled, 1);
    assert_eq!(report.enqueued, 1);

    let executor = production_executor(
        &harness,
        &fixture,
        object_credentialless_resolver(&harness).await,
        target_config,
        "source-store",
    );
    let restarted = harness.runtime(executor);
    assert_ne!(first.worker_incarnation, *restarted.incarnation());
    let outcome = restarted.run_once().await;
    if let Err(error) = &outcome {
        let job: (String, Option<String>, Option<String>) = sqlx::query_as(
            "SELECT phase, authority_attempt_id, checkpoint_owner_attempt_id
             FROM jobs WHERE job_id = $1",
        )
        .bind(job_id.as_str())
        .fetch_one(&harness.pool)
        .await
        .unwrap();
        let attempts: Vec<(i64, String, String)> = sqlx::query_as(
            "SELECT fence, attempt_id, state FROM attempts
             WHERE job_id = $1 ORDER BY fence",
        )
        .bind(job_id.as_str())
        .fetch_all(&harness.pool)
        .await
        .unwrap();
        panic!("resume failed with {error:?}; job={job:?}; attempts={attempts:?}");
    }
    assert_eq!(outcome.unwrap(), WorkerRunOutcome::Succeeded);
    let attempts: Vec<(String, String)> = sqlx::query_as(
        "SELECT attempt_id, staging_origin_attempt_id
         FROM attempts WHERE job_id = $1 ORDER BY fence",
    )
    .bind(job_id.as_str())
    .fetch_all(&harness.pool)
    .await
    .unwrap();
    assert_eq!(attempts.len(), 2);
    assert_eq!(attempts[1].1, attempts[0].0);
    let member: (String, String) = sqlx::query_as(
        "SELECT attempt_id, object_key
         FROM result_inventory_members WHERE path = 'http.bin'",
    )
    .fetch_one(&harness.pool)
    .await
    .unwrap();
    assert_eq!(member.0, attempts[1].0);
    assert!(
        member
            .1
            .starts_with(&format!("users/alice/_remora/a/{}/o/", attempts[0].0))
    );
    let cleanup_owner: String =
        sqlx::query_scalar("SELECT attempt_id FROM attempt_staging_cleanup WHERE object_key = $1")
            .bind(&member.1)
            .fetch_one(&harness.pool)
            .await
            .unwrap();
    assert_eq!(cleanup_owner, attempts[1].0);
}

#[tokio::test]
async fn production_http_adapter_resumes_a_partial_multipart_after_process_loss() {
    let Some(harness) = Harness::new("production-partial-resume").await else {
        return;
    };
    let fixture = SourceFixture::start().await;
    let target_config = TargetStoreConfig::from_env().unwrap();
    let executor = production_executor(
        &harness,
        &fixture,
        object_credentialless_resolver(&harness).await,
        target_config,
        "source-store",
    );
    let first = harness.runtime(executor.clone());
    let first_incarnation = first.incarnation().clone();
    let job_id = harness
        .submit(SourceSubmissionDto::Http(HttpSourceSubmissionDto {
            url: format!("{}/large.bin", fixture.base_url),
            result_path: "large.bin".into(),
            credential_id: None,
        }))
        .await;
    let task = tokio::spawn(async move { first.run_once().await });
    let partial = loop {
        let checkpoint: Option<serde_json::Value> =
            sqlx::query_scalar("SELECT checkpoint FROM jobs WHERE job_id = $1")
                .bind(job_id.as_str())
                .fetch_one(&harness.pool)
                .await
                .unwrap();
        if let Some(checkpoint) = checkpoint.filter(|value| {
            value
                .pointer("/checkpoint/validated_members")
                .and_then(|value| value.as_u64())
                == Some(0)
        }) {
            break checkpoint;
        }
        assert!(
            !task.is_finished(),
            "transfer finished before a partial checkpoint"
        );
        tokio::time::sleep(Duration::from_millis(10)).await;
    };
    assert_eq!(
        partial
            .pointer("/checkpoint/validated_bytes")
            .and_then(|value| value.as_u64()),
        Some(MULTIPART_BOUNDARY as u64)
    );
    harness
        .coordination
        .drop_job_state(&[job_id.as_str()])
        .await
        .unwrap();
    assert_eq!(
        task.await.unwrap().unwrap(),
        WorkerRunOutcome::AuthorityLost
    );
    let recovery = RecoveryConfig {
        retry_delay: Duration::ZERO,
        ..RecoveryConfig::default()
    };
    assert_eq!(
        harness
            .attempts
            .reconcile_once(&recovery)
            .await
            .unwrap()
            .attempts_reaped,
        1
    );
    let report = harness.attempts.reconcile_once(&recovery).await.unwrap();
    assert_eq!(report.retries_scheduled, 1);
    assert_eq!(report.enqueued, 1);

    let restarted = harness.runtime(executor);
    assert_ne!(first_incarnation, *restarted.incarnation());
    assert_eq!(
        restarted.run_once().await.unwrap(),
        WorkerRunOutcome::Succeeded
    );
    assert_eq!(fixture.range_requests.load(Ordering::Acquire), 1);
    let attempts: Vec<(String, String)> = sqlx::query_as(
        "SELECT attempt_id, staging_origin_attempt_id
         FROM attempts WHERE job_id = $1 ORDER BY fence",
    )
    .bind(job_id.as_str())
    .fetch_all(&harness.pool)
    .await
    .unwrap();
    assert_eq!(attempts.len(), 2);
    assert_eq!(attempts[1].1, attempts[0].0);
    let member: (String, String, i64) = sqlx::query_as(
        "SELECT attempt_id, object_key, size
         FROM result_inventory_members WHERE path = 'large.bin'",
    )
    .fetch_one(&harness.pool)
    .await
    .unwrap();
    assert_eq!(member.0, attempts[1].0);
    assert_eq!(member.2, LARGE_SOURCE_SIZE as i64);
    assert!(
        member
            .1
            .starts_with(&format!("users/alice/_remora/a/{}/o/", attempts[0].0))
    );
    let multipart_owner: String = sqlx::query_scalar(
        "SELECT attempt_id FROM attempt_staging_cleanup
         WHERE kind = 'multipart' ORDER BY cleanup_id DESC LIMIT 1",
    )
    .fetch_one(&harness.pool)
    .await
    .unwrap();
    assert_eq!(multipart_owner, attempts[1].0);
}

#[tokio::test]
async fn lease_loss_cancels_io_and_stale_authority_cannot_report() {
    let Some(harness) = Harness::new("lease-loss").await else {
        return;
    };
    let executor = Arc::new(ProtocolExecutor::blocking(harness.pool.clone()));
    let runtime = harness.runtime(executor.clone());
    let job_id = harness
        .submit(SourceSubmissionDto::Http(HttpSourceSubmissionDto {
            url: "https://example.invalid/slow".into(),
            result_path: "slow.bin".into(),
            credential_id: None,
        }))
        .await;
    let task = tokio::spawn(async move { runtime.run_once().await.unwrap() });
    executor.started.notified().await;
    harness
        .coordination
        .drop_job_state(&[job_id.as_str()])
        .await
        .unwrap();
    assert_eq!(task.await.unwrap(), WorkerRunOutcome::AuthorityLost);
    assert!(executor.stale_report_failed.load(Ordering::Acquire) > 0);
}

#[tokio::test]
async fn cancellation_stops_the_attempt_and_credential_outage_is_retryable() {
    let Some(harness) = Harness::new("cancel-outage").await else {
        return;
    };
    let blocking = Arc::new(ProtocolExecutor::blocking(harness.pool.clone()));
    let runtime = harness.runtime(blocking.clone());
    let job_id = harness
        .submit(SourceSubmissionDto::Http(HttpSourceSubmissionDto {
            url: "https://example.invalid/cancel".into(),
            result_path: "cancel.bin".into(),
            credential_id: None,
        }))
        .await;
    let task = tokio::spawn(async move { runtime.run_once().await.unwrap() });
    blocking.started.notified().await;
    harness.attempts.request_cancel(&job_id).await.unwrap();
    harness
        .coordination
        .publish_cancel(job_id.as_str())
        .await
        .unwrap();
    assert_eq!(task.await.unwrap(), WorkerRunOutcome::Cancelled);

    let fixture = SourceFixture::start().await;
    let target = TargetStoreConfig::from_env().unwrap();
    let (credential_id, resolver) = object_credential(&harness, true).await;
    let runtime = harness.runtime(production_executor(
        &harness,
        &fixture,
        resolver,
        target,
        "source-store",
    ));
    let job_id = harness
        .submit(SourceSubmissionDto::ObjectStore(
            ObjectStoreSourceSubmissionDto {
                store: "source-store".into(),
                bucket: std::env::var("REMORA_TARGET_BUCKET").unwrap(),
                key: "credential-outage.bin".into(),
                version_id: None,
                result_path: "credential.bin".into(),
                credential_id: Some(credential_id),
            },
        ))
        .await;
    assert_eq!(
        runtime.run_once().await.unwrap(),
        WorkerRunOutcome::Failed(JobPhase::RetryWait)
    );
    let phase: String = sqlx::query_scalar("SELECT phase FROM jobs WHERE job_id = $1")
        .bind(job_id.as_str())
        .fetch_one(&harness.pool)
        .await
        .unwrap();
    assert_eq!(phase, "retry_wait");
}

#[tokio::test]
async fn crashed_process_is_reaped_and_restarted_with_fresh_authority() {
    let Some(harness) = Harness::new("restart").await else {
        return;
    };
    let crashed = Arc::new(ProtocolExecutor::crash_blocking(harness.pool.clone()));
    let first = harness.runtime(crashed.clone());
    let first_incarnation = first.incarnation().clone();
    let job_id = harness
        .submit(SourceSubmissionDto::Http(HttpSourceSubmissionDto {
            url: "https://example.invalid/restart".into(),
            result_path: "restart.bin".into(),
            credential_id: None,
        }))
        .await;
    let task = tokio::spawn(async move { first.run_once().await });
    crashed.started.notified().await;
    task.abort();
    assert!(task.await.unwrap_err().is_cancelled());
    harness
        .coordination
        .drop_job_state(&[job_id.as_str()])
        .await
        .unwrap();
    let recovery = RecoveryConfig {
        retry_delay: Duration::ZERO,
        ..RecoveryConfig::default()
    };
    assert_eq!(
        harness
            .attempts
            .reconcile_once(&recovery)
            .await
            .unwrap()
            .attempts_reaped,
        1
    );
    assert_eq!(
        harness
            .attempts
            .reconcile_once(&recovery)
            .await
            .unwrap()
            .retries_scheduled,
        1
    );

    let restarted = harness.runtime(Arc::new(ProtocolExecutor::success(harness.pool.clone())));
    assert_ne!(first_incarnation, *restarted.incarnation());
    assert_eq!(
        restarted.run_once().await.unwrap(),
        WorkerRunOutcome::Succeeded
    );
    let attempts: Vec<(i64, String)> = sqlx::query_as(
        "SELECT fence, worker_incarnation FROM attempts WHERE job_id = $1 ORDER BY fence",
    )
    .bind(job_id.as_str())
    .fetch_all(&harness.pool)
    .await
    .unwrap();
    assert_eq!(attempts.len(), 2);
    assert_eq!(attempts[0].0, 1);
    assert_eq!(attempts[1].0, 2);
    assert_ne!(attempts[0].1, attempts[1].1);
}

#[tokio::test]
async fn local_limit_bounds_concurrent_attempt_tasks() {
    let Some(harness) = Harness::new("local-limit").await else {
        return;
    };
    let executor = Arc::new(ConcurrencyExecutor::default());
    for suffix in ["one", "two", "three"] {
        harness
            .submit(SourceSubmissionDto::Http(HttpSourceSubmissionDto {
                url: format!("https://example.invalid/{suffix}"),
                result_path: format!("{suffix}.bin"),
                credential_id: None,
            }))
            .await;
    }
    let shutdown = CancelToken::new();
    let runtime = harness.runtime(executor.clone());
    let task = tokio::spawn({
        let shutdown = shutdown.clone();
        async move { runtime.run(shutdown).await }
    });
    tokio::time::timeout(Duration::from_secs(5), async {
        while executor.active.load(Ordering::Acquire) < 2 {
            tokio::task::yield_now().await;
        }
    })
    .await
    .unwrap();
    tokio::time::sleep(Duration::from_millis(100)).await;
    assert_eq!(executor.active.load(Ordering::Acquire), 2);
    let executing: i64 =
        sqlx::query_scalar("SELECT count(*) FROM attempts WHERE state = 'executing'")
            .fetch_one(&harness.pool)
            .await
            .unwrap();
    assert_eq!(executing, 2);
    shutdown.cancel();
    task.await.unwrap().unwrap();
}

#[derive(Clone, Copy)]
enum ExecutorMode {
    Success,
    Blocking,
    CrashBlocking,
}

#[derive(Clone)]
struct ProtocolExecutor {
    pool: PgPool,
    mode: ExecutorMode,
    started: Arc<tokio::sync::Notify>,
    stale_report_failed: Arc<AtomicUsize>,
}

impl ProtocolExecutor {
    fn success(pool: PgPool) -> Self {
        Self::new(pool, ExecutorMode::Success)
    }

    fn blocking(pool: PgPool) -> Self {
        Self::new(pool, ExecutorMode::Blocking)
    }

    fn crash_blocking(pool: PgPool) -> Self {
        Self::new(pool, ExecutorMode::CrashBlocking)
    }

    fn new(pool: PgPool, mode: ExecutorMode) -> Self {
        Self {
            pool,
            mode,
            started: Arc::new(tokio::sync::Notify::new()),
            stale_report_failed: Arc::new(AtomicUsize::new(0)),
        }
    }

    async fn run(
        &self,
        context: AttemptContext,
        envelope: remora_contracts::ExecutionEnvelopeDto,
    ) -> AttemptExecution {
        if matches!(self.mode, ExecutorMode::CrashBlocking) {
            self.started.notify_one();
            std::future::pending().await
        }
        let descriptors = DescriptorStore::new(self.pool.clone());
        let resolved = descriptors
            .store_resolved(
                &envelope.attempt_id,
                &serde_json::json!({ "kind": source_kind(&envelope.source) }),
                Some(1),
            )
            .await
            .unwrap();
        context.pin_resolved_source(resolved.clone()).await.unwrap();
        let checkpoint = descriptors
            .store_checkpoint(
                &envelope.attempt_id,
                &serde_json::json!({
                    "resolved_source_id": resolved.id,
                    "validated_bytes": 1,
                    "validated_members": 1
                }),
                Some(0),
            )
            .await
            .unwrap();
        context
            .persist_checkpoint(checkpoint, resolved.id, safe(1), safe(1), None)
            .await
            .unwrap();
        context
            .report_progress(ProgressValueDto {
                completed_bytes: safe(64),
                total_bytes: Some(safe(64)),
                completed_members: safe(1),
                total_members: Some(safe(1)),
            })
            .await
            .unwrap();
        if matches!(self.mode, ExecutorMode::Blocking) {
            self.started.notify_one();
            context.cancellation().cancelled().await;
            if context
                .report_progress(ProgressValueDto {
                    completed_bytes: safe(65),
                    total_bytes: Some(safe(65)),
                    completed_members: safe(1),
                    total_members: Some(safe(1)),
                })
                .await
                .is_err()
            {
                self.stale_report_failed.fetch_add(1, Ordering::Release);
            }
            return AttemptExecution::Cancelled;
        }
        AttemptExecution::Succeeded(seal_result(&self.pool, &envelope).await)
    }
}

#[derive(Default)]
struct ConcurrencyExecutor {
    active: AtomicUsize,
}

impl AttemptExecutor for ConcurrencyExecutor {
    fn execute(
        &self,
        _context: AttemptContext,
        _envelope: remora_contracts::ExecutionEnvelopeDto,
    ) -> Pin<Box<dyn Future<Output = AttemptExecution> + Send>> {
        self.active.fetch_add(1, Ordering::AcqRel);
        Box::pin(std::future::pending())
    }
}

impl AttemptExecutor for ProtocolExecutor {
    fn execute(
        &self,
        context: AttemptContext,
        envelope: remora_contracts::ExecutionEnvelopeDto,
    ) -> Pin<Box<dyn Future<Output = AttemptExecution> + Send>> {
        let executor = self.clone();
        Box::pin(async move { executor.run(context, envelope).await })
    }
}

async fn seal_result(
    pool: &PgPool,
    envelope: &remora_contracts::ExecutionEnvelopeDto,
) -> ResultDescriptorDto {
    let manifest_id = unique_id("manifest");
    let checksum = format!("sha256:{}", "11".repeat(32));
    let key = attempt_manifest_key_in(&envelope.target.prefix, &envelope.attempt_id, &manifest_id);
    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, $3, $4, $5,
                   'version-1', $6, 128, 1, 64, $1, $7)",
    )
    .bind(envelope.attempt_id.as_str())
    .bind(manifest_id.as_str())
    .bind(&envelope.target.store)
    .bind(&envelope.target.bucket)
    .bind(&key)
    .bind(&checksum)
    .bind(&envelope.target.prefix)
    .execute(pool)
    .await
    .unwrap();
    ResultDescriptorDto {
        kind: ResultKindDto::Manifest,
        manifest_id,
        schema_version: positive(1),
        manifest: ObjectIdentityDto {
            store: envelope.target.store.clone(),
            bucket: envelope.target.bucket.clone(),
            key,
            version_id: Some("version-1".into()),
            checksum,
            size: safe(128),
        },
        member_count: positive(1),
        total_bytes: safe(64),
    }
}

fn source_kind(source: &SourceProjectionDto) -> &'static str {
    match source {
        SourceProjectionDto::Http(_) => "http",
        SourceProjectionDto::ObjectStore(_) => "object_store",
        SourceProjectionDto::HuggingFace(_) => "hugging_face",
    }
}

const SOURCE_BODY: &[u8] = b"source-body";
const MULTIPART_BOUNDARY: usize = 5 * 1024 * 1024;
const LARGE_SOURCE_SIZE: usize = MULTIPART_BOUNDARY * 2 + 123;
const HF_BODY: &[u8] = b"hf-body";
const HF_COMMIT: &str = "0123456789abcdef0123456789abcdef01234567";

struct SourceFixture {
    base_url: String,
    port: u16,
    root_der: Vec<u8>,
    server: actix_web::dev::ServerHandle,
    range_requests: Arc<AtomicUsize>,
}

impl SourceFixture {
    async fn start() -> Self {
        let mut ca_params = CertificateParams::new(Vec::new()).unwrap();
        ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
        ca_params.key_usages = vec![
            KeyUsagePurpose::DigitalSignature,
            KeyUsagePurpose::KeyCertSign,
            KeyUsagePurpose::CrlSign,
        ];
        let ca_key = KeyPair::generate().unwrap();
        let ca = ca_params.self_signed(&ca_key).unwrap();
        let issuer = Issuer::new(ca_params, ca_key);
        let server_key = KeyPair::generate().unwrap();
        let mut server_params = CertificateParams::new(vec!["127.0.0.1".into()]).unwrap();
        server_params.key_usages = vec![KeyUsagePurpose::DigitalSignature];
        server_params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth];
        let server = server_params.signed_by(&server_key, &issuer).unwrap();
        let root_der = ca.der().as_ref().to_vec();
        let key = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(server_key.serialize_der()));
        let server_config = ServerConfig::builder()
            .with_no_client_auth()
            .with_single_cert(vec![server.der().clone()], key)
            .unwrap();
        let listener = TcpListener::bind("127.0.0.1:443").unwrap();
        let address = listener.local_addr().unwrap();
        let range_requests = Arc::new(AtomicUsize::new(0));
        let fixture = web::Data::new(SourceFixtureData {
            large_body: Arc::new(vec![0x5a; LARGE_SOURCE_SIZE]),
            range_requests: range_requests.clone(),
        });
        let server = HttpServer::new(move || {
            App::new()
                .app_data(fixture.clone())
                .default_service(web::to(source_response))
        })
        .listen_rustls_0_23(listener, server_config)
        .unwrap()
        .run();
        let handle = server.handle();
        tokio::spawn(server);
        Self {
            base_url: "https://127.0.0.1".into(),
            port: address.port(),
            root_der,
            server: handle,
            range_requests,
        }
    }
}

impl Drop for SourceFixture {
    fn drop(&mut self) {
        let server = self.server.clone();
        tokio::spawn(async move { server.stop(true).await });
    }
}

struct SourceFixtureData {
    large_body: Arc<Vec<u8>>,
    range_requests: Arc<AtomicUsize>,
}

async fn source_response(
    request: HttpRequest,
    fixture: web::Data<SourceFixtureData>,
) -> HttpResponse {
    let path = request.path();
    if path == "/http.bin" {
        return HttpResponse::Ok()
            .insert_header((header::ETAG, "\"source-v1\""))
            .body(SOURCE_BODY);
    }
    if path == "/large.bin" {
        if let Some(range) = request
            .headers()
            .get(header::RANGE)
            .and_then(|value| value.to_str().ok())
        {
            let offset = range
                .strip_prefix("bytes=")
                .and_then(|value| value.strip_suffix('-'))
                .and_then(|value| value.parse::<usize>().ok())
                .filter(|offset| *offset < fixture.large_body.len());
            let valid_if_range = request
                .headers()
                .get(header::IF_MATCH)
                .and_then(|value| value.to_str().ok())
                == Some("\"large-v1\"");
            let Some(offset) = offset.filter(|_| valid_if_range) else {
                return HttpResponse::PreconditionFailed().finish();
            };
            fixture.range_requests.fetch_add(1, Ordering::AcqRel);
            return HttpResponse::PartialContent()
                .insert_header((header::ETAG, "\"large-v1\""))
                .insert_header((
                    header::CONTENT_RANGE,
                    format!(
                        "bytes {offset}-{}/{}",
                        fixture.large_body.len() - 1,
                        fixture.large_body.len()
                    ),
                ))
                .body(fixture.large_body[offset..].to_vec());
        }
        let body = fixture.large_body.clone();
        let chunks = stream::unfold(0_usize, move |offset| {
            let body = body.clone();
            async move {
                if offset >= body.len() {
                    return None;
                }
                if offset == MULTIPART_BOUNDARY {
                    tokio::time::sleep(Duration::from_secs(2)).await;
                }
                let end = (offset + 64 * 1024).min(body.len());
                Some((
                    Ok::<_, std::io::Error>(bytes::Bytes::copy_from_slice(&body[offset..end])),
                    end,
                ))
            }
        });
        return HttpResponse::Ok()
            .insert_header((header::ETAG, "\"large-v1\""))
            .insert_header((header::CONTENT_LENGTH, fixture.large_body.len()))
            .streaming(chunks);
    }
    if path.ends_with(&format!("/revision/{HF_COMMIT}")) {
        return HttpResponse::Ok().json(serde_json::json!({
            "id": "owner/model",
            "sha": HF_COMMIT
        }));
    }
    if path.contains(&format!("/tree/{HF_COMMIT}")) {
        return HttpResponse::Ok().json(serde_json::json!([
            {
                "type": "file",
                "oid": "e69de29bb2d1d6434b8b29ae775ad8c2e48c5391",
                "size": 0,
                "path": "empty.bin",
                "lfs": null
            },
            {
                "type": "file",
                "oid": "af0f6ad1a70708b53ec3774c084bee9ad41f1984",
                "size": HF_BODY.len(),
                "path": "model.bin",
                "lfs": null
            }
        ]));
    }
    if path.ends_with(&format!("/resolve/{HF_COMMIT}/empty.bin")) {
        return HttpResponse::Ok().body("");
    }
    if path.ends_with(&format!("/resolve/{HF_COMMIT}/model.bin")) {
        return HttpResponse::Ok().body(HF_BODY);
    }
    HttpResponse::NotFound().finish()
}

async fn object_credential(
    harness: &Harness,
    resolver_outage: bool,
) -> (OpaqueId, GrantBoundResolver) {
    let vault_config = VaultConfig::from_env().unwrap();
    let vault = VaultClient::new(vault_config.clone()).unwrap();
    let registry = CredentialRegistry::new(harness.pool.clone(), vault.clone());
    let credential = registry
        .create(
            &harness.owner,
            CreateCredential {
                label: unique_id("label").to_string(),
                material: SecretMaterial::ObjectStore {
                    access_key_id: std::env::var("REMORA_TARGET_ACCESS_KEY").unwrap(),
                    secret_access_key: std::env::var("REMORA_TARGET_SECRET_KEY").unwrap(),
                    session_token: None,
                },
            },
        )
        .await
        .unwrap();
    let resolver_vault = if resolver_outage {
        VaultClient::new(VaultConfig {
            address: "http://127.0.0.1:9".into(),
            timeout: Duration::from_millis(100),
            ..vault_config
        })
        .unwrap()
    } else {
        vault
    };
    (
        credential.id,
        GrantBoundResolver::new(registry, resolver_vault),
    )
}

async fn object_credentialless_resolver(harness: &Harness) -> GrantBoundResolver {
    let vault = VaultClient::new(VaultConfig::from_env().unwrap()).unwrap();
    GrantBoundResolver::new(
        CredentialRegistry::new(harness.pool.clone(), vault.clone()),
        vault,
    )
}

fn production_executor(
    harness: &Harness,
    fixture: &SourceFixture,
    resolver: GrantBoundResolver,
    target: TargetStoreConfig,
    source_store: &str,
) -> Arc<dyn AttemptExecutor> {
    let object_sources = BTreeMap::from([(
        source_store.to_owned(),
        ObjectSourceEndpoint {
            endpoint: target.endpoint.clone(),
            region: target.region.clone(),
        },
    )]);
    let outbound = OutboundPolicyBuilder::new()
        .allow_scheme_port("https", fixture.port)
        .allow_non_global("127.0.0.0/8".parse().unwrap())
        .allow_non_global("::1/128".parse().unwrap())
        .build();
    let copy = CopyConfig {
        progress_interval_bytes: 1,
        checkpoint_interval_bytes: 1,
        ..CopyConfig::default()
    };
    Arc::new(
        AdapterExecutor::new(
            harness.pool.clone(),
            resolver,
            target,
            object_sources,
            outbound,
            RetryTiming::default_bounds(),
            copy,
            &fixture.base_url,
        )
        .unwrap()
        .with_outbound_root_certificate_der(fixture.root_der.clone()),
    )
}

async fn start_server(state: WorkerApiState) -> (String, actix_web::dev::ServerHandle, Client) {
    let certificates = certificates();
    let client = trusted_client(&certificates);
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let address = listener.local_addr().unwrap();
    let server = HttpServer::new(move || {
        App::new()
            .app_data(web::Data::new(state.clone()))
            .configure(configure_worker_api)
    })
    .on_connect(authenticate_rustls_worker)
    .listen_rustls_0_23(listener, certificates.server_config)
    .unwrap()
    .run();
    let handle = server.handle();
    tokio::spawn(server);
    (
        format!("https://localhost:{}", address.port()),
        handle,
        client,
    )
}

fn trusted_client(certificates: &Certificates) -> Client {
    let root = ClientCertificate::from_der(certificates.ca.der()).unwrap();
    let mut client_pem = certificates.client.pem();
    client_pem.push_str(&certificates.client_key.serialize_pem());
    Client::builder()
        .add_root_certificate(root)
        .identity(Identity::from_pem(client_pem.as_bytes()).unwrap())
        .build()
        .unwrap()
}

struct Certificates {
    ca: Certificate,
    client: Certificate,
    client_key: KeyPair,
    server_config: ServerConfig,
}

fn certificates() -> Certificates {
    let _ = rustls::crypto::ring::default_provider().install_default();
    let mut ca_params = CertificateParams::new(Vec::new()).unwrap();
    ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
    ca_params.key_usages = vec![
        KeyUsagePurpose::DigitalSignature,
        KeyUsagePurpose::KeyCertSign,
        KeyUsagePurpose::CrlSign,
    ];
    let ca_key = KeyPair::generate().unwrap();
    let ca = ca_params.self_signed(&ca_key).unwrap();
    let issuer = Issuer::new(ca_params, ca_key);

    let server_key = KeyPair::generate().unwrap();
    let mut server_params = CertificateParams::new(vec!["localhost".into()]).unwrap();
    server_params.key_usages = vec![KeyUsagePurpose::DigitalSignature];
    server_params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth];
    let server = server_params.signed_by(&server_key, &issuer).unwrap();

    let client_key = KeyPair::generate().unwrap();
    let mut client_params = CertificateParams::new(Vec::new()).unwrap();
    client_params.key_usages = vec![KeyUsagePurpose::DigitalSignature];
    client_params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ClientAuth];
    let client = client_params.signed_by(&client_key, &issuer).unwrap();

    let mut roots = RootCertStore::empty();
    roots.add(ca.der().clone()).unwrap();
    let verifier = WebPkiClientVerifier::builder(Arc::new(roots))
        .build()
        .unwrap();
    let server_key_der = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(server_key.serialize_der()));
    let server_config = ServerConfig::builder()
        .with_client_cert_verifier(verifier)
        .with_single_cert(vec![server.der().clone()], server_key_der)
        .unwrap();
    Certificates {
        ca,
        client,
        client_key,
        server_config,
    }
}

fn required_env(name: &str) -> Option<String> {
    match std::env::var(name) {
        Ok(value) if !value.is_empty() => Some(value),
        _ if std::env::var_os("REMORA_WORKER_REQUIRE").is_some() => {
            panic!("{name} is required")
        }
        _ => None,
    }
}

fn coordination_namespace() -> String {
    std::env::var("REMORA_COORDINATION_NAMESPACE").unwrap_or_else(|_| "remora:test:".into())
}

fn nonce() -> u128 {
    use std::sync::atomic::{AtomicU64, Ordering};
    static COUNTER: AtomicU64 = AtomicU64::new(0);
    (u128::from(std::process::id()) << 64) | u128::from(COUNTER.fetch_add(1, Ordering::Relaxed))
}

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

fn positive(value: u64) -> PositiveIntegerDto {
    PositiveIntegerDto::new(value).unwrap()
}

fn safe(value: u64) -> SafeIntegerDto {
    SafeIntegerDto::new(value).unwrap()
}

fn resolved_descriptor(value: &ResolvedSourceDescriptorRefDto) -> DurableDescriptor {
    DurableDescriptor {
        id: value.id.clone(),
        schema_version: u32::try_from(value.schema_version.get()).unwrap(),
        byte_length: value.byte_length.get(),
        item_count: value.item_count.map(SafeIntegerDto::get),
        checksum: value.checksum.clone(),
    }
}

fn checkpoint_descriptor(value: &CheckpointDescriptorRefDto) -> DurableDescriptor {
    DurableDescriptor {
        id: value.id.clone(),
        schema_version: u32::try_from(value.schema_version.get()).unwrap(),
        byte_length: value.byte_length.get(),
        item_count: value.item_count.map(SafeIntegerDto::get),
        checksum: value.checksum.clone(),
    }
}

fn digest_text(value: &str) -> String {
    Sha256::digest(value.as_bytes())
        .iter()
        .fold(String::from("sha256:"), |mut output, byte| {
            use std::fmt::Write as _;
            write!(output, "{byte:02x}").unwrap();
            output
        })
}