diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 7fe3b02f4..2ea198113 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,153 @@ 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, + message: &str, +) -> bool { + match method { + Method::Shutdown { .. } | Method::ListPendingSubmissions => false, + Method::Submit { + submission_request_id, + .. + } + | Method::SubmitTracked { + submission_request_id, + .. + } => { + let _ = working_event_tx.send(Event::SubmissionRejected { + submission_request_id: submission_request_id.clone(), + message: message.to_owned(), + }); + true + } + Method::Notify { .. } | Method::NotifyTracked { .. } => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: message.to_owned(), + }); + true + } + _ => { + let _ = working_event_tx.send(Event::Error { + code: ErrorCode::InvalidRequest, + message: message.to_owned(), + }); + true + } + } +} + /// Idle/Paused event loop. Each iteration either fires a staged /// `PendingRun` (delegating to [`drive_turn`] for the Running phase) or /// waits for the next `Method`. Method handlers stop at "update state + @@ -1551,6 +1700,12 @@ async fn controller_loop( let mut deferred_methods = VecDeque::new(); 'controller: loop { + if worker.has_failed_segment_activation() || worker.has_pending_compaction_cleanup() { + // A queued activation may have been accepted while the failing turn + // was still running. Preserve its durable queue entry, but never + // start it while repair/cleanup attention is required. + pending = None; + } // here so the status flip → drive_turn → finish sequence lives // in one place, regardless of which Method caused it. if let Some(run) = pending.take() { @@ -1675,6 +1830,16 @@ async fn controller_loop( ) .await; if shutdown { + while let Err(error) = worker + .finish_pending_compaction_cleanup_for_shutdown() + .await + { + let _ = working_event_tx.send(Event::Error { + code: worker_error_code(&error), + message: error.to_string(), + }); + tokio::time::sleep(Duration::from_millis(50)).await; + } let _ = working_event_tx.send(Event::Shutdown); break; } @@ -1700,7 +1865,15 @@ async fn controller_loop( tokio::select! { method = method_rx.recv() => match method { Some(method) => method, - None => break, + None => { + while let Err(error) = + worker.finish_pending_compaction_cleanup_for_shutdown().await + { + tracing::warn!(error = %error, "shutdown cleanup retry failed"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + break 'controller; + } }, _ = notify_buffer.notified() => continue, _ = tokio::time::sleep(remaining) => { @@ -1726,12 +1899,71 @@ async fn controller_loop( tokio::select! { method = method_rx.recv() => match method { Some(method) => method, - None => break, + None => { + while let Err(error) = + worker.finish_pending_compaction_cleanup_for_shutdown().await + { + tracing::warn!(error = %error, "shutdown cleanup retry failed"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + break 'controller; + } }, _ = notify_buffer.notified() => continue, } }; + if worker.has_pending_compaction_cleanup() { + match worker.retry_pending_compaction_cleanup().await { + Ok(()) => { + set_controller_state( + &shared_state, + &runtime_dir, + &working_event_tx, + WorkerState::Idle, + ) + .await; + } + Err(error) => { + if matches!(&method, Method::Shutdown { .. }) { + if let Err(error) = worker + .finish_pending_compaction_cleanup_for_shutdown() + .await + { + let _ = working_event_tx.send(Event::Error { + code: worker_error_code(&error), + message: error.to_string(), + }); + continue; + } + } else { + let message = error.to_string(); + let _ = working_event_tx.send(Event::Error { + code: worker_error_code(&error), + message: message.clone(), + }); + if reject_method_while_attention_locked( + &method, + &working_event_tx, + &message, + ) { + continue; + } + } + } + } + } + + if worker.has_failed_segment_activation() + && reject_method_while_attention_locked( + &method, + &working_event_tx, + "Segment activation is incomplete; restore or shut down the Worker", + ) + { + continue; + } + match method { Method::Submit { submission_request_id, @@ -2132,7 +2364,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); @@ -2147,6 +2387,8 @@ async fn controller_loop( Err(WorkerError::Store(_)) | Err(WorkerError::WorkerStore(_)) | Err(WorkerError::InvalidState(_)) + | Err(WorkerError::SegmentActivationIncomplete { .. }) + | Err(WorkerError::CompactionCleanupPending { .. }) ) { set_controller_state( &shared_state, @@ -2163,6 +2405,16 @@ async fn controller_loop( }); } if shutdown_after_compaction { + while let Err(error) = worker + .finish_pending_compaction_cleanup_for_shutdown() + .await + { + let _ = working_event_tx.send(Event::Error { + code: worker_error_code(&error), + message: error.to_string(), + }); + tokio::time::sleep(Duration::from_millis(50)).await; + } let _ = working_event_tx.send(Event::Shutdown); break 'controller; } @@ -2565,6 +2817,19 @@ where (WorkerStatus::Paused, shutdown_requested, false) } Err(e) => { + let status = if matches!( + &e, + WorkerError::SegmentActivationIncomplete { .. } + | WorkerError::CompactionCleanupPending { .. } + ) { + // Metadata has already committed the replacement, but + // one live authority failed to publish. Keep the + // controller visibly busy/fail-closed; publishing Idle + // would admit another turn against inconsistent state. + WorkerStatus::Running + } else { + WorkerStatus::Idle + }; let code = worker_error_code(&e); let message = e.to_string(); let _ = working_event_tx.send(Event::Error { @@ -2580,7 +2845,7 @@ where }, ); } - (WorkerStatus::Idle, shutdown_requested, false) + (status, shutdown_requested, false) } }; } @@ -3318,6 +3583,137 @@ 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); + let method = Method::Submit { + submission_request_id: "request-1".into(), + input: Vec::new(), + }; + + assert!(reject_method_while_attention_locked( + &method, + &event_tx, + "repair required", + )); + assert!(matches!( + event_rx.try_recv().unwrap(), + Event::SubmissionRejected { + submission_request_id, + message, + } if submission_request_id == "request-1" && message == "repair required" + )); + } + + #[tokio::test] + async fn segment_activation_failure_from_run_never_publishes_idle() { + let mut env = make_env().await; + let worker_future = async { + Err::(WorkerError::SegmentActivationIncomplete { + source: crate::runtime::worker_allocation::ScopeLockError::UnknownWorker( + "worker".into(), + ), + }) + }; + + let (status, shutdown, may_drain_pending) = drive_turn( + worker_future, + &mut env.method_rx, + &env.working_event_tx, + &env.cancel_tx, + &env.pause_tx, + &env.shared_state, + &env.runtime_dir, + None, + &env.notify_buffer, + &env.pending_submissions, + None, + "worker", + &env.spawned_registry, + false, + ) + .await; + + assert_eq!(status, WorkerStatus::Running); + assert!(!shutdown); + assert!(!may_drain_pending); + } + + #[tokio::test] + async fn pending_compaction_cleanup_preserves_shutdown_request_for_outer_barrier() { + let mut env = make_env().await; + let method_tx = env._method_tx.clone(); + env.shared_state + .transition(WorkerState::Busy(WorkerBusyState::Run( + WorkerRunState::Running, + ))); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(10)).await; + method_tx + .send(Method::Shutdown { + command: WorkerCommandEnvelope::new(1), + }) + .await + .expect("send shutdown"); + }); + let worker_future = async { + tokio::time::sleep(Duration::from_millis(50)).await; + Err::(WorkerError::CompactionCleanupPending { + source: crate::runtime::worker_allocation::ScopeLockError::UnknownWorker( + "worker".into(), + ), + }) + }; + + let (status, shutdown, may_drain_pending) = drive_turn( + worker_future, + &mut env.method_rx, + &env.working_event_tx, + &env.cancel_tx, + &env.pause_tx, + &env.shared_state, + &env.runtime_dir, + None, + &env.notify_buffer, + &env.pending_submissions, + None, + "worker", + &env.spawned_registry, + false, + ) + .await; + + assert_eq!(status, WorkerStatus::Running); + assert!(shutdown); + assert!(!may_drain_pending); + } + #[tokio::test] async fn pause_waits_for_run_boundary_and_uses_safe_pause_channel() { let mut env = make_env().await; diff --git a/crates/worker/src/runtime/worker_allocation.rs b/crates/worker/src/runtime/worker_allocation.rs index 1eb3874bb..c7445cdaa 100644 --- a/crates/worker/src/runtime/worker_allocation.rs +++ b/crates/worker/src/runtime/worker_allocation.rs @@ -18,12 +18,16 @@ pub use conflict::{ ConflictOwner, find_conflict_owner, find_conflict_owners, is_within_effective_write, }; pub use error::ScopeLockError; +#[cfg(test)] +pub(crate) use lifecycle::install_top_level_at_for_test; pub use lifecycle::{ - ScopeAllocationGuard, SegmentLockInfo, adopt_allocation, install_top_level, - install_top_level_with_deny, lookup_segment, update_segment, + ScopeAllocationGuard, SegmentActivationCommitError, SegmentActivationGuard, SegmentLockInfo, + adopt_allocation, install_top_level, install_top_level_with_deny, lookup_segment, }; pub use mutate::{ delegate_scope, reclaim_delegated_scope, reclaim_stale, reclaim_stale_with, register_worker, register_worker_with_deny, release_worker, }; +#[cfg(test)] +pub(crate) use table::fail_next_save_for_test; pub use table::{Allocation, LockFile, LockFileGuard, default_allocation_path}; diff --git a/crates/worker/src/runtime/worker_allocation/error.rs b/crates/worker/src/runtime/worker_allocation/error.rs index a0bcb78ef..601202693 100644 --- a/crates/worker/src/runtime/worker_allocation/error.rs +++ b/crates/worker/src/runtime/worker_allocation/error.rs @@ -37,4 +37,12 @@ pub enum ScopeLockError { worker_name: String, socket: PathBuf, }, + #[error( + "worker `{worker_name}` allocation moved from expected Segment {expected} to {actual:?}" + )] + SegmentChanged { + worker_name: String, + expected: SegmentId, + actual: Option, + }, } diff --git a/crates/worker/src/runtime/worker_allocation/lifecycle.rs b/crates/worker/src/runtime/worker_allocation/lifecycle.rs index 480d895af..a596c1aa9 100644 --- a/crates/worker/src/runtime/worker_allocation/lifecycle.rs +++ b/crates/worker/src/runtime/worker_allocation/lifecycle.rs @@ -20,6 +20,41 @@ pub struct ScopeAllocationGuard { lock_path: PathBuf, } +/// Exclusive authority to activate a replacement Segment for one allocated Worker. +/// +/// The allocation lock remains held across the durable metadata CAS and the final +/// allocation-table update, so restore admission cannot observe an intermediate +/// ownership state. +pub struct SegmentActivationGuard { + guard: LockFileGuard, + worker_name: String, + expected: SegmentId, + replacement: SegmentId, +} + +/// A failed allocation-table publication together with the still-held machine-wide +/// activation authority. Callers must retain `authority` until recovery or process +/// teardown so restore admission cannot register a competing writer. +pub struct SegmentActivationCommitError { + source: ScopeLockError, + authority: SegmentActivationGuard, +} + +impl std::fmt::Debug for SegmentActivationCommitError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("SegmentActivationCommitError") + .field("source", &self.source) + .finish_non_exhaustive() + } +} + +impl SegmentActivationCommitError { + pub fn into_parts(self) -> (ScopeLockError, SegmentActivationGuard) { + (self.source, self.authority) + } +} + impl ScopeAllocationGuard { pub fn worker_name(&self) -> &str { &self.worker_name @@ -28,6 +63,86 @@ impl ScopeAllocationGuard { pub fn lock_path(&self) -> &Path { &self.lock_path } + + /// Lock the machine-wide table and validate a Segment replacement before + /// its durable metadata is changed. + pub fn begin_segment_activation( + &self, + expected: SegmentId, + replacement: SegmentId, + ) -> Result { + let guard = LockFileGuard::open(&self.lock_path)?; + let actual = guard + .data() + .allocations + .iter() + .find(|allocation| allocation.worker_name == self.worker_name) + .and_then(|allocation| allocation.segment_id); + + if actual != Some(expected) { + return Err(ScopeLockError::SegmentChanged { + worker_name: self.worker_name.clone(), + expected, + actual, + }); + } + + if let Some(existing) = guard.data().allocations.iter().find(|allocation| { + allocation.segment_id == Some(replacement) && allocation.worker_name != self.worker_name + }) { + return Err(ScopeLockError::SegmentConflict { + segment_id: replacement, + worker_name: existing.worker_name.clone(), + socket: existing.socket.clone(), + }); + } + + Ok(SegmentActivationGuard { + guard, + worker_name: self.worker_name.clone(), + expected, + replacement, + }) + } +} + +impl SegmentActivationGuard { + /// Persist the replacement Segment while retaining exclusive allocation + /// authority. This is intentionally consumed so a successful metadata CAS + /// cannot accidentally be followed by an unlocked update. + pub fn commit(mut self) -> Result<(), SegmentActivationCommitError> { + let actual = self + .guard + .data() + .allocations + .iter() + .find(|allocation| allocation.worker_name == self.worker_name) + .and_then(|allocation| allocation.segment_id); + let Some(allocation) = self + .guard + .data_mut() + .allocations + .iter_mut() + .find(|allocation| allocation.worker_name == self.worker_name) + else { + return Err(SegmentActivationCommitError { + source: ScopeLockError::SegmentChanged { + worker_name: self.worker_name.clone(), + expected: self.expected, + actual, + }, + authority: self, + }); + }; + allocation.segment_id = Some(self.replacement); + if let Err(source) = self.guard.save() { + return Err(SegmentActivationCommitError { + source: ScopeLockError::Io(source), + authority: self, + }); + } + Ok(()) + } } impl Drop for ScopeAllocationGuard { @@ -85,6 +200,31 @@ pub fn install_top_level_with_deny( }) } +#[cfg(test)] +pub(crate) fn install_top_level_at_for_test( + lock_path: PathBuf, + worker_name: String, + pid: u32, + socket: PathBuf, + scope_allow: Vec, + segment_id: SegmentId, +) -> Result { + let mut guard = LockFileGuard::open(&lock_path)?; + super::mutate::register_worker_with_deny( + &mut guard, + worker_name.clone(), + pid, + socket, + scope_allow, + Vec::new(), + segment_id, + )?; + Ok(ScopeAllocationGuard { + worker_name, + lock_path, + }) +} + /// Take ownership of an existing allocation that was pre-registered by /// a spawning Worker. /// @@ -114,46 +254,6 @@ pub fn adopt_allocation( }) } -/// Rewrite the `segment_id` recorded for `worker_name` to -/// `new_segment_id`. -/// -/// The Worker's in-memory `segment_id` can change underneath the -/// allocation in two normal places: -/// -/// - `Worker::compact` mints a fresh Segment in the same Session. -/// - `session_store::ensure_head_or_fork` auto-forks within that Session when another -/// writer has advanced the store head behind our back. -/// -/// Both paths must call this so subsequent [`lookup_segment`] queries -/// find the live Segment id, not the old one. Without this update a -/// concurrent `restore_from_manifest(new_id)` would see "no live -/// writer" and proceed to register a competing allocation on the -/// Segment lineage this Worker just moved into. -/// -/// The lock is opened once and the allocation is rewritten inside the -/// guard, so the segment_id collision check is atomic with the -/// rewrite. -pub fn update_segment(worker_name: &str, new_segment_id: SegmentId) -> Result<(), ScopeLockError> { - let lock_path = default_allocation_path()?; - let mut guard = LockFileGuard::open(&lock_path)?; - if let Some(other) = guard.data().find_by_segment(new_segment_id) { - if other.worker_name != worker_name { - return Err(ScopeLockError::SegmentConflict { - segment_id: new_segment_id, - worker_name: other.worker_name.clone(), - socket: other.socket.clone(), - }); - } - } - let alloc = guard - .data_mut() - .find_mut(worker_name) - .ok_or_else(|| ScopeLockError::UnknownWorker(worker_name.into()))?; - alloc.segment_id = Some(new_segment_id); - guard.save()?; - Ok(()) -} - /// Information about a Worker that currently holds an allocation for a /// given session. #[derive(Debug, Clone)] @@ -184,7 +284,7 @@ pub fn lookup_segment(segment_id: SegmentId) -> Result, #[cfg(test)] mod tests { - use super::super::table::Allocation; + use super::super::table::{Allocation, LockFile}; use super::super::test_util::*; use super::*; use tempfile::TempDir; @@ -283,32 +383,12 @@ mod tests { } #[test] - fn update_session_rewrites_allocation_session_id() { - let dir = TempDir::new().unwrap(); - let _sandbox = RuntimeDirSandbox::new(dir.path()); - let original = sid(); - let updated = sid(); - let _guard = install_top_level( - "p".into(), - std::process::id(), - sock("p"), - vec![write_rule("/work", true)], - original, - ) - .unwrap(); - update_segment("p", updated).unwrap(); - // lookup against the original is now empty, the updated id wins. - assert!(lookup_segment(original).unwrap().is_none()); - assert_eq!(lookup_segment(updated).unwrap().unwrap().worker_name, "p"); - } - - #[test] - fn update_session_rejects_when_target_already_held() { + fn segment_activation_rejects_when_target_already_held() { let dir = TempDir::new().unwrap(); let _sandbox = RuntimeDirSandbox::new(dir.path()); let s_a = sid(); let s_b = sid(); - let _g_a = install_top_level( + let g_a = install_top_level( "a".into(), std::process::id(), sock("a"), @@ -324,8 +404,11 @@ mod tests { s_b, ) .unwrap(); - // `a` cannot adopt b's live session id. - let err = update_segment("a", s_b).unwrap_err(); + // `a` cannot activate b's live Segment id. + let err = match g_a.begin_segment_activation(s_a, s_b) { + Ok(_) => panic!("expected SegmentConflict"), + Err(error) => error, + }; match err { ScopeLockError::SegmentConflict { worker_name, @@ -338,4 +421,191 @@ mod tests { other => panic!("expected SegmentConflict, got {other:?}"), } } + + #[test] + fn segment_activation_commits_verified_replacement() { + let dir = TempDir::new().unwrap(); + let _sandbox = RuntimeDirSandbox::new(dir.path()); + let lock_path = dir.path().join("workers.json"); + let old_segment = sid(); + let replacement = sid(); + let guard = install_top_level( + "activation".into(), + std::process::id(), + sock("activation"), + vec![write_rule("/work", true)], + old_segment, + ) + .unwrap(); + + guard + .begin_segment_activation(old_segment, replacement) + .unwrap() + .commit() + .unwrap(); + + let table = LockFileGuard::open(&lock_path).unwrap(); + assert_eq!( + table + .data() + .find("activation") + .and_then(|allocation| allocation.segment_id), + Some(replacement) + ); + } + + #[test] + fn failed_segment_activation_returns_locked_recovery_authority() { + let dir = TempDir::new().unwrap(); + let _sandbox = RuntimeDirSandbox::new(dir.path()); + let old_segment = sid(); + let replacement = sid(); + let guard = install_top_level( + "activation".into(), + std::process::id(), + sock("activation"), + vec![write_rule("/work", true)], + old_segment, + ) + .unwrap(); + let mut activation = guard + .begin_segment_activation(old_segment, replacement) + .unwrap(); + activation + .guard + .data_mut() + .allocations + .retain(|allocation| allocation.worker_name != "activation"); + + let error = activation.commit().unwrap_err(); + let (source, authority) = error.into_parts(); + assert!(matches!( + source, + ScopeLockError::SegmentChanged { + expected, + actual: None, + .. + } if expected == old_segment + )); + // The authority is returned rather than dropped on failure. Production + // stores it on Worker to keep restore admission blocked. + let (sent, received) = std::sync::mpsc::channel(); + let lookup = std::thread::spawn(move || { + sent.send(lookup_segment(replacement).unwrap()).unwrap(); + }); + assert!( + received + .recv_timeout(std::time::Duration::from_millis(50)) + .is_err() + ); + drop(authority); + assert!( + received + .recv_timeout(std::time::Duration::from_secs(1)) + .unwrap() + .is_none() + ); + lookup.join().unwrap(); + } + + #[test] + fn allocation_save_failure_preserves_source_table_and_returns_authority() { + let dir = TempDir::new().unwrap(); + let _sandbox = RuntimeDirSandbox::new(dir.path()); + let lock_path = dir.path().join("workers.json"); + let old_segment = sid(); + let replacement = sid(); + let guard = install_top_level( + "activation".into(), + std::process::id(), + sock("activation"), + vec![write_rule("/work", true)], + old_segment, + ) + .unwrap(); + let activation = guard + .begin_segment_activation(old_segment, replacement) + .unwrap(); + crate::runtime::worker_allocation::table::fail_next_save_for_test(&lock_path); + + let error = activation.commit().unwrap_err(); + let (source, authority) = error.into_parts(); + assert!(matches!(source, ScopeLockError::Io(_))); + let persisted: LockFile = + serde_json::from_slice(&std::fs::read(&lock_path).unwrap()).unwrap(); + assert_eq!( + persisted + .find("activation") + .and_then(|allocation| allocation.segment_id), + Some(old_segment) + ); + drop(authority); + } + + #[test] + fn dropping_segment_activation_preserves_source_allocation() { + let dir = TempDir::new().unwrap(); + let _sandbox = RuntimeDirSandbox::new(dir.path()); + let lock_path = dir.path().join("workers.json"); + let old_segment = sid(); + let replacement = sid(); + let guard = install_top_level( + "activation".into(), + std::process::id(), + sock("activation"), + vec![write_rule("/work", true)], + old_segment, + ) + .unwrap(); + + drop( + guard + .begin_segment_activation(old_segment, replacement) + .unwrap(), + ); + + let table = LockFileGuard::open(&lock_path).unwrap(); + assert_eq!( + table + .data() + .find("activation") + .and_then(|allocation| allocation.segment_id), + Some(old_segment) + ); + } + + #[test] + fn segment_activation_rejects_stale_expected_segment_without_mutation() { + let dir = TempDir::new().unwrap(); + let _sandbox = RuntimeDirSandbox::new(dir.path()); + let lock_path = dir.path().join("workers.json"); + let allocated = sid(); + let stale = sid(); + let replacement = sid(); + let guard = install_top_level( + "activation".into(), + std::process::id(), + sock("activation"), + vec![write_rule("/work", true)], + allocated, + ) + .unwrap(); + + assert!(matches!( + guard.begin_segment_activation(stale, replacement), + Err(ScopeLockError::SegmentChanged { + expected, + actual: Some(actual), + .. + }) if expected == stale && actual == allocated + )); + let table = LockFileGuard::open(&lock_path).unwrap(); + assert_eq!( + table + .data() + .find("activation") + .and_then(|allocation| allocation.segment_id), + Some(allocated) + ); + } } diff --git a/crates/worker/src/runtime/worker_allocation/table.rs b/crates/worker/src/runtime/worker_allocation/table.rs index 3056074f1..af7c6159a 100644 --- a/crates/worker/src/runtime/worker_allocation/table.rs +++ b/crates/worker/src/runtime/worker_allocation/table.rs @@ -1,12 +1,48 @@ //! On-disk allocation table and the `flock`-protected guard. -use std::fs::{DirBuilder, File, OpenOptions}; +use std::fs::{self, DirBuilder, File, OpenOptions}; use std::io::{self, Read, Seek, SeekFrom, Write}; use std::os::unix::fs::{DirBuilderExt, OpenOptionsExt}; use std::path::{Path, PathBuf}; use std::thread; use std::time::{Duration, Instant}; +#[cfg(test)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SaveFaultPoint { + BeforeRename, + AfterRename, +} + +#[cfg(test)] +static SAVE_FAULTS: std::sync::LazyLock< + std::sync::Mutex>, +> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new())); + +#[cfg(test)] +pub(crate) fn fail_next_save_for_test(path: &Path) { + fail_next_save_at_for_test(path, SaveFaultPoint::BeforeRename); +} + +#[cfg(test)] +pub(crate) fn fail_next_save_at_for_test(path: &Path, point: SaveFaultPoint) { + SAVE_FAULTS + .lock() + .unwrap() + .insert(path.to_path_buf(), point); +} + +#[cfg(test)] +fn take_save_fault(path: &Path, point: SaveFaultPoint) -> bool { + let mut faults = SAVE_FAULTS.lock().unwrap(); + if faults.get(path) == Some(&point) { + faults.remove(path); + true + } else { + false + } +} + use fs4::fs_std::FileExt; use manifest::{ScopeRule, paths}; use serde::{Deserialize, Serialize}; @@ -101,7 +137,9 @@ pub fn default_allocation_path() -> io::Result { /// never call `save` leave the table unchanged, which is the right /// behaviour for error paths. pub struct LockFileGuard { - file: File, + data_file: File, + lock_file: File, + data_path: PathBuf, data: LockFile, } @@ -121,16 +159,19 @@ impl LockFileGuard { .mode(0o700) .create(parent)?; } - let file = OpenOptions::new() + let mut lock_os = path.as_os_str().to_owned(); + lock_os.push(".lock"); + let lock_path = PathBuf::from(lock_os); + let lock_file = OpenOptions::new() .read(true) .write(true) .create(true) .truncate(false) .mode(0o600) - .open(path)?; + .open(&lock_path)?; let started = Instant::now(); loop { - match FileExt::try_lock_exclusive(&file) { + match FileExt::try_lock_exclusive(&lock_file) { Ok(true) => break, Ok(false) => { if started.elapsed() >= LOCK_WAIT_TIMEOUT { @@ -159,8 +200,17 @@ impl LockFileGuard { Err(error) => return Err(error), } } + let data_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .mode(0o600) + .open(path)?; let mut this = Self { - file, + data_file, + lock_file, + data_path: path.to_path_buf(), data: LockFile::default(), }; this.reload()?; @@ -168,9 +218,9 @@ impl LockFileGuard { } fn reload(&mut self) -> io::Result<()> { - self.file.seek(SeekFrom::Start(0))?; + self.data_file.seek(SeekFrom::Start(0))?; let mut buf = String::new(); - self.file.read_to_string(&mut buf)?; + self.data_file.read_to_string(&mut buf)?; self.data = if buf.trim().is_empty() { LockFile::default() } else { @@ -192,20 +242,49 @@ impl LockFileGuard { &mut self.data } - /// Serialise `self.data` back to the file (truncate + rewrite). + /// Persist with atomic replacement while the separate allocation lock stays + /// held. A crash exposes either the old complete table or the new one. pub fn save(&mut self) -> io::Result<()> { let json = serde_json::to_vec_pretty(&self.data).map_err(io::Error::other)?; - self.file.seek(SeekFrom::Start(0))?; - self.file.set_len(0)?; - self.file.write_all(&json)?; - self.file.sync_data()?; + let mut temp_os = self.data_path.as_os_str().to_owned(); + temp_os.push(".tmp"); + let temp_path = PathBuf::from(temp_os); + let mut temp_file = OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .mode(0o600) + .open(&temp_path)?; + temp_file.write_all(&json)?; + temp_file.sync_all()?; + #[cfg(test)] + if take_save_fault(&self.data_path, SaveFaultPoint::BeforeRename) { + return Err(io::Error::other( + "injected allocation save failure before rename", + )); + } + fs::rename(&temp_path, &self.data_path)?; + #[cfg(test)] + if take_save_fault(&self.data_path, SaveFaultPoint::AfterRename) { + return Err(io::Error::other( + "injected allocation save failure after rename", + )); + } + if let Some(parent) = self.data_path.parent() { + File::open(parent)?.sync_all()?; + } + self.data_file = OpenOptions::new() + .read(true) + .write(true) + .mode(0o600) + .open(&self.data_path)?; Ok(()) } } impl Drop for LockFileGuard { fn drop(&mut self) { - let _ = FileExt::unlock(&self.file); + let _ = FileExt::unlock(&self.lock_file); } } @@ -259,6 +338,76 @@ mod tests { assert_eq!(guard.data().allocations[0].worker_name, "a"); } + #[test] + fn atomic_save_reopens_complete_table_at_both_rename_boundaries() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("workers.json"); + let old_segment = sid(); + let replacement = sid(); + { + let mut guard = open_empty(&path); + register_worker( + &mut guard, + "a".into(), + std::process::id(), + sock("a"), + vec![write_rule("/src", true)], + old_segment, + ) + .unwrap(); + } + + { + let mut guard = LockFileGuard::open(&path).unwrap(); + guard.data_mut().find_mut("a").unwrap().segment_id = Some(replacement); + fail_next_save_at_for_test(&path, SaveFaultPoint::BeforeRename); + assert!(guard.save().is_err()); + } + let guard = LockFileGuard::open(&path).unwrap(); + assert_eq!( + guard + .data() + .find("a") + .and_then(|allocation| allocation.segment_id), + Some(old_segment) + ); + drop(guard); + + { + let mut guard = LockFileGuard::open(&path).unwrap(); + guard.data_mut().find_mut("a").unwrap().segment_id = Some(replacement); + fail_next_save_at_for_test(&path, SaveFaultPoint::AfterRename); + assert!(guard.save().is_err()); + } + let guard = LockFileGuard::open(&path).unwrap(); + assert_eq!( + guard + .data() + .find("a") + .and_then(|allocation| allocation.segment_id), + Some(replacement) + ); + } + + #[test] + fn allocation_lock_remains_exclusive_across_atomic_save() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("workers.json"); + let mut guard = open_empty(&path); + guard.save().unwrap(); + + let (sent, received) = std::sync::mpsc::channel(); + let contender_path = path.clone(); + let contender = std::thread::spawn(move || { + let contender = LockFileGuard::open(&contender_path).unwrap(); + sent.send(contender.data().allocations.len()).unwrap(); + }); + assert!(received.recv_timeout(Duration::from_millis(50)).is_err()); + drop(guard); + assert_eq!(received.recv_timeout(Duration::from_secs(1)).unwrap(), 0); + contender.join().unwrap(); + } + #[test] fn find_by_session_skips_none_placeholders() { let dir = TempDir::new().unwrap(); diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index f328838d0..7c2606f74 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -8,6 +8,8 @@ //! Parent registry drop closes all session handles and synchronously returns delegated Write deny //! rules to the parent scope. +#[cfg(test)] +use std::collections::HashMap; use std::collections::{BTreeMap, HashSet}; use std::io; use std::sync::{ @@ -311,6 +313,8 @@ impl Drop for InternalSpawnReservation { pub struct SpawnedWorkerRegistry { internal_records: std::sync::Mutex>, service_records: std::sync::Mutex>, + #[cfg(test)] + fail_service_stops: std::sync::Mutex>, internal_names: std::sync::Mutex>, internal_shutting_down: AtomicBool, pending_internal_spawns: AtomicUsize, @@ -331,6 +335,8 @@ impl SpawnedWorkerRegistry { Arc::new(Self { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), + #[cfg(test)] + fail_service_stops: std::sync::Mutex::new(HashMap::new()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), pending_internal_spawns: AtomicUsize::new(0), @@ -346,6 +352,8 @@ impl SpawnedWorkerRegistry { Arc::new(Self { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), + #[cfg(test)] + fail_service_stops: std::sync::Mutex::new(HashMap::new()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), pending_internal_spawns: AtomicUsize::new(0), @@ -360,6 +368,8 @@ impl SpawnedWorkerRegistry { Arc::new(Self { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), + #[cfg(test)] + fail_service_stops: std::sync::Mutex::new(HashMap::new()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), pending_internal_spawns: AtomicUsize::new(0), @@ -443,6 +453,8 @@ impl SpawnedWorkerRegistry { registry: Arc::new(Self { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), + #[cfg(test)] + fail_service_stops: std::sync::Mutex::new(HashMap::new()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), pending_internal_spawns: AtomicUsize::new(0), @@ -535,9 +547,53 @@ impl SpawnedWorkerRegistry { Ok(worker_ref) } - /// Stop and remove one parent-owned service Worker. This is host-only and is - /// intentionally separate from the SubWorker control surface. + #[cfg(test)] + pub(crate) fn install_service_for_test( + &self, + ) -> (String, tokio::sync::broadcast::Sender) { + let (session, status_tx) = crate::internal_worker::test_internal_worker_session( + crate::internal_worker::InternalWorkerVisibility::ServicePrivate, + ); + let session_id = session.session_id_string(); + self.service_records + .lock() + .unwrap() + .push(InternalServiceWorkerRecord::new( + "test-service", + "test-service", + session, + )); + (session_id, status_tx) + } + + #[cfg(test)] + pub(crate) fn has_service_for_test(&self, session_id: &str) -> bool { + self.service_records + .lock() + .unwrap() + .iter() + .any(|record| record.session.session_id_string() == session_id) + } + + #[cfg(test)] + pub(crate) fn fail_service_stops_for_test(&self, session_id: &str, count: usize) { + self.fail_service_stops + .lock() + .unwrap() + .insert(session_id.to_owned(), count); + } + pub(crate) async fn stop_service(&self, session_id: &str) -> io::Result { + #[cfg(test)] + { + let mut failures = self.fail_service_stops.lock().unwrap(); + if let Some(remaining) = failures.get_mut(session_id) { + if *remaining > 0 { + *remaining -= 1; + return Err(io::Error::other("injected service stop failure")); + } + } + } let record = self .service_records .lock() diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 9acbdded3..a8e9722bd 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -272,17 +272,18 @@ impl ParentNotificationTarget { }; match self { Self::Controller { sender, fallback } => { + // Persist through the Worker-owned pending authority before the + // child completion path returns. The Controller delivery is only + // a wake-up hint; replay is idempotent by notification_request_id. + fallback(method.clone()); let Some(parent_method_tx) = sender.upgrade() else { - fallback(method); return; }; - let fallback = fallback.clone(); tokio::spawn(async move { - if let Err(error) = parent_method_tx.send(method).await { + if parent_method_tx.send(method).await.is_err() { tracing::warn!( - "failed to notify parent Controller; using durable pending authority" + "failed to wake parent Controller after durable notification acceptance" ); - fallback(error.0); } }); } @@ -1201,6 +1202,30 @@ enabled = false assert!(*captured.lock().unwrap()); } + #[tokio::test] + async fn parent_controller_notification_is_durable_before_wakeup_delivery() { + let (parent_method_tx, mut parent_method_rx) = mpsc::channel(1); + let durable = Arc::new(AtomicBool::new(false)); + let durable_for_fallback = Arc::clone(&durable); + let target = ParentNotificationTarget::with_controller_fallback( + parent_method_tx.downgrade(), + ParentNotificationTarget::Durable(Arc::new(move |_| { + durable_for_fallback.store(true, Ordering::Release); + })), + ); + + target.notify("child-session".into(), "completed".into()); + let wakeup = parent_method_rx.recv().await.unwrap(); + assert!(durable.load(Ordering::Acquire)); + assert!(matches!( + wakeup, + Method::NotifyTracked { + source: protocol::AuthenticatedInputSource::SubWorker { session_id }, + .. + } if session_id == "child-session" + )); + } + #[test] fn durable_parent_notification_target_preserves_child_source() { let captured = Arc::new(std::sync::Mutex::new(None)); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 216826f0b..d5c193f20 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -2,7 +2,7 @@ use std::collections::VecDeque; #[cfg(test)] use std::path::Path; use std::path::PathBuf; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -360,7 +360,9 @@ use crate::prompt::catalog::{CatalogError, PromptCatalog, WorkspacePromptProject use crate::prompt::source::PromptCatalogSource; use crate::prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTemplate}; use crate::runtime::dir; -use crate::runtime::worker_allocation::{self, ScopeAllocationGuard, ScopeLockError}; +use crate::runtime::worker_allocation::{ + self, ScopeAllocationGuard, ScopeLockError, SegmentActivationGuard, +}; use crate::skill::{SkillActivationResponse, SkillClientError}; #[cfg(test)] use async_trait::async_trait; @@ -1055,6 +1057,7 @@ pub struct SegmentState { location: ArcSwap, entries_written: AtomicUsize, append_lock: Mutex<()>, + activation_failed: AtomicBool, } impl SegmentState { @@ -1066,6 +1069,7 @@ impl SegmentState { }), entries_written: AtomicUsize::new(entries_written), append_lock: Mutex::new(()), + activation_failed: AtomicBool::new(false), }) } @@ -1093,6 +1097,19 @@ impl SegmentState { self.entries_written.store(n, Ordering::Release); } + fn ensure_append_allowed(&self) -> Result<(), StoreError> { + if self.activation_failed.load(Ordering::Acquire) { + return Err(StoreError::Io(std::io::Error::other( + "Segment activation is incomplete; restore the Worker before accepting more input", + ))); + } + Ok(()) + } + + fn fail_closed_activation(&self) { + self.activation_failed.store(true, Ordering::Release); + } + fn increment_entries(&self) { self.entries_written.fetch_add(1, Ordering::Release); } @@ -1223,6 +1240,7 @@ where } fn append_entry_locked(&self, entry: LogEntry) -> Result<(), StoreError> { + self.state.ensure_append_allowed()?; let loc = self.state.location(); self.store.append(loc.session_id, loc.segment_id, &entry)?; self.state.increment_entries(); @@ -1439,6 +1457,7 @@ where .append_lock .lock() .expect("segment append lock poisoned"); + self.writer.state.ensure_append_allowed()?; let mut current = self .state .lock() @@ -1580,6 +1599,7 @@ where .append_lock .lock() .expect("segment append lock poisoned"); + self.writer.state.ensure_append_allowed()?; let mut state = self .state .lock() @@ -1910,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 @@ -2068,6 +2097,12 @@ pub(crate) trait SystemPromptContributionSource: Send + Sync { async fn load(&self) -> Option; } +#[derive(Debug, Clone)] +struct PendingCompactionCleanup { + session_id: String, + compaction_id: String, +} + /// An independent agent execution unit. /// /// Holds a [`Engine`] directly and persists session state via @@ -2166,6 +2201,8 @@ pub struct Worker { /// Parent-owned projection/control boundary for observable Internal service Workers. /// Service Workers are never exposed through the model-facing SubWorker control surface. internal_worker_registry: Option>, + /// Cleanup authority survives cancellation of an individual compaction future. + pending_compaction_cleanup: Mutex>, in_flight: Option, /// Monotonic counter incremented by worker event bridges when an /// assistant-side execution artifact becomes visible to clients before @@ -2190,6 +2227,10 @@ pub struct Worker { /// interceptor before Agen applies them to live typed history. pending_committed_history: Arc>>>, + /// Machine-wide lock retained after a post-CAS allocation publication + /// failure. Declared before `scope_allocation` so drop releases this nested + /// authority before the outer allocation guard reacquires the same lock. + failed_activation_authority: Mutex>, /// Scope allocation in the machine-wide lock file. `Some` for /// Workers built via `from_manifest` / `from_manifest_spawned` / /// `restore_from_manifest` (production paths); `None` for the @@ -2251,6 +2292,19 @@ pub struct Worker { log_writer: Option>, } +impl Drop for Worker { + fn drop(&mut self) { + // A failed activation authority nests the same machine-wide lock used by + // `ScopeAllocationGuard::drop`. Release the nested guard first so Worker + // teardown cannot self-deadlock while removing its outer allocation. + let authority = match self.failed_activation_authority.get_mut() { + Ok(authority) => authority.take(), + Err(poisoned) => poisoned.into_inner().take(), + }; + drop(authority); + } +} + impl Worker { pub async fn stop_feature_runtime(&mut self, reason: impl Into) { if let Some(hooks) = self.hook_registry.clone() { @@ -2443,12 +2497,14 @@ impl Worker { alerter: None, working_event_tx: None, internal_worker_registry: None, + pending_compaction_cleanup: Mutex::new(None), in_flight: None, ai_activity_counter: Arc::new(AtomicUsize::new(0)), pending_notifies: NotifyBuffer::new(), notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY, pending_attachments: Arc::new(Mutex::new(Vec::::new())), pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), + failed_activation_authority: Mutex::new(None), scope_allocation: None, callback_socket: None, runtime_ticket_role: None, @@ -3611,7 +3667,7 @@ impl Worker { self.ensure_system_prompt_materialized().await?; self.ensure_segment_head().await?; if self.should_pre_run_compact() { - self.try_pre_run_compact().await; + self.try_pre_run_compact().await?; } Ok(()) } @@ -4414,8 +4470,15 @@ impl Worker { let _rewrite_guard = self .prepare_session_rewrite(SessionRewriteKind::Fork) .await?; - let w = self.engine.as_ref().unwrap(); + let segment_state = Arc::clone(&self.segment_state); + let _append_guard = segment_state + .append_lock + .lock() + .expect("segment append lock poisoned"); + segment_state.ensure_append_allowed()?; + let loc = segment_state.location(); let fork_segment_id = session_store::new_segment_id(); + let w = self.engine.as_ref().unwrap(); let entry = LogEntry::AnnotatedSegmentStart { ts: segment_log::now_millis(), session_id: loc.session_id, @@ -4446,24 +4509,47 @@ impl Worker { { initial_entries.push(checkpoint); } + initial_entries.push(LogEntry::Extension { + ts: segment_log::now_millis(), + domain: SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN.to_owned(), + payload: serde_json::to_value( + &*self + .session + .pending_activations + .lock() + .expect("pending activation state poisoned"), + ) + .map_err(|error| { + WorkerError::InvalidState(format!( + "serialize pending submissions during automatic fork: {error}" + )) + })?, + }); self.store .create_segment(loc.session_id, fork_segment_id, &initial_entries) .map_err(WorkerError::from)?; - self.segment_state.set_location(SegmentLocation { + let new_location = SegmentLocation { session_id: loc.session_id, segment_id: fork_segment_id, - }); + }; + let allocation_activation = self + .scope_allocation + .as_ref() + .map(|allocation| allocation.begin_segment_activation(loc.segment_id, fork_segment_id)) + .transpose()?; + self.compare_and_swap_worker_metadata_segment(loc, new_location)?; + if let Some(allocation_activation) = allocation_activation { + if let Err(error) = allocation_activation.commit() { + let (source, authority) = error.into_parts(); + self.fail_closed_segment_activation(authority); + return Err(WorkerError::SegmentActivationIncomplete { source }); + } + } + self.segment_state.set_location(new_location); self.segment_state .set_entries_written(initial_entries.len()); self.sink .reset_with_initial_entries(initial_entries.clone()); - if self.scope_allocation.is_some() { - worker_allocation::update_segment(&self.manifest.worker.name, fork_segment_id)?; - } - self.write_worker_metadata_active(SegmentLocation { - session_id: loc.session_id, - segment_id: fork_segment_id, - })?; Ok(()) } @@ -4662,16 +4748,18 @@ impl Worker { /// This used to run in the controller's post-run path. Keeping it here /// preserves the ordering requirement that the next turn starts with a /// compacted history, without introducing a separate Busy controller state. - /// Best-effort: failures are logged and surfaced, but do not abort the - /// user turn that triggered the check. - pub async fn try_pre_run_compact(&mut self) { + /// Fatal activation failures are returned to the controller so it cannot + /// publish Idle or start the triggering turn with inconsistent authority. + pub async fn try_pre_run_compact(&mut self) -> Result<(), WorkerError> { let Some(state) = self.compact_state.clone() else { - return; + return Ok(()); }; let current_tokens = self.total_tokens().tokens; match state.evaluate_pre_run(current_tokens) { AutomaticCompactDecision::Start(AutomaticCompactTrigger::PreRun) => {} - AutomaticCompactDecision::Continue | AutomaticCompactDecision::Block(_) => return, + AutomaticCompactDecision::Continue | AutomaticCompactDecision::Block(_) => { + return Ok(()); + } AutomaticCompactDecision::Start(AutomaticCompactTrigger::RequestThreshold) => { unreachable!("pre-run evaluation returned request-threshold trigger") } @@ -4692,6 +4780,7 @@ impl Worker { completed, "automatic compaction must complete a claimed attempt" ); + Ok(()) } Err(e) => { warn!(error = %e, "Proactive pre-run compaction failed"); @@ -4710,6 +4799,15 @@ impl Worker { completed, "automatic compaction must complete a claimed attempt" ); + if matches!( + &e, + WorkerError::SegmentActivationIncomplete { .. } + | WorkerError::CompactionCleanupPending { .. } + ) { + Err(e) + } else { + Ok(()) + } } } } @@ -4995,6 +5093,7 @@ impl Worker { mut cancel: Option>, trigger: CompactionTrigger, ) -> Result { + self.release_pending_compaction_service().await?; let _rewrite_guard = self .prepare_session_rewrite(SessionRewriteKind::Compact) .await?; @@ -5058,15 +5157,23 @@ impl Worker { match outcome { Ok((new_segment_id, _summary, stats)) => { debug_assert_eq!(lifecycle.state, CompactionLifecycleState::Done); + self.set_compaction_progress(None); for metric in attempt.success_metrics(new_segment_id, started.elapsed(), &stats) { self.try_record_metric(&metric); } self.usage_tracker .note_compaction_correlation_id(attempt.correlation_id().to_string()); - self.release_compaction_service(&lifecycle).await; Ok(new_segment_id) } Err(error) => { + let error = if matches!(&error, WorkerError::CompactionCleanupPending { .. }) { + error + } else { + match self.release_compaction_service(&mut lifecycle).await { + Ok(()) => error, + Err(cleanup_error) => cleanup_error, + } + }; let observed_segment_id = self.segment_state.location().segment_id; for metric in attempt.failure_metrics( observed_segment_id, @@ -5076,40 +5183,112 @@ impl Worker { self.try_record_metric(&metric); } lifecycle.revision = lifecycle.revision.saturating_add(1); - lifecycle.state = if matches!(error, WorkerError::CompactCancelled) { + lifecycle.state = if matches!(&error, WorkerError::CompactCancelled) { CompactionLifecycleState::Interrupted } else { CompactionLifecycleState::Failed }; lifecycle.ended_at_ms = Some(segment_log::now_millis()); - self.set_compaction_progress(None); - self.release_compaction_service(&lifecycle).await; + if !matches!(&error, WorkerError::CompactionCleanupPending { .. }) { + self.set_compaction_progress(None); + } Err(error) } } } - async fn release_compaction_service(&self, lifecycle: &CompactionLifecycle) { - let Some(session_id) = lifecycle - .internal_worker - .as_ref() - .map(|worker| worker.session_id.as_str()) - else { - return; + pub(crate) fn has_pending_compaction_cleanup(&self) -> bool { + self.pending_compaction_cleanup + .lock() + .expect("pending compaction cleanup mutex poisoned") + .is_some() + } + + fn fail_closed_segment_activation(&self, authority: SegmentActivationGuard) { + *self + .failed_activation_authority + .lock() + .expect("failed activation authority mutex poisoned") = Some(authority); + self.segment_state.fail_closed_activation(); + } + + pub(crate) fn has_failed_segment_activation(&self) -> bool { + self.segment_state.activation_failed.load(Ordering::Acquire) + } + + pub(crate) async fn retry_pending_compaction_cleanup(&self) -> Result<(), WorkerError> { + self.release_pending_compaction_service().await + } + + pub(crate) async fn finish_pending_compaction_cleanup_for_shutdown( + &self, + ) -> Result<(), WorkerError> { + loop { + match self.retry_pending_compaction_cleanup().await { + Ok(()) => return Ok(()), + Err(error) => { + warn!( + error = %error, + "shutdown remains blocked on compaction service cleanup" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + } + } + + async fn release_pending_compaction_service(&self) -> Result<(), WorkerError> { + let cleanup = self + .pending_compaction_cleanup + .lock() + .expect("pending compaction cleanup mutex poisoned") + .clone(); + let Some(cleanup) = cleanup else { + return Ok(()); }; let Some(registry) = &self.internal_worker_registry else { - return; + return Err(WorkerError::CompactionCleanupPending { + source: ScopeLockError::Io(std::io::Error::other( + "compaction registry unavailable during cleanup", + )), + }); }; - if let Err(error) = registry.stop_service(session_id).await { - warn!( - compaction_id = %lifecycle.compaction_id, - internal_worker_session_id = %session_id, - error = %error, - "failed to release terminal compaction Internal Worker" - ); + match registry.stop_service(&cleanup.session_id).await { + Ok(_) => { + let mut pending = self + .pending_compaction_cleanup + .lock() + .expect("pending compaction cleanup mutex poisoned"); + if pending + .as_ref() + .is_some_and(|pending| pending.session_id == cleanup.session_id) + { + *pending = None; + } + Ok(()) + } + Err(source) => { + warn!( + compaction_id = %cleanup.compaction_id, + error = %source, + "compaction service cleanup failed; retaining cleanup authority" + ); + Err(WorkerError::CompactionCleanupPending { + source: ScopeLockError::Io(source), + }) + } } } + async fn release_compaction_service( + &self, + lifecycle: &mut CompactionLifecycle, + ) -> Result<(), WorkerError> { + self.release_pending_compaction_service().await?; + lifecycle.internal_worker = None; + Ok(()) + } + async fn compact_impl( &mut self, retained_tokens: u64, @@ -5345,6 +5524,14 @@ impl Worker { .map_err(|error| WorkerError::InvalidState(error.to_string()))?; lifecycle.revision = lifecycle.revision.saturating_add(1); lifecycle.internal_worker = Some(internal_ref); + let cleanup = PendingCompactionCleanup { + session_id: handle.session_id_string(), + compaction_id: lifecycle.compaction_id.clone(), + }; + *self + .pending_compaction_cleanup + .lock() + .expect("pending compaction cleanup mutex poisoned") = Some(cleanup); self.set_compaction_progress(Some(InFlightCompaction { phase: CompactionPhase::Summarizing, started_at_ms: lifecycle.started_at_ms, @@ -5352,7 +5539,8 @@ impl Worker { })); if let Err(error) = handle.send(summary_input.text).await { - let _ = registry.remove_service(&handle.session_id_string()); + // Keep the registry record and Worker-owned cleanup authority. The + // outer compaction boundary performs the single stop attempt. return Err(WorkerError::InvalidState(error.to_string())); } match handle.wait_until_idle().await { @@ -5643,8 +5831,31 @@ impl Worker { // rotate: create on disk, swap location, reset the broadcast // sink so existing subscribers see the new `SegmentStart // { compacted_from }` and reset their view. + // The Compactor has produced and validated its output. Retire it before + // activation so no internal process or registry authority survives into + // the replacement's live lifetime. Cleanup failures retain this future + // (and Busy state) until a retry succeeds. + self.release_compaction_service(lifecycle).await?; + + self.set_compaction_progress(Some(InFlightCompaction { + phase: CompactionPhase::Committing, + started_at_ms: lifecycle.started_at_ms, + trigger, + })); + + // Activation lock order is append barrier -> allocation table -> metadata. + // From this point through live publication, every accepted Submit/Notify + // either completed on the source Segment or waits to observe the fully + // activated replacement. + let segment_state = Arc::clone(&self.segment_state); + let _append_guard = segment_state + .append_lock + .lock() + .expect("segment append lock poisoned"); + segment_state.ensure_append_allowed()?; + let new_segment_id = session_store::new_segment_id(); - let old_loc = self.segment_state.location(); + let old_loc = segment_state.location(); let source_turn_count = self.engine.as_ref().unwrap().turn_count(); let w = self.engine.as_ref().unwrap(); let entry = LogEntry::AnnotatedSegmentStart { @@ -5706,11 +5917,6 @@ impl Worker { })?, }); } - self.set_compaction_progress(Some(InFlightCompaction { - phase: CompactionPhase::Committing, - started_at_ms: lifecycle.started_at_ms, - trigger, - })); self.store .create_segment(old_loc.session_id, new_segment_id, &initial_entries)?; @@ -5724,21 +5930,34 @@ impl Worker { session_id: old_loc.session_id, segment_id: new_segment_id, }; - // The writer lease is stable for the Worker lifetime and is keyed by - // worker name. Compaction must not transfer or rewrite that lease: the - // active Segment is derived exclusively from the CAS-protected Worker - // metadata pointer below. A lost CAS therefore leaves every live and - // durable authority on the previous Segment. + // The machine-wide allocation lock is acquired only after staging, but + // before the durable commit point. Restore admission therefore cannot + // race the metadata CAS and allocation hand-off. + let allocation_activation = self + .scope_allocation + .as_ref() + .map(|allocation| { + allocation.begin_segment_activation(old_loc.segment_id, new_segment_id) + }) + .transpose()?; self.compare_and_swap_worker_metadata_segment(old_loc, new_location)?; + if let Some(allocation_activation) = allocation_activation { + if let Err(error) = allocation_activation.commit() { + let (source, authority) = error.into_parts(); + // Metadata already names the complete replacement. Do not allow + // the old in-memory writer to accept another byte while durable + // authorities disagree; retain the allocation lock so restore + // admission cannot register a competing writer. + self.fail_closed_segment_activation(authority); + return Err(WorkerError::SegmentActivationIncomplete { source }); + } + } // All live mutations after the durable commit are infallible and happen // before the replacement SegmentStart is broadcast. This keeps the // append destination, Session projection, Engine cache identity, and // reconnect snapshot on one side of the same activation boundary. - self.segment_state.set_location(SegmentLocation { - session_id: old_loc.session_id, - segment_id: new_segment_id, - }); + self.segment_state.set_location(new_location); self.segment_state .set_entries_written(initial_entries.len()); self.user_segments = retained_user_segments; @@ -5764,7 +5983,6 @@ impl Worker { lifecycle.revision = lifecycle.revision.saturating_add(1); lifecycle.state = CompactionLifecycleState::Done; lifecycle.ended_at_ms = Some(segment_log::now_millis()); - self.set_compaction_progress(None); Ok(( new_segment_id, summary_text, @@ -5942,12 +6160,14 @@ where alerter: None, working_event_tx: None, internal_worker_registry: None, + pending_compaction_cleanup: Mutex::new(None), in_flight: None, ai_activity_counter: Arc::new(AtomicUsize::new(0)), pending_notifies: NotifyBuffer::new(), notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY, pending_attachments: Arc::new(Mutex::new(Vec::::new())), pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), + failed_activation_authority: Mutex::new(None), scope_allocation: Some(scope_allocation), callback_socket: None, runtime_ticket_role: None, @@ -6028,12 +6248,14 @@ where alerter: None, working_event_tx: None, internal_worker_registry: None, + pending_compaction_cleanup: Mutex::new(None), in_flight: None, ai_activity_counter: Arc::new(AtomicUsize::new(0)), pending_notifies: NotifyBuffer::new(), notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY, pending_attachments: Arc::new(Mutex::new(Vec::::new())), pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), + failed_activation_authority: Mutex::new(None), scope_allocation: None, callback_socket: None, runtime_ticket_role: None, @@ -6149,12 +6371,14 @@ where alerter: None, working_event_tx: None, internal_worker_registry: None, + pending_compaction_cleanup: Mutex::new(None), in_flight: None, ai_activity_counter: Arc::new(AtomicUsize::new(0)), pending_notifies: NotifyBuffer::new(), notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY, pending_attachments: Arc::new(Mutex::new(Vec::::new())), pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), + failed_activation_authority: Mutex::new(None), scope_allocation: Some(scope_allocation), callback_socket: Some(callback_socket), runtime_ticket_role: None, @@ -6527,12 +6751,14 @@ where alerter: None, working_event_tx: None, internal_worker_registry: None, + pending_compaction_cleanup: Mutex::new(None), in_flight: None, ai_activity_counter: Arc::new(AtomicUsize::new(0)), pending_notifies: NotifyBuffer::new(), notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY, pending_attachments: Arc::new(Mutex::new(Vec::::new())), pending_committed_history: Arc::new(Mutex::new(std::collections::VecDeque::new())), + failed_activation_authority: Mutex::new(None), scope_allocation: Some(scope_allocation), callback_socket: None, runtime_ticket_role: None, @@ -7169,7 +7395,9 @@ fn compact_failure_category(error: &WorkerError) -> CompactFailureCategory { WorkerError::CompactResultContextTooLarge { .. } => { CompactFailureCategory::ResultContextTooLarge } - WorkerError::WorkerStore(_) | WorkerError::CompactActiveSegmentChanged => { + WorkerError::WorkerStore(_) + | WorkerError::CompactActiveSegmentChanged + | WorkerError::SegmentActivationIncomplete { .. } => { CompactFailureCategory::ActiveSegmentCommit } WorkerError::Store(_) => CompactFailureCategory::Storage, @@ -7275,6 +7503,20 @@ pub enum WorkerError { #[error(transparent)] ScopeLock(#[from] ScopeLockError), + #[error( + "Segment activation could not be completed; the Worker must be restored before accepting more input" + )] + SegmentActivationIncomplete { + #[source] + source: ScopeLockError, + }, + + #[error("compaction cleanup is pending; retry or shut down the Worker")] + CompactionCleanupPending { + #[source] + source: ScopeLockError, + }, + #[error(transparent)] PromptCatalog(#[from] CatalogError), @@ -8322,6 +8564,231 @@ mod build_summary_prompt_tests { assert_eq!(restored.active_run_turn_count, None); } + #[tokio::test] + async fn post_metadata_allocation_failure_fences_appends_and_restart_converges() { + 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_segment = worker.segment_id(); + let replacement = session_store::new_segment_id(); + let old_location = SegmentLocation { + session_id: worker.session_id(), + segment_id: old_segment, + }; + let replacement_location = SegmentLocation { + session_id: worker.session_id(), + segment_id: replacement, + }; + 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_segment, + ) + .unwrap(); + let activation = allocation + .begin_segment_activation(old_segment, replacement) + .unwrap(); + worker.scope_allocation = Some(allocation); + worker + .compare_and_swap_worker_metadata_segment(old_location, replacement_location) + .unwrap(); + assert_eq!( + *persisted_metadata.lock().unwrap(), + WorkerActiveSegmentRef::active_segment( + replacement_location.session_id, + replacement_location.segment_id, + ) + ); + worker_allocation::fail_next_save_for_test(&lock_path); + let (source, authority) = activation.commit().unwrap_err().into_parts(); + assert!(matches!(source, ScopeLockError::Io(_))); + worker.fail_closed_segment_activation(authority); + assert!(worker.segment_state.ensure_append_allowed().is_err()); + + let (sent, received) = std::sync::mpsc::channel(); + let teardown = std::thread::spawn(move || { + drop(worker); + sent.send(()).unwrap(); + }); + received + .recv_timeout(Duration::from_secs(1)) + .expect("Worker teardown must not deadlock on nested allocation locks"); + teardown.join().unwrap(); + + // Process-stop releases the retained authority and the old allocation; + // restart follows the replacement named by committed metadata. + 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, + ) + .unwrap(); + drop( + recovered + .begin_segment_activation(replacement, session_store::new_segment_id()) + .unwrap(), + ); + 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(); + 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(); + let registry = crate::spawn::registry::SpawnedWorkerRegistry::new_for_internal_services(); + let (cleanup_session_id, _events) = registry.install_service_for_test(); + registry.fail_service_stops_for_test(&cleanup_session_id, 4); + worker.internal_worker_registry = Some(Arc::clone(®istry)); + *worker + .pending_compaction_cleanup + .lock() + .expect("pending cleanup lock poisoned") = Some(PendingCompactionCleanup { + session_id: cleanup_session_id.clone(), + compaction_id: "compaction".into(), + }); + + let cleanup = worker.finish_pending_compaction_cleanup_for_shutdown(); + tokio::pin!(cleanup); + tokio::select! { + result = &mut cleanup => panic!("shutdown cleanup completed before injected failures: {result:?}"), + _ = tokio::time::sleep(Duration::from_millis(75)) => {} + } + assert!(worker.has_pending_compaction_cleanup()); + assert!(registry.has_service_for_test(&cleanup_session_id)); + tokio::time::timeout(Duration::from_secs(1), cleanup) + .await + .expect("shutdown cleanup must converge after retry") + .unwrap(); + assert!(!worker.has_pending_compaction_cleanup()); + assert!(!registry.has_service_for_test(&cleanup_session_id)); + } + #[tokio::test] async fn auto_fork_checkpoints_interrupted_run_budget_for_restore() { let dir = tempfile::tempdir().unwrap(); @@ -8373,8 +8840,9 @@ mod build_summary_prompt_tests { active_turn_count: 3, total_turn_count: 7, .. - } - ] + }, + LogEntry::Extension { domain, .. } + ] if domain == SESSION_PENDING_ACTIVATIONS_EXTENSION_DOMAIN )); let restored = session_store::collect_state(&fork_entries); assert!(restored.last_run_interrupted); @@ -10068,6 +10536,82 @@ mod build_summary_prompt_tests { ); } + #[test] + fn submit_and_notification_wait_for_activation_barrier_then_append_to_replacement() { + let temp = tempfile::tempdir().unwrap(); + let handle = PendingSubmissionHandle::for_test(temp.path()); + let old_location = handle.writer.state.location(); + let replacement = session_store::new_segment_id(); + handle + .writer + .store + .create_segment(old_location.session_id, replacement, &[]) + .unwrap(); + let old_entries = handle + .writer + .store + .read_all(old_location.session_id, old_location.segment_id) + .unwrap() + .len(); + let append_guard = handle.writer.state.append_lock.lock().unwrap(); + let submit_handle = handle.clone(); + let notify_handle = handle.clone(); + let (started, starts) = std::sync::mpsc::channel(); + let (finished, finishes) = std::sync::mpsc::channel(); + let submit_started = started.clone(); + let submit_finished = finished.clone(); + let submit = std::thread::spawn(move || { + submit_started.send(()).unwrap(); + submit_handle + .accept( + "barrier-submit".into(), + vec![Segment::text("queued")], + false, + ) + .unwrap(); + submit_finished.send(()).unwrap(); + }); + let notify = std::thread::spawn(move || { + started.send(()).unwrap(); + notify_handle + .accept_notification("barrier-notify".into(), "completed".into()) + .unwrap(); + finished.send(()).unwrap(); + }); + starts.recv_timeout(Duration::from_secs(1)).unwrap(); + starts.recv_timeout(Duration::from_secs(1)).unwrap(); + assert!(finishes.recv_timeout(Duration::from_millis(50)).is_err()); + + handle.writer.state.set_location(SegmentLocation { + session_id: old_location.session_id, + segment_id: replacement, + }); + drop(append_guard); + finishes.recv_timeout(Duration::from_secs(1)).unwrap(); + finishes.recv_timeout(Duration::from_secs(1)).unwrap(); + submit.join().unwrap(); + notify.join().unwrap(); + + assert_eq!( + handle + .writer + .store + .read_all(old_location.session_id, old_location.segment_id) + .unwrap() + .len(), + old_entries + ); + assert_eq!( + handle + .writer + .store + .read_all(old_location.session_id, replacement) + .unwrap() + .len(), + 2 + ); + } + #[test] fn pending_submission_queue_is_durable_idempotent_and_bounded() { let temp = tempfile::tempdir().unwrap(); diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index 416437931..c6c37367c 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -344,7 +344,10 @@ target = "./" permission = "write" "#; -async fn make_worker_with_manifest(manifest_toml: &str, client: C) -> Worker +async fn make_worker_with_manifest_and_store( + manifest_toml: &str, + client: C, +) -> (Worker, TestStore) where C: LlmClient + Clone + Send + Sync + 'static, { @@ -364,6 +367,7 @@ where let worker = Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client); + let observed_store = store.clone(); let mut worker = Worker::new( manifest, worker, @@ -375,7 +379,16 @@ where .await .unwrap(); worker.enable_worker_metadata_write_through().unwrap(); - worker + (worker, observed_store) +} + +async fn make_worker_with_manifest(manifest_toml: &str, client: C) -> Worker +where + C: LlmClient + Clone + Send + Sync + 'static, +{ + make_worker_with_manifest_and_store(manifest_toml, client) + .await + .0 } async fn make_worker(client: MockClient) -> Worker { @@ -922,7 +935,7 @@ async fn pre_run_compact_publishes_runtime_progress_phases() { let session_before = worker.session_id(); let segment_before = worker.segment_id(); - worker.try_pre_run_compact().await; + worker.try_pre_run_compact().await.unwrap(); assert_eq!(worker.session_id(), session_before); assert_ne!(worker.segment_id(), segment_before); @@ -1083,7 +1096,7 @@ async fn pre_run_compact_failure_clears_runtime_progress() { let _ = drain(&mut rx); // Best-effort: returns Ok(()) even on failure and clears runtime progress. - worker.try_pre_run_compact().await; + worker.try_pre_run_compact().await.unwrap(); let events = drain(&mut rx); assert!(events.iter().any(|event| matches!( @@ -1104,8 +1117,11 @@ async fn pre_run_compact_failure_clears_runtime_progress() { #[tokio::test] async fn manual_compact_cancel_clears_progress_before_returning_idle() { - let worker = - make_worker_with_manifest(POST_RUN_MANIFEST_TOML, BlockingCompactClient::new()).await; + let (worker, observed_store) = + make_worker_with_manifest_and_store(POST_RUN_MANIFEST_TOML, BlockingCompactClient::new()) + .await; + let source_session = worker.session_id(); + let source_segment = worker.segment_id(); let runtime_tmp = tempfile::tempdir().unwrap(); let bash_output_dir = runtime_tmp.path().join("bash-output"); let (handle, shutdown_receiver) = @@ -1201,6 +1217,15 @@ async fn manual_compact_cancel_clears_progress_before_returning_idle() { break; } } + let completion_request_id = "subworker-completion-during-compact"; + handle + .send(Method::NotifyTracked { + notification_request_id: completion_request_id.into(), + message: "subworker completed".into(), + source: protocol::AuthenticatedInputSource::UntrustedWire, + }) + .await + .expect("send completion notification during compact"); let shutdown = protocol::WorkerCommandEnvelope::new(4); handle .send(Method::Shutdown { command: shutdown }) @@ -1230,6 +1255,18 @@ async fn manual_compact_cancel_clears_progress_before_returning_idle() { .await .expect("controller shutdown timeout") .expect("shutdown confirmation"); + + let persisted = observed_store + .read_all(source_session, source_segment) + .expect("read source Segment after shutdown"); + assert!(persisted.iter().any(|entry| { + matches!( + entry, + LogEntry::Extension { domain, payload, .. } + if domain == "worker.pending_activations.v1" + && payload.to_string().contains(completion_request_id) + ) + })); } #[tokio::test] diff --git a/docs/design/compaction.md b/docs/design/compaction.md index e6dd2b195..6d33041b5 100644 --- a/docs/design/compaction.md +++ b/docs/design/compaction.md @@ -32,6 +32,44 @@ Compaction should preserve persisted reasoning history and avoid serializing unv The important property is explainability: after compaction, records should still show what summary replaced which older context and why future turns can rely on it. +## Activation boundary + +A compacted Segment is staged completely before it can become authoritative. The +activation boundary uses this lock order everywhere a live Worker changes Segment: + +1. stop the compaction service and remove its registry record; +2. acquire the Worker's append barrier; +3. write the final pending-submission checkpoint and stage the replacement Segment; +4. acquire and preflight the machine-wide Worker allocation lock; +5. compare-and-swap `metadata.json` from the source Segment to the replacement; +6. atomically replace the allocation table while still holding its separate lock; +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. 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 +unusable and retains the machine-wide allocation lock: further acceptance fails +closed and restore admission cannot register a competing writer until process +teardown, after which restore converges from the replacement named by metadata. +The Worker must likewise remain non-idle while terminal compaction-service cleanup +is pending. Cleanup authority is stored on the Worker (not only on one compaction +future). A failed stop returns a typed cleanup-pending outcome without clearing that +authority; the next compaction boundary retries it. The authority and lifecycle +reference are cleared exactly once after the registry record is gone, before +activation or a terminal idle-capable result. While either activation repair or +cleanup attention is outstanding, the controller fences Submit/Notify dispatch and +does not start queued work. Shutdown is itself a cleanup barrier: it retries the +Worker-owned stop authority with backoff and does not emit its terminal event until +the registry record is gone. + ## Metrics and comparison procedure Compaction measurements stay out of the ordinary transcript. They are appended as