remora/crates/remora-worker/src/application.rs

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(),
            });
        Ok(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",
            ),
        })
    }
}

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

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