krata: improvements to event handling during reconciliation

crates/kratactl/src/console.rs

@@ -7,9 +7,9 @@ use anyhow::Result;
 use async_stream::stream;
 use krata::{
     common::GuestStatus,
-    control::{watch_events_reply::Event, ConsoleDataReply, ConsoleDataRequest, WatchEventsReply},
+    control::{watch_events_reply::Event, ConsoleDataReply, ConsoleDataRequest},
 };
-use log::{debug, error, warn};
+use log::{debug, warn};
 use termion::raw::IntoRawMode;
 use tokio::{
     fs::File,
@@ -19,6 +19,8 @@ use tokio::{
 use tokio_stream::{Stream, StreamExt};
 use tonic::Streaming;
 
+use crate::events::EventStream;
+
 pub struct StdioConsoleStream;
 
 impl StdioConsoleStream {
@@ -59,46 +61,31 @@ impl StdioConsoleStream {
         Ok(())
     }
 
-    pub async fn guest_exit_hook(
-        id: String,
-        mut events: Streaming<WatchEventsReply>,
-    ) -> Result<JoinHandle<()>> {
+    pub async fn guest_exit_hook(id: String, events: EventStream) -> 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 {
+            let mut stream = events.subscribe();
+            while let Ok(event) = stream.recv().await {
+                match event {
+                    Event::GuestChanged(changed) => {
+                        let Some(guest) = changed.guest else {
                             continue;
                         };
 
-                        match event {
-                            Event::GuestChanged(changed) => {
-                                let Some(guest) = changed.guest else {
-                                    continue;
-                                };
+                        let Some(state) = guest.state else {
+                            continue;
+                        };
 
-                                let Some(state) = guest.state else {
-                                    continue;
-                                };
+                        if guest.id != id {
+                            continue;
+                        }
 
-                                if guest.id != id {
-                                    continue;
-                                }
+                        if let Some(exit_info) = state.exit_info {
+                            std::process::exit(exit_info.code);
+                        }
 
-                                if let Some(exit_info) = state.exit_info {
-                                    std::process::exit(exit_info.code);
-                                }
-
-                                if state.status() == GuestStatus::Destroyed {
-                                    warn!("attached guest was destroyed");
-                                    std::process::exit(1);
-                                }
-                            }
+                        if state.status() == GuestStatus::Destroy {
+                            warn!("attached guest was destroyed");
+                            std::process::exit(1);
                         }
                     }
                 }

crates/kratactl/src/events.rs

@@ -0,0 +1,57 @@
+use std::sync::Arc;
+
+use anyhow::Result;
+use krata::control::{watch_events_reply::Event, WatchEventsReply};
+use log::trace;
+use tokio::{sync::broadcast, task::JoinHandle};
+use tokio_stream::StreamExt;
+use tonic::Streaming;
+
+#[derive(Clone)]
+pub struct EventStream {
+    sender: Arc<broadcast::Sender<Event>>,
+    task: Arc<JoinHandle<()>>,
+}
+
+impl EventStream {
+    pub async fn open(mut events: Streaming<WatchEventsReply>) -> Result<Self> {
+        let (sender, _) = broadcast::channel(1000);
+        let emit = sender.clone();
+        let task = tokio::task::spawn(async move {
+            loop {
+                let Some(result) = events.next().await else {
+                    break;
+                };
+
+                let reply = match result {
+                    Ok(reply) => reply,
+                    Err(error) => {
+                        trace!("event stream processing failed: {}", error);
+                        break;
+                    }
+                };
+
+                let Some(event) = reply.event else {
+                    continue;
+                };
+                let _ = emit.send(event);
+            }
+        });
+        Ok(Self {
+            sender: Arc::new(sender),
+            task: Arc::new(task),
+        })
+    }
+
+    pub fn subscribe(&self) -> broadcast::Receiver<Event> {
+        self.sender.subscribe()
+    }
+}
+
+impl Drop for EventStream {
+    fn drop(&mut self) {
+        if Arc::strong_count(&self.task) <= 1 {
+            self.task.abort();
+        }
+    }
+}

crates/kratactl/src/lib.rs

@@ -1,2 +1,3 @@
 pub mod client;
 pub mod console;
+pub mod events;