Internal `grpc-timeout` can cancel receiver shard initialization after `local.take()`, leaving the target replica without a local shard slot

Root Cause

During receiver-side shard-transfer setup, Qdrant can run ShardReplicaSet::init_empty_local_shard() under an internal gRPC request future. That method removes the current local slot with local.take() and then awaits several fallible async operations before installing either a new Shard::Local or a fallback Shard::Dummy.

This creates the following state-machine gap:

self.local = Some(Shard::Dummy(_)) or Some(other local receiver shard)
    -> local.take()
self.local = None
    -> await stop_gracefully / LocalShard::clear / LocalShard::build
    -> request future is cancelled or dropped
    -> no Local or Dummy replacement is installed

Analyzed revision: 44ad62f8cd69642be5afa6441612525e24a0d063.

The receiver initialization path is reached from the internal CollectionsInternal/Initiate handler:

// src/tonic/api/collections_internal_api.rs:64-80
async fn initiate(
    &self,
    request: Request<InitiateShardTransferRequest>,
) -> Result<Response<CollectionOperationResponse>, Status> {
    // TODO: Ensure cancel safety!

    validate_and_log(request.get_ref());
    let timing = Instant::now();
    let InitiateShardTransferRequest {
        collection_name,
        shard_id,
    } = request.into_inner();

    // TODO: Ensure cancel safety!
    self.toc
        .initiate_receiving_shard(collection_name, shard_id)
        .await?;

Collection::initiate_shard_transfer() then handles a dummy receiver by initializing an empty local shard:

// lib/collection/src/collection/shard_transfer.rs:574-586
if replica_set.is_dummy().await {
    // We can reach here because of either of these:
    // 1. Qdrant is in recovery mode, and user intentionally triggered a transfer
    // 2. Shard is dirty (shard initializing flag), and Qdrant triggered a transfer to recover from Dead state after an update fails
    //
    // In both cases, it's safe to drop existing local shard data
    log::debug!(
        "Initiating transfer to dummy shard {}. Initializing empty local shard first",
        replica_set.shard_id,
    );
    replica_set.init_empty_local_shard().await?;

    let shard_flag = shard_initializing_flag_path(&collection_path, shard_id);

The vulnerable transition is inside init_empty_local_shard():

// lib/collection/src/shards/replica_set/mod.rs:585-623
/// Clears the local shard data and loads an empty local shard
pub async fn init_empty_local_shard(&self) -> CollectionResult<()> {
    let mut local = self.local.write().await;

    let current_shard = local.take();
    if let Some(current_shard) = current_shard {
        current_shard.stop_gracefully().await;
    }
    LocalShard::clear(&self.shard_path).await?;

    let local_shard_res = LocalShard::build(
        self.shard_id,
        self.collection_id.clone(),
        &self.shard_path,
        self.collection_config.clone(),
        self.shared_storage_config.clone(),
        self.payload_index_schema.clone(),
        self.update_runtime.clone(),
        self.search_runtime.clone(),
        self.optimizer_resource_budget.clone(),
        self.optimizers_config.clone(),
    )
    .await;

    match local_shard_res {
        Ok(local_shard) => {
            *local = Some(Shard::Local(local_shard));
            Ok(())
        }
        Err(err) => {
            let error = format!(
                "Failed to initialize local shard at {:?}: {err}",
                self.shard_path
            );
            log::error!("{error}");
            *local = Some(Shard::Dummy(DummyShard::new(error)));
            Err(err)
        }
    }
}

The method holds a tokio::RwLockWriteGuard across the awaited operations. Dropping the future releases the guard, but it does not repair the protected value. If the future is dropped after local.take() and before either assignment back to Some(...), the protected slot remains None.

This future can be cancelled by the internal gRPC timeout path. Qdrant's internal channel pool installs grpc-timeout metadata for internal requests when the request has none:

// lib/api/src/grpc/transport_channel_pool.rs:98-102
impl Interceptor for PoolInterceptor {
    fn call(&mut self, mut request: Request<()>) -> Result<Request<()>, Status> {
        if request.metadata().get("grpc-timeout").is_none() {
            request.set_timeout(self.default_timeout);
        }

On the server side, tonic's GrpcTimeout wrapper returns a timeout error when the timeout sleep wins before the inner service future completes. Once the wrapper future returns, the unfinished inner service future is dropped. In this path, the dropped inner future contains the generated unary service future and ultimately CollectionsInternalService::initiate -> initiate_receiving_shard -> initiate_shard_transfer -> init_empty_local_shard. The relevant tonic implementation is in tonic/transport/service/grpc_timeout.rs: https://docs.rs/tonic/latest/src/tonic/transport/service/grpc_timeout.rs.html.

This is therefore not just a slow operation. It is a cancellation boundary around a side-effecting receiver-initialization future.

Current Behavior

If timeout or request cancellation lands after local.take() in init_empty_local_shard(), the target replica set can be left with no local shard object:

replica_set.has_local_shard().await == false

At the same time, replica metadata can still say that this peer is local:

// lib/collection/src/shards/replica_set/mod.rs:493-498
/// Wait for a local shard to be initialized.
///
/// Uses a blocking thread internally.
pub async fn wait_for_local(&self, timeout: Duration) -> CollectionResult<()> {
    self.wait_for(|replica_set_state| replica_set_state.is_local, timeout)
        .await
}

The mismatch matters because the retry path can pass wait_for_local(...) using metadata, then skip both receiver repair branches because None is neither a proxy nor a dummy:

// lib/collection/src/shards/replica_set/mod.rs:389-418
pub async fn has_local_shard(&self) -> bool {
    self.local.read().await.is_some()
}

/// Checks if the shard exists locally and not a proxy.
pub async fn is_local(&self) -> bool {
    let local_read = self.local.read().await;
    matches!(*local_read, Some(Shard::Local(_) | Shard::Dummy(_)))
}

pub async fn is_proxy(&self) -> bool {
    let local_read = self.local.read().await;
    match *local_read {
        None => false,
        Some(Shard::Local(_)) => false,
        Some(Shard::Proxy(_)) => true,
        Some(Shard::ForwardProxy(_)) => true,
        Some(Shard::QueueProxy(_)) => true,
        Some(Shard::Dummy(_)) => false,
    }
}

pub async fn is_dummy(&self) -> bool {
    let local_read = self.local.read().await;
    matches!(*local_read, Some(Shard::Dummy(_)))
}

So the observable broken state is:

• self.local == None;
• wait_for_local(...) can still succeed because it checks replica metadata, not the local slot;
• is_dummy().await == false, so a later receiver retry can skip init_empty_local_shard();
• is_proxy().await == false, so proxy cleanup is not attempted either;
• the receiver has neither the controlled error behavior of DummyShard nor a new empty LocalShard ready to receive transferred data.

This is primarily an availability and shard-transfer recovery issue. It does not require claiming permanent on-disk data loss. The concrete problem is that an in-memory receiver replica can enter an unmodeled None state even though transfer/recovery logic expects a valid local receiver slot.

Steps to Reproduce

A minimal source-level scenario is:

1. Have a target ShardReplicaSet whose local slot is Some(Shard::Dummy(_)) and whose replica metadata marks the current peer as local/partial.
2. Start ShardReplicaSet::init_empty_local_shard().
3. Let the future pass let current_shard = local.take();.
4. Drop the future before it writes Some(Shard::Local(_)) or Some(Shard::Dummy(_)) back into self.local.
5. Observe that has_local_shard() is false while wait_for_local(...) still succeeds from metadata.

The following whitebox test demonstrates the issue deterministically. It is a bug-existence validation test: it passes on the current implementation because it asserts the reachable bad state. It uses a test-only hook immediately after local.take() and does not rely on sleeps or scheduling probability.

diff --git a/lib/collection/src/shards/replica_set/mod.rs b/lib/collection/src/shards/replica_set/mod.rs
index e3bc8de59..2c41801f5 100644
--- a/lib/collection/src/shards/replica_set/mod.rs
+++ b/lib/collection/src/shards/replica_set/mod.rs
@@ -131,6 +131,27 @@ pub type AbortShardTransfer = Arc<dyn Fn(ShardTransfer, &str) + Send + Sync>;
 pub type ChangePeerState = Arc<dyn Fn(PeerId, ShardId) + Send + Sync>;
 pub type ChangePeerFromState = Arc<dyn Fn(PeerId, ShardId, Option<ReplicaState>) + Send + Sync>;
 
+#[cfg(test)]
+static INIT_EMPTY_LOCAL_SHARD_AFTER_TAKE_HOOK: std::sync::Mutex<
+    Option<(
+        tokio::sync::oneshot::Sender<()>,
+        tokio::sync::oneshot::Receiver<()>,
+    )>,
+> = std::sync::Mutex::new(None);
+
+#[cfg(test)]
+async fn wait_init_empty_local_shard_after_take() {
+    let hook = INIT_EMPTY_LOCAL_SHARD_AFTER_TAKE_HOOK
+        .lock()
+        .unwrap()
+        .take();
+
+    if let Some((reached, release)) = hook {
+        let _ = reached.send(());
+        let _ = release.await;
+    }
+}
+
 const REPLICA_STATE_FILE: &str = "replica_state.json";
 
 impl ShardReplicaSet {
@@ -587,6 +608,9 @@ impl ShardReplicaSet {
         let mut local = self.local.write().await;
 
         let current_shard = local.take();
+        #[cfg(test)]
+        wait_init_empty_local_shard_after_take().await;
+
         if let Some(current_shard) = current_shard {
             current_shard.stop_gracefully().await;
         }
@@ -1525,3 +1549,172 @@ impl ShardReplicaSet {
 pub enum Change {
     Remove(ShardId, PeerId),
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+    use std::num::NonZeroU32;
+    use std::sync::Arc;
+    use std::time::Duration;
+
+    use common::budget::ResourceBudget;
+    use common::save_on_disk::SaveOnDisk;
+    use segment::types::Distance;
+    use tempfile::{Builder, TempDir};
+    use tokio::runtime::Handle;
+    use tokio::sync::{RwLock, oneshot};
+
+    use super::*;
+    use crate::collection::payload_index_schema::PayloadIndexSchema;
+    use crate::common::adaptive_handle::AdaptiveSearchHandle;
+    use crate::config::{CollectionConfigInternal, CollectionParams, WalConfig};
+    use crate::operations::shared_storage_config::SharedStorageConfig;
+    use crate::operations::types::VectorsConfig;
+    use crate::operations::vector_params_builder::VectorParamsBuilder;
+    use crate::optimizers_builder::OptimizersConfig;
+    use crate::shards::channel_service::ChannelService;
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_cancel_init_empty_local_shard_after_local_take_leaves_no_local_slot() {
+        let collection_dir = Builder::new()
+            .prefix("init-empty-local-cancel")
+            .tempdir()
+            .unwrap();
+        let replica_set = new_dummy_receiver_replica_set(&collection_dir).await;
+
+        assert!(
+            replica_set.has_local_shard().await,
+            "test fixture must start with a local receiver slot"
+        );
+        assert!(
+            replica_set.is_dummy().await,
+            "test fixture must start with a dummy receiver shard"
+        );
+        replica_set
+            .wait_for_local(Duration::from_secs(1))
+            .await
+            .expect("test fixture metadata must mark this peer as local");
+
+        let (reached_tx, reached_rx) = oneshot::channel();
+        let (release_tx, release_rx) = oneshot::channel();
+        install_init_empty_local_shard_after_take_hook(reached_tx, release_rx);
+
+        let mut init_empty_local_shard = Box::pin(replica_set.init_empty_local_shard());
+
+        tokio::select! {
+            _ = reached_rx => {}
+            result = &mut init_empty_local_shard => {
+                panic!("empty local shard initialization completed before the cancellation window: {result:?}");
+            }
+        }
+
+        drop(init_empty_local_shard);
+        drop(release_tx);
+
+        assert!(
+            !replica_set.has_local_shard().await,
+            "dropping init_empty_local_shard after local.take() leaves no local receiver slot"
+        );
+        assert!(
+            !replica_set.is_dummy().await,
+            "the cancelled future dropped the dummy shard before it could install a replacement"
+        );
+        assert!(
+            !replica_set.is_proxy().await,
+            "None is not a proxy either, so receiver retry logic can skip both repair branches"
+        );
+        replica_set
+            .wait_for_local(Duration::from_secs(1))
+            .await
+            .expect("replica metadata still says this peer is local even though the slot is gone");
+    }
+
+    fn install_init_empty_local_shard_after_take_hook(
+        reached: oneshot::Sender<()>,
+        release: oneshot::Receiver<()>,
+    ) {
+        let mut hook = INIT_EMPTY_LOCAL_SHARD_AFTER_TAKE_HOOK.lock().unwrap();
+        assert!(
+            hook.is_none(),
+            "init-empty-local-shard test hook is already installed"
+        );
+        *hook = Some((reached, release));
+    }
+
+    async fn new_dummy_receiver_replica_set(collection_dir: &TempDir) -> ShardReplicaSet {
+        let update_runtime = Handle::current();
+        let search_runtime = AdaptiveSearchHandle::current_for_tests();
+
+        let wal_config = WalConfig {
+            wal_capacity_mb: 1,
+            wal_segments_ahead: 0,
+            wal_retain_closed: 1,
+        };
+
+        let collection_params = CollectionParams {
+            vectors: VectorsConfig::Single(VectorParamsBuilder::new(4, Distance::Dot).build()),
+            shard_number: NonZeroU32::new(1).unwrap(),
+            replication_factor: NonZeroU32::new(1).unwrap(),
+            write_consistency_factor: NonZeroU32::new(1).unwrap(),
+            ..CollectionParams::empty()
+        };
+
+        let optimizers_config = OptimizersConfig::fixture();
+        let config = CollectionConfigInternal {
+            params: collection_params,
+            optimizer_config: optimizers_config.clone(),
+            wal_config,
+            hnsw_config: Default::default(),
+            quantization_config: None,
+            strict_mode_config: None,
+            uuid: None,
+            metadata: None,
+        };
+
+        let payload_index_schema_file = collection_dir.path().join("payload-schema.json");
+        let payload_index_schema: Arc<SaveOnDisk<PayloadIndexSchema>> = Arc::new(
+            SaveOnDisk::load_or_init_default(payload_index_schema_file).unwrap(),
+        );
+        let shared_config = Arc::new(RwLock::new(config));
+
+        let replica_set = ShardReplicaSet::build(
+            1,
+            None,
+            "test_collection".to_string(),
+            1,
+            true,
+            HashSet::new(),
+            dummy_on_replica_failure(),
+            dummy_abort_shard_transfer(),
+            collection_dir.path(),
+            shared_config,
+            optimizers_config,
+            Arc::new(SharedStorageConfig::default()),
+            payload_index_schema,
+            ChannelService::default(),
+            update_runtime,
+            search_runtime,
+            ResourceBudget::default(),
+            Some(ReplicaState::Partial),
+        )
+        .await
+        .unwrap();
+
+        {
+            let mut local = replica_set.local.write().await;
+            *local = Some(Shard::Dummy(DummyShard::new(
+                "dummy receiver awaiting shard transfer",
+            )));
+        }
+
+        replica_set
+    }
+
+    fn dummy_on_replica_failure() -> ChangePeerFromState {
+        Arc::new(move |_peer_id, _shard_id, _from_state| {})
+    }
+
+    fn dummy_abort_shard_transfer() -> AbortShardTransfer {
+        Arc::new(|_shard_transfer, _reason| {})
+    }
+}

Run the targeted test:

cargo +nightly test -p collection --lib test_cancel_init_empty_local_shard_after_local_take_leaves_no_local_slot -- --nocapture

Observed result on the current implementation:

Finished `test` profile [unoptimized + debuginfo] target(s) in 2m 40s
warning: the following packages contain code that will be rejected by a future version of Rust: proc-macro-error2 v2.0.1
note: to see what the problems were, use the option `--future-incompat-report`, or run `cargo report future-incompatibilities --id 1`
     Running unittests src/lib.rs (target/debug/deps/collection-0998863105b688e9)

running 1 test
test shards::replica_set::tests::test_cancel_init_empty_local_shard_after_local_take_leaves_no_local_slot ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 216 filtered out; finished in 0.21s

The important assertion is:

assert!(
    !replica_set.has_local_shard().await,
    "dropping init_empty_local_shard after local.take() leaves no local receiver slot"
);

The test also checks the retry-relevant mismatch:

assert!(!replica_set.is_dummy().await);
assert!(!replica_set.is_proxy().await);
replica_set
    .wait_for_local(Duration::from_secs(1))
    .await
    .expect("replica metadata still says this peer is local even though the slot is gone");

This means the test is not merely proving that a hook fired. It proves the full bad state: the future was dropped after the local slot was taken, the slot remained empty, the dummy/proxy repair predicates became false, and metadata still allowed wait_for_local(...) to succeed.

For a post-fix regression test, the core assertion should be inverted or strengthened so that cancellation at this same point preserves a valid receiver slot:

assert!(
    replica_set.has_local_shard().await,
    "receiver initialization cancellation must not leave the local shard slot empty"
);

Expected Behavior

Cancelling or timing out receiver-side shard initialization should not leave ShardReplicaSet.local as None after a real local receiver slot has already been removed.

After any cancellation point in receiver initialization, one of these should remain true:

• the old valid receiver slot is still installed;
• a Shard::Dummy or equivalent explicit receiver-initializing/error state is installed;
• a new empty Shard::Local has been fully built and installed;
• the operation is deliberately moved into a must-complete task whose progress is not dropped with the gRPC request future and whose errors are still observed.

In particular, replica_state.is_local == true should not coexist indefinitely with self.local == None as a result of request timeout.

Possible Solution

The safest local invariant is:

self.local must contain a valid receiver state at every await boundary.

One practical direction is to replace the slot with a DummyShard or an explicit receiver-initializing shard state before performing awaited stop/clear/build work, instead of leaving the slot as None:

let current_shard = {
    let mut local = self.local.write().await;
    local.replace(Shard::Dummy(DummyShard::new(
        "Initializing empty local shard for incoming shard transfer",
    )))
};

if let Some(current_shard) = current_shard {
    current_shard.stop_gracefully().await;
}

LocalShard::clear(&self.shard_path).await?;

let local_shard_res = LocalShard::build(...).await;

let mut local = self.local.write().await;
match local_shard_res {
    Ok(local_shard) => *local = Some(Shard::Local(local_shard)),
    Err(err) => *local = Some(Shard::Dummy(DummyShard::new(...))),
}

That shape has two advantages:

• the self.local write lock is not held across long async filesystem/build operations;
• if the future is dropped at any await point, the replica set still has a modeled receiver state instead of None.

Another defensive option is a synchronous cancellation guard for the critical take -> commit section. The guard would install a valid Shard::Dummy fallback on Drop unless it has been disarmed by a successful commit. The guard must not require async cleanup in Drop.

A broader architectural option is to move receiver initialization into a runtime-owned task that is not cancelled when the request future is dropped. The gRPC handler can wait for the task result, but timing out the request must not destroy the in-progress state transition. This option should include explicit error reporting and shutdown behavior; simply detaching the task would risk hiding failures.

Increasing cluster.grpc_timeout_ms alone is not a correctness fix. It reduces the chance of hitting the window but does not remove the invalid intermediate state.

Context (Environment)

This affects distributed shard-transfer receiver preparation when the target replica is in a dummy/partial/recovery receiver state and the internal Initiate request times out or is otherwise cancelled after receiver initialization has started.

The default internal gRPC timeout is relatively large, so the issue is conditional. However, the timeout is configurable, the receiver path already waits for consensus state and performs filesystem-heavy shard cleanup/build work, and tonic timeout is a real cancellation source for the handler future. The dangerous window is small but sits after a shared in-memory state transition has already made the local slot invalid.

The bounded impact is an invalid in-memory receiver state. This issue does not claim that existing on-disk shard data is always permanently lost. The problem is that the running replica set can lose the local receiver object that should represent either a controlled dummy state or a new empty local shard. Later retry/recovery logic can then observe metadata-local state while the local slot itself is empty.

This is especially concerning for recovery flows where the dummy shard is the deliberate representation of an unhealthy or not-yet-restored local replica. Replacing that modeled state with None removes the signal that later code uses to enter the dummy-repair branch.

Detailed Description

The cancellation chain has three required links:

1. A cancel-unsafe future exists. init_empty_local_shard() mutates shared replica-set state with local.take() and then awaits before restoring a valid value.
2. That future can actually be cancelled. Internal requests carry grpc-timeout, and tonic's timeout wrapper completes with an error and drops the unfinished inner service future when the timeout wins.
3. Cancelling it violates a system property. The receiver can be left with self.local == None even though replica metadata still says the peer is local and transfer/recovery code expects a modeled local receiver state.

The issue is not that local.take() is always wrong. The issue is the combination of take() with awaited work before any replacement is installed. In async Rust, a timeout is implemented by dropping or no longer polling the unfinished future. Drop for a future cannot run async repair code, and ordinary lock release does not restore the protected value.

The problematic receiver retry shape is:

retry initiate_shard_transfer
    -> wait_for_local(...) succeeds from replica metadata
    -> is_proxy() is false because self.local is None
    -> is_dummy() is false because self.local is None
    -> init_empty_local_shard() is skipped
    -> handler can return Ok without recreating a local receiver slot

A DummyShard is a valid modeled state: it represents a local replica that cannot serve normally and returns controlled service errors. None is different. It means there is no local shard object at all, while the surrounding state machine can still consider this peer a local receiver.

This is distinct from source-side queue-proxy cancellation problems. Here the affected role is the target receiver, the trigger is framework request timeout of the internal Initiate service future, and the broken invariant is the receiver local-slot invariant during dummy-to-empty-local initialization.

Possible Implementation

A robust implementation should make receiver initialization follow this invariant:

No await is allowed while self.local is temporarily None because a receiver shard was taken out.

Concretely:

• avoid holding the self.local write lock across stop_gracefully, LocalShard::clear, and LocalShard::build;
• publish a valid fallback state, such as Shard::Dummy("initializing empty local shard for incoming transfer"), before the first awaited operation;
• perform async stop/clear/build outside the slot or behind a guard;
• commit Shard::Local only after the new local shard has been fully built;
• on build error, keep or replace the fallback dummy with a dummy containing the concrete error;
• add a regression test that drops the future at the same local.take() window and asserts the slot remains valid.

The fixed state machine should look like:

Some(Dummy or Local receiver)
    -> publish Some(Dummy initializing receiver)
    -> await stop/clear/build
    -> publish Some(Local) or Some(Dummy error)

rather than:

Some(Dummy or Local receiver)
    -> take slot, leaving None
    -> await stop/clear/build
    -> publish Some(Local) or Some(Dummy error)

After a fix, the whitebox test above can be reused as a regression test by preserving the deterministic pause and replacing the current bad-state assertions with:

assert!(
    replica_set.has_local_shard().await,
    "receiver initialization cancellation must preserve a modeled local receiver state"
);

If the chosen fix uses a dummy fallback, an even stronger regression assertion would be:

assert!(
    replica_set.is_dummy().await,
    "cancelling receiver initialization should leave a retry-visible dummy receiver state"
);

Start with init_empty_local_shard in lib/collection/src/shards/replica_set/mod.rs, then trace the cancellation path through lib/tonic/api/collections_internal_api.rs, lib/collection/src/collection/shard_transfer.rs, and lib/api/src/grpc/transport_channel_pool.rs. Run the deterministic whitebox test described in mod.rs and verify that cancelling after local.take() no longer leaves the replica without a local shard slot.