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

fix: hold feature task barrier across rewrites

Browse Source
This commit is contained in:
Hare
2026-09-04 16:34:15 +09:00
parent f1dc90621c
commit 60a5495ccd
2 changed files with 181 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/feature/background.rs
+169 -14
View File
@@ -6,7 +6,7 @@
//! feature was explicitly granted at install time.
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
@@ -25,6 +25,9 @@ const MAX_TASK_ATTEMPTS: u16 = 16;
const MAX_TASK_TIMEOUT_MS: u64 = 24 * 60 * 60 * 1_000;
const MAX_RETAINED_DIAGNOSTICS: usize = 128;
const TASK_SETTLE_TIMEOUT_MS: u64 = 30_000;
const TASK_SCOPE_ACCEPTING: u8 = 0;
const TASK_SCOPE_REWRITING: u8 = 1;
const TASK_SCOPE_STOPPING: u8 = 2;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
@@ -258,7 +261,7 @@ impl FeatureBackgroundTaskRegistryBuilder {
diagnostics: Mutex::new(Vec::new()),
next_execution_id: AtomicU64::new(1),
session_generation: Arc::new(AtomicU64::new(1)),
accepting: AtomicBool::new(true),
lifecycle: AtomicU8::new(TASK_SCOPE_ACCEPTING),
}),
}
}
@@ -277,7 +280,7 @@ struct RegistryInner {
diagnostics: Mutex<Vec<BackgroundTaskDiagnostic>>,
next_execution_id: AtomicU64,
session_generation: Arc<AtomicU64>,
accepting: AtomicBool,
lifecycle: AtomicU8,
}
impl Drop for RegistryInner {
@@ -297,6 +300,31 @@ pub struct FeatureBackgroundTaskRegistry {
inner: Arc<RegistryInner>,
}
/// Holds the task registry quiescent across the entire Session rewrite. Drop
/// reopens starts only when shutdown has not moved the scope to Stopping.
pub struct BackgroundTaskRewriteGuard {
inner: Arc<RegistryInner>,
}
impl std::fmt::Debug for BackgroundTaskRewriteGuard {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("BackgroundTaskRewriteGuard")
.finish_non_exhaustive()
}
}
impl Drop for BackgroundTaskRewriteGuard {
fn drop(&mut self) {
let _ = self.inner.lifecycle.compare_exchange(
TASK_SCOPE_REWRITING,
TASK_SCOPE_ACCEPTING,
Ordering::AcqRel,
Ordering::Acquire,
);
}
}
impl Default for FeatureBackgroundTaskRegistry {
fn default() -> Self {
FeatureBackgroundTaskRegistryBuilder::default().build()
@@ -361,7 +389,7 @@ impl FeatureBackgroundTaskRegistry {
task_name: &str,
invocation: HookInvocationContext,
) -> Result<BackgroundTaskStart, HookError> {
if !self.inner.accepting.load(Ordering::Acquire) {
if self.inner.lifecycle.load(Ordering::Acquire) != TASK_SCOPE_ACCEPTING {
return Err(HookError::new(
HookErrorCategory::ScopeDisposed,
"background task scope is stopping",
@@ -380,10 +408,11 @@ impl FeatureBackgroundTaskRegistry {
.running
.lock()
.expect("background tasks poisoned");
// `shutdown()` flips accepting before taking this same lock. Recheck
// while holding the spawn/drain serialization boundary so a starter
// that passed the optimistic check cannot publish a task after drain.
if !self.inner.accepting.load(Ordering::Acquire) {
// `shutdown()` and `begin_session_rewrite()` change lifecycle before
// taking this same lock. Recheck while holding the spawn/drain
// serialization boundary so a starter that passed the optimistic
// check cannot publish after either barrier.
if self.inner.lifecycle.load(Ordering::Acquire) != TASK_SCOPE_ACCEPTING {
return Err(HookError::new(
HookErrorCategory::ScopeDisposed,
"background task scope is stopping",
@@ -484,14 +513,45 @@ impl FeatureBackgroundTaskRegistry {
.clone()
}
pub async fn before_session_rewrite(&self) -> Result<(), HookError> {
pub async fn begin_session_rewrite(&self) -> Result<BackgroundTaskRewriteGuard, HookError> {
match self.inner.lifecycle.compare_exchange(
TASK_SCOPE_ACCEPTING,
TASK_SCOPE_REWRITING,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {}
Err(TASK_SCOPE_REWRITING) => {
return Err(HookError::new(
HookErrorCategory::Dependency,
"another Session rewrite already owns the background task barrier",
));
}
Err(_) => {
return Err(HookError::new(
HookErrorCategory::ScopeDisposed,
"background task scope is stopping",
));
}
}
let guard = BackgroundTaskRewriteGuard {
inner: Arc::clone(&self.inner),
};
self.settle(false).await?;
if self.inner.lifecycle.load(Ordering::Acquire) != TASK_SCOPE_REWRITING {
return Err(HookError::new(
HookErrorCategory::ScopeDisposed,
"background task scope stopped during Session rewrite preparation",
));
}
self.inner.session_generation.fetch_add(1, Ordering::AcqRel);
Ok(())
Ok(guard)
}
pub async fn shutdown(&self) -> Result<(), HookError> {
self.inner.accepting.store(false, Ordering::Release);
self.inner
.lifecycle
.store(TASK_SCOPE_STOPPING, Ordering::Release);
self.settle(true).await
}
@@ -668,7 +728,7 @@ async fn execute_task(
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
fn invocation() -> HookInvocationContext {
HookInvocationContext {
@@ -893,6 +953,73 @@ mod tests {
);
}
struct StartsTaskFromRewriteHook {
registry: FeatureBackgroundTaskRegistry,
feature: FeatureId,
rejected: Arc<AtomicBool>,
}
#[async_trait]
impl crate::hook::Hook<crate::hook::BeforeSessionRewrite> for StartsTaskFromRewriteHook {
async fn call(
&self,
_input: &crate::hook::BeforeSessionRewriteContext,
) -> Result<crate::hook::BeforeSessionRewriteAction, HookError> {
let error = self
.registry
.start(&self.feature, "hook-start", invocation())
.unwrap_err();
self.rejected.store(
error.category == HookErrorCategory::ScopeDisposed,
Ordering::Release,
);
Ok(crate::hook::BeforeSessionRewriteAction::Continue)
}
}
#[tokio::test]
async fn rewrite_hook_cannot_start_task_while_quiescence_guard_is_held() {
let feature = FeatureId::builtin("hook-start-test");
let declaration = BackgroundTaskDeclaration::worker_managed("hook-start", "hook-start");
let mut task_builder = FeatureBackgroundTaskRegistryBuilder::default();
task_builder
.register(
feature.clone(),
BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(1)),
WaitForCancellation,
)
.unwrap();
let tasks = task_builder.build();
let rejected = Arc::new(AtomicBool::new(false));
let mut hook_builder = crate::hook::HookRegistryBuilder::new();
hook_builder
.add_named_before_session_rewrite(
"hook-start-test",
crate::hook::HookExecutionPolicy::fail_closed(),
StartsTaskFromRewriteHook {
registry: tasks.clone(),
feature,
rejected: Arc::clone(&rejected),
},
)
.unwrap();
let hooks = hook_builder.build();
let guard = tasks.begin_session_rewrite().await.unwrap();
let context = crate::hook::BeforeSessionRewriteContext {
invocation: invocation(),
kind: crate::hook::SessionRewriteKind::Compact,
current_history: crate::hook::HookHistoryRange::default(),
};
assert_eq!(
hooks.before_session_rewrite(&context).await.unwrap(),
crate::hook::BeforeSessionRewriteAction::Continue
);
assert!(rejected.load(Ordering::Acquire));
drop(guard);
tasks.shutdown().await.unwrap();
}
#[tokio::test]
async fn rewrite_joins_old_tasks_before_advancing_session_generation() {
let feature = FeatureId::builtin("generation-test");
@@ -915,9 +1042,16 @@ mod tests {
.unwrap();
assert_eq!(registry.inner.session_generation.load(Ordering::Acquire), 1);
registry.before_session_rewrite().await.unwrap();
let rewrite_guard = registry.begin_session_rewrite().await.unwrap();
assert_eq!(registry.inner.session_generation.load(Ordering::Acquire), 2);
assert!(matches!(
registry.start(&feature, "generation", invocation()),
Err(HookError {
category: HookErrorCategory::ScopeDisposed,
..
})
));
assert_eq!(
stale_fence.ensure_current().unwrap_err().category,
HookErrorCategory::Cancelled
@@ -927,6 +1061,27 @@ mod tests {
registry.diagnostics()[0].outcome,
BackgroundTaskOutcome::Cancelled
);
drop(rewrite_guard);
assert!(matches!(
registry
.start(&feature, "generation", invocation())
.unwrap(),
BackgroundTaskStart::Started { .. }
));
registry.shutdown().await.unwrap();
}
#[tokio::test]
async fn shutdown_during_rewrite_prevents_guard_drop_from_reopening_starts() {
let registry = FeatureBackgroundTaskRegistry::default();
let guard = registry.begin_session_rewrite().await.unwrap();
registry.shutdown().await.unwrap();
drop(guard);
let error = registry
.start(&FeatureId::builtin("missing"), "missing", invocation())
.unwrap_err();
assert_eq!(error.category, HookErrorCategory::ScopeDisposed);
}
#[tokio::test]
@@ -942,7 +1097,7 @@ mod tests {
let registry = builder.build();
registry.start(&feature, "rewrite", invocation()).unwrap();
let error = registry.before_session_rewrite().await.unwrap_err();
let error = registry.begin_session_rewrite().await.unwrap_err();
assert_eq!(error.category, HookErrorCategory::Dependency);
registry.shutdown().await.unwrap();
assert_eq!(registry.diagnostics().len(), 1);
crates/worker/src/worker.rs
+12 -8
View File
@@ -40,7 +40,7 @@ use manifest::{
use crate::compact::state::CompactState;
use crate::compact::usage_tracker::UsageTracker;
use crate::feature::background::FeatureBackgroundTaskRegistry;
use crate::feature::background::{BackgroundTaskRewriteGuard, FeatureBackgroundTaskRegistry};
use crate::feature::builtin::memory::WorkspaceMemoryBackendError;
use crate::feature::builtin::{
MemoryExtractFeature, MemoryExtractState, SessionExploreFeature, SessionExploreState,
@@ -2031,7 +2031,8 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
target: RewindTargetId,
expected_head_entries: usize,
) -> Result<RewindAppliedState, RewindError> {
self.prepare_session_rewrite(SessionRewriteKind::Rewind)
let _rewrite_guard = self
.prepare_session_rewrite(SessionRewriteKind::Rewind)
.await
.map_err(|error| RewindError::Invalid(error.to_string()))?;
let loc = self.segment_state.location();
@@ -3393,7 +3394,8 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
// state up to that turn). The new SegmentStart replaces the mirror
// and is broadcast through the sink so existing subscribers reset
// their view.
self.prepare_session_rewrite(SessionRewriteKind::Fork)
let _rewrite_guard = self
.prepare_session_rewrite(SessionRewriteKind::Fork)
.await?;
let w = self.engine.as_ref().unwrap();
let fork_segment_id = session_store::new_segment_id();
@@ -3693,9 +3695,10 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
async fn prepare_session_rewrite(
&mut self,
kind: SessionRewriteKind,
) -> Result<(), WorkerError> {
self.feature_background_tasks
.before_session_rewrite()
) -> Result<BackgroundTaskRewriteGuard, WorkerError> {
let rewrite_guard = self
.feature_background_tasks
.begin_session_rewrite()
.await
.map_err(|error| WorkerError::FeatureLifecycle(error.to_string()))?;
if let Some(hooks) = self.hook_registry.clone() {
@@ -3723,7 +3726,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
}
}
}
Ok(())
Ok(rewrite_guard)
}
pub async fn manual_compact(&mut self) -> Result<ManualCompactResult, WorkerError> {
@@ -3941,7 +3944,8 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// Runs one parent-owned observable compaction service and returns the new
/// Segment ID. Lifecycle revisions are committed before they are broadcast.
pub async fn compact(&mut self, retained_tokens: u64) -> Result<SegmentId, WorkerError> {
self.prepare_session_rewrite(SessionRewriteKind::Compact)
let _rewrite_guard = self
.prepare_session_rewrite(SessionRewriteKind::Compact)
.await?;
let mut lifecycle = CompactionLifecycle {
schema_version: 2,