diff --git a/binaries/daemon/src/node_communication/tcp.rs b/binaries/daemon/src/node_communication/tcp.rs index e324305d7..057bcfc71 100644 --- a/binaries/daemon/src/node_communication/tcp.rs +++ b/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, +} impl Connection for TcpConnection { async fn receive_message(&mut self) -> eyre::Result>> { // 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(()) } } diff --git a/libraries/message/src/lib.rs b/libraries/message/src/lib.rs index 34fe6415c..bb838f097 100644 --- a/libraries/message/src/lib.rs +++ b/libraries/message/src/lib.rs @@ -43,6 +43,25 @@ pub fn encode_presized( ) } +/// [`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( + value: &T, + bulk_bytes: usize, + mut buf: Vec, +) -> postcard::Result> { + 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::(&out).expect("result must still decode"); + } + } }