Repository: Migorithm/duva
Pull request #255 (2025-02-16 15:26:28)

feat: raft happy path
- ReqConsensus : consensus_con
- VoteConsensus : to handle follower votes

- `replicas` methods


Patch:
diff --git a/src/client_utils/mod.rs b/src/client_utils/mod.rs
index e47d4ec3..9da27690 100644
--- a/src/client_utils/mod.rs
+++ b/src/client_utils/mod.rs
@@ -28,6 +28,9 @@ impl ClientStreamHandler {
     /// # Panics
     /// Panics if the connection to `bind_addr` fails.
     pub async fn new(bind_addr: String) -> Self {
+        while tokio::net::TcpStream::connect(&bind_addr).await.is_err() {
+            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+        }
         let stream = tokio::net::TcpStream::connect(bind_addr).await.unwrap();
 
         let (read_half, write_half) = stream.into_split();
diff --git a/src/presentation/client_in/manager.rs b/src/presentation/client_in/manager.rs
index c93fce94..c06099da 100644
--- a/src/presentation/client_in/manager.rs
+++ b/src/presentation/client_in/manager.rs
@@ -148,16 +148,9 @@ impl ClientManager {
                     continue;
                 };
 
+                // apply state change
                 let res = match self.handle(request).await {
-                    Ok(response) => {
-                        // ! SAFETY: at this point, majority votes was made already. as long as at least
-                        // ! one node has the commit log, it will be committed.
-                        tokio::spawn(Self::commit_log(
-                            optional_log_offset,
-                            self.cluster_communication_manager.clone(),
-                        ));
-                        response
-                    }
+                    Ok(response) => response,
                     Err(e) => QueryIO::Err(e.to_string().into()),
                 };
 
@@ -177,17 +170,8 @@ impl ClientManager {
         };
         let (tx, rx) = tokio::sync::oneshot::channel();
         self.cluster_communication_manager
-            .send(ClusterCommand::Consensus { log, sender: tx })
+            .send(ClusterCommand::ReqConsensus { log, sender: tx })
             .await?;
-        Ok(Some(rx.await?))
-    }
-
-    async fn commit_log(
-        optional_log_offset: Option<u64>,
-        cluster_manager: ClusterCommunicationManager,
-    ) {
-        if let Some(commit_log) = optional_log_offset {
-            let _ = cluster_manager.send(ClusterCommand::CommitLog(commit_log)).await;
-        }
+        Ok(rx.await?)
     }
 }
diff --git a/src/presentation/cluster_in/communication_manager.rs b/src/presentation/cluster_in/communication_manager.rs
index 906e821e..1fc9e3fe 100644
--- a/src/presentation/cluster_in/communication_manager.rs
+++ b/src/presentation/cluster_in/communication_manager.rs
@@ -10,8 +10,6 @@ use crate::{
 use std::time::Duration;
 use tokio::{sync::mpsc::Sender, time::interval};
 
-use super::connection_manager::ClusterConnectionManager;
-
 #[derive(Clone)]
 pub struct ClusterCommunicationManager(pub(crate) Sender<ClusterCommand>);
 
@@ -85,8 +83,4 @@ impl ClusterCommunicationManager {
         let Some(_) = rx.await? else { return Ok(false) };
         Ok(true)
     }
-
-    pub(crate) fn to_connection_manager(self) -> ClusterConnectionManager {
-        ClusterConnectionManager(self.clone())
-    }
 }
diff --git a/src/presentation/cluster_in/peer_listeners/peer_listener.rs b/src/presentation/cluster_in/peer_listeners/peer_listener.rs
index 054eaadf..617dda1c 100644
--- a/src/presentation/cluster_in/peer_listeners/peer_listener.rs
+++ b/src/presentation/cluster_in/peer_listeners/peer_listener.rs
@@ -77,6 +77,9 @@ impl PeerListener {
                     RequestFromSlave::HeartBeat(state) => {
                         self.receive_heartbeat(state).await;
                     }
+                    RequestFromSlave::Acks(items) => {
+                        let _ = self.cluster_handler.send(ClusterCommand::ReceiveAcks(items)).await;
+                    }
                 }
             }
         }
@@ -91,7 +94,8 @@ impl PeerListener {
             for cmd in cmds {
                 match cmd {
                     RequestFromMaster::HeartBeat(mut state) => {
-                        self.log_entries(&mut state);
+                        self.log_entries(&mut state).await;
+
                         self.receive_heartbeat(state).await;
                     }
                     RequestFromMaster::FullSync(data) => {
@@ -109,6 +113,15 @@ impl PeerListener {
         }
     }
 
+    async fn log_entries(&self, state: &mut HeartBeatMessage) {
+        let append_entries = state.append_entries.drain(..).collect::<Vec<_>>();
+        if append_entries.is_empty() {
+            return;
+        }
+
+        let _ = self.cluster_handler.send(ClusterCommand::ReceiveLogEntries(append_entries)).await;
+    }
+
     async fn read_command<T>(&mut self) -> anyhow::Result<Vec<T>>
     where
         T: std::convert::TryFrom<QueryIO>,
@@ -123,13 +136,6 @@ impl PeerListener {
             .collect::<Result<_, _>>()
             .map_err(Into::into)
     }
-
-    fn log_entries(&self, state: &mut HeartBeatMessage) {
-        let append_entries = state.append_entries.drain(..).collect::<Vec<_>>();
-        for entry in append_entries {
-            println!("[INFO] Append entry {:?}", entry);
-        }
-    }
 }
 
 pub(super) type ReactorKillSwitch = tokio::sync::oneshot::Receiver<()>;
diff --git a/src/presentation/cluster_in/peer_listeners/requests.rs b/src/presentation/cluster_in/peer_listeners/requests.rs
index 9b850454..d545bdc5 100644
--- a/src/presentation/cluster_in/peer_listeners/requests.rs
+++ b/src/presentation/cluster_in/peer_listeners/requests.rs
@@ -10,6 +10,7 @@ pub enum RequestFromMaster {
 
 pub enum RequestFromSlave {
     HeartBeat(HeartBeatMessage),
+    Acks(Vec<u64>),
 }
 
 impl TryFrom<QueryIO> for RequestFromMaster {
@@ -29,6 +30,7 @@ impl TryFrom<QueryIO> for RequestFromSlave {
     fn try_from(query: QueryIO) -> anyhow::Result<Self> {
         match query {
             QueryIO::HeartBeat(peer_state) => Ok(RequestFromSlave::HeartBeat(peer_state)),
+            QueryIO::Acks(acks) => Ok(RequestFromSlave::Acks(acks)),
             _ => todo!(),
         }
     }
diff --git a/src/services/cluster/actors/actor.rs b/src/services/cluster/actors/actor.rs
index c283cf97..d0963a72 100644
--- a/src/services/cluster/actors/actor.rs
+++ b/src/services/cluster/actors/actor.rs
@@ -11,9 +11,11 @@ use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use std::collections::BTreeMap;
 use tokio::sync::mpsc::Receiver;
+
 use tokio::time::Instant;
 
 use super::commands::{AddPeer, ClusterCommand};
+use super::consensus::ConsensusTracker;
 
 const FANOUT: usize = 2;
 
@@ -33,6 +35,8 @@ impl ClusterActor {
         mut cluster_message_listener: Receiver<ClusterCommand>,
         notifier: tokio::sync::watch::Sender<bool>,
     ) {
+        let mut consensus_con = ConsensusTracker::default();
+
         while let Some(command) = cluster_message_listener.recv().await {
             // TODO notifier will be used when election process is implemented
             let _ = notifier.clone();
@@ -46,7 +50,6 @@ impl ClusterActor {
                     let _ = callback.send(self.members.keys().cloned().collect::<Vec<_>>().into());
                 }
 
-                ClusterCommand::Replicate { query: _ } => todo!(),
                 ClusterCommand::SendHeartBeat => {
                     let hop_count = self.hop_count(FANOUT, self.members.len());
                     self.send_liveness_heartbeat(hop_count).await;
@@ -77,25 +80,27 @@ impl ClusterActor {
                         let _ = sender.send(None);
                     }
                 }
-                ClusterCommand::Consensus { log, sender } => {
-                    self.consensus(log).await;
-
-                    // ! We won't wait for the result to come.
-                    // ! Replica will reject consensus if the log's index is higher than it's own watermark +1
-                    // ! and the leader will retry the consensus
-                    // ! This means that we need an consensus response actor to handle the response from the replicas
-                    // ! The actor will take (sender + log index) and
-
-                    // self.response_actor.send(ResponseActorCommand{
-                    //     log_index: log.index,
-                    //     // TODO if any operations failed, it's okay to drop sender
-                    //     sender
-                    // }).await;
-
-                    let _ = sender.send(self.replication.master_repl_offset);
+                ClusterCommand::ReqConsensus { log, sender } => {
+                    // ! If there are no replicas, don't send the request
+                    if self.replicas().count() == 0 {
+                        let _ = sender.send(None);
+                        continue;
+                    }
+                    self.req_consensus(log).await;
+
+                    consensus_con.add(
+                        self.replication.master_repl_offset,
+                        sender,
+                        self.replicas().count(),
+                    );
                 }
-                ClusterCommand::CommitLog(commit_log) => {
-                    // TODO send commit to all replicas
+                ClusterCommand::ReceiveAcks(offsets) => {
+                    self.apply_acks(&mut consensus_con, offsets);
+                }
+                ClusterCommand::ReceiveLogEntries(write_operations) => {
+                    // TODO handle the log entries
+                    println!("[INFO] Received log entries: {:?}", write_operations);
+                    self.receive_log_entries_from_master(write_operations).await;
                 }
             }
         }
@@ -208,14 +213,13 @@ impl ClusterActor {
         }
     }
 
-    async fn consensus(&mut self, req: WriteRequest) {
+    async fn req_consensus(&mut self, req: WriteRequest) {
         // TODO when are we going to increase offset?
-        let write_op = WriteOperation { op: req, offset: self.replication.master_repl_offset };
 
-        let heartbeat = self.replication.append_entry(0, write_op);
+        let heartbeat = self.replication.append_entry(0, req);
 
         let mut tasks = self
-            .replicas()
+            .replicas_mut()
             .into_iter()
             .map(|peer| peer.write_io(heartbeat.clone()))
             .collect::<FuturesUnordered<_>>();
@@ -224,13 +228,44 @@ impl ClusterActor {
         while let Some(_) = tasks.next().await {}
     }
 
-    fn replicas(&mut self) -> Vec<&mut Peer> {
+    fn master_mut(&mut self) -> Option<&mut Peer> {
+        self.members.values_mut().find(|peer| matches!(peer.w_conn.kind, PeerKind::Master))
+    }
+
+    fn replicas(&self) -> impl Iterator<Item = &Peer> {
+        self.members
+            .values()
+            .into_iter()
+            .filter(|peer| matches!(peer.w_conn.kind, PeerKind::Replica))
+    }
+
+    fn replicas_mut(&mut self) -> Vec<&mut Peer> {
         self.members
             .values_mut()
             .into_iter()
             .filter(|peer| matches!(peer.w_conn.kind, PeerKind::Replica))
             .collect::<Vec<_>>()
     }
+
+    fn apply_acks(&self, consensus_con: &mut ConsensusTracker, offsets: Vec<u64>) {
+        offsets.into_iter().for_each(|offset| {
+            if let Some(mut consensus) = consensus_con.take(&offset) {
+                println!("Received acks for offset: {}", offset);
+                consensus.apply_vote();
+
+                if let Some(consensus) = consensus.maybe_not_finished(offset) {
+                    consensus_con.insert(offset, consensus);
+                }
+            }
+        });
+    }
+
+    async fn receive_log_entries_from_master(&mut self, write_operations: Vec<WriteOperation>) {
+        let offsets = write_operations.iter().map(|op| op.offset).collect::<Vec<_>>();
+        if let Some(master) = self.master_mut() {
+            let _ = master.write_io(QueryIO::Acks(offsets)).await;
+        }
+    }
 }
 
 #[test]
diff --git a/src/services/cluster/actors/commands.rs b/src/services/cluster/actors/commands.rs
index 4b3677cc..b3ac0271 100644
--- a/src/services/cluster/actors/commands.rs
+++ b/src/services/cluster/actors/commands.rs
@@ -1,10 +1,9 @@
-use crate::services::aof::WriteRequest;
+use crate::services::aof::{WriteOperation, WriteRequest};
 use crate::services::cluster::peers::address::PeerAddrs;
 use crate::services::cluster::peers::identifier::PeerIdentifier;
 
 use crate::services::cluster::peers::peer::Peer;
 use crate::services::cluster::replications::replication::{HeartBeatMessage, ReplicationInfo};
-use crate::services::query_io::QueryIO;
 
 pub enum ClusterCommand {
     AddPeer(AddPeer),
@@ -12,11 +11,14 @@ pub enum ClusterCommand {
     ReplicationInfo(tokio::sync::oneshot::Sender<ReplicationInfo>),
     SetReplicationInfo { master_repl_id: String, offset: u64 },
     SendHeartBeat,
-    Replicate { query: QueryIO },
-    ReceiveHeartBeat(HeartBeatMessage),
+
     ForgetPeer(PeerIdentifier, tokio::sync::oneshot::Sender<Option<()>>),
-    Consensus { log: WriteRequest, sender: tokio::sync::oneshot::Sender<u64> },
-    CommitLog(u64),
+    ReqConsensus { log: WriteRequest, sender: tokio::sync::oneshot::Sender<Option<u64>> },
+    ReceiveAcks(Vec<u64>),
+
+    // peer listener commands
+    ReceiveHeartBeat(HeartBeatMessage),
+    ReceiveLogEntries(Vec<WriteOperation>),
 }
 
 pub struct AddPeer {
diff --git a/src/services/cluster/actors/consensus.rs b/src/services/cluster/actors/consensus.rs
new file mode 100644
index 00000000..a78f8020
--- /dev/null
+++ b/src/services/cluster/actors/consensus.rs
@@ -0,0 +1,44 @@
+use std::collections::HashMap;
+
+use tokio::sync::oneshot::Sender;
+
+use crate::make_smart_pointer;
+
+pub struct ConsensusVoting {
+    sender: Sender<Option<u64>>,
+    vote_count: u8,
+    required_votes: u8,
+}
+impl ConsensusVoting {
+    pub fn apply_vote(&mut self) {
+        self.vote_count += 1;
+    }
+
+    pub fn maybe_not_finished(self, offset: u64) -> Option<Self> {
+        if self.vote_count >= self.required_votes {
+            let _ = self.sender.send(Some(offset));
+            None
+        } else {
+            Some(self)
+        }
+    }
+}
+
+#[derive(Default)]
+pub struct ConsensusTracker(HashMap<u64, ConsensusVoting>);
+impl ConsensusTracker {
+    pub fn add(&mut self, key: u64, value: Sender<Option<u64>>, replica_count: usize) {
+        self.0.insert(
+            key,
+            ConsensusVoting {
+                sender: value,
+                vote_count: 0,
+                required_votes: (replica_count as f64 / 2.0).ceil() as u8,
+            },
+        );
+    }
+    pub fn take(&mut self, offset: &u64) -> Option<ConsensusVoting> {
+        self.0.remove(offset)
+    }
+}
+make_smart_pointer!(ConsensusTracker, HashMap<u64, ConsensusVoting>);
diff --git a/src/services/cluster/actors/mod.rs b/src/services/cluster/actors/mod.rs
index 4a9546a5..afba956c 100644
--- a/src/services/cluster/actors/mod.rs
+++ b/src/services/cluster/actors/mod.rs
@@ -1,2 +1,3 @@
 pub(crate) mod actor;
 pub mod commands;
+mod consensus;
diff --git a/src/services/cluster/replications/replication.rs b/src/services/cluster/replications/replication.rs
index 08bb391b..876a7964 100644
--- a/src/services/cluster/replications/replication.rs
+++ b/src/services/cluster/replications/replication.rs
@@ -1,4 +1,4 @@
-use crate::services::aof::WriteOperation;
+use crate::services::aof::{WriteOperation, WriteRequest};
 use crate::services::cluster::peers::identifier::PeerIdentifier;
 use crate::services::cluster::replications::replid_generator::generate_replid;
 use bytes::Bytes;
@@ -80,12 +80,14 @@ impl ReplicationInfo {
         }
     }
 
-    pub fn append_entry(&mut self, hop_count: u8, entry: WriteOperation) -> HeartBeatMessage {
+    pub fn append_entry(&mut self, hop_count: u8, req: WriteRequest) -> HeartBeatMessage {
         self.master_repl_offset += 1;
+        let entry = WriteOperation { op: req, offset: self.master_repl_offset };
         let mut heartbeat = self.default_heartbeat(hop_count);
         heartbeat.append_entries.push(entry);
         heartbeat
     }
+
     pub fn default_heartbeat(&self, hop_count: u8) -> HeartBeatMessage {
         HeartBeatMessage {
             heartbeat_from: self.self_identifier.clone(),
diff --git a/src/services/query_io.rs b/src/services/query_io.rs
index ac098c63..658cab76 100644
--- a/src/services/query_io.rs
+++ b/src/services/query_io.rs
@@ -13,6 +13,7 @@ const ARRAY_PREFIX: char = '*';
 const ERROR_PREFIX: char = '-';
 const PEERSTATE_PREFIX: char = '^';
 const REPLICATE_PREFIX: char = '#';
+const ACKS_PREFIX: char = '@';
 
 #[macro_export]
 macro_rules! write_array {
@@ -32,6 +33,7 @@ pub enum QueryIO {
     File(Bytes),
     HeartBeat(HeartBeatMessage),
     ReplicateLog(WriteOperation),
+    Acks(Vec<u64>),
 }
 
 impl QueryIO {
@@ -115,6 +117,15 @@ impl QueryIO {
                 .into();
                 [message, op.to_array().serialize()].concat().into()
             }
+            QueryIO::Acks(items) => [
+                format!("{}{}\r\n", ACKS_PREFIX, items.len()).into_bytes(),
+                items
+                    .into_iter()
+                    .flat_map(|item| QueryIO::BulkString(item.to_string().into()).serialize())
+                    .collect(),
+            ]
+            .concat()
+            .into(),
         }
     }
 
@@ -196,6 +207,7 @@ pub fn deserialize(buffer: BytesMut) -> Result<(QueryIO, usize)> {
         }
         PEERSTATE_PREFIX => parse_heartbeat(buffer),
         REPLICATE_PREFIX => parse_replicate(buffer),
+        ACKS_PREFIX => parse_acks(buffer),
 
         _ => Err(anyhow::anyhow!("Not a known value type {:?}", buffer)),
     }
@@ -289,6 +301,11 @@ fn parse_replicate(buffer: BytesMut) -> std::result::Result<(QueryIO, usize), an
     ))
 }
 
+fn parse_acks(buffer: BytesMut) -> std::result::Result<(QueryIO, usize), anyhow::Error> {
+    let (acks, usize) = parse_array(buffer)?;
+    Ok((QueryIO::Acks(acks.unpack_array()?), usize))
+}
+
 fn parse_bulk_string(buffer: BytesMut) -> Result<(Bytes, usize)> {
     let (line, mut len) =
         read_until_crlf(&buffer[1..].into()).ok_or(anyhow::anyhow!("Invalid bulk string"))?;
@@ -642,3 +659,29 @@ fn test_from_binary_to_replicate_log() {
         })
     );
 }
+
+#[test]
+fn test_from_binary_to_acks() {
+    // GIVEN
+    let data = "@2\r\n$1\r\n1\r\n$1\r\n2\r\n";
+    let buffer = BytesMut::from(data);
+
+    // WHEN
+    let (value, _) = deserialize(buffer).unwrap();
+
+    // THEN
+    assert_eq!(value, QueryIO::Acks(vec![1, 2]));
+}
+
+#[test]
+fn test_from_acks_to_binary() {
+    // GIVEN
+    let acks = vec![1, 2];
+    let replicate = QueryIO::Acks(acks);
+
+    // WHEN
+    let serialized = replicate.clone().serialize();
+
+    // THEN
+    assert_eq!("@2\r\n$1\r\n1\r\n$1\r\n2\r\n", serialized);
+}


Test patch:
diff --git a/tests/test_config_get_dir.rs b/tests/test_config_get_dir.rs
index 59815289..64a5efb9 100644
--- a/tests/test_config_get_dir.rs
+++ b/tests/test_config_get_dir.rs
@@ -1,15 +1,19 @@
 /// if the value of dir is /tmp, then the expected response to CONFIG GET dir is:
 /// *2\r\n$3\r\ndir\r\n$4\r\n/tmp\r\n
 mod common;
+use std::time::Duration;
+
 use crate::common::array;
 use common::spawn_server_process;
 use duva::client_utils::ClientStreamHandler;
+use tokio::time::sleep;
 
 #[tokio::test]
 async fn test_config_get_dir() {
     // GIVEN
     let process = spawn_server_process();
 
+    sleep(Duration::from_millis(500)).await;
     let mut h = ClientStreamHandler::new(process.bind_addr()).await;
 
     // WHEN