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

worker: make segment activation atomic

Browse Source
This commit is contained in:
Hare
2026-09-18 06:08:35 +09:00
parent 195aeb69be
commit 34085b121d
6 changed files with 361 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+1
View File
@@ -2147,6 +2147,7 @@ async fn controller_loop<C, St>(
Err(WorkerError::Store(_))
| Err(WorkerError::WorkerStore(_))
| Err(WorkerError::InvalidState(_))
| Err(WorkerError::SegmentActivationIncomplete { .. })
) {
set_controller_state(
&shared_state,
crates/worker/src/runtime/worker_allocation.rs
+2 -2
View File
@@ -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,
crates/worker/src/runtime/worker_allocation/error.rs
+8
View File
@@ -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<SegmentId>,
},
}
crates/worker/src/runtime/worker_allocation/lifecycle.rs
+191 -64
View File
@@ -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<SegmentActivationGuard, ScopeLockError> {
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)
);
}
}
crates/worker/src/worker.rs
+132 -39
View File
@@ -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<SegmentLocation>,
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
{
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
// 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<C: LlmClient + 'static, St: Store> Worker<C, St> {
})?,
});
}
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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);