Skip memory balance for preloaded until pause

Balloon is allowed before paused and can be canceled by touching the
file.

For: QubesOS/qubes-issues#1512

Author: ben-grande

qubes/qmemman/systemstate.py

@@ -73,6 +73,18 @@ def init(self) -> None:
     def get_xs_path(self, domid, key) -> str:
         return "/local/domain/" + str(domid) + "/memory/" + key
 
+    def can_membalance(self, domid) -> bool:
+        stop_file = "/var/run/qubes/do-not-membalance-domid-" + str(domid)
+        curr_stop_file = (
+            "/var/run/qubes/do-not-membalance-cancel-current-domid-"
+            + str(domid)
+        )
+        for file in [stop_file, curr_stop_file]:
+            if os.path.isfile(file):
+                self.log.debug("File present: %r" % file)
+                return False
+        return True
+
     def add_domain(self, domid) -> None:
         self.log.debug("add_domain(domid={!r})".format(domid))
         self.dom_dict[domid] = DomainState(domid)
@@ -208,6 +220,7 @@ def inhibit_balloon_up(self) -> None:
             if (
                 dom.mem_actual is not None
                 and dom.mem_actual + 200 * 1024 < dom.last_target
+                and self.can_membalance(domid)
             ):
                 self.log.info(
                     "Preventing balloon up to {}".format(dom.last_target)
@@ -247,7 +260,7 @@ def do_balloon(self, mem_size) -> bool:
             xenfree_ring[ring_slot] = xenfree
             for domid, prev_mem in prev_mem_actual.items():
                 dom = self.dom_dict[domid]
-                if prev_mem == dom.mem_actual:
+                if prev_mem == dom.mem_actual or not self.can_membalance(domid):
                     # domain not responding to memset requests, remove it
                     #  from donors
                     dom.no_progress = True
@@ -289,6 +302,9 @@ def do_balloon_dom(self, dom_memset: dict) -> bool:
         memset_reqs = {}
         for domid, memset in dom_memset.items():
             dom = dom_dict[domid]
+            if not self.can_membalance(domid):
+                succeeded.append(domid)
+                continue
             if memset == 0:
                 # Domain hasn't sent meminfo back to the server, it is still at
                 # the initial amount. pref_mem can't be calculated yet.
@@ -327,6 +343,8 @@ def do_balloon_dom(self, dom_memset: dict) -> bool:
         while True:
             self.refresh_mem_actual(domid_list)
             for domid, dom in dom_dict.items():
+                if not self.can_membalance(domid):
+                    succeeded.append(domid)
                 if domid in succeeded:
                     continue
                 assert isinstance(dom.mem_actual, int)
@@ -346,6 +364,12 @@ def do_balloon_dom(self, dom_memset: dict) -> bool:
                     diff,
                 )
             if all(dom in succeeded for dom in domid_list):
+                for domid in dom_dict.keys():
+                    curr_stop_file = (
+                        "/var/run/qubes/do-not-membalance-cancel-current-domid-"
+                        + str(domid)
+                    )
+                    Path(curr_stop_file).unlink(missing_ok=True)
                 return True
             slow = []
             for domid, dom in dom_dict.items():
@@ -365,6 +389,12 @@ def do_balloon_dom(self, dom_memset: dict) -> bool:
                         )
                         slow.append(domid)
             if any(dom in slow for dom in domid_list):
+                for domid in dom_dict.keys():
+                    curr_stop_file = (
+                        "/var/run/qubes/do-not-membalance-cancel-current-domid-"
+                        + str(domid)
+                    )
+                    Path(curr_stop_file).unlink(missing_ok=True)
                 return False
             for domid, memset in memset_reqs.items():
                 if domid in succeeded:

qubes/vm/dispvm.py

@@ -26,6 +26,7 @@
 import copy
 import subprocess
 from typing import Optional, Any
+from pathlib import Path
 
 import qubes.config
 import qubes.vm.appvm
@@ -622,16 +623,24 @@ async def on_domain_pre_paused(self, event, **kwargs) -> None:
         result = None
         cancelled = False
         self.log.info("Setting qube memory to pref mem")
+        do_not_membalance_file = Path(
+            "/var/run/qubes/do-not-membalance-domid-" + str(self.xid)
+        )
         try:
             # CI uses Python 3.12 and asynchronous iterator requires >=3.13
             # pylint: disable=not-an-iterable
             async for earliest_task in asyncio.as_completed(tasks):
                 await earliest_task
+                do_not_membalance_file.touch(exist_ok=True)
                 if earliest_task == break_task:
                     self.log.info(
                         "Canceling ballooning task, server might continue"
                     )
                     cancelled = True
+                    Path(
+                        "/var/run/qubes/do-not-membalance-cancel-current-domid-"
+                        + str(self.xid)
+                    ).touch(exist_ok=True)
                     qmemman_task.cancel()
                 else:
                     result = qmemman_task.result()
@@ -898,6 +907,9 @@ async def use_preload(self) -> None:
         """
         if not self.is_preload:
             raise qubes.exc.QubesException("Disposable is not preloaded")
+        Path("/var/run/qubes/do-not-membalance-domid-" + str(self.xid)).unlink(
+            missing_ok=True
+        )
         appvm = self.template
         if self.preload_requested:
             self.log.info("Using preloaded qube")