Repository: Migorithm/duva (rust)

Issue:
feat: consensus-based sets on migration during rebalancing
### Description


Currently, sets operations are being handled as cache writes on migration during rebalancing. This is incorrect and causes issues during log replay and state restoration.



### Acceptance Criteria


- [ ]  modify `receive_batch`  so it goes through consensus
- [ ] Relevant unit tests are updated/added.





Pull request:
feat: set consensus on rebalancing
resolved #614 

Patch:
diff --git a/duva/src/domains/caches/cache_manager.rs b/duva/src/domains/caches/cache_manager.rs
index f5fa6ca4..7edb877b 100644
--- a/duva/src/domains/caches/cache_manager.rs
+++ b/duva/src/domains/caches/cache_manager.rs
@@ -66,11 +66,6 @@ impl CacheManager {
         Ok(format!("s:{}|idx:{}", value, current_idx))
     }
 
-    pub(crate) async fn route_direct_set(&self, cache_entry: CacheEntry) -> Result<()> {
-        self.select_shard(cache_entry.key()).send(CacheCommand::Set { cache_entry }).await?;
-        Ok(())
-    }
-
     pub(crate) async fn route_save(
         &self,
         save_target: SaveTarget,
@@ -113,7 +108,7 @@ impl CacheManager {
             | WriteRequest::Incr { key, delta } => {
                 self.route_numeric_delta(key, delta, log_index).await?;
             },
-            | WriteRequest::BulkSet { entries } => {
+            | WriteRequest::MSet { entries } => {
                 self.route_bulk_set(entries).await?;
             },
         };
@@ -288,12 +283,20 @@ impl CacheManager {
     async fn route_bulk_set(&self, entries: Vec<CacheEntry>) -> Result<()> {
         let closure = |entry| -> CacheCommand { CacheCommand::Set { cache_entry: entry } };
         FuturesOrdered::from_iter(entries.into_iter().map(|entry| async move {
-            let _ = self.select_shard(&entry.key()).send(closure(entry)).await;
+            let _ = self.select_shard(entry.key()).send(closure(entry)).await;
         }))
         .collect::<Vec<_>>()
         .await;
         Ok(())
     }
+
+    pub(crate) async fn route_mset(&self, cache_entries: Vec<CacheEntry>) {
+        join_all(cache_entries.into_iter().map(|entry| async move {
+            let _ =
+                self.select_shard(entry.key()).send(CacheCommand::Set { cache_entry: entry }).await;
+        }))
+        .await;
+    }
 }
 
 #[cfg(test)]
diff --git a/duva/src/domains/cluster_actors/actor.rs b/duva/src/domains/cluster_actors/actor.rs
index bd7c3c66..9cb1b32f 100644
--- a/duva/src/domains/cluster_actors/actor.rs
+++ b/duva/src/domains/cluster_actors/actor.rs
@@ -12,7 +12,6 @@ use super::replication::ReplicationState;
 use super::replication::time_in_secs;
 use super::*;
 use crate::domains::QueryIO;
-use crate::domains::caches::cache_manager;
 use crate::domains::caches::cache_manager::CacheManager;
 use crate::domains::caches::cache_objects::CacheEntry;
 use crate::domains::cluster_actors::hash_ring::BatchId;
@@ -36,7 +35,6 @@ use crate::domains::peers::peer::NodeKind;
 use crate::domains::peers::peer::PeerState;
 use crate::err;
 use client_sessions::ClientSessions;
-use futures::future::join_all;
 use heartbeat_scheduler::HeartBeatScheduler;
 
 use std::collections::HashMap;
@@ -1087,25 +1085,46 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
             return;
         }
 
-        let write_results = join_all(
-            migrate_batch
-                .cache_entries
-                .into_iter()
-                .map(|entry| cache_manager.route_direct_set(entry)),
-        )
-        .await;
+        // If cache entries are empty, skip consensus and directly send success ack
+        if migrate_batch.cache_entries.is_empty() {
+            let _ = peer.send(MigrationBatchAck::with_success(migrate_batch.batch_id)).await;
+            return;
+        }
 
-        let all_successful = write_results.iter().all(|r| r.is_ok());
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        self.req_consensus(ConsensusRequest::new(
+            WriteRequest::MSet { entries: migrate_batch.cache_entries.clone() },
+            tx,
+            None,
+        ))
+        .await;
 
-        if all_successful {
-            let _ = peer.send(MigrationBatchAck::with_success(migrate_batch.batch_id)).await;
-        } else {
-            error!(
-                "Failed to write some keys during migration for batch {}",
-                migrate_batch.batch_id.0
-            );
-            let _ = peer.send(MigrationBatchAck::with_reject(migrate_batch.batch_id)).await;
+        // * If there are no replicas, we can directly apply the state
+        if self.replicas().count() == 0 {
+            let _ = cache_manager.route_mset(migrate_batch.cache_entries.clone()).await;
         }
+
+        // * If there are replicas, we need to wait for the consensus to be applied which should be done in the background
+        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_mset(migrate_batch.cache_entries.clone()).await; // reflect state change
+                    let _ = handler
+                        .send(SchedulerMessage::SendBatchAck {
+                            batch_id: migrate_batch.batch_id,
+                            to: from,
+                        })
+                        .await;
+                    return;
+                }
+                error!(
+                    "Failed to write some keys during migration for batch {}",
+                    migrate_batch.batch_id.0
+                );
+            }
+        });
     }
 
     pub(crate) async fn handle_migration_ack(
@@ -1172,4 +1191,11 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
             }
         }
     }
+
+    pub(crate) async fn send_batch_ack(&mut self, batch_id: BatchId, to: PeerIdentifier) {
+        let Some(peer) = self.members.get_mut(&to) else {
+            return;
+        };
+        let _ = peer.send(MigrationBatchAck::with_success(batch_id)).await;
+    }
 }
diff --git a/duva/src/domains/cluster_actors/command.rs b/duva/src/domains/cluster_actors/command.rs
index 7c66263a..8b9170f3 100644
--- a/duva/src/domains/cluster_actors/command.rs
+++ b/duva/src/domains/cluster_actors/command.rs
@@ -1,5 +1,5 @@
 use crate::ReplicationState;
-use crate::domains::cluster_actors::hash_ring::MigrationBatch;
+use crate::domains::cluster_actors::hash_ring::{BatchId, MigrationBatch};
 use crate::domains::cluster_actors::replication::{ReplicationId, ReplicationRole};
 use crate::domains::operation_logs::WriteRequest;
 use crate::domains::peers::command::PeerCommand;
@@ -25,6 +25,7 @@ pub enum SchedulerMessage {
     RebalanceRequest { request_to: PeerIdentifier, lazy_option: LazyOption },
     ScheduleMigrationBatch(MigrationBatch, tokio::sync::oneshot::Sender<anyhow::Result<()>>),
     TryUnblockWriteReqs,
+    SendBatchAck { batch_id: BatchId, to: PeerIdentifier },
 }
 impl From<SchedulerMessage> for ClusterCommand {
     fn from(msg: SchedulerMessage) -> Self {
diff --git a/duva/src/domains/cluster_actors/service.rs b/duva/src/domains/cluster_actors/service.rs
index ad0ed7f8..90de5982 100644
--- a/duva/src/domains/cluster_actors/service.rs
+++ b/duva/src/domains/cluster_actors/service.rs
@@ -57,6 +57,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
                 self.migrate_batch(tasks, cache_manager, callback).await;
             },
             | TryUnblockWriteReqs => self.unblock_write_reqs_if_done(),
+            | SendBatchAck { batch_id, to } => self.send_batch_ack(batch_id, to).await,
         }
     }
 
diff --git a/duva/src/domains/operation_logs/operation.rs b/duva/src/domains/operation_logs/operation.rs
index d07da1ae..4b109ec6 100644
--- a/duva/src/domains/operation_logs/operation.rs
+++ b/duva/src/domains/operation_logs/operation.rs
@@ -13,7 +13,7 @@ pub struct WriteOperation {
 #[derive(Debug, Clone, PartialEq, Eq, bincode::Encode, bincode::Decode)]
 pub enum WriteRequest {
     Set { key: String, value: String, expires_at: Option<u64> },
-    BulkSet { entries: Vec<CacheEntry> },
+    MSet { entries: Vec<CacheEntry> },
     Delete { keys: Vec<String> },
     Append { key: String, value: String },
     Decr { key: String, delta: i64 },
@@ -53,7 +53,7 @@ impl WriteRequest {
             | WriteRequest::Incr { key, .. } => Some(key),
             | WriteRequest::Decr { key, .. } => Some(key),
             | WriteRequest::Delete { .. } => None, // Multiple keys, no single representative
-            | WriteRequest::BulkSet { .. } => None, // Multiple entries, no single representative
+            | WriteRequest::MSet { .. } => None,   // Multiple entries, no single representative
         }
     }
 
@@ -66,7 +66,7 @@ impl WriteRequest {
             | WriteRequest::Incr { key, .. } => vec![key],
             | WriteRequest::Decr { key, .. } => vec![key],
             | WriteRequest::Delete { keys, .. } => keys.iter().map(|k| k.as_str()).collect(),
-            | WriteRequest::BulkSet { entries } => entries.iter().map(|e| e.key()).collect(),
+            | WriteRequest::MSet { entries } => entries.iter().map(|e| e.key()).collect(),
         }
     }
 }
diff --git a/duva/src/domains/peers/command.rs b/duva/src/domains/peers/command.rs
index 81f018d6..387f11f7 100644
--- a/duva/src/domains/peers/command.rs
+++ b/duva/src/domains/peers/command.rs
@@ -123,7 +123,7 @@ mod peer_messages {
         }
     }
 
-    #[derive(Debug, Clone, PartialEq, bincode::Encode, bincode::Decode)]
+    #[derive(Debug, Clone, PartialEq, bincode::Encode, bincode::Decode, Default)]
     pub struct HeartBeat {
         pub(crate) from: PeerIdentifier,
         pub(crate) term: u64,


Test patch:
diff --git a/duva/src/domains/cluster_actors/actor/tests/mod.rs b/duva/src/domains/cluster_actors/actor/tests/mod.rs
index eee9bf8d..a82dca82 100644
--- a/duva/src/domains/cluster_actors/actor/tests/mod.rs
+++ b/duva/src/domains/cluster_actors/actor/tests/mod.rs
@@ -31,9 +31,7 @@ use std::time::Duration;
 use tempfile::TempDir;
 use tokio::fs::OpenOptions;
 use tokio::net::TcpListener;
-
 use tokio::sync::mpsc::channel;
-use tokio::time::timeout;
 use uuid::Uuid;
 
 use std::sync::Arc;
@@ -233,21 +231,16 @@ pub(crate) fn cache_entries_create_helper(keys_values: &[(&str, &str)]) -> Vec<C
 }
 
 // Helper function to assert migration batch ack
-pub(crate) async fn assert_migration_batch_ack(
+pub(crate) async fn assert_expected_queryio(
     message_buf: &FakeReadWrite,
-    expected_batch_id: &str,
-    expected_success: bool,
+    expected_query_io: QueryIO,
 ) {
     let sent_messages = message_buf.lock().await;
     assert_eq!(sent_messages.len(), 1);
 
     let message = sent_messages.front().unwrap();
-    let QueryIO::MigrationBatchAck(ack) = message else {
-        panic!("Expected MigrationBatchAck message");
-    };
 
-    assert_eq!(ack.batch_id.0, expected_batch_id);
-    assert_eq!(ack.success, expected_success);
+    assert_eq!(message, &expected_query_io);
 }
 
 #[cfg(test)]
diff --git a/duva/src/domains/cluster_actors/actor/tests/partitionings.rs b/duva/src/domains/cluster_actors/actor/tests/partitionings.rs
index 970ffc11..fbfe8bc6 100644
--- a/duva/src/domains/cluster_actors/actor/tests/partitionings.rs
+++ b/duva/src/domains/cluster_actors/actor/tests/partitionings.rs
@@ -397,13 +397,16 @@ async fn test_migrate_keys_retrieves_actual_data() {
 }
 
 #[tokio::test]
-async fn test_receive_batch_success_path() {
+async fn test_receive_batch_success_path_when_consensus_is_required() {
     // GIVEN
     let mut cluster_actor = cluster_actor_create_helper(ReplicationRole::Leader).await;
     let (_hwm, cache_manager) = cache_manager_create_helper();
-
+    let current_index = cluster_actor.logger.last_log_index;
     let peer_replid = ReplicationId::Key("repl_id_for_other_node".to_string());
-    let (message_buf, sender_peer_id) =
+
+    let (repl_buf, _) = cluster_actor.test_add_peer(6579, NodeKind::Replica, None);
+
+    let (_, sender_peer_id) =
         cluster_actor.test_add_peer(6567, NodeKind::NonData, Some(peer_replid.clone()));
     cluster_actor.hash_ring = cluster_actor
         .hash_ring
@@ -411,22 +414,56 @@ async fn test_receive_batch_success_path() {
         .unwrap();
 
     let cache_entries = cache_entries_create_helper(&[("success_key3", "value2")]);
-    let migrate_batch = migration_batch_create_helper("success_test", cache_entries);
+    let migrate_batch = migration_batch_create_helper("success_test", cache_entries.clone());
 
     // WHEN
     cluster_actor.receive_batch(migrate_batch, &cache_manager, sender_peer_id).await;
 
-    // THEN
-    tokio::time::sleep(Duration::from_millis(100)).await;
-    assert_migration_batch_ack(&message_buf, "success_test", true).await;
+    // THEN - verify that the log index is incremented
+    assert_eq!(cluster_actor.logger.last_log_index, current_index + 1);
+
+    assert_expected_queryio(
+        &repl_buf,
+        QueryIO::AppendEntriesRPC(HeartBeat {
+            from: cluster_actor.replication.self_identifier(),
+            replid: cluster_actor.replication.replid.clone(),
+            append_entries: vec![WriteOperation {
+                request: WriteRequest::MSet { entries: cache_entries.clone() },
+                log_index: 1,
+                term: 0,
+            }],
+            ..Default::default()
+        }),
+    )
+    .await;
+}
 
-    // Verify the keys were actually stored in the cache
+#[tokio::test]
+async fn test_receive_batch_success_path_when_noreplica_found() {
+    // GIVEN
+    let mut cluster_actor = cluster_actor_create_helper(ReplicationRole::Leader).await;
+    let (_hwm, cache_manager) = cache_manager_create_helper();
+    let current_index = cluster_actor.logger.last_log_index;
+    let peer_replid = ReplicationId::Key("repl_id_for_other_node".to_string());
+
+    let (_, sender_peer_id) =
+        cluster_actor.test_add_peer(6567, NodeKind::NonData, Some(peer_replid.clone()));
+    cluster_actor.hash_ring = cluster_actor
+        .hash_ring
+        .add_partition_if_not_exists(peer_replid.clone(), sender_peer_id.clone())
+        .unwrap();
 
-    let retrieved_value2 = cache_manager.route_get("success_key3").await.unwrap();
+    let cache_entries =
+        cache_entries_create_helper(&[("success_key3", "value2"), ("success_key4", "value4")]);
+    let migrate_batch = migration_batch_create_helper("success_test", cache_entries.clone());
 
-    assert!(retrieved_value2.is_some());
+    // WHEN
+    cluster_actor.receive_batch(migrate_batch, &cache_manager, sender_peer_id).await;
 
-    assert_eq!(retrieved_value2.unwrap().value(), "value2");
+    // THEN - verify that the log index is incremented
+    assert_eq!(cluster_actor.logger.last_log_index, current_index + 1);
+    let keys = cache_manager.route_keys(None).await;
+    assert_eq!(keys.len(), 2);
 }
 
 #[tokio::test]
@@ -453,23 +490,14 @@ async fn test_receive_batch_validation_failure_keys_not_belonging_to_node() {
     cluster_actor.receive_batch(migrate_batch, &cache_manager, sender_peer_id).await;
 
     // THEN
-    assert_migration_batch_ack(&message_buf, "validation_test", false).await;
-}
-
-#[tokio::test]
-async fn test_receive_batch_empty_cache_entries() {
-    // GIVEN
-    let mut cluster_actor = cluster_actor_create_helper(ReplicationRole::Leader).await;
-    let (_hwm, cache_manager) = cache_manager_create_helper();
-    let (message_buf, sender_peer_id) = cluster_actor.test_add_peer(6568, NodeKind::NonData, None);
-
-    let migrate_batch = migration_batch_create_helper("empty_test", vec![]);
-
-    // WHEN
-    cluster_actor.receive_batch(migrate_batch, &cache_manager, sender_peer_id).await;
-
-    // THEN
-    assert_migration_batch_ack(&message_buf, "empty_test", true).await;
+    assert_expected_queryio(
+        &message_buf,
+        QueryIO::MigrationBatchAck(MigrationBatchAck {
+            batch_id: BatchId("validation_test".into()),
+            success: false,
+        }),
+    )
+    .await;
 }
 
 #[tokio::test]