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

worker: retain failed activation authority

Browse Source
This commit is contained in:
Hare
2026-09-18 06:27:11 +09:00
parent 34085b121d
commit fc24cb672e
5 changed files with 187 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/runtime/worker_allocation.rs
+2 -2
View File
@@ -19,8 +19,8 @@ pub use conflict::{
};
pub use error::ScopeLockError;
pub use lifecycle::{
ScopeAllocationGuard, SegmentActivationGuard, SegmentLockInfo, adopt_allocation,
install_top_level, install_top_level_with_deny, lookup_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,
crates/worker/src/runtime/worker_allocation/lifecycle.rs
+87 -6
View File
@@ -32,6 +32,29 @@ pub struct SegmentActivationGuard {
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
@@ -88,7 +111,7 @@ impl SegmentActivationGuard {
/// Persist the replacement Segment while retaining exclusive allocation
/// authority. This is intentionally consumed so a successful metadata CAS
/// cannot accidentally be followed by an unlocked update.
pub fn commit(mut self) -> Result<(), ScopeLockError> {
pub fn commit(mut self) -> Result<(), SegmentActivationCommitError> {
let actual = self
.guard
.data()
@@ -103,14 +126,22 @@ impl SegmentActivationGuard {
.iter_mut()
.find(|allocation| allocation.worker_name == self.worker_name)
else {
return Err(ScopeLockError::SegmentChanged {
worker_name: self.worker_name,
expected: self.expected,
actual,
return Err(SegmentActivationCommitError {
source: ScopeLockError::SegmentChanged {
worker_name: self.worker_name.clone(),
expected: self.expected,
actual,
},
authority: self,
});
};
allocation.segment_id = Some(self.replacement);
self.guard.save()?;
if let Err(source) = self.guard.save() {
return Err(SegmentActivationCommitError {
source: ScopeLockError::Io(source),
authority: self,
});
}
Ok(())
}
}
@@ -399,6 +430,56 @@ mod tests {
);
}
#[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 SegmentState 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 dropping_segment_activation_preserves_source_allocation() {
let dir = TempDir::new().unwrap();
crates/worker/src/worker.rs
+89 -23
View File
@@ -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;
@@ -1056,6 +1058,7 @@ pub struct SegmentState {
entries_written: AtomicUsize,
append_lock: Mutex<()>,
activation_failed: AtomicBool,
activation_failure_authority: Mutex<Option<SegmentActivationGuard>>,
}
impl SegmentState {
@@ -1068,6 +1071,7 @@ impl SegmentState {
entries_written: AtomicUsize::new(entries_written),
append_lock: Mutex::new(()),
activation_failed: AtomicBool::new(false),
activation_failure_authority: Mutex::new(None),
})
}
@@ -1104,7 +1108,11 @@ impl SegmentState {
Ok(())
}
fn fail_closed_activation(&self) {
fn fail_closed_activation(&self, authority: SegmentActivationGuard) {
*self
.activation_failure_authority
.lock()
.expect("Segment activation failure authority mutex poisoned") = Some(authority);
self.activation_failed.store(true, Ordering::Release);
}
@@ -2086,6 +2094,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
@@ -2184,6 +2198,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
@@ -2461,6 +2477,7 @@ 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(),
@@ -3629,7 +3646,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(())
}
@@ -4503,8 +4520,9 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.transpose()?;
self.compare_and_swap_worker_metadata_segment(loc, new_location)?;
if let Some(allocation_activation) = allocation_activation {
if let Err(source) = allocation_activation.commit() {
segment_state.fail_closed_activation();
if let Err(error) = allocation_activation.commit() {
let (source, authority) = error.into_parts();
segment_state.fail_closed_activation(authority);
return Err(WorkerError::SegmentActivationIncomplete { source });
}
}
@@ -4711,16 +4729,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")
}
@@ -4741,6 +4761,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");
@@ -4759,6 +4780,11 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
completed,
"automatic compaction must complete a claimed attempt"
);
if matches!(e, WorkerError::SegmentActivationIncomplete { .. }) {
Err(e)
} else {
Ok(())
}
}
}
}
@@ -5044,6 +5070,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?;
@@ -5131,30 +5158,43 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
CompactionLifecycleState::Failed
};
lifecycle.ended_at_ms = Some(segment_log::now_millis());
self.release_compaction_service(&lifecycle).await;
self.release_compaction_service(&mut lifecycle).await;
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 {
async fn release_pending_compaction_service(&self) {
let cleanup = self
.pending_compaction_cleanup
.lock()
.expect("pending compaction cleanup mutex poisoned")
.clone();
let Some(cleanup) = cleanup else {
return;
};
let Some(registry) = &self.internal_worker_registry else {
return;
};
loop {
match registry.stop_service(session_id).await {
Ok(_) => return,
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;
}
return;
}
Err(error) => {
warn!(
compaction_id = %lifecycle.compaction_id,
compaction_id = %cleanup.compaction_id,
error = %error,
"compaction service cleanup failed; retaining cleanup authority and retrying"
);
@@ -5164,6 +5204,18 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
}
}
async fn release_compaction_service(&self, lifecycle: &mut CompactionLifecycle) {
self.release_pending_compaction_service().await;
if self
.pending_compaction_cleanup
.lock()
.expect("pending compaction cleanup mutex poisoned")
.is_none()
{
lifecycle.internal_worker = None;
}
}
async fn compact_impl(
&mut self,
retained_tokens: u64,
@@ -5399,6 +5451,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,
@@ -5808,11 +5868,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.transpose()?;
self.compare_and_swap_worker_metadata_segment(old_loc, new_location)?;
if let Some(allocation_activation) = allocation_activation {
if let Err(source) = allocation_activation.commit() {
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; restore will converge from metadata.
segment_state.fail_closed_activation();
// authorities disagree; retain the allocation lock so restore
// admission cannot register a competing writer.
segment_state.fail_closed_activation(authority);
return Err(WorkerError::SegmentActivationIncomplete { source });
}
}
@@ -6024,6 +6086,7 @@ 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(),
@@ -6110,6 +6173,7 @@ 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(),
@@ -6231,6 +6295,7 @@ 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(),
@@ -6609,6 +6674,7 @@ 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(),
crates/worker/tests/compact_events_test.rs
+2 -2
View File
@@ -922,7 +922,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 +1083,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!(
docs/design/compaction.md
+7 -4
View File
@@ -54,10 +54,13 @@ metadata CAS and allocation hand-off.
The metadata CAS is the commit point. Failures before it leave the source Segment
active. A failure updating allocation after it marks the old in-memory writer
unusable: further acceptance fails closed and process restart/restore converges from
the replacement named by metadata. The Worker must likewise remain non-idle while
terminal compaction-service cleanup is pending; cleanup authority is retained and
retried until the registry record is gone.
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), retried until the registry record is gone, and cleared exactly once before
activation or a terminal result.
## Metrics and comparison procedure