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

feat: add revisioned worker execution state

Browse Source
This commit is contained in:
Hare
2026-09-06 06:54:29 +09:00
parent 3ed1545c3c
commit e8b9adcde4
32 changed files with 2168 additions and 750 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker-runtime/src/execution.rs
+14 -28
View File
@@ -15,18 +15,6 @@ use std::fmt;
use std::sync::Arc;
use workdir::WorkdirSessionHandle;
/// Current execution-side run state for a Worker.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkerExecutionRunState {
#[default]
Stopped,
Idle,
Busy,
Rejected,
Errored,
}
/// Execution operation that produced a result.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
@@ -55,7 +43,8 @@ pub struct WorkerSubmissionAck {
pub struct WorkerExecutionResult {
pub operation: WorkerExecutionOperation,
pub outcome: WorkerExecutionOutcome,
pub run_state: WorkerExecutionRunState,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub worker_state: Option<protocol::WorkerStateSnapshot>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub message: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
@@ -74,22 +63,23 @@ pub enum WorkerExecutionOutcome {
}
impl WorkerExecutionResult {
pub fn accepted(
operation: WorkerExecutionOperation,
run_state: WorkerExecutionRunState,
) -> Self {
pub fn accepted(operation: WorkerExecutionOperation) -> Self {
Self {
operation,
outcome: WorkerExecutionOutcome::Accepted,
run_state,
worker_state: None,
message: None,
submission: None,
}
}
pub fn with_worker_state(mut self, worker_state: protocol::WorkerStateSnapshot) -> Self {
self.worker_state = Some(worker_state);
self
}
pub fn accepted_submission(
operation: WorkerExecutionOperation,
run_state: WorkerExecutionRunState,
submission_request_id: impl Into<String>,
submission_id: impl Into<String>,
disposition: protocol::SubmissionDisposition,
@@ -97,7 +87,7 @@ impl WorkerExecutionResult {
Self {
operation,
outcome: WorkerExecutionOutcome::Accepted,
run_state,
worker_state: None,
message: None,
submission: Some(WorkerSubmissionAck {
submission_request_id: submission_request_id.into(),
@@ -111,7 +101,7 @@ impl WorkerExecutionResult {
Self {
operation,
outcome: WorkerExecutionOutcome::Busy,
run_state: WorkerExecutionRunState::Busy,
worker_state: None,
message: Some(message.into()),
submission: None,
}
@@ -121,7 +111,7 @@ impl WorkerExecutionResult {
Self {
operation,
outcome: WorkerExecutionOutcome::Rejected,
run_state: WorkerExecutionRunState::Stopped,
worker_state: None,
message: Some(message.into()),
submission: None,
}
@@ -131,7 +121,7 @@ impl WorkerExecutionResult {
Self {
operation,
outcome: WorkerExecutionOutcome::Errored,
run_state: WorkerExecutionRunState::Errored,
worker_state: None,
message: Some(message.into()),
submission: None,
}
@@ -141,7 +131,7 @@ impl WorkerExecutionResult {
Self {
operation,
outcome: WorkerExecutionOutcome::Unsupported,
run_state: WorkerExecutionRunState::Stopped,
worker_state: None,
message: Some(message.into()),
submission: None,
}
@@ -280,7 +270,6 @@ pub struct WorkerExecutionRestoreRequest {
pub enum WorkerExecutionSpawnResult {
Connected {
handle: WorkerExecutionHandle,
run_state: WorkerExecutionRunState,
working_directory: Option<WorkingDirectoryStatus>,
},
Rejected(WorkerExecutionResult),
@@ -290,12 +279,10 @@ pub enum WorkerExecutionSpawnResult {
impl WorkerExecutionSpawnResult {
pub fn connected(
handle: WorkerExecutionHandle,
run_state: WorkerExecutionRunState,
working_directory: Option<WorkingDirectoryStatus>,
) -> Self {
Self::Connected {
handle,
run_state,
working_directory,
}
}
@@ -623,7 +610,6 @@ mod tests {
fn submission_ack_survives_json_round_trip() {
let result = WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
"request-1",
"submission-1",
protocol::SubmissionDisposition::Started,
crates/worker-runtime/src/http_server.rs
+11 -29
View File
@@ -2206,8 +2206,8 @@ mod tests {
};
use crate::execution::{
WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation,
WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionRunState,
WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult,
WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest,
WorkerExecutionSpawnResult,
};
use crate::management::RuntimeOptions;
use axum::body::to_bytes;
@@ -2979,7 +2979,6 @@ mod tests {
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request
.working_directory
.as_ref()
@@ -2993,7 +2992,6 @@ mod tests {
) -> WorkerExecutionSpawnResult {
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request.previous_working_directory,
}
}
@@ -3006,24 +3004,17 @@ mod tests {
if let Some(submission_id) = input.submission_request_id {
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
submission_id.clone(),
submission_id,
protocol::SubmissionDisposition::Started,
)
} else {
WorkerExecutionResult::accepted(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
)
WorkerExecutionResult::accepted(WorkerExecutionOperation::Input)
}
}
fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
WorkerExecutionResult::accepted(
WorkerExecutionOperation::Stop,
WorkerExecutionRunState::Stopped,
)
WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop)
}
}
@@ -3295,8 +3286,7 @@ mod ws_tests {
};
use crate::execution::{
WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation,
WorkerExecutionResult, WorkerExecutionRunState, WorkerExecutionSpawnRequest,
WorkerExecutionSpawnResult,
WorkerExecutionResult, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult,
};
use crate::management::RuntimeOptions;
use futures::{SinkExt, StreamExt};
@@ -3316,7 +3306,6 @@ mod ws_tests {
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request
.working_directory
.as_ref()
@@ -3332,16 +3321,12 @@ mod ws_tests {
if let Some(submission_id) = input.submission_request_id {
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
submission_id.clone(),
submission_id,
protocol::SubmissionDisposition::Started,
)
} else {
WorkerExecutionResult::accepted(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
)
WorkerExecutionResult::accepted(WorkerExecutionOperation::Input)
}
}
@@ -3350,10 +3335,7 @@ mod ws_tests {
_handle: &WorkerExecutionHandle,
_method: protocol::Method,
) -> WorkerExecutionResult {
WorkerExecutionResult::accepted(
WorkerExecutionOperation::ProtocolMethod,
WorkerExecutionRunState::Idle,
)
WorkerExecutionResult::accepted(WorkerExecutionOperation::ProtocolMethod)
}
}
@@ -3564,16 +3546,16 @@ mod ws_tests {
runtime
.observe_worker_event(
&other.worker_ref,
protocol::Event::Status {
status: protocol::WorkerStatus::Running,
protocol::Event::WorkerState {
snapshot: protocol::WorkerStatus::Running.into(),
},
)
.unwrap();
runtime
.observe_worker_event(
&worker_ref,
protocol::Event::Status {
status: protocol::WorkerStatus::Running,
protocol::Event::WorkerState {
snapshot: protocol::WorkerStatus::Running.into(),
},
)
.unwrap();
crates/worker-runtime/src/runtime.rs
+109 -107
View File
@@ -13,8 +13,8 @@ use crate::error::RuntimeError;
use crate::execution::WorkerExecutionRestoreRequest;
use crate::execution::{
WorkerExecutionBackend, WorkerExecutionBackendRef, WorkerExecutionHandle,
WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionRunState,
WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult,
WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionSpawnRequest,
WorkerExecutionSpawnResult,
};
#[cfg(feature = "fs-store")]
use crate::fs_store::{
@@ -725,12 +725,11 @@ impl Runtime {
};
let spawn_result = backend.spawn_worker(spawn_request);
let (handle, run_state, working_directory) = match spawn_result {
let (handle, working_directory) = match spawn_result {
WorkerExecutionSpawnResult::Connected {
handle,
run_state,
working_directory,
} => (handle, run_state, working_directory),
} => (handle, working_directory),
WorkerExecutionSpawnResult::Rejected(result)
| WorkerExecutionSpawnResult::Errored(result) => {
self.rollback_failed_create(&worker_ref)?;
@@ -785,11 +784,10 @@ impl Runtime {
result,
});
}
let initial_run_state = dispatch_result.run_state;
let detail = self.commit_created_worker(
&worker_ref,
handle,
initial_run_state,
WorkerStatus::Running,
working_directory,
dispatch_result,
)?;
@@ -799,9 +797,9 @@ impl Runtime {
self.commit_created_worker(
&worker_ref,
handle,
run_state,
WorkerStatus::Idle,
working_directory,
WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn, run_state),
WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn),
)
}
}
@@ -1086,13 +1084,12 @@ impl Runtime {
match backend.restore_worker(request) {
WorkerExecutionSpawnResult::Connected {
handle,
run_state,
working_directory,
} => {
self.commit_restored_worker_execution(
worker_ref,
handle,
run_state,
WorkerStatus::Idle,
working_directory,
)?;
self.worker_detail(worker_ref)
@@ -1222,7 +1219,19 @@ impl Runtime {
let mut state = self.lock()?;
state.ensure_running()?;
let worker = state.worker_mut(worker_ref)?;
worker.status = worker_status_from_run_state(dispatch_result.run_state);
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 status = worker.status;
#[cfg(feature = "ws-server")]
if let Some(payload) = input_protocol_event(&input) {
@@ -1431,7 +1440,7 @@ impl Runtime {
let entries = self.worker_completions(worker_ref, kind, &prefix)?;
return Ok(vec![Event::Completions { kind, entries }]);
}
if matches!(&method, Method::Shutdown) {
if matches!(&method, Method::Shutdown { .. }) {
self.stop_worker(worker_ref, Some("worker protocol shutdown".to_string()))?;
return Ok(Vec::new());
}
@@ -1481,7 +1490,7 @@ impl Runtime {
&self,
worker_ref: &WorkerRef,
handle: WorkerExecutionHandle,
run_state: WorkerExecutionRunState,
status: WorkerStatus,
working_directory: Option<CatalogWorkingDirectoryStatus>,
_result: WorkerExecutionResult,
) -> Result<WorkerDetail, RuntimeError> {
@@ -1490,7 +1499,7 @@ impl Runtime {
let worker = state.worker_mut(worker_ref)?;
worker.execution_handle = Some(handle);
worker.execution_bound = true;
worker.status = worker_status_from_run_state(run_state);
worker.status = status;
worker.restore_intent = restore_intent_for_status(worker.status);
worker.working_directory = working_directory;
worker.detail()
@@ -1518,16 +1527,28 @@ impl Runtime {
worker_ref: &WorkerRef,
result: WorkerExecutionResult,
) -> Result<(), RuntimeError> {
let mut state = self.lock()?;
if result.is_accepted() {
let status = worker_status_from_run_state(result.run_state);
let worker = state.worker_mut(worker_ref)?;
worker.status = status;
worker.restore_intent = restore_intent_for_status(status);
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
// Accepted dispatch without a state snapshot is transport evidence only;
// the revisioned protocol stream remains live authority. Test/detached
// backends may return an exact full snapshot as their acknowledgement.
if !result.is_accepted() {
return Ok(());
}
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);
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
Ok(())
}
@@ -1730,7 +1751,7 @@ impl Runtime {
context_window: 0,
context_tokens: 0,
},
status: protocol::WorkerStatus::Idle,
state: protocol::WorkerStateSnapshot::initial(1),
in_flight: protocol::InFlightSnapshot {
blocks: Vec::new(),
commands: Vec::new(),
@@ -1968,12 +1989,11 @@ impl Runtime {
match backend.restore_worker(request) {
WorkerExecutionSpawnResult::Connected {
handle,
run_state,
working_directory,
} => self.commit_restored_worker_execution(
&candidate.worker_ref,
handle,
run_state,
WorkerStatus::Idle,
working_directory,
)?,
WorkerExecutionSpawnResult::Rejected(result)
@@ -1990,7 +2010,7 @@ impl Runtime {
&self,
worker_ref: &WorkerRef,
handle: WorkerExecutionHandle,
run_state: WorkerExecutionRunState,
status: WorkerStatus,
working_directory: Option<CatalogWorkingDirectoryStatus>,
) -> Result<(), RuntimeError> {
let mut state = self.lock()?;
@@ -1999,7 +2019,7 @@ impl Runtime {
let worker = state.worker_mut(worker_ref)?;
worker.execution_handle = Some(handle);
worker.execution_bound = true;
worker.status = worker_status_from_run_state(run_state);
worker.status = status;
worker.restore_intent = restore_intent_for_status(worker.status);
worker.working_directory = working_directory;
}
@@ -2867,7 +2887,7 @@ impl RuntimeState {
) {
match event {
protocol::Event::Snapshot {
status,
state,
internal_workers,
..
} => {
@@ -2875,7 +2895,7 @@ impl RuntimeState {
statuses.insert(
worker.session_id.clone(),
InternalWorkerActivity {
status: *status,
status: state.catalog_status(),
parent_session_id: worker.parent_session_id.clone(),
},
);
@@ -2888,26 +2908,17 @@ impl RuntimeState {
event,
..
} => Self::project_internal_worker_event(statuses, nested_worker, event),
protocol::Event::Status { status } => {
statuses.insert(
worker.session_id.clone(),
InternalWorkerActivity {
status: *status,
parent_session_id: worker.parent_session_id.clone(),
protocol::Event::WorkerState { snapshot }
| protocol::Event::CommandAcknowledged {
acknowledgement:
protocol::WorkerCommandAcknowledgement {
state: snapshot, ..
},
);
}
protocol::Event::RunEnd { result } => {
let status = match result {
protocol::RunResult::Paused => protocol::WorkerStatus::Paused,
protocol::RunResult::Finished
| protocol::RunResult::LimitReached
| protocol::RunResult::RolledBack => protocol::WorkerStatus::Idle,
};
} => {
statuses.insert(
worker.session_id.clone(),
InternalWorkerActivity {
status,
status: snapshot.catalog_status(),
parent_session_id: worker.parent_session_id.clone(),
},
);
@@ -2963,28 +2974,21 @@ impl RuntimeState {
return false;
};
let next_status = match event {
protocol::Event::Status {
status: protocol::WorkerStatus::Running,
} => Some(WorkerStatus::Running),
protocol::Event::Status {
status: protocol::WorkerStatus::Idle,
} => Some(WorkerStatus::Idle),
protocol::Event::Status {
status: protocol::WorkerStatus::Paused,
} => Some(WorkerStatus::Paused),
protocol::Event::Snapshot { status, .. } => match status {
protocol::WorkerStatus::Running => Some(WorkerStatus::Running),
protocol::WorkerStatus::Idle => Some(WorkerStatus::Idle),
protocol::WorkerStatus::Paused => Some(WorkerStatus::Paused),
protocol::WorkerStatus::Stopped => Some(WorkerStatus::Stopped),
},
protocol::Event::RunEnd { result } => match result {
protocol::RunResult::Finished | protocol::RunResult::RolledBack => {
Some(WorkerStatus::Idle)
}
protocol::RunResult::Paused => Some(WorkerStatus::Paused),
protocol::RunResult::LimitReached => Some(WorkerStatus::Idle),
},
protocol::Event::WorkerState { snapshot }
| protocol::Event::Snapshot {
state: snapshot, ..
}
| protocol::Event::CommandAcknowledged {
acknowledgement:
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,
};
if let Some(next_status) = next_status {
@@ -3081,16 +3085,6 @@ fn restore_intent_for_status(status: WorkerStatus) -> WorkerRestoreIntent {
}
}
fn worker_status_from_run_state(run_state: WorkerExecutionRunState) -> WorkerStatus {
match run_state {
WorkerExecutionRunState::Idle => WorkerStatus::Idle,
WorkerExecutionRunState::Busy => WorkerStatus::Running,
WorkerExecutionRunState::Stopped
| WorkerExecutionRunState::Rejected
| WorkerExecutionRunState::Errored => WorkerStatus::Stopped,
}
}
fn repository_resource_error(error: BackendResourceError) -> RuntimeError {
let (code, message) = match error {
BackendResourceError::Expired => (
@@ -3304,7 +3298,7 @@ mod tests {
};
use crate::execution::{
WorkerExecutionBackend, WorkerExecutionContext, WorkerExecutionHandle,
WorkerExecutionRestoreRequest, WorkerExecutionRunState,
WorkerExecutionRestoreRequest,
};
use crate::working_directory::WorkingDirectoryDiagnostic;
use async_trait::async_trait;
@@ -3313,6 +3307,14 @@ mod tests {
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
fn test_command() -> protocol::WorkerCommandEnvelope {
protocol::WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 1,
expected_worker_state_revision: 0,
}
}
#[test]
fn repository_resource_failures_keep_typed_credential_diagnostics() {
let cases = [
@@ -3359,7 +3361,9 @@ mod tests {
protocol::Event::InternalWorker {
worker,
revision: 1,
event: Box::new(protocol::Event::Status { status }),
event: Box::new(protocol::Event::WorkerState {
snapshot: status.into(),
}),
}
}
@@ -3452,7 +3456,7 @@ mod tests {
context_window: 0,
context_tokens: 0,
},
status: protocol::WorkerStatus::Idle,
state: protocol::WorkerStatus::Idle.into(),
in_flight: protocol::InFlightSnapshot::default(),
internal_workers: Vec::new(),
};
@@ -3967,7 +3971,6 @@ mod tests {
.insert(request.worker_ref.worker_id.clone(), request.context);
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request
.working_directory
.as_ref()
@@ -3997,7 +4000,6 @@ mod tests {
.insert(request.worker_ref.worker_id.clone(), request.context);
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request
.working_directory
.as_ref()
@@ -4020,7 +4022,6 @@ mod tests {
.unwrap_or_else(|| {
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
"request-test",
"test-submission",
protocol::SubmissionDisposition::Started,
@@ -4038,17 +4039,11 @@ mod tests {
}
fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
WorkerExecutionResult::accepted(
WorkerExecutionOperation::Stop,
WorkerExecutionRunState::Stopped,
)
WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop)
}
fn cancel_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
WorkerExecutionResult::accepted(
WorkerExecutionOperation::Cancel,
WorkerExecutionRunState::Stopped,
)
WorkerExecutionResult::accepted(WorkerExecutionOperation::Cancel)
}
#[cfg(feature = "ws-server")]
@@ -4374,7 +4369,9 @@ mod tests {
.send_protocol_method_scoped(
&scope("workspace-a", "server-a"),
&workspace_b.worker_ref,
Method::Shutdown,
Method::Shutdown {
command: test_command(),
},
)
.unwrap_err();
assert!(matches!(
@@ -4722,11 +4719,10 @@ mod tests {
}
#[test]
fn create_worker_uses_committed_input_ack_run_state() {
fn create_worker_uses_started_submission_ack_for_initial_running_status() {
let (runtime, backend) = runtime_and_backend();
backend.set_dispatch_result(WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
"request-test",
"test-submission",
protocol::SubmissionDisposition::Started,
@@ -4736,7 +4732,7 @@ mod tests {
let detail = runtime.create_worker(request).unwrap();
assert_eq!(detail.status, WorkerStatus::Idle);
assert_eq!(detail.status, WorkerStatus::Running);
}
#[test]
@@ -4745,7 +4741,6 @@ mod tests {
backend.preserve_commit_ack_submission_id();
backend.set_dispatch_result(WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
"request-test",
"forged-submission",
protocol::SubmissionDisposition::Started,
@@ -4770,7 +4765,6 @@ mod tests {
let (runtime, backend) = runtime_and_backend();
backend.set_dispatch_result(WorkerExecutionResult::accepted(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Busy,
));
let mut request = task_request("missing initial input commit ack");
request.initial_input = Some(WorkerInput::user("start the ticket"));
@@ -4898,7 +4892,7 @@ mod tests {
context_window: 128,
context_tokens: 64,
},
status: protocol::WorkerStatus::Running,
state: protocol::WorkerStatus::Running.into(),
in_flight: protocol::InFlightSnapshot {
blocks: Vec::new(),
commands: Vec::new(),
@@ -4914,13 +4908,13 @@ mod tests {
protocol::Event::Snapshot {
session,
greeting,
status,
state,
..
} => {
assert_eq!(session.entries.len(), 1);
assert_eq!(session.entries[0].entry_id, "restored-log-entry");
assert_eq!(greeting.worker_name, "live-worker");
assert_eq!(status, protocol::WorkerStatus::Running);
assert_eq!(state.catalog_status(), protocol::WorkerStatus::Running);
}
other => panic!("expected snapshot, got {other:?}"),
}
@@ -4936,7 +4930,6 @@ mod tests {
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request
.working_directory
.as_ref()
@@ -4951,7 +4944,6 @@ mod tests {
) -> WorkerExecutionResult {
WorkerExecutionResult::accepted_submission(
WorkerExecutionOperation::Input,
WorkerExecutionRunState::Idle,
"request-test",
input
.submission_request_id
@@ -4993,7 +4985,12 @@ mod tests {
.unwrap();
runtime
.send_protocol_method(&detail.worker_ref, Method::Shutdown)
.send_protocol_method(
&detail.worker_ref,
Method::Shutdown {
command: test_command(),
},
)
.unwrap();
assert_eq!(
@@ -5009,7 +5006,12 @@ mod tests {
.create_worker(task_request("restore explicitly"))
.unwrap();
runtime
.send_protocol_method(&detail.worker_ref, Method::Shutdown)
.send_protocol_method(
&detail.worker_ref,
Method::Shutdown {
command: test_command(),
},
)
.unwrap();
assert!(matches!(
@@ -5027,7 +5029,7 @@ mod tests {
assert_eq!(*backend.run_generations.lock().unwrap(), vec![1, 2]);
assert_eq!(
runtime.worker_detail(&detail.worker_ref).unwrap().status,
WorkerStatus::Idle
WorkerStatus::Running
);
}
crates/worker-runtime/src/worker_backend.rs
+249 -166
View File
@@ -10,8 +10,8 @@
use std::collections::HashMap;
use std::future::Future;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock, mpsc};
use std::time::Duration;
use crate::auth::{
@@ -25,8 +25,8 @@ use crate::catalog::{
};
use crate::execution::{
WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation,
WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionRunState,
WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult,
WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionSpawnRequest,
WorkerExecutionSpawnResult,
};
use crate::identity::WorkerRef;
use crate::interaction::{WorkerInput, WorkerInputKind};
@@ -38,7 +38,26 @@ use crate::working_directory::{
WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer,
};
use async_trait::async_trait;
use protocol::{ErrorCode, Event, Method, Segment, WorkerStatus};
use protocol::{Event, Method, Segment, WorkerCommandEnvelope, WorkerStatus};
static NEXT_INTERNAL_COMMAND_ID: AtomicU64 = AtomicU64::new(1);
fn next_internal_command(
state: &RwLock<protocol::WorkerStateSnapshot>,
) -> Result<WorkerCommandEnvelope, String> {
let snapshot = state
.read()
.map_err(|_| "worker state lock is poisoned".to_string())?
.clone();
let floor = snapshot.last_command_id.saturating_add(1);
let command_id = NEXT_INTERNAL_COMMAND_ID
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
Some(current.max(floor).saturating_add(1))
})
.unwrap_or(floor)
.max(floor);
Ok(WorkerCommandEnvelope::for_snapshot(command_id, &snapshot))
}
use session_store::{CombinedStore, WorkerAggregateStore, WorkerSessionStore};
#[cfg(test)]
use session_store::{FsStore, FsWorkerStore};
@@ -172,7 +191,7 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider {
},
display_name: grant.worker_id.clone(),
relation: "granted_peer".to_string(),
status: format!("{:?}", state.get_status()).to_lowercase(),
status: format!("{:?}", state.catalog_status()).to_lowercase(),
});
}
subjects.sort_by(|left, right| left.subject.cmp(&right.subject));
@@ -1174,10 +1193,12 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
}
}
#[derive(Clone)]
struct RuntimeWorkerExecution {
handle: WorkerHandle,
shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
busy: Arc<AtomicBool>,
worker_state: Arc<RwLock<protocol::WorkerStateSnapshot>>,
workspace_client: Option<Arc<dyn WorkspaceClient>>,
}
@@ -1276,6 +1297,7 @@ where
(
WorkerHandle,
Arc<AtomicBool>,
Arc<RwLock<protocol::WorkerStateSnapshot>>,
Option<Arc<dyn WorkspaceClient>>,
),
WorkerExecutionResult,
@@ -1302,6 +1324,7 @@ where
(
execution.handle.clone(),
execution.busy.clone(),
execution.worker_state.clone(),
execution.workspace_client.clone(),
)
})
@@ -1318,7 +1341,6 @@ where
operation: WorkerExecutionOperation,
worker: WorkerHandle,
method: Method,
accepted_run_state: WorkerExecutionRunState,
) -> WorkerExecutionResult {
self.run_on_adapter_runtime(async move {
worker
@@ -1326,7 +1348,7 @@ where
.await
.map_err(|err| format!("failed to send Worker method: {err}"))
})
.map(|_| WorkerExecutionResult::accepted(operation, accepted_run_state))
.map(|_| WorkerExecutionResult::accepted(operation))
.unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message))
}
@@ -1336,7 +1358,6 @@ where
worker: WorkerHandle,
method: Method,
submission_request_id: String,
accepted_run_state: WorkerExecutionRunState,
) -> WorkerExecutionResult {
let request_id = submission_request_id.clone();
self.run_on_adapter_runtime(async move {
@@ -1395,7 +1416,6 @@ where
.map(|(submission_id, disposition)| {
WorkerExecutionResult::accepted_submission(
operation,
accepted_run_state,
submission_request_id,
submission_id,
disposition,
@@ -1415,38 +1435,45 @@ where
workspace_client: Option<Arc<dyn WorkspaceClient>>,
) -> WorkerExecutionSpawnResult {
let busy = Arc::new(AtomicBool::new(false));
let worker_state = Arc::new(RwLock::new(handle.shared_state.snapshot()));
#[cfg(feature = "ws-server")]
{
let streams = subscribe_worker_protocol_session(&handle);
let mut events = streams.events;
let mut entry_events = streams.log_entries;
let bridge_busy = busy.clone();
let bridge_worker_state = worker_state.clone();
if let Err(message) = self.spawn_on_adapter_runtime(async move {
loop {
tokio::select! {
event = events.recv() => {
match event {
Ok(event) => {
let next_busy = match &event {
Event::InvokeStart { .. }
| Event::Status {
status: WorkerStatus::Running,
} => Some(true),
Event::RunEnd { .. }
| Event::Error {
code: ErrorCode::NotPaused,
..
let next_state = match &event {
Event::WorkerState { snapshot }
| Event::Snapshot { state: snapshot, .. } => {
Some(snapshot.clone())
}
| Event::Status {
status:
WorkerStatus::Idle
| WorkerStatus::Paused
| WorkerStatus::Stopped,
Event::CommandAcknowledged { acknowledgement } => {
Some(acknowledgement.state.clone())
}
| Event::Shutdown => Some(false),
_ => None,
};
let next_busy = next_state
.as_ref()
.map(worker_state_is_executing)
.or_else(|| matches!(event, Event::Shutdown).then_some(false));
let _ = bridge_context.publish_protocol_event(event);
if let Some(next_state) = next_state {
if let Ok(mut current) = bridge_worker_state.write() {
if next_state.execution_generation > current.execution_generation
|| (next_state.execution_generation == current.execution_generation
&& next_state.revision >= current.revision)
{
*current = next_state;
}
}
}
if let Some(next_busy) = next_busy {
bridge_busy.store(next_busy, Ordering::SeqCst);
}
@@ -1494,13 +1521,13 @@ where
handle,
shutdown,
busy,
worker_state,
workspace_client,
},
);
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: working_directory.map(|binding| binding.status()),
}
}
@@ -1516,6 +1543,17 @@ impl<F> Drop for WorkerRuntimeExecutionBackend<F> {
}
}
fn worker_state_is_executing(snapshot: &protocol::WorkerStateSnapshot) -> bool {
matches!(
snapshot.state,
protocol::WorkerState::Busy(protocol::WorkerBusyState::Run(
protocol::WorkerRunState::Running
| protocol::WorkerRunState::Pausing
| protocol::WorkerRunState::Cancelling
)) | protocol::WorkerState::Busy(protocol::WorkerBusyState::Maintenance(_))
)
}
fn method_starts_turn(method: &Method) -> bool {
matches!(
method,
@@ -1523,41 +1561,17 @@ fn method_starts_turn(method: &Method) -> bool {
| Method::SubmitTracked { .. }
| Method::Notify { auto_run: true, .. }
| Method::NotifyTracked { auto_run: true, .. }
| Method::Resume
| Method::Compact
| Method::Resume { .. }
)
}
fn method_can_start_turn_from_status(method: &Method, status: WorkerStatus) -> bool {
match method {
Method::Resume => matches!(status, WorkerStatus::Idle | WorkerStatus::Paused),
Method::Resume { .. } => matches!(status, WorkerStatus::Idle | WorkerStatus::Paused),
_ => status == WorkerStatus::Idle,
}
}
fn accepted_notify_run_state(status: WorkerStatus, auto_run: bool) -> WorkerExecutionRunState {
match status {
WorkerStatus::Running => WorkerExecutionRunState::Busy,
WorkerStatus::Idle if auto_run => WorkerExecutionRunState::Busy,
WorkerStatus::Idle | WorkerStatus::Paused | WorkerStatus::Stopped => {
WorkerExecutionRunState::Idle
}
}
}
fn accepted_run_state_for_method(method: &Method) -> WorkerExecutionRunState {
match method {
Method::Submit { .. }
| Method::SubmitTracked { .. }
| Method::Notify { auto_run: true, .. }
| Method::NotifyTracked { auto_run: true, .. }
| Method::Resume
| Method::Compact => WorkerExecutionRunState::Busy,
Method::Shutdown => WorkerExecutionRunState::Stopped,
_ => WorkerExecutionRunState::Idle,
}
}
impl<F> WorkerExecutionBackend for WorkerRuntimeExecutionBackend<F>
where
F: RuntimeWorkerFactory,
@@ -1883,7 +1897,7 @@ where
handle: &WorkerExecutionHandle,
input: WorkerInput,
) -> WorkerExecutionResult {
let (worker, busy, _workspace_client) = match self.get_execution(handle) {
let (worker, busy, worker_state, _workspace_client) = match self.get_execution(handle) {
Ok(execution) => execution,
Err(mut result) => {
result.operation = WorkerExecutionOperation::Input;
@@ -1892,8 +1906,7 @@ where
};
if input.kind == WorkerInputKind::Notify {
let status = worker.shared_state.get_status();
let accepted_run_state = accepted_notify_run_state(status, true);
let status = worker.shared_state.catalog_status();
let claimed_here = status == WorkerStatus::Idle
&& busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
@@ -1912,7 +1925,6 @@ where
operation_id: notification_request_id,
},
},
accepted_run_state,
);
if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted
{
@@ -1921,8 +1933,22 @@ where
return result;
}
if input.kind == WorkerInputKind::Compact {
let command = match next_internal_command(&worker_state) {
Ok(command) => command,
Err(error) => {
return WorkerExecutionResult::errored(WorkerExecutionOperation::Input, error);
}
};
return self.send_method(
WorkerExecutionOperation::Input,
worker,
Method::Compact { command },
);
}
let is_user_submit = input.kind == WorkerInputKind::User;
let status = worker.shared_state.get_status();
let status = worker.shared_state.catalog_status();
let claimed_here = status == WorkerStatus::Idle
&& busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
@@ -1962,7 +1988,7 @@ where
WorkerInputKind::Notify => {
unreachable!("Notify input is dispatched before the turn-start busy guard")
}
WorkerInputKind::Compact => (Method::Compact, None),
WorkerInputKind::Compact => unreachable!("compact input is dispatched above"),
WorkerInputKind::ListRewindTargets => (Method::ListRewindTargets, None),
WorkerInputKind::RegisterPeer => (
Method::RegisterPeer {
@@ -1971,15 +1997,6 @@ where
None,
),
};
let accepted_run_state = match method {
Method::Submit { .. }
| Method::SubmitTracked { .. }
| Method::Notify { .. }
| Method::NotifyTracked { .. }
| Method::Compact => WorkerExecutionRunState::Busy,
_ => WorkerExecutionRunState::Idle,
};
let accepted_is_idle = accepted_run_state == WorkerExecutionRunState::Idle;
let waits_for_submission_acceptance = submission_request_id.is_some();
let result = if waits_for_submission_acceptance {
@@ -1988,20 +2005,11 @@ where
worker,
method,
submission_request_id.expect("Submit must have a submission request id"),
accepted_run_state,
)
} else {
self.send_method(
WorkerExecutionOperation::Input,
worker,
method,
accepted_run_state,
)
self.send_method(WorkerExecutionOperation::Input, worker, method)
};
if accepted_is_idle
|| (claimed_here
&& result.outcome != crate::execution::WorkerExecutionOutcome::Accepted)
{
if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted {
busy.store(false, Ordering::SeqCst);
}
result
@@ -2015,7 +2023,7 @@ where
content: &[u8],
context: Option<&session_store::UploadedFileUploadContext>,
) -> Result<protocol::UploadedFileRef, WorkerExecutionResult> {
let (worker, _, _) = self.get_execution(handle).map_err(|mut result| {
let (worker, _, _, _) = self.get_execution(handle).map_err(|mut result| {
result.operation = WorkerExecutionOperation::UploadFile;
result
})?;
@@ -2038,7 +2046,7 @@ where
handle: &WorkerExecutionHandle,
artifact_id: &str,
) -> WorkerExecutionResult {
let (worker, _, _) = match self.get_execution(handle) {
let (worker, _, _, _) = match self.get_execution(handle) {
Ok(execution) => execution,
Err(mut result) => {
result.operation = WorkerExecutionOperation::DeleteUploadedFile;
@@ -2046,10 +2054,7 @@ where
}
};
match worker.delete_uploaded_file(artifact_id) {
Ok(_) => WorkerExecutionResult::accepted(
WorkerExecutionOperation::DeleteUploadedFile,
WorkerExecutionRunState::Idle,
),
Ok(_) => WorkerExecutionResult::accepted(WorkerExecutionOperation::DeleteUploadedFile),
Err(error) => WorkerExecutionResult::rejected(
WorkerExecutionOperation::DeleteUploadedFile,
format!("uploaded_file_delete_rejected: {error}"),
@@ -2062,7 +2067,7 @@ where
handle: &WorkerExecutionHandle,
method: Method,
) -> WorkerExecutionResult {
let (worker, busy, _workspace_client) = match self.get_execution(handle) {
let (worker, busy, _worker_state, _workspace_client) = match self.get_execution(handle) {
Ok(execution) => execution,
Err(mut result) => {
result.operation = WorkerExecutionOperation::ProtocolMethod;
@@ -2076,19 +2081,13 @@ where
}
_ => None,
} {
let status = worker.shared_state.get_status();
let accepted_run_state = accepted_notify_run_state(status, auto_run);
let status = worker.shared_state.catalog_status();
let claimed_here = status == WorkerStatus::Idle
&& auto_run
&& busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_ok();
let result = self.send_method(
WorkerExecutionOperation::ProtocolMethod,
worker,
method,
accepted_run_state,
);
let result = self.send_method(WorkerExecutionOperation::ProtocolMethod, worker, method);
if claimed_here && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted
{
busy.store(false, Ordering::SeqCst);
@@ -2098,7 +2097,7 @@ where
let starts_turn = method_starts_turn(&method);
if starts_turn
&& (!method_can_start_turn_from_status(&method, worker.shared_state.get_status())
&& (!method_can_start_turn_from_status(&method, worker.shared_state.catalog_status())
|| busy
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_err())
@@ -2109,17 +2108,8 @@ where
);
}
let accepted_run_state = accepted_run_state_for_method(&method);
let accepted_is_idle = accepted_run_state == WorkerExecutionRunState::Idle;
let result = self.send_method(
WorkerExecutionOperation::ProtocolMethod,
worker,
method,
accepted_run_state,
);
if (starts_turn && accepted_is_idle)
|| (starts_turn && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted)
{
let result = self.send_method(WorkerExecutionOperation::ProtocolMethod, worker, method);
if starts_turn && result.outcome != crate::execution::WorkerExecutionOutcome::Accepted {
busy.store(false, Ordering::SeqCst);
}
result
@@ -2137,7 +2127,7 @@ where
);
}
let execution = match self.workers.lock() {
Ok(mut workers) => workers.remove(handle.worker_ref()),
Ok(workers) => workers.get(handle.worker_ref()).cloned(),
Err(_) => {
return WorkerExecutionResult::errored(
WorkerExecutionOperation::Stop,
@@ -2153,48 +2143,73 @@ where
};
let artifact_cleanup = execution.handle.clone();
let shutdown = execution.shutdown.clone();
let command = match next_internal_command(&execution.worker_state) {
Ok(command) => command,
Err(error) => {
return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, error);
}
};
let result = self.send_method(
WorkerExecutionOperation::Stop,
execution.handle,
Method::Shutdown,
WorkerExecutionRunState::Stopped,
execution.handle.clone(),
Method::Shutdown { command },
);
if result.outcome != crate::execution::WorkerExecutionOutcome::Accepted {
return result;
}
match self.run_on_adapter_runtime(async move {
let receiver = shutdown.lock().await.take();
if let Some(receiver) = receiver {
receiver
.await
.map_err(|_| "Worker shutdown completion channel closed".to_string())?;
let shutdown_wait = self.run_on_adapter_runtime(async move {
let mut guard = shutdown.lock().await;
let Some(mut receiver) = guard.take() else {
return Ok(());
};
match tokio::time::timeout(Duration::from_secs(5), &mut receiver).await {
Ok(Ok(())) => Ok(()),
Ok(Err(_)) => Err("Worker shutdown completion channel closed".to_string()),
Err(_) => {
*guard = Some(receiver);
Err("Worker shutdown confirmation timed out; stop remains retryable".into())
}
}
Ok(())
}) {
Ok(()) => match artifact_cleanup.delete_uncommitted_uploaded_files() {
Ok(_) => result,
Err(error) => WorkerExecutionResult::errored(
WorkerExecutionOperation::Stop,
format!("uploaded_file_cleanup_failed: {error}"),
),
},
Err(message) => WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message),
});
if let Err(message) = shutdown_wait {
return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message);
}
if let Err(error) = artifact_cleanup.delete_uncommitted_uploaded_files() {
return WorkerExecutionResult::errored(
WorkerExecutionOperation::Stop,
format!("uploaded_file_cleanup_failed: {error}"),
);
}
match self.workers.lock() {
Ok(mut workers) => {
workers.remove(handle.worker_ref());
result
}
Err(_) => WorkerExecutionResult::errored(
WorkerExecutionOperation::Stop,
"worker adapter registry lock is poisoned after shutdown",
),
}
}
fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
let (worker, _busy, _workspace_client) = match self.get_execution(handle) {
let (worker, _busy, worker_state, _workspace_client) = match self.get_execution(handle) {
Ok(execution) => execution,
Err(mut result) => {
result.operation = WorkerExecutionOperation::Cancel;
return result;
}
};
let command = match next_internal_command(&worker_state) {
Ok(command) => command,
Err(error) => {
return WorkerExecutionResult::errored(WorkerExecutionOperation::Cancel, error);
}
};
self.send_method(
WorkerExecutionOperation::Cancel,
worker,
Method::Cancel,
WorkerExecutionRunState::Idle,
Method::Cancel { command },
)
}
@@ -2259,6 +2274,29 @@ mod tests {
use manifest::{Scope, WorkerManifest};
use session_store::{LogEntry, WorkerMetadataStore};
fn test_command() -> WorkerCommandEnvelope {
WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 1,
expected_worker_state_revision: 0,
}
}
fn adapter_command(
backend: &WorkerRuntimeExecutionBackend<MockFactory>,
worker_ref: &WorkerRef,
) -> WorkerCommandEnvelope {
let workers = backend.workers.lock().unwrap();
let state = workers
.get(worker_ref)
.expect("worker execution")
.worker_state
.read()
.unwrap()
.clone();
WorkerCommandEnvelope::for_snapshot(state.last_command_id.saturating_add(1), &state)
}
#[test]
fn workspace_prompt_projection_notification_advances_shared_cache() {
let cache = WorkspacePromptProjectionCache::default();
@@ -2406,41 +2444,39 @@ mod tests {
}
#[test]
fn notify_run_state_allows_running_worker_inbox_delivery() {
assert_eq!(
accepted_notify_run_state(WorkerStatus::Running, true),
WorkerExecutionRunState::Busy
);
assert_eq!(
accepted_notify_run_state(WorkerStatus::Idle, true),
WorkerExecutionRunState::Busy
);
assert_eq!(
accepted_notify_run_state(WorkerStatus::Idle, false),
WorkerExecutionRunState::Idle
);
assert_eq!(
accepted_notify_run_state(WorkerStatus::Paused, true),
WorkerExecutionRunState::Idle
);
fn compact_is_maintenance_not_a_turn_start() {
assert!(!method_starts_turn(&Method::Compact {
command: test_command(),
}));
assert!(method_starts_turn(&Method::Resume {
command: test_command(),
}));
}
#[test]
fn resume_turn_claim_accepts_paused_and_idle_but_not_running_status() {
assert!(method_can_start_turn_from_status(
&Method::Resume,
&Method::Resume {
command: test_command()
},
WorkerStatus::Paused
));
assert!(method_can_start_turn_from_status(
&Method::Resume,
&Method::Resume {
command: test_command()
},
WorkerStatus::Idle
));
assert!(!method_can_start_turn_from_status(
&Method::Resume,
&Method::Resume {
command: test_command()
},
WorkerStatus::Running
));
assert!(!method_can_start_turn_from_status(
&Method::Compact,
&Method::Compact {
command: test_command()
},
WorkerStatus::Paused
));
}
@@ -2656,19 +2692,22 @@ mod tests {
let observed = {
let workers = backend.workers.lock().unwrap();
let execution = workers.get(worker_ref).expect("live Worker execution");
let projected = execution.worker_state.read().unwrap().catalog_status();
(
execution.handle.shared_state.get_status(),
execution.handle.shared_state.catalog_status(),
projected,
execution.busy.load(Ordering::SeqCst),
)
};
if observed == (expected_status, expected_busy) {
if observed == (expected_status, expected_status, expected_busy) {
return;
}
assert!(
std::time::Instant::now() < deadline,
"timed out waiting for adapter state {expected_status:?}, busy={expected_busy}; last observed status={:?}, busy={}",
"timed out waiting for adapter state {expected_status:?}, busy={expected_busy}; last observed controller={:?}, projected={:?}, busy={}",
observed.0,
observed.1,
observed.2,
);
std::thread::sleep(Duration::from_millis(10));
}
@@ -3169,13 +3208,19 @@ mod tests {
.expect("in-process restore must not bind the overlong Unix socket path");
assert_eq!(
controller.handle.shared_state.get_status(),
controller.handle.shared_state.catalog_status(),
WorkerStatus::Idle
);
assert!(!socket_path.exists());
assert!(run_dir.join("worker.out.log").is_file());
assert!(run_dir.join("worker.err.log").is_file());
controller.handle.send(Method::Shutdown).await.unwrap();
controller
.handle
.send(Method::Shutdown {
command: test_command(),
})
.await
.unwrap();
if let Some(receiver) = controller.shutdown.lock().await.take() {
receiver.await.unwrap();
}
@@ -3289,7 +3334,9 @@ mod tests {
backend
.run_on_adapter_runtime(async move {
handle
.send(Method::Shutdown)
.send(Method::Shutdown {
command: test_command(),
})
.await
.map_err(|error| error.to_string())?;
if let Some(receiver) = shutdown.lock().await.take() {
@@ -3614,6 +3661,7 @@ mod tests {
#[test]
#[cfg(feature = "ws-server")]
#[serial_test::serial(worker_allocation)]
fn adapter_resumes_paused_turn_once_and_preserves_idle_not_paused_error() {
let hanging_events = || simple_text_events().into_iter().take(2).collect::<Vec<_>>();
let client = MockClient::sequential(vec![
@@ -3649,7 +3697,12 @@ mod tests {
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running, true);
let running_resume = runtime
.send_protocol_method(&detail.worker_ref, Method::Resume)
.send_protocol_method(
&detail.worker_ref,
Method::Resume {
command: adapter_command(&backend, &detail.worker_ref),
},
)
.expect_err("Resume while Running must be rejected");
assert!(
running_resume
@@ -3659,17 +3712,32 @@ mod tests {
);
runtime
.send_protocol_method(&detail.worker_ref, Method::Pause)
.send_protocol_method(
&detail.worker_ref,
Method::Pause {
command: adapter_command(&backend, &detail.worker_ref),
},
)
.expect("pause initial turn");
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused, false);
runtime
.send_protocol_method(&detail.worker_ref, Method::Resume)
.send_protocol_method(
&detail.worker_ref,
Method::Resume {
command: adapter_command(&backend, &detail.worker_ref),
},
)
.expect("resume paused turn");
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Running, true);
let duplicate_resume = runtime
.send_protocol_method(&detail.worker_ref, Method::Resume)
.send_protocol_method(
&detail.worker_ref,
Method::Resume {
command: adapter_command(&backend, &detail.worker_ref),
},
)
.expect_err("duplicate Resume must be rejected");
assert!(
duplicate_resume
@@ -3679,17 +3747,32 @@ mod tests {
);
runtime
.send_protocol_method(&detail.worker_ref, Method::Pause)
.send_protocol_method(
&detail.worker_ref,
Method::Pause {
command: adapter_command(&backend, &detail.worker_ref),
},
)
.expect("pause resumed turn");
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Paused, false);
runtime
.send_protocol_method(&detail.worker_ref, Method::Resume)
.send_protocol_method(
&detail.worker_ref,
Method::Resume {
command: adapter_command(&backend, &detail.worker_ref),
},
)
.expect("resume paused turn a second time");
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle, false);
assert_eq!(call_count.load(Ordering::SeqCst), 3);
runtime
.send_protocol_method(&detail.worker_ref, Method::Resume)
.send_protocol_method(
&detail.worker_ref,
Method::Resume {
command: adapter_command(&backend, &detail.worker_ref),
},
)
.expect("Idle Resume preserves controller NotPaused semantics");
wait_for_adapter_state(&backend, &detail.worker_ref, WorkerStatus::Idle, false);
let events = runtime
@@ -3698,10 +3781,10 @@ mod tests {
assert!(events.iter().any(|event| {
matches!(
&event.payload,
Event::Error {
code: protocol::ErrorCode::NotPaused,
..
}
Event::CommandAcknowledged { acknowledgement }
if acknowledgement.command == protocol::WorkerCommandKind::Resume
&& acknowledgement.disposition
== protocol::WorkerCommandDisposition::InvalidState
)
}));
assert_eq!(call_count.load(Ordering::SeqCst), 3);