fix: keep recovery marker cleanup under lock

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

@@ -57,6 +57,7 @@ use crate::shards::shard_trait::WaitUntil;
 #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
 enum DummyInitPhase {
     BeforeInit,
+    BeforeMarkerRemoval,
 }
 
 #[cfg(test)]
@@ -667,13 +668,18 @@ impl ShardReplicaSet {
         }
         self.init_empty_local_shard_locked(&mut local).await?;
 
-        drop(local);
+        #[cfg(test)]
+        pause_dummy_init(&self.shard_path, DummyInitPhase::BeforeMarkerRemoval).await;
 
         let shard_flag =
             crate::shards::shard_initializing_flag_path(collection_path, self.shard_id);
-        if fs_err::tokio::try_exists(&shard_flag).await.is_ok() {
-            fs_err::tokio::remove_file(&shard_flag).await?;
+        // A newer recovery must not replace this marker before we remove it.
+        match fs_err::tokio::remove_file(&shard_flag).await {
+            Ok(()) => {}
+            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
+            Err(err) => return Err(err.into()),
         }
+        drop(local);
         Ok(())
     }
 

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

@@ -559,6 +559,55 @@ async fn test_delayed_dummy_init_preserves_replacement_data() {
     replica_set.stop_gracefully().await;
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_dummy_init_preserves_new_recovery_marker() {
+    let dir = Builder::new()
+        .prefix("dummy-init-marker")
+        .tempdir()
+        .unwrap();
+    let replica_set = new_shard_replica_set(&dir, TEST_TARGET_SHARD_ID).await;
+    replica_set
+        .set_replica_state(TEST_PEER_ID, ReplicaState::Recovery)
+        .await
+        .unwrap();
+    replica_set
+        .clear_local_for_snapshot_recovery(dir.path(), || Ok(()))
+        .await
+        .unwrap();
+    let (reached, resume) = pause_dummy_initialization(
+        &replica_set,
+        crate::shards::replica_set::DummyInitPhase::BeforeMarkerRemoval,
+    );
+    let init = replica_set.init_dummy_local_shard(dir.path(), || Ok(()));
+    tokio::pin!(init);
+    tokio::select! {
+        result = reached => result.unwrap(),
+        result = &mut init => panic!("init missed the marker pause: {result:?}"),
+    }
+
+    let (clear_started_tx, mut clear_started_rx) = oneshot::channel();
+    let clear = replica_set.clear_local_for_snapshot_recovery(dir.path(), || {
+        clear_started_tx.send(()).unwrap();
+        Ok(())
+    });
+    tokio::pin!(clear);
+    assert!(futures::poll!(&mut clear).is_pending());
+    // The newer clear must wait for initialization to remove the old marker.
+    assert!(matches!(
+        clear_started_rx.try_recv(),
+        Err(oneshot::error::TryRecvError::Empty)
+    ));
+    resume.send(()).unwrap();
+    init.await.unwrap();
+    clear.await.unwrap();
+    assert!(replica_set.is_dummy().await);
+    assert!(
+        shard_initializing_flag_path(dir.path(), TEST_TARGET_SHARD_ID).exists(),
+        "old Initiate removed the newer recovery marker"
+    );
+    replica_set.stop_gracefully().await;
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_dummy_init_rechecks_sender_after_waiting() {
     use crate::shards::shard_holder::ShardHolder;
@@ -641,6 +690,37 @@ async fn test_dummy_init_rechecks_state_after_waiting() {
     replica_set.stop_gracefully().await;
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_dummy_init_without_recovery_marker() {
+    let dir = Builder::new()
+        .prefix("dummy-init-no-marker")
+        .tempdir()
+        .unwrap();
+    let replica_set = new_shard_replica_set(&dir, TEST_TARGET_SHARD_ID).await;
+    replica_set
+        .set_replica_state(TEST_PEER_ID, ReplicaState::Recovery)
+        .await
+        .unwrap();
+    replica_set
+        .clear_local_for_snapshot_recovery(dir.path(), || Ok(()))
+        .await
+        .unwrap();
+    // A dummy loaded after a shard load failure need not have a recovery marker.
+    fs_err::tokio::remove_file(shard_initializing_flag_path(
+        dir.path(),
+        TEST_TARGET_SHARD_ID,
+    ))
+    .await
+    .unwrap();
+    replica_set
+        .init_dummy_local_shard(dir.path(), || Ok(()))
+        .await
+        .unwrap();
+    assert!(!replica_set.is_dummy().await);
+    assert_eq!(count_points(&replica_set).await, 0);
+    replica_set.stop_gracefully().await;
+}
+
 /// Build a valid unpacked shard snapshot to recover from.
 ///
 /// Returns the temp dir - which the caller must keep alive - and the replica path