feature(krata): dynamic resource allocation (closes #298) (#333)

crates/ctl/src/cli/zone/launch.rs

@@ -6,8 +6,9 @@ use krata::{
     events::EventStream,
     v1::{
         common::{
-            zone_image_spec::Image, OciImageFormat, ZoneImageSpec, ZoneOciImageSpec, ZoneSpec,
-            ZoneSpecDevice, ZoneState, ZoneTaskSpec, ZoneTaskSpecEnvVar,
+            zone_image_spec::Image, OciImageFormat, ZoneImageSpec, ZoneOciImageSpec,
+            ZoneResourceSpec, ZoneSpec, ZoneSpecDevice, ZoneState, ZoneTaskSpec,
+            ZoneTaskSpecEnvVar,
         },
         control::{
             control_service_client::ControlServiceClient, watch_events_reply::Event,
@@ -41,12 +42,19 @@ pub struct ZoneLaunchCommand {
     #[arg(short, long, default_value_t = 1, help = "vCPUs available to the zone")]
     cpus: u32,
     #[arg(
-        short,
-        long,
+        short = 'M',
+        long = "max-memory",
         default_value_t = 512,
-        help = "Memory available to the zone, in megabytes"
+        help = "Maximum memory available to the zone, in megabytes"
     )]
-    mem: u64,
+    max_memory: u64,
+    #[arg(
+        short = 'm',
+        long = "target-memory",
+        default_value_t = 512,
+        help = "Memory target for the zone, in megabytes"
+    )]
+    target_memory: u64,
     #[arg[short = 'D', long = "device", help = "Devices to request for the zone"]]
     device: Vec<String>,
     #[arg[short, long, help = "Environment variables set in the zone"]]
@@ -120,8 +128,11 @@ impl ZoneLaunchCommand {
                 image: Some(image),
                 kernel,
                 initrd,
-                cpus: self.cpus,
-                mem: self.mem,
+                initial_resources: Some(ZoneResourceSpec {
+                    cpus: self.cpus,
+                    max_memory: self.max_memory,
+                    target_memory: self.target_memory,
+                }),
                 task: Some(ZoneTaskSpec {
                     environment: env_map(&self.env.unwrap_or_default())
                         .iter()

crates/ctl/src/cli/zone/mod.rs

@@ -14,6 +14,7 @@ use crate::cli::zone::logs::ZoneLogsCommand;
 use crate::cli::zone::metrics::ZoneMetricsCommand;
 use crate::cli::zone::resolve::ZoneResolveCommand;
 use crate::cli::zone::top::ZoneTopCommand;
+use crate::cli::zone::update_resources::ZoneUpdateResourcesCommand;
 use crate::cli::zone::watch::ZoneWatchCommand;
 
 pub mod attach;
@@ -25,6 +26,7 @@ pub mod logs;
 pub mod metrics;
 pub mod resolve;
 pub mod top;
+mod update_resources;
 pub mod watch;
 
 #[derive(Parser)]
@@ -56,6 +58,7 @@ pub enum ZoneCommands {
     Resolve(ZoneResolveCommand),
     Top(ZoneTopCommand),
     Watch(ZoneWatchCommand),
+    UpdateResources(ZoneUpdateResourcesCommand),
 }
 
 impl ZoneCommands {
@@ -84,6 +87,8 @@ impl ZoneCommands {
             ZoneCommands::Top(top) => top.run(client, events).await,
 
             ZoneCommands::Exec(exec) => exec.run(client).await,
+
+            ZoneCommands::UpdateResources(update_resources) => update_resources.run(client).await,
         }
     }
 }

crates/ctl/src/cli/zone/update_resources.rs

@@ -0,0 +1,76 @@
+use anyhow::Result;
+use clap::Parser;
+use krata::v1::{
+    common::ZoneResourceSpec,
+    control::{control_service_client::ControlServiceClient, UpdateZoneResourcesRequest},
+};
+
+use crate::cli::resolve_zone;
+use krata::v1::control::GetZoneRequest;
+use tonic::{transport::Channel, Request};
+
+#[derive(Parser)]
+#[command(about = "Update the available resources to a zone")]
+pub struct ZoneUpdateResourcesCommand {
+    #[arg(help = "Zone to update resources of, either the name or the uuid")]
+    zone: String,
+    #[arg(short, long, default_value_t = 0, help = "vCPUs available to the zone")]
+    cpus: u32,
+    #[arg(
+        short = 'M',
+        long = "max-memory",
+        default_value_t = 0,
+        help = "Maximum memory available to the zone, in megabytes"
+    )]
+    max_memory: u64,
+    #[arg(
+        short = 'm',
+        long = "target-memory",
+        default_value_t = 0,
+        help = "Memory target for the zone, in megabytes"
+    )]
+    target_memory: u64,
+}
+
+impl ZoneUpdateResourcesCommand {
+    pub async fn run(self, mut client: ControlServiceClient<Channel>) -> Result<()> {
+        let zone_id = resolve_zone(&mut client, &self.zone).await?;
+        let zone = client
+            .get_zone(GetZoneRequest { zone_id })
+            .await?
+            .into_inner()
+            .zone
+            .unwrap_or_default();
+        let active_resources = zone
+            .status
+            .clone()
+            .unwrap_or_default()
+            .resource_status
+            .unwrap_or_default()
+            .active_resources
+            .unwrap_or_default();
+        client
+            .update_zone_resources(Request::new(UpdateZoneResourcesRequest {
+                zone_id: zone.id.clone(),
+                resources: Some(ZoneResourceSpec {
+                    max_memory: if self.max_memory == 0 {
+                        active_resources.max_memory
+                    } else {
+                        self.max_memory
+                    },
+                    target_memory: if self.target_memory == 0 {
+                        active_resources.target_memory
+                    } else {
+                        self.target_memory
+                    },
+                    cpus: if self.cpus == 0 {
+                        active_resources.cpus
+                    } else {
+                        self.cpus
+                    },
+                }),
+            }))
+            .await?;
+        Ok(())
+    }
+}