Merge pull request #787 from Migorithm/refactor/remove-batchid

refactor: batch

Author: Migorithm

duva/src/domains/cluster_actors/actor.rs

@@ -15,10 +15,10 @@ use crate::domains::QueryIO;
 use crate::domains::TAsyncReadWrite;
 use crate::domains::caches::cache_manager::CacheManager;
 use crate::domains::cluster_actors::consensus::election::ElectionVoting;
-use crate::domains::cluster_actors::hash_ring::BatchId;
 use crate::domains::cluster_actors::hash_ring::MigrationBatch;
 use crate::domains::cluster_actors::hash_ring::PendingMigration;
 use crate::domains::cluster_actors::hash_ring::PendingMigrationBatch;
+
 use crate::domains::cluster_actors::topology::{NodeReplInfo, Topology};
 use crate::domains::operation_logs::WriteRequest;
 use crate::domains::operation_logs::interfaces::TWriteAheadLog;
@@ -346,22 +346,14 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
             return;
         };
 
-        match self.hash_ring.list_replids_for_keys(&req.request.all_keys()) {
-            | Ok(replids) if replids.keys().all(|replid| *replid == self.replication.replid) => {
+        match self.hash_ring.key_ownership(&req.request.all_keys()) {
+            | Ok(replids) if replids.all_belongs_to(&self.replication.replid) => {
                 self.req_consensus(req).await;
             },
             | Ok(replids) => {
                 // To notify client's of what keys have been moved.
                 // ! Still, client won't know where the key has been moved. The assumption here is client SHOULD have correct hashring information.
-                let moved_keys = replids
-                    .iter()
-                    .filter_map(|(replid, keys)| {
-                        if *replid != self.replication.replid { Some(keys.iter()) } else { None }
-                    })
-                    .flatten()
-                    .copied()
-                    .collect::<Vec<_>>()
-                    .join(" ");
+                let moved_keys = replids.except(&self.replication.replid).join(" ");
                 req.callback.send(format!("MOVED {moved_keys}").into());
             },
             | Err(err) => {
@@ -1278,11 +1270,11 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
             return;
         };
 
-        self.pending_migrations
-            .as_mut()
-            .map(|p| p.add_batch(target.id.clone(), PendingMigrationBatch::new(callback, keys)));
+        self.pending_migrations.as_mut().map(|p| {
+            p.add_batch(target.batch_id.clone(), PendingMigrationBatch::new(callback, keys))
+        });
 
-        let _ = target_peer.send(MigrateBatch { batch_id: target.id, cache_entries }).await;
+        let _ = target_peer.send(MigrateBatch { batch_id: target.batch_id, cache_entries }).await;
     }
 
     pub(crate) async fn receive_batch(
@@ -1327,7 +1319,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
                 }
                 error!(
                     "Failed to write some keys during migration for batch {}",
-                    migrate_batch.batch_id.0
+                    migrate_batch.batch_id
                 );
             }
         });
@@ -1410,7 +1402,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
         }
     }
 
-    pub(crate) async fn send_batch_ack(&mut self, batch_id: BatchId, to: PeerIdentifier) {
+    pub(crate) async fn send_batch_ack(&mut self, batch_id: String, to: PeerIdentifier) {
         let Some(peer) = self.members.get_mut(&to) else {
             return;
         };

duva/src/domains/cluster_actors/actor/tests/partitionings.rs

@@ -1,6 +1,6 @@
 use crate::domains::QueryIO;
 use crate::domains::caches::cache_objects::{CacheValue, TypedValue};
-use crate::domains::cluster_actors::hash_ring::BatchId;
+
 use crate::domains::cluster_actors::hash_ring::{HashRing, tests::migration_task_create_helper};
 
 use std::time::Duration;
@@ -376,7 +376,7 @@ async fn test_migrate_batch_send_migrate_batch_peer_message() {
     // THEN
     assert_expected_queryio(
         &buf,
-        QueryIO::MigrateBatch(MigrateBatch { batch_id: batch.id, cache_entries: vec![] }),
+        QueryIO::MigrateBatch(MigrateBatch { batch_id: batch.batch_id, cache_entries: vec![] }),
     )
     .await;
 }
@@ -390,7 +390,7 @@ async fn test_receive_batch_when_empty_cache_entries() {
     let (buf, _id) = cluster_actor.test_add_peer(6909, Some(replid.clone()), true);
 
     // WHEN
-    let batch = MigrateBatch { batch_id: BatchId("empty_test".into()), cache_entries: vec![] };
+    let batch = MigrateBatch { batch_id: "empty_test".into(), cache_entries: vec![] };
     cluster_actor.receive_batch(batch.clone(), &cache_manager, _id).await;
 
     // THEN - verify that no log index is incremented
@@ -416,10 +416,8 @@ async fn test_receive_batch_when_consensus_is_required() {
 
     let cache_entries = vec![CacheEntry::new("success_key3", "value2")];
 
-    let batch = MigrateBatch {
-        batch_id: BatchId("success_test".into()),
-        cache_entries: cache_entries.clone(),
-    };
+    let batch =
+        MigrateBatch { batch_id: "success_test".into(), cache_entries: cache_entries.clone() };
 
     // WHEN
     cluster_actor.receive_batch(batch, &cache_manager, ack_to.clone()).await;
@@ -456,10 +454,8 @@ async fn test_receive_batch_when_noreplica_found() {
     let cache_entries =
         vec![CacheEntry::new("success_key3", "value2"), CacheEntry::new("success_key4", "value4")];
 
-    let batch = MigrateBatch {
-        batch_id: BatchId("success_test".into()),
-        cache_entries: cache_entries.clone(),
-    };
+    let batch =
+        MigrateBatch { batch_id: "success_test".into(), cache_entries: cache_entries.clone() };
 
     // WHEN
     let task = tokio::spawn(recv.wait_message(SchedulerMessage::SendBatchAck {
@@ -495,7 +491,7 @@ async fn test_unblock_write_reqs_if_done_when_migrations_still_pending() {
 
     // Add pending migration (simulating migration still in progress)
     let (callback, _migration_rx) = Callback::create();
-    let batch_id = BatchId("test_batch".into());
+    let batch_id = "test_batch".into();
     cluster_actor
         .pending_migrations
         .as_mut()
@@ -566,7 +562,7 @@ async fn test_handle_migration_ack_failure() {
     let mut cluster_actor = setup_blocked_cluster_actor_with_requests(1).await;
     let (_hwm, _cache_manager) = Helper::cache_manager();
     let (callback, callback_rx) = Callback::create();
-    let batch_id = BatchId("failure_batch".into());
+    let batch_id = "failure_batch".to_string();
 
     // Ensure pending_migrations is set up
     assert!(cluster_actor.pending_migrations.is_some());
@@ -606,9 +602,9 @@ async fn test_handle_migration_ack_batch_id_not_found() {
         .pending_migrations
         .as_mut()
         .unwrap()
-        .add_batch(BatchId("existing_batch".into()), PendingMigrationBatch::new(callback, vec![]));
+        .add_batch("existing_batch".into(), PendingMigrationBatch::new(callback, vec![]));
 
-    let non_existent_batch_id = BatchId("non_existent_batch".into());
+    let non_existent_batch_id = "non_existent_batch".into();
     let ack = MigrationBatchAck { batch_id: non_existent_batch_id, success: true };
 
     // WHEN
@@ -648,7 +644,7 @@ async fn test_handle_migration_ack_success_case_with_pending_reqs_and_migration(
 
     // Add the last pending migration with the test keys
     let (callback, callback_rx) = Callback::create();
-    let batch_id = BatchId("last_batch".into());
+    let batch_id = "last_batch".to_string();
     cluster_actor
         .pending_migrations
         .as_mut()

duva/src/domains/cluster_actors/command.rs

@@ -1,5 +1,5 @@
 use crate::ReplicationState;
-use crate::domains::cluster_actors::hash_ring::{BatchId, MigrationBatch};
+use crate::domains::cluster_actors::hash_ring::MigrationBatch;
 use crate::domains::cluster_actors::replication::{ReplicationId, ReplicationRole};
 use crate::domains::cluster_actors::topology::Topology;
 use crate::domains::operation_logs::WriteRequest;
@@ -29,7 +29,7 @@ pub enum SchedulerMessage {
     RebalanceRequest { request_to: PeerIdentifier, lazy_option: LazyOption },
     ScheduleMigrationBatch(MigrationBatch, Callback<anyhow::Result<()>>),
     TryUnblockWriteReqs,
-    SendBatchAck { batch_id: BatchId, to: PeerIdentifier },
+    SendBatchAck { batch_id: String, to: PeerIdentifier },
 }
 
 impl From<SchedulerMessage> for ClusterCommand {

duva/src/domains/cluster_actors/hash_ring.rs

@@ -91,20 +91,6 @@ impl HashRing {
             .map(|(_, node_id)| node_id.as_ref())
     }
 
-    fn find_node(&self, hash: u64) -> Option<&PeerIdentifier> {
-        // Find the first vnode with hash >= target hash
-        self.pnodes.get(self.find_replid(hash)?)
-    }
-
-    /// Verifies that all given keys belong to the specified node according to the hash ring
-    pub(crate) fn verify_key_belongs_to_node(
-        &self,
-        keys: &[&str],
-        expected_node: &PeerIdentifier,
-    ) -> bool {
-        keys.iter().all(|key| self.find_node(fnv_1a_hash(key)) == Some(expected_node))
-    }
-
     pub(crate) fn create_migration_tasks(
         &self,
         new_ring: &HashRing,
@@ -120,23 +106,23 @@ impl HashRing {
 
         // Check each partition for ownership changes
         for (i, &token) in tokens.iter().enumerate() {
-            let prev_token = if i == 0 { tokens[tokens.len() - 1] } else { tokens[i - 1] };
-            let (start, end) = (prev_token.wrapping_add(1), token);
-
-            if let (Some(old_owner), Some(new_owner)) =
-                (self.find_replid(token), new_ring.find_replid(token))
-                && (old_owner != new_owner)
+            if let Some(old_owner) = self.find_replid(token)
+                && let Some(new_owner) = new_ring.find_replid(token)
             {
-                // If both old and new owners exist, we need to check if ownership changed
+                if old_owner == new_owner {
+                    continue;
+                }
 
                 // Node ownership changed for this partition
                 // Need to migrate data from old node to new node
-                let affected_keys = filter_keys_in_partition(&keys, start, end);
-                if !affected_keys.is_empty() {
-                    migration_tasks.entry(new_owner.clone()).or_default().push(MigrationTask {
-                        task_id: (start, end),
-                        keys_to_migrate: affected_keys,
-                    });
+                let prev_token = if i == 0 { tokens[tokens.len() - 1] } else { tokens[i - 1] };
+                let (start, end) = (prev_token.wrapping_add(1), token);
+                let keys_to_migrate = filter_keys_in_partition(&keys, start, end);
+                if !keys_to_migrate.is_empty() {
+                    migration_tasks
+                        .entry(new_owner.clone())
+                        .or_default()
+                        .push(MigrationTask { range: (start, end), keys_to_migrate });
                 }
             }
         }
@@ -158,10 +144,7 @@ impl HashRing {
         self.vnodes.len()
     }
 
-    pub(crate) fn list_replids_for_keys<'a>(
-        &self,
-        keys: &[&'a str],
-    ) -> anyhow::Result<HashMap<ReplicationId, Vec<&'a str>>> {
+    pub(crate) fn key_ownership<'a>(&self, keys: &[&'a str]) -> anyhow::Result<KeyOwnership<'a>> {
         let mut map: HashMap<ReplicationId, Vec<&str>> = HashMap::new();
 
         for key in keys {
@@ -174,7 +157,7 @@ impl HashRing {
             v.push(key);
         }
 
-        Ok(map)
+        Ok(KeyOwnership(map))
     }
 
     #[cfg(test)]
@@ -197,6 +180,21 @@ impl HashRing {
     }
 }
 
+pub(crate) struct KeyOwnership<'a>(HashMap<ReplicationId, Vec<&'a str>>);
+impl<'a> KeyOwnership<'a> {
+    pub(crate) fn all_belongs_to(&self, target: &ReplicationId) -> bool {
+        self.0.keys().all(|replid| replid == target)
+    }
+    pub(crate) fn except(self, target: &ReplicationId) -> Vec<&'a str> {
+        self.0
+            .iter()
+            .filter_map(|(replid, keys)| if replid != target { Some(keys.iter()) } else { None })
+            .flatten()
+            .copied()
+            .collect::<Vec<_>>()
+    }
+}
+
 fn filter_keys_in_partition(
     keys: &[String],
     partition_start: u64,

duva/src/domains/cluster_actors/hash_ring/migration_task.rs

@@ -4,7 +4,7 @@ use crate::{ReplicationId, domains::cluster_actors::ConsensusRequest, types::Cal
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) struct MigrationTask {
-    pub(crate) task_id: (u64, u64),          // (start_hash, end_hash)
+    pub(crate) range: (u64, u64),            // (start_hash, end_hash)
     pub(crate) keys_to_migrate: Vec<String>, // actual keys in this range
 }
 
@@ -14,19 +14,16 @@ impl MigrationTask {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash, bincode::Encode, bincode::Decode)]
-pub(crate) struct BatchId(pub(crate) String);
-
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) struct MigrationBatch {
-    pub(crate) id: BatchId,
+    pub(crate) batch_id: String,
     pub(crate) target_repl: ReplicationId,
     pub(crate) tasks: Vec<MigrationTask>,
 }
 
 impl MigrationBatch {
     pub(crate) fn new(target_repl: ReplicationId, tasks: Vec<MigrationTask>) -> Self {
-        Self { id: BatchId(uuid::Uuid::now_v7().to_string()), target_repl, tasks }
+        Self { batch_id: uuid::Uuid::now_v7().to_string(), target_repl, tasks }
     }
 }
 
@@ -48,16 +45,16 @@ impl PendingMigrationBatch {
 #[derive(Debug, Default)]
 pub(crate) struct PendingMigration {
     requests: VecDeque<ConsensusRequest>,
-    batches: HashMap<BatchId, PendingMigrationBatch>,
+    batches: HashMap<String, PendingMigrationBatch>,
 }
 impl PendingMigration {
     pub(crate) fn add_req(&mut self, req: ConsensusRequest) {
         self.requests.push_back(req);
     }
-    pub(crate) fn add_batch(&mut self, id: BatchId, batch: PendingMigrationBatch) {
+    pub(crate) fn add_batch(&mut self, id: String, batch: PendingMigrationBatch) {
         self.batches.insert(id, batch);
     }
-    pub(crate) fn pop_batch(&mut self, id: &BatchId) -> Option<PendingMigrationBatch> {
+    pub(crate) fn pop_batch(&mut self, id: &String) -> Option<PendingMigrationBatch> {
         self.batches.remove(id)
     }
     pub(crate) fn pending_requests(self) -> VecDeque<ConsensusRequest> {

duva/src/domains/cluster_actors/hash_ring/tests/mod.rs

@@ -9,7 +9,7 @@ mod migration;
 
 pub(crate) fn migration_task_create_helper(start_hash: u64, end_hash: u64) -> MigrationTask {
     MigrationTask {
-        task_id: (start_hash, end_hash),
+        range: (start_hash, end_hash),
         keys_to_migrate: (start_hash..end_hash).map(|i| format!("key_{i}")).collect(),
     }
 }

duva/src/domains/peers/command.rs

@@ -53,10 +53,7 @@ mod peer_messages {
     use super::*;
     use crate::domains::{
         caches::cache_objects::CacheEntry,
-        cluster_actors::{
-            hash_ring::{BatchId, HashRing},
-            replication::ReplicationId,
-        },
+        cluster_actors::{hash_ring::HashRing, replication::ReplicationId},
         operation_logs::{WriteOperation, logger::ReplicatedLogs},
         peers::peer::PeerState,
     };
@@ -171,22 +168,22 @@ mod peer_messages {
 
     #[derive(Debug, Clone, PartialEq, Eq, bincode::Encode, bincode::Decode)]
     pub struct MigrateBatch {
-        pub(crate) batch_id: BatchId,
+        pub(crate) batch_id: String,
         pub(crate) cache_entries: Vec<CacheEntry>,
     }
 
     #[derive(Debug, Clone, PartialEq, Eq, bincode::Encode, bincode::Decode)]
     pub struct MigrationBatchAck {
-        pub(crate) batch_id: BatchId,
+        pub(crate) batch_id: String,
         pub(crate) success: bool,
     }
 
     impl MigrationBatchAck {
-        pub(crate) fn with_reject(batch_id: BatchId) -> Self {
+        pub(crate) fn with_reject(batch_id: String) -> Self {
             Self { batch_id, success: false }
         }
 
-        pub(crate) fn with_success(batch_id: BatchId) -> Self {
+        pub(crate) fn with_success(batch_id: String) -> Self {
             Self { success: true, batch_id }
         }
     }