krata: implement auto-exit handling

crates/kratactl/bin/control.rs

@@ -52,6 +52,10 @@ async fn main() -> Result<()> {
 
     let args = ControllerArgs::parse();
     let mut client = ControlClientProvider::dial(args.connection.parse()?).await?;
+    let events = client
+        .watch_events(WatchEventsRequest {})
+        .await?
+        .into_inner();
 
     match args.command {
         Commands::Launch {
@@ -82,9 +86,12 @@ async fn main() -> Result<()> {
             };
             println!("launched guest: {}", guest.id);
             if attach {
-                let input = StdioConsoleStream::stdin_stream(guest.id).await;
+                let input = StdioConsoleStream::stdin_stream(guest.id.clone()).await;
                 let output = client.console_data(input).await?.into_inner();
+                let exit_hook_task =
+                    StdioConsoleStream::guest_exit_hook(guest.id.clone(), events).await?;
                 StdioConsoleStream::stdout(output).await?;
+                exit_hook_task.abort();
             }
         }
 
@@ -99,9 +106,11 @@ async fn main() -> Result<()> {
         }
 
         Commands::Console { guest } => {
-            let input = StdioConsoleStream::stdin_stream(guest).await;
+            let input = StdioConsoleStream::stdin_stream(guest.clone()).await;
             let output = client.console_data(input).await?.into_inner();
+            let exit_hook_task = StdioConsoleStream::guest_exit_hook(guest.clone(), events).await?;
             StdioConsoleStream::stdout(output).await?;
+            exit_hook_task.abort();
         }
 
         Commands::List { .. } => {

crates/kratactl/src/console.rs

@@ -5,12 +5,15 @@ use std::{
 
 use anyhow::Result;
 use async_stream::stream;
-use krata::control::{ConsoleDataReply, ConsoleDataRequest};
-use log::debug;
+use krata::control::{
+    watch_events_reply::Event, ConsoleDataReply, ConsoleDataRequest, WatchEventsReply,
+};
+use log::{debug, error, warn};
 use termion::raw::IntoRawMode;
 use tokio::{
     fs::File,
     io::{stdin, AsyncReadExt, AsyncWriteExt},
+    task::JoinHandle,
 };
 use tokio_stream::{Stream, StreamExt};
 use tonic::Streaming;
@@ -54,4 +57,45 @@ impl StdioConsoleStream {
         }
         Ok(())
     }
+
+    pub async fn guest_exit_hook(
+        id: String,
+        mut events: Streaming<WatchEventsReply>,
+    ) -> Result<JoinHandle<()>> {
+        Ok(tokio::task::spawn(async move {
+            while let Some(result) = events.next().await {
+                match result {
+                    Err(error) => {
+                        error!("failed to handle events for exit hook: {}", error);
+                        break;
+                    }
+
+                    Ok(reply) => {
+                        let Some(event) = reply.event else {
+                            continue;
+                        };
+
+                        match event {
+                            Event::GuestExited(exit) => {
+                                if exit.guest_id == id {
+                                    std::process::exit(exit.code);
+                                }
+                            }
+
+                            Event::GuestDestroyed(destroy) => {
+                                if destroy.guest_id == id {
+                                    warn!("attached guest destroyed");
+                                    std::process::exit(1);
+                                }
+                            }
+
+                            _ => {
+                                continue;
+                            }
+                        }
+                    }
+                }
+            }
+        }))
+    }
 }