feat: guest metrics support (#46)

* feat: initial support for idm send in daemon

* feat: implement IdmClient backend support

* feat: daemon idm now uses IdmClient

* fix: implement channel destruction propagation

* feat: implement request response idm system

* feat: implement metrics support

* proto: move metrics into GuestMetrics for reusability

* fix: log level of guest agent was trace

* feat: metrics tree with process information

crates/ctl/Cargo.toml

@@ -16,11 +16,15 @@ comfy-table = { workspace = true }
 crossterm = { workspace = true }
 ctrlc = { workspace = true, features = ["termination"] }
 env_logger = { workspace = true }
+fancy-duration = { workspace = true }
+human_bytes = { workspace = true }
 krata = { path = "../krata", version = "^0.0.8" }
 log = { workspace = true }
 prost-reflect = { workspace = true, features = ["serde"] }
+prost-types = { workspace = true }
 serde_json = { workspace = true }
 serde_yaml = { workspace = true }
+termtree = { workspace = true }
 tokio = { workspace = true }
 tokio-stream = { workspace = true }
 tonic = { workspace = true }

crates/ctl/src/cli/metrics.rs

@@ -0,0 +1,83 @@
+use anyhow::Result;
+use clap::{Parser, ValueEnum};
+use krata::{
+    events::EventStream,
+    v1::{
+        common::GuestMetricNode,
+        control::{control_service_client::ControlServiceClient, ReadGuestMetricsRequest},
+    },
+};
+
+use tonic::transport::Channel;
+
+use crate::format::{kv2line, metrics_flat, metrics_tree, proto2dynamic};
+
+use super::resolve_guest;
+
+#[derive(ValueEnum, Clone, Debug, PartialEq, Eq)]
+enum MetricsFormat {
+    Tree,
+    Json,
+    JsonPretty,
+    Yaml,
+    KeyValue,
+}
+
+#[derive(Parser)]
+#[command(about = "Read metrics from the guest")]
+pub struct MetricsCommand {
+    #[arg(short, long, default_value = "tree", help = "Output format")]
+    format: MetricsFormat,
+    #[arg(help = "Guest to read metrics for, either the name or the uuid")]
+    guest: String,
+}
+
+impl MetricsCommand {
+    pub async fn run(
+        self,
+        mut client: ControlServiceClient<Channel>,
+        _events: EventStream,
+    ) -> Result<()> {
+        let guest_id: String = resolve_guest(&mut client, &self.guest).await?;
+        let root = client
+            .read_guest_metrics(ReadGuestMetricsRequest { guest_id })
+            .await?
+            .into_inner()
+            .root
+            .unwrap_or_default();
+        match self.format {
+            MetricsFormat::Tree => {
+                self.print_metrics_tree(root)?;
+            }
+
+            MetricsFormat::Json | MetricsFormat::JsonPretty | MetricsFormat::Yaml => {
+                let value = serde_json::to_value(proto2dynamic(root)?)?;
+                let encoded = if self.format == MetricsFormat::JsonPretty {
+                    serde_json::to_string_pretty(&value)?
+                } else if self.format == MetricsFormat::Yaml {
+                    serde_yaml::to_string(&value)?
+                } else {
+                    serde_json::to_string(&value)?
+                };
+                println!("{}", encoded.trim());
+            }
+
+            MetricsFormat::KeyValue => {
+                self.print_key_value(root)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn print_metrics_tree(&self, root: GuestMetricNode) -> Result<()> {
+        print!("{}", metrics_tree(root));
+        Ok(())
+    }
+
+    fn print_key_value(&self, metrics: GuestMetricNode) -> Result<()> {
+        let kvs = metrics_flat(metrics);
+        println!("{}", kv2line(kvs));
+        Ok(())
+    }
+}

crates/ctl/src/cli/mod.rs

@@ -3,6 +3,7 @@ pub mod destroy;
 pub mod launch;
 pub mod list;
 pub mod logs;
+pub mod metrics;
 pub mod resolve;
 pub mod watch;
 
@@ -17,7 +18,7 @@ use tonic::{transport::Channel, Request};
 
 use self::{
     attach::AttachCommand, destroy::DestroyCommand, launch::LauchCommand, list::ListCommand,
-    logs::LogsCommand, resolve::ResolveCommand, watch::WatchCommand,
+    logs::LogsCommand, metrics::MetricsCommand, resolve::ResolveCommand, watch::WatchCommand,
 };
 
 #[derive(Parser)]
@@ -47,6 +48,7 @@ pub enum Commands {
     Logs(LogsCommand),
     Watch(WatchCommand),
     Resolve(ResolveCommand),
+    Metrics(MetricsCommand),
 }
 
 impl ControlCommand {
@@ -82,6 +84,10 @@ impl ControlCommand {
             Commands::Resolve(resolve) => {
                 resolve.run(client).await?;
             }
+
+            Commands::Metrics(metrics) => {
+                metrics.run(client, events).await?;
+            }
         }
         Ok(())
     }

crates/ctl/src/format.rs

@@ -1,8 +1,12 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, time::Duration};
 
 use anyhow::Result;
-use krata::v1::common::{Guest, GuestStatus};
-use prost_reflect::{DynamicMessage, ReflectMessage, Value};
+use fancy_duration::FancyDuration;
+use human_bytes::human_bytes;
+use krata::v1::common::{Guest, GuestMetricFormat, GuestMetricNode, GuestStatus};
+use prost_reflect::{DynamicMessage, FieldDescriptor, ReflectMessage, Value as ReflectValue};
+use prost_types::Value;
+use termtree::Tree;
 
 pub fn proto2dynamic(proto: impl ReflectMessage) -> Result<DynamicMessage> {
     Ok(DynamicMessage::decode(
@@ -15,38 +19,56 @@ 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);
+    fn crawl(
+        prefix: String,
+        field: Option<&FieldDescriptor>,
+        map: &mut HashMap<String, String>,
+        value: &ReflectValue,
+    ) {
+        match value {
+            ReflectValue::Message(child) => {
+                for (field, field_value) in child.fields() {
+                    let path = if prefix.is_empty() {
+                        field.json_name().to_string()
+                    } else {
+                        format!("{}.{}", prefix, field.json_name())
+                    };
+                    crawl(path, Some(&field), map, field_value);
                 }
+            }
 
-                Value::EnumNumber(number) => {
-                    if let Some(e) = field.kind().as_enum() {
+            ReflectValue::EnumNumber(number) => {
+                if let Some(kind) = field.map(|x| x.kind()) {
+                    if let Some(e) = kind.as_enum() {
                         if let Some(value) = e.get_value(*number) {
-                            map.insert(path, value.name().to_string());
+                            map.insert(prefix, value.name().to_string());
                         }
                     }
                 }
+            }
 
-                Value::String(value) => {
-                    map.insert(path, value.clone());
-                }
+            ReflectValue::String(value) => {
+                map.insert(prefix.to_string(), value.clone());
+            }
 
-                _ => {
-                    map.insert(path, value.to_string());
+            ReflectValue::List(value) => {
+                for (x, value) in value.iter().enumerate() {
+                    crawl(format!("{}.{}", prefix, x), field, map, value);
                 }
             }
+
+            _ => {
+                map.insert(prefix.to_string(), value.to_string());
+            }
         }
     }
 
-    crawl("", &mut map, &message);
+    crawl(
+        "".to_string(),
+        None,
+        &mut map,
+        &ReflectValue::Message(message),
+    );
 
     Ok(map)
 }
@@ -85,3 +107,63 @@ pub fn guest_simple_line(guest: &Guest) -> String {
     let ipv6 = network.map(|x| x.guest_ipv6.as_str()).unwrap_or("");
     format!("{}\t{}\t{}\t{}\t{}", guest.id, state, name, ipv4, ipv6)
 }
+
+fn metrics_value_string(value: Value) -> String {
+    proto2dynamic(value)
+        .map(|x| serde_json::to_string(&x).ok())
+        .ok()
+        .flatten()
+        .unwrap_or_default()
+}
+
+fn metrics_value_numeric(value: Value) -> f64 {
+    let string = metrics_value_string(value);
+    string.parse::<f64>().ok().unwrap_or(f64::NAN)
+}
+
+fn metrics_value_pretty(value: Value, format: GuestMetricFormat) -> String {
+    match format {
+        GuestMetricFormat::Bytes => human_bytes(metrics_value_numeric(value)),
+        GuestMetricFormat::Integer => (metrics_value_numeric(value) as u64).to_string(),
+        GuestMetricFormat::DurationSeconds => {
+            FancyDuration(Duration::from_secs_f64(metrics_value_numeric(value))).to_string()
+        }
+        _ => metrics_value_string(value),
+    }
+}
+
+fn metrics_flat_internal(prefix: &str, node: GuestMetricNode, map: &mut HashMap<String, String>) {
+    if let Some(value) = node.value {
+        map.insert(prefix.to_string(), metrics_value_string(value));
+    }
+
+    for child in node.children {
+        let path = if prefix.is_empty() {
+            child.name.to_string()
+        } else {
+            format!("{}.{}", prefix, child.name)
+        };
+        metrics_flat_internal(&path, child, map);
+    }
+}
+
+pub fn metrics_flat(root: GuestMetricNode) -> HashMap<String, String> {
+    let mut map = HashMap::new();
+    metrics_flat_internal("", root, &mut map);
+    map
+}
+
+pub fn metrics_tree(node: GuestMetricNode) -> Tree<String> {
+    let mut name = node.name.to_string();
+    let format = node.format();
+    if let Some(value) = node.value {
+        let value_string = metrics_value_pretty(value, format);
+        name.push_str(&format!(": {}", value_string));
+    }
+
+    let mut tree = Tree::new(name);
+    for child in node.children {
+        tree.push(metrics_tree(child));
+    }
+    tree
+}