diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 7fe3b02f4..2dda1ad49 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -2147,6 +2147,7 @@ async fn controller_loop( Err(WorkerError::Store(_)) | Err(WorkerError::WorkerStore(_)) | Err(WorkerError::InvalidState(_)) + | Err(WorkerError::SegmentActivationIncomplete { .. }) ) { set_controller_state( &shared_state, diff --git a/crates/worker/src/runtime/worker_allocation.rs b/crates/worker/src/runtime/worker_allocation.rs index 1eb3874bb..692d1ed82 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, SegmentLockInfo, adopt_allocation, install_top_level, - install_top_level_with_deny, lookup_segment, update_segment, + ScopeAllocationGuard, 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/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..15e02fd10 100644 --- a/crates/worker/src/runtime/worker_allocation/lifecycle.rs +++ b/crates/worker/src/runtime/worker_allocation/lifecycle.rs @@ -20,6 +20,18 @@ 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, +} + impl ScopeAllocationGuard { pub fn worker_name(&self) -> &str { &self.worker_name @@ -28,6 +40,79 @@ 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<(), ScopeLockError> { + 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(ScopeLockError::SegmentChanged { + worker_name: self.worker_name, + expected: self.expected, + actual, + }); + }; + allocation.segment_id = Some(self.replacement); + self.guard.save()?; + Ok(()) + } } impl Drop for ScopeAllocationGuard { @@ -114,46 +199,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)] @@ -283,32 +328,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 +349,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 +366,103 @@ 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 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/worker.rs b/crates/worker/src/worker.rs index af46ab494..459850acf 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}; @@ -1055,6 +1055,7 @@ pub struct SegmentState { location: ArcSwap, entries_written: AtomicUsize, append_lock: Mutex<()>, + activation_failed: AtomicBool, } impl SegmentState { @@ -1066,6 +1067,7 @@ impl SegmentState { }), entries_written: AtomicUsize::new(entries_written), append_lock: Mutex::new(()), + activation_failed: AtomicBool::new(false), }) } @@ -1093,6 +1095,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 +1238,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 +1455,7 @@ where .append_lock .lock() .expect("segment append lock poisoned"); + self.writer.state.ensure_append_allowed()?; let mut current = self .state .lock() @@ -1580,6 +1597,7 @@ where .append_lock .lock() .expect("segment append lock poisoned"); + self.writer.state.ensure_append_allowed()?; let mut state = self .state .lock() @@ -4414,8 +4432,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 +4471,48 @@ 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(source) = allocation_activation.commit() { + segment_state.fail_closed_activation(); + 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(()) } @@ -5058,12 +5107,12 @@ 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) => { @@ -5082,8 +5131,8 @@ impl Worker { CompactionLifecycleState::Failed }; lifecycle.ended_at_ms = Some(segment_log::now_millis()); - self.set_compaction_progress(None); self.release_compaction_service(&lifecycle).await; + self.set_compaction_progress(None); Err(error) } } @@ -5100,13 +5149,18 @@ impl Worker { let Some(registry) = &self.internal_worker_registry else { return; }; - 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" - ); + loop { + match registry.stop_service(session_id).await { + Ok(_) => return, + Err(error) => { + warn!( + compaction_id = %lifecycle.compaction_id, + error = %error, + "compaction service cleanup failed; retaining cleanup authority and retrying" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + } } } @@ -5643,8 +5697,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 +5783,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 +5796,32 @@ 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(source) = allocation_activation.commit() { + // 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(); + 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 +5847,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, @@ -7169,7 +7251,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 +7359,14 @@ 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(transparent)] PromptCatalog(#[from] CatalogError), @@ -8373,8 +8465,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); diff --git a/docs/design/compaction.md b/docs/design/compaction.md index e6dd2b195..25a82c5cf 100644 --- a/docs/design/compaction.md +++ b/docs/design/compaction.md @@ -32,6 +32,33 @@ 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. update the allocation table while still holding its 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. 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: 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. + ## Metrics and comparison procedure Compaction measurements stay out of the ordinary transcript. They are appended as