From 3715b0b8488d76ac91246f62bcd435e97a438a8a Mon Sep 17 00:00:00 2001
From: Alex Zenla <alex@edera.dev>
Date: Wed, 10 Apr 2024 13:05:08 +0000
Subject: [PATCH] proto: move metrics into GuestMetrics for reusability

---
 crates/ctl/src/cli/metrics.rs             | 21 +++++++++++----------
 crates/daemon/src/control.rs              | 11 +++++++----
 crates/krata/proto/krata/v1/control.proto |  4 ++++
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/crates/ctl/src/cli/metrics.rs b/crates/ctl/src/cli/metrics.rs
index cbb2223..42c9d7f 100644
--- a/crates/ctl/src/cli/metrics.rs
+++ b/crates/ctl/src/cli/metrics.rs
@@ -4,8 +4,7 @@ use comfy_table::{presets::UTF8_FULL_CONDENSED, Table};
 use krata::{
     events::EventStream,
     v1::control::{
-        control_service_client::ControlServiceClient, ReadGuestMetricsReply,
-        ReadGuestMetricsRequest,
+        control_service_client::ControlServiceClient, GuestMetrics, ReadGuestMetricsRequest,
     },
 };
 
@@ -40,17 +39,19 @@ impl MetricsCommand {
         _events: EventStream,
     ) -> Result<()> {
         let guest_id: String = resolve_guest(&mut client, &self.guest).await?;
-        let reply = client
+        let metrics = client
             .read_guest_metrics(ReadGuestMetricsRequest { guest_id })
             .await?
-            .into_inner();
+            .into_inner()
+            .metrics
+            .unwrap_or_default();
         match self.format {
             MetricsFormat::Table => {
-                self.print_metrics_table(reply)?;
+                self.print_metrics_table(metrics)?;
             }
 
             MetricsFormat::Json | MetricsFormat::JsonPretty | MetricsFormat::Yaml => {
-                let value = serde_json::to_value(proto2dynamic(reply)?)?;
+                let value = serde_json::to_value(proto2dynamic(metrics)?)?;
                 let encoded = if self.format == MetricsFormat::JsonPretty {
                     serde_json::to_string_pretty(&value)?
                 } else if self.format == MetricsFormat::Yaml {
@@ -62,19 +63,19 @@ impl MetricsCommand {
             }
 
             MetricsFormat::KeyValue => {
-                self.print_key_value(reply)?;
+                self.print_key_value(metrics)?;
             }
         }
 
         Ok(())
     }
 
-    fn print_metrics_table(&self, reply: ReadGuestMetricsReply) -> Result<()> {
+    fn print_metrics_table(&self, metrics: GuestMetrics) -> Result<()> {
         let mut table = Table::new();
         table.load_preset(UTF8_FULL_CONDENSED);
         table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
         table.set_header(vec!["metric", "value"]);
-        let kvs = proto2kv(reply)?;
+        let kvs = proto2kv(metrics)?;
         for (key, value) in kvs {
             table.add_row(vec![key, value]);
         }
@@ -82,7 +83,7 @@ impl MetricsCommand {
         Ok(())
     }
 
-    fn print_key_value(&self, metrics: ReadGuestMetricsReply) -> Result<()> {
+    fn print_key_value(&self, metrics: GuestMetrics) -> Result<()> {
         let kvs = proto2kv(metrics)?;
         println!("{}", kv2line(kvs),);
         Ok(())
diff --git a/crates/daemon/src/control.rs b/crates/daemon/src/control.rs
index fe32131..2eb1f42 100644
--- a/crates/daemon/src/control.rs
+++ b/crates/daemon/src/control.rs
@@ -12,8 +12,9 @@ use krata::{
         control::{
             control_service_server::ControlService, ConsoleDataReply, ConsoleDataRequest,
             CreateGuestReply, CreateGuestRequest, DestroyGuestReply, DestroyGuestRequest,
-            ListGuestsReply, ListGuestsRequest, ReadGuestMetricsReply, ReadGuestMetricsRequest,
-            ResolveGuestReply, ResolveGuestRequest, WatchEventsReply, WatchEventsRequest,
+            GuestMetrics, ListGuestsReply, ListGuestsRequest, ReadGuestMetricsReply,
+            ReadGuestMetricsRequest, ResolveGuestReply, ResolveGuestRequest, WatchEventsReply,
+            WatchEventsRequest,
         },
     },
 };
@@ -327,8 +328,10 @@ impl ControlService for RuntimeControlService {
 
         let mut reply = ReadGuestMetricsReply::default();
         if let IdmResponseType::Metrics(metrics) = response {
-            reply.total_memory_bytes = metrics.total_memory_bytes;
-            reply.used_memory_bytes = metrics.used_memory_bytes;
+            reply.metrics = Some(GuestMetrics {
+                total_memory_bytes: metrics.total_memory_bytes,
+                used_memory_bytes: metrics.used_memory_bytes,
+            });
         }
         Ok(Response::new(reply))
     }
diff --git a/crates/krata/proto/krata/v1/control.proto b/crates/krata/proto/krata/v1/control.proto
index 6634cd1..0c91af3 100644
--- a/crates/krata/proto/krata/v1/control.proto
+++ b/crates/krata/proto/krata/v1/control.proto
@@ -73,6 +73,10 @@ message ReadGuestMetricsRequest {
 }
 
 message ReadGuestMetricsReply {
+    GuestMetrics metrics = 1;
+}
+
+message GuestMetrics {
     uint64 total_memory_bytes = 1;
     uint64 used_memory_bytes = 2;
 }