mirror of
https://github.com/openai/codex.git
synced 2026-11-02 19:43:24 +00:00
Prevent Unix PTY I/O from blocking runtime shutdown (#40460)
## Why Blocking PTY reads and output-channel sends can keep Tokio runtime shutdown waiting when a detached child retains the terminal or output backpressure fills the channel. ## What changed - Drive Unix PTY reads and writes through nonblocking `AsyncFd` readiness so their tasks can be cancelled during shutdown. - Keep draining child output after its receiver closes, and preserve queued input plus EOF delivery when portable PTY stdin closes. - Return a descriptive error when PTY spawning uses a Tokio runtime without an I/O driver. ## Testing Add Unix coverage for detached children, full and dropped output channels, large input with EOF, inherited file descriptors, and runtimes without I/O. GitOrigin-RevId: f7f1f58abebf8bf1d39285cd61b8d69946d54155
This commit is contained in:
@@ -10,7 +10,7 @@ workspace = true
|
||||
[dependencies]
|
||||
anyhow = { workspace = true }
|
||||
portable-pty = { workspace = true }
|
||||
tokio = { workspace = true, features = ["io-util", "macros", "process", "rt-multi-thread", "sync", "time"] }
|
||||
tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt-multi-thread", "sync", "time"] }
|
||||
|
||||
[dev-dependencies]
|
||||
pretty_assertions = { workspace = true }
|
||||
|
||||
@@ -4,6 +4,8 @@ pub mod process_group;
|
||||
pub mod pty;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
#[cfg(unix)]
|
||||
mod unix_io;
|
||||
#[cfg(windows)]
|
||||
mod win;
|
||||
#[cfg(windows)]
|
||||
|
||||
@@ -18,6 +18,7 @@ use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex as StdMutex;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
#[cfg(not(unix))]
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
@@ -159,6 +160,12 @@ async fn spawn_process_portable(
|
||||
) -> Result<SpawnedProcess> {
|
||||
let pty_system = platform_native_pty_system();
|
||||
let pair = pty_system.openpty(size.into())?;
|
||||
#[cfg(unix)]
|
||||
let io = crate::unix_io::PtyIo::new(
|
||||
pair.master
|
||||
.as_raw_fd()
|
||||
.ok_or_else(|| anyhow::anyhow!("PTY master has no file descriptor"))?,
|
||||
)?;
|
||||
|
||||
let mut command_builder = CommandBuilder::new(arg0.as_ref().unwrap_or(&program.to_string()));
|
||||
command_builder.cwd(cwd);
|
||||
@@ -178,45 +185,56 @@ async fn spawn_process_portable(
|
||||
let process_group_id = child.process_id();
|
||||
let killer = child.clone_killer();
|
||||
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
|
||||
let (writer_tx, writer_rx) = mpsc::channel::<Vec<u8>>(128);
|
||||
let (stdout_tx, stdout_rx) = mpsc::channel::<Vec<u8>>(128);
|
||||
let (_stderr_tx, stderr_rx) = mpsc::channel::<Vec<u8>>(1);
|
||||
let mut reader = pair.master.try_clone_reader()?;
|
||||
let reader_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || {
|
||||
let mut buf = [0u8; 8_192];
|
||||
loop {
|
||||
match reader.read(&mut buf) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => {
|
||||
let _ = stdout_tx.blocking_send(buf[..n].to_vec());
|
||||
#[cfg(unix)]
|
||||
let (reader_handle, writer_handle) = io.spawn(
|
||||
stdout_tx,
|
||||
writer_rx,
|
||||
crate::unix_io::StdinCloseBehavior::SendEof,
|
||||
);
|
||||
#[cfg(not(unix))]
|
||||
let (reader_handle, writer_handle) = {
|
||||
let mut reader = pair.master.try_clone_reader()?;
|
||||
let reader_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || {
|
||||
let mut buf = [0u8; 8_192];
|
||||
loop {
|
||||
match reader.read(&mut buf) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => {
|
||||
let _ = stdout_tx.blocking_send(buf[..n].to_vec());
|
||||
}
|
||||
Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
|
||||
Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
continue;
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
|
||||
Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
continue;
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
let writer = pair.master.take_writer()?;
|
||||
let writer = Arc::new(tokio::sync::Mutex::new(writer));
|
||||
let writer_handle: JoinHandle<()> = tokio::spawn({
|
||||
let writer = Arc::clone(&writer);
|
||||
async move {
|
||||
#[cfg(windows)]
|
||||
let mut windows_input = crate::WindowsTtyInputNormalizer::default();
|
||||
while let Some(bytes) = writer_rx.recv().await {
|
||||
let mut writer_rx = writer_rx;
|
||||
let writer = pair.master.take_writer()?;
|
||||
let writer = Arc::new(tokio::sync::Mutex::new(writer));
|
||||
let writer_handle: JoinHandle<()> = tokio::spawn({
|
||||
let writer = Arc::clone(&writer);
|
||||
async move {
|
||||
#[cfg(windows)]
|
||||
let bytes = windows_input.normalize(&bytes);
|
||||
let mut guard = writer.lock().await;
|
||||
use std::io::Write;
|
||||
let _ = guard.write_all(&bytes);
|
||||
let _ = guard.flush();
|
||||
let mut windows_input = crate::WindowsTtyInputNormalizer::default();
|
||||
while let Some(bytes) = writer_rx.recv().await {
|
||||
#[cfg(windows)]
|
||||
let bytes = windows_input.normalize(&bytes);
|
||||
let mut guard = writer.lock().await;
|
||||
use std::io::Write;
|
||||
let _ = guard.write_all(&bytes);
|
||||
let _ = guard.flush();
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
(reader_handle, writer_handle)
|
||||
};
|
||||
|
||||
let (exit_tx, exit_rx) = oneshot::channel::<i32>();
|
||||
let exit_status = Arc::new(AtomicBool::new(false));
|
||||
@@ -280,6 +298,7 @@ async fn spawn_process_preserving_fds(
|
||||
inherited_fds: &[RawFd],
|
||||
) -> Result<SpawnedProcess> {
|
||||
let (master, slave) = open_unix_pty(size)?;
|
||||
let io = crate::unix_io::PtyIo::new(master.as_raw_fd())?;
|
||||
let mut command = StdCommand::new(program);
|
||||
if let Some(arg0) = arg0 {
|
||||
command.arg0(arg0);
|
||||
@@ -342,40 +361,14 @@ async fn spawn_process_preserving_fds(
|
||||
drop(slave);
|
||||
let process_group_id = child.id();
|
||||
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
|
||||
let (writer_tx, writer_rx) = mpsc::channel::<Vec<u8>>(128);
|
||||
let (stdout_tx, stdout_rx) = mpsc::channel::<Vec<u8>>(128);
|
||||
let (_stderr_tx, stderr_rx) = mpsc::channel::<Vec<u8>>(1);
|
||||
let mut reader = master.try_clone()?;
|
||||
let reader_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || {
|
||||
let mut buf = [0u8; 8_192];
|
||||
loop {
|
||||
match std::io::Read::read(&mut reader, &mut buf) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => {
|
||||
let _ = stdout_tx.blocking_send(buf[..n].to_vec());
|
||||
}
|
||||
Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
|
||||
Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
continue;
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let writer = Arc::new(tokio::sync::Mutex::new(master.try_clone()?));
|
||||
let writer_handle: JoinHandle<()> = tokio::spawn({
|
||||
let writer = Arc::clone(&writer);
|
||||
async move {
|
||||
while let Some(bytes) = writer_rx.recv().await {
|
||||
let mut guard = writer.lock().await;
|
||||
use std::io::Write;
|
||||
let _ = guard.write_all(&bytes);
|
||||
let _ = guard.flush();
|
||||
}
|
||||
}
|
||||
});
|
||||
let (reader_handle, writer_handle) = io.spawn(
|
||||
stdout_tx,
|
||||
writer_rx,
|
||||
crate::unix_io::StdinCloseBehavior::NoEof,
|
||||
);
|
||||
|
||||
let (exit_tx, exit_rx) = oneshot::channel::<i32>();
|
||||
let exit_status = Arc::new(AtomicBool::new(false));
|
||||
|
||||
@@ -873,6 +873,314 @@ async fn pipe_drop_reaps_child() -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
enum PtyShutdown {
|
||||
Terminate,
|
||||
Drop,
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn assert_pty_shutdown_with_detached_child(
|
||||
inherited_fds: &[i32],
|
||||
shutdown: PtyShutdown,
|
||||
) -> anyhow::Result<()> {
|
||||
let Some(python) = find_python() else {
|
||||
eprintln!("python not found; skipping detached PTY shutdown test");
|
||||
return Ok(());
|
||||
};
|
||||
let runtime = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()?;
|
||||
let (session, child_pid) = runtime.block_on(async {
|
||||
let marker = "__codex_detached_pid:";
|
||||
// The detached child retains the slave descriptors but outlives the
|
||||
// original process group. Its alarm bounds a regressed runtime drop.
|
||||
let script = format!(
|
||||
r"import os, select, signal, time, tty
|
||||
signal.alarm(10)
|
||||
if os.fork() == 0:
|
||||
os.setsid()
|
||||
signal.signal(signal.SIGHUP, signal.SIG_IGN)
|
||||
signal.alarm(10)
|
||||
tty.setraw(0)
|
||||
print('{marker}' + str(os.getpid()), flush=True)
|
||||
select.select([0], [], [])
|
||||
print('__codex_input_ready__', flush=True)
|
||||
time.sleep(10)
|
||||
os._exit(0)
|
||||
time.sleep(10)
|
||||
"
|
||||
);
|
||||
let env_map: HashMap<String, String> = std::env::vars().collect();
|
||||
let spawned = spawn_pty_process(
|
||||
&python,
|
||||
&["-c".to_string(), script],
|
||||
Path::new("."),
|
||||
&env_map,
|
||||
&None,
|
||||
TerminalSize::default(),
|
||||
inherited_fds,
|
||||
)
|
||||
.await?;
|
||||
let (session, mut output_rx, exit_rx) = combine_spawned_output(spawned);
|
||||
let child_pid = wait_for_marker_pid(&mut output_rx, marker, /*timeout_ms*/ 2_000).await?;
|
||||
let writer = session.writer_sender();
|
||||
writer.send(vec![b'x'; 1024 * 1024]).await?;
|
||||
writer.send(b"pending".to_vec()).await?;
|
||||
// The slave reports readable input without consuming it; the first write
|
||||
// exceeds terminal capacity, so the second chunk must remain queued.
|
||||
wait_for_output_contains(
|
||||
&mut output_rx,
|
||||
"__codex_input_ready__",
|
||||
/*timeout_ms*/ 2_000,
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(writer.capacity(), writer.max_capacity() - 1);
|
||||
drop(writer);
|
||||
let session = match shutdown {
|
||||
PtyShutdown::Terminate => {
|
||||
session.terminate();
|
||||
Some(session)
|
||||
}
|
||||
PtyShutdown::Drop => {
|
||||
drop(session);
|
||||
None
|
||||
}
|
||||
};
|
||||
tokio::time::timeout(std::time::Duration::from_secs(2), exit_rx).await??;
|
||||
Ok::<_, anyhow::Error>((session, child_pid))
|
||||
})?;
|
||||
|
||||
let started = std::time::Instant::now();
|
||||
drop(runtime);
|
||||
let elapsed = started.elapsed();
|
||||
let detached_child_alive = process_exists(child_pid)?;
|
||||
let _ = unsafe { libc::kill(child_pid, libc::SIGKILL) };
|
||||
drop(session);
|
||||
assert!(
|
||||
elapsed < std::time::Duration::from_secs(2),
|
||||
"runtime shutdown waited for detached PTY child: {elapsed:?}"
|
||||
);
|
||||
assert!(
|
||||
detached_child_alive,
|
||||
"detached child exited before runtime shutdown"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn pty_spawn_without_io_driver_returns_error() -> anyhow::Result<()> {
|
||||
use std::os::fd::AsRawFd;
|
||||
|
||||
let runtime = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_time()
|
||||
.build()?;
|
||||
let preserved_fd = std::fs::File::open("/dev/null")?;
|
||||
let env_map: HashMap<String, String> = std::env::vars().collect();
|
||||
let (program, args) = shell_command("true");
|
||||
for inherited_fds in [&[][..], &[preserved_fd.as_raw_fd()][..]] {
|
||||
let Err(error) = runtime.block_on(spawn_pty_process(
|
||||
&program,
|
||||
&args,
|
||||
Path::new("."),
|
||||
&env_map,
|
||||
&None,
|
||||
TerminalSize::default(),
|
||||
inherited_fds,
|
||||
)) else {
|
||||
anyhow::bail!("PTY spawn succeeded without a Tokio I/O driver");
|
||||
};
|
||||
assert_eq!(
|
||||
error.to_string(),
|
||||
"PTY I/O requires a Tokio runtime with I/O enabled"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn pty_terminate_allows_runtime_shutdown_with_detached_child() -> anyhow::Result<()> {
|
||||
assert_pty_shutdown_with_detached_child(&[], PtyShutdown::Terminate)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn pty_drop_allows_runtime_shutdown_with_detached_child() -> anyhow::Result<()> {
|
||||
assert_pty_shutdown_with_detached_child(&[], PtyShutdown::Drop)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn pty_preserving_fds_terminate_allows_runtime_shutdown_with_detached_child() -> anyhow::Result<()>
|
||||
{
|
||||
use std::os::fd::AsRawFd;
|
||||
|
||||
let preserved_fd = std::fs::File::open("/dev/null")?;
|
||||
assert_pty_shutdown_with_detached_child(&[preserved_fd.as_raw_fd()], PtyShutdown::Terminate)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn pty_preserving_fds_drop_allows_runtime_shutdown_with_detached_child() -> anyhow::Result<()> {
|
||||
use std::os::fd::AsRawFd;
|
||||
|
||||
let preserved_fd = std::fs::File::open("/dev/null")?;
|
||||
assert_pty_shutdown_with_detached_child(&[preserved_fd.as_raw_fd()], PtyShutdown::Drop)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn pty_terminate_allows_runtime_shutdown_with_full_output_channel() -> anyhow::Result<()> {
|
||||
use std::os::fd::AsRawFd;
|
||||
|
||||
let Some(python) = find_python() else {
|
||||
eprintln!("python not found; skipping PTY output backpressure test");
|
||||
return Ok(());
|
||||
};
|
||||
let preserved_fd = std::fs::File::open("/dev/null")?;
|
||||
for inherited_fds in [&[][..], &[preserved_fd.as_raw_fd()][..]] {
|
||||
let runtime = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()?;
|
||||
let (session, stdout_rx) = runtime.block_on(async {
|
||||
let env_map: HashMap<String, String> = std::env::vars().collect();
|
||||
let script =
|
||||
"import os, signal\nsignal.alarm(10)\nwhile True:\n os.write(1, b'x' * 8192)\n";
|
||||
let SpawnedProcess {
|
||||
session,
|
||||
stdout_rx,
|
||||
exit_rx,
|
||||
..
|
||||
} = spawn_pty_process(
|
||||
&python,
|
||||
&["-c".to_string(), script.to_string()],
|
||||
Path::new("."),
|
||||
&env_map,
|
||||
&None,
|
||||
TerminalSize::default(),
|
||||
inherited_fds,
|
||||
)
|
||||
.await?;
|
||||
tokio::time::timeout(std::time::Duration::from_secs(2), async {
|
||||
while stdout_rx.len() < stdout_rx.max_capacity() {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await?;
|
||||
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
|
||||
session.terminate();
|
||||
tokio::time::timeout(std::time::Duration::from_secs(2), exit_rx).await??;
|
||||
Ok::<_, anyhow::Error>((session, stdout_rx))
|
||||
})?;
|
||||
|
||||
// Keep the receiver alive through runtime shutdown. An OS-thread
|
||||
// watchdog releases a regressed blocking_send without a Tokio timer.
|
||||
let (shutdown_tx, shutdown_rx) = std::sync::mpsc::channel();
|
||||
let receiver_guard = std::thread::spawn(move || {
|
||||
let _ = shutdown_rx.recv_timeout(std::time::Duration::from_secs(3));
|
||||
drop(stdout_rx);
|
||||
});
|
||||
let started = std::time::Instant::now();
|
||||
drop(runtime);
|
||||
let elapsed = started.elapsed();
|
||||
let _ = shutdown_tx.send(());
|
||||
receiver_guard
|
||||
.join()
|
||||
.map_err(|_| anyhow::anyhow!("output receiver watchdog panicked"))?;
|
||||
drop(session);
|
||||
assert!(
|
||||
elapsed < std::time::Duration::from_secs(2),
|
||||
"runtime shutdown waited on a full PTY output channel: {elapsed:?}"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn pty_dropped_output_receiver_keeps_draining_child() -> anyhow::Result<()> {
|
||||
use std::os::fd::AsRawFd;
|
||||
|
||||
let Some(python) = find_python() else {
|
||||
eprintln!("python not found; skipping PTY dropped output receiver test");
|
||||
return Ok(());
|
||||
};
|
||||
let preserved_fd = std::fs::File::open("/dev/null")?;
|
||||
let env_map: HashMap<String, String> = std::env::vars().collect();
|
||||
let script =
|
||||
"import signal, sys\nsignal.alarm(10)\nsys.stdout.buffer.write(b'x' * (4 * 1024 * 1024))";
|
||||
for inherited_fds in [&[][..], &[preserved_fd.as_raw_fd()][..]] {
|
||||
let SpawnedProcess {
|
||||
session: _session,
|
||||
stdout_rx,
|
||||
exit_rx,
|
||||
..
|
||||
} = spawn_pty_process(
|
||||
&python,
|
||||
&["-c".to_string(), script.to_string()],
|
||||
Path::new("."),
|
||||
&env_map,
|
||||
&None,
|
||||
TerminalSize::default(),
|
||||
inherited_fds,
|
||||
)
|
||||
.await?;
|
||||
drop(stdout_rx);
|
||||
let code = tokio::time::timeout(std::time::Duration::from_secs(2), exit_rx).await??;
|
||||
assert_eq!(code, 0, "child should finish even when output is discarded");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn pty_close_stdin_preserves_large_input_and_delivers_eof() -> anyhow::Result<()> {
|
||||
let Some(python) = find_python() else {
|
||||
eprintln!("python not found; skipping PTY input backpressure and EOF test");
|
||||
return Ok(());
|
||||
};
|
||||
let script = r"import signal, sys, termios, time
|
||||
signal.alarm(10)
|
||||
attrs = termios.tcgetattr(0)
|
||||
attrs[3] &= ~termios.ECHO
|
||||
termios.tcsetattr(0, termios.TCSANOW, attrs)
|
||||
print('__ready__', flush=True)
|
||||
time.sleep(0.2)
|
||||
data = sys.stdin.buffer.read()
|
||||
# Closing the portable writer appends a newline before VEOF.
|
||||
assert data == b'0123456789abcdefghijklmnopqrstuvwxyz\n' * 32768 + b'\n', len(data)
|
||||
print('__complete__', flush=True)
|
||||
";
|
||||
let env_map: HashMap<String, String> = std::env::vars().collect();
|
||||
let spawned = spawn_pty_process(
|
||||
&python,
|
||||
&["-c".to_string(), script.to_string()],
|
||||
Path::new("."),
|
||||
&env_map,
|
||||
&None,
|
||||
TerminalSize::default(),
|
||||
&[],
|
||||
)
|
||||
.await?;
|
||||
let (session, mut output_rx, exit_rx) = combine_spawned_output(spawned);
|
||||
wait_for_output_contains(&mut output_rx, "__ready__", /*timeout_ms*/ 2_000).await?;
|
||||
let writer = session.writer_sender();
|
||||
writer
|
||||
.send(b"0123456789abcdefghijklmnopqrstuvwxyz\n".repeat(32_768))
|
||||
.await?;
|
||||
drop(writer);
|
||||
session.close_stdin();
|
||||
|
||||
let (output, code) = collect_output_until_exit(output_rx, exit_rx, /*timeout_ms*/ 5_000).await;
|
||||
assert_eq!(
|
||||
(code, String::from_utf8_lossy(&output).trim()),
|
||||
(0, "__complete__")
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn pty_terminate_reaps_child_when_waiter_is_queued() -> anyhow::Result<()> {
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
use std::io::Read;
|
||||
use std::io::Write;
|
||||
use std::os::fd::AsRawFd;
|
||||
use std::os::fd::BorrowedFd;
|
||||
use std::os::fd::RawFd;
|
||||
use std::sync::Arc;
|
||||
|
||||
use tokio::io::Interest;
|
||||
use tokio::io::unix::AsyncFd;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
pub(crate) enum StdinCloseBehavior {
|
||||
SendEof,
|
||||
NoEof,
|
||||
}
|
||||
|
||||
pub(crate) struct PtyIo {
|
||||
fd: Arc<AsyncFd<File>>,
|
||||
}
|
||||
|
||||
impl PtyIo {
|
||||
pub(crate) fn new(master_fd: RawFd) -> io::Result<Self> {
|
||||
// The PTY owner remains alive while its descriptor is duplicated.
|
||||
let fd = unsafe { BorrowedFd::borrow_raw(master_fd) }.try_clone_to_owned()?;
|
||||
let flags = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_GETFL) };
|
||||
if flags == -1
|
||||
|| unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_SETFL, flags | libc::O_NONBLOCK) } == -1
|
||||
{
|
||||
return Err(io::Error::last_os_error());
|
||||
}
|
||||
// Duplicates share O_NONBLOCK, so both reads and writes must use readiness.
|
||||
// Async tasks can be aborted even when a detached child keeps the PTY open
|
||||
// or the output channel is full; blocking read/send tasks cannot.
|
||||
// AsyncFd panics if the current runtime has no I/O driver.
|
||||
let fd = std::panic::catch_unwind(move || AsyncFd::new(File::from(fd)))
|
||||
.map_err(|_| io::Error::other("PTY I/O requires a Tokio runtime with I/O enabled"))??;
|
||||
Ok(Self { fd: Arc::new(fd) })
|
||||
}
|
||||
|
||||
pub(crate) fn spawn(
|
||||
self,
|
||||
stdout_tx: mpsc::Sender<Vec<u8>>,
|
||||
mut writer_rx: mpsc::Receiver<Vec<u8>>,
|
||||
stdin_close: StdinCloseBehavior,
|
||||
) -> (JoinHandle<()>, JoinHandle<()>) {
|
||||
let fd = self.fd;
|
||||
let reader = Arc::clone(&fd);
|
||||
let reader_handle = tokio::spawn(async move {
|
||||
let mut buf = [0u8; 8_192];
|
||||
loop {
|
||||
match reader
|
||||
.async_io(Interest::READABLE, |mut file| file.read(&mut buf))
|
||||
.await
|
||||
{
|
||||
Ok(0) => break,
|
||||
Ok(count) => {
|
||||
// Output caps may close the receiver before the child exits.
|
||||
let _ = stdout_tx.send(buf[..count].to_vec()).await;
|
||||
}
|
||||
Err(err) if err.kind() == io::ErrorKind::Interrupted => continue,
|
||||
// Unix PTYs may report EIO rather than EOF when the slave closes.
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
let writer_handle = tokio::spawn(async move {
|
||||
while let Some(bytes) = writer_rx.recv().await {
|
||||
if write_all(&fd, &bytes).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
match stdin_close {
|
||||
StdinCloseBehavior::SendEof => {
|
||||
// Preserve portable-pty's newline + current VEOF on stdin close,
|
||||
// including when the terminal input queue requires waiting.
|
||||
let mut termios = std::mem::MaybeUninit::<libc::termios>::uninit();
|
||||
if unsafe { libc::tcgetattr(fd.get_ref().as_raw_fd(), termios.as_mut_ptr()) }
|
||||
== 0
|
||||
{
|
||||
let eof = unsafe { termios.assume_init() }.c_cc[libc::VEOF];
|
||||
if eof != 0 {
|
||||
let _ = write_all(&fd, &[b'\n', eof]).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
StdinCloseBehavior::NoEof => {}
|
||||
}
|
||||
});
|
||||
(reader_handle, writer_handle)
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_all(fd: &AsyncFd<File>, mut bytes: &[u8]) -> io::Result<()> {
|
||||
while !bytes.is_empty() {
|
||||
match fd
|
||||
.async_io(Interest::WRITABLE, |mut file| file.write(bytes))
|
||||
.await
|
||||
{
|
||||
Ok(0) => return Err(io::ErrorKind::WriteZero.into()),
|
||||
Ok(count) => bytes = &bytes[count..],
|
||||
Err(err) if err.kind() == io::ErrorKind::Interrupted => continue,
|
||||
Err(err) => return Err(err),
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user