147 lines
4.5 KiB
Rust
147 lines
4.5 KiB
Rust
use core::fmt;
|
|
use std::io;
|
|
use std::sync::atomic::AtomicBool;
|
|
use std::sync::Arc;
|
|
use std::sync::Mutex as StdMutex;
|
|
|
|
use portable_pty::MasterPty;
|
|
use portable_pty::SlavePty;
|
|
use tokio::sync::broadcast;
|
|
use tokio::sync::mpsc;
|
|
use tokio::sync::oneshot;
|
|
use tokio::task::AbortHandle;
|
|
use tokio::task::JoinHandle;
|
|
|
|
pub(crate) trait ChildTerminator: Send + Sync {
|
|
fn kill(&mut self) -> io::Result<()>;
|
|
}
|
|
|
|
pub struct PtyHandles {
|
|
pub _slave: Option<Box<dyn SlavePty + Send>>,
|
|
pub _master: Box<dyn MasterPty + Send>,
|
|
}
|
|
|
|
impl fmt::Debug for PtyHandles {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
f.debug_struct("PtyHandles").finish()
|
|
}
|
|
}
|
|
|
|
/// Handle for driving an interactive process (PTY or pipe).
|
|
pub struct ProcessHandle {
|
|
writer_tx: mpsc::Sender<Vec<u8>>,
|
|
output_tx: broadcast::Sender<Vec<u8>>,
|
|
killer: StdMutex<Option<Box<dyn ChildTerminator>>>,
|
|
reader_handle: StdMutex<Option<JoinHandle<()>>>,
|
|
reader_abort_handles: StdMutex<Vec<AbortHandle>>,
|
|
writer_handle: StdMutex<Option<JoinHandle<()>>>,
|
|
wait_handle: StdMutex<Option<JoinHandle<()>>>,
|
|
exit_status: Arc<AtomicBool>,
|
|
exit_code: Arc<StdMutex<Option<i32>>>,
|
|
// PtyHandles must be preserved because the process will receive Control+C if the
|
|
// slave is closed
|
|
_pty_handles: StdMutex<Option<PtyHandles>>,
|
|
}
|
|
|
|
impl fmt::Debug for ProcessHandle {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
f.debug_struct("ProcessHandle").finish()
|
|
}
|
|
}
|
|
|
|
impl ProcessHandle {
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub(crate) fn new(
|
|
writer_tx: mpsc::Sender<Vec<u8>>,
|
|
output_tx: broadcast::Sender<Vec<u8>>,
|
|
initial_output_rx: broadcast::Receiver<Vec<u8>>,
|
|
killer: Box<dyn ChildTerminator>,
|
|
reader_handle: JoinHandle<()>,
|
|
reader_abort_handles: Vec<AbortHandle>,
|
|
writer_handle: JoinHandle<()>,
|
|
wait_handle: JoinHandle<()>,
|
|
exit_status: Arc<AtomicBool>,
|
|
exit_code: Arc<StdMutex<Option<i32>>>,
|
|
pty_handles: Option<PtyHandles>,
|
|
) -> (Self, broadcast::Receiver<Vec<u8>>) {
|
|
(
|
|
Self {
|
|
writer_tx,
|
|
output_tx,
|
|
killer: StdMutex::new(Some(killer)),
|
|
reader_handle: StdMutex::new(Some(reader_handle)),
|
|
reader_abort_handles: StdMutex::new(reader_abort_handles),
|
|
writer_handle: StdMutex::new(Some(writer_handle)),
|
|
wait_handle: StdMutex::new(Some(wait_handle)),
|
|
exit_status,
|
|
exit_code,
|
|
_pty_handles: StdMutex::new(pty_handles),
|
|
},
|
|
initial_output_rx,
|
|
)
|
|
}
|
|
|
|
/// Returns a channel sender for writing raw bytes to the child stdin.
|
|
pub fn writer_sender(&self) -> mpsc::Sender<Vec<u8>> {
|
|
self.writer_tx.clone()
|
|
}
|
|
|
|
/// Returns a broadcast receiver that yields stdout/stderr chunks.
|
|
pub fn output_receiver(&self) -> broadcast::Receiver<Vec<u8>> {
|
|
self.output_tx.subscribe()
|
|
}
|
|
|
|
/// True if the child process has exited.
|
|
pub fn has_exited(&self) -> bool {
|
|
self.exit_status.load(std::sync::atomic::Ordering::SeqCst)
|
|
}
|
|
|
|
/// Returns the exit code if known.
|
|
pub fn exit_code(&self) -> Option<i32> {
|
|
self.exit_code.lock().ok().and_then(|guard| *guard)
|
|
}
|
|
|
|
/// Attempts to kill the child and abort helper tasks.
|
|
pub fn terminate(&self) {
|
|
if let Ok(mut killer_opt) = self.killer.lock() {
|
|
if let Some(mut killer) = killer_opt.take() {
|
|
let _ = killer.kill();
|
|
}
|
|
}
|
|
|
|
if let Ok(mut h) = self.reader_handle.lock() {
|
|
if let Some(handle) = h.take() {
|
|
handle.abort();
|
|
}
|
|
}
|
|
if let Ok(mut handles) = self.reader_abort_handles.lock() {
|
|
for handle in handles.drain(..) {
|
|
handle.abort();
|
|
}
|
|
}
|
|
if let Ok(mut h) = self.writer_handle.lock() {
|
|
if let Some(handle) = h.take() {
|
|
handle.abort();
|
|
}
|
|
}
|
|
if let Ok(mut h) = self.wait_handle.lock() {
|
|
if let Some(handle) = h.take() {
|
|
handle.abort();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for ProcessHandle {
|
|
fn drop(&mut self) {
|
|
self.terminate();
|
|
}
|
|
}
|
|
|
|
/// Return value from spawn helpers (PTY or pipe).
|
|
#[derive(Debug)]
|
|
pub struct SpawnedProcess {
|
|
pub session: ProcessHandle,
|
|
pub output_rx: broadcast::Receiver<Vec<u8>>,
|
|
pub exit_rx: oneshot::Receiver<i32>,
|
|
}
|