fix: make compaction Segment activation fail closed

This commit is contained in:
2026-09-18 19:45:12 +09:00
10 changed files with 1683 additions and 156 deletions
+401 -5
View File
@@ -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<St>(
method: Method,
pending_submissions: &PendingSubmissionHandle<St>,
working_event_tx: &broadcast::Sender<Event>,
) -> Option<Method>
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<Event>,
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<C, St>(
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<C, St>(
)
.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<C, St>(
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<C, St>(
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<C, St>(
);
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<C, St>(
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<C, St>(
});
}
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::<WorkerRunResult, _>(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::<WorkerRunResult, _>(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;
@@ -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};
@@ -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>,
},
}
@@ -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<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<(), 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<ScopeRule>,
segment_id: SegmentId,
) -> Result<ScopeAllocationGuard, ScopeLockError> {
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<Option<SegmentLockInfo>,
#[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)
);
}
}
@@ -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::collections::HashMap<PathBuf, SaveFaultPoint>>,
> = 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<PathBuf> {
/// 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();
+58 -2
View File
@@ -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<Vec<InternalSpawnedWorkerRecord>>,
service_records: std::sync::Mutex<Vec<InternalServiceWorkerRecord>>,
#[cfg(test)]
fail_service_stops: std::sync::Mutex<HashMap<String, usize>>,
internal_names: std::sync::Mutex<HashSet<String>>,
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<protocol::Event>) {
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<bool> {
#[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()
+30 -5
View File
@@ -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));
+601 -57
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};
@@ -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<SegmentLocation>,
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<session_store::FsStore> {
},
}
}
#[cfg(test)]
pub(crate) fn persisted_entries_for_test(&self) -> Vec<LogEntry> {
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<String>;
}
#[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<C: LlmClient, St: Store> {
/// 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<Arc<crate::spawn::registry::SpawnedWorkerRegistry>>,
/// Cleanup authority survives cancellation of an individual compaction future.
pending_compaction_cleanup: Mutex<Option<PendingCompactionCleanup>>,
in_flight: Option<InFlightEvents>,
/// 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<C: LlmClient, St: Store> {
/// interceptor before Agen applies them to live typed history.
pending_committed_history:
Arc<Mutex<std::collections::VecDeque<HistoryEntry<SessionHistoryMetadata>>>>,
/// 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<Option<SegmentActivationGuard>>,
/// 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<C: LlmClient, St: Store> {
log_writer: Option<Arc<dyn SystemItemCommitter>>,
}
impl<C: LlmClient, St: Store> Drop for Worker<C, St> {
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<C: LlmClient + 'static, St: Store + 'static> Worker<C, St> {
pub async fn stop_feature_runtime(&mut self, reason: impl Into<String>) {
if let Some(hooks) = self.hook_registry.clone() {
@@ -2443,12 +2497,14 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
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::<SystemItem>::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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<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 +4509,47 @@ 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(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<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// 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<C: LlmClient + 'static, St: Store> Worker<C, St> {
completed,
"automatic compaction must complete a claimed attempt"
);
Ok(())
}
Err(e) => {
warn!(error = %e, "Proactive pre-run compaction failed");
@@ -4710,6 +4799,15 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
completed,
"automatic compaction must complete a claimed attempt"
);
if matches!(
&e,
WorkerError::SegmentActivationIncomplete { .. }
| WorkerError::CompactionCleanupPending { .. }
) {
Err(e)
} else {
Ok(())
}
}
}
}
@@ -4995,6 +5093,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
mut cancel: Option<tokio::sync::watch::Receiver<bool>>,
trigger: CompactionTrigger,
) -> Result<SegmentId, WorkerError> {
self.release_pending_compaction_service().await?;
let _rewrite_guard = self
.prepare_session_rewrite(SessionRewriteKind::Compact)
.await?;
@@ -5058,15 +5157,23 @@ 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) => {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
.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<C: LlmClient + 'static, St: Store> Worker<C, St> {
}));
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<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 +5917,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 +5930,34 @@ 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(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<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,
@@ -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::<SystemItem>::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::<SystemItem>::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::<SystemItem>::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::<SystemItem>::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(&registry));
*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();
+43 -6
View File
@@ -344,7 +344,10 @@ target = "./"
permission = "write"
"#;
async fn make_worker_with_manifest<C>(manifest_toml: &str, client: C) -> Worker<C, TestStore>
async fn make_worker_with_manifest_and_store<C>(
manifest_toml: &str,
client: C,
) -> (Worker<C, TestStore>, 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<C>(manifest_toml: &str, client: C) -> Worker<C, TestStore>
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<MockClient, TestStore> {
@@ -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]
+38
View File
@@ -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