Recreate workers/optimizers async to not block consensus (#9121)

* Recreate optimizers in non-blocking fashion from consensus calls

* Update comments

* On optimizer config update failure, report error status to local shard

* Add a test to confirm we don't block consensus

* Rerun recreation if called multiple times

* Use atomics instead

* Move to the bottom

* Reformat

lib/collection/src/collection/collection_ops.rs

@@ -16,6 +16,7 @@ use crate::operations::types::*;
 use crate::shards::replica_set::Change;
 use crate::shards::replica_set::replica_set_state::ReplicaState;
 use crate::shards::shard::PeerId;
+use crate::shards::shard_holder::SharedShardHolder;
 
 /// Old logic for aborting shard transfers on shard drop, had a bug: it dropped all transfers
 /// regardless of the shard id. In order to keep consensus consistent, we can only
@@ -30,7 +31,7 @@ impl Collection {
     /// Updates collection params:
     /// Saves new params on disk
     ///
-    /// After this, `recreate_optimizers_blocking` must be called to create new optimizers using
+    /// After this, `recreate_optimizers_background` must be called to create new optimizers using
     /// the updated configuration.
     pub async fn update_params_from_diff(
         &self,
@@ -47,7 +48,7 @@ impl Collection {
     /// Updates HNSW config:
     /// Saves new params on disk
     ///
-    /// After this, `recreate_optimizers_blocking` must be called to create new optimizers using
+    /// After this, `recreate_optimizers_background` must be called to create new optimizers using
     /// the updated configuration.
     pub async fn update_hnsw_config_from_diff(
         &self,
@@ -64,7 +65,7 @@ impl Collection {
     /// Updates vectors config:
     /// Saves new params on disk
     ///
-    /// After this, `recreate_optimizers_blocking` must be called to create new optimizers using
+    /// After this, `recreate_optimizers_background` must be called to create new optimizers using
     /// the updated configuration.
     pub async fn update_vectors_from_diff(
         &self,
@@ -82,7 +83,7 @@ impl Collection {
     /// Updates sparse vectors config:
     /// Saves new params on disk
     ///
-    /// After this, `recreate_optimizers_blocking` must be called to create new optimizers using
+    /// After this, `recreate_optimizers_background` must be called to create new optimizers using
     /// the updated configuration.
     pub async fn update_sparse_vectors_from_other(
         &self,
@@ -100,7 +101,7 @@ impl Collection {
     /// Updates shard optimization params:
     /// Saves new params on disk
     ///
-    /// After this, `recreate_optimizers_blocking` must be called to create new optimizers using
+    /// After this, `recreate_optimizers_background` must be called to create new optimizers using
     /// the updated configuration.
     pub async fn update_optimizer_params_from_diff(
         &self,
@@ -117,7 +118,7 @@ impl Collection {
     /// Updates quantization config:
     /// Saves new params on disk
     ///
-    /// After this, `recreate_optimizers_blocking` must be called to create new optimizers using
+    /// After this, `recreate_optimizers_background` must be called to create new optimizers using
     /// the updated configuration.
     pub async fn update_quantization_config_from_diff(
         &self,
@@ -293,28 +294,70 @@ impl Collection {
         Ok(())
     }
 
-    /// Recreate the optimizers on all shards for this collection
+    /// Recreate the optimizers on all shards for this collection, in the background.
     ///
-    /// This will stop existing optimizers, and start new ones with new configurations.
+    /// Returns immediately and performs all the work - stopping the existing workers and starting
+    /// new ones - in a detached task. Stopping the existing workers waits for in-flight
+    /// optimizations to finish, which can take a long time. This is why it runs in the background:
+    /// it is reached from paths that go through consensus, where blocking the caller stalls the
+    /// whole consensus loop and can take down a cluster.
     ///
-    /// # Blocking
+    /// At most one recreation runs at a time. If one is already running, this records that another
+    /// run is needed and returns; the running task then runs once more when it finishes, picking up
+    /// the latest config. Any number of requests that arrive while a task is running collapse into a
+    /// single additional run (recreation always rebuilds from the current config, so coalescing is
+    /// safe - the last run reflects the latest state).
     ///
-    /// Partially blocking. Stopping existing optimizers is blocking. Starting new optimizers is
-    /// not blocking.
-    ///
-    /// ## Cancel safety
-    ///
-    /// This function is cancel safe, and will always run to completion.
-    pub async fn recreate_optimizers_blocking(&self) -> CollectionResult<()> {
+    /// Errors are logged rather than returned: the configuration change that triggers the
+    /// recreation has already been applied and persisted by the time we get here, so there is no
+    /// caller left to propagate them to. Failures are also surfaced as optimizer errors per shard
+    /// (see `LocalShard::on_optimizer_config_update`).
+    pub fn recreate_optimizers_background(&self) {
+        // Single-flight: only spawn a task if none is running. Otherwise the request is coalesced
+        // into a queued re-run handled by the task that is already running.
+        if !self.recreate_optimizers_state.request() {
+            return;
+        }
+
         let shards_holder = self.shards_holder.clone();
+        let collection_id = self.id.clone();
+        let recreate_state = self.recreate_optimizers_state.clone();
         tokio::task::spawn(async move {
-            let shard_holder = shards_holder.read().await;
-            let updates = shard_holder
-                .all_shards()
-                .map(|replica_set| replica_set.on_optimizer_config_update());
-            future::try_join_all(updates).await
-        })
-        .await??;
+            loop {
+                // Run the recreation as a child task, so a panic is contained (surfaced as a
+                // `JoinError`) and never unwinds the coordinator loop below - which would otherwise
+                // leave `running` stuck and wedge all future recreations.
+                match tokio::task::spawn(Self::recreate_optimizers(shards_holder.clone())).await {
+                    Ok(Ok(())) => {}
+                    Ok(Err(err)) => log::error!(
+                        "Failed to recreate optimizers for collection {collection_id} in background: {err}",
+                    ),
+                    Err(err) => log::error!(
+                        "Optimizer recreation task for collection {collection_id} failed: {err}",
+                    ),
+                }
+
+                // Run again if a request arrived while we were running, otherwise stop. The
+                // decision is a single atomic compare-and-swap (in `finish_run`), so a request
+                // arriving exactly now is never lost: it either still sees us running (and queues a
+                // re-run, handled by the next iteration) or sees us idle (and spawns a fresh task).
+                if !recreate_state.finish_run() {
+                    break;
+                }
+            }
+        });
+    }
+
+    /// Stop the existing optimizers on all shards and start new ones using the current config.
+    ///
+    /// Implementation behind [`Collection::recreate_optimizers_background`]. Takes an owned
+    /// [`SharedShardHolder`] so the returned future is `'static` and can be spawned as a task.
+    async fn recreate_optimizers(shards_holder: SharedShardHolder) -> CollectionResult<()> {
+        let shard_holder = shards_holder.read().await;
+        let updates = shard_holder
+            .all_shards()
+            .map(|replica_set| replica_set.on_optimizer_config_update());
+        future::try_join_all(updates).await?;
         Ok(())
     }
 

lib/collection/src/collection/mod.rs

@@ -19,6 +19,7 @@ use std::collections::HashMap;
 use std::ops::Deref;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU8, Ordering};
 use std::time::Duration;
 
 use clean::ShardCleanTasks;
@@ -91,6 +92,9 @@ pub struct Collection {
     collection_stats_cache: CollectionSizeStatsCache,
     // Background tasks to clean shards
     shard_clean_tasks: ShardCleanTasks,
+    // Coordinates background optimizer recreation: at most one runs at a time, and requests that
+    // arrive while one is running are coalesced into a single re-run.
+    recreate_optimizers_state: Arc<RecreateOptimizersState>,
 }
 
 pub type RequestShardTransfer = Arc<dyn Fn(ShardTransfer) + Send + Sync>;
@@ -200,6 +204,7 @@ impl Collection {
             optimizer_resource_budget,
             collection_stats_cache,
             shard_clean_tasks: Default::default(),
+            recreate_optimizers_state: Default::default(),
         })
     }
 
@@ -319,6 +324,7 @@ impl Collection {
             optimizer_resource_budget,
             collection_stats_cache,
             shard_clean_tasks: Default::default(),
+            recreate_optimizers_state: Default::default(),
         }
     }
 
@@ -967,3 +973,104 @@ impl StorageVersion for CollectionVersion {
         env!("CARGO_PKG_VERSION")
     }
 }
+
+/// Coordinates background optimizer recreation so at most one recreation runs at a time.
+///
+/// This tracks two things - "a task is running" and "another run is queued" - packed into a single
+/// atomic. They cannot be two independent atomics: deciding to stop (in [`Self::finish_run`]) and
+/// deciding to coalesce (in [`Self::request`]) each read *and* write a combination of the two, so
+/// the transition must be one atomic compare-and-swap. With separate atomics a request could be
+/// lost - a task stops while a concurrent request believes it coalesced into it.
+///
+/// See [`Collection::recreate_optimizers_background`].
+#[derive(Default)]
+struct RecreateOptimizersState {
+    state: AtomicU8,
+}
+
+/// No recreation task is running.
+const RECREATE_IDLE: u8 = 0;
+/// A recreation task is running; no further run is queued.
+const RECREATE_RUNNING: u8 = 1;
+/// A recreation task is running and another run is queued for when it finishes.
+const RECREATE_RUNNING_PENDING: u8 = 2;
+
+impl RecreateOptimizersState {
+    /// Register a recreation request.
+    ///
+    /// Returns `true` if the caller must spawn a new recreation task, or `false` if a task is
+    /// already running - in which case the request is coalesced into a single queued re-run.
+    fn request(&self) -> bool {
+        let prev =
+            self.state
+                .fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| match state {
+                    RECREATE_IDLE => Some(RECREATE_RUNNING),
+                    RECREATE_RUNNING => Some(RECREATE_RUNNING_PENDING),
+                    // Already queued, nothing to change.
+                    _ => None,
+                });
+        // We must spawn exactly when we moved the state out of idle.
+        prev == Ok(RECREATE_IDLE)
+    }
+
+    /// Record that the running task finished one recreation.
+    ///
+    /// Returns `true` if it must run again (a request arrived while it was running), or `false` if
+    /// it should stop. When it stops, the state returns to idle so the next request spawns anew.
+    fn finish_run(&self) -> bool {
+        let prev =
+            self.state
+                .fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| match state {
+                    RECREATE_RUNNING_PENDING => Some(RECREATE_RUNNING),
+                    RECREATE_RUNNING => Some(RECREATE_IDLE),
+                    // Idle would mean finish_run was called without a running task.
+                    _ => None,
+                });
+        // We run again exactly when we consumed a queued re-run.
+        prev == Ok(RECREATE_RUNNING_PENDING)
+    }
+}
+
+#[cfg(test)]
+mod recreate_optimizers_state_tests {
+    use super::RecreateOptimizersState;
+
+    #[test]
+    fn first_request_spawns() {
+        let state = RecreateOptimizersState::default();
+        assert!(state.request(), "first request must spawn a task");
+    }
+
+    #[test]
+    fn requests_while_running_are_coalesced() {
+        let state = RecreateOptimizersState::default();
+        assert!(state.request());
+
+        // Several requests arrive while the task runs: none spawns, all collapse into one re-run.
+        assert!(!state.request());
+        assert!(!state.request());
+    }
+
+    #[test]
+    fn finish_without_pending_stops_then_next_request_spawns() {
+        let state = RecreateOptimizersState::default();
+        state.request();
+
+        assert!(!state.finish_run(), "no queued request, so it must stop");
+        // Back to idle: a fresh request spawns a new task again.
+        assert!(
+            state.request(),
+            "a request after stopping must spawn a new task"
+        );
+    }
+
+    #[test]
+    fn finish_with_pending_runs_once_more_then_stops() {
+        let state = RecreateOptimizersState::default();
+        state.request();
+        state.request(); // coalesced -> queued re-run
+
+        assert!(state.finish_run(), "queued request, so it must run again");
+        assert!(!state.finish_run(), "no further requests, so it stops");
+    }
+}

lib/collection/src/collection/state_management.rs

@@ -193,8 +193,11 @@ impl Collection {
 
         self.print_warnings().await;
 
+        // Recreate optimizers in the background: this path is reached from consensus (Raft snapshot
+        // application), and stopping the existing optimizers can take a long time, which would
+        // otherwise stall the consensus loop.
         if recreate_optimizers {
-            self.recreate_optimizers_blocking().await?;
+            self.recreate_optimizers_background();
         }
 
         Ok(())

lib/collection/src/collection/vector_name_schema.rs

@@ -40,7 +40,11 @@ impl Collection {
         // Refresh shard optimizers so the cached `SegmentOptimizerConfig` picks up the
         // new vector schema. Otherwise optimization-built destination segments use the
         // stale dim and panic on dim mismatch when copying from sources with the new dim.
-        self.recreate_optimizers_blocking().await?;
+        //
+        // Done in the background: this path is reached from consensus, where blocking can stall the
+        // whole cluster. The refresh itself is unchanged - it already ran on a spawned task either
+        // way - so this does not widen the window in which the stale config is in effect.
+        self.recreate_optimizers_background();
 
         Ok(())
     }
@@ -67,7 +71,10 @@ impl Collection {
         // Refresh shard optimizers so the cached `SegmentOptimizerConfig` drops the
         // removed vector. Without this, optimization keeps allocating storage for it
         // using the old config.
-        self.recreate_optimizers_blocking().await?;
+        //
+        // Done in the background: this path is reached from consensus, where blocking can stall the
+        // whole cluster.
+        self.recreate_optimizers_background();
 
         Ok(())
     }

lib/collection/src/shards/local_shard/updaters.rs

@@ -33,18 +33,39 @@ impl LocalShard {
     /// Handles updates to the optimizer configuration by rebuilding optimizers
     /// and restarting the update handler's workers with the new configuration.
     ///
+    /// On failure, the error is recorded as an optimizer error on this shard, so it is exposed in
+    /// the collection's optimizer status. This matters especially for background recreation, where
+    /// the error would otherwise only be logged with no caller left to propagate it to.
+    ///
     /// ## Cancel safety
     ///
     /// This function is **not** cancel safe.
     pub async fn on_optimizer_config_update(&self) -> CollectionResult<()> {
-        let config = self.collection_config.read().await;
+        let result = self.on_optimizer_config_update_impl().await;
+
+        // Surface a failed recreation as an optimizer error, so it shows up in the optimizer status
+        // instead of being silently lost on the background recreation path.
+        if let Err(err) = &result {
+            self.segments
+                .write()
+                .report_optimizer_error(format!("Failed to recreate optimizers: {err}"));
+        }
+
+        result
+    }
+
+    async fn on_optimizer_config_update_impl(&self) -> CollectionResult<()> {
         let mut update_handler = self.update_handler.lock().await;
 
         // Signal all workers to stop
         update_handler.stop_flush_worker();
         update_handler.stop_update_worker();
 
-        // Wait for workers to finish and get pending operations from the old channel
+        // Wait for workers to finish and get pending operations from the old channel.
+        //
+        // This can take a long time, because in-flight optimizations are awaited before they stop.
+        // We deliberately do not hold the `collection_config` read lock across this wait, so a
+        // concurrent collection config update is not blocked on acquiring the config write lock.
         let update_receiver = update_handler.wait_workers_stops().await?;
 
         let update_receiver = match update_receiver {
@@ -65,6 +86,11 @@ impl LocalShard {
             }
         };
 
+        // Read the latest config now that the workers have stopped, and build new optimizers from
+        // it. Reading it here (rather than before the wait) also ensures we pick up the most recent
+        // config if it changed again while we were waiting.
+        let config = self.collection_config.read().await;
+
         let new_optimizers = build_optimizers(
             &self.path,
             &config.params,
@@ -81,6 +107,9 @@ impl LocalShard {
             .optimizer_config
             .prevent_unoptimized
             .unwrap_or_default();
+
+        drop(config);
+
         update_handler.run_workers(update_receiver);
 
         self.optimizers.store(new_optimizers);

lib/storage/src/content_manager/toc/collection_meta_ops.rs

@@ -202,8 +202,12 @@ impl TableOfContent {
         collection.print_warnings().await;
 
         // Recreate optimizers
+        //
+        // This runs in the background and does not block: this path is reached from consensus, and
+        // stopping the existing optimizers can take a long time (in-flight optimizations are
+        // awaited), which would otherwise stall the consensus loop and can take down a cluster.
         if recreate_optimizers {
-            collection.recreate_optimizers_blocking().await?;
+            collection.recreate_optimizers_background();
         }
         Ok(true)
     }

tests/consensus_tests/test_collection_update_not_blocked_by_busy_worker.py

@@ -0,0 +1,130 @@
+"""
+Test that a collection update applied through consensus is not blocked by a busy
+update worker.
+
+A collection config update (PATCH /collections/{name}) triggers recreation of the
+shard optimizers. Recreating optimizers stops the update worker, which first has to
+finish whatever operation it is currently running. If that recreation runs
+*synchronously* on the consensus apply thread, a long-running update operation stalls
+the whole consensus loop and the collection update hangs - which can take down a
+cluster. The recreation must instead happen in the background, so the update returns
+right away.
+
+Repro:
+  1. Keep the shard's update worker busy for a long time with a staging `delay`
+     operation.
+  2. Send a collection config update through consensus.
+  3. The update must return promptly, not wait for the delay to finish.
+
+Requires the `staging` feature (for the `delay` debug operation); skipped otherwise.
+
+Usage:
+    pytest tests/consensus_tests/test_collection_update_not_blocked_by_busy_worker.py -s -v
+"""
+
+import os
+import pathlib
+import time
+
+import requests
+
+os.environ.setdefault(
+    "PYTEST_CURRENT_TEST", "test_collection_update_not_blocked_by_busy_worker"
+)
+
+from .utils import (
+    assert_http_ok,
+    assert_project_root,
+    check_collection_green,
+    skip_if_no_feature,
+    start_cluster,
+    wait_for,
+)
+
+COLLECTION = "update_not_blocked"
+DIM = 4
+
+# How long the update worker is kept busy by the staging delay operation.
+DELAY_SEC = 20.0
+# Time to let the worker actually pick up and start processing the delay before we send
+# the collection update. The worker only blocks on an operation it has already started;
+# a queued-but-not-yet-started delay would simply be dropped when the worker stops, so
+# we must make sure it is genuinely in-flight to exercise the blocking path.
+WORKER_BUSY_WAIT_SEC = 4.0
+# Generous client-side timeout: with the fix the PATCH returns near-instantly, without
+# it the request blocks for ~(DELAY_SEC - WORKER_BUSY_WAIT_SEC) (or errors out earlier).
+PATCH_CLIENT_TIMEOUT_SEC = 30.0
+# The collection update must return well before the delay would finish. With the fix it
+# is sub-second; without it, it is ~16s. 10s cleanly separates the two.
+MAX_PATCH_SEC = 10.0
+
+
+def create_collection(peer_uri: str):
+    r = requests.put(
+        f"{peer_uri}/collections/{COLLECTION}",
+        json={
+            "vectors": {"size": DIM, "distance": "Cosine"},
+            "shard_number": 1,
+            "replication_factor": 1,
+        },
+    )
+    assert_http_ok(r)
+
+
+def submit_delay(peer_uri: str, duration_sec: float):
+    """Submit a staging delay operation to keep the shard's update worker busy.
+
+    Sent with wait=false so the call returns immediately while the worker processes the
+    delay in the background.
+    """
+    r = requests.post(
+        f"{peer_uri}/collections/{COLLECTION}/debug?wait=false",
+        json={"delay": {"duration_sec": duration_sec}},
+    )
+    assert_http_ok(r)
+
+
+def test_collection_update_not_blocked_by_busy_worker(tmp_path: pathlib.Path):
+    assert_project_root()
+
+    peer_uris, _, _ = start_cluster(tmp_path, 1)
+    peer_uri = peer_uris[0]
+
+    # The `delay` debug operation is only available with the staging feature.
+    skip_if_no_feature(peer_uri, "staging")
+
+    create_collection(peer_uri)
+    wait_for(check_collection_green, peer_uri, COLLECTION, wait_for_timeout=30)
+
+    # Keep the update worker busy for a long time, then let it actually start the delay.
+    submit_delay(peer_uri, DELAY_SEC)
+    time.sleep(WORKER_BUSY_WAIT_SEC)
+
+    # Send a collection config update through consensus. Any optimizers_config change
+    # triggers optimizer recreation, which must not block the consensus apply thread on
+    # the busy worker.
+    start = time.time()
+    try:
+        r = requests.patch(
+            f"{peer_uri}/collections/{COLLECTION}",
+            json={"optimizers_config": {"default_segment_number": 2}},
+            timeout=PATCH_CLIENT_TIMEOUT_SEC,
+        )
+    except requests.exceptions.Timeout:
+        elapsed = time.time() - start
+        raise AssertionError(
+            f"Collection update did not return within {PATCH_CLIENT_TIMEOUT_SEC}s "
+            f"(waited {elapsed:.1f}s) - consensus apply was blocked by the busy update worker"
+        )
+    elapsed = time.time() - start
+
+    assert_http_ok(r)
+    assert elapsed < MAX_PATCH_SEC, (
+        f"Collection update took {elapsed:.1f}s (>= {MAX_PATCH_SEC}s) while the update worker "
+        f"was busy for {DELAY_SEC}s - consensus apply was blocked instead of recreating "
+        f"optimizers in the background"
+    )
+    print(
+        f"Collection update returned in {elapsed:.2f}s while the update worker was busy "
+        f"for {DELAY_SEC}s"
+    )