From e25cbf087d2ae814e61f66037cfa3a6e72aa2d96 Mon Sep 17 00:00:00 2001 From: Alex Zenla Date: Sat, 23 Mar 2024 09:48:53 +0000 Subject: [PATCH] kratactl: implement guest resolution and rename console to attach --- crates/krata/proto/krata/control.proto | 9 ++++ .../src/cli/{console.rs => attach.rs} | 12 +++-- crates/kratactl/src/cli/destroy.rs | 13 ++---- crates/kratactl/src/cli/list.rs | 4 +- crates/kratactl/src/cli/mod.rs | 46 +++++++++++++++---- crates/kratactl/src/cli/pretty.rs | 28 ----------- crates/kratactl/src/cli/resolve.rs | 28 +++++++++++ crates/kratactl/src/cli/watch.rs | 3 +- crates/kratactl/src/format.rs | 28 +++++++++++ crates/kratad/src/control.rs | 23 +++++++++- 10 files changed, 136 insertions(+), 58 deletions(-) rename crates/kratactl/src/cli/{console.rs => attach.rs} (72%) delete mode 100644 crates/kratactl/src/cli/pretty.rs create mode 100644 crates/kratactl/src/cli/resolve.rs diff --git a/crates/krata/proto/krata/control.proto b/crates/krata/proto/krata/control.proto index 799e861..090b7a3 100644 --- a/crates/krata/proto/krata/control.proto +++ b/crates/krata/proto/krata/control.proto @@ -22,6 +22,14 @@ message ListGuestsReply { repeated krata.common.Guest guests = 1; } +message ResolveGuestRequest { + string name = 1; +} + +message ResolveGuestReply { + krata.common.Guest guest = 1; +} + message DestroyGuestRequest { string guest_id = 1; } @@ -53,6 +61,7 @@ service ControlService { rpc CreateGuest(CreateGuestRequest) returns (CreateGuestReply); rpc DestroyGuest(DestroyGuestRequest) returns (DestroyGuestReply); rpc ListGuests(ListGuestsRequest) returns (ListGuestsReply); + rpc ResolveGuest(ResolveGuestRequest) returns (ResolveGuestReply); rpc ConsoleData(stream ConsoleDataRequest) returns (stream ConsoleDataReply); rpc WatchEvents(WatchEventsRequest) returns (stream WatchEventsReply); } diff --git a/crates/kratactl/src/cli/console.rs b/crates/kratactl/src/cli/attach.rs similarity index 72% rename from crates/kratactl/src/cli/console.rs rename to crates/kratactl/src/cli/attach.rs index 83e45a3..def05e7 100644 --- a/crates/kratactl/src/cli/console.rs +++ b/crates/kratactl/src/cli/attach.rs @@ -7,24 +7,26 @@ use tonic::transport::Channel; use crate::{console::StdioConsoleStream, events::EventStream}; +use super::resolve_guest; + #[derive(Parser)] -pub struct ConsoleCommand { +pub struct AttachCommand { #[arg()] guest: String, } -impl ConsoleCommand { +impl AttachCommand { pub async fn run( self, mut client: ControlServiceClient, events: EventStream, ) -> Result<()> { - let input = StdioConsoleStream::stdin_stream(self.guest.clone()).await; + let guest_id: String = resolve_guest(&mut client, &self.guest).await?; + let input = StdioConsoleStream::stdin_stream(guest_id.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?; + let exit_hook_task = StdioConsoleStream::guest_exit_hook(guest_id.clone(), events).await?; let code = select! { x = stdout_handle => { x??; diff --git a/crates/kratactl/src/cli/destroy.rs b/crates/kratactl/src/cli/destroy.rs index 1ebc445..6b237da 100644 --- a/crates/kratactl/src/cli/destroy.rs +++ b/crates/kratactl/src/cli/destroy.rs @@ -4,7 +4,7 @@ use krata::control::{control_service_client::ControlServiceClient, DestroyGuestR use tonic::{transport::Channel, Request}; -use crate::events::EventStream; +use crate::cli::resolve_guest; #[derive(Parser)] pub struct DestroyCommand { @@ -13,15 +13,10 @@ pub struct DestroyCommand { } impl DestroyCommand { - pub async fn run( - self, - mut client: ControlServiceClient, - _events: EventStream, - ) -> Result<()> { + pub async fn run(self, mut client: ControlServiceClient) -> Result<()> { + let guest_id: String = resolve_guest(&mut client, &self.guest).await?; let _ = client - .destroy_guest(Request::new(DestroyGuestRequest { - guest_id: self.guest.clone(), - })) + .destroy_guest(Request::new(DestroyGuestRequest { guest_id })) .await? .into_inner(); println!("destroyed guest: {}", self.guest); diff --git a/crates/kratactl/src/cli/list.rs b/crates/kratactl/src/cli/list.rs index 74d90ab..1ba2546 100644 --- a/crates/kratactl/src/cli/list.rs +++ b/crates/kratactl/src/cli/list.rs @@ -11,11 +11,9 @@ use tonic::{transport::Channel, Request}; use crate::{ events::EventStream, - format::{kv2line, proto2dynamic, proto2kv}, + format::{guest_state_text, kv2line, proto2dynamic, proto2kv}, }; -use super::pretty::guest_state_text; - #[derive(ValueEnum, Clone, Debug, PartialEq, Eq)] enum ListFormat { CliTable, diff --git a/crates/kratactl/src/cli/mod.rs b/crates/kratactl/src/cli/mod.rs index fc4b1b7..63bf5e6 100644 --- a/crates/kratactl/src/cli/mod.rs +++ b/crates/kratactl/src/cli/mod.rs @@ -1,19 +1,22 @@ -pub mod console; +pub mod attach; pub mod destroy; pub mod launch; pub mod list; -pub mod pretty; +pub mod resolve; pub mod watch; -use anyhow::Result; +use anyhow::{anyhow, Result}; use clap::{Parser, Subcommand}; -use krata::control::WatchEventsRequest; +use krata::control::{ + control_service_client::ControlServiceClient, ResolveGuestRequest, WatchEventsRequest, +}; +use tonic::{transport::Channel, Request}; use crate::{client::ControlClientProvider, events::EventStream}; use self::{ - console::ConsoleCommand, destroy::DestroyCommand, launch::LauchCommand, list::ListCommand, - watch::WatchCommand, + attach::AttachCommand, destroy::DestroyCommand, launch::LauchCommand, list::ListCommand, + resolve::ResolveCommand, watch::WatchCommand, }; #[derive(Parser)] @@ -31,8 +34,9 @@ pub enum Commands { Launch(LauchCommand), Destroy(DestroyCommand), List(ListCommand), - Console(ConsoleCommand), + Attach(AttachCommand), Watch(WatchCommand), + Resolve(ResolveCommand), } impl ControlCommand { @@ -52,11 +56,11 @@ impl ControlCommand { } Commands::Destroy(destroy) => { - destroy.run(client, events).await?; + destroy.run(client).await?; } - Commands::Console(console) => { - console.run(client, events).await?; + Commands::Attach(attach) => { + attach.run(client, events).await?; } Commands::List(list) => { @@ -66,7 +70,29 @@ impl ControlCommand { Commands::Watch(watch) => { watch.run(events).await?; } + + Commands::Resolve(resolve) => { + resolve.run(client).await?; + } } Ok(()) } } + +pub async fn resolve_guest( + client: &mut ControlServiceClient, + name: &str, +) -> Result { + let reply = client + .resolve_guest(Request::new(ResolveGuestRequest { + name: name.to_string(), + })) + .await? + .into_inner(); + + if let Some(guest) = reply.guest { + Ok(guest.id) + } else { + Err(anyhow!("unable to resolve guest {}", name)) + } +} diff --git a/crates/kratactl/src/cli/pretty.rs b/crates/kratactl/src/cli/pretty.rs deleted file mode 100644 index 3e8b39b..0000000 --- a/crates/kratactl/src/cli/pretty.rs +++ /dev/null @@ -1,28 +0,0 @@ -use krata::common::{GuestState, GuestStatus}; - -pub fn guest_status_text(status: GuestStatus) -> String { - match status { - GuestStatus::Starting => "starting", - GuestStatus::Started => "started", - GuestStatus::Destroying => "destroying", - GuestStatus::Destroyed => "destroyed", - GuestStatus::Exited => "exited", - GuestStatus::Failed => "failed", - _ => "unknown", - } - .to_string() -} - -pub fn guest_state_text(state: Option<&GuestState>) -> String { - let state = state.cloned().unwrap_or_default(); - let mut text = guest_status_text(state.status()); - - if let Some(exit) = state.exit_info { - text.push_str(&format!(" (exit code: {})", exit.code)); - } - - if let Some(error) = state.error_info { - text.push_str(&format!(" (error: {})", error.message)); - } - text -} diff --git a/crates/kratactl/src/cli/resolve.rs b/crates/kratactl/src/cli/resolve.rs new file mode 100644 index 0000000..f39d8f8 --- /dev/null +++ b/crates/kratactl/src/cli/resolve.rs @@ -0,0 +1,28 @@ +use anyhow::Result; +use clap::Parser; +use krata::control::{control_service_client::ControlServiceClient, ResolveGuestRequest}; + +use tonic::{transport::Channel, Request}; + +#[derive(Parser)] +pub struct ResolveCommand { + #[arg()] + guest: String, +} + +impl ResolveCommand { + pub async fn run(self, mut client: ControlServiceClient) -> Result<()> { + let reply = client + .resolve_guest(Request::new(ResolveGuestRequest { + name: self.guest.clone(), + })) + .await? + .into_inner(); + if let Some(guest) = reply.guest { + println!("{}", guest.id); + } else { + std::process::exit(1); + } + Ok(()) + } +} diff --git a/crates/kratactl/src/cli/watch.rs b/crates/kratactl/src/cli/watch.rs index 5e9b04c..e5029bd 100644 --- a/crates/kratactl/src/cli/watch.rs +++ b/crates/kratactl/src/cli/watch.rs @@ -5,9 +5,8 @@ use prost_reflect::ReflectMessage; use serde_json::Value; use crate::{ - cli::pretty::guest_state_text, events::EventStream, - format::{kv2line, proto2dynamic, proto2kv}, + format::{guest_state_text, kv2line, proto2dynamic, proto2kv}, }; #[derive(ValueEnum, Clone, Debug, PartialEq, Eq)] diff --git a/crates/kratactl/src/format.rs b/crates/kratactl/src/format.rs index e01915a..eeb087e 100644 --- a/crates/kratactl/src/format.rs +++ b/crates/kratactl/src/format.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use anyhow::Result; +use krata::common::{GuestState, GuestStatus}; use prost_reflect::{DynamicMessage, ReflectMessage, Value}; pub fn proto2dynamic(proto: impl ReflectMessage) -> Result { @@ -56,3 +57,30 @@ pub fn kv2line(map: HashMap) -> String { .collect::>() .join(" ") } + +pub fn guest_status_text(status: GuestStatus) -> String { + match status { + GuestStatus::Starting => "starting", + GuestStatus::Started => "started", + GuestStatus::Destroying => "destroying", + GuestStatus::Destroyed => "destroyed", + GuestStatus::Exited => "exited", + GuestStatus::Failed => "failed", + _ => "unknown", + } + .to_string() +} + +pub fn guest_state_text(state: Option<&GuestState>) -> String { + let state = state.cloned().unwrap_or_default(); + let mut text = guest_status_text(state.status()); + + if let Some(exit) = state.exit_info { + text.push_str(&format!(" (exit code: {})", exit.code)); + } + + if let Some(error) = state.error_info { + text.push_str(&format!(" (error: {})", error.message)); + } + text +} diff --git a/crates/kratad/src/control.rs b/crates/kratad/src/control.rs index 2e1cb20..62739ca 100644 --- a/crates/kratad/src/control.rs +++ b/crates/kratad/src/control.rs @@ -7,7 +7,8 @@ use krata::{ control::{ control_service_server::ControlService, ConsoleDataReply, ConsoleDataRequest, CreateGuestReply, CreateGuestRequest, DestroyGuestReply, DestroyGuestRequest, - ListGuestsReply, ListGuestsRequest, WatchEventsReply, WatchEventsRequest, + ListGuestsReply, ListGuestsRequest, ResolveGuestReply, ResolveGuestRequest, + WatchEventsReply, WatchEventsRequest, }, }; use kratart::Runtime; @@ -179,6 +180,26 @@ impl ControlService for RuntimeControlService { Ok(Response::new(ListGuestsReply { guests })) } + async fn resolve_guest( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let guests = self.guests.list().await.map_err(ApiError::from)?; + let guests = guests + .into_values() + .filter_map(|entry| entry.guest) + .filter(|x| { + let comparison_spec = x.spec.as_ref().cloned().unwrap_or_default(); + (!request.name.is_empty() && comparison_spec.name == request.name) + || x.id == request.name + }) + .collect::>(); + Ok(Response::new(ResolveGuestReply { + guest: guests.first().cloned(), + })) + } + async fn console_data( &self, request: Request>,