From 79c7392424b41b17c6e03af02530e672273f4414 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Mon, 20 Apr 2020 02:16:48 +0200
Subject: [PATCH] Cache power state when caching is enabled

Power state changes are signaled with events too, so it is possible to
cache it and update/invalidate cache with events.
Additionally, admin.vm.List returns a power state, so the cache can be
populated early. This in particular greatly improves qvm-ls performance -
eliminate admin.vm.CurrentState call at all.

QubesOS/qubes-issues#3293
---
 qubesadmin/app.py             | 52 +++++++++++++++++++++++++++++++++--
 qubesadmin/events/__init__.py |  3 ++
 qubesadmin/tests/app.py       | 27 ++++++++++++++++++
 qubesadmin/vm/__init__.py     | 10 +++++--
 4 files changed, 88 insertions(+), 4 deletions(-)

diff --git a/qubesadmin/app.py b/qubesadmin/app.py
index de83840..041cd1c 100644
--- a/qubesadmin/app.py
+++ b/qubesadmin/app.py
@@ -70,6 +70,12 @@ class VMCollection(object):
             props = props.split(' ')
             new_vm_list[vm_name] = dict(
                 [vm_prop.split('=', 1) for vm_prop in props])
+            # if cache not enabled, drop power state
+            if not self.app.cache_enabled:
+                try:
+                    del new_vm_list[vm_name]['state']
+                except KeyError:
+                    pass
 
         self._vm_list = new_vm_list
         for name, vm in list(self._vm_objects.items()):
@@ -103,9 +109,12 @@ class VMCollection(object):
             # done by 'item not in self' check above, unless blind_mode is
             # enabled
             klass = None
+            power_state = None
             if self._vm_list and item in self._vm_list:
                 klass = self._vm_list[item]['class']
-            self._vm_objects[item] = cls(self.app, item, klass=klass)
+                power_state = self._vm_list[item].get('state')
+            self._vm_objects[item] = cls(self.app, item, klass=klass,
+                                         power_state=power_state)
         return self._vm_objects[item]
 
     def __contains__(self, item):
@@ -598,7 +607,7 @@ class QubesBase(qubesadmin.base.PropertyHolder):
         :param name: name of the property
         :param kwargs: other arguments
         :return: none
-        """ # pylint: disable=unused-argument
+        """  # pylint: disable=unused-argument
         if subject is None:
             subject = self
 
@@ -608,6 +617,45 @@ class QubesBase(qubesadmin.base.PropertyHolder):
         except KeyError:
             pass
 
+    def _update_power_state_cache(self, subject, event, **kwargs):
+        """ Update cached VM power state.
+
+        This method is designed to be hooed as an event handler for:
+        - domain-pre-start
+        - domain-start
+        - domain-shutdown
+        - domain-paused
+        - domain-unpaused
+
+        This is done in :py:class:`qubesadmin.events.EventsDispatcher` class
+        directly, before calling other handlers.
+
+        :param subject: a VM object
+        :param event: name of the event
+        :param kwargs: other arguments
+        :return:
+        """  # pylint: disable=unused-argument,no-self-use
+
+        if not self.app.cache_enabled:
+            return
+
+        if event == 'domain-pre-start':
+            power_state = 'Transient'
+        elif event == 'domain-start':
+            power_state = 'Running'
+        elif event == 'domain-shutdown':
+            power_state = 'Halted'
+        elif event == 'domain-paused':
+            power_state = 'Paused'
+        elif event == 'domain-unpaused':
+            power_state = 'Running'
+        else:
+            # unknown power state change, drop cached power state
+            power_state = None
+
+        # pylint: disable=protected-access
+        subject._power_state_cache = power_state
+
 
 class QubesLocal(QubesBase):
     """Application object communicating through local socket.
diff --git a/qubesadmin/events/__init__.py b/qubesadmin/events/__init__.py
index 17eb9f3..9cd8b08 100644
--- a/qubesadmin/events/__init__.py
+++ b/qubesadmin/events/__init__.py
@@ -215,6 +215,9 @@ class EventsDispatcher(object):
         if event.startswith('property-set:') or \
                 event.startswith('property-reset:'):
             self.app._invalidate_cache(subject, event, **kwargs)
+        elif event in ('domain-pre-start', 'domain-start', 'domain-shutdown',
+                       'domain-paused', 'domain-unpaused'):
+            self.app._update_power_state_cache(subject, event, **kwargs)
 
         handlers = [h_func for h_name, h_func_set in self.handlers.items()
             for h_func in h_func_set
diff --git a/qubesadmin/tests/app.py b/qubesadmin/tests/app.py
index 743819f..4ab6e21 100644
--- a/qubesadmin/tests/app.py
+++ b/qubesadmin/tests/app.py
@@ -131,6 +131,33 @@ class TC_00_VMCollection(qubesadmin.tests.QubesTestCase):
             self.fail('VM not found in collection')
         self.assertAllCalled()
 
+    def test_010_getitem_cache_power_state(self):
+        self.app.cache_enabled = True
+        self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
+            b'0\x00test-vm class=AppVM state=Running\n'
+        try:
+            vm = self.app.domains['test-vm']
+            self.assertEqual(vm.name, 'test-vm')
+            self.assertEqual(vm.klass, 'AppVM')
+            self.assertEqual(vm.get_power_state(), 'Running')
+        except KeyError:
+            self.fail('VM not found in collection')
+        self.assertAllCalled()
+
+    def test_011_getitem_non_cache_power_state(self):
+        self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
+            b'0\x00test-vm class=AppVM state=Running\n'
+        self.app.expected_calls[('test-vm', 'admin.vm.CurrentState', None, None)] = \
+            b'0\x00power_state=Running mem=1024'
+        try:
+            vm = self.app.domains['test-vm']
+            self.assertEqual(vm.name, 'test-vm')
+            self.assertEqual(vm.klass, 'AppVM')
+            self.assertEqual(vm.get_power_state(), 'Running')
+        except KeyError:
+            self.fail('VM not found in collection')
+        self.assertAllCalled()
+
 
 class TC_10_QubesBase(qubesadmin.tests.QubesTestCase):
     def setUp(self):
diff --git a/qubesadmin/vm/__init__.py b/qubesadmin/vm/__init__.py
index 1670cf3..2cbcaeb 100644
--- a/qubesadmin/vm/__init__.py
+++ b/qubesadmin/vm/__init__.py
@@ -52,10 +52,11 @@ class QubesVM(qubesadmin.base.PropertyHolder):
 
     firewall = None
 
-    def __init__(self, app, name, klass=None):
+    def __init__(self, app, name, klass=None, power_state=None):
         super(QubesVM, self).__init__(app, 'admin.vm.property.', name)
         self._volumes = None
         self._klass = klass
+        self._power_state_cache = power_state
         self.log = logging.getLogger(name)
         self.tags = qubesadmin.tags.Tags(self)
         self.features = qubesadmin.features.Features(self)
@@ -181,8 +182,13 @@ class QubesVM(qubesadmin.base.PropertyHolder):
 
         '''
 
+        if self._power_state_cache is not None:
+            return self._power_state_cache
         try:
-            return self._get_current_state()['power_state']
+            power_state = self._get_current_state()['power_state']
+            if self.app.cache_enabled:
+                self._power_state_cache = power_state
+            return power_state
         except qubesadmin.exc.QubesDaemonNoResponseError:
             return 'NA'