Show CPU/memory usage in Qube manager again

fixes: QubesOS/qubes-issues#7817

Author: alimirjamali

qubesmanager/qube_manager.py

@@ -244,10 +244,18 @@ def __init__(self, vm):
         self.state = {'power': "", 'outdated': ""}
         self.updateable = getattr(vm, 'updateable', False)
         self.update(True)
+        self.CPU_usage = None
+        self.RAM_usage = None
+
+    def update_resource_usage(self, RAM_usage, CPU_usage):
+        self.RAM_usage = RAM_usage
+        self.CPU_usage = CPU_usage
 
     def update_power_state(self):
         try:
             self.state['power'] = self.vm.get_power_state()
+            if self.state["power"] in ["Halted", "Paused", "Suspended"]:
+                self.update_resource_usage(None, None)
             if self.state['power'] == "Halted" and \
                     self.vm.klass != "AdminVM" and \
                     manager_utils.get_feature(
@@ -259,6 +267,7 @@ def update_power_state(self):
         except exc.QubesDaemonAccessError:
             self.state['power'] = ""
 
+
         self.state['outdated'] = ""
         try:
             if manager_utils.is_running(self.vm, False):
@@ -422,6 +431,8 @@ def __init__(self, qubes_cache):
                 "Label",
                 "Name",
                 "State",
+                "CPU",
+                "MEM",
                 "Template",
                 "NetVM",
                 "Disk Usage",
@@ -453,6 +464,12 @@ def data(self, index, role):
         col_name = self.columns_indices[col]
         vm = self.qubes_cache.get_vm(row)
 
+        if role == Qt.ItemDataRole.SizeHintRole:
+            if col_name in ["CPU", "MEM"]:
+                return QSize(100, 22)
+        if role == Qt.ItemDataRole.TextAlignmentRole:
+            if col_name in ["CPU", "MEM", "Disk Usage"]:
+                return Qt.AlignmentFlag.AlignCenter
         if role == Qt.ItemDataRole.DisplayRole:
             if col_name == "Name":
                 return vm.name
@@ -464,6 +481,14 @@ def data(self, index, role):
                 return vm.template
             if col_name == "NetVM":
                 return vm.netvm
+            if col_name == "CPU":
+                if not vm.CPU_usage:
+                    return "-"
+                return vm.CPU_usage + " %"
+            if col_name == "MEM":
+                if not vm.RAM_usage:
+                    return "-"
+                return str(int(int(vm.RAM_usage) / 1024)) + " MiB"
             if col_name == "Disk Usage":
                 return vm.disk
             if col_name == "Internal":
@@ -508,6 +533,13 @@ def data(self, index, role):
         if role == Qt.ItemDataRole.UserRole + 1:
             if vm.klass == 'AdminVM':
                 return ""
+            # Consider allowing dom0 to be sorted for CPU & MEM usage
+            if col_name == "CPU":
+                if vm.CPU_usage:
+                    return int(vm.CPU_usage)
+            if col_name == "MEM":
+                if vm.RAM_usage:
+                    return int(vm.RAM_usage)
             if col_name == "Label":
                 vmtype, vmcolor = vm.icon.split("-", 1)
                 try:
@@ -746,7 +778,14 @@ class VmManagerWindow(ui_qubemanager.Ui_VmManagerWindow, QMainWindow):
     # suppress saving settings while initializing widgets
     settings_loaded = False
 
-    def __init__(self, qt_app, qubes_app, dispatcher, _parent=None):
+    def __init__(
+            self,
+            qt_app,
+            qubes_app,
+            dispatcher,
+            stats_dispatcher=None,
+            _parent=None
+        ):
         # pylint: disable=too-many-statements
         super().__init__()
         self.setupUi(self)
@@ -772,6 +811,7 @@ def __init__(self, qt_app, qubes_app, dispatcher, _parent=None):
 
         self.frame_width = 0
         self.frame_height = 0
+        self.foreground = True
 
         self.init_template_menu()
         self.init_network_menu()
@@ -908,6 +948,9 @@ def __init__(self, qt_app, qubes_app, dispatcher, _parent=None):
         dispatcher.add_handler('domain-feature-delete:skip-update',
                                self.on_domain_updates_available)
 
+        if stats_dispatcher:
+            stats_dispatcher.add_handler("vm-stats", self.on_vm_stats)
+
         self.installEventFilter(self)
 
         # It needs to store threads until they finish
@@ -927,8 +970,10 @@ def eventFilter(self, _object, event):
         if event.type() == QEvent.Type.WindowActivate:
             self.update_running_size()
             self.size_on_disk_timer.setInterval(1000 * 60)
+            self.foreground = True
         elif event.type() == QEvent.Type.WindowDeactivate:
             self.size_on_disk_timer.setInterval(1000 * 60 * 5)
+            self.foreground = False
         return False
 
     def scroll_to_top(self):
@@ -1191,6 +1236,16 @@ def update_running_size(self, *_args):
                 self.qubes_cache.get_vm(qid=vm.qid).update(
                     update_size_on_disk=True, event='disk_size')
 
+    def on_vm_stats(self, vm, _event, **kwargs):
+        if not self.foreground:
+            return
+        domain = self.qubes_app.domains[vm]
+        self.qubes_cache.get_vm(qid=domain.qid).update_resource_usage(
+            kwargs["memory_kb"],
+            kwargs["cpu_usage"],
+        )
+        self.proxy.invalidate()
+
     def on_domain_added(self, _submitter, _event, vm, **_kwargs):
         try:
             domain = self.qubes_app.domains[vm]

qubesmanager/tests/test_qube_manager.py

@@ -23,7 +23,7 @@
 import time
 from datetime import datetime
 
-from PyQt6.QtCore import Qt, QSettings, QItemSelectionModel
+from PyQt6.QtCore import Qt, QSettings, QItemSelectionModel, QEvent
 from PyQt6.QtGui import QPixmap, QIcon
 from PyQt6.QtWidgets import QMessageBox
 
@@ -168,7 +168,12 @@ def test_000_window_loads(qapp, test_qubes_app):
 
 def test_001_model_correctness(qapp, test_qubes_app):
     dispatcher = MockDispatcher(test_qubes_app)
-    qm = qube_manager.VmManagerWindow(qapp, test_qubes_app, dispatcher)
+    qm = qube_manager.VmManagerWindow(
+        qapp,
+        test_qubes_app,
+        dispatcher,
+        stats_dispatcher=dispatcher,
+    )
 
     model = qm.qubes_model
 
@@ -177,6 +182,13 @@ def test_001_model_correctness(qapp, test_qubes_app):
     # number of domains
     assert model.rowCount(None) == len(domains)
 
+    # emulating CPU/MEM usage update for dom0
+    qm.on_vm_stats('dom0', _event=None, memory_kb="67108864", cpu_usage="100")
+    # emulating skipping of CPU/MEM usage update for personal
+    qm.eventFilter(_object=None, event=QEvent(QEvent.Type.WindowDeactivate))
+    qm.on_vm_stats('personal', _event=None, memory_kb="524288", cpu_usage="1")
+    qm.eventFilter(_object=None, event=QEvent(QEvent.Type.WindowActivate))
+
     # domain data
     for row in range(model.rowCount(None)):
         # name
@@ -222,6 +234,24 @@ def test_001_model_correctness(qapp, test_qubes_app):
         if getattr(vm_object, 'internal', False):
             assert internal_data == "Yes"
 
+        # cpu usage
+        cpu_index = model.index(row, model.columns_indices.index(
+            "CPU"))
+        cpu_data = model.data(cpu_index, Qt.ItemDataRole.DisplayRole)
+        if vm_object.klass == "AdminVM":
+            assert str(cpu_data) == "100 %"
+        else:
+            assert str(cpu_data) == "-"
+
+        # mem usage
+        mem_index = model.index(row, model.columns_indices.index(
+            "MEM"))
+        mem_data = model.data(mem_index, Qt.ItemDataRole.DisplayRole)
+        if vm_object.klass == "AdminVM":
+            assert str(mem_data) == "65536 MiB"
+        else:
+            assert str(mem_data) == "-"
+
         # disk usage
         du_index = model.index(row, model.columns_indices.index(
             "Disk Usage"))

qubesmanager/utils.py

@@ -562,10 +562,14 @@ def run_asynchronous(window_class):
     loop = qasync.QEventLoop(qt_app)
     asyncio.set_event_loop(loop)
 
+    stats_dispatcher = events.EventsDispatcher(
+        qubes_app,
+        api_method = "admin.vm.Stats"
+    )
     async def setup():
         dispatcher = events.EventsDispatcher(qubes_app)
 
-        window = window_class(qt_app, qubes_app, dispatcher)
+        window = window_class(qt_app, qubes_app, dispatcher, stats_dispatcher)
 
         if hasattr(window, "setup_application"):
             window.setup_application()
@@ -575,6 +579,7 @@ async def setup():
         await dispatcher.listen_for_events()
 
     try:
+        loop.create_task(stats_dispatcher.listen_for_events())
         loop.run_until_complete(asyncio.ensure_future(setup()))
     except asyncio.CancelledError:
         pass