feature(krata): first pass on cpu hotplug support

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

@@ -39,20 +39,32 @@ pub struct ZoneLaunchCommand {
     pull_update: bool,
     #[arg(short, long, help = "Name of the zone")]
     name: Option<String>,
-    #[arg(short, long, default_value_t = 1, help = "vCPUs available to the zone")]
+    #[arg(
-    cpus: u32,
+        short = 'C',
+        long = "max-cpus",
+        default_value_t = 4,
+        help = "Maximum vCPUs available for the zone"
+    )]
+    max_cpus: u32,
+    #[arg(
+        short = 'c',
+        long = "target-cpus",
+        default_value_t = 1,
+        help = "Target vCPUs for the zone to use"
+    )]
+    target_cpus: u32,
     #[arg(
         short = 'M',
         long = "max-memory",
-        default_value_t = 512,
+        default_value_t = 1024,
         help = "Maximum memory available to the zone, in megabytes"
     )]
     max_memory: u64,
     #[arg(
         short = 'm',
         long = "target-memory",
-        default_value_t = 512,
+        default_value_t = 1024,
-        help = "Memory target for the zone, in megabytes"
+        help = "Target memory for the zone to use, in megabytes"
     )]
     target_memory: u64,
     #[arg[short = 'D', long = "device", help = "Devices to request for the zone"]]
@@ -131,9 +143,10 @@ impl ZoneLaunchCommand {
                 kernel,
                 initrd,
                 initial_resources: Some(ZoneResourceSpec {
-                    cpus: self.cpus,
                     max_memory: self.max_memory,
                     target_memory: self.target_memory,
+                    max_cpus: self.max_cpus,
+                    target_cpus: self.target_cpus,
                 }),
                 task: Some(ZoneTaskSpec {
                     environment: env_map(&self.env.unwrap_or_default())

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

@@ -14,20 +14,32 @@ use tonic::{transport::Channel, Request};
 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")]
+    #[arg(
-    cpus: u32,
+        short = 'C',
+        long = "max-cpus",
+        default_value_t = 0,
+        help = "Maximum vCPUs available to the zone (0 means previous value)"
+    )]
+    max_cpus: u32,
+    #[arg(
+        short = 'c',
+        long = "target-cpus",
+        default_value_t = 0,
+        help = "Target vCPUs for the zone to use (0 means previous value)"
+    )]
+    target_cpus: u32,
     #[arg(
         short = 'M',
         long = "max-memory",
         default_value_t = 0,
-        help = "Maximum memory available to the zone, in megabytes"
+        help = "Maximum memory available to the zone, in megabytes (0 means previous value)"
     )]
     max_memory: u64,
     #[arg(
         short = 'm',
         long = "target-memory",
         default_value_t = 0,
-        help = "Memory target for the zone, in megabytes"
+        help = "Target memory for the zone to use, in megabytes (0 means previous value)"
     )]
     target_memory: u64,
 }
@@ -63,10 +75,15 @@ impl ZoneUpdateResourcesCommand {
                     } else {
                         self.target_memory
                     },
-                    cpus: if self.cpus == 0 {
+                    max_cpus: if self.max_cpus == 0 {
-                        active_resources.cpus
+                        active_resources.max_cpus
                     } else {
-                        self.cpus
+                        self.max_cpus
+                    },
+                    target_cpus: if self.target_cpus == 0 {
+                        active_resources.target_cpus
+                    } else {
+                        self.target_cpus
                     },
                 }),
             }))

crates/daemon/src/control.rs

@@ -669,6 +669,21 @@ impl ControlService for DaemonControlService {
             resources.max_memory = resources.target_memory;
         }
 
+        if resources.target_cpus < 1 {
+            resources.target_cpus = 1;
+        }
+
+        let initial_resources = zone
+            .spec
+            .clone()
+            .unwrap_or_default()
+            .initial_resources
+            .unwrap_or_default();
+        if resources.target_cpus > initial_resources.max_cpus {
+            resources.target_cpus = initial_resources.max_cpus;
+        }
+        resources.max_cpus = initial_resources.max_cpus;
+
         self.runtime
             .set_memory_resources(
                 status.domid,
@@ -679,6 +694,12 @@ impl ControlService for DaemonControlService {
             .map_err(|error| ApiError {
                 message: format!("failed to set memory resources: {}", error),
             })?;
+        self.runtime
+            .set_cpu_resources(status.domid, resources.target_cpus)
+            .await
+            .map_err(|error| ApiError {
+                message: format!("failed to set cpu resources: {}", error),
+            })?;
         status.resource_status = Some(ZoneResourceStatus {
             active_resources: Some(resources),
         });

crates/daemon/src/reconcile/zone/create.rs

@@ -76,7 +76,7 @@ impl ZoneCreator<'_> {
     }
 
     pub async fn create(&self, uuid: Uuid, zone: &mut Zone) -> Result<ZoneReconcilerResult> {
-        let Some(ref spec) = zone.spec else {
+        let Some(ref mut spec) = zone.spec else {
             return Err(anyhow!("zone spec not specified"));
         };
 
@@ -176,7 +176,14 @@ impl ZoneCreator<'_> {
 
         let reservation = self.ip_assignment.assign(uuid).await?;
 
-        let initial_resources = spec.initial_resources.unwrap_or_default();
+        let mut initial_resources = spec.initial_resources.unwrap_or_default();
+        if initial_resources.target_cpus < 1 {
+            initial_resources.target_cpus = 1;
+        }
+        if initial_resources.target_cpus > initial_resources.max_cpus {
+            initial_resources.max_cpus = initial_resources.target_cpus;
+        }
+        spec.initial_resources = Some(initial_resources);
         let info = self
             .runtime
             .launch(ZoneLaunchRequest {
@@ -190,7 +197,8 @@ impl ZoneCreator<'_> {
                 image,
                 kernel,
                 initrd,
-                cpus: initial_resources.cpus,
+                target_cpus: initial_resources.target_cpus,
+                max_cpus: initial_resources.max_cpus,
                 max_memory: initial_resources.max_memory,
                 target_memory: initial_resources.target_memory,
                 pcis,

crates/krata/proto/krata/v1/common.proto

@@ -30,7 +30,8 @@ message ZoneSpec {
 message ZoneResourceSpec {
     uint64 max_memory = 1;
     uint64 target_memory = 2;
-    uint32 cpus = 3;
+    uint32 max_cpus = 3;
+    uint32 target_cpus = 4;
 }
 
 message ZoneImageSpec {

crates/runtime/src/launch.rs

@@ -30,7 +30,8 @@ pub struct ZoneLaunchRequest {
     pub initrd: Vec<u8>,
     pub uuid: Option<Uuid>,
     pub name: Option<String>,
-    pub cpus: u32,
+    pub target_cpus: u32,
+    pub max_cpus: u32,
     pub target_memory: u64,
     pub max_memory: u64,
     pub env: HashMap<String, String>,
@@ -195,7 +196,8 @@ impl ZoneLauncher {
 
         let config = DomainConfig {
             base: BaseDomainConfig {
-                max_vcpus: request.cpus,
+                max_vcpus: request.max_cpus,
+                target_vcpus: request.target_cpus,
                 max_mem_mb: request.max_memory,
                 target_mem_mb: request.target_memory,
                 kernel: request.kernel,

crates/runtime/src/lib.rs

@@ -1,7 +1,6 @@
-use std::{fs, path::PathBuf, str::FromStr, sync::Arc};
-
 use anyhow::{anyhow, Result};
 use krataloopdev::LoopControl;
+use std::{fs, path::PathBuf, str::FromStr, sync::Arc};
 use tokio::sync::Semaphore;
 use uuid::Uuid;
 
@@ -258,6 +257,34 @@ impl Runtime {
         Ok(())
     }
 
+    pub async fn set_cpu_resources(&self, domid: u32, target_cpus: u32) -> Result<()> {
+        let domain_path = self.context.xen.store.get_domain_path(domid).await?;
+        let cpus = self
+            .context
+            .xen
+            .store
+            .list(&format!("{}/cpu", domain_path))
+            .await?;
+        let tx = self.context.xen.store.transaction().await?;
+        for cpu in cpus {
+            let Some(id) = cpu.parse::<u32>().ok() else {
+                continue;
+            };
+            let available = if id >= target_cpus {
+                "offline"
+            } else {
+                "online"
+            };
+            tx.write_string(
+                format!("{}/cpu/{}/availability", domain_path, id),
+                available,
+            )
+            .await?;
+        }
+        tx.commit().await?;
+        Ok(())
+    }
+
     pub async fn list(&self) -> Result<Vec<ZoneInfo>> {
         self.context.list().await
     }

crates/xen/xenclient/examples/boot.rs

@@ -27,6 +27,7 @@ async fn main() -> Result<()> {
         base: BaseDomainConfig {
             uuid: Uuid::new_v4(),
             max_vcpus: 1,
+            target_vcpus: 1,
             max_mem_mb: 512,
             target_mem_mb: 512,
             enable_iommu: true,

crates/xen/xenclient/src/tx.rs

@@ -194,7 +194,16 @@ impl ClientTransaction {
             self.tx.mkdir(&path).await?;
             self.tx.set_perms(&path, ro_perm).await?;
             let path = format!("{}/cpu/{}/availability", self.dom_path, i);
-            self.tx.write_string(&path, "online").await?;
+            self.tx
+                .write_string(
+                    &path,
+                    if i < base.target_vcpus {
+                        "online"
+                    } else {
+                        "offline"
+                    },
+                )
+                .await?;
             self.tx.set_perms(&path, ro_perm).await?;
         }
         Ok(())

crates/xen/xenplatform/src/domain.rs

@@ -66,6 +66,7 @@ pub struct BaseDomainConfig {
     pub uuid: Uuid,
     pub owner_domid: u32,
     pub max_vcpus: u32,
+    pub target_vcpus: u32,
     pub max_mem_mb: u64,
     pub target_mem_mb: u64,
     pub kernel: Vec<u8>,