Allow interacting with stdin

compiler/base/orchestrator/src/lib.rs

@@ -4,7 +4,7 @@ pub mod coordinator;
 mod message;
 pub mod worker;
 
-trait DropErrorDetailsExt<T> {
+pub trait DropErrorDetailsExt<T> {
     fn drop_error_details(self) -> Result<T, tokio::sync::mpsc::error::SendError<()>>;
 }
 

compiler/base/orchestrator/src/message.rs

@@ -26,6 +26,8 @@ pub enum CoordinatorMessage {
     ReadFile(ReadFileRequest),
     ExecuteCommand(ExecuteCommandRequest),
     StdinPacket(String),
+    StdinClose,
+    Kill,
 }
 
 impl_narrow_to_broad!(

compiler/base/orchestrator/src/worker.rs

@@ -46,6 +46,7 @@ use tokio::{
     sync::mpsc,
     task::JoinSet,
 };
+use tokio_util::sync::CancellationToken;
 
 use crate::{
     bincode_input_closed,
@@ -57,25 +58,21 @@ use crate::{
     DropErrorDetailsExt,
 };
 
-type CommandRequest = (Multiplexed<ExecuteCommandRequest>, MultiplexingSender);
-
 pub async fn listen(project_dir: impl Into<PathBuf>) -> Result<(), Error> {
     let project_dir = project_dir.into();
 
     let (coordinator_msg_tx, coordinator_msg_rx) = mpsc::channel(8);
     let (worker_msg_tx, worker_msg_rx) = mpsc::channel(8);
     let mut io_tasks = spawn_io_queue(coordinator_msg_tx, worker_msg_rx);
 
-    let (cmd_tx, cmd_rx) = mpsc::channel(8);
-    let (stdin_tx, stdin_rx) = mpsc::channel(8);
-    let process_task = tokio::spawn(manage_processes(stdin_rx, cmd_rx, project_dir.clone()));
+    let (process_tx, process_rx) = mpsc::channel(8);
+    let process_task = tokio::spawn(manage_processes(process_rx, project_dir.clone()));
 
     let handler_task = tokio::spawn(handle_coordinator_message(
         coordinator_msg_rx,
         worker_msg_tx,
         project_dir,
-        cmd_tx,
-        stdin_tx,
+        process_tx,
     ));
 
     select! {
@@ -122,8 +119,7 @@ async fn handle_coordinator_message(
     mut coordinator_msg_rx: mpsc::Receiver<Multiplexed<CoordinatorMessage>>,
     worker_msg_tx: mpsc::Sender<Multiplexed<WorkerMessage>>,
     project_dir: PathBuf,
-    cmd_tx: mpsc::Sender<CommandRequest>,
-    stdin_tx: mpsc::Sender<Multiplexed<String>>,
+    process_tx: mpsc::Sender<Multiplexed<ProcessCommand>>,
 ) -> Result<(), HandleCoordinatorMessageError> {
     use handle_coordinator_message_error::*;
 
@@ -177,20 +173,36 @@ async fn handle_coordinator_message(
                     }
 
                     CoordinatorMessage::ExecuteCommand(req) => {
-                        cmd_tx
-                            .send((Multiplexed(job_id, req), worker_msg_tx()))
+                        process_tx
+                            .send(Multiplexed(job_id, ProcessCommand::Start(req, worker_msg_tx())))
                             .await
                             .drop_error_details()
                             .context(UnableToSendCommandExecutionRequestSnafu)?;
                     }
 
                     CoordinatorMessage::StdinPacket(data) => {
-                        stdin_tx
-                            .send(Multiplexed(job_id, data))
+                        process_tx
+                            .send(Multiplexed(job_id, ProcessCommand::Stdin(data)))
                             .await
                             .drop_error_details()
                             .context(UnableToSendStdinPacketSnafu)?;
                     }
+
+                    CoordinatorMessage::StdinClose => {
+                        process_tx
+                            .send(Multiplexed(job_id, ProcessCommand::StdinClose))
+                            .await
+                            .drop_error_details()
+                            .context(UnableToSendStdinCloseSnafu)?;
+                    }
+
+                    CoordinatorMessage::Kill => {
+                        process_tx
+                        .send(Multiplexed(job_id, ProcessCommand::Kill))
+                        .await
+                        .drop_error_details()
+                        .context(UnableToSendKillSnafu)?;
+                    }
                 }
             }
 
@@ -221,6 +233,12 @@ pub enum HandleCoordinatorMessageError {
     #[snafu(display("Failed to send stdin packet to the command task"))]
     UnableToSendStdinPacket { source: mpsc::error::SendError<()> },
 
+    #[snafu(display("Failed to send stdin close request to the command task"))]
+    UnableToSendStdinClose { source: mpsc::error::SendError<()> },
+
+    #[snafu(display("Failed to send kill request to the command task"))]
+    UnableToSendKill { source: mpsc::error::SendError<()> },
+
     #[snafu(display("A coordinator command handler background task panicked"))]
     TaskPanicked { source: tokio::task::JoinError },
 }
@@ -373,63 +391,144 @@ fn parse_working_dir(cwd: Option<String>, project_path: impl Into<PathBuf>) -> P
     final_path
 }
 
+enum ProcessCommand {
+    Start(ExecuteCommandRequest, MultiplexingSender),
+    Stdin(String),
+    StdinClose,
+    Kill,
+}
+
+struct ProcessState {
+    project_path: PathBuf,
+    processes: JoinSet<Result<(), ProcessError>>,
+    stdin_senders: HashMap<JobId, mpsc::Sender<String>>,
+    stdin_shutdown_tx: mpsc::Sender<JobId>,
+    kill_tokens: HashMap<JobId, CancellationToken>,
+}
+
+impl ProcessState {
+    fn new(project_path: PathBuf, stdin_shutdown_tx: mpsc::Sender<JobId>) -> Self {
+        Self {
+            project_path,
+            processes: Default::default(),
+            stdin_senders: Default::default(),
+            stdin_shutdown_tx,
+            kill_tokens: Default::default(),
+        }
+    }
+
+    async fn start(
+        &mut self,
+        job_id: JobId,
+        req: ExecuteCommandRequest,
+        worker_msg_tx: MultiplexingSender,
+    ) -> Result<(), ProcessError> {
+        use process_error::*;
+
+        let token = CancellationToken::new();
+
+        let RunningChild {
+            child,
+            stdin_rx,
+            stdin,
+            stdout,
+            stderr,
+        } = match process_begin(req, &self.project_path, &mut self.stdin_senders, job_id) {
+            Ok(v) => v,
+            Err(e) => {
+                // Should we add a message for process started
+                // in addition to the current message which
+                // indicates that the process has ended?
+                worker_msg_tx
+                    .send_err(e)
+                    .await
+                    .context(UnableToSendExecuteCommandStartedResponseSnafu)?;
+                return Ok(());
+            }
+        };
+
+        let task_set = stream_stdio(worker_msg_tx.clone(), stdin_rx, stdin, stdout, stderr);
+
+        self.kill_tokens.insert(job_id, token.clone());
+
+        self.processes.spawn({
+            let stdin_shutdown_tx = self.stdin_shutdown_tx.clone();
+            async move {
+                worker_msg_tx
+                    .send(process_end(token, child, task_set, stdin_shutdown_tx, job_id).await)
+                    .await
+                    .context(UnableToSendExecuteCommandResponseSnafu)
+            }
+        });
+
+        Ok(())
+    }
+
+    async fn stdin(&mut self, job_id: JobId, packet: String) -> Result<(), ProcessError> {
+        use process_error::*;
+
+        if let Some(stdin_tx) = self.stdin_senders.get(&job_id) {
+            stdin_tx
+                .send(packet)
+                .await
+                .drop_error_details()
+                .context(UnableToSendStdinDataSnafu)?;
+        }
+
+        Ok(())
+    }
+
+    fn stdin_close(&mut self, job_id: JobId) {
+        self.stdin_senders.remove(&job_id);
+        // Should we care if we remove a sender that's already removed?
+    }
+
+    async fn join_process(&mut self) -> Option<Result<(), ProcessError>> {
+        use process_error::*;
+
+        let process = self.processes.join_next().await?;
+        Some(process.context(ProcessTaskPanickedSnafu).and_then(|e| e))
+    }
+
+    fn kill(&mut self, job_id: JobId) {
+        if let Some(token) = self.kill_tokens.get(&job_id) {
+            token.cancel();
+        }
+    }
+}
+
 async fn manage_processes(
-    mut stdin_rx: mpsc::Receiver<Multiplexed<String>>,
-    mut cmd_rx: mpsc::Receiver<CommandRequest>,
+    mut rx: mpsc::Receiver<Multiplexed<ProcessCommand>>,
     project_path: PathBuf,
 ) -> Result<(), ProcessError> {
     use process_error::*;
 
-    let mut processes = JoinSet::new();
-    let mut stdin_senders = HashMap::new();
     let (stdin_shutdown_tx, mut stdin_shutdown_rx) = mpsc::channel(8);
+    let mut state = ProcessState::new(project_path, stdin_shutdown_tx);
 
     loop {
         select! {
-            cmd_req = cmd_rx.recv() => {
-                let Some((Multiplexed(job_id, req), worker_msg_tx)) = cmd_req else { break };
-
-                let RunningChild { child, stdin_rx, stdin, stdout, stderr } = match process_begin(req, &project_path, &mut stdin_senders, job_id) {
-                    Ok(v) => v,
-                    Err(e) => {
-                        // Should we add a message for process started
-                        // in addition to the current message which
-                        // indicates that the process has ended?
-                        worker_msg_tx.send_err(e).await.context(UnableToSendExecuteCommandStartedResponseSnafu)?;
-                        continue;
-                    }
-                };
+            cmd = rx.recv() => {
+                let Some(Multiplexed(job_id, cmd)) = cmd else { break };
 
-                let task_set = stream_stdio(worker_msg_tx.clone(), stdin_rx, stdin, stdout, stderr);
+                match cmd {
+                    ProcessCommand::Start(req, worker_msg_tx) => state.start(job_id, req, worker_msg_tx).await?,
 
-                processes.spawn({
-                    let stdin_shutdown_tx = stdin_shutdown_tx.clone();
-                    async move {
-                        worker_msg_tx
-                            .send(process_end(child, task_set, stdin_shutdown_tx, job_id).await)
-                            .await
-                            .context(UnableToSendExecuteCommandResponseSnafu)
-                    }
-                });
-            }
+                    ProcessCommand::Stdin(packet) => state.stdin(job_id, packet).await?,
 
-            stdin_packet = stdin_rx.recv() => {
-                // Dispatch stdin packet to different child by attached command id.
-                let Some(Multiplexed(job_id, packet)) = stdin_packet else { break };
+                    ProcessCommand::StdinClose => state.stdin_close(job_id),
 
-                if let Some(stdin_tx) = stdin_senders.get(&job_id) {
-                    stdin_tx.send(packet).await.drop_error_details().context(UnableToSendStdinDataSnafu)?;
+                    ProcessCommand::Kill => state.kill(job_id),
                 }
             }
 
             job_id = stdin_shutdown_rx.recv() => {
                 let job_id = job_id.context(StdinShutdownReceiverEndedSnafu)?;
-                stdin_senders.remove(&job_id);
-                // Should we care if we remove a sender that's already removed?
+                state.stdin_close(job_id);
             }
 
-            Some(process) = processes.join_next() => {
-                process.context(ProcessTaskPanickedSnafu)??;
+            Some(process) = state.join_process() => {
+                process?;
             }
         }
     }
@@ -488,13 +587,19 @@ fn process_begin(
 }
 
 async fn process_end(
+    token: CancellationToken,
     mut child: Child,
     mut task_set: JoinSet<Result<(), StdioError>>,
     stdin_shutdown_tx: mpsc::Sender<JobId>,
     job_id: JobId,
 ) -> Result<ExecuteCommandResponse, ProcessError> {
     use process_error::*;
 
+    select! {
+        () = token.cancelled() => child.kill().await.context(KillChildSnafu)?,
+        _ = child.wait() => {},
+    };
+
     let status = child.wait().await.context(WaitChildSnafu)?;
 
     stdin_shutdown_tx
@@ -634,6 +739,9 @@ pub enum ProcessError {
     #[snafu(display("Failed to send stdin data"))]
     UnableToSendStdinData { source: mpsc::error::SendError<()> },
 
+    #[snafu(display("Failed to kill the child process"))]
+    KillChild { source: std::io::Error },
+
     #[snafu(display("Failed to wait for child process exiting"))]
     WaitChild { source: std::io::Error },
 
@@ -671,10 +779,7 @@ fn stream_stdio(
     let mut set = JoinSet::new();
 
     set.spawn(async move {
-        loop {
-            let Some(data) = stdin_rx.recv().await else {
-                break;
-            };
+        while let Some(data) = stdin_rx.recv().await {
             stdin
                 .write_all(data.as_bytes())
                 .await