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;
+                            }
+                        }
+                    }
+                }
+            }
+        }))
+    }
 }

crates/kratad/src/event.rs

@@ -2,7 +2,7 @@ use std::{collections::HashMap, time::Duration};
 
 use anyhow::Result;
 use krata::control::{GuestDestroyedEvent, GuestExitedEvent, GuestLaunchedEvent};
-use log::error;
+use log::{error, info, warn};
 use tokio::{sync::broadcast, task::JoinHandle, time};
 use uuid::Uuid;
 
@@ -52,6 +52,7 @@ impl DaemonEventGenerator {
         };
 
         let mut events: Vec<DaemonEvent> = Vec::new();
+        let mut exits: Vec<GuestExitedEvent> = Vec::new();
 
         for uuid in guests.keys() {
             if !self.last.contains_key(uuid) {
@@ -82,10 +83,13 @@ impl DaemonEventGenerator {
                 continue;
             };
 
-            events.push(DaemonEvent::GuestExited(GuestExitedEvent {
+            let exit = GuestExitedEvent {
                 guest_id: uuid.to_string(),
                 code,
-            }));
+            };
+
+            exits.push(exit.clone());
+            events.push(DaemonEvent::GuestExited(exit));
         }
 
         self.last = guests;
@@ -94,6 +98,8 @@ impl DaemonEventGenerator {
             let _ = self.sender.send(event);
         }
 
+        self.process_exit_auto_destroy(exits).await?;
+
         Ok(())
     }
 
@@ -109,4 +115,21 @@ impl DaemonEventGenerator {
             }
         }))
     }
+
+    async fn process_exit_auto_destroy(&mut self, exits: Vec<GuestExitedEvent>) -> Result<()> {
+        for exit in exits {
+            if let Err(error) = self.runtime.destroy(&exit.guest_id).await {
+                warn!(
+                    "failed to auto-destroy exited guest {}: {}",
+                    exit.guest_id, error
+                );
+            } else {
+                info!(
+                    "auto-destroyed guest {}: exited with status {}",
+                    exit.guest_id, exit.code
+                );
+            }
+        }
+        Ok(())
+    }
 }