implement sending batches of messages

broker/src/lib.rs

@@ -245,8 +245,6 @@ where
             verify_broker!(connection, inner);
             authenticate_with_broker!(connection, inner)
         };
-
-
     }
 
     /// This function handles a user (public) connection. We take the following steps:
@@ -310,6 +308,7 @@ where
                     Ok(brokers) => {
                         // Calculate the difference, spawn tasks to connect to them
                         for broker in brokers.difference(&inner.brokers_connected.read()) {
+                            // TODO: make this into a separate function
                             // Extrapolate the address to connect to
                             let to_connect_address = broker.broker_advertise_address.clone();
 

client/src/retry.rs

@@ -244,11 +244,11 @@ where
     /// - If the message sending failed
     pub async fn send_message(&self, message: Message) -> Result<()> {
         // Serialize the message
-        let message = bail!(
+        let message = Arc::from(bail!(
             message.serialize(),
             Serialize,
             "failed to serialize message"
-        );
+        ));
 
         // Try to send the message, reconnecting if needed
         Ok(try_with_reconnect!(self, send_message_raw, message,))

proto/src/connection/protocols/mod.rs

@@ -1,6 +1,6 @@
 //! This module defines connections, listeners, and their implementations.
 
-use std::net::SocketAddr;
+use std::{net::SocketAddr, sync::Arc};
 
 use async_trait::async_trait;
 
@@ -36,7 +36,14 @@ pub trait Connection: Send + Sync + 'static {
     ///
     /// # Errors
     /// - If we fail to deliver the message. This usually means a connection problem.
-    async fn send_message_raw(&self, message: Vec<u8>) -> Result<()>;
+    async fn send_message_raw(&self, message: Arc<Vec<u8>>) -> Result<()>;
+
+    /// Send a vector of pre-formed messages over the connection.
+    ///
+    /// # Errors
+    /// - If we fail to deliver any of the messages. This usually means a connection problem.
+    async fn send_messages_raw(&self, messages: Vec<Arc<Vec<u8>>>) -> Result<()>;
+
 
     /// Connect to a remote address, returning an instance of `Self`.
     ///

proto/src/connection/protocols/quic.rs

@@ -13,7 +13,7 @@ use crate::{
     MAX_MESSAGE_SIZE,
 };
 use core::hash::Hash;
-use std::net::ToSocketAddrs;
+use std::{net::ToSocketAddrs, sync::Arc};
 
 use super::{Connection, Listener, Protocol};
 
@@ -114,14 +114,14 @@ impl Connection for QuicConnection {
         );
 
         // Send the message
-        self.send_message_raw(message_bytes).await
+        self.send_message_raw(Arc::from(message_bytes)).await
     }
 
     /// Send a pre-formed message over the connection.
     ///
     /// # Errors
     /// - If we fail to deliver the message. This usually means a connection problem.
-    async fn send_message_raw(&self, message: Vec<u8>) -> Result<()> {
+    async fn send_message_raw(&self, message: Arc<Vec<u8>>) -> Result<()> {
         // Open the outgoing unidirectional stream
         let mut stream = bail!(
             self.0.open_uni().await,
@@ -145,6 +145,26 @@ impl Connection for QuicConnection {
         ))
     }
 
+    /// Send a vector of pre-formed message over the connection.
+    ///
+    /// TODO: FIGURE OUT IF WE WANT TO FRAME LIKE THIS. it may be more performant with batching
+    /// to not do it this way.
+    ///
+    /// # Errors
+    /// - If we fail to deliver any of the messages. This usually means a connection problem.
+    async fn send_messages_raw(&self, messages: Vec<Arc<Vec<u8>>>) -> Result<()> {
+        // Send each message over the connection
+        for message in messages {
+            bail!(
+                self.send_message_raw(message).await,
+                Connection,
+                "failed to send message"
+            );
+        }
+
+        Ok(())
+    }
+
     /// Connect to a remote endpoint, returning an instance of `Self`. With QUIC,
     /// this requires creating an endpoint, binding to it, and then attempting
     /// a connection.

proto/src/connection/protocols/tcp.rs

@@ -101,14 +101,14 @@ impl Connection for TcpConnection {
         );
 
         // Send the serialized message
-        self.send_message_raw(serialized_message).await
+        self.send_message_raw(Arc::from(serialized_message)).await
     }
 
     /// Send a pre-formed message over the connection.
     ///
     /// # Errors
     /// - If we fail to deliver the message. This usually means a connection problem.
-    async fn send_message_raw(&self, message: Vec<u8>) -> Result<()> {
+    async fn send_message_raw(&self, message: Arc<Vec<u8>>) -> Result<()> {
         // Lock the stream so we don't send message/message sizes interleaved
         let mut sender_guard = self.sender.lock().await;
 
@@ -130,6 +130,36 @@ impl Connection for TcpConnection {
         Ok(())
     }
 
+    /// Send a vector pre-formed messages over the connection.
+    ///
+    /// # Errors
+    /// - If we fail to deliver the message. This usually means a connection problem.
+    async fn send_messages_raw(&self, messages: Vec<Arc<Vec<u8>>>) -> Result<()> {
+        // Lock the stream so we don't send message/message sizes interleaved
+        let mut sender_guard = self.sender.lock().await;
+
+        // For each message:
+        for message in messages {
+            // Write the message size to the stream
+            bail!(
+                sender_guard.write_u64(message.len() as u64).await,
+                Connection,
+                "failed to send message size"
+            );
+
+            // Write the message to the stream
+            bail!(
+                sender_guard.write_all(&message).await,
+                Connection,
+                "failed to send message"
+            );
+        }
+
+        drop(sender_guard);
+
+        Ok(())
+    }
+
     /// Connect to a remote endpoint, returning an instance of `Self`.
     /// With TCP, this requires just connecting to the remote endpoint.
     ///