Cache block devices

Author: ben-grande

qubes/ext/block.py

@@ -233,10 +233,8 @@ def _is_attached_to(self, vm):
 
         for disk in xml_desc.findall("devices/disk[source][@type='block']"):
             backend_domain, port_id = _try_get_block_device_info(vm.app, disk)
-
             if backend_domain.name != self.backend_domain.name:
                 continue
-
             if self.port_id == port_id:
                 return True
         return False
@@ -307,6 +305,11 @@ def _try_get_block_device_info(app, disk):
 
 class BlockDeviceExtension(qubes.ext.Extension):
 
+    def __init__(self):
+        super().__init__()
+        self._attachments = {}
+        self._attachments_with_options = {}
+
     @qubes.ext.handler("domain-init", "domain-load")
     def on_domain_init_load(self, vm, event):
         """Initialize watching for changes"""
@@ -327,6 +330,26 @@ def on_domain_init_load(self, vm, event):
         else:
             self.devices_cache[vm.name] = {}
 
+    @qubes.ext.handler("device-attach:block")
+    def on_device_block_attach(self, vm, event, device, options):
+        # pylint: disable=unused-argument
+        self._attachments[device.port_id] = device.backend_domain
+
+    @qubes.ext.handler("device-detach:block")
+    def on_device_block_detach(self, vm, event, port):
+        # pylint: disable=unused-argument
+        if port in self._attachments:
+            self._attachments.pop(port)
+        if vm not in self._attachments_with_options:
+            return
+        if port not in self._attachments_with_options[vm]:
+            return
+        for index, (device, options) in enumerate(
+            self._attachments_with_options[vm]
+        ):
+            if device.port_id == port:
+                self._attachments_with_options[vm].pop(device)
+
     @qubes.ext.handler("domain-qdb-change:/qubes-block-devices")
     def on_qdb_change(self, vm, event, path):
         """A change in QubesDB means a change in a device list."""
@@ -338,8 +361,7 @@ def on_qdb_change(self, vm, event, path):
         )
         utils.device_list_change(self, current_devices, vm, path, BlockDevice)
 
-    @staticmethod
-    def get_device_attachments(vm_):
+    def get_device_attachments(self, vm_):
         result = {}
         for vm in vm_.app.domains:
             if not vm.is_running() or isinstance(vm, RemoteVM):
@@ -348,7 +370,10 @@ def get_device_attachments(vm_):
             if vm.app.vmm.offline_mode:
                 return result
 
-            # TODO: ben: cache
+            if vm in self._attachments:
+                result = self._attachments[vm]
+                continue
+
             xml_desc = lxml.etree.fromstring(vm.libvirt_domain.XMLDesc())
 
             for disk in xml_desc.findall("devices/disk[source][@type='block']"):
@@ -359,6 +384,7 @@ def get_device_attachments(vm_):
                     continue
 
                 result[port_id] = vm
+            self._attachments[vm] = result
         return result
 
     @staticmethod
@@ -422,10 +448,14 @@ def on_device_list_attached(self, vm, event, **kwargs):
         if not vm.is_running():
             return
 
+        if vm in self._attachments_with_options:
+            yield from self._attachments_with_options[vm]
+            return
+        self._attachments_with_options[vm] = []
+
         system_disks = SYSTEM_DISKS
         if getattr(vm, "kernel", None):
             system_disks = SYSTEM_DISKS_DOM0_KERNEL
-        # TODO: ben: cache
         xml_desc = lxml.etree.fromstring(vm.libvirt_domain.XMLDesc())
 
         for disk in xml_desc.findall(
@@ -450,7 +480,10 @@ def on_device_list_attached(self, vm, event, **kwargs):
             if devtype != "disk":
                 options["devtype"] = devtype
 
-            yield BlockDevice(Port(backend_domain, port_id, "block")), options
+            dev = BlockDevice(Port(backend_domain, port_id, "block")), options
+            yield dev
+            self._attachments_with_options[vm].append(dev)
+
 
     @staticmethod
     def _is_block(vm, block: str) -> bool: