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

refactor: make worker state snapshots authoritative

Browse Source
This commit is contained in:
Hare
2026-09-17 02:02:51 +09:00
parent d6bef5d1c7
commit 10ebac142e
31 changed files with 473 additions and 1132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+11 -66
View File
@@ -190,12 +190,6 @@ fn command_admission_disposition(
Err(WorkerCommandDisposition::StaleCommandId)
}
WorkerCommandAdmission::Conflict => Err(WorkerCommandDisposition::Conflict),
WorkerCommandAdmission::ExecutionGenerationMismatch => {
Err(WorkerCommandDisposition::StaleExecutionGeneration)
}
WorkerCommandAdmission::StateRevisionMismatch => {
Err(WorkerCommandDisposition::StaleWorkerStateRevision)
}
}
}
@@ -204,14 +198,14 @@ fn validate_command(
kind: WorkerCommandKind,
shared_state: &WorkerSharedState,
) -> Result<(), WorkerCommandDisposition> {
command_admission_disposition(shared_state.admit_command(envelope, kind, true))
command_admission_disposition(shared_state.admit_command(envelope, kind))
}
fn validate_shutdown_command(
envelope: WorkerCommandEnvelope,
shared_state: &WorkerSharedState,
) -> Result<(), WorkerCommandDisposition> {
match shared_state.admit_command(envelope, WorkerCommandKind::Shutdown, false) {
match shared_state.admit_command(envelope, WorkerCommandKind::Shutdown) {
WorkerCommandAdmission::Accepted | WorkerCommandAdmission::Retry => Ok(()),
admission => command_admission_disposition(admission),
}
@@ -793,19 +787,11 @@ impl WorkerController {
.await
.map_err(|error| std::io::Error::other(error.to_string()))?;
let greeting = build_greeting(&worker);
let execution_generation = runtime_dir
.path()
.file_name()
.and_then(|name| name.to_str())
.and_then(|name| name.parse::<u64>().ok())
.filter(|generation| *generation > 0)
.unwrap_or(1);
let shared_state = Arc::new(WorkerSharedState::new_with_generation(
let shared_state = Arc::new(WorkerSharedState::new(
worker.manifest().worker.name.clone(),
worker.segment_id(),
manifest_toml.clone(),
greeting,
execution_generation,
));
if let Some(fs_for_view) = fs_for_view {
shared_state.set_fs_view(crate::fs_view::WorkerFsView::new(fs_for_view));
@@ -3436,7 +3422,7 @@ mod tests {
.transition(WorkerState::Busy(WorkerBusyState::Run(
WorkerRunState::Running,
)));
let command = WorkerCommandEnvelope::for_snapshot(1, &env.shared_state.snapshot());
let command = WorkerCommandEnvelope::new(1);
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(10)).await;
method_tx
@@ -3719,7 +3705,7 @@ mod tests {
.transition(WorkerState::Busy(WorkerBusyState::Run(
WorkerRunState::Running,
)));
let command = WorkerCommandEnvelope::for_snapshot(1, &env.shared_state.snapshot());
let command = WorkerCommandEnvelope::new(1);
env._method_tx
.send(Method::Compact { command })
.await
@@ -3770,8 +3756,8 @@ mod tests {
}
#[test]
fn command_admission_rejects_stale_generation_revision_and_order() {
let shared = WorkerSharedState::new_with_generation(
fn command_admission_rejects_stale_ids_and_reuse_conflicts() {
let shared = WorkerSharedState::new(
"worker".into(),
session_store::new_segment_id(),
String::new(),
@@ -3785,39 +3771,10 @@ mod tests {
context_window: 1,
context_tokens: 0,
},
9,
);
assert_eq!(
validate_command(
WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 8,
expected_worker_state_revision: 0,
},
WorkerCommandKind::Pause,
&shared,
),
Err(WorkerCommandDisposition::StaleExecutionGeneration)
);
assert_eq!(
validate_command(
WorkerCommandEnvelope {
command_id: 2,
expected_execution_generation: 9,
expected_worker_state_revision: 1,
},
WorkerCommandKind::Pause,
&shared,
),
Err(WorkerCommandDisposition::StaleWorkerStateRevision)
);
assert!(
validate_command(
WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 9,
expected_worker_state_revision: 0,
},
WorkerCommandEnvelope { command_id: 1 },
WorkerCommandKind::Pause,
&shared,
)
@@ -3825,11 +3782,7 @@ mod tests {
);
assert_eq!(
validate_command(
WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 9,
expected_worker_state_revision: 0,
},
WorkerCommandEnvelope { command_id: 1 },
WorkerCommandKind::Pause,
&shared,
),
@@ -3837,11 +3790,7 @@ mod tests {
);
assert_eq!(
validate_command(
WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 9,
expected_worker_state_revision: 0,
},
WorkerCommandEnvelope { command_id: 1 },
WorkerCommandKind::Cancel,
&shared,
),
@@ -3849,11 +3798,7 @@ mod tests {
);
assert!(
validate_command(
WorkerCommandEnvelope {
command_id: 2,
expected_execution_generation: 9,
expected_worker_state_revision: 1,
},
WorkerCommandEnvelope { command_id: 2 },
WorkerCommandKind::Pause,
&shared,
)
crates/worker/src/internal_worker.rs
+3 -21
View File
@@ -296,7 +296,6 @@ impl InternalWorkerSessionStatus {
fn send_internal_worker_state(
event_tx: &broadcast::Sender<Event>,
state_revision: &std::sync::atomic::AtomicU64,
status: InternalWorkerSessionStatus,
) {
let state = match status {
@@ -313,13 +312,8 @@ fn send_internal_worker_state(
protocol::WorkerBusyState::Run(protocol::WorkerRunState::Cancelling),
),
};
let revision = state_revision
.fetch_add(1, std::sync::atomic::Ordering::AcqRel)
.saturating_add(1);
let _ = event_tx.send(Event::WorkerState {
snapshot: protocol::WorkerStateSnapshot {
execution_generation: 1,
revision,
last_command_id: 0,
state,
},
@@ -382,7 +376,6 @@ pub(crate) struct InternalWorkerSessionSnapshot {
pub(crate) struct InternalWorkerSessionHandle {
command_tx: tokio::sync::mpsc::Sender<InternalWorkerSessionCommand>,
status: Arc<std::sync::atomic::AtomicU8>,
state_revision: Arc<std::sync::atomic::AtomicU64>,
store: EphemeralSessionStore,
session_id: SessionId,
segment_id: SegmentId,
@@ -433,7 +426,7 @@ impl InternalWorkerSessionHandle {
}
fn emit_worker_state(&self, status: InternalWorkerSessionStatus) {
send_internal_worker_state(&self.event_tx, &self.state_revision, status);
send_internal_worker_state(&self.event_tx, status);
}
pub(crate) fn protocol_snapshot(&self) -> InternalWorkerSessionSnapshot {
@@ -803,13 +796,11 @@ pub(crate) async fn prepare_internal_worker_session(
let status = Arc::new(std::sync::atomic::AtomicU8::new(
InternalWorkerSessionStatus::Idle.encode(),
));
let state_revision = Arc::new(std::sync::atomic::AtomicU64::new(0));
let state_changed = Arc::new(tokio::sync::Notify::new());
let last_error = Arc::new(Mutex::new(None));
let handle = InternalWorkerSessionHandle {
command_tx,
status: status.clone(),
state_revision: state_revision.clone(),
store,
session_id,
segment_id,
@@ -845,11 +836,7 @@ pub(crate) async fn prepare_internal_worker_session(
message,
});
}
send_internal_worker_state(
&event_tx,
&state_revision,
turn_status,
);
send_internal_worker_state(&event_tx, turn_status);
if let Some(callback) = &on_turn_end {
callback(turn_status);
}
@@ -891,11 +878,7 @@ pub(crate) async fn prepare_internal_worker_session(
InternalWorkerSessionStatus::Stopped.encode(),
std::sync::atomic::Ordering::Release,
);
send_internal_worker_state(
&event_tx,
&state_revision,
InternalWorkerSessionStatus::Stopped,
);
send_internal_worker_state(&event_tx, InternalWorkerSessionStatus::Stopped);
let _ = event_tx.send(Event::Shutdown);
state_changed.notify_waiters();
if let Some(done) = stop_done {
@@ -1147,7 +1130,6 @@ pub(crate) fn test_internal_worker_session(
status: Arc::new(std::sync::atomic::AtomicU8::new(
InternalWorkerSessionStatus::Idle.encode(),
)),
state_revision: Arc::new(std::sync::atomic::AtomicU64::new(0)),
store,
session_id,
segment_id,
crates/worker/src/shared_state.rs
+14 -45
View File
@@ -28,8 +28,6 @@ pub(crate) enum WorkerCommandAdmission {
Retry,
Conflict,
StaleCommandId,
ExecutionGenerationMismatch,
StateRevisionMismatch,
}
/// Shared state between WorkerController and runtime directory.
@@ -59,23 +57,13 @@ impl WorkerSharedState {
segment_id: SegmentId,
manifest_toml: String,
greeting: protocol::Greeting,
) -> Self {
Self::new_with_generation(worker_name, segment_id, manifest_toml, greeting, 1)
}
pub fn new_with_generation(
worker_name: String,
segment_id: SegmentId,
manifest_toml: String,
greeting: protocol::Greeting,
execution_generation: u64,
) -> Self {
Self {
worker_name,
segment_id,
manifest_toml,
greeting,
state: RwLock::new(WorkerStateSnapshot::initial(execution_generation)),
state: RwLock::new(WorkerStateSnapshot::initial()),
accepted_commands: RwLock::new(VecDeque::new()),
fs_view: OnceLock::new(),
flow_transition_enabled: AtomicBool::new(false),
@@ -108,7 +96,6 @@ impl WorkerSharedState {
.write()
.expect("worker state lock poisoned; refusing an inferred fallback state");
if snapshot.state != state {
snapshot.revision = snapshot.revision.saturating_add(1);
snapshot.state = state;
}
snapshot.clone()
@@ -118,7 +105,6 @@ impl WorkerSharedState {
&self,
envelope: WorkerCommandEnvelope,
kind: WorkerCommandKind,
require_state_revision: bool,
) -> WorkerCommandAdmission {
let mut snapshot = self
.state
@@ -138,18 +124,11 @@ impl WorkerSharedState {
WorkerCommandAdmission::Conflict
};
}
if envelope.expected_execution_generation != snapshot.execution_generation {
return WorkerCommandAdmission::ExecutionGenerationMismatch;
}
if require_state_revision && envelope.expected_worker_state_revision != snapshot.revision {
return WorkerCommandAdmission::StateRevisionMismatch;
}
if envelope.command_id <= snapshot.last_command_id {
return WorkerCommandAdmission::StaleCommandId;
}
snapshot.last_command_id = envelope.command_id;
snapshot.revision = snapshot.revision.saturating_add(1);
accepted.push_back(AcceptedWorkerCommand {
envelope,
kind,
@@ -248,12 +227,11 @@ mod tests {
use super::*;
fn test_state() -> WorkerSharedState {
WorkerSharedState::new_with_generation(
WorkerSharedState::new(
"test-worker".into(),
session_store::new_segment_id(),
"[engine]\nname = \"test-worker\"".into(),
test_greeting(),
7,
)
}
@@ -273,43 +251,34 @@ mod tests {
#[test]
fn initial_snapshot_is_idle() {
let state = test_state();
assert_eq!(state.snapshot(), WorkerStateSnapshot::initial(7));
assert_eq!(state.snapshot(), WorkerStateSnapshot::initial());
assert_eq!(state.catalog_status(), WorkerStatus::Idle);
}
#[test]
fn transitions_increment_revision_only_when_state_changes() {
fn transitions_publish_full_state() {
let state = test_state();
let running = WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Running));
let snapshot = state.transition(running.clone());
assert_eq!(snapshot.revision, 1);
assert_eq!(snapshot.state, running);
assert_eq!(state.transition(running).revision, 1);
assert_eq!(state.transition(running.clone()).state, running);
let paused = WorkerState::Busy(WorkerBusyState::Run(WorkerRunState::Paused));
let snapshot = state.transition(paused.clone());
assert_eq!(snapshot.revision, 2);
assert_eq!(snapshot.state, paused);
assert_eq!(snapshot.last_command_id, 0);
assert_eq!(state.catalog_status(), WorkerStatus::Paused);
}
#[test]
fn accepted_command_identity_advances_revision_and_detects_reuse_conflicts() {
fn accepted_command_identity_advances_last_id_and_detects_reuse_conflicts() {
let state = test_state();
let envelope = WorkerCommandEnvelope {
command_id: 9,
expected_execution_generation: 7,
expected_worker_state_revision: 0,
};
let envelope = WorkerCommandEnvelope { command_id: 9 };
assert_eq!(
state.admit_command(envelope, WorkerCommandKind::Pause, true),
state.admit_command(envelope, WorkerCommandKind::Pause),
WorkerCommandAdmission::Accepted
);
assert_eq!(
state.snapshot(),
WorkerStateSnapshot {
execution_generation: 7,
revision: 1,
last_command_id: 9,
state: WorkerState::Idle,
}
@@ -325,14 +294,14 @@ mod tests {
Some(Some(WorkerCommandDisposition::Accepted))
);
assert_eq!(
state.admit_command(envelope, WorkerCommandKind::Pause, true),
state.admit_command(envelope, WorkerCommandKind::Pause),
WorkerCommandAdmission::Retry
);
assert_eq!(
state.admit_command(envelope, WorkerCommandKind::Cancel, true),
state.admit_command(envelope, WorkerCommandKind::Cancel),
WorkerCommandAdmission::Conflict
);
assert_eq!(state.snapshot().revision, 1);
assert_eq!(state.snapshot().last_command_id, 9);
}
#[test]
@@ -343,8 +312,8 @@ mod tests {
)));
let parsed: serde_json::Value = serde_json::from_str(&state.status_json()).unwrap();
assert_eq!(parsed["state"], "running");
assert_eq!(parsed["worker_state"]["execution_generation"], 7);
assert_eq!(parsed["worker_state"]["revision"], 1);
assert!(parsed["worker_state"].get("execution_generation").is_none());
assert!(parsed["worker_state"].get("revision").is_none());
assert_eq!(parsed["worker_state"]["state"]["kind"], "busy");
assert_eq!(parsed["worker_name"], "test-worker");
assert!(parsed["segment_id"].is_string());
crates/worker/src/spawn/comm_tools.rs
+1 -4
View File
@@ -138,10 +138,7 @@ mod tests {
);
let method = Method::Shutdown {
command: protocol::WorkerCommandEnvelope::for_snapshot(
1,
&protocol::WorkerStateSnapshot::initial(1),
),
command: protocol::WorkerCommandEnvelope::new(1),
};
connect_and_send(&socket, &method).await.unwrap();
crates/worker/tests/compact_events_test.rs
+6 -7
View File
@@ -1135,7 +1135,7 @@ async fn manual_compact_cancel_clears_progress_before_returning_idle() {
}
}
let compact = protocol::WorkerCommandEnvelope::for_snapshot(1, &handle.shared_state.snapshot());
let compact = protocol::WorkerCommandEnvelope::new(1);
handle
.send(Method::Compact { command: compact })
.await
@@ -1154,7 +1154,7 @@ async fn manual_compact_cancel_clears_progress_before_returning_idle() {
}
}
let cancel = protocol::WorkerCommandEnvelope::for_snapshot(2, &handle.shared_state.snapshot());
let cancel = protocol::WorkerCommandEnvelope::new(2);
handle
.send(Method::Cancel { command: cancel })
.await
@@ -1183,7 +1183,7 @@ async fn manual_compact_cancel_clears_progress_before_returning_idle() {
}
}
let compact = protocol::WorkerCommandEnvelope::for_snapshot(3, &handle.shared_state.snapshot());
let compact = protocol::WorkerCommandEnvelope::new(3);
handle
.send(Method::Compact { command: compact })
.await
@@ -1201,8 +1201,7 @@ async fn manual_compact_cancel_clears_progress_before_returning_idle() {
break;
}
}
let shutdown =
protocol::WorkerCommandEnvelope::for_snapshot(4, &handle.shared_state.snapshot());
let shutdown = protocol::WorkerCommandEnvelope::new(4);
handle
.send(Method::Shutdown { command: shutdown })
.await
@@ -1269,7 +1268,7 @@ async fn controller_compact_method_publishes_progress_and_clear() {
}
}
let command = protocol::WorkerCommandEnvelope::for_snapshot(1, &handle.shared_state.snapshot());
let command = protocol::WorkerCommandEnvelope::new(1);
handle
.send(Method::Compact { command })
.await
@@ -1316,6 +1315,6 @@ async fn controller_compact_method_publishes_progress_and_clear() {
protocol::WorkerStatus::Idle,
"successful manual compaction must release the execution fence"
);
let command = protocol::WorkerCommandEnvelope::for_snapshot(2, &handle.shared_state.snapshot());
let command = protocol::WorkerCommandEnvelope::new(2);
let _ = handle.send(Method::Shutdown { command }).await;
}
crates/worker/tests/controller_test.rs
+2 -5
View File
@@ -27,11 +27,8 @@ type TestStore = CombinedStore<FsStore, FsWorkerStore>;
static NEXT_COMMAND_ID: AtomicU64 = AtomicU64::new(1);
fn worker_command(handle: &WorkerHandle) -> protocol::WorkerCommandEnvelope {
protocol::WorkerCommandEnvelope::for_snapshot(
NEXT_COMMAND_ID.fetch_add(1, Ordering::Relaxed),
&handle.shared_state.snapshot(),
)
fn worker_command(_handle: &WorkerHandle) -> protocol::WorkerCommandEnvelope {
protocol::WorkerCommandEnvelope::new(NEXT_COMMAND_ID.fetch_add(1, Ordering::Relaxed))
}
/// Reconstruct a worker-history-like `Vec<Item>` from the live session