Bound forwarded stream buffering
ci / test (push) Has been cancelled
ci / remote-bench (push) Has been cancelled

This commit is contained in:
DuProcess
2026-06-13 15:43:42 -04:00
parent 6a0f12c59e
commit 64b00783fd
2 changed files with 19 additions and 14 deletions
+11 -6
View File
@@ -1449,7 +1449,7 @@ async fn run_terminal(
let mut last_size = size().unwrap_or((80, 24)); let mut last_size = size().unwrap_or((80, 24));
let mut frame_buffer = FrameBuffer::default(); let mut frame_buffer = FrameBuffer::default();
let mut predictor = Predictor::new(predict && cred.mode != "view-only" && !forward_only); let mut predictor = Predictor::new(predict && cred.mode != "view-only" && !forward_only);
let (forward_tx, mut forward_rx) = mpsc::unbounded_channel::<ForwardEvent>(); let (forward_tx, mut forward_rx) = mpsc::channel::<ForwardEvent>(1024);
let _forward_keepalive = if local_forwards.is_empty() { let _forward_keepalive = if local_forwards.is_empty() {
Some(forward_tx.clone()) Some(forward_tx.clone())
} else { } else {
@@ -1502,6 +1502,7 @@ async fn run_terminal(
let mut recv_buf = vec![0u8; 65535]; let mut recv_buf = vec![0u8; 65535];
loop { loop {
tokio::select! { tokio::select! {
biased;
stdin_msg = stdin_rx.recv() => { stdin_msg = stdin_rx.recv() => {
match stdin_msg { match stdin_msg {
Some(bytes) => { Some(bytes) => {
@@ -1624,6 +1625,7 @@ async fn run_terminal(
stream_id: open.stream_id, stream_id: open.stream_id,
bytes: buf[..n].to_vec(), bytes: buf[..n].to_vec(),
}) })
.await
.is_err() .is_err()
{ {
return; return;
@@ -1634,7 +1636,7 @@ async fn run_terminal(
} }
let _ = forward_tx.send(ForwardEvent::Close { let _ = forward_tx.send(ForwardEvent::Close {
stream_id: open.stream_id, stream_id: open.stream_id,
}); }).await;
}); });
} }
Err(err) => { Err(err) => {
@@ -1761,7 +1763,7 @@ async fn run_terminal(
async fn start_local_forwards( async fn start_local_forwards(
local_forwards: &[LocalForward], local_forwards: &[LocalForward],
forward_tx: mpsc::UnboundedSender<ForwardEvent>, forward_tx: mpsc::Sender<ForwardEvent>,
stream_ids: Arc<AtomicU64>, stream_ids: Arc<AtomicU64>,
) -> Result<()> { ) -> Result<()> {
for forward in local_forwards { for forward in local_forwards {
@@ -1779,12 +1781,14 @@ async fn start_local_forwards(
}; };
let stream_id = stream_ids.fetch_add(1, Ordering::Relaxed); let stream_id = stream_ids.fetch_add(1, Ordering::Relaxed);
let (mut reader, writer) = stream.into_split(); let (mut reader, writer) = stream.into_split();
let _ = forward_tx.send(ForwardEvent::Open { let _ = forward_tx
.send(ForwardEvent::Open {
stream_id, stream_id,
target_host: forward.target_host.clone(), target_host: forward.target_host.clone(),
target_port: forward.target_port, target_port: forward.target_port,
writer, writer,
}); })
.await;
let forward_tx = forward_tx.clone(); let forward_tx = forward_tx.clone();
tokio::spawn(async move { tokio::spawn(async move {
let mut buf = [0u8; 16 * 1024]; let mut buf = [0u8; 16 * 1024];
@@ -1797,6 +1801,7 @@ async fn start_local_forwards(
stream_id, stream_id,
bytes: buf[..n].to_vec(), bytes: buf[..n].to_vec(),
}) })
.await
.is_err() .is_err()
{ {
return; return;
@@ -1805,7 +1810,7 @@ async fn start_local_forwards(
Err(_) => break, Err(_) => break,
} }
} }
let _ = forward_tx.send(ForwardEvent::Close { stream_id }); let _ = forward_tx.send(ForwardEvent::Close { stream_id }).await;
}); });
} }
}); });
+4 -4
View File
@@ -189,7 +189,7 @@ struct ClientState {
pending: VecDeque<PendingFrame>, pending: VecDeque<PendingFrame>,
last_screen: Option<vt100::Screen>, last_screen: Option<vt100::Screen>,
allowed_forwardings: Vec<ForwardingRequest>, allowed_forwardings: Vec<ForwardingRequest>,
stream_writers: HashMap<u64, mpsc::UnboundedSender<Vec<u8>>>, stream_writers: HashMap<u64, mpsc::Sender<Vec<u8>>>,
} }
#[derive(Clone)] #[derive(Clone)]
@@ -1014,7 +1014,7 @@ async fn handle_stream_open(
} }
}; };
let (mut reader, mut writer) = stream.into_split(); let (mut reader, mut writer) = stream.into_split();
let (writer_tx, mut writer_rx) = mpsc::unbounded_channel::<Vec<u8>>(); let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
{ {
let mut locked = state.lock().expect("server state poisoned"); let mut locked = state.lock().expect("server state poisoned");
let session = locked let session = locked
@@ -1096,7 +1096,7 @@ async fn handle_stream_data(
client.stream_writers.get(&data.stream_id).cloned() client.stream_writers.get(&data.stream_id).cloned()
}; };
if let Some(writer) = writer { if let Some(writer) = writer {
let _ = writer.send(data.bytes); let _ = writer.send(data.bytes).await;
} }
Ok(()) Ok(())
} }
@@ -1227,7 +1227,7 @@ async fn start_remote_forwards(
}; };
let stream_id = allocate_server_stream_id(&state); let stream_id = allocate_server_stream_id(&state);
let (mut reader, mut writer) = stream.into_split(); let (mut reader, mut writer) = stream.into_split();
let (writer_tx, mut writer_rx) = mpsc::unbounded_channel::<Vec<u8>>(); let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
{ {
let mut locked = state.lock().expect("server state poisoned"); let mut locked = state.lock().expect("server state poisoned");
if let Some(session) = locked.sessions.get_mut(&session_name) if let Some(session) = locked.sessions.get_mut(&session_name)