From 13bea70c0d7c4e37eb6e19eee8914c05ed7f902b Mon Sep 17 00:00:00 2001
From: Alex Zenla <alex@edera.dev>
Date: Wed, 13 Mar 2024 11:34:52 +0000
Subject: [PATCH] krata: implement auto-exit handling

---
 crates/kratactl/bin/control.rs | 13 +++++++--
 crates/kratactl/src/console.rs | 48 ++++++++++++++++++++++++++++++++--
 crates/kratad/src/event.rs     | 29 +++++++++++++++++---
 3 files changed, 83 insertions(+), 7 deletions(-)

diff --git a/crates/kratactl/bin/control.rs b/crates/kratactl/bin/control.rs
index c44e31e..875edc7 100644
--- a/crates/kratactl/bin/control.rs
+++ b/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 { .. } => {
diff --git a/crates/kratactl/src/console.rs b/crates/kratactl/src/console.rs
index 9387b26..6afc265 100644
--- a/crates/kratactl/src/console.rs
+++ b/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;
+                            }
+                        }
+                    }
+                }
+            }
+        }))
+    }
 }
diff --git a/crates/kratad/src/event.rs b/crates/kratad/src/event.rs
index 3ff2796..ceb1bc0 100644
--- a/crates/kratad/src/event.rs
+++ b/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(())
+    }
 }