Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: separate worker state from runtime lifecycle

Browse Source
This commit is contained in:
Hare
2026-09-06 10:05:54 +09:00
parent 282a8d31b5
commit 2d1956b653
21 changed files with 602 additions and 188 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker-runtime/src/runtime.rs
+131 -78
View File
@@ -700,6 +700,7 @@ impl Runtime {
worker_ref: worker_ref.clone(),
worker_id: worker_id.clone(),
status: WorkerStatus::Stopped,
worker_state: None,
workspace_id: scope.map(|scope| scope.workspace_id.clone()),
request: durable_request,
run_generation: 1,
@@ -787,7 +788,6 @@ impl Runtime {
let detail = self.commit_created_worker(
&worker_ref,
handle,
WorkerStatus::Running,
working_directory,
dispatch_result,
)?;
@@ -797,7 +797,6 @@ impl Runtime {
self.commit_created_worker(
&worker_ref,
handle,
WorkerStatus::Idle,
working_directory,
WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn),
)
@@ -1220,17 +1219,7 @@ impl Runtime {
state.ensure_running()?;
let worker = state.worker_mut(worker_ref)?;
if let Some(snapshot) = dispatch_result.worker_state.as_ref() {
worker.status = match snapshot.catalog_status() {
protocol::WorkerStatus::Idle => WorkerStatus::Idle,
protocol::WorkerStatus::Running => WorkerStatus::Running,
protocol::WorkerStatus::Paused => WorkerStatus::Paused,
protocol::WorkerStatus::Stopped => WorkerStatus::Stopped,
};
} else if matches!(
submission.as_ref().map(|ack| ack.disposition),
Some(protocol::SubmissionDisposition::Started)
) {
worker.status = WorkerStatus::Running;
let _ = worker.apply_worker_state(snapshot);
}
let status = worker.status;
#[cfg(feature = "ws-server")]
@@ -1490,16 +1479,19 @@ impl Runtime {
&self,
worker_ref: &WorkerRef,
handle: WorkerExecutionHandle,
status: WorkerStatus,
working_directory: Option<CatalogWorkingDirectoryStatus>,
_result: WorkerExecutionResult,
result: WorkerExecutionResult,
) -> Result<WorkerDetail, RuntimeError> {
let mut state = self.lock()?;
let detail = {
let worker = state.worker_mut(worker_ref)?;
worker.execution_handle = Some(handle);
worker.execution_bound = true;
worker.status = status;
worker.status = WorkerStatus::Idle;
worker.worker_state = None;
if let Some(snapshot) = result.worker_state.as_ref() {
let _ = worker.apply_worker_state(snapshot);
}
worker.restore_intent = restore_intent_for_status(worker.status);
worker.working_directory = working_directory;
worker.detail()
@@ -1536,16 +1528,14 @@ impl Runtime {
let Some(snapshot) = result.worker_state else {
return Ok(());
};
let status = match snapshot.catalog_status() {
protocol::WorkerStatus::Idle => WorkerStatus::Idle,
protocol::WorkerStatus::Running => WorkerStatus::Running,
protocol::WorkerStatus::Paused => WorkerStatus::Paused,
protocol::WorkerStatus::Stopped => WorkerStatus::Stopped,
};
let mut state = self.lock()?;
let worker = state.worker_mut(worker_ref)?;
worker.status = status;
worker.restore_intent = restore_intent_for_status(status);
let applied = worker
.apply_worker_state(&snapshot)
.is_ok_and(|result| matches!(result, protocol::WorkerStateSnapshotApply::Applied));
if !applied {
return Ok(());
}
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
@@ -1636,20 +1626,26 @@ impl Runtime {
worker_ref: &WorkerRef,
reason: Option<String>,
) -> Result<WorkerLifecycleAck, RuntimeError> {
let current = {
{
let state = self.lock()?;
state.ensure_running()?;
state.worker(worker_ref)?.status
};
if matches!(current, WorkerStatus::Idle | WorkerStatus::Stopped) {
return Ok(WorkerLifecycleAck {
worker_ref: worker_ref.clone(),
status: current,
});
if state.worker(worker_ref)?.status == WorkerStatus::Stopped {
return Ok(WorkerLifecycleAck {
worker_ref: worker_ref.clone(),
status: WorkerStatus::Stopped,
worker_state: None,
});
}
}
self.dispatch_lifecycle_to_backend(worker_ref, WorkerExecutionOperation::Cancel)?;
let _ = reason;
self.transition_worker_preserving_execution(worker_ref, WorkerStatus::Idle)
let state = self.lock()?;
let worker = state.worker(worker_ref)?;
Ok(WorkerLifecycleAck {
worker_ref: worker_ref.clone(),
status: worker.status,
worker_state: worker.worker_state.clone(),
})
}
/// Delete a non-running Worker through a workspace-scoped Runtime authorization context.
@@ -1795,12 +1791,13 @@ impl Runtime {
) -> Result<WorkerObservationEvent, RuntimeError> {
let mut state = self.lock()?;
state.ensure_worker_ref(worker_ref)?;
let status_changed = state.project_protocol_event_to_status(worker_ref, &payload);
let worker_state_changed =
state.project_protocol_event_to_worker_state(worker_ref, &payload);
let activity_changed = state.project_internal_worker_activity(worker_ref, &payload);
if status_changed || activity_changed {
if worker_state_changed || activity_changed {
state.publish_worker_upsert(worker_ref.worker_id)?;
}
if status_changed {
if worker_state_changed {
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
}
@@ -1836,26 +1833,6 @@ impl Runtime {
Ok(())
}
fn transition_worker_preserving_execution(
&self,
worker_ref: &WorkerRef,
status: WorkerStatus,
) -> Result<WorkerLifecycleAck, RuntimeError> {
let mut state = self.lock()?;
state.ensure_running()?;
let worker = state.worker_mut(worker_ref)?;
worker.status = status;
worker.restore_intent = restore_intent_for_status(status);
let status = worker.status;
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
Ok(WorkerLifecycleAck {
worker_ref: worker_ref.clone(),
status,
})
}
fn transition_worker(
&self,
worker_ref: &WorkerRef,
@@ -1867,6 +1844,7 @@ impl Runtime {
let worker = state.worker_mut(worker_ref)?;
worker.status = status;
worker.worker_state = None;
worker.restore_intent = restore_intent_for_status(status);
worker.execution_handle = None;
worker.internal_workers.clear();
@@ -1877,6 +1855,7 @@ impl Runtime {
Ok(WorkerLifecycleAck {
worker_ref: worker_ref.clone(),
status,
worker_state: None,
})
}
@@ -2301,6 +2280,7 @@ impl RuntimeState {
worker_ref: worker.worker_ref,
worker_id: worker.worker_id,
status: worker.status,
worker_state: None,
workspace_id: worker.workspace_id,
request: worker.request,
run_generation,
@@ -2630,6 +2610,7 @@ impl RuntimeState {
.get(&worker.worker_id)
.copied()
.unwrap_or(0),
worker_state: worker.worker_state.clone(),
state: subscription_worker_state(worker.status),
has_running_internal_workers: worker
.internal_workers
@@ -2965,7 +2946,7 @@ impl RuntimeState {
Self::update_internal_worker_activity(&mut worker.internal_workers, event)
}
fn project_protocol_event_to_status(
fn project_protocol_event_to_worker_state(
&mut self,
worker_ref: &WorkerRef,
event: &protocol::Event,
@@ -2973,7 +2954,7 @@ impl RuntimeState {
let Some(worker) = self.workers.get_mut(&worker_ref.worker_id) else {
return false;
};
let next_status = match event {
let incoming = match event {
protocol::Event::WorkerState { snapshot }
| protocol::Event::Snapshot {
state: snapshot, ..
@@ -2983,21 +2964,16 @@ impl RuntimeState {
protocol::WorkerCommandAcknowledgement {
state: snapshot, ..
},
} => Some(match snapshot.catalog_status() {
protocol::WorkerStatus::Idle => WorkerStatus::Idle,
protocol::WorkerStatus::Running => WorkerStatus::Running,
protocol::WorkerStatus::Paused => WorkerStatus::Paused,
protocol::WorkerStatus::Stopped => WorkerStatus::Stopped,
}),
_ => None,
} => snapshot,
_ => return false,
};
if let Some(next_status) = next_status {
let changed = worker.status != next_status;
worker.status = next_status;
worker.restore_intent = restore_intent_for_status(next_status);
changed
} else {
false
match worker.apply_worker_state(incoming) {
Ok(protocol::WorkerStateSnapshotApply::Applied) => true,
Ok(
protocol::WorkerStateSnapshotApply::Duplicate
| protocol::WorkerStateSnapshotApply::Stale,
)
| Err(_) => false,
}
}
}
@@ -3013,6 +2989,7 @@ struct WorkerRecord {
worker_ref: WorkerRef,
worker_id: WorkerId,
status: WorkerStatus,
worker_state: Option<protocol::WorkerStateSnapshot>,
workspace_id: Option<String>,
request: CreateWorkerRequest,
run_generation: u64,
@@ -3024,6 +3001,19 @@ struct WorkerRecord {
}
impl WorkerRecord {
fn apply_worker_state(
&mut self,
incoming: &protocol::WorkerStateSnapshot,
) -> Result<protocol::WorkerStateSnapshotApply, protocol::WorkerStateSnapshotConflict> {
match self.worker_state.as_mut() {
Some(current) => protocol::apply_worker_state_snapshot(current, incoming),
None => {
self.worker_state = Some(incoming.clone());
Ok(protocol::WorkerStateSnapshotApply::Applied)
}
}
}
fn belongs_to_workspace(&self, workspace_id: &str) -> bool {
self.workspace_id.as_deref() == Some(workspace_id)
}
@@ -3033,6 +3023,7 @@ impl WorkerRecord {
worker_ref: self.worker_ref.clone(),
worker_id: self.worker_id,
status: self.status,
worker_state: self.worker_state.clone(),
workspace_id: self.workspace_id.clone(),
working_directory: self.working_directory.clone(),
profile: self.request.profile.clone(),
@@ -3047,6 +3038,7 @@ impl WorkerRecord {
worker_ref: self.worker_ref.clone(),
worker_id: self.worker_id,
status: self.status,
worker_state: self.worker_state.clone(),
workspace_id: self.workspace_id.clone(),
working_directory: self.working_directory.clone(),
profile: self.request.profile.clone(),
@@ -4719,7 +4711,7 @@ mod tests {
}
#[test]
fn create_worker_uses_started_submission_ack_for_initial_running_status() {
fn create_worker_does_not_infer_state_from_started_submission_ack() {
let (runtime, backend) = runtime_and_backend();
backend.set_dispatch_result(WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
@@ -4732,7 +4724,69 @@ mod tests {
let detail = runtime.create_worker(request).unwrap();
assert_eq!(detail.status, WorkerStatus::Running);
assert_eq!(detail.status, WorkerStatus::Idle);
assert_eq!(detail.worker_state, None);
}
#[test]
fn runtime_applies_only_newer_worker_state_snapshots() {
let (runtime, _) = runtime_and_backend();
let detail = runtime
.create_worker(task_request("state ordering"))
.unwrap();
let running = protocol::WorkerStateSnapshot {
execution_generation: 7,
revision: 3,
state: protocol::WorkerState::Busy(protocol::WorkerBusyState::Run(
protocol::WorkerRunState::Running,
)),
last_command_id: 2,
};
assert!({
let mut state = runtime.lock().unwrap();
state.project_protocol_event_to_worker_state(
&detail.worker_ref,
&protocol::Event::WorkerState {
snapshot: running.clone(),
},
)
});
assert_eq!(
runtime
.worker_detail(&detail.worker_ref)
.unwrap()
.worker_state,
Some(running.clone())
);
assert!({
let mut state = runtime.lock().unwrap();
!state.project_protocol_event_to_worker_state(
&detail.worker_ref,
&protocol::Event::WorkerState {
snapshot: protocol::WorkerStateSnapshot {
revision: 2,
state: protocol::WorkerState::Idle,
..running.clone()
},
},
)
});
assert!({
let mut state = runtime.lock().unwrap();
!state.project_protocol_event_to_worker_state(
&detail.worker_ref,
&protocol::Event::WorkerState {
snapshot: protocol::WorkerStateSnapshot {
state: protocol::WorkerState::Idle,
..running.clone()
},
},
)
});
let after = runtime.worker_detail(&detail.worker_ref).unwrap();
assert_eq!(after.status, WorkerStatus::Idle);
assert_eq!(after.worker_state, Some(running));
}
#[test]
@@ -5027,10 +5081,9 @@ mod tests {
assert_eq!(*backend.restore_count.lock().unwrap(), 1);
assert_eq!(*backend.run_generations.lock().unwrap(), vec![1, 2]);
assert_eq!(
runtime.worker_detail(&detail.worker_ref).unwrap().status,
WorkerStatus::Running
);
let restored = runtime.worker_detail(&detail.worker_ref).unwrap();
assert_eq!(restored.status, WorkerStatus::Idle);
assert_eq!(restored.worker_state, None);
}
#[test]