fix(idm): process all idm messages in the same frame and use childwait exit notification for exec (fixes #290) (#302)

crates/zone/src/background.rs

@@ -39,6 +39,7 @@ impl ZoneBackground {
         let mut event_subscription = self.idm.subscribe().await?;
         let mut requests_subscription = self.idm.requests().await?;
         let mut request_streams_subscription = self.idm.request_streams().await?;
+        let mut wait_subscription = self.wait.subscribe().await?;
         loop {
             select! {
                 x = event_subscription.recv() => match x {
@@ -85,9 +86,9 @@ impl ZoneBackground {
                     }
                 },
 
-                event = self.wait.recv() => match event {
-                    Some(event) => self.child_event(event).await?,
-                    None => {
+                event = wait_subscription.recv() => match event {
+                    Ok(event) => self.child_event(event).await?,
+                    Err(_) => {
                         break;
                     }
                 }
@@ -128,9 +129,10 @@ impl ZoneBackground {
         &mut self,
         handle: IdmClientStreamResponseHandle<Request>,
     ) -> Result<()> {
+        let wait = self.wait.clone();
         if let Some(RequestType::ExecStream(_)) = &handle.initial.request {
             tokio::task::spawn(async move {
-                let exec = ZoneExecTask { handle };
+                let exec = ZoneExecTask { wait, handle };
                 if let Err(error) = exec.run().await {
                     let _ = exec
                         .handle

crates/zone/src/childwait.rs

@@ -11,7 +11,7 @@ use anyhow::Result;
 use libc::{c_int, waitpid, WEXITSTATUS, WIFEXITED};
 use log::warn;
 use nix::unistd::Pid;
-use tokio::sync::mpsc::{channel, Receiver, Sender};
+use tokio::sync::broadcast::{channel, Receiver, Sender};
 
 const CHILD_WAIT_QUEUE_LEN: usize = 10;
 
@@ -21,18 +21,19 @@ pub struct ChildEvent {
     pub status: c_int,
 }
 
+#[derive(Clone)]
 pub struct ChildWait {
-    receiver: Receiver<ChildEvent>,
+    sender: Sender<ChildEvent>,
     signal: Arc<AtomicBool>,
-    _task: JoinHandle<()>,
+    _task: Arc<JoinHandle<()>>,
 }
 
 impl ChildWait {
     pub fn new() -> Result<ChildWait> {
-        let (sender, receiver) = channel(CHILD_WAIT_QUEUE_LEN);
+        let (sender, _) = channel(CHILD_WAIT_QUEUE_LEN);
         let signal = Arc::new(AtomicBool::new(false));
         let mut processor = ChildWaitTask {
-            sender,
+            sender: sender.clone(),
             signal: signal.clone(),
         };
         let task = thread::spawn(move || {
@@ -41,14 +42,14 @@ impl ChildWait {
             }
         });
         Ok(ChildWait {
-            receiver,
+            sender,
             signal,
-            _task: task,
+            _task: Arc::new(task),
         })
     }
 
-    pub async fn recv(&mut self) -> Option<ChildEvent> {
-        self.receiver.recv().await
+    pub async fn subscribe(&self) -> Result<Receiver<ChildEvent>> {
+        Ok(self.sender.subscribe())
     }
 }
 
@@ -68,7 +69,7 @@ impl ChildWaitTask {
                     pid: Pid::from_raw(pid),
                     status: WEXITSTATUS(status),
                 };
-                let _ = self.sender.try_send(event);
+                let _ = self.sender.send(event);
 
                 if self.signal.load(Ordering::Acquire) {
                     return Ok(());
@@ -80,6 +81,8 @@ impl ChildWaitTask {
 
 impl Drop for ChildWait {
     fn drop(&mut self) {
-        self.signal.store(true, Ordering::Release);
+        if Arc::strong_count(&self.signal) <= 1 {
+            self.signal.store(true, Ordering::Release);
+        }
     }
 }

crates/zone/src/exec.rs

@@ -1,6 +1,12 @@
 use std::{collections::HashMap, process::Stdio};
 
 use anyhow::{anyhow, Result};
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    join,
+    process::Command,
+};
+
 use krata::idm::{
     client::IdmClientStreamResponseHandle,
     internal::{
@@ -9,13 +15,11 @@ use krata::idm::{
     },
     internal::{response::Response as ResponseType, Request, Response},
 };
-use tokio::{
-    io::{AsyncReadExt, AsyncWriteExt},
-    join,
-    process::Command,
-};
+
+use crate::childwait::ChildWait;
 
 pub struct ZoneExecTask {
+    pub wait: ChildWait,
     pub handle: IdmClientStreamResponseHandle<Request>,
 }
 
@@ -58,6 +62,7 @@ impl ZoneExecTask {
             start.working_directory.clone()
         };
 
+        let mut wait_subscription = self.wait.subscribe().await?;
         let mut child = Command::new(exe)
             .args(cmd)
             .envs(env)
@@ -69,6 +74,7 @@ impl ZoneExecTask {
             .spawn()
             .map_err(|error| anyhow!("failed to spawn: {}", error))?;
 
+        let pid = child.id().ok_or_else(|| anyhow!("pid is not provided"))?;
         let mut stdin = child
             .stdin
             .take()
@@ -150,12 +156,19 @@ impl ZoneExecTask {
             }
         });
 
-        let exit = child.wait().await?;
-        let code = exit.code().unwrap_or(-1);
-
-        let _ = join!(stdout_task, stderr_task);
-        stdin_task.abort();
+        let data_task = tokio::task::spawn(async move {
+            let _ = join!(stdout_task, stderr_task);
+            stdin_task.abort();
+        });
 
+        let code = loop {
+            if let Ok(event) = wait_subscription.recv().await {
+                if event.pid.as_raw() as u32 == pid {
+                    break event.status;
+                }
+            }
+        };
+        data_task.await?;
         let response = Response {
             response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
                 exited: true,