diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index a6717f442..842b9a4a6 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -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( + method: Method, + pending_submissions: &PendingSubmissionHandle, + working_event_tx: &broadcast::Sender, +) -> Option +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, @@ -2250,7 +2361,15 @@ async fn controller_loop( ); 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); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 27f7cd354..adf4de779 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -1930,6 +1930,15 @@ impl PendingSubmissionHandle { }, } } + + #[cfg(test)] + pub(crate) fn persisted_entries_for_test(&self) -> Vec { + 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(); diff --git a/docs/design/compaction.md b/docs/design/compaction.md index f12740c07..6d33041b5 100644 --- a/docs/design/compaction.md +++ b/docs/design/compaction.md @@ -46,11 +46,13 @@ activation boundary uses this lock order everywhere a live Worker changes Segmen 7. publish the append destination, session projection, sink, and in-memory history; 8. release the append barrier. -Submit and Notify acceptance take the same append barrier. Consequently an accepted -item is either durably included in the replacement checkpoint, or waits until every -live authority points at the replacement and is then appended there. Restore -admission takes the allocation lock, so it cannot observe the interval between the -metadata CAS and allocation hand-off. +Submit and Notify acceptance take the same append barrier. The controller persists +requests received while manual compaction is running immediately rather than holding +volatile deferred methods, so an overlapping Shutdown cannot discard a SubWorker +completion. Consequently an accepted item is either durably included in the +replacement checkpoint, or waits until every live authority points at the replacement +and is then appended there. Restore admission takes the allocation lock, so it cannot +observe the interval between the metadata CAS and allocation hand-off. The metadata CAS is the commit point. Failures before it leave the source Segment active. A failure updating allocation after it marks the old in-memory writer