fix(rust-protocol): switch tunnel frame buffers from Vec<u8> to Bytes to reduce copying and memory overhead
This commit is contained in:
@@ -11,6 +11,7 @@ use tokio_rustls::TlsConnector;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use bytes::Bytes;
|
||||
use remoteingress_protocol::*;
|
||||
|
||||
type EdgeTlsStream = tokio_rustls::client::TlsStream<TcpStream>;
|
||||
@@ -26,7 +27,7 @@ enum EdgeFrameAction {
|
||||
struct EdgeStreamState {
|
||||
/// Unbounded channel to deliver FRAME_DATA_BACK payloads to the hub_to_client task.
|
||||
/// Unbounded because flow control (WINDOW_UPDATE) already limits bytes-in-flight.
|
||||
back_tx: mpsc::UnboundedSender<Vec<u8>>,
|
||||
back_tx: mpsc::UnboundedSender<Bytes>,
|
||||
/// Send window for FRAME_DATA (upload direction).
|
||||
/// Decremented by the client reader, incremented by FRAME_WINDOW_UPDATE_BACK from hub.
|
||||
send_window: Arc<AtomicU32>,
|
||||
@@ -290,8 +291,8 @@ async fn handle_edge_frame(
|
||||
client_writers: &Arc<Mutex<HashMap<u32, EdgeStreamState>>>,
|
||||
listen_ports: &Arc<RwLock<Vec<u16>>>,
|
||||
event_tx: &mpsc::Sender<EdgeEvent>,
|
||||
tunnel_writer_tx: &mpsc::Sender<Vec<u8>>,
|
||||
tunnel_data_tx: &mpsc::Sender<Vec<u8>>,
|
||||
tunnel_writer_tx: &mpsc::Sender<Bytes>,
|
||||
tunnel_data_tx: &mpsc::Sender<Bytes>,
|
||||
port_listeners: &mut HashMap<u16, JoinHandle<()>>,
|
||||
active_streams: &Arc<AtomicU32>,
|
||||
next_stream_id: &Arc<AtomicU32>,
|
||||
@@ -496,8 +497,8 @@ async fn connect_to_hub_and_run(
|
||||
|
||||
// QoS dual-channel: ctrl frames have priority over data frames.
|
||||
// Stream handlers send through these channels → TunnelIo drains them.
|
||||
let (tunnel_ctrl_tx, mut tunnel_ctrl_rx) = mpsc::channel::<Vec<u8>>(256);
|
||||
let (tunnel_data_tx, mut tunnel_data_rx) = mpsc::channel::<Vec<u8>>(4096);
|
||||
let (tunnel_ctrl_tx, mut tunnel_ctrl_rx) = mpsc::channel::<Bytes>(256);
|
||||
let (tunnel_data_tx, mut tunnel_data_rx) = mpsc::channel::<Bytes>(4096);
|
||||
let tunnel_writer_tx = tunnel_ctrl_tx.clone();
|
||||
|
||||
// Start TCP listeners for initial ports
|
||||
@@ -612,8 +613,8 @@ async fn connect_to_hub_and_run(
|
||||
fn apply_port_config(
|
||||
new_ports: &[u16],
|
||||
port_listeners: &mut HashMap<u16, JoinHandle<()>>,
|
||||
tunnel_ctrl_tx: &mpsc::Sender<Vec<u8>>,
|
||||
tunnel_data_tx: &mpsc::Sender<Vec<u8>>,
|
||||
tunnel_ctrl_tx: &mpsc::Sender<Bytes>,
|
||||
tunnel_data_tx: &mpsc::Sender<Bytes>,
|
||||
client_writers: &Arc<Mutex<HashMap<u32, EdgeStreamState>>>,
|
||||
active_streams: &Arc<AtomicU32>,
|
||||
next_stream_id: &Arc<AtomicU32>,
|
||||
@@ -727,8 +728,8 @@ async fn handle_client_connection(
|
||||
stream_id: u32,
|
||||
dest_port: u16,
|
||||
edge_id: &str,
|
||||
tunnel_ctrl_tx: mpsc::Sender<Vec<u8>>,
|
||||
tunnel_data_tx: mpsc::Sender<Vec<u8>>,
|
||||
tunnel_ctrl_tx: mpsc::Sender<Bytes>,
|
||||
tunnel_data_tx: mpsc::Sender<Bytes>,
|
||||
client_writers: Arc<Mutex<HashMap<u32, EdgeStreamState>>>,
|
||||
client_token: CancellationToken,
|
||||
active_streams: Arc<AtomicU32>,
|
||||
@@ -753,7 +754,7 @@ async fn handle_client_connection(
|
||||
// Per-stream unbounded back-channel. Flow control (WINDOW_UPDATE) limits
|
||||
// bytes-in-flight, so this won't grow unbounded. Unbounded avoids killing
|
||||
// streams due to channel overflow — backpressure slows streams, never kills them.
|
||||
let (back_tx, mut back_rx) = mpsc::unbounded_channel::<Vec<u8>>();
|
||||
let (back_tx, mut back_rx) = mpsc::unbounded_channel::<Bytes>();
|
||||
// Adaptive initial window: scale with current stream count to keep total in-flight
|
||||
// data within the 32MB budget. Prevents burst flooding when many streams open.
|
||||
let initial_window = remoteingress_protocol::compute_window_for_stream_count(
|
||||
@@ -871,7 +872,7 @@ async fn handle_client_connection(
|
||||
Ok(n) => {
|
||||
send_window.fetch_sub(n as u32, Ordering::Release);
|
||||
encode_frame_header(&mut buf, stream_id, FRAME_DATA, n);
|
||||
let data_frame = buf[..FRAME_HEADER_SIZE + n].to_vec();
|
||||
let data_frame = Bytes::copy_from_slice(&buf[..FRAME_HEADER_SIZE + n]);
|
||||
let sent = tokio::select! {
|
||||
result = tunnel_data_tx.send(data_frame) => result.is_ok(),
|
||||
_ = client_token.cancelled() => false,
|
||||
|
||||
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user