From 3930269125e26587250c8f3d6dff839b97285299 Mon Sep 17 00:00:00 2001
From: Ben Grande <ben@invisiblethingslab.com>
Date: Tue, 7 Apr 2026 00:37:19 +0200
Subject: [PATCH] Close qmemman client with correct handler

QMemmanClient is instantiated from the same method that is calling it
as a task, so it guarantees access to the close() method, else, it may
not be able to close the connection properly if the task is cancelled,
as the result() will not contain the instance, but CancelledError.

For: https://github.com/QubesOS/qubes-issues/issues/1512
---
 qubes/vm/dispvm.py  | 29 +++++++++++++++++++++--------
 qubes/vm/qubesvm.py | 21 ---------------------
 2 files changed, 21 insertions(+), 29 deletions(-)

diff --git a/qubes/vm/dispvm.py b/qubes/vm/dispvm.py
index 2f8edcd4f..a3411226b 100644
--- a/qubes/vm/dispvm.py
+++ b/qubes/vm/dispvm.py
@@ -611,32 +611,45 @@ async def on_domain_pre_paused(self, event, **kwargs) -> None:
             return
         if self.preload_requested:
             return
-        qmemman_client = None
         break_task = asyncio.create_task(self.preload_requested_event.wait())
-        qmemman_task = asyncio.get_event_loop().run_in_executor(
-            None, self.set_mem
+        qmemman_client = qubes.qmemman.client.QMemmanClient()
+        qmemman_task = asyncio.get_running_loop().run_in_executor(
+            None,
+            qmemman_client.set_mem,  # type: ignore[arg-type]
+            {self.xid: 0},
         )
-        tasks = [break_task, qmemman_task]
+        tasks: list = [break_task, qmemman_task]
+        result = None
+        cancelled = False
+        self.log.info("Setting qube memory to pref mem")
         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
                 if earliest_task == break_task:
+                    self.log.info(
+                        "Canceling ballooning task, server might continue"
+                    )
+                    cancelled = True
                     qmemman_task.cancel()
                 else:
+                    result = qmemman_task.result()
                     break_task.cancel()
         except asyncio.CancelledError:
-            if qmemman_client:
-                qmemman_client.close()
+            pass
+        except IOError as e:
+            raise IOError("Failed to connect to qmemman: {!s}".format(e))
         except Exception as exc:
             self.log.warning(
                 "Preload memory request before pause failed: %s", str(exc)
             )
-            if qmemman_client:
-                qmemman_client.close()
             raise
         finally:
+            if qmemman_client.sock:
+                qmemman_client.close()
+            if not result or cancelled:
+                self.log.warning("Failed to set memory")
             if self.preload_requested:
                 raise qubes.exc.QubesVMCancelledPauseError(
                     self,
diff --git a/qubes/vm/qubesvm.py b/qubes/vm/qubesvm.py
index 9c8d67cca..ee0124120 100644
--- a/qubes/vm/qubesvm.py
+++ b/qubes/vm/qubesvm.py
@@ -2097,27 +2097,6 @@ def request_mem(self, mem_required=None):
 
         return qmemman_client
 
-    def set_mem(self):
-        """
-        Balloon qube to preferred memory, which is just enough for it to be
-        running without problems.
-        """
-        if not qmemman_present or self.maxmem == 0:
-            return None
-
-        self.log.info("Setting qube memory to pref mem")
-        qmemman_client = qubes.qmemman.client.QMemmanClient()
-        try:
-            result = qmemman_client.set_mem({self.xid: 0})
-        except IOError as e:
-            raise IOError("Failed to connect to qmemman: {!s}".format(e))
-
-        if not result:
-            qmemman_client.close()
-            self.log.warning("Failed to set memory")
-
-        return qmemman_client
-
     @staticmethod
     async def start_daemon(*command, input=None, **kwargs):
         """Start a daemon for the VM