Merge pull request #892 from Migorithm/refactor/seperate-out-peer-message

feat: seperate out peer message from query io

Author: Migorithm

duva-client/src/broker/mod.rs

@@ -21,7 +21,7 @@ use duva::prelude::{Topology, anyhow};
 use duva::presentation::clients::request::{ClientAction, NonMutatingAction, ServerResponse};
 use futures::future::try_join_all;
 
-use duva::prelude::anyhow::anyhow;
+use duva::prelude::anyhow::bail;
 use node_connections::NodeConnection;
 use read_stream::ServerStreamReader;
 use tracing::error;
@@ -37,8 +37,10 @@ pub struct Broker {
 
 impl Broker {
     pub(crate) async fn new(server_addr: &PeerIdentifier) -> anyhow::Result<Self> {
-        let (r, w, auth_response) = Broker::authenticate(&server_addr.clone(), None).await?;
+        let stream =
+            TcpStream::connect(server_addr.as_str()).await.map_err(|_| IoError::NotConnected)?;
 
+        let (r, w, auth_response) = Broker::authenticate(stream, Default::default()).await?;
         let (broker_tx, rx) = tokio::sync::mpsc::channel::<BrokerMessage>(2000);
 
         let seed_replid = auth_response.replication_id;
@@ -109,21 +111,17 @@ impl Broker {
     }
 
     pub(crate) async fn authenticate(
-        server_addr: &PeerIdentifier,
-        auth_request: Option<ConnectionRequest>,
-    ) -> Result<(ServerStreamReader, ServerStreamWriter, ConnectionResponse), IoError> {
-        let mut stream =
-            TcpStream::connect(server_addr.as_str()).await.map_err(|_| IoError::NotConnected)?;
-        let request = auth_request.unwrap_or_default();
-        let connection_req = ConnectionRequests::Authenticate(request);
-        stream.serialized_write(connection_req).await?; // client_id not exist
-        let connection_res: ConnectionResponses = stream.deserialized_read().await?;
-        let auth_response = match connection_res {
-            ConnectionResponses::Authenticated(response) => response,
-            _ => return Err(IoError::Custom("Authentication failed".into())),
+        mut stream: TcpStream,
+        conn_req: ConnectionRequest,
+    ) -> anyhow::Result<(ServerStreamReader, ServerStreamWriter, ConnectionResponse)> {
+        stream.serialized_write(ConnectionRequests::Authenticate(conn_req)).await?; // client_id not exist
+
+        let ConnectionResponses::Authenticated(response) = stream.deserialized_read().await? else {
+            bail!("Authentication failed");
         };
+
         let (r, w) = stream.into_split();
-        Ok((ServerStreamReader(r), ServerStreamWriter(w), auth_response))
+        Ok((ServerStreamReader(r), ServerStreamWriter(w), response))
     }
 
     // pull-based leader discovery
@@ -161,12 +159,12 @@ impl Broker {
     }
 
     async fn discover_leader_from(&mut self, follower: PeerIdentifier) -> anyhow::Result<()> {
-        let mut stream =
-            TcpStream::connect(follower.as_str()).await.map_err(|_| IoError::NotConnected)?;
+        let mut stream = TcpStream::connect(follower.as_str()).await?;
         stream.serialized_write(ConnectionRequests::Discovery).await?;
         let ConnectionResponses::Discovery { leader_id } = stream.deserialized_read().await? else {
-            return Err(anyhow!("Discovery failed"));
+            bail!("Discovery failed!");
         };
+
         self.add_node_connection(leader_id.clone()).await
     }
 
@@ -190,14 +188,14 @@ impl Broker {
         let auth_req =
             ConnectionRequest { client_id: Some(self.client_id.to_string()), request_id: 0 };
 
-        let Ok((server_stream_reader, server_stream_writer, auth_response)) =
-            Self::authenticate(&peer_id, Some(auth_req)).await
-        else {
-            return Err(anyhow::anyhow!("Authentication failed!"));
-        };
+        let stream =
+            TcpStream::connect(peer_id.as_str()).await.map_err(|_| IoError::NotConnected)?;
+
+        let (server_stream_reader, server_stream_writer, auth_response) =
+            Self::authenticate(stream, auth_req).await?;
 
         if !auth_response.is_leader_node {
-            return Err(anyhow::anyhow!("Only Leader connection is allowed!"));
+            bail!("Only Leader connection is allowed!");
         }
 
         self.node_connections.insert(

duva/src/adapters/io/tokio_stream.rs

@@ -1,7 +1,11 @@
 use crate::domains::interface::{TRead, TWrite};
 use crate::domains::peers::connections::connection_types::{ReadConnected, WriteConnected};
 use crate::domains::query_io::SERDE_CONFIG;
-use crate::domains::{IoError, TAsyncReadWrite, TSerdeRead, TSerdeWrite};
+use crate::domains::replications::messages::PeerMessage;
+use crate::domains::{
+    IoError, TAsyncReadWrite, TReadBytes, TSerdeDynamicRead, TSerdeDynamicWrite, TSerdeRead,
+    TSerdeWrite,
+};
 use crate::domains::{QueryIO, deserialize};
 use bytes::BytesMut;
 use std::fmt::Debug;
@@ -20,7 +24,7 @@ impl<T: AsyncWriteExt + std::marker::Unpin + Sync + Send + Debug + 'static> TWri
 }
 
 #[async_trait::async_trait]
-impl<T: AsyncReadExt + std::marker::Unpin + Sync + Send + Debug + 'static> TRead for T {
+impl<T: AsyncReadExt + std::marker::Unpin + Sync + Send + Debug + 'static> TReadBytes for T {
     // TCP doesn't inherently delimit messages.
     // The data arrives in a continuous stream of bytes. And
     // we might not receive all the data in one go.
@@ -52,7 +56,10 @@ impl<T: AsyncReadExt + std::marker::Unpin + Sync + Send + Debug + 'static> TRead
         }
         Ok(())
     }
+}
 
+#[async_trait::async_trait]
+impl<T: AsyncReadExt + std::marker::Unpin + Sync + Send + Debug + 'static> TRead for T {
     async fn read_values(&mut self) -> Result<Vec<QueryIO>, IoError> {
         let mut buffer = BytesMut::with_capacity(INITIAL_CAPACITY);
         self.read_bytes(&mut buffer).await?;
@@ -75,6 +82,22 @@ impl<T: AsyncReadExt + std::marker::Unpin + Sync + Send + Debug + 'static> TRead
     }
 }
 
+#[async_trait::async_trait]
+impl<T: AsyncReadExt + std::marker::Unpin + Sync + Send + Debug + 'static> TSerdeDynamicRead for T {
+    async fn receive_peer_msgs(&mut self) -> Result<Vec<PeerMessage>, IoError> {
+        let mut buffer = BytesMut::with_capacity(INITIAL_CAPACITY);
+        self.read_bytes(&mut buffer).await?;
+        let mut parsed_values = Vec::new();
+        while !buffer.is_empty() {
+            let (request, size) = bincode::decode_from_slice(&buffer, SERDE_CONFIG)
+                .map_err(|e| IoError::Custom(e.to_string()))?;
+            parsed_values.push(request);
+            buffer = buffer.split_off(size);
+        }
+        Ok(parsed_values)
+    }
+}
+
 impl<T: AsyncWriteExt + std::marker::Unpin + Sync + Send + Debug + 'static> TSerdeWrite for T {
     async fn serialized_write(&mut self, buf: impl bincode::Encode + Send) -> Result<(), IoError> {
         let encoded = bincode::encode_to_vec(buf, SERDE_CONFIG)
@@ -83,6 +106,17 @@ impl<T: AsyncWriteExt + std::marker::Unpin + Sync + Send + Debug + 'static> TSer
     }
 }
 
+#[async_trait::async_trait]
+impl<T: AsyncWriteExt + std::marker::Unpin + Sync + Send + Debug + 'static> TSerdeDynamicWrite
+    for T
+{
+    async fn send(&mut self, msg: PeerMessage) -> Result<(), IoError> {
+        let encoded = bincode::encode_to_vec(msg, SERDE_CONFIG)
+            .map_err(|e| IoError::Custom(e.to_string()))?;
+        self.write_all(&encoded).await.map_err(|e| io_error_from_kind(e.kind()))
+    }
+}
+
 impl<T: AsyncReadExt + std::marker::Unpin + Sync + Send + Debug + 'static> TSerdeRead for T {
     async fn deserialized_read<U>(&mut self) -> Result<U, IoError>
     where

duva/src/adapters/op_logs/disk_based.rs

@@ -1,5 +1,5 @@
+use crate::domains::query_io::SERDE_CONFIG;
 use crate::domains::query_io::serialized_len_with_bincode;
-use crate::domains::query_io::{SERDE_CONFIG, WRITE_OP_PREFIX};
 use crate::domains::replications::LogEntry;
 use crate::domains::replications::WriteOperation;
 use crate::domains::replications::interfaces::TWriteAheadLog;
@@ -13,6 +13,7 @@ use std::io::{ErrorKind, Read};
 use std::path::{Path, PathBuf};
 
 const SEGMENT_SIZE: usize = 1024 * 1024; // 1MB per segment
+const WRITE_OP_PREFIX: char = '#';
 
 /// A local write-ahead-log (WAL) file (op_logs) implementation using segmented logs.
 pub struct FileOpLogs {
@@ -587,7 +588,7 @@ mod tests {
         file.read_to_end(&mut buf).unwrap();
 
         let (encoded, _): (WriteOperation, usize) =
-            bincode::decode_from_slice(&buf[1..], bincode::config::standard()).unwrap();
+            bincode::decode_from_slice(&buf[1..], SERDE_CONFIG).unwrap();
 
         assert_eq!(encoded.entry, request);
     }

duva/src/domains/caches/cache_objects/entry.rs

@@ -70,6 +70,8 @@ impl CacheEntry {
 
 #[cfg(test)]
 mod tests {
+    use crate::domains::query_io::SERDE_CONFIG;
+
     use super::*;
     use bincode::{decode_from_slice, encode_to_vec};
     use chrono::{DateTime, Utc};
@@ -82,10 +84,10 @@ mod tests {
         let original_value = CacheValue::new("test_value").with_expiry(expiry_time);
 
         // WHEN
-        let encoded = encode_to_vec(&original_value, bincode::config::standard()).unwrap();
+        let encoded = encode_to_vec(&original_value, SERDE_CONFIG).unwrap();
 
         let (decoded_value, _): (CacheValue, usize) =
-            decode_from_slice(&encoded, bincode::config::standard()).unwrap();
+            decode_from_slice(&encoded, SERDE_CONFIG).unwrap();
 
         // THEN - Verify the decoded value matches the original
         assert_eq!(decoded_value.value, original_value.value);
@@ -99,11 +101,11 @@ mod tests {
         let original_value = CacheValue::new("test_value_no_expiry");
 
         // Encode the value
-        let encoded = encode_to_vec(&original_value, bincode::config::standard()).unwrap();
+        let encoded = encode_to_vec(&original_value, SERDE_CONFIG).unwrap();
 
         // Decode the value back
         let (decoded_value, _): (CacheValue, usize) =
-            decode_from_slice(&encoded, bincode::config::standard()).unwrap();
+            decode_from_slice(&encoded, SERDE_CONFIG).unwrap();
 
         // Verify the decoded value matches the original
         assert_eq!(decoded_value.value, original_value.value);
@@ -120,11 +122,11 @@ mod tests {
             .with_expiry(DateTime::from_timestamp_millis(expiry_millis).unwrap());
 
         // Encode the entry
-        let encoded = encode_to_vec(&original_entry, bincode::config::standard()).unwrap();
+        let encoded = encode_to_vec(&original_entry, SERDE_CONFIG).unwrap();
 
         // Decode the entry back
         let (decoded_entry, _): (CacheEntry, usize) =
-            decode_from_slice(&encoded, bincode::config::standard()).unwrap();
+            decode_from_slice(&encoded, SERDE_CONFIG).unwrap();
 
         // Verify the decoded entry matches the original
         assert_eq!(decoded_entry.key(), original_entry.key());

duva/src/domains/caches/cache_objects/value.rs

@@ -74,15 +74,16 @@ impl THasExpiry for CacheValue {
 
 #[cfg(test)]
 mod tests {
+    use crate::domains::query_io::SERDE_CONFIG;
+
     use super::*;
     use bincode::{decode_from_slice, encode_to_vec};
 
     #[test]
     fn test_string_variant_encode_decode() {
         let original = CacheValue::new("hello");
-        let encoded = encode_to_vec(&original, bincode::config::standard()).unwrap();
-        let (decoded, _): (CacheValue, usize) =
-            decode_from_slice(&encoded, bincode::config::standard()).unwrap();
+        let encoded = encode_to_vec(&original, SERDE_CONFIG).unwrap();
+        let (decoded, _): (CacheValue, usize) = decode_from_slice(&encoded, SERDE_CONFIG).unwrap();
         assert_eq!(decoded, original);
         assert_eq!(decoded.value, TypedValue::String(Bytes::copy_from_slice(b"hello").into()));
     }
@@ -95,9 +96,8 @@ mod tests {
             q_list.rpush(i);
         }
         let original = CacheValue { value: TypedValue::List(q_list.clone()), expiry: None };
-        let encoded = encode_to_vec(&original, bincode::config::standard()).unwrap();
-        let (decoded, _): (CacheValue, usize) =
-            decode_from_slice(&encoded, bincode::config::standard()).unwrap();
+        let encoded = encode_to_vec(&original, SERDE_CONFIG).unwrap();
+        let (decoded, _): (CacheValue, usize) = decode_from_slice(&encoded, SERDE_CONFIG).unwrap();
         assert_eq!(decoded, original);
         assert_eq!(decoded.value, TypedValue::List(q_list));
     }

duva/src/domains/cluster_actors/actor.rs

@@ -18,6 +18,7 @@ use crate::domains::peers::command::BannedPeer;
 use crate::domains::peers::command::BatchEntries;
 use crate::domains::peers::command::BatchId;
 use crate::domains::replications::messages::ElectionVote;
+use crate::domains::replications::messages::PeerMessage;
 use crate::domains::replications::messages::RejectionReason;
 use crate::domains::replications::messages::ReplicationAck;
 use crate::domains::replications::messages::RequestVote;
@@ -654,7 +655,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
             }
 
             let peer = self.members.get_mut(&request_to).unwrap();
-            let _ = peer.send(QueryIO::StartRebalance).await;
+            let _ = peer.send(PeerMessage::StartRebalance).await;
         }
     }
 
@@ -861,7 +862,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
 
     async fn send_rpc_to_replicas(&mut self) {
         self.iter_follower_append_entries()
-            .map(|(peer, hb)| peer.send(QueryIO::AppendEntriesRPC(hb)))
+            .map(|(peer, hb)| peer.send(PeerMessage::AppendEntriesRPC(hb)))
             .collect::<FuturesUnordered<_>>()
             .for_each(|_| async {})
             .await;
@@ -1338,7 +1339,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
                 warn!("No Member Found");
                 return;
             };
-            let _ = peer.send(QueryIO::MigrationBatchAck(migrate_batch.batch_id)).await;
+            let _ = peer.send(PeerMessage::MigrationBatchAck(migrate_batch.batch_id)).await;
             return;
         }
 
@@ -1446,7 +1447,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
         let Some(peer) = self.members.get_mut(&to) else {
             return;
         };
-        let _ = peer.send(QueryIO::MigrationBatchAck(batch_id)).await;
+        let _ = peer.send(PeerMessage::MigrationBatchAck(batch_id)).await;
     }
 
     async fn update_cluster_members(
@@ -1468,7 +1469,7 @@ impl<T: TWriteAheadLog> ClusterActor<T> {
     pub(crate) async fn process_graceful_shutdown(&mut self) {
         self.members
             .iter_mut()
-            .map(|(_, p)| p.send(QueryIO::CloseConnection))
+            .map(|(_, p)| p.send(PeerMessage::CloseConnection))
             .collect::<FuturesUnordered<_>>()
             .for_each(|_| async {})
             .await;

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

@@ -194,7 +194,7 @@ async fn test_receive_election_vote_candidate_wins_election() {
         ..Default::default()
     };
 
-    assert_expected_queryio(&replica1_fake_buf, QueryIO::AppendEntriesRPC(hb.clone())).await;
+    assert_expected_queryio(&replica1_fake_buf, PeerMessage::AppendEntriesRPC(hb.clone())).await;
 
     // hb that sends to cluster
     hb.append_entries = vec![];
@@ -203,7 +203,7 @@ async fn test_receive_election_vote_candidate_wins_election() {
 
     assert_expected_queryio(
         &replica1_fake_buf,
-        QueryIO::ClusterHeartBeat(hb.set_hashring(candidate_actor.hash_ring.clone())),
+        PeerMessage::ClusterHeartBeat(hb.set_hashring(candidate_actor.hash_ring.clone())),
     )
     .await;
 }