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

use std::{future::Future, pin::Pin, sync::Arc, time::Duration};

use remora_contracts::{
    AttemptActivateRequestDto, AttemptAuthorityDto, AttemptFailureOutcomeDto, AttemptOutcomeDto,
    AttemptOutcomeRequestDto, AttemptRenewRequestDto, AttemptRuntimeFailureDto,
    AttemptStagingMultipartRequestDto, AttemptStagingObjectRequestDto,
    AttemptStagingWriteIntentRequestDto, AttemptStopRequestDto, AttemptSuccessOutcomeDto,
    ExecutionEnvelopeDto, ResultDescriptorDto, StagingMultipartIdentityDto,
    StagingObjectIdentityDto, StagingWriteIntentDto,
};
use remora_coordination::CoordinationStore;
use remora_domain::{JobPhase, OpaqueId, RuntimeFailureClass};
use remora_stream::CancelToken;
use tokio::{sync::Semaphore, task::JoinSet, time};

use crate::{WorkerApiClient, WorkerError, incarnation_id, operation_id};

#[derive(Clone, Debug)]
pub struct WorkerConfig {
    pub local_limit: usize,
    pub idle_delay: Duration,
    pub renew_interval: Duration,
    pub cancel_poll_interval: Duration,
}

impl Default for WorkerConfig {
    fn default() -> Self {
        Self {
            local_limit: 4,
            idle_delay: Duration::from_millis(250),
            renew_interval: Duration::from_secs(10),
            cancel_poll_interval: Duration::from_millis(500),
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum WorkerRunOutcome {
    Idle,
    Succeeded,
    Failed(JobPhase),
    Cancelled,
    AuthorityLost,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ExecutionFailure {
    pub class: RuntimeFailureClass,
    pub message: &'static str,
}

impl ExecutionFailure {
    #[must_use]
    pub const fn new(class: RuntimeFailureClass, message: &'static str) -> Self {
        Self { class, message }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum AttemptExecution {
    Succeeded(ResultDescriptorDto),
    Failed(ExecutionFailure),
    Cancelled,
}

pub trait AttemptExecutor: Send + Sync + 'static {
    fn execute(
        &self,
        context: AttemptContext,
        envelope: ExecutionEnvelopeDto,
    ) -> Pin<Box<dyn Future<Output = AttemptExecution> + Send>>;
}

#[derive(Clone)]
pub struct AttemptContext {
    api: WorkerApiClient,
    coordination: CoordinationStore,
    authority: AttemptAuthorityDto,
    job_id: OpaqueId,
    cancel: CancelToken,
}

impl AttemptContext {
    #[must_use]
    pub fn authority(&self) -> &AttemptAuthorityDto {
        &self.authority
    }

    #[must_use]
    pub fn job_id(&self) -> &OpaqueId {
        &self.job_id
    }

    #[must_use]
    pub fn cancellation(&self) -> &CancelToken {
        &self.cancel
    }

    pub async fn require_live_authority(&self) -> Result<(), WorkerError> {
        let lease = self
            .coordination
            .get_live_lease(self.job_id.as_str())
            .await
            .map_err(|_| WorkerError::CoordinationUnavailable)?;
        let live = lease.is_some_and(|lease| {
            lease.record.attempt_id == self.authority.attempt_id.as_str()
                && lease.record.fence == self.authority.fence.get()
                && lease.record.worker_id == self.authority.worker_incarnation.as_str()
        });
        if !live || self.cancel.is_cancelled() {
            return Err(WorkerError::PreconditionFailed);
        }
        Ok(())
    }

    pub async fn report_progress(
        &self,
        progress: remora_contracts::ProgressValueDto,
    ) -> Result<(), WorkerError> {
        self.require_live_authority().await?;
        self.api
            .report_progress(
                &self.authority.attempt_id,
                operation_id(),
                self.authority.fence,
                self.authority.worker_incarnation.clone(),
                progress,
            )
            .await
    }

    pub async fn pin_resolved_source(
        &self,
        descriptor: remora_contracts::ResolvedSourceDescriptorRefDto,
    ) -> Result<(), WorkerError> {
        self.require_live_authority().await?;
        self.api
            .pin_resolved_source(
                &self.authority.attempt_id,
                operation_id(),
                self.authority.fence,
                self.authority.worker_incarnation.clone(),
                descriptor,
            )
            .await
    }

    pub async fn persist_checkpoint(
        &self,
        checkpoint_ref: remora_contracts::CheckpointDescriptorRefDto,
        resolved_source_id: OpaqueId,
        validated_bytes: remora_contracts::SafeIntegerDto,
        validated_members: remora_contracts::SafeIntegerDto,
        staging_multipart: Option<StagingMultipartIdentityDto>,
    ) -> Result<(), WorkerError> {
        self.require_live_authority().await?;
        self.api
            .persist_checkpoint(
                &self.authority.attempt_id,
                remora_contracts::AttemptCheckpointRequestDto {
                    operation_id: operation_id(),
                    fence: self.authority.fence,
                    worker_incarnation: self.authority.worker_incarnation.clone(),
                    checkpoint_ref,
                    resolved_source_id,
                    validated_bytes,
                    validated_members,
                    staging_multipart,
                },
            )
            .await
    }

    pub async fn register_staging_object(
        &self,
        object: StagingObjectIdentityDto,
        completed_upload_id: Option<String>,
    ) -> Result<(), WorkerError> {
        self.require_live_authority().await?;
        self.api
            .register_staging_object(
                &self.authority.attempt_id,
                AttemptStagingObjectRequestDto {
                    operation_id: operation_id(),
                    fence: self.authority.fence,
                    worker_incarnation: self.authority.worker_incarnation.clone(),
                    object,
                    completed_upload_id,
                },
            )
            .await
    }

    pub async fn register_staging_write_intent(
        &self,
        intent: StagingWriteIntentDto,
    ) -> Result<(), WorkerError> {
        self.require_live_authority().await?;
        self.api
            .register_staging_write_intent(
                &self.authority.attempt_id,
                AttemptStagingWriteIntentRequestDto {
                    operation_id: operation_id(),
                    fence: self.authority.fence,
                    worker_incarnation: self.authority.worker_incarnation.clone(),
                    intent,
                },
            )
            .await
    }

    pub async fn register_staging_multipart(
        &self,
        multipart: StagingMultipartIdentityDto,
    ) -> Result<(), WorkerError> {
        self.require_live_authority().await?;
        self.api
            .register_staging_multipart(
                &self.authority.attempt_id,
                AttemptStagingMultipartRequestDto {
                    operation_id: operation_id(),
                    fence: self.authority.fence,
                    worker_incarnation: self.authority.worker_incarnation.clone(),
                    multipart,
                },
            )
            .await
    }
}

#[derive(Clone)]
pub struct WorkerRuntime {
    incarnation: OpaqueId,
    coordination: CoordinationStore,
    api: WorkerApiClient,
    executor: Arc<dyn AttemptExecutor>,
    permits: Arc<Semaphore>,
    config: WorkerConfig,
}

impl WorkerRuntime {
    pub fn new(
        coordination: CoordinationStore,
        api: WorkerApiClient,
        executor: Arc<dyn AttemptExecutor>,
        config: WorkerConfig,
    ) -> Result<Self, WorkerError> {
        if config.local_limit == 0
            || config.idle_delay.is_zero()
            || config.renew_interval.is_zero()
            || config.renew_interval >= coordination.config().lease_ttl
            || config.cancel_poll_interval.is_zero()
        {
            return Err(WorkerError::InvalidConfiguration);
        }
        Ok(Self {
            incarnation: incarnation_id(),
            coordination,
            api,
            executor,
            permits: Arc::new(Semaphore::new(config.local_limit)),
            config,
        })
    }

    #[must_use]
    pub fn incarnation(&self) -> &OpaqueId {
        &self.incarnation
    }

    pub async fn run_once(&self) -> Result<WorkerRunOutcome, WorkerError> {
        let permit = self
            .permits
            .clone()
            .try_acquire_owned()
            .map_err(|_| WorkerError::PreconditionFailed)?;
        let Some(job) = self
            .coordination
            .dequeue()
            .await
            .map_err(|_| WorkerError::CoordinationUnavailable)?
        else {
            return Ok(WorkerRunOutcome::Idle);
        };
        let job_id = OpaqueId::parse(job).map_err(|_| WorkerError::InvalidResponse)?;
        let outcome = self.process(job_id).await;
        drop(permit);
        outcome
    }

    pub async fn run(&self, shutdown: CancelToken) -> Result<(), WorkerError> {
        let mut tasks = JoinSet::new();
        while !shutdown.is_cancelled() {
            while tasks.len() < self.config.local_limit {
                let runtime = self.clone();
                tasks.spawn(async move {
                    let outcome = runtime.run_once().await;
                    if matches!(outcome, Ok(WorkerRunOutcome::Idle)) {
                        time::sleep(runtime.config.idle_delay).await;
                    }
                    outcome
                });
            }
            tokio::select! {
                () = shutdown.cancelled() => break,
                result = tasks.join_next() => {
                    if let Some(result) = result {
                        let _ = result.map_err(|_| WorkerError::ApiUnavailable)?;
                    }
                }
            }
        }
        shutdown.cancel();
        tasks.abort_all();
        while let Some(result) = tasks.join_next().await {
            match result {
                Err(error) if error.is_cancelled() => {}
                Err(_) => return Err(WorkerError::ApiUnavailable),
                Ok(outcome) => {
                    outcome?;
                }
            }
        }
        Ok(())
    }

    async fn process(&self, job_id: OpaqueId) -> Result<WorkerRunOutcome, WorkerError> {
        let claim = self
            .api
            .claim(&job_id, operation_id(), self.incarnation.clone())
            .await?;
        validate_claim(&job_id, &self.incarnation, &claim)?;
        let activation = self
            .api
            .activate(
                &claim.authority.attempt_id,
                AttemptActivateRequestDto {
                    operation_id: operation_id(),
                    fence: claim.authority.fence,
                    worker_incarnation: self.incarnation.clone(),
                },
            )
            .await?;
        if activation.authority != claim.authority || activation.execution != claim.execution {
            return Err(WorkerError::InvalidResponse);
        }
        let cancel = CancelToken::new();
        let context = AttemptContext {
            api: self.api.clone(),
            coordination: self.coordination.clone(),
            authority: activation.authority.clone(),
            job_id: job_id.clone(),
            cancel: cancel.clone(),
        };
        context.require_live_authority().await?;
        let mut execution = self.executor.execute(context.clone(), activation.execution);
        let mut renewal = time::interval(self.config.renew_interval);
        renewal.reset();
        let mut cancellation = time::interval(self.config.cancel_poll_interval);
        cancellation.reset();
        let mut authority_lost = false;
        let completed = loop {
            tokio::select! {
                result = &mut execution => break result,
                _ = renewal.tick() => {
                    if self.api.renew(
                        &claim.authority.attempt_id,
                        AttemptRenewRequestDto {
                            operation_id: operation_id(),
                            fence: claim.authority.fence,
                            worker_incarnation: self.incarnation.clone(),
                        },
                    ).await.is_err() {
                        authority_lost = true;
                        cancel.cancel();
                    }
                }
                _ = cancellation.tick() => {
                    let requested = self.coordination.cancel_notice(job_id.as_str())
                        .await
                        .map_err(|_| WorkerError::CoordinationUnavailable)?;
                    if requested {
                        cancel.cancel();
                    }
                }
            }
        };
        if authority_lost {
            let _ = self.stop_attempt(&claim.authority).await;
            return Ok(WorkerRunOutcome::AuthorityLost);
        }
        if cancel.is_cancelled() || matches!(completed, AttemptExecution::Cancelled) {
            self.stop_attempt(&claim.authority).await?;
            return Ok(WorkerRunOutcome::Cancelled);
        }
        if context.require_live_authority().await.is_err() {
            let _ = self.stop_attempt(&claim.authority).await;
            return Ok(WorkerRunOutcome::AuthorityLost);
        }
        let outcome = match completed {
            AttemptExecution::Succeeded(result) => {
                AttemptOutcomeDto::Success(AttemptSuccessOutcomeDto { result })
            }
            AttemptExecution::Failed(failure) => {
                AttemptOutcomeDto::Failure(AttemptFailureOutcomeDto {
                    failure: AttemptRuntimeFailureDto {
                        class: failure.class,
                        message: failure.message.into(),
                        safe_details: None,
                    },
                })
            }
            AttemptExecution::Cancelled => unreachable!(),
        };
        let phase = self
            .api
            .report_outcome(
                &claim.authority.attempt_id,
                AttemptOutcomeRequestDto {
                    operation_id: operation_id(),
                    fence: claim.authority.fence,
                    worker_incarnation: self.incarnation.clone(),
                    outcome,
                },
            )
            .await?
            .phase;
        Ok(if phase == JobPhase::Succeeded {
            WorkerRunOutcome::Succeeded
        } else {
            WorkerRunOutcome::Failed(phase)
        })
    }

    async fn stop_attempt(&self, authority: &AttemptAuthorityDto) -> Result<(), WorkerError> {
        self.api
            .stop(
                &authority.attempt_id,
                AttemptStopRequestDto {
                    operation_id: operation_id(),
                    fence: authority.fence,
                    worker_incarnation: authority.worker_incarnation.clone(),
                },
            )
            .await
    }
}

fn validate_claim(
    job_id: &OpaqueId,
    incarnation: &OpaqueId,
    claim: &remora_contracts::ClaimResponseDto,
) -> Result<(), WorkerError> {
    if claim.authority.worker_incarnation != *incarnation
        || claim.execution.job_id != *job_id
        || claim.execution.attempt_id != claim.authority.attempt_id
        || claim.execution.fence != claim.authority.fence
        || claim.execution.worker_incarnation != claim.authority.worker_incarnation
        || claim.execution.id != claim.execution.execution_grant_ref
        || claim.execution.staging_target.store != claim.execution.target.store
        || claim.execution.staging_target.bucket != claim.execution.target.bucket
        || !valid_staging_target(&claim.execution)
    {
        return Err(WorkerError::InvalidResponse);
    }
    Ok(())
}

fn valid_staging_target(execution: &remora_contracts::ExecutionEnvelopeDto) -> bool {
    let prefix = &execution.staging_target.prefix;
    let root = if execution.target.prefix.is_empty() {
        format!("{}/a/", remora_target::RESERVED_NAMESPACE_SEGMENT)
    } else {
        format!(
            "{}/{}/a/",
            execution.target.prefix,
            remora_target::RESERVED_NAMESPACE_SEGMENT
        )
    };
    let Some(origin) = prefix
        .strip_prefix(&root)
        .and_then(|value| value.strip_suffix('/'))
        .and_then(|value| OpaqueId::parse(value).ok())
    else {
        return false;
    };
    origin == execution.attempt_id || execution.checkpoint_ref.0.is_some()
}