From 82b2b976b49e90211ea2fd1b0aaefe418ef312ed Mon Sep 17 00:00:00 2001
From: Ariadne Conill <ariadne@ariadne.space>
Date: Sat, 29 Jun 2024 00:11:38 -0700
Subject: [PATCH] feat(power-management-core): add cpu-topology command

Signed-off-by: Ariadne Conill <ariadne@ariadne.space>
---
 crates/ctl/src/cli/cpu_topology.rs | 37 ++++++++++++++++++++++++++++++
 crates/ctl/src/cli/mod.rs          |  8 ++++++-
 2 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 crates/ctl/src/cli/cpu_topology.rs

diff --git a/crates/ctl/src/cli/cpu_topology.rs b/crates/ctl/src/cli/cpu_topology.rs
new file mode 100644
index 0000000..ca9eeed
--- /dev/null
+++ b/crates/ctl/src/cli/cpu_topology.rs
@@ -0,0 +1,37 @@
+use anyhow::Result;
+use clap::Parser;
+use krata::v1::control::{control_service_client::ControlServiceClient, HostCpuTopologyRequest};
+
+use tonic::{transport::Channel, Request};
+
+fn class_to_str(input: i32) -> String {
+    match input {
+        0 => "Standard".to_string(),
+        1 => "Performance".to_string(),
+        2 => "Efficiency".to_string(),
+        _ => "???".to_string()
+    }
+}
+
+#[derive(Parser)]
+#[command(about = "Display information about a host's CPU topology")]
+pub struct CpuTopologyCommand {}
+
+impl CpuTopologyCommand {
+    pub async fn run(self, mut client: ControlServiceClient<Channel>) -> Result<()> {
+        println!("{0:<10} {1:<10} {2:<10} {3:<10} {4:<10} {5:<10}", "CPUID", "Node", "Socket", "Core", "Thread", "Class");
+
+        let response = client
+            .get_host_cpu_topology(Request::new(HostCpuTopologyRequest {}))
+            .await?
+            .into_inner();
+
+        let mut i = 0;
+        for cpu in response.cpus {
+            println!("{0:<10} {1:<10} {2:<10} {3:<10} {4:<10} {5:<10}", i, cpu.node, cpu.socket, cpu.core, cpu.thread, class_to_str(cpu.class));
+            i += 1;
+        }
+        
+        Ok(())
+    }
+}
diff --git a/crates/ctl/src/cli/mod.rs b/crates/ctl/src/cli/mod.rs
index 2f13db5..93d368d 100644
--- a/crates/ctl/src/cli/mod.rs
+++ b/crates/ctl/src/cli/mod.rs
@@ -1,4 +1,5 @@
 pub mod attach;
+pub mod cpu_topology;
 pub mod destroy;
 pub mod exec;
 pub mod identify_host;
@@ -23,7 +24,7 @@ use krata::{
 use tonic::{transport::Channel, Request};
 
 use self::{
-    attach::AttachCommand, destroy::DestroyCommand, exec::ExecCommand,
+    attach::AttachCommand, cpu_topology::CpuTopologyCommand, destroy::DestroyCommand, exec::ExecCommand,
     identify_host::IdentifyHostCommand, idm_snoop::IdmSnoopCommand, launch::LaunchCommand,
     list::ListCommand, list_devices::ListDevicesCommand, logs::LogsCommand,
     metrics::MetricsCommand, pull::PullCommand, resolve::ResolveCommand, top::TopCommand,
@@ -61,6 +62,7 @@ pub enum Commands {
     Top(TopCommand),
     IdentifyHost(IdentifyHostCommand),
     Exec(ExecCommand),
+    CpuTopology(CpuTopologyCommand),
 }
 
 impl ControlCommand {
@@ -124,6 +126,10 @@ impl ControlCommand {
             Commands::ListDevices(list) => {
                 list.run(client, events).await?;
             }
+
+            Commands::CpuTopology(cpu_topology) => {
+                cpu_topology.run(client).await?;
+            }
         }
         Ok(())
     }