428 lines
16 KiB
Rust
428 lines
16 KiB
Rust
use std::{collections::BTreeMap, fs, path::PathBuf, sync::Arc, time::Duration};
|
|
|
|
use ipnet::IpNet;
|
|
use remora_coordination::{CoordinationStore, StoreConfig};
|
|
use remora_credentials::{CredentialRegistry, GrantBoundResolver, VaultClient, VaultConfig};
|
|
use remora_durable::{PoolConfig, connect};
|
|
use remora_outbound::OutboundPolicyBuilder;
|
|
use remora_runtime::{
|
|
ManagementConfig, RuntimeError, RuntimeStatus, env_duration, env_parse, env_string,
|
|
required_env, start_management_server,
|
|
};
|
|
use remora_stream::{CancelToken, CopyConfig, RetryTiming};
|
|
use remora_target::TargetStoreConfig;
|
|
use reqwest::{Certificate, Client, Identity};
|
|
use serde::Deserialize;
|
|
|
|
use crate::{AdapterExecutor, ObjectSourceEndpoint, WorkerApiClient, WorkerConfig, WorkerRuntime};
|
|
|
|
#[derive(Clone)]
|
|
pub struct WorkerApplicationConfig {
|
|
pub management: ManagementConfig,
|
|
pub database_url: String,
|
|
pub database_connections: u32,
|
|
pub coordination_url: String,
|
|
pub coordination_namespace: String,
|
|
pub coordination: StoreConfig,
|
|
pub worker: WorkerConfig,
|
|
pub worker_api_url: String,
|
|
pub worker_tls: WorkerClientTlsConfig,
|
|
pub target: TargetStoreConfig,
|
|
pub vault: VaultConfig,
|
|
pub object_sources: BTreeMap<String, ObjectSourceEndpoint>,
|
|
pub allowed_non_global: Vec<IpNet>,
|
|
pub allowed_scheme_ports: Vec<(String, u16)>,
|
|
pub max_redirects: usize,
|
|
pub retry_timing: RetryTiming,
|
|
pub copy: CopyConfig,
|
|
pub hugging_face_endpoint: String,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct WorkerClientTlsConfig {
|
|
pub certificate: PathBuf,
|
|
pub private_key: PathBuf,
|
|
pub server_ca: PathBuf,
|
|
}
|
|
|
|
pub struct WorkerApplication {
|
|
config: WorkerApplicationConfig,
|
|
}
|
|
|
|
impl WorkerApplicationConfig {
|
|
pub fn from_env() -> Result<Self, RuntimeError> {
|
|
let global_limit = env_parse("REMORA_GLOBAL_ATTEMPT_LIMIT", "64")?;
|
|
let lease_ttl = env_duration("REMORA_LEASE_TTL_SECONDS", 30)?;
|
|
let target = TargetStoreConfig::from_env().map_err(runtime_error)?;
|
|
let mut object_sources = object_source_endpoints()?;
|
|
object_sources
|
|
.entry(target.target_store_id.clone())
|
|
.or_insert_with(|| ObjectSourceEndpoint {
|
|
endpoint: target.endpoint.clone(),
|
|
region: target.region.clone(),
|
|
});
|
|
let config = Self {
|
|
management: ManagementConfig::from_env(crate::COMPONENT_NAME)?,
|
|
database_url: required_env("DATABASE_URL")?,
|
|
database_connections: env_parse("REMORA_DATABASE_MAX_CONNECTIONS", "16")?,
|
|
coordination_url: required_env("REMORA_COORDINATION_URL")?,
|
|
coordination_namespace: required_env("REMORA_COORDINATION_NAMESPACE")?,
|
|
coordination: StoreConfig {
|
|
global_limit,
|
|
lease_ttl,
|
|
reserve_ttl: env_duration("REMORA_RESERVE_TTL_SECONDS", 30)?,
|
|
progress_ttl: env_duration("REMORA_PROGRESS_TTL_SECONDS", 60)?,
|
|
cancel_ttl: None,
|
|
},
|
|
worker: WorkerConfig {
|
|
local_limit: env_parse("REMORA_WORKER_LOCAL_ATTEMPT_LIMIT", "4")?,
|
|
idle_delay: duration_millis("REMORA_WORKER_IDLE_MILLIS", 250)?,
|
|
renew_interval: env_duration("REMORA_WORKER_RENEW_SECONDS", 10)?,
|
|
cancel_poll_interval: duration_millis("REMORA_WORKER_CANCEL_POLL_MILLIS", 500)?,
|
|
},
|
|
worker_api_url: required_env("REMORA_WORKER_API_URL")?,
|
|
worker_tls: WorkerClientTlsConfig {
|
|
certificate: PathBuf::from(required_env("REMORA_WORKER_TLS_CERT_FILE")?),
|
|
private_key: PathBuf::from(required_env("REMORA_WORKER_TLS_KEY_FILE")?),
|
|
server_ca: PathBuf::from(required_env("REMORA_WORKER_TLS_SERVER_CA_FILE")?),
|
|
},
|
|
vault: VaultConfig::from_env().map_err(runtime_error)?,
|
|
target,
|
|
object_sources,
|
|
allowed_non_global: parse_cidrs(&env_string("REMORA_OUTBOUND_ALLOWED_CIDRS", ""))?,
|
|
allowed_scheme_ports: parse_scheme_ports(&env_string(
|
|
"REMORA_OUTBOUND_ALLOWED_SCHEME_PORTS",
|
|
"https:443",
|
|
))?,
|
|
max_redirects: env_parse("REMORA_OUTBOUND_MAX_REDIRECTS", "5")?,
|
|
retry_timing: RetryTiming {
|
|
connect_timeout: env_duration("REMORA_SOURCE_CONNECT_TIMEOUT_SECONDS", 10)?,
|
|
response_start_timeout: env_duration("REMORA_SOURCE_RESPONSE_TIMEOUT_SECONDS", 30)?,
|
|
no_progress_timeout: env_duration("REMORA_SOURCE_NO_PROGRESS_TIMEOUT_SECONDS", 60)?,
|
|
base_backoff: duration_millis("REMORA_SOURCE_BASE_BACKOFF_MILLIS", 100)?,
|
|
max_backoff: env_duration("REMORA_SOURCE_MAX_BACKOFF_SECONDS", 30)?,
|
|
max_server_hint: env_duration("REMORA_SOURCE_MAX_SERVER_HINT_SECONDS", 60)?,
|
|
},
|
|
copy: CopyConfig {
|
|
buffer_size: env_parse(
|
|
"REMORA_STREAM_BUFFER_BYTES",
|
|
&CopyConfig::DEFAULT_BUFFER_SIZE.to_string(),
|
|
)?,
|
|
progress_interval_bytes: env_parse("REMORA_PROGRESS_INTERVAL_BYTES", "1048576")?,
|
|
checkpoint_interval_bytes: env_parse(
|
|
"REMORA_CHECKPOINT_INTERVAL_BYTES",
|
|
"16777216",
|
|
)?,
|
|
no_progress_timeout: env_duration("REMORA_SOURCE_NO_PROGRESS_TIMEOUT_SECONDS", 60)?,
|
|
initial_completed_bytes: 0,
|
|
total_bytes: None,
|
|
completed_members: 0,
|
|
total_members: None,
|
|
}
|
|
.sanitized(),
|
|
hugging_face_endpoint: env_string(
|
|
"REMORA_HUGGING_FACE_ENDPOINT",
|
|
"https://huggingface.co",
|
|
),
|
|
};
|
|
validate_worker_timing(&config.coordination, &config.worker)?;
|
|
Ok(config)
|
|
}
|
|
}
|
|
|
|
impl WorkerApplication {
|
|
#[must_use]
|
|
pub fn new(config: WorkerApplicationConfig) -> Self {
|
|
Self { config }
|
|
}
|
|
|
|
pub async fn run(self) -> Result<(), RuntimeError> {
|
|
let _ = rustls::crypto::ring::default_provider().install_default();
|
|
let pool = connect(
|
|
&self.config.database_url,
|
|
PoolConfig {
|
|
max_connections: self.config.database_connections,
|
|
acquire_timeout: Duration::from_secs(10),
|
|
idle_timeout: Some(Duration::from_secs(60)),
|
|
},
|
|
)
|
|
.await
|
|
.map_err(runtime_error)?;
|
|
wait_for_schema(&pool).await?;
|
|
let coordination = CoordinationStore::connect(
|
|
&self.config.coordination_url,
|
|
&self.config.coordination_namespace,
|
|
self.config.coordination.clone(),
|
|
)
|
|
.await
|
|
.map_err(runtime_error)?;
|
|
let vault = VaultClient::new(self.config.vault.clone()).map_err(runtime_error)?;
|
|
let readiness_vault = vault.clone();
|
|
let credentials = CredentialRegistry::new(pool.clone(), vault.clone());
|
|
let resolver = GrantBoundResolver::new(credentials, vault);
|
|
let outbound_policy = outbound_policy(&self.config);
|
|
let executor = AdapterExecutor::new(
|
|
pool.clone(),
|
|
resolver,
|
|
self.config.target.clone(),
|
|
self.config.object_sources,
|
|
outbound_policy,
|
|
self.config.retry_timing,
|
|
self.config.copy,
|
|
self.config.hugging_face_endpoint,
|
|
)
|
|
.map_err(runtime_error)?;
|
|
let readiness_target = remora_target::TargetStore::connect(self.config.target.clone())
|
|
.map_err(runtime_error)?;
|
|
let api = WorkerApiClient::new(
|
|
self.config.worker_api_url,
|
|
worker_http_client(&self.config.worker_tls)?,
|
|
)
|
|
.map_err(runtime_error)?;
|
|
let readiness_api = api.clone();
|
|
let runtime = WorkerRuntime::new(
|
|
coordination.clone(),
|
|
api,
|
|
Arc::new(executor),
|
|
self.config.worker,
|
|
)
|
|
.map_err(runtime_error)?;
|
|
let status = RuntimeStatus::new(crate::COMPONENT_NAME, pool.clone(), coordination)
|
|
.with_readiness_check(move || {
|
|
let target = readiness_target.clone();
|
|
async move { target.check().await.is_ok() }
|
|
})
|
|
.with_readiness_check(move || {
|
|
let vault = readiness_vault.clone();
|
|
async move { vault.check().await.is_ok() }
|
|
})
|
|
.with_readiness_check(move || {
|
|
let api = readiness_api.clone();
|
|
async move { api.check().await.is_ok() }
|
|
});
|
|
let management = start_management_server(self.config.management, status.clone())
|
|
.map_err(runtime_error)?;
|
|
let shutdown = CancelToken::new();
|
|
let worker_shutdown = shutdown.clone();
|
|
let worker = tokio::spawn(async move { runtime.run(worker_shutdown).await });
|
|
status.set_accepting_traffic(true);
|
|
tracing::info!("worker application ready");
|
|
wait_for_shutdown().await?;
|
|
status.set_accepting_traffic(false);
|
|
shutdown.cancel();
|
|
worker
|
|
.await
|
|
.map_err(runtime_error)?
|
|
.map_err(runtime_error)?;
|
|
management.stop(true).await;
|
|
pool.close().await;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
async fn wait_for_schema(pool: &sqlx::PgPool) -> Result<(), RuntimeError> {
|
|
for _ in 0..120 {
|
|
if remora_schema::is_ready(pool).await.unwrap_or(false) {
|
|
return Ok(());
|
|
}
|
|
tokio::time::sleep(Duration::from_secs(1)).await;
|
|
}
|
|
Err(RuntimeError::message(
|
|
"database schema did not reach the required generation",
|
|
))
|
|
}
|
|
|
|
fn outbound_policy(config: &WorkerApplicationConfig) -> remora_outbound::OutboundPolicy {
|
|
let mut builder = OutboundPolicyBuilder::new().max_redirects(config.max_redirects);
|
|
for (scheme, port) in &config.allowed_scheme_ports {
|
|
builder = builder.allow_scheme_port(scheme, *port);
|
|
}
|
|
for network in &config.allowed_non_global {
|
|
builder = builder.allow_non_global(*network);
|
|
}
|
|
builder.build()
|
|
}
|
|
|
|
fn worker_http_client(config: &WorkerClientTlsConfig) -> Result<Client, RuntimeError> {
|
|
let ca_pem = fs::read(&config.server_ca)
|
|
.map_err(|error| file_error("REMORA_WORKER_TLS_SERVER_CA_FILE", error))?;
|
|
let certificate_pem = fs::read(&config.certificate)
|
|
.map_err(|error| file_error("REMORA_WORKER_TLS_CERT_FILE", error))?;
|
|
let private_key_pem = fs::read(&config.private_key)
|
|
.map_err(|error| file_error("REMORA_WORKER_TLS_KEY_FILE", error))?;
|
|
let mut identity_pem = certificate_pem;
|
|
identity_pem.extend_from_slice(&private_key_pem);
|
|
Client::builder()
|
|
.tls_certs_only([Certificate::from_pem(&ca_pem).map_err(runtime_error)?])
|
|
.identity(Identity::from_pem(&identity_pem).map_err(runtime_error)?)
|
|
.redirect(reqwest::redirect::Policy::none())
|
|
.no_proxy()
|
|
.build()
|
|
.map_err(runtime_error)
|
|
}
|
|
|
|
fn object_source_endpoints() -> Result<BTreeMap<String, ObjectSourceEndpoint>, RuntimeError> {
|
|
let input = env_string("REMORA_OBJECT_SOURCE_ENDPOINTS_JSON", "{}");
|
|
let parsed: BTreeMap<String, SourceEndpointConfig> =
|
|
serde_json::from_str(&input).map_err(|_| {
|
|
RuntimeError::message("REMORA_OBJECT_SOURCE_ENDPOINTS_JSON must be a JSON object")
|
|
})?;
|
|
parsed
|
|
.into_iter()
|
|
.map(|(store, endpoint)| {
|
|
if store.is_empty() || endpoint.endpoint.is_empty() || endpoint.region.is_empty() {
|
|
return Err(RuntimeError::message(
|
|
"REMORA_OBJECT_SOURCE_ENDPOINTS_JSON entries must be non-empty",
|
|
));
|
|
}
|
|
Ok((
|
|
store,
|
|
ObjectSourceEndpoint {
|
|
endpoint: endpoint.endpoint,
|
|
region: endpoint.region,
|
|
},
|
|
))
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn parse_cidrs(input: &str) -> Result<Vec<IpNet>, RuntimeError> {
|
|
if input.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
input
|
|
.split(',')
|
|
.map(str::trim)
|
|
.map(|value| {
|
|
value.parse().map_err(|_| {
|
|
RuntimeError::message("REMORA_OUTBOUND_ALLOWED_CIDRS contains an invalid CIDR")
|
|
})
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn parse_scheme_ports(input: &str) -> Result<Vec<(String, u16)>, RuntimeError> {
|
|
input
|
|
.split(',')
|
|
.map(str::trim)
|
|
.map(|value| {
|
|
let (scheme, port) = value.split_once(':').ok_or_else(|| {
|
|
RuntimeError::message(
|
|
"REMORA_OUTBOUND_ALLOWED_SCHEME_PORTS entries must be scheme:port",
|
|
)
|
|
})?;
|
|
let port = port.parse().map_err(|_| {
|
|
RuntimeError::message(
|
|
"REMORA_OUTBOUND_ALLOWED_SCHEME_PORTS contains an invalid port",
|
|
)
|
|
})?;
|
|
if scheme.is_empty() || port == 0 {
|
|
return Err(RuntimeError::message(
|
|
"REMORA_OUTBOUND_ALLOWED_SCHEME_PORTS entries must be non-empty",
|
|
));
|
|
}
|
|
Ok((scheme.to_ascii_lowercase(), port))
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn validate_worker_timing(
|
|
coordination: &StoreConfig,
|
|
worker: &WorkerConfig,
|
|
) -> Result<(), RuntimeError> {
|
|
if coordination.global_limit == 0
|
|
|| coordination.lease_ttl.is_zero()
|
|
|| coordination.reserve_ttl.is_zero()
|
|
|| coordination.progress_ttl.is_zero()
|
|
|| worker.local_limit == 0
|
|
|| worker.idle_delay.is_zero()
|
|
|| worker.renew_interval.is_zero()
|
|
|| worker.renew_interval >= coordination.lease_ttl
|
|
|| worker.cancel_poll_interval.is_zero()
|
|
{
|
|
return Err(RuntimeError::message(
|
|
"worker limits and timing must be positive, and renewal must precede lease expiry",
|
|
));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn wait_for_shutdown() -> Result<(), RuntimeError> {
|
|
#[cfg(unix)]
|
|
{
|
|
let mut terminate =
|
|
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
|
|
.map_err(runtime_error)?;
|
|
tokio::select! {
|
|
result = tokio::signal::ctrl_c() => result.map_err(runtime_error),
|
|
_ = terminate.recv() => Ok(()),
|
|
}
|
|
}
|
|
#[cfg(not(unix))]
|
|
tokio::signal::ctrl_c().await.map_err(runtime_error)
|
|
}
|
|
|
|
fn duration_millis(name: &str, default: u64) -> Result<Duration, RuntimeError> {
|
|
let milliseconds = env_parse(name, &default.to_string())?;
|
|
if milliseconds == 0 {
|
|
return Err(RuntimeError::message(format!(
|
|
"{name} must be greater than zero"
|
|
)));
|
|
}
|
|
Ok(Duration::from_millis(milliseconds))
|
|
}
|
|
|
|
fn runtime_error(error: impl std::fmt::Display) -> RuntimeError {
|
|
RuntimeError::message(error.to_string())
|
|
}
|
|
|
|
fn file_error(name: &str, error: std::io::Error) -> RuntimeError {
|
|
RuntimeError::message(format!("cannot read {name}: {}", error.kind()))
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct SourceEndpointConfig {
|
|
endpoint: String,
|
|
region: String,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn parses_explicit_source_reachability() {
|
|
assert_eq!(
|
|
parse_scheme_ports("https:443,http:8080").unwrap(),
|
|
vec![("https".into(), 443), ("http".into(), 8080)]
|
|
);
|
|
assert_eq!(
|
|
parse_cidrs("10.0.0.0/8,fd00::/8").unwrap(),
|
|
vec![
|
|
"10.0.0.0/8".parse::<IpNet>().unwrap(),
|
|
"fd00::/8".parse::<IpNet>().unwrap()
|
|
]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn rejects_ambiguous_source_reachability() {
|
|
assert!(parse_scheme_ports("https").is_err());
|
|
assert!(parse_cidrs("private").is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn validates_worker_timing_before_external_connections() {
|
|
let coordination = StoreConfig {
|
|
global_limit: 1,
|
|
lease_ttl: Duration::from_secs(30),
|
|
reserve_ttl: Duration::from_secs(30),
|
|
progress_ttl: Duration::from_secs(60),
|
|
cancel_ttl: None,
|
|
};
|
|
let mut worker = WorkerConfig::default();
|
|
assert!(validate_worker_timing(&coordination, &worker).is_ok());
|
|
worker.renew_interval = coordination.lease_ttl;
|
|
assert!(validate_worker_timing(&coordination, &worker).is_err());
|
|
}
|
|
}
|