Implement crash-safe streaming worker runtime
This commit is contained in:
parent
16cba3c317
commit
cdeef3cc1a
1 changed files with 7448 additions and 194 deletions
488
crates/remora-worker/src/runtime.rs
Normal file
488
crates/remora-worker/src/runtime.rs
Normal file
|
|
@ -0,0 +1,488 @@
|
|||
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()
|
||||
}
|
||||
Loading…
Reference in a new issue