1287 lines
46 KiB
Rust
1287 lines
46 KiB
Rust
use std::collections::{HashMap, VecDeque};
|
|
use std::path::Path;
|
|
use std::sync::{Arc, Mutex, Weak};
|
|
use std::time::Duration;
|
|
|
|
use crate::config::AgentConfig;
|
|
use crate::protocol::{AgentTerminalSession, TerminalAgentHandshake, TerminalControl, TerminalId};
|
|
use anyhow::{Context, Result};
|
|
use futures_util::{SinkExt, StreamExt};
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::sync::{mpsc, oneshot};
|
|
use tokio_tungstenite::tungstenite::Message;
|
|
use tracing::{info, warn};
|
|
|
|
const MAX_TERMINAL_FRAME_BYTES: usize = 64 * 1024;
|
|
const TERMINAL_SCROLLBACK_ROWS: usize = 1_000;
|
|
const RELAY_INPUT_QUEUE: usize = 32;
|
|
const RELAY_OUTPUT_QUEUE: usize = 32;
|
|
const PTY_EXIT_DRAIN_TIMEOUT: Duration = Duration::from_secs(1);
|
|
const PROCESS_SIGNAL_GRACE: Duration = Duration::from_secs(1);
|
|
|
|
/// Tracks the terminal's current rendered state while the browser is detached.
|
|
///
|
|
/// A snapshot is terminal escape output, so the browser can restore the screen
|
|
/// through its normal parser without a second state protocol.
|
|
struct TerminalState {
|
|
parser: vt100::Parser,
|
|
}
|
|
|
|
/// Buffers ordered terminal output without letting a slow relay grow memory
|
|
/// without bound. Once a frame is pending, the PTY read branch is paused until
|
|
/// the relay channel has capacity; control input remains independently active.
|
|
#[derive(Default)]
|
|
struct RelayOutputBuffer {
|
|
sender: Option<mpsc::Sender<Message>>,
|
|
pending: VecDeque<Message>,
|
|
snapshot_pending: bool,
|
|
backpressure_wait: Option<wakey_core::telemetry::Operation>,
|
|
}
|
|
|
|
impl RelayOutputBuffer {
|
|
fn connect(&mut self, sender: mpsc::Sender<Message>) {
|
|
self.disconnect();
|
|
self.sender = Some(sender);
|
|
}
|
|
|
|
fn disconnect(&mut self) {
|
|
if let Some(mut wait) = self.backpressure_wait.take() {
|
|
wait.outcome("disconnected");
|
|
}
|
|
self.sender = None;
|
|
self.pending.clear();
|
|
self.snapshot_pending = false;
|
|
}
|
|
|
|
fn is_connected(&self) -> bool {
|
|
self.sender.is_some()
|
|
}
|
|
|
|
fn has_pending(&self) -> bool {
|
|
!self.pending.is_empty()
|
|
}
|
|
|
|
fn pending_sender(&self) -> Option<mpsc::Sender<Message>> {
|
|
self.has_pending().then(|| self.sender.clone()).flatten()
|
|
}
|
|
|
|
fn queue_live(&mut self, frame: Message) {
|
|
debug_assert!(
|
|
self.pending.is_empty(),
|
|
"PTY reads must pause while relay output is pending"
|
|
);
|
|
self.queue_frame(frame);
|
|
}
|
|
|
|
fn queue_frame(&mut self, frame: Message) {
|
|
let Some(sender) = self.sender.as_ref() else {
|
|
return;
|
|
};
|
|
if !self.pending.is_empty() {
|
|
self.pending.push_back(frame);
|
|
return;
|
|
}
|
|
match sender.try_send(frame) {
|
|
Ok(()) => {}
|
|
Err(mpsc::error::TrySendError::Full(frame)) => {
|
|
// One measurement per backlog episode, not per select poll.
|
|
self.backpressure_wait = Some(wakey_core::telemetry::Operation::new(
|
|
"terminal_backpressure_wait",
|
|
));
|
|
wakey_core::telemetry::count(
|
|
wakey_core::telemetry::Counter::TerminalBackpressure,
|
|
1,
|
|
);
|
|
self.pending.push_back(frame);
|
|
}
|
|
Err(mpsc::error::TrySendError::Closed(_)) => self.disconnect(),
|
|
}
|
|
}
|
|
|
|
fn queue_snapshot(&mut self, snapshot: Vec<u8>) {
|
|
if self.sender.is_none() || self.snapshot_pending {
|
|
return;
|
|
}
|
|
for chunk in snapshot.chunks(MAX_TERMINAL_FRAME_BYTES) {
|
|
self.queue_frame(Message::Binary(chunk.to_vec().into()));
|
|
if self.sender.is_none() {
|
|
return;
|
|
}
|
|
}
|
|
self.snapshot_pending = self.has_pending();
|
|
}
|
|
|
|
fn send_reserved(&mut self, permit: mpsc::OwnedPermit<Message>) {
|
|
let frame = self
|
|
.pending
|
|
.pop_front()
|
|
.expect("relay permit requires pending output");
|
|
permit.send(frame);
|
|
if self.pending.is_empty() {
|
|
self.snapshot_pending = false;
|
|
if let Some(mut wait) = self.backpressure_wait.take() {
|
|
wait.outcome("ok");
|
|
}
|
|
}
|
|
}
|
|
|
|
fn live_sender(&self) -> Option<mpsc::Sender<Message>> {
|
|
self.sender.clone()
|
|
}
|
|
}
|
|
|
|
impl TerminalState {
|
|
fn new(rows: u16, cols: u16) -> Self {
|
|
Self {
|
|
parser: vt100::Parser::new(rows, cols, TERMINAL_SCROLLBACK_ROWS),
|
|
}
|
|
}
|
|
|
|
fn process(&mut self, bytes: &[u8]) {
|
|
self.parser.process(bytes);
|
|
}
|
|
|
|
fn resize(&mut self, rows: u16, cols: u16) {
|
|
self.parser.screen_mut().set_size(rows, cols);
|
|
}
|
|
|
|
#[tracing::instrument(name = "terminal_snapshot_generate", skip_all, fields(bytes = tracing::field::Empty))]
|
|
fn snapshot(&self) -> Vec<u8> {
|
|
let snapshot = self
|
|
.parser
|
|
.screen()
|
|
.snapshot_formatted(TERMINAL_SCROLLBACK_ROWS);
|
|
tracing::Span::current().record("bytes", snapshot.len());
|
|
snapshot
|
|
}
|
|
}
|
|
|
|
struct ActiveTerminal {
|
|
cancel: Option<oneshot::Sender<()>>,
|
|
relay_credentials: mpsc::UnboundedSender<String>,
|
|
created_at_unix: u64,
|
|
session_ttl_seconds: u64,
|
|
}
|
|
|
|
/// Owns PTY workers independently of any individual control-plane connection.
|
|
pub struct TerminalManager {
|
|
active: Arc<Mutex<HashMap<String, ActiveTerminal>>>,
|
|
max_sessions: usize,
|
|
events: mpsc::UnboundedSender<TerminalManagerEvent>,
|
|
}
|
|
|
|
pub struct TerminalManagerEvent {
|
|
pub terminal_id: TerminalId,
|
|
pub error: String,
|
|
}
|
|
|
|
impl TerminalManager {
|
|
pub fn new(config: &AgentConfig) -> (Self, mpsc::UnboundedReceiver<TerminalManagerEvent>) {
|
|
let (events, event_rx) = mpsc::unbounded_channel();
|
|
(
|
|
Self {
|
|
active: Arc::new(Mutex::new(HashMap::new())),
|
|
max_sessions: config.terminal.max_sessions.max(1),
|
|
events,
|
|
},
|
|
event_rx,
|
|
)
|
|
}
|
|
|
|
#[tracing::instrument(name = "terminal_open", skip_all, fields(terminal_id = %terminal_id))]
|
|
pub fn open(
|
|
&self,
|
|
config: &AgentConfig,
|
|
terminal_id: TerminalId,
|
|
relay_token: String,
|
|
rows: u16,
|
|
cols: u16,
|
|
) -> Result<()> {
|
|
if !config.terminal.enabled {
|
|
anyhow::bail!("terminal capability is disabled");
|
|
}
|
|
validate_size(rows, cols)?;
|
|
let session_ttl = config.terminal.session_ttl()?;
|
|
|
|
let terminal_key = terminal_id.to_string();
|
|
let (cancel_tx, cancel_rx) = oneshot::channel();
|
|
let (relay_tx, relay_rx) = mpsc::unbounded_channel();
|
|
let created_at_unix = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap_or_default()
|
|
.as_secs();
|
|
{
|
|
let mut active = self.active.lock().expect("terminal manager poisoned");
|
|
if active.contains_key(&terminal_key) {
|
|
anyhow::bail!("terminal session {terminal_key} is already active");
|
|
}
|
|
if active.len() >= self.max_sessions {
|
|
anyhow::bail!("agent terminal session limit reached");
|
|
}
|
|
active.insert(
|
|
terminal_key.clone(),
|
|
ActiveTerminal {
|
|
cancel: Some(cancel_tx),
|
|
relay_credentials: relay_tx.clone(),
|
|
created_at_unix,
|
|
session_ttl_seconds: config.terminal.session_ttl_seconds,
|
|
},
|
|
);
|
|
}
|
|
let _ = relay_tx.send(relay_token);
|
|
|
|
let config = config.clone();
|
|
let active = Arc::downgrade(&self.active);
|
|
let events = self.events.clone();
|
|
tokio::spawn(async move {
|
|
let _active_metric = wakey_core::telemetry::ActiveTerminal::new();
|
|
let result = run_terminal(
|
|
&config,
|
|
&terminal_id,
|
|
rows,
|
|
cols,
|
|
session_ttl,
|
|
cancel_rx,
|
|
relay_rx,
|
|
)
|
|
.await;
|
|
// Inventory is authoritative. Remove the stopped worker before
|
|
// notifying the control session, which may immediately resync it.
|
|
remove_completed(&active, terminal_id.as_str());
|
|
if let Err(err) = result {
|
|
warn!(terminal_id = %terminal_id, error = %err, "terminal worker failed");
|
|
let _ = events.send(TerminalManagerEvent {
|
|
terminal_id: terminal_id.clone(),
|
|
error: err.to_string(),
|
|
});
|
|
}
|
|
});
|
|
Ok(())
|
|
}
|
|
|
|
#[tracing::instrument(name = "terminal_close", skip_all, fields(terminal_id = %terminal_id))]
|
|
pub fn close(&self, terminal_id: &TerminalId) -> bool {
|
|
let mut active = self.active.lock().expect("terminal manager poisoned");
|
|
let Some(terminal) = active.get_mut(terminal_id.as_str()) else {
|
|
return false;
|
|
};
|
|
let Some(cancel) = terminal.cancel.take() else {
|
|
return true;
|
|
};
|
|
// A failed send means the worker has already dropped its receiver.
|
|
// Its wrapper still owns final registry removal, so this close request
|
|
// is successful either way.
|
|
let _ = cancel.send(());
|
|
true
|
|
}
|
|
|
|
#[tracing::instrument(name = "terminal_resume", skip_all, fields(terminal_id = %terminal_id))]
|
|
pub fn resume(&self, terminal_id: &TerminalId, relay_token: String) -> Result<()> {
|
|
let mut operation = wakey_core::telemetry::Operation::new("terminal_resume");
|
|
let active = self.active.lock().expect("terminal manager poisoned");
|
|
let session = active
|
|
.get(terminal_id.as_str())
|
|
.with_context(|| format!("terminal session {terminal_id} is not active"))?;
|
|
if session.relay_credentials.send(relay_token).is_err() {
|
|
anyhow::bail!("terminal session {terminal_id} has stopped");
|
|
}
|
|
operation.outcome("ok");
|
|
info!("terminal resume credentials queued for worker");
|
|
Ok(())
|
|
}
|
|
|
|
pub fn sessions(&self) -> Vec<AgentTerminalSession> {
|
|
let active = self.active.lock().expect("terminal manager poisoned");
|
|
active
|
|
.iter()
|
|
.filter(|(_, terminal)| {
|
|
terminal.cancel.is_some() && !terminal.relay_credentials.is_closed()
|
|
})
|
|
.filter_map(|(terminal_id, active)| {
|
|
TerminalId::new(terminal_id.clone())
|
|
.ok()
|
|
.map(|terminal_id| AgentTerminalSession {
|
|
terminal_id,
|
|
created_at_unix: active.created_at_unix,
|
|
session_ttl_seconds: active.session_ttl_seconds,
|
|
})
|
|
})
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
impl Drop for TerminalManager {
|
|
fn drop(&mut self) {
|
|
if Arc::strong_count(&self.active) == 1
|
|
&& let Ok(mut active) = self.active.lock()
|
|
{
|
|
for (_, active) in active.drain() {
|
|
if let Some(cancel) = active.cancel {
|
|
let _ = cancel.send(());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn remove_completed(active: &Weak<Mutex<HashMap<String, ActiveTerminal>>>, terminal_id: &str) {
|
|
if let Some(active) = active.upgrade()
|
|
&& let Ok(mut active) = active.lock()
|
|
{
|
|
active.remove(terminal_id);
|
|
}
|
|
}
|
|
|
|
#[cfg(unix)]
|
|
async fn run_terminal(
|
|
config: &AgentConfig,
|
|
terminal_id: &TerminalId,
|
|
rows: u16,
|
|
cols: u16,
|
|
session_ttl: Option<Duration>,
|
|
mut cancel: oneshot::Receiver<()>,
|
|
mut relay_credentials: mpsc::UnboundedReceiver<String>,
|
|
) -> Result<()> {
|
|
let terminal = match wakey::wakey_linux::terminal::TerminalPty::spawn(
|
|
Path::new(&config.terminal.shell),
|
|
&config.terminal.args,
|
|
config.terminal.current_dir.as_deref(),
|
|
rows,
|
|
cols,
|
|
) {
|
|
Ok(terminal) => terminal,
|
|
Err(err) => {
|
|
return Err(err);
|
|
}
|
|
};
|
|
let (mut reader, mut writer, mut child) = terminal.into_parts();
|
|
let process_group = child.id();
|
|
info!(
|
|
terminal_id = %terminal_id,
|
|
program = %config.terminal.shell.display(),
|
|
args = ?config.terminal.args,
|
|
current_dir = ?config.terminal.current_dir,
|
|
"terminal PTY ready"
|
|
);
|
|
|
|
let (relay_input_tx, mut relay_input_rx) = mpsc::channel(RELAY_INPUT_QUEUE);
|
|
let mut relay_output = RelayOutputBuffer::default();
|
|
let mut relay_task: Option<tokio::task::JoinHandle<()>> = None;
|
|
let mut relay_generation = 0_u64;
|
|
let mut terminal_state = TerminalState::new(rows, cols);
|
|
let mut output = [0_u8; 16 * 1024];
|
|
let mut requested_close = false;
|
|
let mut ttl_expired = false;
|
|
let mut observed_status = None;
|
|
let mut drain_deadline: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;
|
|
let mut ttl_deadline = session_ttl.map(|ttl| Box::pin(tokio::time::sleep(ttl)));
|
|
loop {
|
|
let pending_output_sender = relay_output.pending_sender();
|
|
tokio::select! {
|
|
biased;
|
|
_ = &mut cancel => {
|
|
requested_close = true;
|
|
break;
|
|
}
|
|
_ = async {
|
|
match ttl_deadline.as_mut() {
|
|
Some(deadline) => deadline.as_mut().await,
|
|
None => std::future::pending().await,
|
|
}
|
|
} => {
|
|
ttl_expired = true;
|
|
break;
|
|
}
|
|
status = child.wait(), if observed_status.is_none() => {
|
|
observed_status = Some(status.context("failed waiting for terminal child")?);
|
|
drain_deadline = Some(Box::pin(tokio::time::sleep(PTY_EXIT_DRAIN_TIMEOUT)));
|
|
}
|
|
_ = async {
|
|
match drain_deadline.as_mut() {
|
|
Some(deadline) => deadline.as_mut().await,
|
|
None => std::future::pending().await,
|
|
}
|
|
} => {
|
|
break;
|
|
}
|
|
incoming = relay_input_rx.recv() => {
|
|
let Some(message) = incoming else { break; };
|
|
match message {
|
|
RelayInput::Binary { generation, bytes } if generation == relay_generation => {
|
|
if bytes.len() > MAX_TERMINAL_FRAME_BYTES {
|
|
anyhow::bail!("terminal input frame exceeds size limit");
|
|
}
|
|
writer.write_all(&bytes).await.context("failed to write PTY input")?;
|
|
}
|
|
RelayInput::Resize { generation, rows, cols } if generation == relay_generation => {
|
|
validate_size(rows, cols)?;
|
|
writer.resize(rows, cols)?;
|
|
terminal_state.resize(rows, cols);
|
|
}
|
|
RelayInput::Snapshot { generation } if generation == relay_generation => {
|
|
relay_output.queue_snapshot(terminal_state.snapshot());
|
|
if let Err(err) = writer.refresh() {
|
|
warn!(terminal_id = %terminal_id, error = %err, "terminal redraw signal failed");
|
|
}
|
|
}
|
|
RelayInput::Close { generation } if generation == relay_generation => {
|
|
requested_close = true;
|
|
break;
|
|
}
|
|
RelayInput::Connected { generation, output } if generation == relay_generation => {
|
|
relay_output.connect(output);
|
|
}
|
|
RelayInput::Disconnected { generation } if generation == relay_generation => {
|
|
relay_output.disconnect();
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
credential = relay_credentials.recv() => {
|
|
let Some(relay_token) = credential else { break; };
|
|
if let Some(task) = relay_task.take() {
|
|
task.abort();
|
|
}
|
|
relay_generation = relay_generation.wrapping_add(1);
|
|
let generation = relay_generation;
|
|
let (output_tx, output_rx) = mpsc::channel(RELAY_OUTPUT_QUEUE);
|
|
relay_output.disconnect();
|
|
let config = config.clone();
|
|
let terminal_id = terminal_id.clone();
|
|
let relay_input_tx = relay_input_tx.clone();
|
|
relay_task = Some(tokio::spawn(async move {
|
|
if let Err(err) = run_terminal_relay(RelayConnection {
|
|
config,
|
|
terminal_id: terminal_id.clone(),
|
|
relay_token,
|
|
generation,
|
|
output_tx,
|
|
output_rx,
|
|
input: relay_input_tx.clone(),
|
|
}).await {
|
|
warn!(terminal_id = %terminal_id, error = %err, "terminal relay disconnected");
|
|
}
|
|
let _ = relay_input_tx.send(RelayInput::Disconnected { generation }).await;
|
|
}));
|
|
}
|
|
permit = reserve_relay_output(pending_output_sender), if relay_output.has_pending() => {
|
|
match permit {
|
|
Ok(permit) => relay_output.send_reserved(permit),
|
|
Err(()) => {
|
|
relay_output.disconnect();
|
|
if let Some(task) = relay_task.take() {
|
|
task.abort();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
read = reader.read(&mut output), if !relay_output.has_pending() => {
|
|
match read {
|
|
Ok(0) => break,
|
|
Ok(count) => {
|
|
wakey_core::telemetry::count(wakey_core::telemetry::Counter::TerminalOutputBytes, count as u64);
|
|
terminal_state.process(&output[..count]);
|
|
relay_output.queue_live(Message::Binary(output[..count].to_vec().into()));
|
|
}
|
|
// Linux PTY masters commonly report EIO after the slave closes.
|
|
Err(err) if err.raw_os_error() == Some(5) => break,
|
|
Err(err) => return Err(err).context("failed to read PTY output"),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let status = match observed_status {
|
|
Some(status) => status,
|
|
None => terminate_process_group(&mut child, process_group).await?,
|
|
};
|
|
if relay_output.is_connected() {
|
|
let exit_frame = serde_json::to_string(&TerminalControl::Exited {
|
|
exit_code: status.code(),
|
|
})
|
|
.ok()
|
|
.map(|text| Message::Text(text.into()));
|
|
flush_relay_output(&mut relay_output, exit_frame, PTY_EXIT_DRAIN_TIMEOUT).await;
|
|
}
|
|
info!(terminal_id = %terminal_id, exit_code = ?status.code(), requested_close, ttl_expired, "terminal worker exited");
|
|
Ok(())
|
|
}
|
|
|
|
enum RelayInput {
|
|
Connected {
|
|
generation: u64,
|
|
output: mpsc::Sender<Message>,
|
|
},
|
|
Binary {
|
|
generation: u64,
|
|
bytes: Vec<u8>,
|
|
},
|
|
Resize {
|
|
generation: u64,
|
|
rows: u16,
|
|
cols: u16,
|
|
},
|
|
Snapshot {
|
|
generation: u64,
|
|
},
|
|
Close {
|
|
generation: u64,
|
|
},
|
|
Disconnected {
|
|
generation: u64,
|
|
},
|
|
}
|
|
|
|
#[cfg(unix)]
|
|
struct RelayConnection {
|
|
config: AgentConfig,
|
|
terminal_id: TerminalId,
|
|
relay_token: String,
|
|
generation: u64,
|
|
output_tx: mpsc::Sender<Message>,
|
|
output_rx: mpsc::Receiver<Message>,
|
|
input: mpsc::Sender<RelayInput>,
|
|
}
|
|
|
|
#[cfg(unix)]
|
|
async fn run_terminal_relay(relay: RelayConnection) -> Result<()> {
|
|
let ws_url = terminal_websocket_url(&relay.config.server_url, &relay.terminal_id)?;
|
|
let (stream, _) = tokio_tungstenite::connect_async(ws_url.as_str())
|
|
.await
|
|
.context("failed to connect terminal relay websocket")?;
|
|
let (mut sink, source) = stream.split();
|
|
send_json(
|
|
&mut sink,
|
|
&TerminalAgentHandshake::Auth {
|
|
agent_id: relay.config.agent_id.clone(),
|
|
relay_token: relay.relay_token,
|
|
},
|
|
)
|
|
.await?;
|
|
send_json(&mut sink, &TerminalControl::Ready).await?;
|
|
relay
|
|
.input
|
|
.send(RelayInput::Connected {
|
|
generation: relay.generation,
|
|
output: relay.output_tx,
|
|
})
|
|
.await
|
|
.map_err(|_| anyhow::anyhow!("terminal worker stopped"))?;
|
|
|
|
relay_loop(source, sink, relay.input, relay.generation, relay.output_rx).await
|
|
}
|
|
|
|
async fn relay_loop<S, St>(
|
|
mut source: St,
|
|
mut sink: S,
|
|
input: mpsc::Sender<RelayInput>,
|
|
generation: u64,
|
|
mut output: mpsc::Receiver<Message>,
|
|
) -> Result<()>
|
|
where
|
|
S: futures_util::Sink<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin,
|
|
St: futures_util::Stream<
|
|
Item = std::result::Result<Message, tokio_tungstenite::tungstenite::Error>,
|
|
> + Unpin,
|
|
{
|
|
let mut worker_stopped = false;
|
|
let mut worker_stop_deadline: Option<tokio::time::Instant> = None;
|
|
// Reserve at most one control so a full worker queue cannot block output.
|
|
let mut pending_input = None;
|
|
let mut close_after_input = false;
|
|
loop {
|
|
tokio::select! {
|
|
biased;
|
|
_ = async {
|
|
match worker_stop_deadline {
|
|
Some(deadline) => tokio::time::sleep_until(deadline).await,
|
|
None => std::future::pending().await,
|
|
}
|
|
} => break,
|
|
incoming = source.next(), if !worker_stopped && pending_input.is_none() => {
|
|
let Some(message) = incoming else { break; };
|
|
match message.context("terminal relay websocket receive failed")? {
|
|
Message::Binary(bytes) => {
|
|
wakey_core::telemetry::count(wakey_core::telemetry::Counter::TerminalInputBytes, bytes.len() as u64);
|
|
pending_input = Some(RelayInput::Binary {
|
|
generation,
|
|
bytes: bytes.to_vec(),
|
|
});
|
|
}
|
|
Message::Text(text) => match serde_json::from_str::<TerminalControl>(&text)
|
|
.context("invalid terminal control frame")?
|
|
{
|
|
TerminalControl::Resize { rows, cols } => {
|
|
pending_input = Some(RelayInput::Resize {
|
|
generation,
|
|
rows,
|
|
cols,
|
|
});
|
|
}
|
|
TerminalControl::Snapshot => {
|
|
pending_input = Some(RelayInput::Snapshot { generation });
|
|
}
|
|
TerminalControl::Close => {
|
|
pending_input = Some(RelayInput::Close { generation });
|
|
close_after_input = true;
|
|
}
|
|
_ => anyhow::bail!("terminal control frame has invalid direction"),
|
|
},
|
|
Message::Ping(payload) => sink.send(Message::Pong(payload)).await?,
|
|
Message::Pong(_) => {}
|
|
// Transport closure only detaches the relay. The agent-owned
|
|
// PTY remains alive and waits for replacement credentials.
|
|
Message::Close(_) => break,
|
|
Message::Frame(_) => {}
|
|
}
|
|
}
|
|
permit = input.reserve(), if !worker_stopped && pending_input.is_some() => {
|
|
match permit {
|
|
Ok(permit) => {
|
|
permit.send(pending_input.take().expect("pending terminal input"));
|
|
if close_after_input {
|
|
break;
|
|
}
|
|
}
|
|
Err(_) => {
|
|
worker_stopped = true;
|
|
pending_input = None;
|
|
worker_stop_deadline =
|
|
Some(tokio::time::Instant::now() + PTY_EXIT_DRAIN_TIMEOUT);
|
|
}
|
|
}
|
|
}
|
|
outgoing = output.recv() => {
|
|
let Some(message) = outgoing else { break; };
|
|
if let Some(deadline) = worker_stop_deadline {
|
|
tokio::time::timeout_at(deadline, sink.send(message))
|
|
.await
|
|
.context("terminal relay output drain timed out")?
|
|
.context("failed to send terminal relay output")?;
|
|
} else {
|
|
sink.send(message).await.context("failed to send terminal relay output")?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn reserve_relay_output(
|
|
sender: Option<mpsc::Sender<Message>>,
|
|
) -> Result<mpsc::OwnedPermit<Message>, ()> {
|
|
match sender {
|
|
Some(sender) => sender.reserve_owned().await.map_err(|_| ()),
|
|
None => std::future::pending().await,
|
|
}
|
|
}
|
|
|
|
/// Flushes frames the select loop left pending, then sends the exit lifecycle
|
|
/// frame behind them, bounded by `deadline`: a relay that stays saturated past
|
|
/// it drops the backlog rather than holding worker shutdown open. Queuing the
|
|
/// exit frame without this would drop it together with the pending buffer when
|
|
/// the worker returns.
|
|
async fn flush_relay_output(
|
|
relay_output: &mut RelayOutputBuffer,
|
|
exit_frame: Option<Message>,
|
|
deadline: Duration,
|
|
) {
|
|
let _ = tokio::time::timeout(deadline, async {
|
|
while relay_output.has_pending() {
|
|
let sender = relay_output.pending_sender();
|
|
match reserve_relay_output(sender).await {
|
|
Ok(permit) => relay_output.send_reserved(permit),
|
|
Err(()) => {
|
|
relay_output.disconnect();
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
let Some(frame) = exit_frame else {
|
|
return;
|
|
};
|
|
let sender = relay_output
|
|
.pending_sender()
|
|
.or_else(|| relay_output.live_sender());
|
|
match reserve_relay_output(sender).await {
|
|
Ok(permit) => {
|
|
permit.send(frame);
|
|
}
|
|
Err(()) => relay_output.disconnect(),
|
|
}
|
|
})
|
|
.await;
|
|
}
|
|
|
|
#[cfg(not(unix))]
|
|
async fn run_terminal(
|
|
_config: &AgentConfig,
|
|
_terminal_id: &TerminalId,
|
|
_rows: u16,
|
|
_cols: u16,
|
|
_session_ttl: Option<Duration>,
|
|
_cancel: oneshot::Receiver<()>,
|
|
_relay_credentials: mpsc::UnboundedReceiver<String>,
|
|
) -> Result<()> {
|
|
anyhow::bail!("terminal sessions are unsupported on this platform")
|
|
}
|
|
|
|
#[cfg(unix)]
|
|
async fn terminate_process_group(
|
|
child: &mut tokio::process::Child,
|
|
process_group: Option<u32>,
|
|
) -> Result<std::process::ExitStatus> {
|
|
use nix::sys::signal::{Signal, killpg};
|
|
use nix::unistd::Pid;
|
|
|
|
let Some(process_group) = process_group else {
|
|
child
|
|
.start_kill()
|
|
.context("failed to kill terminal child")?;
|
|
return child
|
|
.wait()
|
|
.await
|
|
.context("failed waiting for terminal child");
|
|
};
|
|
let pid = Pid::from_raw(process_group as i32);
|
|
let _ = killpg(pid, Signal::SIGHUP);
|
|
if let Ok(status) = tokio::time::timeout(PROCESS_SIGNAL_GRACE, child.wait()).await {
|
|
return status.context("failed waiting for terminal child after SIGHUP");
|
|
}
|
|
let _ = killpg(pid, Signal::SIGTERM);
|
|
if let Ok(status) = tokio::time::timeout(PROCESS_SIGNAL_GRACE, child.wait()).await {
|
|
return status.context("failed waiting for terminal child after SIGTERM");
|
|
}
|
|
let _ = killpg(pid, Signal::SIGKILL);
|
|
child
|
|
.wait()
|
|
.await
|
|
.context("failed waiting for terminal child after SIGKILL")
|
|
}
|
|
|
|
async fn send_json<S, T>(sink: &mut S, value: &T) -> Result<()>
|
|
where
|
|
S: futures_util::Sink<Message> + Unpin,
|
|
S::Error: std::error::Error + Send + Sync + 'static,
|
|
T: serde::Serialize,
|
|
{
|
|
let json = serde_json::to_string(value).context("failed to encode terminal frame")?;
|
|
sink.send(Message::Text(json.into()))
|
|
.await
|
|
.context("failed to send terminal frame")
|
|
}
|
|
|
|
fn terminal_websocket_url(server_url: &str, terminal_id: &TerminalId) -> Result<url::Url> {
|
|
let mut url = url::Url::parse(server_url).context("invalid server_url")?;
|
|
let scheme = match url.scheme() {
|
|
"http" => "ws",
|
|
"https" => "wss",
|
|
"ws" => "ws",
|
|
"wss" => "wss",
|
|
other => anyhow::bail!("unsupported server_url scheme `{other}`"),
|
|
};
|
|
url.set_scheme(scheme)
|
|
.map_err(|_| anyhow::anyhow!("failed to convert server_url scheme"))?;
|
|
url.set_path(&format!(
|
|
"/api/v1/agent/terminals/{}/ws",
|
|
terminal_id.as_str()
|
|
));
|
|
url.set_query(None);
|
|
url.set_fragment(None);
|
|
Ok(url)
|
|
}
|
|
|
|
fn validate_size(rows: u16, cols: u16) -> Result<()> {
|
|
if !(1..=300).contains(&rows) || !(1..=500).contains(&cols) {
|
|
anyhow::bail!("terminal size is outside supported bounds");
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn manager_with_stopped_terminal(terminal_id: &str) -> TerminalManager {
|
|
let (cancel, cancel_rx) = oneshot::channel();
|
|
drop(cancel_rx);
|
|
let (relay_credentials, relay_rx) = mpsc::unbounded_channel();
|
|
drop(relay_rx);
|
|
let (events, event_rx) = mpsc::unbounded_channel();
|
|
drop(event_rx);
|
|
TerminalManager {
|
|
active: Arc::new(Mutex::new(HashMap::from([(
|
|
terminal_id.to_string(),
|
|
ActiveTerminal {
|
|
cancel: Some(cancel),
|
|
relay_credentials,
|
|
created_at_unix: 42,
|
|
session_ttl_seconds: 600,
|
|
},
|
|
)]))),
|
|
max_sessions: 2,
|
|
events,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn terminal_url_uses_dedicated_agent_path() {
|
|
let id = TerminalId::new("term-1").expect("terminal id");
|
|
let url = terminal_websocket_url("https://example.com/base", &id).expect("url");
|
|
assert_eq!(
|
|
url.as_str(),
|
|
"wss://example.com/api/v1/agent/terminals/term-1/ws"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn failed_resume_hides_stopped_terminal_from_inventory() {
|
|
let manager = manager_with_stopped_terminal("stopped-terminal");
|
|
let terminal_id = TerminalId::new("stopped-terminal").expect("terminal id");
|
|
|
|
assert!(manager.resume(&terminal_id, "replacement".into()).is_err());
|
|
assert!(manager.sessions().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn inventory_prunes_terminal_with_stopped_worker() {
|
|
let manager = manager_with_stopped_terminal("stopped-terminal");
|
|
|
|
assert!(manager.sessions().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn close_succeeds_when_worker_already_dropped_cancel_receiver() {
|
|
let manager = manager_with_stopped_terminal("stopped-terminal");
|
|
let terminal_id = TerminalId::new("stopped-terminal").expect("terminal id");
|
|
|
|
assert!(manager.close(&terminal_id));
|
|
assert!(manager.sessions().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn close_keeps_session_registered_until_worker_completion() {
|
|
let (cancel, mut cancel_rx) = oneshot::channel();
|
|
let (relay_credentials, _relay_rx) = mpsc::unbounded_channel();
|
|
let (events, _event_rx) = mpsc::unbounded_channel();
|
|
let active = Arc::new(Mutex::new(HashMap::from([(
|
|
"closing-terminal".to_string(),
|
|
ActiveTerminal {
|
|
cancel: Some(cancel),
|
|
relay_credentials,
|
|
created_at_unix: 42,
|
|
session_ttl_seconds: 600,
|
|
},
|
|
)])));
|
|
let manager = TerminalManager {
|
|
active: active.clone(),
|
|
max_sessions: 2,
|
|
events,
|
|
};
|
|
let terminal_id = TerminalId::new("closing-terminal").expect("terminal id");
|
|
|
|
assert!(manager.close(&terminal_id));
|
|
assert!(cancel_rx.try_recv().is_ok());
|
|
assert!(
|
|
active
|
|
.lock()
|
|
.expect("terminal manager")
|
|
.contains_key("closing-terminal")
|
|
);
|
|
assert!(manager.sessions().is_empty());
|
|
|
|
remove_completed(&Arc::downgrade(&active), terminal_id.as_str());
|
|
assert!(
|
|
!active
|
|
.lock()
|
|
.expect("terminal manager")
|
|
.contains_key("closing-terminal")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn exit_frame_is_delivered_after_saturated_relay_drains() {
|
|
let (tx, mut rx) = mpsc::channel(RELAY_OUTPUT_QUEUE);
|
|
let mut relay_output = RelayOutputBuffer::default();
|
|
relay_output.connect(tx);
|
|
|
|
let total = RELAY_OUTPUT_QUEUE * 2;
|
|
for i in 0..total {
|
|
relay_output.queue_frame(Message::Text(format!("out-{i}").into()));
|
|
}
|
|
assert!(relay_output.has_pending());
|
|
|
|
let collected = tokio::spawn(async move {
|
|
let mut frames = Vec::new();
|
|
while let Some(frame) = rx.recv().await {
|
|
frames.push(frame);
|
|
}
|
|
frames
|
|
});
|
|
|
|
let exit_frame = serde_json::to_string(&TerminalControl::Exited { exit_code: Some(0) })
|
|
.expect("exit frame encodes");
|
|
flush_relay_output(
|
|
&mut relay_output,
|
|
Some(Message::Text(exit_frame.clone().into())),
|
|
PTY_EXIT_DRAIN_TIMEOUT,
|
|
)
|
|
.await;
|
|
relay_output.disconnect();
|
|
|
|
let mut frames = collected.await.expect("relay reader completes");
|
|
assert_eq!(frames.len(), total + 1);
|
|
let Some(Message::Text(last)) = frames.pop() else {
|
|
panic!("last frame should be the exit control frame");
|
|
};
|
|
assert_eq!(last.as_str(), exit_frame);
|
|
for (i, frame) in frames.iter().enumerate() {
|
|
let Message::Text(payload) = frame else {
|
|
panic!("expected text frame at position {i}, got {frame:?}");
|
|
};
|
|
assert_eq!(payload.as_str(), format!("out-{i}"));
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn relay_drains_final_output_after_worker_stops_during_concurrent_controls() {
|
|
let (input, input_rx) = mpsc::channel(RELAY_INPUT_QUEUE);
|
|
drop(input_rx);
|
|
let (output_tx, output_rx) = mpsc::channel(RELAY_OUTPUT_QUEUE);
|
|
output_tx
|
|
.send(Message::Binary(vec![1, 2, 3].into()))
|
|
.await
|
|
.expect("live output queues");
|
|
let exit_frame = serde_json::to_string(&TerminalControl::Exited { exit_code: Some(0) })
|
|
.expect("exit frame encodes");
|
|
output_tx
|
|
.send(Message::Text(exit_frame.clone().into()))
|
|
.await
|
|
.expect("exit output queues");
|
|
drop(output_tx);
|
|
|
|
let source = futures_util::stream::iter([
|
|
Ok(Message::Binary(vec![4].into())),
|
|
Ok(Message::Text(
|
|
serde_json::to_string(&TerminalControl::Resize {
|
|
rows: 30,
|
|
cols: 100,
|
|
})
|
|
.expect("resize encodes")
|
|
.into(),
|
|
)),
|
|
Ok(Message::Text(
|
|
serde_json::to_string(&TerminalControl::Snapshot)
|
|
.expect("snapshot encodes")
|
|
.into(),
|
|
)),
|
|
]);
|
|
let delivered = Arc::new(Mutex::new(Vec::new()));
|
|
let sink = futures_util::sink::unfold(delivered.clone(), |delivered, frame| async move {
|
|
delivered.lock().expect("delivered frames").push(frame);
|
|
Ok::<_, tokio_tungstenite::tungstenite::Error>(delivered)
|
|
});
|
|
|
|
relay_loop(source, Box::pin(sink), input, 7, output_rx)
|
|
.await
|
|
.expect("relay drains after worker completion");
|
|
|
|
let delivered = delivered.lock().expect("delivered frames");
|
|
assert_eq!(delivered.len(), 2);
|
|
assert_eq!(delivered[0].clone().into_data(), vec![1, 2, 3]);
|
|
let Message::Text(exit) = &delivered[1] else {
|
|
panic!("exit control should follow final bytes");
|
|
};
|
|
assert_eq!(exit.as_str(), exit_frame);
|
|
}
|
|
|
|
#[test]
|
|
fn open_rejects_invalid_initial_size_before_spawning() {
|
|
let mut config = crate::config::DEFAULT_CONFIG.clone();
|
|
config.terminal.enabled = true;
|
|
let (manager, _events) = TerminalManager::new(&config);
|
|
let terminal_id = TerminalId::new("invalid-size").expect("terminal id");
|
|
|
|
let error = manager
|
|
.open(&config, terminal_id, "relay".into(), 0, 80)
|
|
.expect_err("invalid rows must fail");
|
|
assert!(error.to_string().contains("outside supported bounds"));
|
|
assert!(manager.active.lock().expect("terminal manager").is_empty());
|
|
}
|
|
|
|
#[cfg(unix)]
|
|
#[tokio::test]
|
|
async fn local_ttl_terminates_pty_without_control_plane_close() {
|
|
let mut config = crate::config::DEFAULT_CONFIG.clone();
|
|
config.terminal.enabled = true;
|
|
config.terminal.shell = "/bin/sh".into();
|
|
config.terminal.args = vec!["-c".into(), "sleep 30".into()];
|
|
let terminal_id = TerminalId::new("ttl-test").expect("terminal id");
|
|
let (_cancel, cancel_rx) = oneshot::channel();
|
|
let (_relay, relay_rx) = mpsc::unbounded_channel();
|
|
|
|
tokio::time::timeout(
|
|
Duration::from_secs(3),
|
|
run_terminal(
|
|
&config,
|
|
&terminal_id,
|
|
24,
|
|
80,
|
|
Some(Duration::from_millis(20)),
|
|
cancel_rx,
|
|
relay_rx,
|
|
),
|
|
)
|
|
.await
|
|
.expect("local TTL should stop the PTY promptly")
|
|
.expect("TTL cleanup should succeed");
|
|
}
|
|
|
|
#[test]
|
|
fn snapshot_reconstructs_screen_content_and_cursor() {
|
|
let mut state = TerminalState::new(24, 80);
|
|
state.process(b"hello\r\n\x1b[31mred\x1b[0m\x1b[10;20Hcursor");
|
|
|
|
let mut restored = vt100::Parser::new(24, 80, 0);
|
|
restored.process(b"stale browser contents\x1b[24;80Hjunk");
|
|
restored.process(&state.snapshot());
|
|
|
|
assert_eq!(
|
|
restored.screen().contents(),
|
|
state.parser.screen().contents()
|
|
);
|
|
assert_eq!(
|
|
restored.screen().cursor_position(),
|
|
state.parser.screen().cursor_position()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn snapshot_reconstructs_terminal_input_modes() {
|
|
let mut state = TerminalState::new(24, 80);
|
|
state.process(b"\x1b[?1h\x1b[?1000h\x1b[?2004h");
|
|
|
|
let mut restored = vt100::Parser::new(24, 80, 0);
|
|
restored.process(&state.snapshot());
|
|
|
|
assert!(restored.screen().application_cursor());
|
|
assert!(restored.screen().bracketed_paste());
|
|
assert_eq!(
|
|
restored.screen().mouse_protocol_mode(),
|
|
state.parser.screen().mouse_protocol_mode()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn terminal_state_tracks_resize() {
|
|
let mut state = TerminalState::new(24, 80);
|
|
state.resize(40, 120);
|
|
assert_eq!(state.parser.screen().size(), (40, 120));
|
|
}
|
|
|
|
#[test]
|
|
fn snapshot_reconstructs_agent_scrollback() {
|
|
let mut state = TerminalState::new(3, 20);
|
|
for line in 0..10 {
|
|
if line == 0 {
|
|
state.process(b"\x1b[31mhistory 0\x1b[0m\r\n");
|
|
} else {
|
|
state.process(format!("history {line}\r\n").as_bytes());
|
|
}
|
|
}
|
|
let expected_screen = state.parser.screen().contents();
|
|
|
|
let mut restored = vt100::Parser::new(3, 20, TERMINAL_SCROLLBACK_ROWS);
|
|
restored.process(&state.snapshot());
|
|
|
|
assert_eq!(restored.screen().contents(), expected_screen);
|
|
restored.screen_mut().set_scrollback(usize::MAX);
|
|
assert!(restored.screen().scrollback() >= 8);
|
|
assert!(restored.screen().contents().contains("history 0"));
|
|
assert_eq!(
|
|
restored
|
|
.screen()
|
|
.cell(0, 0)
|
|
.expect("first history cell")
|
|
.fgcolor(),
|
|
vt100::Color::Idx(1)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn snapshot_restores_alternate_screen_mode() {
|
|
let mut state = TerminalState::new(3, 20);
|
|
state.process(b"shell\r\n\x1b[?1049hfull-screen");
|
|
|
|
let mut restored = vt100::Parser::new(3, 20, TERMINAL_SCROLLBACK_ROWS);
|
|
restored.process(&state.snapshot());
|
|
|
|
assert!(restored.screen().alternate_screen());
|
|
assert_eq!(restored.screen().contents(), "full-screen");
|
|
}
|
|
|
|
#[test]
|
|
fn alternate_screen_snapshot_restores_hidden_primary_on_exit() {
|
|
let mut state = TerminalState::new(3, 20);
|
|
state.process(b"\x1b[31mshell history\r\n$ btop\x1b[?1049h\x1b[34mbtop dashboard");
|
|
|
|
let mut restored = vt100::Parser::new(3, 20, TERMINAL_SCROLLBACK_ROWS);
|
|
restored.process(&state.snapshot());
|
|
|
|
assert!(restored.screen().alternate_screen());
|
|
assert_eq!(restored.screen().contents(), "btop dashboard");
|
|
assert_eq!(restored.screen().fgcolor(), vt100::Color::Idx(4));
|
|
|
|
// Compare the reattached terminal with the original parser after both
|
|
// receive the application's real alternate-screen exit sequence.
|
|
state.process(b"\x1b[?1049l");
|
|
restored.process(b"\x1b[?1049l");
|
|
assert_eq!(
|
|
restored.screen().contents(),
|
|
state.parser.screen().contents()
|
|
);
|
|
assert_eq!(
|
|
restored.screen().cursor_position(),
|
|
state.parser.screen().cursor_position()
|
|
);
|
|
assert_eq!(restored.screen().fgcolor(), vt100::Color::Idx(1));
|
|
|
|
restored.screen_mut().set_scrollback(usize::MAX);
|
|
assert!(restored.screen().contents().contains("shell history"));
|
|
assert!(!restored.screen().contents().contains("btop dashboard"));
|
|
}
|
|
|
|
#[test]
|
|
fn snapshot_does_not_change_the_source_viewport_or_state() {
|
|
let mut state = TerminalState::new(3, 20);
|
|
for line in 0..10 {
|
|
state.process(format!("history {line}\r\n").as_bytes());
|
|
}
|
|
state.parser.screen_mut().set_scrollback(4);
|
|
|
|
let before_contents = state.parser.screen().contents();
|
|
let before_state = state.parser.screen().state_formatted();
|
|
let before_scrollback = state.parser.screen().scrollback();
|
|
let _ = state.snapshot();
|
|
|
|
assert_eq!(state.parser.screen().contents(), before_contents);
|
|
assert_eq!(state.parser.screen().state_formatted(), before_state);
|
|
assert_eq!(state.parser.screen().scrollback(), before_scrollback);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn snapshot_is_split_at_the_terminal_frame_limit() {
|
|
let snapshot = vec![7; MAX_TERMINAL_FRAME_BYTES * 2 + 1];
|
|
let (tx, mut rx) = mpsc::channel(1);
|
|
let mut relay = RelayOutputBuffer::default();
|
|
relay.connect(tx);
|
|
relay.queue_snapshot(snapshot.clone());
|
|
|
|
let mut restored = Vec::new();
|
|
for _ in 0..3 {
|
|
if rx.is_empty() {
|
|
let permit = reserve_relay_output(relay.pending_sender())
|
|
.await
|
|
.expect("relay remains connected");
|
|
relay.send_reserved(permit);
|
|
}
|
|
let Message::Binary(chunk) = rx.recv().await.expect("snapshot chunk") else {
|
|
panic!("snapshot chunks must be binary");
|
|
};
|
|
assert!(chunk.len() <= MAX_TERMINAL_FRAME_BYTES);
|
|
restored.extend_from_slice(&chunk);
|
|
}
|
|
assert_eq!(restored, snapshot);
|
|
assert!(!relay.has_pending());
|
|
assert!(relay.is_connected());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn saturated_relay_output_stays_attached_without_dropping() {
|
|
let (tx, mut rx) = mpsc::channel(1);
|
|
let mut relay = RelayOutputBuffer::default();
|
|
relay.connect(tx);
|
|
relay.queue_live(Message::Binary(vec![1].into()));
|
|
// This read discovers saturation. The worker retains it, then pauses
|
|
// further PTY reads until relay capacity becomes available.
|
|
relay.queue_live(Message::Binary(vec![2].into()));
|
|
|
|
assert!(relay.is_connected());
|
|
assert!(relay.has_pending());
|
|
assert!(relay.backpressure_wait.is_some());
|
|
assert_eq!(
|
|
rx.recv().await.expect("original queued frame").into_data(),
|
|
vec![1]
|
|
);
|
|
let permit = reserve_relay_output(relay.pending_sender())
|
|
.await
|
|
.expect("relay remains connected");
|
|
relay.send_reserved(permit);
|
|
assert_eq!(
|
|
rx.recv().await.expect("backpressured frame").into_data(),
|
|
vec![2]
|
|
);
|
|
assert!(!relay.has_pending());
|
|
assert!(relay.backpressure_wait.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn disconnect_finishes_backpressure_episode() {
|
|
let (tx, _rx) = mpsc::channel(1);
|
|
let mut relay = RelayOutputBuffer::default();
|
|
relay.connect(tx);
|
|
relay.queue_live(Message::Binary(vec![1].into()));
|
|
relay.queue_live(Message::Binary(vec![2].into()));
|
|
assert!(relay.backpressure_wait.is_some());
|
|
relay.disconnect();
|
|
assert!(relay.backpressure_wait.is_none());
|
|
assert!(!relay.has_pending());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn snapshot_waits_behind_pending_live_output() {
|
|
let (tx, mut rx) = mpsc::channel(1);
|
|
let mut relay = RelayOutputBuffer::default();
|
|
relay.connect(tx);
|
|
relay.queue_live(Message::Binary(vec![1].into()));
|
|
relay.queue_live(Message::Binary(vec![2].into()));
|
|
relay.queue_snapshot(vec![3]);
|
|
|
|
let mut received = Vec::new();
|
|
for _ in 0..3 {
|
|
if rx.is_empty() {
|
|
let permit = reserve_relay_output(relay.pending_sender())
|
|
.await
|
|
.expect("relay remains connected");
|
|
relay.send_reserved(permit);
|
|
}
|
|
received.extend(rx.recv().await.expect("ordered relay frame").into_data());
|
|
}
|
|
assert_eq!(received, vec![1, 2, 3]);
|
|
}
|
|
|
|
#[test]
|
|
fn replacement_relay_does_not_receive_old_pending_output() {
|
|
let (old_tx, mut old_rx) = mpsc::channel(1);
|
|
let (new_tx, mut new_rx) = mpsc::channel(1);
|
|
let mut relay = RelayOutputBuffer::default();
|
|
relay.connect(old_tx);
|
|
relay.queue_live(Message::Binary(vec![1].into()));
|
|
relay.queue_live(Message::Binary(vec![2].into()));
|
|
|
|
relay.connect(new_tx);
|
|
relay.queue_live(Message::Binary(vec![3].into()));
|
|
|
|
assert_eq!(
|
|
old_rx.try_recv().expect("old queued frame").into_data(),
|
|
vec![1]
|
|
);
|
|
assert!(old_rx.try_recv().is_err());
|
|
assert_eq!(
|
|
new_rx.try_recv().expect("new relay frame").into_data(),
|
|
vec![3]
|
|
);
|
|
assert!(!relay.has_pending());
|
|
}
|
|
}
|