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(())
+    }
+}

crates/daemon/src/control.rs

@@ -6,6 +6,7 @@ use crate::{
 };
 use async_stream::try_stream;
 use futures::Stream;
+use krata::v1::common::ZoneResourceStatus;
 use krata::v1::control::{
     GetZoneReply, GetZoneRequest, SetHostPowerManagementPolicyReply,
     SetHostPowerManagementPolicyRequest,
@@ -25,8 +26,8 @@ use krata::{
             HostCpuTopologyInfo, HostStatusReply, HostStatusRequest, ListDevicesReply,
             ListDevicesRequest, ListZonesReply, ListZonesRequest, PullImageReply, PullImageRequest,
             ReadZoneMetricsReply, ReadZoneMetricsRequest, ResolveZoneIdReply, ResolveZoneIdRequest,
-            SnoopIdmReply, SnoopIdmRequest, WatchEventsReply, WatchEventsRequest, ZoneConsoleReply,
-            ZoneConsoleRequest,
+            SnoopIdmReply, SnoopIdmRequest, UpdateZoneResourcesReply, UpdateZoneResourcesRequest,
+            WatchEventsReply, WatchEventsRequest, ZoneConsoleReply, ZoneConsoleRequest,
         },
     },
 };
@@ -165,6 +166,7 @@ impl ControlService for DaemonControlService {
                         network_status: None,
                         exit_status: None,
                         error_status: None,
+                        resource_status: None,
                         host: self.glt.host_uuid().to_string(),
                         domid: u32::MAX,
                     }),
@@ -623,4 +625,67 @@ impl ControlService for DaemonControlService {
             zone: zone.cloned(),
         }))
     }
+
+    async fn update_zone_resources(
+        &self,
+        request: Request<UpdateZoneResourcesRequest>,
+    ) -> Result<Response<UpdateZoneResourcesReply>, Status> {
+        let request = request.into_inner();
+        let uuid = Uuid::from_str(&request.zone_id).map_err(|error| ApiError {
+            message: error.to_string(),
+        })?;
+        let Some(mut zone) = self.zones.read(uuid).await.map_err(ApiError::from)? else {
+            return Err(ApiError {
+                message: "zone not found".to_string(),
+            }
+            .into());
+        };
+
+        let Some(ref mut status) = zone.status else {
+            return Err(ApiError {
+                message: "zone state not available".to_string(),
+            }
+            .into());
+        };
+
+        if status.state() != ZoneState::Created {
+            return Err(ApiError {
+                message: "zone is in an invalid state".to_string(),
+            }
+            .into());
+        }
+
+        if status.domid == 0 || status.domid == u32::MAX {
+            return Err(ApiError {
+                message: "zone domid is invalid".to_string(),
+            }
+            .into());
+        }
+
+        let resources = request.resources.unwrap_or_default();
+
+        self.runtime
+            .set_max_memory(status.domid, resources.max_memory * 1024 * 1024)
+            .await
+            .map_err(|error| ApiError {
+                message: format!("failed to set maximum memory: {}", error),
+            })?;
+
+        self.runtime
+            .set_target_memory(status.domid, resources.target_memory * 1024 * 1024)
+            .await
+            .map_err(|error| ApiError {
+                message: format!("failed to set target memory: {}", error),
+            })?;
+
+        status.resource_status = Some(ZoneResourceStatus {
+            active_resources: Some(resources),
+        });
+
+        self.zones
+            .update(uuid, zone)
+            .await
+            .map_err(ApiError::from)?;
+        Ok(Response::new(UpdateZoneResourcesReply {}))
+    }
 }

crates/daemon/src/event.rs

@@ -137,6 +137,7 @@ impl DaemonEventGenerator {
                 network_status: zone.status.clone().unwrap_or_default().network_status,
                 exit_status: Some(ZoneExitStatus { code }),
                 error_status: None,
+                resource_status: zone.status.clone().unwrap_or_default().resource_status,
                 host: zone.status.clone().map(|x| x.host).unwrap_or_default(),
                 domid: zone.status.clone().map(|x| x.domid).unwrap_or(u32::MAX),
             });

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

@@ -1,8 +1,8 @@
 use anyhow::{anyhow, Result};
 use futures::StreamExt;
 use krata::launchcfg::LaunchPackedFormat;
-use krata::v1::common::ZoneOciImageSpec;
 use krata::v1::common::{OciImageFormat, Zone, ZoneState, ZoneStatus};
+use krata::v1::common::{ZoneOciImageSpec, ZoneResourceStatus};
 use krataoci::packer::{service::OciPackerService, OciPackedFormat};
 use kratart::launch::{PciBdf, PciDevice, PciRdmReservePolicy, ZoneLaunchNetwork};
 use kratart::{launch::ZoneLaunchRequest, Runtime};
@@ -176,6 +176,7 @@ impl ZoneCreator<'_> {
 
         let reservation = self.ip_assignment.assign(uuid).await?;
 
+        let initial_resources = spec.initial_resources.unwrap_or_default();
         let info = self
             .runtime
             .launch(ZoneLaunchRequest {
@@ -189,8 +190,9 @@ impl ZoneCreator<'_> {
                 image,
                 kernel,
                 initrd,
-                vcpus: spec.cpus,
-                mem: spec.mem,
+                cpus: initial_resources.cpus,
+                max_memory: initial_resources.max_memory,
+                target_memory: initial_resources.target_memory,
                 pcis,
                 env: task
                     .environment
@@ -219,6 +221,9 @@ impl ZoneCreator<'_> {
             network_status: Some(ip_reservation_to_network_status(&reservation)),
             exit_status: None,
             error_status: None,
+            resource_status: Some(ZoneResourceStatus {
+                active_resources: Some(initial_resources),
+            }),
             host: self.zlt.host_uuid().to_string(),
             domid: info.domid,
         });

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

@@ -328,6 +328,7 @@ impl ZoneReconciler {
             network_status: None,
             exit_status: None,
             error_status: None,
+            resource_status: None,
             host: self.zlt.host_uuid().to_string(),
             domid: domid.unwrap_or(u32::MAX),
         });

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

@@ -21,11 +21,16 @@ message ZoneSpec {
     ZoneImageSpec kernel = 3;
     // If not specified, defaults to the daemon default initrd.
     ZoneImageSpec initrd = 4;
-    uint32 cpus = 5;
-    uint64 mem = 6;
-    ZoneTaskSpec task = 7;
-    repeated ZoneSpecAnnotation annotations = 8;
-    repeated ZoneSpecDevice devices = 9;
+    ZoneResourceSpec initial_resources = 5;
+    ZoneTaskSpec task = 6;
+    repeated ZoneSpecAnnotation annotations = 7;
+    repeated ZoneSpecDevice devices = 8;
+}
+
+message ZoneResourceSpec {
+    uint64 max_memory = 1;
+    uint64 target_memory = 2;
+    uint32 cpus = 3;
 }
 
 message ZoneImageSpec {
@@ -74,6 +79,7 @@ message ZoneStatus {
     ZoneErrorStatus error_status = 4;
     string host = 5;
     uint32 domid = 6;
+    ZoneResourceStatus resource_status = 7;
 }
 
 enum ZoneState {
@@ -103,6 +109,10 @@ message ZoneErrorStatus {
     string message = 1;
 }
 
+message ZoneResourceStatus {
+    ZoneResourceSpec active_resources = 1;
+}
+
 message ZoneMetricNode {
     string name = 1;
     google.protobuf.Value value = 2;

crates/krata/proto/krata/v1/control.proto

@@ -26,6 +26,8 @@ service ControlService {
 
     rpc GetZone(GetZoneRequest) returns (GetZoneReply);
 
+    rpc UpdateZoneResources(UpdateZoneResourcesRequest) returns (UpdateZoneResourcesReply);
+
     rpc ListZones(ListZonesRequest) returns (ListZonesReply);
 
     rpc AttachZoneConsole(stream ZoneConsoleRequest) returns (stream ZoneConsoleReply);
@@ -242,3 +244,10 @@ message SetHostPowerManagementPolicyRequest {
 }
 
 message SetHostPowerManagementPolicyReply {}
+
+message UpdateZoneResourcesRequest {
+    string zone_id = 1;
+    krata.v1.common.ZoneResourceSpec resources = 2;
+}
+
+message UpdateZoneResourcesReply {}

crates/runtime/src/launch.rs

@@ -30,8 +30,9 @@ pub struct ZoneLaunchRequest {
     pub initrd: Vec<u8>,
     pub uuid: Option<Uuid>,
     pub name: Option<String>,
-    pub vcpus: u32,
-    pub mem: u64,
+    pub cpus: u32,
+    pub target_memory: u64,
+    pub max_memory: u64,
     pub env: HashMap<String, String>,
     pub run: Option<Vec<String>>,
     pub pcis: Vec<PciDevice>,
@@ -194,8 +195,9 @@ impl ZoneLauncher {
 
         let config = DomainConfig {
             base: BaseDomainConfig {
-                max_vcpus: request.vcpus,
-                mem_mb: request.mem,
+                max_vcpus: request.cpus,
+                max_mem_mb: request.max_memory,
+                target_mem_mb: request.target_memory,
                 kernel: request.kernel,
                 initrd: request.initrd,
                 cmdline,

crates/runtime/src/lib.rs

@@ -226,6 +226,36 @@ impl Runtime {
         Ok(uuid)
     }
 
+    pub async fn set_max_memory(&self, domid: u32, max_memory_bytes: u64) -> Result<()> {
+        self.context
+            .xen
+            .call
+            .set_max_mem(domid, max_memory_bytes / 1024)
+            .await?;
+        let domain_path = self.context.xen.store.get_domain_path(domid).await?;
+        let max_memory_path = format!("{}/memory/static-max", domain_path);
+        self.context
+            .xen
+            .store
+            .write_string(max_memory_path, &(max_memory_bytes / 1024).to_string())
+            .await?;
+        Ok(())
+    }
+
+    pub async fn set_target_memory(&self, domid: u32, target_memory_bytes: u64) -> Result<()> {
+        let domain_path = self.context.xen.store.get_domain_path(domid).await?;
+        let target_memory_path = format!("{}/memory/target", domain_path);
+        self.context
+            .xen
+            .store
+            .write_string(
+                target_memory_path,
+                &(target_memory_bytes / 1024).to_string(),
+            )
+            .await?;
+        Ok(())
+    }
+
     pub async fn list(&self) -> Result<Vec<ZoneInfo>> {
         self.context.list().await
     }

crates/xen/xenclient/examples/boot.rs

@@ -27,7 +27,8 @@ async fn main() -> Result<()> {
         base: BaseDomainConfig {
             uuid: Uuid::new_v4(),
             max_vcpus: 1,
-            mem_mb: 512,
+            max_mem_mb: 512,
+            target_mem_mb: 512,
             enable_iommu: true,
             kernel: fs::read(&kernel_image_path).await?,
             initrd: fs::read(&initrd_path).await?,

crates/xen/xenclient/src/tx.rs

@@ -156,13 +156,13 @@ impl ClientTransaction {
         self.tx
             .write_string(
                 format!("{}/memory/static-max", self.dom_path).as_str(),
-                &(base.mem_mb * 1024).to_string(),
+                &(base.max_mem_mb * 1024).to_string(),
             )
             .await?;
         self.tx
             .write_string(
                 format!("{}/memory/target", self.dom_path).as_str(),
-                &(base.mem_mb * 1024).to_string(),
+                &(base.target_mem_mb * 1024).to_string(),
             )
             .await?;
         self.tx

crates/xen/xenplatform/src/boot.rs

@@ -162,11 +162,13 @@ impl<I: BootImageLoader, P: BootSetupPlatform> BootSetup<I, P> {
     pub async fn initialize(
         &mut self,
         initrd: &[u8],
-        mem_mb: u64,
+        target_mem_mb: u64,
+        max_mem_mb: u64,
         max_vcpus: u32,
         cmdline: &str,
     ) -> Result<BootDomain> {
-        let total_pages = mem_mb << (20 - self.platform.page_shift());
+        let target_pages = target_mem_mb << (20 - self.platform.page_shift());
+        let total_pages = max_mem_mb << (20 - self.platform.page_shift());
         let image_info = self.image_loader.parse(self.platform.hvm()).await?;
         let mut domain = BootDomain {
             domid: self.domid,
@@ -175,7 +177,7 @@ impl<I: BootImageLoader, P: BootSetupPlatform> BootSetup<I, P> {
             virt_pgtab_end: 0,
             pfn_alloc_end: 0,
             total_pages,
-            target_pages: total_pages,
+            target_pages,
             page_size: self.platform.page_size(),
             image_info,
             console_evtchn: 0,

crates/xen/xenplatform/src/domain.rs

@@ -29,7 +29,7 @@ impl<P: BootSetupPlatform> BaseDomainManager<P> {
         let domid = self.call.create_domain(domain).await?;
         self.call.set_max_vcpus(domid, config.max_vcpus).await?;
         self.call
-            .set_max_mem(domid, (config.mem_mb * 1024) + 2048)
+            .set_max_mem(domid, (config.max_mem_mb * 1024) + 2048)
             .await?;
         let loader = ElfImageLoader::load_file_kernel(&config.kernel)?;
         let platform = (*self.platform).clone();
@@ -37,7 +37,8 @@ impl<P: BootSetupPlatform> BaseDomainManager<P> {
         let mut domain = boot
             .initialize(
                 &config.initrd,
-                config.mem_mb,
+                config.target_mem_mb,
+                config.max_mem_mb,
                 config.max_vcpus,
                 &config.cmdline,
             )
@@ -63,7 +64,8 @@ pub struct BaseDomainConfig {
     pub uuid: Uuid,
     pub owner_domid: u32,
     pub max_vcpus: u32,
-    pub mem_mb: u64,
+    pub max_mem_mb: u64,
+    pub target_mem_mb: u64,
     pub kernel: Vec<u8>,
     pub initrd: Vec<u8>,
     pub cmdline: String,