krata: implement fast exit code notification

crates/kratactl/src/cli/console.rs

@@ -2,6 +2,7 @@ use anyhow::Result;
 use clap::Parser;
 use krata::control::control_service_client::ControlServiceClient;
 
+use tokio::select;
 use tonic::transport::Channel;
 
 use crate::{console::StdioConsoleStream, events::EventStream};
@@ -20,10 +21,17 @@ impl ConsoleCommand {
     ) -> Result<()> {
         let input = StdioConsoleStream::stdin_stream(self.guest.clone()).await;
         let output = client.console_data(input).await?.into_inner();
+        let stdout_handle =
+            tokio::task::spawn(async move { StdioConsoleStream::stdout(output).await });
         let exit_hook_task =
             StdioConsoleStream::guest_exit_hook(self.guest.clone(), events).await?;
-        StdioConsoleStream::stdout(output).await?;
-        exit_hook_task.abort();
-        Ok(())
+        let code = select! {
+            x = stdout_handle => {
+                x??;
+                None
+            },
+            x = exit_hook_task => x?
+        };
+        std::process::exit(code.unwrap_or(0));
     }
 }

crates/kratactl/src/cli/launch.rs

@@ -7,6 +7,7 @@ use krata::{
     },
 };
 use log::error;
+use tokio::select;
 use tonic::{transport::Channel, Request};
 
 use crate::{console::StdioConsoleStream, events::EventStream};
@@ -52,17 +53,25 @@ impl LauchCommand {
             .await?
             .into_inner();
         let id = response.guest_id;
-        if self.attach {
+        let code = if self.attach {
             wait_guest_started(&id, events.clone()).await?;
             let input = StdioConsoleStream::stdin_stream(id.clone()).await;
             let output = client.console_data(input).await?.into_inner();
+            let stdout_handle =
+                tokio::task::spawn(async move { StdioConsoleStream::stdout(output).await });
             let exit_hook_task = StdioConsoleStream::guest_exit_hook(id.clone(), events).await?;
-            StdioConsoleStream::stdout(output).await?;
-            exit_hook_task.abort();
+            select! {
+                x = stdout_handle => {
+                    x??;
+                    None
+                },
+                x = exit_hook_task => x?
+            }
         } else {
             println!("created guest: {}", id);
-        }
-        Ok(())
+            None
+        };
+        std::process::exit(code.unwrap_or(0));
     }
 }
 

crates/kratactl/src/console.rs

@@ -9,7 +9,7 @@ use krata::{
     common::GuestStatus,
     control::{watch_events_reply::Event, ConsoleDataReply, ConsoleDataRequest},
 };
-use log::{debug, warn};
+use log::debug;
 use termion::raw::IntoRawMode;
 use tokio::{
     fs::File,
@@ -61,7 +61,10 @@ impl StdioConsoleStream {
         Ok(())
     }
 
-    pub async fn guest_exit_hook(id: String, events: EventStream) -> Result<JoinHandle<()>> {
+    pub async fn guest_exit_hook(
+        id: String,
+        events: EventStream,
+    ) -> Result<JoinHandle<Option<i32>>> {
         Ok(tokio::task::spawn(async move {
             let mut stream = events.subscribe();
             while let Ok(event) = stream.recv().await {
@@ -80,16 +83,17 @@ impl StdioConsoleStream {
                         }
 
                         if let Some(exit_info) = state.exit_info {
-                            std::process::exit(exit_info.code);
+                            return Some(exit_info.code);
                         }
 
-                        if state.status() == GuestStatus::Destroy {
-                            warn!("attached guest was destroyed");
-                            std::process::exit(1);
+                        let status = state.status();
+                        if status == GuestStatus::Destroy || status == GuestStatus::Destroyed {
+                            return Some(10);
                         }
                     }
                 }
             }
+            None
         }))
     }
 }