1460 lines
49 KiB
Rust
1460 lines
49 KiB
Rust
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};
|
|
use remora_api::{WorkerApiState, authenticate_rustls_worker, configure_worker_api};
|
|
use remora_attempt::{
|
|
ActivationOutcome, AttemptConfig, AttemptService, CheckpointRecord, DurableDescriptor,
|
|
RecoveryConfig, StagingObject,
|
|
};
|
|
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,
|
|
};
|
|
use remora_domain::{BucketName, JobPhase, OpaqueId, StorageScope, TargetPrefix, TargetStoreId};
|
|
use remora_durable::{PoolConfig, connect};
|
|
use remora_identity::{DefaultTarget, Principal};
|
|
use remora_outbound::OutboundPolicyBuilder;
|
|
use remora_result::{ResultInventory, ResultMember};
|
|
use remora_schema::migrate as migrate_schema;
|
|
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,
|
|
};
|
|
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();
|
|
migrate_schema(&pool).await.unwrap();
|
|
sqlx::query("TRUNCATE TABLE jobs CASCADE")
|
|
.execute(&pool)
|
|
.await
|
|
.unwrap();
|
|
let coordination = CoordinationStore::connect(
|
|
&coordination_url,
|
|
format!(
|
|
"{}worker-runtime:{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 },
|
|
)
|
|
.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();
|
|
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
|
|
})
|
|
}
|