krata: restructure packages for cleanliness

This commit is contained in:
Alex Zenla
2024-03-30 06:17:30 +00:00
parent da9e6cac14
commit bdb91a6cb3
85 changed files with 35 additions and 1345 deletions

View File

@ -0,0 +1,40 @@
use anyhow::Result;
use clap::Parser;
use krata::{events::EventStream, v1::control::control_service_client::ControlServiceClient};
use tokio::select;
use tonic::transport::Channel;
use crate::console::StdioConsoleStream;
use super::resolve_guest;
#[derive(Parser)]
pub struct AttachCommand {
#[arg()]
guest: String,
}
impl AttachCommand {
pub async fn run(
self,
mut client: ControlServiceClient<Channel>,
events: EventStream,
) -> Result<()> {
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(guest_id.clone(), events).await?;
let code = select! {
x = stdout_handle => {
x??;
None
},
x = exit_hook_task => x?
};
StdioConsoleStream::restore_terminal_mode();
std::process::exit(code.unwrap_or(0));
}
}

View File

@ -0,0 +1,80 @@
use anyhow::Result;
use clap::Parser;
use krata::{
events::EventStream,
v1::{
common::GuestStatus,
control::{
control_service_client::ControlServiceClient, watch_events_reply::Event,
DestroyGuestRequest,
},
},
};
use log::error;
use tonic::{transport::Channel, Request};
use crate::cli::resolve_guest;
#[derive(Parser)]
pub struct DestroyCommand {
#[arg(short = 'W', long)]
wait: bool,
#[arg()]
guest: String,
}
impl DestroyCommand {
pub async fn run(
self,
mut client: ControlServiceClient<Channel>,
events: EventStream,
) -> Result<()> {
let guest_id: String = resolve_guest(&mut client, &self.guest).await?;
let _ = client
.destroy_guest(Request::new(DestroyGuestRequest {
guest_id: guest_id.clone(),
}))
.await?
.into_inner();
if self.wait {
wait_guest_destroyed(&guest_id, events).await?;
}
Ok(())
}
}
async fn wait_guest_destroyed(id: &str, events: EventStream) -> Result<()> {
let mut stream = events.subscribe();
while let Ok(event) = stream.recv().await {
match event {
Event::GuestChanged(changed) => {
let Some(guest) = changed.guest else {
continue;
};
if guest.id != id {
continue;
}
let Some(state) = guest.state else {
continue;
};
if let Some(ref error) = state.error_info {
if state.status() == GuestStatus::Failed {
error!("destroy failed: {}", error.message);
std::process::exit(1);
} else {
error!("guest error: {}", error.message);
}
}
if state.status() == GuestStatus::Destroyed {
std::process::exit(0);
}
}
}
}
Ok(())
}

View File

@ -0,0 +1,151 @@
use std::collections::HashMap;
use anyhow::Result;
use clap::Parser;
use krata::{
events::EventStream,
v1::{
common::{
guest_image_spec::Image, GuestImageSpec, GuestOciImageSpec, GuestSpec, GuestStatus,
GuestTaskSpec, GuestTaskSpecEnvVar,
},
control::{
control_service_client::ControlServiceClient, watch_events_reply::Event,
CreateGuestRequest,
},
},
};
use log::error;
use tokio::select;
use tonic::{transport::Channel, Request};
use crate::console::StdioConsoleStream;
#[derive(Parser)]
pub struct LauchCommand {
#[arg(short, long)]
name: Option<String>,
#[arg(short, long, default_value_t = 1)]
cpus: u32,
#[arg(short, long, default_value_t = 512)]
mem: u64,
#[arg[short, long]]
env: Option<Vec<String>>,
#[arg(short, long)]
attach: bool,
#[arg(short = 'W', long)]
wait: bool,
#[arg()]
oci: String,
#[arg(allow_hyphen_values = true, trailing_var_arg = true)]
run: Vec<String>,
}
impl LauchCommand {
pub async fn run(
self,
mut client: ControlServiceClient<Channel>,
events: EventStream,
) -> Result<()> {
let request = CreateGuestRequest {
spec: Some(GuestSpec {
name: self.name.unwrap_or_default(),
image: Some(GuestImageSpec {
image: Some(Image::Oci(GuestOciImageSpec { image: self.oci })),
}),
vcpus: self.cpus,
mem: self.mem,
task: Some(GuestTaskSpec {
environment: env_map(&self.env.unwrap_or_default())
.iter()
.map(|(key, value)| GuestTaskSpecEnvVar {
key: key.clone(),
value: value.clone(),
})
.collect(),
command: self.run,
}),
annotations: vec![],
}),
};
let response = client
.create_guest(Request::new(request))
.await?
.into_inner();
let id = response.guest_id;
if self.wait || self.attach {
wait_guest_started(&id, events.clone()).await?;
}
let code = if self.attach {
let input = StdioConsoleStream::stdin_stream(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(id.clone(), events).await?;
select! {
x = stdout_handle => {
x??;
None
},
x = exit_hook_task => x?
}
} else {
println!("{}", id);
None
};
StdioConsoleStream::restore_terminal_mode();
std::process::exit(code.unwrap_or(0));
}
}
async fn wait_guest_started(id: &str, events: EventStream) -> Result<()> {
let mut stream = events.subscribe();
while let Ok(event) = stream.recv().await {
match event {
Event::GuestChanged(changed) => {
let Some(guest) = changed.guest else {
continue;
};
if guest.id != id {
continue;
}
let Some(state) = guest.state else {
continue;
};
if let Some(ref error) = state.error_info {
if state.status() == GuestStatus::Failed {
error!("launch failed: {}", error.message);
std::process::exit(1);
} else {
error!("guest error: {}", error.message);
}
}
if state.status() == GuestStatus::Destroyed {
error!("guest destroyed");
std::process::exit(1);
}
if state.status() == GuestStatus::Started {
break;
}
}
}
}
Ok(())
}
fn env_map(env: &[String]) -> HashMap<String, String> {
let mut map = HashMap::<String, String>::new();
for item in env {
if let Some((key, value)) = item.split_once('=') {
map.insert(key.to_string(), value.to_string());
}
}
map
}

182
crates/ctl/src/cli/list.rs Normal file
View File

@ -0,0 +1,182 @@
use anyhow::{anyhow, Result};
use clap::{Parser, ValueEnum};
use cli_tables::Table;
use krata::{
events::EventStream,
v1::{
common::{guest_image_spec::Image, Guest, GuestStatus},
control::{
control_service_client::ControlServiceClient, ListGuestsRequest, ResolveGuestRequest,
},
},
};
use serde_json::Value;
use tonic::{transport::Channel, Request};
use crate::format::{guest_state_text, guest_status_text, kv2line, proto2dynamic, proto2kv};
#[derive(ValueEnum, Clone, Debug, PartialEq, Eq)]
enum ListFormat {
CliTable,
Json,
JsonPretty,
Jsonl,
Yaml,
KeyValue,
Simple,
}
#[derive(Parser)]
pub struct ListCommand {
#[arg(short, long, default_value = "cli-table")]
format: ListFormat,
#[arg()]
guest: Option<String>,
}
impl ListCommand {
pub async fn run(
self,
mut client: ControlServiceClient<Channel>,
_events: EventStream,
) -> Result<()> {
let mut guests = if let Some(ref guest) = self.guest {
let reply = client
.resolve_guest(Request::new(ResolveGuestRequest {
name: guest.clone(),
}))
.await?
.into_inner();
if let Some(guest) = reply.guest {
vec![guest]
} else {
return Err(anyhow!("unable to resolve guest '{}'", guest));
}
} else {
client
.list_guests(Request::new(ListGuestsRequest {}))
.await?
.into_inner()
.guests
};
guests.sort_by(|a, b| {
a.spec
.as_ref()
.map(|x| x.name.as_str())
.unwrap_or("")
.cmp(b.spec.as_ref().map(|x| x.name.as_str()).unwrap_or(""))
});
match self.format {
ListFormat::CliTable => {
self.print_guest_table(guests)?;
}
ListFormat::Simple => {
for guest in guests {
let state = guest_status_text(
guest
.state
.as_ref()
.map(|x| x.status())
.unwrap_or(GuestStatus::Unknown),
);
let name = guest.spec.as_ref().map(|x| x.name.as_str()).unwrap_or("");
let network = guest.state.as_ref().and_then(|x| x.network.as_ref());
let ipv4 = network.map(|x| x.guest_ipv4.as_str()).unwrap_or("");
let ipv6 = network.map(|x| x.guest_ipv6.as_str()).unwrap_or("");
println!("{}\t{}\t{}\t{}\t{}", guest.id, state, name, ipv4, ipv6);
}
}
ListFormat::Json | ListFormat::JsonPretty | ListFormat::Yaml => {
let mut values = Vec::new();
for guest in guests {
let message = proto2dynamic(guest)?;
values.push(serde_json::to_value(message)?);
}
let value = Value::Array(values);
let encoded = if self.format == ListFormat::JsonPretty {
serde_json::to_string_pretty(&value)?
} else if self.format == ListFormat::Yaml {
serde_yaml::to_string(&value)?
} else {
serde_json::to_string(&value)?
};
println!("{}", encoded.trim());
}
ListFormat::Jsonl => {
for guest in guests {
let message = proto2dynamic(guest)?;
println!("{}", serde_json::to_string(&message)?);
}
}
ListFormat::KeyValue => {
self.print_key_value(guests)?;
}
}
Ok(())
}
fn print_guest_table(&self, guests: Vec<Guest>) -> Result<()> {
let mut table = Table::new();
let header = vec!["name", "uuid", "state", "ipv4", "ipv6", "image"];
table.push_row(&header)?;
for guest in guests {
let ipv4 = guest
.state
.as_ref()
.and_then(|x| x.network.as_ref())
.map(|x| x.guest_ipv4.as_str())
.unwrap_or("unknown");
let ipv6 = guest
.state
.as_ref()
.and_then(|x| x.network.as_ref())
.map(|x| x.guest_ipv6.as_str())
.unwrap_or("unknown");
let Some(spec) = guest.spec else {
continue;
};
let image = spec
.image
.map(|x| {
x.image
.map(|y| match y {
Image::Oci(oci) => oci.image,
})
.unwrap_or("unknown".to_string())
})
.unwrap_or("unknown".to_string());
table.push_row_string(&vec![
spec.name,
guest.id,
format!("{}", guest_state_text(guest.state.as_ref())),
ipv4.to_string(),
ipv6.to_string(),
image,
])?;
}
if table.num_records() == 1 {
if self.guest.is_none() {
println!("no guests have been launched");
}
} else {
println!("{}", table.to_string());
}
Ok(())
}
fn print_key_value(&self, guests: Vec<Guest>) -> Result<()> {
for guest in guests {
let kvs = proto2kv(guest)?;
println!("{}", kv2line(kvs),);
}
Ok(())
}
}

100
crates/ctl/src/cli/mod.rs Normal file
View File

@ -0,0 +1,100 @@
pub mod attach;
pub mod destroy;
pub mod launch;
pub mod list;
pub mod resolve;
pub mod watch;
use anyhow::{anyhow, Result};
use clap::{Parser, Subcommand};
use krata::{
client::ControlClientProvider,
events::EventStream,
v1::control::{
control_service_client::ControlServiceClient, ResolveGuestRequest, WatchEventsRequest,
},
};
use tonic::{transport::Channel, Request};
use self::{
attach::AttachCommand, destroy::DestroyCommand, launch::LauchCommand, list::ListCommand,
resolve::ResolveCommand, watch::WatchCommand,
};
#[derive(Parser)]
#[command(version, about)]
pub struct ControlCommand {
#[arg(short, long, default_value = "unix:///var/lib/krata/daemon.socket")]
connection: String,
#[command(subcommand)]
command: Commands,
}
#[derive(Subcommand)]
pub enum Commands {
Launch(LauchCommand),
Destroy(DestroyCommand),
List(ListCommand),
Attach(AttachCommand),
Watch(WatchCommand),
Resolve(ResolveCommand),
}
impl ControlCommand {
pub async fn run(self) -> Result<()> {
let mut client = ControlClientProvider::dial(self.connection.parse()?).await?;
let events = EventStream::open(
client
.watch_events(WatchEventsRequest {})
.await?
.into_inner(),
)
.await?;
match self.command {
Commands::Launch(launch) => {
launch.run(client, events).await?;
}
Commands::Destroy(destroy) => {
destroy.run(client, events).await?;
}
Commands::Attach(attach) => {
attach.run(client, events).await?;
}
Commands::List(list) => {
list.run(client, events).await?;
}
Commands::Watch(watch) => {
watch.run(events).await?;
}
Commands::Resolve(resolve) => {
resolve.run(client).await?;
}
}
Ok(())
}
}
pub async fn resolve_guest(
client: &mut ControlServiceClient<Channel>,
name: &str,
) -> Result<String> {
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))
}
}

View File

@ -0,0 +1,28 @@
use anyhow::Result;
use clap::Parser;
use krata::v1::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<Channel>) -> 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(())
}
}

View File

@ -0,0 +1,74 @@
use anyhow::Result;
use clap::{Parser, ValueEnum};
use krata::{
events::EventStream,
v1::{common::Guest, control::watch_events_reply::Event},
};
use prost_reflect::ReflectMessage;
use serde_json::Value;
use crate::format::{guest_state_text, kv2line, proto2dynamic, proto2kv};
#[derive(ValueEnum, Clone, Debug, PartialEq, Eq)]
enum WatchFormat {
Simple,
Json,
KeyValue,
}
#[derive(Parser)]
pub struct WatchCommand {
#[arg(short, long, default_value = "simple")]
format: WatchFormat,
}
impl WatchCommand {
pub async fn run(self, events: EventStream) -> Result<()> {
let mut stream = events.subscribe();
loop {
let event = stream.recv().await?;
match event {
Event::GuestChanged(changed) => {
let guest = changed.guest.clone();
self.print_event("guest.changed", changed, guest)?;
}
}
}
}
fn print_event(
&self,
typ: &str,
event: impl ReflectMessage,
guest: Option<Guest>,
) -> Result<()> {
match self.format {
WatchFormat::Simple => {
if let Some(guest) = guest {
println!(
"{} guest={} status=\"{}\"",
typ,
guest.id,
guest_state_text(guest.state.as_ref()).replace('"', "\\\"")
);
}
}
WatchFormat::Json => {
let message = proto2dynamic(event)?;
let mut value = serde_json::to_value(&message)?;
if let Value::Object(ref mut map) = value {
map.insert("event.type".to_string(), Value::String(typ.to_string()));
}
println!("{}", serde_json::to_string(&value)?);
}
WatchFormat::KeyValue => {
let mut map = proto2kv(event)?;
map.insert("event.type".to_string(), typ.to_string());
println!("{}", kv2line(map),);
}
}
Ok(())
}
}

115
crates/ctl/src/console.rs Normal file
View File

@ -0,0 +1,115 @@
use anyhow::Result;
use async_stream::stream;
use crossterm::{
terminal::{disable_raw_mode, enable_raw_mode, is_raw_mode_enabled},
tty::IsTty,
};
use krata::{
events::EventStream,
v1::{
common::GuestStatus,
control::{watch_events_reply::Event, ConsoleDataReply, ConsoleDataRequest},
},
};
use log::debug;
use tokio::{
io::{stdin, stdout, AsyncReadExt, AsyncWriteExt},
task::JoinHandle,
};
use tokio_stream::{Stream, StreamExt};
use tonic::Streaming;
pub struct StdioConsoleStream;
impl StdioConsoleStream {
pub async fn stdin_stream(guest: String) -> impl Stream<Item = ConsoleDataRequest> {
let mut stdin = stdin();
stream! {
yield ConsoleDataRequest { guest_id: guest, data: vec![] };
let mut buffer = vec![0u8; 60];
loop {
let size = match stdin.read(&mut buffer).await {
Ok(size) => size,
Err(error) => {
debug!("failed to read stdin: {}", error);
break;
}
};
let data = buffer[0..size].to_vec();
if size == 1 && buffer[0] == 0x1d {
break;
}
yield ConsoleDataRequest { guest_id: String::default(), data };
}
}
}
pub async fn stdout(mut stream: Streaming<ConsoleDataReply>) -> Result<()> {
if stdin().is_tty() {
enable_raw_mode()?;
StdioConsoleStream::register_terminal_restore_hook()?;
}
let mut stdout = stdout();
while let Some(reply) = stream.next().await {
let reply = reply?;
if reply.data.is_empty() {
continue;
}
stdout.write_all(&reply.data).await?;
stdout.flush().await?;
}
Ok(())
}
pub async fn guest_exit_hook(
id: String,
events: EventStream,
) -> Result<JoinHandle<Option<i32>>> {
Ok(tokio::task::spawn(async move {
let mut stream = events.subscribe();
while let Ok(event) = stream.recv().await {
match event {
Event::GuestChanged(changed) => {
let Some(guest) = changed.guest else {
continue;
};
let Some(state) = guest.state else {
continue;
};
if guest.id != id {
continue;
}
if let Some(exit_info) = state.exit_info {
return Some(exit_info.code);
}
let status = state.status();
if status == GuestStatus::Destroying || status == GuestStatus::Destroyed {
return Some(10);
}
}
}
}
None
}))
}
fn register_terminal_restore_hook() -> Result<()> {
if stdin().is_tty() {
ctrlc::set_handler(move || {
StdioConsoleStream::restore_terminal_mode();
})?;
}
Ok(())
}
pub fn restore_terminal_mode() {
if is_raw_mode_enabled().unwrap_or(false) {
let _ = disable_raw_mode();
}
}
}

86
crates/ctl/src/format.rs Normal file
View File

@ -0,0 +1,86 @@
use std::collections::HashMap;
use anyhow::Result;
use krata::v1::common::{GuestState, GuestStatus};
use prost_reflect::{DynamicMessage, ReflectMessage, Value};
pub fn proto2dynamic(proto: impl ReflectMessage) -> Result<DynamicMessage> {
Ok(DynamicMessage::decode(
proto.descriptor(),
proto.encode_to_vec().as_slice(),
)?)
}
pub fn proto2kv(proto: impl ReflectMessage) -> Result<HashMap<String, String>> {
let message = proto2dynamic(proto)?;
let mut map = HashMap::new();
fn crawl(prefix: &str, map: &mut HashMap<String, String>, message: &DynamicMessage) {
for (field, value) in message.fields() {
let path = if prefix.is_empty() {
field.name().to_string()
} else {
format!("{}.{}", prefix, field.name())
};
match value {
Value::Message(child) => {
crawl(&path, map, child);
}
Value::EnumNumber(number) => {
if let Some(e) = field.kind().as_enum() {
if let Some(value) = e.get_value(*number) {
map.insert(path, value.name().to_string());
}
}
}
Value::String(value) => {
map.insert(path, value.clone());
}
_ => {
map.insert(path, value.to_string());
}
}
}
}
crawl("", &mut map, &message);
Ok(map)
}
pub fn kv2line(map: HashMap<String, String>) -> String {
map.iter()
.map(|(k, v)| format!("{}=\"{}\"", k, v.replace('"', "\\\"")))
.collect::<Vec<_>>()
.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
}

3
crates/ctl/src/lib.rs Normal file
View File

@ -0,0 +1,3 @@
pub mod cli;
pub mod console;
pub mod format;