serve.zone/remoteingress
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(remoteingress-core): add per-stream flow control for edge and hub tunnel data transfer

Browse Source
This commit is contained in:
Jürgen Kunz
2026-03-15 17:33:59 +00:00
parent cf2d32bfe7
commit 761551596b
5 changed files with 198 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
rust/crates/remoteingress-core/src/edge.rs
View File

@@ -4,7 +4,7 @@ use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, Mutex, RwLock};
use tokio::sync::{mpsc, Mutex, Notify, RwLock};
use tokio::task::JoinHandle;
use tokio::time::{Instant, sleep_until};
use tokio_rustls::TlsConnector;
@@ -13,6 +13,17 @@ use serde::{Deserialize, Serialize};
use remoteingress_protocol::*;
/// Per-stream state tracked in the edge's client_writers map.
struct EdgeStreamState {
/// Channel to deliver FRAME_DATA_BACK payloads to the hub_to_client task.
back_tx: mpsc::Sender<Vec<u8>>,
/// Send window for FRAME_DATA (upload direction).
/// Decremented by the client reader, incremented by FRAME_WINDOW_UPDATE_BACK from hub.
send_window: Arc<AtomicU32>,
/// Notifier to wake the client reader when the window opens.
window_notify: Arc<Notify>,
}
/// Edge configuration (hub-host + credentials only; ports come from hub).
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
@@ -351,8 +362,8 @@ async fn connect_to_hub_and_run(
}
});
// Client socket map: stream_id -> sender for writing data back to client
let client_writers: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>> =
// Client socket map: stream_id -> per-stream state (back channel + flow control)
let client_writers: Arc<Mutex<HashMap<u32, EdgeStreamState>>> =
Arc::new(Mutex::new(HashMap::new()));
// A5: Channel-based tunnel writer replaces Arc<Mutex<WriteHalf>>
@@ -407,17 +418,31 @@ async fn connect_to_hub_and_run(
match frame.frame_type {
FRAME_DATA_BACK => {
// Non-blocking send to prevent head-of-line blocking in the main dispatch loop.
// If the per-stream channel is full, close the stream rather than silently
// dropping data (which would corrupt the TCP stream).
// Non-blocking dispatch to per-stream channel.
// With flow control, the sender should rarely exceed the channel capacity.
let mut writers = client_writers.lock().await;
if let Some(tx) = writers.get(&frame.stream_id) {
if tx.try_send(frame.payload).is_err() {
log::warn!("Stream {} back-channel full, closing stream to prevent data corruption", frame.stream_id);
if let Some(state) = writers.get(&frame.stream_id) {
if state.back_tx.try_send(frame.payload).is_err() {
log::warn!("Stream {} back-channel full, closing stream", frame.stream_id);
writers.remove(&frame.stream_id);
}
}
}
FRAME_WINDOW_UPDATE_BACK => {
// Hub consumed data — increase our send window for this stream (upload direction)
if let Some(increment) = decode_window_update(&frame.payload) {
if increment > 0 {
let writers = client_writers.lock().await;
if let Some(state) = writers.get(&frame.stream_id) {
let prev = state.send_window.fetch_add(increment, Ordering::Release);
if prev + increment > MAX_WINDOW_SIZE {
state.send_window.store(MAX_WINDOW_SIZE, Ordering::Release);
}
state.window_notify.notify_one();
}
}
}
}
FRAME_CLOSE_BACK => {
let mut writers = client_writers.lock().await;
writers.remove(&frame.stream_id);
@@ -495,7 +520,7 @@ fn apply_port_config(
new_ports: &[u16],
port_listeners: &mut HashMap<u16, JoinHandle<()>>,
tunnel_writer_tx: &mpsc::Sender<Vec<u8>>,
client_writers: &Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
client_writers: &Arc<Mutex<HashMap<u32, EdgeStreamState>>>,
active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>,
edge_id: &str,
@@ -583,7 +608,7 @@ async fn handle_client_connection(
dest_port: u16,
edge_id: &str,
tunnel_writer_tx: mpsc::Sender<Vec<u8>>,
client_writers: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
client_writers: Arc<Mutex<HashMap<u32, EdgeStreamState>>>,
client_token: CancellationToken,
) {
let client_ip = client_addr.ip().to_string();
@@ -599,26 +624,44 @@ async fn handle_client_connection(
return;
}
// Set up channel for data coming back from hub
let (back_tx, mut back_rx) = mpsc::channel::<Vec<u8>>(256);
// Set up channel for data coming back from hub (capacity 16 is sufficient with flow control)
let (back_tx, mut back_rx) = mpsc::channel::<Vec<u8>>(16);
let send_window = Arc::new(AtomicU32::new(INITIAL_STREAM_WINDOW));
let window_notify = Arc::new(Notify::new());
{
let mut writers = client_writers.lock().await;
writers.insert(stream_id, back_tx);
writers.insert(stream_id, EdgeStreamState {
back_tx,
send_window: Arc::clone(&send_window),
window_notify: Arc::clone(&window_notify),
});
}
let (mut client_read, mut client_write) = client_stream.into_split();
// Task: hub -> client
// Task: hub -> client (download direction)
// After writing to client TCP, send WINDOW_UPDATE to hub so it can send more
let hub_to_client_token = client_token.clone();
let wu_tx = tunnel_writer_tx.clone();
let hub_to_client = tokio::spawn(async move {
let mut consumed_since_update: u32 = 0;
loop {
tokio::select! {
data = back_rx.recv() => {
match data {
Some(data) => {
let len = data.len() as u32;
if client_write.write_all(&data).await.is_err() {
break;
}
// Track consumption for flow control
consumed_since_update += len;
if consumed_since_update >= WINDOW_UPDATE_THRESHOLD {
let increment = consumed_since_update;
consumed_since_update = 0;
let frame = encode_window_update(stream_id, FRAME_WINDOW_UPDATE, increment);
let _ = wu_tx.try_send(frame);
}
}
None => break,
}
@@ -626,21 +669,39 @@ async fn handle_client_connection(
_ = hub_to_client_token.cancelled() => break,
}
}
// Send final window update for any remaining consumed bytes
if consumed_since_update > 0 {
let frame = encode_window_update(stream_id, FRAME_WINDOW_UPDATE, consumed_since_update);
let _ = wu_tx.try_send(frame);
}
let _ = client_write.shutdown().await;
});
// Task: client -> hub (via writer channel)
// Task: client -> hub (upload direction) with per-stream flow control
let mut buf = vec![0u8; 32768];
loop {
// Wait for send window to have capacity
loop {
let w = send_window.load(Ordering::Acquire);
if w > 0 { break; }
tokio::select! {
_ = window_notify.notified() => continue,
_ = client_token.cancelled() => break,
}
}
if client_token.is_cancelled() { break; }
// Limit read size to available window
let w = send_window.load(Ordering::Acquire) as usize;
let max_read = w.min(buf.len());
tokio::select! {
read_result = client_read.read(&mut buf) => {
read_result = client_read.read(&mut buf[..max_read]) => {
match read_result {
Ok(0) => break,
Ok(n) => {
send_window.fetch_sub(n as u32, Ordering::Release);
let data_frame = encode_frame(stream_id, FRAME_DATA, &buf[..n]);
// Use send().await for backpressure — this is a per-stream task so
// blocking only stalls this stream, not others. Prevents data loss
// for large transfers (e.g. 352MB Docker layers).
if tunnel_writer_tx.send(data_frame).await.is_err() {
log::warn!("Stream {} tunnel writer closed, closing", stream_id);
break;