Repository: Migorithm/duva (rust)

Issue:
bug: late update of hashring when partitioning / disable key validation during migration
Early update of hashring makes other peers irresponsive. 

catch 22
- If we update hash ring early: Migration batches get rejected due to race condition
- If we update hash ring late: verify_key_belongs_to_node fails because it's checking against the old ring


solution to these are moving update logic at the end of migration and disabling verify keys during migration 

Pull request:
fix: apply optimistic migration, resolve hashring update timing issue
resolved #682 

Patch:
diff --git a/diagrams/consistent-hashing-server.mmd b/diagrams/consistent-hashing-server.mmd
index ef0727ee..79e07ed4 100644
--- a/diagrams/consistent-hashing-server.mmd
+++ b/diagrams/consistent-hashing-server.mmd
@@ -43,6 +43,9 @@ sequenceDiagram
 
         par 
             A -->> A: Delete write operation (or Snaphsot and drop logs for optimization)
+            critical 
+                A -->> A: Update hashring
+            end 
             A -->> A: Unlock pending requests
             alt Request's key is in A's range
                 A ->> A: Process request
@@ -50,9 +53,14 @@ sequenceDiagram
             else Request's key is in B's range
                 A -->> C1: (Moved) error
             end
+
+            
             A --) C1 : Notify token map update
         and
             B -->> B: Delete write operation (or Snaphsot and drop logs for optimization)
+            critical 
+                B -->> B: Update hashring
+            end 
             B -->> B: Unlock pending requests
             B -->> C2 : Notify token map update
         end
diff --git a/duva-client/src/broker/mod.rs b/duva-client/src/broker/mod.rs
index c9bb98a9..f9453366 100644
--- a/duva-client/src/broker/mod.rs
+++ b/duva-client/src/broker/mod.rs
@@ -1,9 +1,7 @@
 mod input_queue;
 mod read_stream;
 mod write_stream;
-
 use crate::command::Input;
-
 use duva::domains::{IoError, query_io::QueryIO};
 use duva::prelude::tokio;
 use duva::prelude::tokio::net::TcpStream;
diff --git a/duva/src/domains/cluster_actors/actor.rs b/duva/src/domains/cluster_actors/actor.rs
index 4ddba907..1ae68754 100644
--- a/duva/src/domains/cluster_actors/actor.rs
+++ b/duva/src/domains/cluster_actors/actor.rs
@@ -1064,11 +1064,9 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
         let keys = cache_manager.route_keys(None).await;
         let migration_plans = self.hash_ring.create_migration_tasks(&new_ring, keys);
 
-        // Update the hash ring after creating migration plans
-        self.hash_ring = *new_ring;
-
         if migration_plans.is_empty() {
             info!("No migration tasks to schedule");
+            self.hash_ring = *new_ring;
             return;
         }
 
@@ -1174,18 +1172,6 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
     ) -> anyhow::Result<()> {
         let peer = self.members.get_mut(&from).context("No Member Found")?;
 
-        // Validation against the now-updated hash ring - check if all keys belong to this node
-        let keys_to_validate: Vec<&str> =
-            migrate_batch.cache_entries.iter().map(|entry| entry.key()).collect();
-        if !self
-            .hash_ring
-            .verify_key_belongs_to_node(&keys_to_validate, &self.replication.self_identifier())
-        {
-            error!("Received batch contains keys that do not belong to this node");
-            peer.send(MigrationBatchAck::with_reject(migrate_batch.batch_id)).await?;
-            return Ok(());
-        }
-
         // If cache entries are empty, skip consensus and directly send success ack
         if migrate_batch.cache_entries.is_empty() {
             peer.send(MigrationBatchAck::with_success(migrate_batch.batch_id)).await?;
@@ -1259,6 +1245,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
         tokio::spawn({
             let handler = self.self_handler.clone();
             let cache_manager = cache_manager.clone();
+
             async move {
                 if rx.await.is_ok() {
                     let _ = cache_manager.route_delete(pending_migration_batch.keys).await; // reflect state change
@@ -1282,6 +1269,10 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
                 if pending_reqs.is_empty() {
                     return;
                 }
+                if let Some(new_ring) = self.hash_ring.set_partitions(self.shard_leaders()) {
+                    self.hash_ring = new_ring;
+                }
+
                 let handler = self.self_handler.clone();
                 tokio::spawn(async move {
                     while let Some(req) = pending_reqs.pop_front() {


Test patch:
diff --git a/duva/src/domains/cluster_actors/actor/tests/partitionings.rs b/duva/src/domains/cluster_actors/actor/tests/partitionings.rs
index 6e52b71e..61b9c429 100644
--- a/duva/src/domains/cluster_actors/actor/tests/partitionings.rs
+++ b/duva/src/domains/cluster_actors/actor/tests/partitionings.rs
@@ -427,37 +427,6 @@ async fn test_receive_batch_success_path_when_noreplica_found() {
     assert_eq!(keys.len(), 2);
 }
 
-#[tokio::test]
-async fn test_receive_batch_validation_failure_keys_not_belonging_to_node() {
-    // GIVEN
-    let mut cluster_actor = cluster_actor_create_helper(ReplicationRole::Leader).await;
-    // Create a hash ring where another node is responsible for the keys
-    let hash_ring = HashRing::default();
-    let other_node_replid = ReplicationId::Key("other_node".to_string());
-    let hash_ring = hash_ring
-        .set_partitions(vec![(other_node_replid, PeerIdentifier("127.0.0.1:5000".into()))])
-        .unwrap();
-    cluster_actor.hash_ring = hash_ring;
-
-    // ! this is the one for receiving actor
-
-    let (_hwm, cache_manager) = cache_manager_create_helper();
-    let (message_buf, sender_peer_id) = cluster_actor.test_add_peer(6565, None, true);
-
-    let cache_entries = cache_entries_create_helper(&[("key1", "value1"), ("key2", "value2")]);
-    let migrate_batch = migration_batch_create_helper("validation_test", cache_entries);
-
-    // WHEN
-    cluster_actor.receive_batch(migrate_batch, &cache_manager, sender_peer_id).await.unwrap();
-
-    // THEN
-    assert_expected_queryio(
-        &message_buf,
-        MigrationBatchAck { batch_id: BatchId("validation_test".into()), success: false },
-    )
-    .await;
-}
-
 #[tokio::test]
 async fn test_unblock_write_reqs_if_done_when_no_pending_migrations() {
     // GIVEN
@@ -613,6 +582,12 @@ async fn test_handle_migration_ack_batch_id_not_found() {
 async fn test_handle_migration_ack_success_case_with_pending_reqs_and_migration() {
     // GIVEN
     let mut cluster_actor = setup_blocked_cluster_actor_with_requests(2).await;
+    cluster_actor.test_add_peer(
+        6055,
+        Some(ReplicationId::Key(uuid::Uuid::now_v7().to_string())),
+        true,
+    );
+
     let (_hwm, cache_manager) = cache_manager_create_helper();
 
     // Set up test keys in cache that will be part of the migration
@@ -681,16 +656,14 @@ async fn test_start_rebalance_schedules_migration_batches() {
 
     // THEN
     // 1. Verify heartbeat was sent immediately (synchronous part)
-    assert_expected_queryio(
-        &buf,
-        QueryIO::ClusterHeartBeat(HeartBeat {
-            from: cluster_actor.replication.self_identifier(),
-            hashring: Some(Box::new(cluster_actor.hash_ring.clone())),
-            replid: cluster_actor.replication.replid.clone(),
-            ..Default::default()
-        }),
-    )
-    .await;
+    let QueryIO::ClusterHeartBeat(HeartBeat { hashring, .. }) =
+        buf.lock().await.pop_front().unwrap()
+    else {
+        panic!()
+    };
+
+    assert!(hashring.is_some());
+    assert_ne!(*hashring.unwrap(), cluster_actor.hash_ring);
 
     // 2. Wait for migration batch message with timeout (asynchronous part)
     let batch = tokio::time::timeout(Duration::from_millis(1000), async {