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

worker: cover activation race boundaries

Browse Source
This commit is contained in:
Hare
2026-09-18 18:35:38 +09:00
parent d1d391e511
commit d59879cc4b
5 changed files with 182 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+2 -2
View File
@@ -1756,7 +1756,7 @@ async fn controller_loop<C, St>(
while let Err(error) =
worker.finish_pending_compaction_cleanup_for_shutdown().await
{
tracing::warn!(error = %error, "shutdown cleanup quarantine failed");
tracing::warn!(error = %error, "shutdown cleanup retry failed");
tokio::time::sleep(Duration::from_millis(50)).await;
}
break 'controller;
@@ -1790,7 +1790,7 @@ async fn controller_loop<C, St>(
while let Err(error) =
worker.finish_pending_compaction_cleanup_for_shutdown().await
{
tracing::warn!(error = %error, "shutdown cleanup quarantine failed");
tracing::warn!(error = %error, "shutdown cleanup retry failed");
tokio::time::sleep(Duration::from_millis(50)).await;
}
break 'controller;
crates/worker/src/runtime/worker_allocation/table.rs
+85 -15
View File
@@ -8,16 +8,39 @@ use std::thread;
use std::time::{Duration, Instant};
#[cfg(test)]
static FAIL_NEXT_SAVE_PATHS: std::sync::LazyLock<
std::sync::Mutex<std::collections::HashSet<PathBuf>>,
> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashSet::new()));
#[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_PATHS
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());
.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;
@@ -222,16 +245,6 @@ impl LockFileGuard {
/// 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<()> {
#[cfg(test)]
{
if FAIL_NEXT_SAVE_PATHS
.lock()
.unwrap()
.remove(&self.data_path)
{
return Err(io::Error::other("injected allocation save failure"));
}
}
let json = serde_json::to_vec_pretty(&self.data).map_err(io::Error::other)?;
let mut temp_os = self.data_path.as_os_str().to_owned();
temp_os.push(".tmp");
@@ -244,7 +257,19 @@ impl LockFileGuard {
.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()?;
}
@@ -313,6 +338,51 @@ 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();
crates/worker/src/spawn/registry.rs
-78
View File
@@ -313,7 +313,6 @@ impl Drop for InternalSpawnReservation {
pub struct SpawnedWorkerRegistry {
internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>,
service_records: std::sync::Mutex<Vec<InternalServiceWorkerRecord>>,
quarantined_service_cleanup: std::sync::Mutex<HashSet<String>>,
#[cfg(test)]
fail_service_stops: std::sync::Mutex<HashMap<String, usize>>,
internal_names: std::sync::Mutex<HashSet<String>>,
@@ -336,7 +335,6 @@ impl SpawnedWorkerRegistry {
Arc::new(Self {
internal_records: std::sync::Mutex::new(Vec::new()),
service_records: std::sync::Mutex::new(Vec::new()),
quarantined_service_cleanup: std::sync::Mutex::new(HashSet::new()),
#[cfg(test)]
fail_service_stops: std::sync::Mutex::new(HashMap::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
@@ -354,7 +352,6 @@ impl SpawnedWorkerRegistry {
Arc::new(Self {
internal_records: std::sync::Mutex::new(Vec::new()),
service_records: std::sync::Mutex::new(Vec::new()),
quarantined_service_cleanup: std::sync::Mutex::new(HashSet::new()),
#[cfg(test)]
fail_service_stops: std::sync::Mutex::new(HashMap::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
@@ -371,7 +368,6 @@ impl SpawnedWorkerRegistry {
Arc::new(Self {
internal_records: std::sync::Mutex::new(Vec::new()),
service_records: std::sync::Mutex::new(Vec::new()),
quarantined_service_cleanup: std::sync::Mutex::new(HashSet::new()),
#[cfg(test)]
fail_service_stops: std::sync::Mutex::new(HashMap::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
@@ -457,7 +453,6 @@ impl SpawnedWorkerRegistry {
registry: Arc::new(Self {
internal_records: std::sync::Mutex::new(Vec::new()),
service_records: std::sync::Mutex::new(Vec::new()),
quarantined_service_cleanup: std::sync::Mutex::new(HashSet::new()),
#[cfg(test)]
fail_service_stops: std::sync::Mutex::new(HashMap::new()),
internal_names: std::sync::Mutex::new(HashSet::new()),
@@ -586,16 +581,6 @@ impl SpawnedWorkerRegistry {
.insert(session_id.to_owned(), count);
}
/// 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 is_service_cleanup_quarantined_for_test(&self, session_id: &str) -> bool {
self.quarantined_service_cleanup
.lock()
.unwrap()
.contains(session_id)
}
pub(crate) async fn stop_service(&self, session_id: &str) -> io::Result<bool> {
#[cfg(test)]
{
@@ -625,45 +610,6 @@ impl SpawnedWorkerRegistry {
self.remove_service(session_id)
}
/// Transfer cleanup authority to a registry-owned quarantine task after the
/// controller has exhausted its bounded shutdown attempts. The service record
/// remains registered until a later stop succeeds.
pub(crate) fn quarantine_service_cleanup(
self: &Arc<Self>,
session_id: String,
) -> io::Result<()> {
let mut quarantined = self
.quarantined_service_cleanup
.lock()
.map_err(|_| io::Error::other("service cleanup quarantine lock poisoned"))?;
if !quarantined.insert(session_id.clone()) {
return Ok(());
}
drop(quarantined);
let registry = Arc::clone(self);
tokio::spawn(async move {
loop {
match registry.stop_service(&session_id).await {
Ok(_) => {
if let Ok(mut quarantined) = registry.quarantined_service_cleanup.lock() {
quarantined.remove(&session_id);
}
return;
}
Err(error) => {
warn!(
error = %error,
"quarantined service cleanup retry failed"
);
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
}
}
}
});
Ok(())
}
/// Remove one service Worker and emit the terminal projection fence.
pub(crate) fn remove_service(&self, session_id: &str) -> io::Result<bool> {
let removed = {
@@ -1643,30 +1589,6 @@ mod tests {
));
}
#[tokio::test]
async fn cleanup_quarantine_owns_record_until_background_stop_finishes() {
let registry = registry();
registry
.quarantine_service_cleanup("missing-service".into())
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(1), async {
loop {
if !registry
.quarantined_service_cleanup
.lock()
.unwrap()
.contains("missing-service")
{
break;
}
tokio::task::yield_now().await;
}
})
.await
.unwrap();
}
#[tokio::test]
async fn read_only_summary_omits_unavailable_change_stat() {
let tracker = tools::Tracker::new();
crates/worker/src/worker.rs
+92 -55
View File
@@ -5216,45 +5216,18 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
pub(crate) async fn finish_pending_compaction_cleanup_for_shutdown(
&self,
) -> Result<(), WorkerError> {
const SHUTDOWN_CLEANUP_ATTEMPTS: usize = 3;
let mut last_error = None;
for attempt in 0..SHUTDOWN_CLEANUP_ATTEMPTS {
loop {
match self.retry_pending_compaction_cleanup().await {
Ok(()) => return Ok(()),
Err(error) => {
last_error = Some(error);
if attempt + 1 < SHUTDOWN_CLEANUP_ATTEMPTS {
tokio::time::sleep(Duration::from_millis(50)).await;
}
warn!(
error = %error,
"shutdown remains blocked on compaction service cleanup"
);
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
}
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 Err(last_error.unwrap_or_else(|| WorkerError::CompactionCleanupPending {
source: ScopeLockError::Io(std::io::Error::other(
"compaction registry unavailable during shutdown quarantine",
)),
}));
};
registry
.quarantine_service_cleanup(cleanup.session_id.clone())
.map_err(|source| WorkerError::CompactionCleanupPending {
source: ScopeLockError::Io(source),
})?;
*self
.pending_compaction_cleanup
.lock()
.expect("pending compaction cleanup mutex poisoned") = None;
Ok(())
}
async fn release_pending_compaction_service(&self) -> Result<(), WorkerError> {
@@ -8682,7 +8655,7 @@ mod build_summary_prompt_tests {
}
#[tokio::test]
async fn shutdown_cleanup_failures_transfer_service_authority_to_quarantine() {
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(
@@ -8708,28 +8681,20 @@ mod build_summary_prompt_tests {
compaction_id: "compaction".into(),
});
tokio::time::timeout(
Duration::from_secs(1),
worker.finish_pending_compaction_cleanup_for_shutdown(),
)
.await
.expect("shutdown cleanup must be bounded")
.unwrap();
assert!(!worker.has_pending_compaction_cleanup());
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));
assert!(registry.is_service_cleanup_quarantined_for_test(&cleanup_session_id));
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if !registry.is_service_cleanup_quarantined_for_test(&cleanup_session_id)
&& !registry.has_service_for_test(&cleanup_session_id)
{
break;
}
tokio::task::yield_now().await;
}
})
.await
.unwrap();
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]
@@ -10479,6 +10444,78 @@ 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();
docs/design/compaction.md
+3 -4
View File
@@ -64,10 +64,9 @@ 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 performs a
bounded set of direct stop attempts and does not emit its terminal event until either
the registry record is gone or cleanup authority has moved to a runtime-owned
quarantine task.
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