fix(rust-protocol): switch tunnel frame buffers from Vec<u8> to Bytes to reduce copying and memory overhead

This commit is contained in:
2026-03-17 23:29:02 +00:00
parent 4cfc518301
commit 93578d7034
8 changed files with 63 additions and 47 deletions

View File

@@ -10,6 +10,7 @@ use tokio_rustls::TlsAcceptor;
use tokio_util::sync::CancellationToken;
use serde::{Deserialize, Serialize};
use bytes::Bytes;
use remoteingress_protocol::*;
type HubTlsStream = tokio_rustls::server::TlsStream<TcpStream>;
@@ -26,7 +27,7 @@ struct HubStreamState {
/// Unbounded channel to deliver FRAME_DATA payloads to the upstream writer task.
/// Unbounded because flow control (WINDOW_UPDATE) already limits bytes-in-flight.
/// A bounded channel would kill streams instead of applying backpressure.
data_tx: mpsc::UnboundedSender<Vec<u8>>,
data_tx: mpsc::UnboundedSender<Bytes>,
/// Cancellation token for this stream.
cancel_token: CancellationToken,
/// Send window for FRAME_DATA_BACK (download direction).
@@ -307,8 +308,8 @@ async fn handle_hub_frame(
edge_stream_count: &Arc<AtomicU32>,
edge_id: &str,
event_tx: &mpsc::Sender<HubEvent>,
ctrl_tx: &mpsc::Sender<Vec<u8>>,
data_tx: &mpsc::Sender<Vec<u8>>,
ctrl_tx: &mpsc::Sender<Bytes>,
data_tx: &mpsc::Sender<Bytes>,
target_host: &str,
edge_token: &CancellationToken,
cleanup_tx: &mpsc::Sender<u32>,
@@ -346,7 +347,7 @@ async fn handle_hub_frame(
});
// Create channel for data from edge to this stream
let (stream_data_tx, mut stream_data_rx) = mpsc::unbounded_channel::<Vec<u8>>();
let (stream_data_tx, mut stream_data_rx) = mpsc::unbounded_channel::<Bytes>();
// Adaptive initial window: scale with current stream count
// to keep total in-flight data within the 32MB budget.
let initial_window = compute_window_for_stream_count(
@@ -496,7 +497,7 @@ async fn handle_hub_frame(
Ok(n) => {
send_window.fetch_sub(n as u32, Ordering::Release);
encode_frame_header(&mut buf, stream_id, FRAME_DATA_BACK, n);
let frame = buf[..FRAME_HEADER_SIZE + n].to_vec();
let frame = Bytes::copy_from_slice(&buf[..FRAME_HEADER_SIZE + n]);
let sent = tokio::select! {
result = data_writer_tx.send(frame) => result.is_ok(),
_ = stream_token.cancelled() => false,
@@ -707,8 +708,8 @@ async fn handle_edge_connection(
// QoS dual-channel: ctrl frames have priority over data frames.
// Stream handlers send through these channels -> TunnelIo drains them.
let (ctrl_tx, mut ctrl_rx) = mpsc::channel::<Vec<u8>>(256);
let (data_tx, mut data_rx) = mpsc::channel::<Vec<u8>>(4096);
let (ctrl_tx, mut ctrl_rx) = mpsc::channel::<Bytes>(256);
let (data_tx, mut data_rx) = mpsc::channel::<Bytes>(4096);
// Spawn task to forward config updates as FRAME_CONFIG frames
let config_writer_tx = ctrl_tx.clone();