diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 2d3e5ad09..6ab3ae64a 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1719,14 +1719,13 @@ async fn controller_loop( ) .await; if shutdown { - while worker.has_pending_compaction_cleanup() { - if let Err(error) = worker.retry_pending_compaction_cleanup().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; - } + 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; @@ -1754,11 +1753,11 @@ async fn controller_loop( method = method_rx.recv() => match method { Some(method) => method, None => { - while worker.has_pending_compaction_cleanup() { - if let Err(error) = worker.retry_pending_compaction_cleanup().await { - tracing::warn!(error = %error, "shutdown waiting for compaction cleanup"); - tokio::time::sleep(Duration::from_millis(50)).await; - } + while let Err(error) = + worker.finish_pending_compaction_cleanup_for_shutdown().await + { + tracing::warn!(error = %error, "shutdown cleanup quarantine failed"); + tokio::time::sleep(Duration::from_millis(50)).await; } break 'controller; } @@ -1788,11 +1787,11 @@ async fn controller_loop( method = method_rx.recv() => match method { Some(method) => method, None => { - while worker.has_pending_compaction_cleanup() { - if let Err(error) = worker.retry_pending_compaction_cleanup().await { - tracing::warn!(error = %error, "shutdown waiting for compaction cleanup"); - tokio::time::sleep(Duration::from_millis(50)).await; - } + while let Err(error) = + worker.finish_pending_compaction_cleanup_for_shutdown().await + { + tracing::warn!(error = %error, "shutdown cleanup quarantine failed"); + tokio::time::sleep(Duration::from_millis(50)).await; } break 'controller; } @@ -1813,22 +1812,29 @@ async fn controller_loop( .await; } Err(error) => { - 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; - } - // Shutdown is an explicit cleanup barrier: it is not allowed - // to complete while the service remains unconfirmed. if matches!(&method, Method::Shutdown { .. }) { - continue; + 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; + } } } } @@ -2277,14 +2283,14 @@ async fn controller_loop( }); } if shutdown_after_compaction { - while worker.has_pending_compaction_cleanup() { - if let Err(error) = worker.retry_pending_compaction_cleanup().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; - } + 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; diff --git a/crates/worker/src/runtime/worker_allocation/table.rs b/crates/worker/src/runtime/worker_allocation/table.rs index 3056074f1..b916cd178 100644 --- a/crates/worker/src/runtime/worker_allocation/table.rs +++ b/crates/worker/src/runtime/worker_allocation/table.rs @@ -1,6 +1,6 @@ //! 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}; @@ -101,7 +101,9 @@ pub fn default_allocation_path() -> io::Result { /// 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 +123,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 +164,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 +182,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 +206,37 @@ 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()?; + fs::rename(&temp_path, &self.data_path)?; + 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 +290,27 @@ mod tests { assert_eq!(guard.data().allocations[0].worker_name, "a"); } + #[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(); diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index f328838d0..19d44b99d 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -311,6 +311,7 @@ impl Drop for InternalSpawnReservation { pub struct SpawnedWorkerRegistry { internal_records: std::sync::Mutex>, service_records: std::sync::Mutex>, + quarantined_service_cleanup: std::sync::Mutex>, internal_names: std::sync::Mutex>, internal_shutting_down: AtomicBool, pending_internal_spawns: AtomicUsize, @@ -331,6 +332,7 @@ 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()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), pending_internal_spawns: AtomicUsize::new(0), @@ -346,6 +348,7 @@ 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()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), pending_internal_spawns: AtomicUsize::new(0), @@ -360,6 +363,7 @@ 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()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), pending_internal_spawns: AtomicUsize::new(0), @@ -443,6 +447,7 @@ 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()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), pending_internal_spawns: AtomicUsize::new(0), @@ -556,6 +561,45 @@ 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, + 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 { let removed = { @@ -1535,6 +1579,30 @@ 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(); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index f0b9dce9e..bc8fc3afb 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -1058,7 +1058,6 @@ pub struct SegmentState { entries_written: AtomicUsize, append_lock: Mutex<()>, activation_failed: AtomicBool, - activation_failure_authority: Mutex>, } impl SegmentState { @@ -1071,7 +1070,6 @@ impl SegmentState { entries_written: AtomicUsize::new(entries_written), append_lock: Mutex::new(()), activation_failed: AtomicBool::new(false), - activation_failure_authority: Mutex::new(None), }) } @@ -1108,11 +1106,7 @@ impl SegmentState { Ok(()) } - fn fail_closed_activation(&self, authority: SegmentActivationGuard) { - *self - .activation_failure_authority - .lock() - .expect("Segment activation failure authority mutex poisoned") = Some(authority); + fn fail_closed_activation(&self) { self.activation_failed.store(true, Ordering::Release); } @@ -2224,6 +2218,10 @@ pub struct Worker { /// interceptor before Agen applies them to live typed history. pending_committed_history: Arc>>>, + /// 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>, /// 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 @@ -2285,6 +2283,19 @@ pub struct Worker { log_writer: Option>, } +impl Drop for Worker { + 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 Worker { pub async fn stop_feature_runtime(&mut self, reason: impl Into) { if let Some(hooks) = self.hook_registry.clone() { @@ -2484,6 +2495,7 @@ impl Worker { notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY, pending_attachments: Arc::new(Mutex::new(Vec::::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, @@ -4522,7 +4534,7 @@ impl Worker { if let Some(allocation_activation) = allocation_activation { if let Err(error) = allocation_activation.commit() { let (source, authority) = error.into_parts(); - segment_state.fail_closed_activation(authority); + self.fail_closed_segment_activation(authority); return Err(WorkerError::SegmentActivationIncomplete { source }); } } @@ -5185,6 +5197,14 @@ impl Worker { .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) } @@ -5193,6 +5213,50 @@ impl Worker { self.release_pending_compaction_service().await } + 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 { + 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; + } + } + } + } + + 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> { let cleanup = self .pending_compaction_cleanup @@ -5904,7 +5968,7 @@ impl Worker { // 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. - segment_state.fail_closed_activation(authority); + self.fail_closed_segment_activation(authority); return Err(WorkerError::SegmentActivationIncomplete { source }); } } @@ -6123,6 +6187,7 @@ where notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY, pending_attachments: Arc::new(Mutex::new(Vec::::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, @@ -6210,6 +6275,7 @@ where notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY, pending_attachments: Arc::new(Mutex::new(Vec::::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, @@ -6332,6 +6398,7 @@ where notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY, pending_attachments: Arc::new(Mutex::new(Vec::::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, @@ -6711,6 +6778,7 @@ where notification_coalesce_delay: DEFAULT_NOTIFICATION_COALESCE_DELAY, pending_attachments: Arc::new(Mutex::new(Vec::::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, diff --git a/docs/design/compaction.md b/docs/design/compaction.md index 4a9e26894..76ca2c9bf 100644 --- a/docs/design/compaction.md +++ b/docs/design/compaction.md @@ -42,7 +42,7 @@ activation boundary uses this lock order everywhere a live Worker changes Segmen 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. update the allocation table while still holding its lock; +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. @@ -64,8 +64,10 @@ 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 is not emitted -until the retained service record has been stopped successfully. +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. ## Metrics and comparison procedure