fix(targets): make persistent queue store crash-safe and replay lifecycle correct (#4505)

* fix(targets): make persistent queue store crash-safe and replay lifecycle correct Harden the target notification persistent queue (store.rs) and the replay worker lifecycle (runtime) against data loss, silent truncation, ordering drift, orphaned tasks, and a few low-risk robustness gaps. store.rs - Atomic, durable writes: write to a per-key temp file, fsync (sync_all), then rename into place; best-effort parent-dir fsync. A crash mid-write can no longer lose an acknowledged event or leave a half-written payload that reads as a valid entry. - open() now removes leftover .tmp residue and zero-byte files, and only indexes files matching the queue extension, so ghosts/foreign files are never replayed. - FIFO ordering is derived from time-ordered UUIDv7 entry names instead of coarse, clock-dependent file mtimes, so replay order is stable and identical after a restart. - Clamp HashMap/Vec pre-allocation derived from untrusted inputs (entry_limit, batch item_count) to avoid capacity-overflow panics / giant allocations. target/mod.rs - QueuedPayload::decode validates body length against the recorded payload_len, rejecting torn/truncated writes instead of delivering a silently truncated body. - send_from_store purges a NotFound/empty entry (index + file) instead of skipping it, so it cannot occupy a queue slot and be replayed forever. - sanitize_queue_dir_component appends a stable hash suffix when the id was lossy, so distinct target ids can no longer collapse onto the same queue directory; path-safe ids are unchanged (no migration). runtime - Replay backoff, idle waits, and inter-scan pauses are now cancel-aware, so reload/shutdown is not blocked for the full retry delay. - ReplayWorkerManager::stop_all signals cancellation and then awaits each worker's exit (bounded, with abort fallback), preventing orphaned tasks and overlapping drain of the same store. - Fix the always-true replay flush condition so batching is real (size/timeout based, one semaphore permit per batch) rather than one permit per entry; dedup keys already pending in the batch. - clear_and_close aggregates and reports per-target close failures instead of swallowing them; explicit shutdown surfaces them. Relates to rustfs/backlog#966 Relates to rustfs/backlog#967 Relates to rustfs/backlog#975 Relates to rustfs/backlog#970 Relates to rustfs/backlog#983 Co-Authored-By: heihutu <heihutu@gmail.com> * style(targets): apply rustfmt to replay batch dedup guard Fixes the Quick Checks rustfmt failure on the multi-line `.iter().any(...)` closure in the replay batch dedup guard. Co-Authored-By: heihutu <heihutu@gmail.com> --------- Co-authored-by: heihutu <heihutu@gmail.com>
2026-08-09 06:39:25 +00:00 · 2026-07-09 00:27:34 +08:00
parent c6d07ffc59
commit 08e44b95f8
5 changed files with 783 additions and 127 deletions
@@ -142,8 +142,15 @@ where
        replay_workers: &mut ReplayWorkerManager,
        activation: RuntimeActivation<E>,
    ) -> Result<(), TargetError> {
+        // Stop (and join) the old replay workers before installing the new set so
+        // no two workers ever drain the same store concurrently, then close the
+        // old targets. A close failure during reload is logged but does not abort
+        // the reload — the new configuration must still take effect.
        self.stop_replay_workers(replay_workers).await;
-        runtime.clear_and_close().await;
+        let close_errors = runtime.clear_and_close().await;
+        if !close_errors.is_empty() {
+            tracing::warn!(failed_targets = close_errors.len(), "Some targets failed to close during runtime reload");
+        }

        for target in activation.targets {
            runtime.add_arc(target);
@@ -174,8 +181,19 @@ where
        runtime: &mut TargetRuntimeManager<E>,
        replay_workers: &mut ReplayWorkerManager,
    ) -> Result<(), TargetError> {
+        // On explicit shutdown, propagate any close/flush failures instead of
+        // swallowing them: the runtime is still fully torn down, but the caller
+        // learns that a target could not be flushed/closed cleanly.
        self.stop_replay_workers(replay_workers).await;
-        runtime.clear_and_close().await;
+        let close_errors = runtime.clear_and_close().await;
+        if !close_errors.is_empty() {
+            let detail = close_errors
+                .into_iter()
+                .map(|(target_id, err)| format!("{target_id}: {err}"))
+                .collect::<Vec<_>>()
+                .join("; ");
+            return Err(TargetError::Storage(format!("Failed to close {detail}")));
+        }
        Ok(())
    }
 }
@@ -31,14 +31,36 @@ use std::sync::Arc;
 use std::{collections::HashMap, fmt::Debug};
 use std::{future::Future, pin::Pin, time::Duration};
 use tokio::sync::{Semaphore, mpsc};
+use tokio::task::JoinHandle;

 /// Shared target trait object used by the runtime manager.
 pub type SharedTarget<E> = Arc<dyn Target<E> + Send + Sync>;
 type ReplayHook<E> = Arc<dyn Fn(ReplayEvent<E>) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;

-#[derive(Debug, Default)]
+/// Upper bound on how long [`ReplayWorkerManager::stop_all`] waits for a single
+/// replay worker to observe its cancel signal and exit before it is forcibly
+/// aborted. Workers observe cancellation promptly (including during retry
+/// backoff), so this only guards against a wedged task.
+const STOP_JOIN_TIMEOUT: Duration = Duration::from_secs(5);
+
+/// Tracks a running replay worker: its cancel channel and, when the worker was
+/// spawned in-process, the [`JoinHandle`] used to await its exit on shutdown.
+struct ReplayWorkerHandle {
+    cancel_tx: mpsc::Sender<()>,
+    join: Option<JoinHandle<()>>,
+}
+
+#[derive(Default)]
 pub struct ReplayWorkerManager {
-    cancellers: HashMap<String, mpsc::Sender<()>>,
+    cancellers: HashMap<String, ReplayWorkerHandle>,
+}
+
+impl Debug for ReplayWorkerManager {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ReplayWorkerManager")
+            .field("worker_count", &self.cancellers.len())
+            .finish()
+    }
 }

 impl ReplayWorkerManager {
@@ -48,8 +70,27 @@ impl ReplayWorkerManager {
        }
    }

+    /// Registers a cancel channel without a join handle.
+    ///
+    /// Used where the worker's lifetime is managed elsewhere (or in tests). Such
+    /// workers are signalled on `stop_all` but not awaited. Prefer
+    /// [`Self::insert_with_handle`] for in-process workers so shutdown can join
+    /// them and avoid orphaned tasks.
    pub fn insert(&mut self, target_id: String, cancel_tx: mpsc::Sender<()>) {
-        self.cancellers.insert(target_id, cancel_tx);
+        self.cancellers
+            .insert(target_id, ReplayWorkerHandle { cancel_tx, join: None });
+    }
+
+    /// Registers a cancel channel together with the worker's join handle so
+    /// `stop_all` can await the worker's exit (bounded by [`STOP_JOIN_TIMEOUT`]).
+    pub fn insert_with_handle(&mut self, target_id: String, cancel_tx: mpsc::Sender<()>, join: JoinHandle<()>) {
+        self.cancellers.insert(
+            target_id,
+            ReplayWorkerHandle {
+                cancel_tx,
+                join: Some(join),
+            },
+        );
    }

    pub fn len(&self) -> usize {
@@ -67,10 +108,39 @@ impl ReplayWorkerManager {
        }
    }

+    /// Stops every replay worker: it first signals cancellation to all of them,
+    /// then awaits each worker's exit (bounded by [`STOP_JOIN_TIMEOUT`], after
+    /// which the task is aborted). Signalling before joining lets all workers
+    /// wind down concurrently, and joining guarantees no worker keeps draining
+    /// the shared store after this returns — preventing duplicate delivery and
+    /// orphaned tasks across reloads and shutdown.
    pub async fn stop_all(&mut self, log_prefix: &str) {
-        for (target_id, cancel_tx) in self.cancellers.drain() {
+        let mut handles: Vec<(String, ReplayWorkerHandle)> = self.cancellers.drain().collect();
+
+        // Phase 1: signal cancellation to all workers.
+        for (target_id, handle) in &handles {
            tracing::info!(target_id = %target_id, "{log_prefix}");
-            let _ = cancel_tx.send(()).await;
+            let _ = handle.cancel_tx.send(()).await;
+        }
+
+        // Phase 2: await each worker's exit, forcibly aborting any that overrun.
+        for (target_id, handle) in handles.drain(..) {
+            let Some(mut join) = handle.join else {
+                continue;
+            };
+            match tokio::time::timeout(STOP_JOIN_TIMEOUT, &mut join).await {
+                Ok(Ok(())) => {}
+                Ok(Err(err)) => {
+                    tracing::warn!(target_id = %target_id, error = %err, "Replay worker terminated abnormally");
+                }
+                Err(_) => {
+                    join.abort();
+                    tracing::warn!(
+                        target_id = %target_id,
+                        "Timed out awaiting replay worker exit; task aborted"
+                    );
+                }
+            }
        }
    }
 }
@@ -232,12 +302,24 @@ where
        self.remove_and_close(&target_id.to_string()).await
    }

-    pub async fn clear_and_close(&mut self) {
+    /// Closes and removes every target, returning the id/error of each target
+    /// whose `close()` failed. Previously these flush/close errors were logged
+    /// and dropped, so a shutdown that failed to flush a target reported success.
+    /// Callers can now surface them (e.g. fail an explicit shutdown) while still
+    /// tearing down the rest of the runtime.
+    pub async fn clear_and_close(&mut self) -> Vec<(String, TargetError)> {
        let target_ids: Vec<String> = self.targets.keys().cloned().collect();
+        let mut errors = Vec::new();
        for target_id in target_ids {
-            let _ = self.remove_and_close(&target_id).await;
+            if let Some(target) = self.targets.remove(&target_id)
+                && let Err(err) = target.close().await
+            {
+                tracing::error!(target_id = %target_id, error = %err, "Failed to close target during shutdown");
+                errors.push((target_id, err));
+            }
        }
        self.targets.clear();
+        errors
    }

    pub fn target_ids(&self) -> Vec<TargetID> {
@@ -317,11 +399,14 @@ pub async fn init_target_and_optionally_start_replay<E, F, G>(
    target: Box<dyn Target<E> + Send + Sync>,
    on_replay_start: F,
    start_replay: G,
-) -> Option<(SharedTarget<E>, Option<mpsc::Sender<()>>)>
+) -> Option<(SharedTarget<E>, Option<(mpsc::Sender<()>, JoinHandle<()>)>)>
 where
    E: PluginEvent,
    F: FnOnce(&str, bool),
-    G: FnOnce(Box<dyn Store<QueuedPayload, Error = StoreError, Key = Key> + Send>, SharedTarget<E>) -> mpsc::Sender<()>,
+    G: FnOnce(
+        Box<dyn Store<QueuedPayload, Error = StoreError, Key = Key> + Send>,
+        SharedTarget<E>,
+    ) -> (mpsc::Sender<()>, JoinHandle<()>),
 {
    let target_id = target.id().to_string();
    let has_store = target.store().is_some();
@@ -350,6 +435,8 @@ where
    Some((shared, cancel))
 }

+type ActivatedTarget<E> = (SharedTarget<E>, Option<(mpsc::Sender<()>, JoinHandle<()>)>);
+
 pub async fn activate_targets_with_replay<E, F, Fut>(
    targets: Vec<Box<dyn Target<E> + Send + Sync>>,
    mut activate_one: F,
@@ -357,16 +444,16 @@ pub async fn activate_targets_with_replay<E, F, Fut>(
 where
    E: PluginEvent,
    F: FnMut(Box<dyn Target<E> + Send + Sync>) -> Fut,
-    Fut: Future<Output = Option<(SharedTarget<E>, Option<mpsc::Sender<()>>)>>,
+    Fut: Future<Output = Option<ActivatedTarget<E>>>,
 {
    let mut replay_workers = ReplayWorkerManager::new();
    let mut shared_targets = Vec::new();

    for target in targets {
-        if let Some((shared_target, cancel_tx)) = activate_one(target).await {
+        if let Some((shared_target, replay)) = activate_one(target).await {
            let target_id = shared_target.id().to_string();
-            if let Some(cancel_tx) = cancel_tx {
-                replay_workers.insert(target_id, cancel_tx);
+            if let Some((cancel_tx, join)) = replay {
+                replay_workers.insert_with_handle(target_id, cancel_tx, join);
            }
            shared_targets.push(shared_target);
        }
@@ -385,17 +472,35 @@ pub fn start_replay_worker<E>(
    semaphore: Option<Arc<Semaphore>>,
    batch_timeout: Duration,
    idle_sleep: Duration,
-) -> mpsc::Sender<()>
+) -> (mpsc::Sender<()>, JoinHandle<()>)
 where
    E: PluginEvent,
 {
    let (cancel_tx, cancel_rx) = mpsc::channel(1);

-    tokio::spawn(async move {
+    let join = tokio::spawn(async move {
        stream_replay_worker(&mut *store, target, cancel_rx, hook, semaphore, batch_timeout, idle_sleep).await;
    });

-    cancel_tx
+    (cancel_tx, join)
+}
+
+/// Number of readable entries accumulated before a replay batch is flushed under
+/// a single semaphore permit. The previous `!batch_keys.is_empty()` flush
+/// condition was always true, so this effectively defaulted to 1 (a permit per
+/// entry) and made `batch_timeout` dead code.
+const REPLAY_BATCH_SIZE: usize = 16;
+
+/// Sleeps for `dur` unless a cancel signal arrives first. Returns `true` if
+/// cancellation was observed. Used so idle waits, inter-scan pauses, and retry
+/// backoff all react promptly to shutdown instead of blocking for the full
+/// duration.
+async fn sleep_or_cancelled(dur: Duration, cancel_rx: &mut mpsc::Receiver<()>) -> bool {
+    tokio::select! {
+        biased;
+        _ = cancel_rx.recv() => true,
+        _ = tokio::time::sleep(dur) => false,
+    }
 }

 async fn stream_replay_worker<E>(
@@ -409,10 +514,7 @@ async fn stream_replay_worker<E>(
 ) where
    E: PluginEvent,
 {
-    const MAX_RETRIES: usize = 5;
-    const BASE_RETRY_DELAY: Duration = Duration::from_secs(2);
-
-    let mut batch_keys = Vec::with_capacity(1);
+    let mut batch_keys = Vec::with_capacity(REPLAY_BATCH_SIZE);
    let mut last_flush = tokio::time::Instant::now();

    loop {
@@ -423,17 +525,21 @@ async fn stream_replay_worker<E>(
        let keys = store.list();
        if keys.is_empty() {
            if !batch_keys.is_empty() && last_flush.elapsed() >= batch_timeout {
-                process_replay_batch(&mut batch_keys, target.clone(), &hook, semaphore.clone()).await;
+                if process_replay_batch(&mut batch_keys, target.clone(), &hook, semaphore.clone(), &mut cancel_rx).await {
+                    return;
+                }
                last_flush = tokio::time::Instant::now();
            }
-            tokio::time::sleep(idle_sleep).await;
+            if sleep_or_cancelled(idle_sleep, &mut cancel_rx).await {
+                return;
+            }
            continue;
        }

        for key in keys {
            if cancel_rx.try_recv().is_ok() {
                if !batch_keys.is_empty() {
-                    process_replay_batch(&mut batch_keys, target.clone(), &hook, semaphore.clone()).await;
+                    process_replay_batch(&mut batch_keys, target.clone(), &hook, semaphore.clone(), &mut cancel_rx).await;
                }
                return;
            }
@@ -452,101 +558,136 @@ async fn stream_replay_worker<E>(
                }
            }

+            // Skip keys already pending in the current batch: an un-flushed
+            // partial batch carries across scans, and `store.list()` keeps
+            // returning not-yet-delivered keys, so without this guard the same
+            // key would be enqueued repeatedly.
+            if batch_keys
+                .iter()
+                .any(|pending: &Key| pending.to_key_string() == key.to_key_string())
+            {
+                continue;
+            }
+
            batch_keys.push(key);
-            if !batch_keys.is_empty() || last_flush.elapsed() >= batch_timeout {
-                process_replay_batch(&mut batch_keys, target.clone(), &hook, semaphore.clone()).await;
+            // Flush once a full batch has accumulated or the batch has aged past
+            // batch_timeout — real size/time-based batching, not once-per-entry.
+            if batch_keys.len() >= REPLAY_BATCH_SIZE || last_flush.elapsed() >= batch_timeout {
+                if process_replay_batch(&mut batch_keys, target.clone(), &hook, semaphore.clone(), &mut cancel_rx).await {
+                    return;
+                }
                last_flush = tokio::time::Instant::now();
            }
        }

-        tokio::time::sleep(Duration::from_millis(100)).await;
-    }
-
-    async fn process_replay_batch<E>(
-        batch_keys: &mut Vec<Key>,
-        target: SharedTarget<E>,
-        hook: &ReplayHook<E>,
-        semaphore: Option<Arc<Semaphore>>,
-    ) where
-        E: PluginEvent,
-    {
-        if batch_keys.is_empty() {
+        if sleep_or_cancelled(Duration::from_millis(100), &mut cancel_rx).await {
            return;
        }
+    }
+}

-        let _permit = match semaphore {
-            Some(ref semaphore) => match semaphore.clone().acquire_owned().await {
-                Ok(permit) => Some(permit),
-                Err(err) => {
-                    tracing::error!(error = %err, "Failed to acquire replay semaphore permit");
-                    return;
+/// Delivers a batch of queued entries under a single semaphore permit.
+///
+/// Returns `true` if a cancel signal was observed while processing (e.g. during
+/// retry backoff), so the caller can stop promptly instead of continuing to
+/// drain a store that a replacement worker may already own.
+async fn process_replay_batch<E>(
+    batch_keys: &mut Vec<Key>,
+    target: SharedTarget<E>,
+    hook: &ReplayHook<E>,
+    semaphore: Option<Arc<Semaphore>>,
+    cancel_rx: &mut mpsc::Receiver<()>,
+) -> bool
+where
+    E: PluginEvent,
+{
+    const MAX_RETRIES: usize = 5;
+    const BASE_RETRY_DELAY: Duration = Duration::from_secs(2);
+
+    if batch_keys.is_empty() {
+        return false;
+    }
+
+    let _permit = match semaphore {
+        Some(ref semaphore) => match semaphore.clone().acquire_owned().await {
+            Ok(permit) => Some(permit),
+            Err(err) => {
+                tracing::error!(error = %err, "Failed to acquire replay semaphore permit");
+                batch_keys.clear();
+                return false;
+            }
+        },
+        None => None,
+    };
+
+    let mut cancelled = false;
+    'keys: for key in batch_keys.iter() {
+        let mut retry_count = 0usize;
+        let mut success = false;
+
+        while retry_count < MAX_RETRIES && !success {
+            match target.send_from_store(key.clone()).await {
+                Ok(_) => {
+                    hook(ReplayEvent::Delivered {
+                        key: key.clone(),
+                        target: target.clone(),
+                    })
+                    .await;
+                    success = true;
                }
-            },
-            None => None,
-        };
+                Err(err) => match err {
+                    TargetError::NotConnected | TargetError::Timeout(_) => {
+                        retry_count += 1;
+                        hook(ReplayEvent::RetryableError {
+                            error: err,
+                            key: key.clone(),
+                            retry_count,
+                            target: target.clone(),
+                        })
+                        .await;

-        for key in batch_keys.iter() {
-            let mut retry_count = 0usize;
-            let mut success = false;
-
-            while retry_count < MAX_RETRIES && !success {
-                match target.send_from_store(key.clone()).await {
-                    Ok(_) => {
-                        hook(ReplayEvent::Delivered {
+                        let jitter = Duration::from_millis(key.to_string().len() as u64 % 500);
+                        let backoff = 1u32 << retry_count as u32;
+                        // Observe cancellation during backoff so shutdown/reload is
+                        // not blocked for the full (potentially many-second) delay.
+                        if sleep_or_cancelled(BASE_RETRY_DELAY * backoff + jitter, cancel_rx).await {
+                            cancelled = true;
+                            break 'keys;
+                        }
+                    }
+                    TargetError::Dropped(reason) => {
+                        hook(ReplayEvent::Dropped {
+                            key: key.clone(),
+                            reason,
+                            target: target.clone(),
+                        })
+                        .await;
+                        break;
+                    }
+                    other => {
+                        hook(ReplayEvent::PermanentFailure {
+                            error: other,
                            key: key.clone(),
                            target: target.clone(),
                        })
                        .await;
-                        success = true;
+                        break;
                    }
-                    Err(err) => match err {
-                        TargetError::NotConnected | TargetError::Timeout(_) => {
-                            retry_count += 1;
-                            hook(ReplayEvent::RetryableError {
-                                error: err,
-                                key: key.clone(),
-                                retry_count,
-                                target: target.clone(),
-                            })
-                            .await;
-
-                            let jitter = Duration::from_millis(key.to_string().len() as u64 % 500);
-                            let backoff = 1u32 << retry_count as u32;
-                            tokio::time::sleep(BASE_RETRY_DELAY * backoff + jitter).await;
-                        }
-                        TargetError::Dropped(reason) => {
-                            hook(ReplayEvent::Dropped {
-                                key: key.clone(),
-                                reason,
-                                target: target.clone(),
-                            })
-                            .await;
-                            break;
-                        }
-                        other => {
-                            hook(ReplayEvent::PermanentFailure {
-                                error: other,
-                                key: key.clone(),
-                                target: target.clone(),
-                            })
-                            .await;
-                            break;
-                        }
-                    },
-                }
-            }
-
-            if retry_count >= MAX_RETRIES && !success {
-                hook(ReplayEvent::RetryExhausted {
-                    key: key.clone(),
-                    target: target.clone(),
-                })
-                .await;
+                },
            }
        }

-        batch_keys.clear();
+        if retry_count >= MAX_RETRIES && !success {
+            hook(ReplayEvent::RetryExhausted {
+                key: key.clone(),
+                target: target.clone(),
+            })
+            .await;
+        }
    }
+
+    batch_keys.clear();
+    cancelled
 }

 #[cfg(test)]
@@ -641,4 +782,57 @@ mod tests {
        assert_eq!(snapshots[0].target_id, "primary:webhook");
        assert_eq!(snapshots[0].target_type, "webhook");
    }
+
+    #[tokio::test]
+    async fn sleep_or_cancelled_returns_immediately_on_cancel() {
+        let (cancel_tx, mut cancel_rx) = tokio::sync::mpsc::channel::<()>(1);
+        cancel_tx.send(()).await.unwrap();
+
+        // A pending cancel signal must short-circuit a long sleep.
+        let start = std::time::Instant::now();
+        let cancelled = super::sleep_or_cancelled(std::time::Duration::from_secs(30), &mut cancel_rx).await;
+        assert!(cancelled);
+        assert!(
+            start.elapsed() < std::time::Duration::from_secs(5),
+            "cancel should not wait for the full sleep"
+        );
+    }
+
+    #[tokio::test]
+    async fn sleep_or_cancelled_returns_false_when_not_cancelled() {
+        let (_cancel_tx, mut cancel_rx) = tokio::sync::mpsc::channel::<()>(1);
+        let cancelled = super::sleep_or_cancelled(std::time::Duration::from_millis(10), &mut cancel_rx).await;
+        assert!(!cancelled);
+    }
+
+    #[tokio::test]
+    async fn stop_all_joins_and_awaits_worker_exit() {
+        use super::ReplayWorkerManager;
+        use std::sync::atomic::AtomicBool;
+
+        let mut manager = ReplayWorkerManager::new();
+        let exited = Arc::new(AtomicBool::new(false));
+
+        let (cancel_tx, mut cancel_rx) = tokio::sync::mpsc::channel::<()>(1);
+        let exited_task = Arc::clone(&exited);
+        let join = tokio::spawn(async move {
+            // Run until cancelled, then record clean exit.
+            loop {
+                if super::sleep_or_cancelled(std::time::Duration::from_millis(50), &mut cancel_rx).await {
+                    break;
+                }
+            }
+            exited_task.store(true, Ordering::SeqCst);
+        });
+
+        manager.insert_with_handle("primary:webhook".to_string(), cancel_tx, join);
+        assert_eq!(manager.len(), 1);
+
+        // stop_all must signal AND await the worker: once it returns, the worker
+        // has actually exited (no orphaned task).
+        manager.stop_all("stopping test worker").await;
+
+        assert!(manager.is_empty());
+        assert!(exited.load(Ordering::SeqCst), "stop_all must await the worker to completion");
+    }
 }