krata: implement fast exit code notification

This commit is contained in:
Alex Zenla
2024-03-15 17:36:26 +00:00
parent e3aa54edea
commit 48123166fe
5 changed files with 189 additions and 66 deletions

View File

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