Mark *new* local replicas as locally-disabled, when recovering Raft snapshot (#7684)

* Remove faulty mark-as-dead condition when recovering Raft snapshot 😅

* Mark new local replicas as locally-disabled when recovering Raft snapshot

* Add test

lib/collection/src/shards/replica_set/mod.rs

@@ -1108,12 +1108,19 @@ impl ShardReplicaSet {
     ///
     /// If `from_state` is given, the peer will only be disabled if the given state matches
     /// consensus.
-    fn add_locally_disabled(
+    pub fn add_locally_disabled(
         &self,
-        state: &ReplicaSetState,
+        state: Option<&ReplicaSetState>,
         peer_id: PeerId,
         from_state: Option<ReplicaState>,
     ) {
+        let mut state_guard = None;
+
+        let state = match state {
+            Some(state) => state,
+            None => state_guard.insert(self.replica_state.read()),
+        };
+
         let other_peers = state
             .active_or_resharding_peers()
             .filter(|id| id != &peer_id);

lib/collection/src/shards/replica_set/snapshots.rs

@@ -434,10 +434,7 @@ impl ShardReplicaSet {
                 )));
 
                 // Mark local replica as Dead since it's dummy and dirty
-                {
-                    let replica_state = self.replica_state.read();
-                    self.add_locally_disabled(&replica_state, self.this_peer_id(), None);
-                }
+                self.add_locally_disabled(None, self.this_peer_id(), None);
 
                 // Remove inner shard data but keep the shard folder with its configuration files.
                 // This way the shard can be read on startup and the user can decide what to do next.

lib/collection/src/shards/replica_set/update.rs

@@ -152,7 +152,7 @@ impl ShardReplicaSet {
                         // Deactivate the peer if forwarding failed with transient error
                         let replica_state = self.replica_state.read();
                         let from_state = replica_state.get_peer_state(leader_peer);
-                        self.add_locally_disabled(&replica_state, leader_peer, from_state);
+                        self.add_locally_disabled(Some(&replica_state), leader_peer, from_state);
 
                         // Return service error
                         CollectionError::service_error(format!(
@@ -647,7 +647,7 @@ impl ShardReplicaSet {
             // Always deactivate the replica if its in a shard transfer related state
             let from_state = Some(peer_state).filter(|state| !state.is_partial_or_recovery());
 
-            self.add_locally_disabled(state, *peer_id, from_state);
+            self.add_locally_disabled(Some(state), *peer_id, from_state);
         }
 
         wait_for_deactivation

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

@@ -244,14 +244,15 @@ impl TableOfContent {
                 // if collection has been created during snapshot application
                 if !collection_exists {
                     for shard_id in collection.get_local_shards().await {
-                        collection
-                            .set_shard_replica_state(
-                                shard_id,
-                                self.this_peer_id,
-                                ReplicaState::Dead,
-                                None,
-                            )
-                            .await?;
+                        let shard_holder = collection.shards_holder().read_owned().await;
+
+                        let Some(replica_set) = shard_holder.get_shard(shard_id) else {
+                            continue;
+                        };
+
+                        if replica_set.is_local().await {
+                            replica_set.add_locally_disabled(None, self.this_peer_id, None);
+                        }
                     }
                 }
             }

tests/consensus_tests/test_consensus_compaction.py

@@ -5,7 +5,8 @@ from time import sleep
 from typing import Any
 
 import requests
-from consensus_tests.fixtures import create_collection, drop_collection
+
+from .fixtures import *
 from .utils import *
 
 N_PEERS = 3
@@ -134,6 +135,43 @@ def test_consensus_compaction_shard_keys(tmp_path: pathlib.Path):
     # - <https://github.com/qdrant/qdrant/pull/6212>
     wait_for_shard_keys(peer_api_uris, "test_collection", SHARD_KEYS)
 
+@pytest.mark.parametrize(
+    "replication_factor",
+    [1, 2]
+)
+def test_consensus_snapshot_create_collection(tmp_path: pathlib.Path, replication_factor: int):
+    assert_project_root()
+
+    N_PEERS = 3
+
+    env = {
+        "QDRANT__LOG_LEVEL": "debug",
+        # Aggressively compact consensus WAL
+        "QDRANT__CLUSTER__CONSENSUS__COMPACT_WAL_ENTRIES": "1",
+    }
+
+    # Start cluster
+    peers, peer_dirs, bootstrap_uri = start_cluster(tmp_path, N_PEERS, extra_env=env)
+
+    # Get last peer ID
+    last_peer_id = get_cluster_info(peers[-1])['peer_id']
+
+    # Kill last peer
+    processes.pop().kill()
+
+    # Bootstrap collection
+    create_collection(peers[0], shard_number=3, replication_factor=replication_factor)
+    wait_collection_on_all_peers("test_collection", peers[:N_PEERS - 1], 10)
+
+    upsert_random_points(peers[0], 1000, fail_on_error=False)
+
+    # Restart last peer
+    peers[-1] = start_peer(peer_dirs[-1], "peer_2_restarted.log", bootstrap_uri)
+    wait_for_peer_online(peers[-1])
+
+    # Wait for last peer recovery
+    wait_collection_exists_and_active_on_all_peers("test_collection", peers, 10)
+
 
 def put_metadata_key(peer_uris: list[str], key: str, value: Any):
     resp = requests.put(f"{peer_uris[0]}/cluster/metadata/keys/{key}?wait=true", json=value)