@feltdb/core 0.6.13 → 0.7.0

package/dist/create/server-source/crates/feltdb/src/tcp_transport.rs

@@ -19,6 +19,21 @@
 /// Deduplication/idempotency/retry logic is the protocol's responsibility via DistributedTransactionExecutor.
 
 use crate::distributed_transactions::ReplicationMessage;
+use crate::replica_acknowledgements::ReplicaAcknowledgement;
+
+/// What actually goes on the wire.
+///
+/// Tagged, so a connection can carry envelopes and acknowledgements without
+/// either being able to masquerade as the other. The alternative -- an
+/// acknowledgement riding inside a `ReplicationMessage` with a placeholder
+/// envelope -- would put a value on the receive path that only a flag
+/// distinguishes from a real operation, and the receive path's whole job is
+/// deciding what to apply.
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub enum Frame {
+    Replication(Box<ReplicationMessage>),
+    Acknowledgement(Box<ReplicaAcknowledgement>),
+}
 use crate::replication_protocol::{ProtocolError, ProtocolMetrics, ProtocolTransport};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
@@ -37,6 +52,12 @@ pub struct TcpTransport {
     metrics: Arc<ProtocolMetrics>,
     is_connected: Arc<AtomicBool>,
     receive_buffer: Arc<Mutex<Vec<ReplicationMessage>>>,
+    /// Acknowledgements read off the socket, waiting to be drained.
+    ///
+    /// They arrive interleaved with envelopes on the same connection, so
+    /// `receive` parks them here rather than changing its return type and the
+    /// trait every caller implements against.
+    ack_buffer: Arc<Mutex<Vec<ReplicaAcknowledgement>>>,
 }
 
 impl TcpTransport {
@@ -49,6 +70,7 @@ impl TcpTransport {
             metrics: Arc::new(ProtocolMetrics::new()),
             is_connected: Arc::new(AtomicBool::new(false)),
             receive_buffer: Arc::new(Mutex::new(Vec::with_capacity(MESSAGE_BUFFER_SIZE))),
+            ack_buffer: Arc::new(Mutex::new(Vec::new())),
         }
     }
 
@@ -76,6 +98,44 @@ impl TcpTransport {
     }
 }
 
+impl TcpTransport {
+    /// Send an acknowledgement over the same connection envelopes use.
+    ///
+    /// Not part of `ProtocolTransport`: that trait is the replication protocol,
+    /// and an acknowledgement is a statement about what has been applied rather
+    /// than something to apply.
+    pub async fn send_acknowledgement(
+        &self,
+        ack: ReplicaAcknowledgement,
+    ) -> Result<(), ProtocolError> {
+        let mut stream_guard = self.stream.lock().await;
+        let stream = stream_guard
+            .as_mut()
+            .ok_or(ProtocolError::ConnectionError("not connected".to_string()))?;
+        let payload = serde_json::to_vec(&Frame::Acknowledgement(Box::new(ack)))
+            .map_err(|e| ProtocolError::SerializationError(e.to_string()))?;
+        stream
+            .write_all(&(payload.len() as u32).to_be_bytes())
+            .await
+            .map_err(|e| ProtocolError::TransportError(e.to_string()))?;
+        stream
+            .write_all(&payload)
+            .await
+            .map_err(|e| ProtocolError::TransportError(e.to_string()))?;
+        stream
+            .flush()
+            .await
+            .map_err(|e| ProtocolError::TransportError(e.to_string()))?;
+        self.metrics.record_send(4 + payload.len());
+        Ok(())
+    }
+
+    /// Take the acknowledgements read off the socket since the last drain.
+    pub async fn take_acknowledgements(&self) -> Vec<ReplicaAcknowledgement> {
+        std::mem::take(&mut *self.ack_buffer.lock().await)
+    }
+}
+
 #[async_trait::async_trait]
 impl ProtocolTransport for TcpTransport {
     async fn send(&self, message: ReplicationMessage) -> Result<(), ProtocolError> {
@@ -85,7 +145,7 @@ impl ProtocolTransport for TcpTransport {
             .ok_or(ProtocolError::ConnectionError("not connected".to_string()))?;
 
         // Serialize message using serde_json (bincode has issues with serde_json::Value)
-        let payload = serde_json::to_vec(&message)
+        let payload = serde_json::to_vec(&Frame::Replication(Box::new(message)))
             .map_err(|e| ProtocolError::SerializationError(e.to_string()))?;
 
         let payload_len = payload.len() as u32;
@@ -152,12 +212,21 @@ impl ProtocolTransport for TcpTransport {
                     .await
                     .map_err(|e| ProtocolError::TransportError(e.to_string()))?;
 
-                let message: ReplicationMessage = serde_json::from_slice(&payload)
+                let frame: Frame = serde_json::from_slice(&payload)
                     .map_err(|e| ProtocolError::SerializationError(e.to_string()))?;
 
                 self.metrics.record_receive(4 + payload_len);
 
-                Ok(vec![message])
+                match frame {
+                    Frame::Replication(message) => Ok(vec![*message]),
+                    Frame::Acknowledgement(ack) => {
+                        // Read off the socket and parked. The bytes crossed the
+                        // network exactly as an envelope's do, and the counters
+                        // above record them.
+                        self.ack_buffer.lock().await.push(*ack);
+                        Ok(Vec::new())
+                    }
+                }
             }
             Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                 // No data available right now