diff --git a/crates/worker/src/runtime/worker_allocation.rs b/crates/worker/src/runtime/worker_allocation.rs index 692d1ed82..726c294ad 100644 --- a/crates/worker/src/runtime/worker_allocation.rs +++ b/crates/worker/src/runtime/worker_allocation.rs @@ -19,8 +19,8 @@ pub use conflict::{ }; pub use error::ScopeLockError; pub use lifecycle::{ - ScopeAllocationGuard, SegmentActivationGuard, SegmentLockInfo, adopt_allocation, - install_top_level, install_top_level_with_deny, lookup_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, diff --git a/crates/worker/src/runtime/worker_allocation/lifecycle.rs b/crates/worker/src/runtime/worker_allocation/lifecycle.rs index 15e02fd10..2f758480c 100644 --- a/crates/worker/src/runtime/worker_allocation/lifecycle.rs +++ b/crates/worker/src/runtime/worker_allocation/lifecycle.rs @@ -32,6 +32,29 @@ pub struct SegmentActivationGuard { 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 @@ -88,7 +111,7 @@ 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<(), ScopeLockError> { + pub fn commit(mut self) -> Result<(), SegmentActivationCommitError> { let actual = self .guard .data() @@ -103,14 +126,22 @@ impl SegmentActivationGuard { .iter_mut() .find(|allocation| allocation.worker_name == self.worker_name) else { - return Err(ScopeLockError::SegmentChanged { - worker_name: self.worker_name, - expected: self.expected, - actual, + return Err(SegmentActivationCommitError { + source: ScopeLockError::SegmentChanged { + worker_name: self.worker_name.clone(), + expected: self.expected, + actual, + }, + authority: self, }); }; allocation.segment_id = Some(self.replacement); - self.guard.save()?; + if let Err(source) = self.guard.save() { + return Err(SegmentActivationCommitError { + source: ScopeLockError::Io(source), + authority: self, + }); + } Ok(()) } } @@ -399,6 +430,56 @@ mod tests { ); } + #[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 SegmentState 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 dropping_segment_activation_preserves_source_allocation() { let dir = TempDir::new().unwrap(); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 459850acf..7fc65102c 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -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; @@ -1056,6 +1058,7 @@ pub struct SegmentState { entries_written: AtomicUsize, append_lock: Mutex<()>, activation_failed: AtomicBool, + activation_failure_authority: Mutex>, } impl SegmentState { @@ -1068,6 +1071,7 @@ impl SegmentState { entries_written: AtomicUsize::new(entries_written), append_lock: Mutex::new(()), activation_failed: AtomicBool::new(false), + activation_failure_authority: Mutex::new(None), }) } @@ -1104,7 +1108,11 @@ impl SegmentState { Ok(()) } - fn fail_closed_activation(&self) { + fn fail_closed_activation(&self, authority: SegmentActivationGuard) { + *self + .activation_failure_authority + .lock() + .expect("Segment activation failure authority mutex poisoned") = Some(authority); self.activation_failed.store(true, Ordering::Release); } @@ -2086,6 +2094,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 @@ -2184,6 +2198,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 @@ -2461,6 +2477,7 @@ 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(), @@ -3629,7 +3646,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(()) } @@ -4503,8 +4520,9 @@ impl Worker { .transpose()?; self.compare_and_swap_worker_metadata_segment(loc, new_location)?; if let Some(allocation_activation) = allocation_activation { - if let Err(source) = allocation_activation.commit() { - segment_state.fail_closed_activation(); + if let Err(error) = allocation_activation.commit() { + let (source, authority) = error.into_parts(); + segment_state.fail_closed_activation(authority); return Err(WorkerError::SegmentActivationIncomplete { source }); } } @@ -4711,16 +4729,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") } @@ -4741,6 +4761,7 @@ impl Worker { completed, "automatic compaction must complete a claimed attempt" ); + Ok(()) } Err(e) => { warn!(error = %e, "Proactive pre-run compaction failed"); @@ -4759,6 +4780,11 @@ impl Worker { completed, "automatic compaction must complete a claimed attempt" ); + if matches!(e, WorkerError::SegmentActivationIncomplete { .. }) { + Err(e) + } else { + Ok(()) + } } } } @@ -5044,6 +5070,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?; @@ -5131,30 +5158,43 @@ impl Worker { CompactionLifecycleState::Failed }; lifecycle.ended_at_ms = Some(segment_log::now_millis()); - self.release_compaction_service(&lifecycle).await; + self.release_compaction_service(&mut lifecycle).await; 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 { + async fn release_pending_compaction_service(&self) { + let cleanup = self + .pending_compaction_cleanup + .lock() + .expect("pending compaction cleanup mutex poisoned") + .clone(); + let Some(cleanup) = cleanup else { return; }; let Some(registry) = &self.internal_worker_registry else { return; }; loop { - match registry.stop_service(session_id).await { - Ok(_) => return, + 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; + } + return; + } Err(error) => { warn!( - compaction_id = %lifecycle.compaction_id, + compaction_id = %cleanup.compaction_id, error = %error, "compaction service cleanup failed; retaining cleanup authority and retrying" ); @@ -5164,6 +5204,18 @@ impl Worker { } } + async fn release_compaction_service(&self, lifecycle: &mut CompactionLifecycle) { + self.release_pending_compaction_service().await; + if self + .pending_compaction_cleanup + .lock() + .expect("pending compaction cleanup mutex poisoned") + .is_none() + { + lifecycle.internal_worker = None; + } + } + async fn compact_impl( &mut self, retained_tokens: u64, @@ -5399,6 +5451,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, @@ -5808,11 +5868,13 @@ impl Worker { .transpose()?; self.compare_and_swap_worker_metadata_segment(old_loc, new_location)?; if let Some(allocation_activation) = allocation_activation { - if let Err(source) = allocation_activation.commit() { + 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; restore will converge from metadata. - segment_state.fail_closed_activation(); + // authorities disagree; retain the allocation lock so restore + // admission cannot register a competing writer. + segment_state.fail_closed_activation(authority); return Err(WorkerError::SegmentActivationIncomplete { source }); } } @@ -6024,6 +6086,7 @@ 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(), @@ -6110,6 +6173,7 @@ 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(), @@ -6231,6 +6295,7 @@ 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(), @@ -6609,6 +6674,7 @@ 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(), diff --git a/crates/worker/tests/compact_events_test.rs b/crates/worker/tests/compact_events_test.rs index 416437931..e88fc223d 100644 --- a/crates/worker/tests/compact_events_test.rs +++ b/crates/worker/tests/compact_events_test.rs @@ -922,7 +922,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 +1083,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!( diff --git a/docs/design/compaction.md b/docs/design/compaction.md index 25a82c5cf..01abe4fd7 100644 --- a/docs/design/compaction.md +++ b/docs/design/compaction.md @@ -54,10 +54,13 @@ 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: further acceptance fails closed and process restart/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 retained and -retried until the registry record is gone. +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), retried until the registry record is gone, and cleared exactly once before +activation or a terminal result. ## Metrics and comparison procedure