Reuse a per-connection encode buffer for daemon replies

binaries/daemon/src/node_communication/tcp.rs

@@ -78,7 +78,10 @@ async fn handle_connection_loop(
     }
 
     Listener::run(
-        TcpConnection(connection),
+        TcpConnection {
+            stream: connection,
+            send_buf: Vec::new(),
+        },
         generation,
         daemon_tx,
         clock,
@@ -87,13 +90,25 @@ async fn handle_connection_loop(
     .await
 }
 
-struct TcpConnection(TcpStream);
+/// Upper bound on the send buffer retained between replies.
+///
+/// The buffer exists to avoid an allocation per reply, which only needs enough
+/// room for an ordinary message. Without a cap, one outsized reply would pin its
+/// full size (up to `MAX_MESSAGE_BYTES`, 64 MiB) for the life of the connection
+/// — and the daemon holds one connection per node.
+const MAX_RETAINED_SEND_BUF: usize = 256 * 1024;
+
+struct TcpConnection {
+    stream: TcpStream,
+    /// Reused across replies; see [`dora_message::encode_into`].
+    send_buf: Vec<u8>,
+}
 
 impl Connection for TcpConnection {
     async fn receive_message(&mut self) -> eyre::Result<Option<Timestamped<DaemonRequest>>> {
         // No header timeout: a node connection may legitimately stay idle
         // between requests, so only mid-frame (body) stalls are faults.
-        let raw = match socket_stream_receive_with_header_timeout(&mut self.0, None).await {
+        let raw = match socket_stream_receive_with_header_timeout(&mut self.stream, None).await {
             Ok(raw) => raw,
             Err(err) => match err.kind() {
                 ErrorKind::UnexpectedEof
@@ -115,11 +130,21 @@ impl Connection for TcpConnection {
             // don't send empty replies
             return Ok(());
         }
-        let serialized = dora_message::encode_presized(&message, message.encode_size_hint())
-            .wrap_err("failed to serialize DaemonReply")?;
-        socket_stream_send(&mut self.0, &serialized)
-            .await
-            .wrap_err("failed to send DaemonReply")?;
+        let mut buf = dora_message::encode_into(
+            &message,
+            message.encode_size_hint(),
+            std::mem::take(&mut self.send_buf),
+        )
+        .wrap_err("failed to serialize DaemonReply")?;
+
+        let sent = socket_stream_send(&mut self.stream, &buf).await;
+
+        // Take the buffer back on the failure path too, so a transient send
+        // error doesn't quietly drop the reuse for the rest of the connection.
+        buf.shrink_to(MAX_RETAINED_SEND_BUF);
+        self.send_buf = buf;
+
+        sent.wrap_err("failed to send DaemonReply")?;
         Ok(())
     }
 }

libraries/message/src/lib.rs

@@ -43,6 +43,25 @@ pub fn encode_presized<T: serde::Serialize>(
     )
 }
 
+/// [`encode_presized`], reusing `buf` rather than allocating a new one.
+///
+/// Returns the buffer so a caller holding a long-lived connection can keep it
+/// and hand it back on the next message, trading one allocation per message for
+/// one retained buffer per connection. `buf`'s existing contents are discarded.
+///
+/// Callers should bound what they retain — see `MAX_RETAINED_SEND_BUF` in the
+/// daemon's TCP connection — or one outsized message pins that much memory for
+/// the life of the connection.
+pub fn encode_into<T: serde::Serialize>(
+    value: &T,
+    bulk_bytes: usize,
+    mut buf: Vec<u8>,
+) -> postcard::Result<Vec<u8>> {
+    buf.clear();
+    buf.reserve(bulk_bytes.saturating_add(ENVELOPE_SIZE_HINT));
+    postcard::to_extend(value, buf)
+}
+
 /// Decode `bytes` in dora's binary wire format, requiring the value to consume
 /// the **entire** slice.
 ///
@@ -216,4 +235,29 @@ mod encoding_tests {
             "error should name the cause, got: {err:#}"
         );
     }
+
+    /// A reused buffer must not leak any of its previous contents into the next
+    /// message — the failure mode would be a silently corrupt frame on a
+    /// long-lived connection, not a crash.
+    #[test]
+    fn encode_into_ignores_the_buffers_previous_contents() {
+        let value = Metadata::new(uhlc::HLC::default().new_timestamp());
+        let expected = crate::encode(&value).expect("baseline");
+
+        let recycled = [
+            Vec::new(),
+            vec![0xAA; 1],
+            vec![0xAA; expected.len()],
+            // Longer than the new message, so a missing `clear()` would leave a
+            // tail behind rather than being overwritten.
+            vec![0xAA; expected.len() * 4],
+            Vec::with_capacity(64 * 1024),
+        ];
+
+        for (i, buf) in recycled.into_iter().enumerate() {
+            let out = crate::encode_into(&value, 0, buf).expect("encode_into");
+            assert_eq!(out, expected, "recycled buffer {i} changed the encoding");
+            crate::decode::<Metadata>(&out).expect("result must still decode");
+        }
+    }
 }