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!(