Repository: Migorithm/duva (rust)

Issue:
Bug: replica triggering rebalancing
## How to reproduce
- Load up leader
- Attach follower to the leader
- Set number of keys to the leader

- load up another leader
- send `CLUSTER MEET {host:post} eager`


## Bugs
Keys are distributed between leaders as expected
Followers on the other hand, lose all the data. 


## Expected reasons
There is no safe guard that protects followers from being involved in rebalancing process.  

Pull request:
feat: Separate leader and replica responsibilities in hash ring updates
## Problem
Previously, both leaders and replicas would initiate migration tasks when receiving hash ring updates, violating the leader-follower principle and potentially causing coordination conflicts and race conditions.


## Solution
Refactored `maybe_update_hashring()` to differentiate between leader and replica behavior:

### For Replicas:
- Only update the local hash ring for routing decisions
- Do not initiate any migration tasks
- Wait for leader to coordinate all migrations

### For Leaders:
- Create migration plans by comparing old vs new hash rings
- Update hash ring after migration planning (fixes critical bug)
- Coordinate all migration tasks across the cluster

## Key Changes
- **Fixed critical bug**: Hash ring was being updated too early for leaders, preventing proper migration task creation, closed #613 
- **Improved distributed systems design**: Clear separation of concerns between leader and replicas
- **Enhanced reliability**: Eliminates potential race conditions from multiple nodes trying to coordinate migrations

## Testing
- Added `test_maybe_update_hashring_replica_only_updates_ring()` to verify replicas only update hash rings
- Existing tests continue to pass, ensuring backward compatibility

This change ensures proper leader-follower semantics while maintaining data consistency during cluster rebalancing operations.

Patch:
diff --git a/duva/src/domains/cluster_actors/actor.rs b/duva/src/domains/cluster_actors/actor.rs
index 9cb1b32f..dc4f77f4 100644
--- a/duva/src/domains/cluster_actors/actor.rs
+++ b/duva/src/domains/cluster_actors/actor.rs
@@ -229,7 +229,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
         self.join_peer_network_if_absent(heartbeat.cluster_nodes).await;
         self.gossip(heartbeat.hop_count).await;
         self.update_on_hertbeat_message(&heartbeat.from, heartbeat.hwm);
-        self.schedule_migration_if_required(heartbeat.hashring, cache_manager, None).await;
+        self.maybe_update_hashring(heartbeat.hashring, cache_manager, None).await;
     }
 
     pub(crate) async fn req_consensus(&mut self, req: ConsensusRequest) {
@@ -418,7 +418,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
         }
     }
 
-    #[instrument(level = tracing::Level::INFO, skip(self), fields(request_from = %request_from))]
+    #[instrument(level = tracing::Level::INFO, skip(self,cache_manager,cluster_handler), fields(request_from = %request_from))]
     pub(crate) async fn start_rebalance(
         &mut self,
         request_from: PeerIdentifier,
@@ -455,12 +455,8 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
             .set_hashring(new_hash_ring.clone());
 
         self.send_heartbeat(hb).await;
-        self.schedule_migration_if_required(
-            Some(new_hash_ring),
-            cache_manager,
-            cluster_handler.clone(),
-        )
-        .await;
+        self.maybe_update_hashring(Some(new_hash_ring), cache_manager, cluster_handler.clone())
+            .await;
     }
 
     fn hop_count(fanout: usize, node_count: usize) -> u8 {
@@ -939,7 +935,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
     }
 
     // * If the hashring is valid, make a plan to migrate data for each paritition
-    async fn schedule_migration_if_required(
+    async fn maybe_update_hashring(
         &mut self,
         hashring: Option<HashRing>,
         cache_manager: &CacheManager,
@@ -957,16 +953,27 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
             return;
         }
 
+        // For replicas, just update the hash ring and wait for leader to coordinate migrations
+        if !self.replication.is_leader_mode {
+            self.hash_ring = new_ring;
+            info!("Replica updated hash ring, waiting for leader to coordinate migrations");
+            return;
+        }
+
+        // Leader-only migration coordination logic below
+        // Keep the old ring to compare with new ring for migration planning
         let keys = cache_manager.route_keys(None).await;
         let migration_plans = self.hash_ring.create_migration_tasks(&new_ring, keys);
 
-        // This is the critical update. The local hash ring is updated *before* migration tasks are sent.
+        // Update the hash ring after creating migration plans
         self.hash_ring = new_ring;
 
         if migration_plans.is_empty() {
             info!("No migration tasks to schedule");
             return;
         }
+
+        info!("Leader scheduling {} migration plan(s)", migration_plans.len());
         self.block_write_reqs();
 
         let batch_handles = FuturesUnordered::new();


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 da493f31..2b0e2ad3 100644
--- a/duva/src/domains/cluster_actors/actor/tests/partitionings.rs
+++ b/duva/src/domains/cluster_actors/actor/tests/partitionings.rs
@@ -146,7 +146,7 @@ async fn test_start_rebalance_happy_path() {
 }
 
 #[tokio::test]
-async fn test_schedule_migration_if_required_when_noplan_is_made() {
+async fn test_maybe_update_hashring_when_noplan_is_made() {
     // GIVEN
     let mut cluster_actor = cluster_actor_create_helper(ReplicationRole::Leader).await;
     let last_modified = cluster_actor.hash_ring.last_modified;
@@ -167,9 +167,7 @@ async fn test_schedule_migration_if_required_when_noplan_is_made() {
 
     // WHEN
     let (_hwm, cache_manager) = cache_manager_create_helper();
-    cluster_actor
-        .schedule_migration_if_required(Some(hash_ring.clone()), &cache_manager, None)
-        .await;
+    cluster_actor.maybe_update_hashring(Some(hash_ring.clone()), &cache_manager, None).await;
 
     // THEN
     assert!(cluster_actor.pending_requests.is_none());
@@ -186,7 +184,7 @@ async fn test_make_migration_plan_when_given_hashring_is_same() {
     // WHEN
     let (_hwm, cache_manager) = cache_manager_create_helper();
     cluster_actor
-        .schedule_migration_if_required(Some(cluster_actor.hash_ring.clone()), &cache_manager, None)
+        .maybe_update_hashring(Some(cluster_actor.hash_ring.clone()), &cache_manager, None)
         .await;
 
     // THEN
@@ -201,7 +199,7 @@ async fn test_make_migration_plan_when_no_hashring_given() {
 
     // WHEN
     let (_hwm, cache_manager) = cache_manager_create_helper();
-    cluster_actor.schedule_migration_if_required(None, &cache_manager, None).await;
+    cluster_actor.maybe_update_hashring(None, &cache_manager, None).await;
 
     // THEN
     assert_eq!(cluster_actor.hash_ring.last_modified, last_modified);
@@ -218,7 +216,7 @@ async fn test_make_migration_plan_when_last_modified_is_lower_than_its_own() {
 
     // WHEN
     let (_hwm, cache_manager) = cache_manager_create_helper();
-    cluster_actor.schedule_migration_if_required(Some(hash_ring), &cache_manager, None).await;
+    cluster_actor.maybe_update_hashring(Some(hash_ring), &cache_manager, None).await;
 
     // THEN
     assert_eq!(cluster_actor.hash_ring.last_modified, last_modified);
@@ -703,7 +701,7 @@ async fn test_handle_migration_ack_success_case_with_pending_reqs_and_migration(
     // assert!(cluster_actor.pending_migrations.is_none());
 }
 
-// Verify that the start_rebalance -> schedule_migration_if_required flow works.
+// Verify that the start_rebalance -> maybe_update_hashring flow works.
 // This test addresses the TODO comment: "need to see if migration batch is scheduled."
 #[tokio::test]
 async fn test_start_rebalance_schedules_migration_batches() {
@@ -764,3 +762,39 @@ async fn test_start_rebalance_schedules_migration_batches() {
     // 3. Verify pending_requests is set (synchronous part)
     assert!(cluster_actor.pending_requests.is_some());
 }
+
+#[tokio::test]
+async fn test_maybe_update_hashring_replica_only_updates_ring() {
+    // GIVEN - Create a replica actor (not leader)
+    let mut cluster_actor = cluster_actor_create_helper(ReplicationRole::Follower).await;
+    let (_hwm, cache_manager) = cache_manager_create_helper_with_keys(vec![
+        "replica_key_1".to_string(),
+        "replica_key_2".to_string(),
+    ])
+    .await;
+
+    let original_ring = cluster_actor.hash_ring.clone();
+    let original_modified = original_ring.last_modified;
+
+    // Create a new hash ring with different configuration
+    let new_node_replid = ReplicationId::Key("new_node".to_string());
+    let new_ring = HashRing::default()
+        .add_partition_if_not_exists(new_node_replid, PeerIdentifier::new("127.0.0.1", 6000))
+        .unwrap()
+        .add_partition_if_not_exists(
+            cluster_actor.replication.replid.clone(),
+            cluster_actor.replication.self_identifier(),
+        )
+        .unwrap();
+
+    // WHEN - Replica receives hash ring update
+    cluster_actor.maybe_update_hashring(Some(new_ring.clone()), &cache_manager, None).await;
+
+    // THEN - Hash ring should be updated
+    assert_eq!(cluster_actor.hash_ring, new_ring);
+    assert_ne!(cluster_actor.hash_ring.last_modified, original_modified);
+
+    // But no migration tasks should be initiated (no pending requests/migrations)
+    assert!(cluster_actor.pending_requests.is_none());
+    assert!(cluster_actor.pending_migrations.is_none());
+}