feature(zone-exec): implement terminal resize support (#363)

This commit is contained in:
Alex Zenla
2024-08-26 00:43:07 -04:00
committed by GitHub
parent f1e3d59b6a
commit 0d2b7a3ae3
7 changed files with 165 additions and 28 deletions

View File

@ -3,11 +3,13 @@ use std::collections::HashMap;
use anyhow::Result;
use clap::Parser;
use crossterm::tty::IsTty;
use krata::v1::{
common::{ZoneTaskSpec, ZoneTaskSpecEnvVar},
common::{TerminalSize, ZoneTaskSpec, ZoneTaskSpecEnvVar},
control::{control_service_client::ControlServiceClient, ExecInsideZoneRequest},
};
use tokio::io::stdin;
use tonic::{transport::Channel, Request};
use crate::console::StdioConsoleStream;
@ -36,6 +38,7 @@ pub struct ZoneExecCommand {
impl ZoneExecCommand {
pub async fn run(self, mut client: ControlServiceClient<Channel>) -> Result<()> {
let zone_id: String = resolve_zone(&mut client, &self.zone).await?;
let should_map_tty = self.tty && stdin().is_tty();
let initial = ExecInsideZoneRequest {
zone_id,
task: Some(ZoneTaskSpec {
@ -52,16 +55,25 @@ impl ZoneExecCommand {
}),
stdin: vec![],
stdin_closed: false,
terminal_size: if should_map_tty {
let size = crossterm::terminal::size().ok();
size.map(|(columns, rows)| TerminalSize {
rows: rows as u32,
columns: columns as u32,
})
} else {
None
},
};
let stream = StdioConsoleStream::stdin_stream_exec(initial).await;
let stream = StdioConsoleStream::input_stream_exec(initial, should_map_tty).await;
let response = client
.exec_inside_zone(Request::new(stream))
.await?
.into_inner();
let code = StdioConsoleStream::exec_output(response, self.tty).await?;
let code = StdioConsoleStream::exec_output(response, should_map_tty).await?;
std::process::exit(code);
}
}

View File

@ -7,6 +7,7 @@ use crossterm::{
use krata::v1::common::ZoneState;
use krata::{
events::EventStream,
v1::common::TerminalSize,
v1::control::{
watch_events_reply::Event, ExecInsideZoneReply, ExecInsideZoneRequest, ZoneConsoleReply,
ZoneConsoleRequest,
@ -15,6 +16,7 @@ use krata::{
use log::debug;
use tokio::{
io::{stderr, stdin, stdout, AsyncReadExt, AsyncWriteExt},
select,
task::JoinHandle,
};
use tokio_stream::{Stream, StreamExt};
@ -22,6 +24,11 @@ use tonic::Streaming;
pub struct StdioConsoleStream;
enum ExecStdinSelect {
DataRead(std::io::Result<usize>),
TerminalResize,
}
impl StdioConsoleStream {
pub async fn stdin_stream(
zone: String,
@ -49,30 +56,106 @@ impl StdioConsoleStream {
}
}
pub async fn stdin_stream_exec(
#[cfg(unix)]
pub async fn input_stream_exec(
initial: ExecInsideZoneRequest,
tty: bool,
) -> impl Stream<Item = ExecInsideZoneRequest> {
let mut stdin = stdin();
stream! {
yield initial;
let mut buffer = vec![0u8; 60];
let mut terminal_size_change = if tty {
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::window_change()).ok()
} else {
None
};
let mut stdin_closed = false;
loop {
let size = match stdin.read(&mut buffer).await {
Ok(size) => size,
Err(error) => {
debug!("failed to read stdin: {}", error);
break;
let selected = if let Some(ref mut terminal_size_change) = terminal_size_change {
if stdin_closed {
select! {
_ = terminal_size_change.recv() => ExecStdinSelect::TerminalResize,
}
} else {
select! {
result = stdin.read(&mut buffer) => ExecStdinSelect::DataRead(result),
_ = terminal_size_change.recv() => ExecStdinSelect::TerminalResize,
}
}
} else {
select! {
result = stdin.read(&mut buffer) => ExecStdinSelect::DataRead(result),
}
};
let stdin = buffer[0..size].to_vec();
if size == 1 && buffer[0] == 0x1d {
break;
match selected {
ExecStdinSelect::DataRead(result) => {
match result {
Ok(size) => {
let stdin = buffer[0..size].to_vec();
if size == 1 && buffer[0] == 0x1d {
break;
}
stdin_closed = size == 0;
yield ExecInsideZoneRequest { zone_id: String::default(), task: None, terminal_size: None, stdin, stdin_closed, };
},
Err(error) => {
debug!("failed to read stdin: {}", error);
break;
}
}
},
ExecStdinSelect::TerminalResize => {
if let Ok((columns, rows)) = crossterm::terminal::size() {
yield ExecInsideZoneRequest { zone_id: String::default(), task: None, terminal_size: Some(TerminalSize {
rows: rows as u32,
columns: columns as u32,
}), stdin: vec![], stdin_closed: false, };
}
}
}
let stdin_closed = size == 0;
yield ExecInsideZoneRequest { zone_id: String::default(), task: None, stdin, stdin_closed, };
if stdin_closed {
break;
}
}
}
#[cfg(not(unix))]
pub async fn input_stream_exec(
initial: ExecInsideZoneRequest,
_tty: bool,
) -> impl Stream<Item = ExecInsideZoneRequest> {
let mut stdin = stdin();
stream! {
yield initial;
let mut buffer = vec![0u8; 60];
let mut stdin_closed = false;
loop {
let selected = select! {
result = stdin.read(&mut buffer) => ExecStdinSelect::DataRead(result),
};
match selected {
ExecStdinSelect::DataRead(result) => {
match result {
Ok(size) => {
let stdin = buffer[0..size].to_vec();
if size == 1 && buffer[0] == 0x1d {
break;
}
stdin_closed = size == 0;
yield ExecInsideZoneRequest { zone_id: String::default(), task: None, terminal_size: None, stdin, stdin_closed, };
},
Err(error) => {
debug!("failed to read stdin: {}", error);
break;
}
}
},
_ => {
continue;
}
}
}
}
@ -96,7 +179,7 @@ impl StdioConsoleStream {
}
pub async fn exec_output(mut stream: Streaming<ExecInsideZoneReply>, raw: bool) -> Result<i32> {
if raw && stdin().is_tty() {
if raw {
enable_raw_mode()?;
StdioConsoleStream::register_terminal_restore_hook()?;
}