QEMU balloon script abstraction improvement

lib/runner.nix

@@ -47,6 +47,8 @@ let
       exit 1
     fi
 
+    # The amount of total memory the guest is requested to release to host, in megabytes.
+    # Expected value between 0 and ${toString microvmConfig.balloonMem}.
     SIZE=$1
     ${hypervisorConfig.setBalloonScript}
   '';

lib/runners/cloud-hypervisor.nix

@@ -216,6 +216,7 @@ in {
     if socket != null
     then ''
       ${pkgs.cloud-hypervisor}/bin/ch-remote --api-socket ${socket} resize --balloon $SIZE"M"
+      ${pkgs.cloud-hypervisor}/bin/ch-remote --api-socket ${socket} info | ${pkgs.jq}/bin/jq '.config.balloon.size / 1024 / 1024 | round'
     ''
     else null;
 

lib/runners/crosvm.nix

@@ -157,6 +157,7 @@ in {
     then ''
       VALUE=$(( $SIZE * 1024 * 1024 ))
       ${pkgs.crosvm}/bin/crosvm balloon $VALUE ${socket}
+      ${pkgs.coreutils}/bin/sleep ${toString microvmConfig.balloonDelay}
       SIZE=$( ${pkgs.crosvm}/bin/crosvm balloon_stats ${socket} | \
         ${pkgs.jq}/bin/jq -r .BalloonStats.balloon_actual \
       )

lib/runners/qemu.nix

@@ -38,7 +38,7 @@ let
   qemu = overrideQemu (if microvmConfig.cpu == null then
     pkgs.qemu_kvm else pkgs.buildPackages.qemu_full);
 
-  inherit (microvmConfig) hostName cpu vcpu mem balloonMem user interfaces shares socket forwardPorts devices vsock graphics storeOnDisk kernel initrdPath storeDisk;
+  inherit (microvmConfig) hostName cpu vcpu mem balloonMem balloonDelay user interfaces shares socket forwardPorts devices vsock graphics storeOnDisk kernel initrdPath storeDisk;
   inherit (microvmConfig.qemu) machine extraArgs serialConsole;
 
   inherit (import ../. { nixpkgs-lib = pkgs.lib; }) withDriveLetters;
@@ -339,16 +339,18 @@ in {
   setBalloonScript =
     if socket != null
     then ''
-      VALUE=$(( $SIZE * 1024 * 1024 ))
+      VALUE=$(( (${toString (mem + balloonMem)} - $SIZE) * 1024 * 1024 ))
       SIZE=$( (
         ${writeQmp { execute = "qmp_capabilities"; }}
-        ${writeQmp { execute = "balloon"; arguments.value = 987; }}
-      ) | sed -e s/987/$VALUE/ | \
+        (${writeQmp { execute = "balloon"; arguments.value = 987; }}) | sed -e s/987/$VALUE/
+        ${pkgs.coreutils}/bin/sleep ${toString balloonDelay}
+        ${writeQmp { execute = "query-balloon"; }}
+      ) | \
         ${pkgs.socat}/bin/socat STDIO UNIX:${socket},shut-none | \
         tail -n 1 | \
-        ${pkgs.jq}/bin/jq -r .data.actual \
+        ${pkgs.jq}/bin/jq -r .return.actual \
       )
-      echo $(( $SIZE / 1024 / 1024 ))
+      echo $(( ${toString (mem + balloonMem)} - $SIZE / 1024 / 1024 ))
     ''
     else null;
 

nixos-modules/microvm/options.nix

@@ -103,6 +103,18 @@ in
       type = types.int;
     };
 
+    balloonDelay = mkOption {
+      description = ''
+        Amount of time to wait for balloon to settle in seconds
+
+        Under some VMMs it takes some time for the reported balloon
+        size to be final. Adding some delay between setting and querying
+        allows microvm-balloon to report correct size.
+      '';
+      default = 2;
+      type = types.int;
+    };
+
     forwardPorts = mkOption {
       type = types.listOf
         (types.submodule {