fix: make compaction Segment activation fail closed
This commit is contained in:
@@ -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();
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
@@ -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(®istry));
|
||||
*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();
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user