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);
docs/design/compaction.md
+27
View File
@@ -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