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

worker: persist compaction-time notifications

Browse Source
This commit is contained in:
Hare
2026-09-18 18:53:47 +09:00
parent d59879cc4b
commit 72a61e4952
3 changed files with 248 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+145 -2
View File
@@ -25,7 +25,9 @@ use crate::shutdown_after_idle::{
};
use crate::spawn::registry::SpawnedWorkerRegistry;
use crate::spawn::tool::sub_worker_spawn_tool;
use crate::worker::{SystemItemCommitter, Worker, WorkerError, WorkerRunResult};
use crate::worker::{
PendingSubmissionHandle, SystemItemCommitter, Worker, WorkerError, WorkerRunResult,
};
use protocol::{
AlertLevel, AlertSource, CommandEvent as ProtocolCommandEvent,
CommandSnapshot as ProtocolCommandSnapshot, CommandStatus as ProtocolCommandStatus,
@@ -1489,6 +1491,115 @@ where
Ok(workdir_for_view)
}
fn durably_accept_method_while_busy<St>(
method: Method,
pending_submissions: &PendingSubmissionHandle<St>,
working_event_tx: &broadcast::Sender<Event>,
) -> Option<Method>
where
St: session_store::Store + Clone,
{
match method {
Method::Submit {
submission_request_id,
input,
} => {
let request_id = submission_request_id.clone();
match pending_submissions.accept_from_source(
submission_request_id,
input,
pending_submissions.direct_client_namespace(),
session_store::LoggedSessionHistoryOrigin::LegacyUnknown,
false,
) {
Ok(acceptance) => {
let _ = working_event_tx.send(Event::SubmissionAccepted {
submission_request_id: acceptance.submission_request_id,
submission_id: acceptance.submission_id,
disposition: acceptance.disposition,
});
}
Err(error) => {
let _ = working_event_tx.send(Event::SubmissionRejected {
submission_request_id: request_id,
message: error.to_string(),
});
}
}
None
}
Method::SubmitTracked {
submission_request_id,
input,
source,
} => {
let request_id = submission_request_id.clone();
let (source_namespace, provenance) =
resolved_input_source(pending_submissions, &source);
match pending_submissions.accept_from_source(
submission_request_id,
input,
source_namespace,
provenance,
false,
) {
Ok(acceptance) => {
let _ = working_event_tx.send(Event::SubmissionAccepted {
submission_request_id: acceptance.submission_request_id,
submission_id: acceptance.submission_id,
disposition: acceptance.disposition,
});
}
Err(error) => {
let _ = working_event_tx.send(Event::SubmissionRejected {
submission_request_id: request_id,
message: error.to_string(),
});
}
}
None
}
Method::Notify {
notification_request_id,
message,
} => {
if let Err(error) = pending_submissions.accept_notification_from_source(
notification_request_id,
message,
pending_submissions.direct_client_namespace(),
session_store::LoggedSessionHistoryOrigin::LegacyUnknown,
) {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
}
None
}
Method::NotifyTracked {
notification_request_id,
message,
source,
} => {
let (source_namespace, provenance) =
resolved_input_source(pending_submissions, &source);
if let Err(error) = pending_submissions.accept_notification_from_source(
notification_request_id,
message,
source_namespace,
provenance,
) {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
}
None
}
method => Some(method),
}
}
fn reject_method_while_attention_locked(
method: &Method,
working_event_tx: &broadcast::Sender<Event>,
@@ -2250,7 +2361,15 @@ async fn controller_loop<C, St>(
);
let _ = cancel_tx.send(true);
}
Some(method) => deferred_methods.push_back(method),
Some(method) => {
if let Some(method) = durably_accept_method_while_busy(
method,
&pending_submissions,
&working_event_tx,
) {
deferred_methods.push_back(method);
}
}
None => {
shutdown_after_compaction = true;
let _ = cancel_tx.send(true);
@@ -3460,6 +3579,30 @@ mod tests {
}
}
#[test]
fn busy_notification_is_durable_before_shutdown_can_observe_it() {
let dir = tempfile::tempdir().unwrap();
let pending = PendingSubmissionHandle::for_test(dir.path());
let (event_tx, _event_rx) = broadcast::channel(4);
assert!(durably_accept_method_while_busy(
Method::NotifyTracked {
notification_request_id: "completion-1".into(),
message: "subworker completed".into(),
source: protocol::AuthenticatedInputSource::UntrustedWire,
},
&pending,
&event_tx,
)
.is_none());
let entries = pending.persisted_entries_for_test();
assert!(entries.iter().any(|entry| matches!(
entry,
LogEntry::Extension { domain, .. }
if domain == "worker.pending_activations.v1"
)));
}
#[test]
fn attention_gate_rejects_submit_before_dispatch() {
let (event_tx, mut event_rx) = broadcast::channel(4);
crates/worker/src/worker.rs
+96
View File
@@ -1930,6 +1930,15 @@ impl PendingSubmissionHandle<session_store::FsStore> {
},
}
}
#[cfg(test)]
pub(crate) fn persisted_entries_for_test(&self) -> Vec<LogEntry> {
let location = self.writer.state.location();
self.writer
.store
.read_all(location.session_id, location.segment_id)
.expect("read test pending entries")
}
}
/// Type-erased commit handle for the interceptor. Lets the interceptor commit `SystemItem`s without being generic over the
@@ -8654,6 +8663,93 @@ mod build_summary_prompt_tests {
drop(recovered);
}
#[tokio::test]
async fn crash_after_allocation_before_live_switch_recovers_committed_replacement() {
let dir = tempfile::tempdir().unwrap();
let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap();
let mut worker = Worker::new(
minimal_manifest(),
Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient),
store,
WorkerWorkspaceContext::no_workspace(),
WorkerFilesystemAuthority::None,
Scope::empty(),
)
.await
.unwrap();
worker.ensure_segment_head().await.unwrap();
let old_location = worker.segment_state.location();
let replacement_location = SegmentLocation {
session_id: old_location.session_id,
segment_id: session_store::new_segment_id(),
};
let persisted_metadata = Arc::new(Mutex::new(WorkerActiveSegmentRef::active_segment(
old_location.session_id,
old_location.segment_id,
)));
let metadata_for_cas = Arc::clone(&persisted_metadata);
worker.worker_metadata_segment_cas = Some(Arc::new(move |_worker_name, expected, replacement| {
let mut persisted = metadata_for_cas.lock().unwrap();
if &*persisted != expected {
return Ok(false);
}
*persisted = replacement;
Ok(true)
}));
let lock_path = dir.path().join("workers.json");
let allocation = worker_allocation::install_top_level_at_for_test(
lock_path.clone(),
"test-worker".into(),
std::process::id(),
dir.path().join("worker.sock"),
Vec::new(),
old_location.segment_id,
)
.unwrap();
let activation = allocation
.begin_segment_activation(
old_location.segment_id,
replacement_location.segment_id,
)
.unwrap();
worker.scope_allocation = Some(allocation);
worker
.compare_and_swap_worker_metadata_segment(old_location, replacement_location)
.unwrap();
activation.commit().unwrap();
// Simulated crash point: durable metadata and allocation moved, while the
// in-memory append destination still names the source Segment.
assert_eq!(worker.segment_id(), old_location.segment_id);
assert_eq!(
*persisted_metadata.lock().unwrap(),
WorkerActiveSegmentRef::active_segment(
replacement_location.session_id,
replacement_location.segment_id,
)
);
drop(worker);
let recovered = worker_allocation::install_top_level_at_for_test(
lock_path,
"test-worker".into(),
std::process::id(),
dir.path().join("worker-restored.sock"),
Vec::new(),
replacement_location.segment_id,
)
.unwrap();
drop(
recovered
.begin_segment_activation(
replacement_location.segment_id,
session_store::new_segment_id(),
)
.unwrap(),
);
drop(recovered);
}
#[tokio::test]
async fn shutdown_cleanup_stays_blocked_until_registered_service_stops() {
let dir = tempfile::tempdir().unwrap();