Adjust code for possibly coroutine Volume.export() and Volume.export_end()

Now Volume.export() may be a coroutine and also may be accompanied by
Volume.export_end() cleaning up after it.

See previous commits for building blocks for this.

This commit adjusts usage of Volume.export() and adds matching
Volume.export_end() throughout the code base.

Fixes QubesOS/qubes-issues#5935

qubes/storage/__init__.py

@@ -203,6 +203,9 @@ class Volume:
             volume data. If extracting volume data require something more
             than just reading from file (for example connecting to some other
             domain, or decompressing the data), the returned path may be a pipe.
+
+            This can be implemented as a coroutine.
+
         '''
         raise self._not_implemented("export")
 
@@ -646,14 +649,14 @@ class Storage:
                     for target in parsed_xml.xpath(
                         "//domain/devices/disk/target")}
 
+    @asyncio.coroutine
     def export(self, volume):
         ''' Helper function to export volume (pool.export(volume))'''
         assert isinstance(volume, (Volume, str)), \
             "You need to pass a Volume or pool name as str"
-        if isinstance(volume, Volume):
-            return volume.export()
-
-        return self.vm.volumes[volume].export()
+        if not isinstance(volume, Volume):
+            volume = self.vm.volumes[volume]
+        return (yield from qubes.utils.coro_maybe(volume.export()))
 
     @asyncio.coroutine
     def export_end(self, volume, export_path):

qubes/storage/file.py

@@ -21,7 +21,7 @@
 #
 
 ''' This module contains pool implementations backed by file images'''
-
+import asyncio
 import os
 import os.path
 import re
@@ -29,6 +29,7 @@ import subprocess
 from contextlib import suppress
 
 import qubes.storage
+import qubes.utils
 
 BLKSIZE = 512
 
@@ -268,6 +269,7 @@ class FileVolume(qubes.storage.Volume):
         #  if domain is running
         return self.path
 
+    @asyncio.coroutine
     def import_volume(self, src_volume):
         if src_volume.snap_on_start:
             raise qubes.storage.StoragePoolException(
@@ -275,7 +277,11 @@ class FileVolume(qubes.storage.Volume):
                     src_volume, self))
         if self.save_on_stop:
             _remove_if_exists(self.path)
-            copy_file(src_volume.export(), self.path)
+            path = yield from qubes.utils.coro_maybe(src_volume.export())
+            try:
+                copy_file(path, self.path)
+            finally:
+                yield from qubes.utils.coro_maybe(src_volume.export_end(path))
         return self
 
     def import_data(self, size):

qubes/storage/lvm.py

@@ -518,15 +518,19 @@ class ThinVolume(qubes.storage.Volume):
                    self._vid_import.split('/')[1],
                    str(src_volume.size)]
             yield from qubes_lvm_coro(cmd, self.log)
-            src_path = src_volume.export()
-            cmd = ['dd', 'if=' + src_path, 'of=/dev/' + self._vid_import,
-                'conv=sparse', 'status=none', 'bs=128K']
-            if not os.access('/dev/' + self._vid_import, os.W_OK) or \
-                    not os.access(src_path, os.R_OK):
-                cmd.insert(0, 'sudo')
-
-            p = yield from asyncio.create_subprocess_exec(*cmd)
-            yield from p.wait()
+            src_path = yield from qubes.utils.coro_maybe(src_volume.export())
+            try:
+                cmd = ['dd', 'if=' + src_path, 'of=/dev/' + self._vid_import,
+                    'conv=sparse', 'status=none', 'bs=128K']
+                if not os.access('/dev/' + self._vid_import, os.W_OK) or \
+                        not os.access(src_path, os.R_OK):
+                    cmd.insert(0, 'sudo')
+
+                p = yield from asyncio.create_subprocess_exec(*cmd)
+                yield from p.wait()
+            finally:
+                yield from qubes.utils.coro_maybe(
+                    src_volume.export_end(src_path))
             if p.returncode != 0:
                 cmd = ['remove', self._vid_import]
                 yield from qubes_lvm_coro(cmd, self.log)

qubes/storage/reflink.py

@@ -35,6 +35,7 @@ import tempfile
 from contextlib import contextmanager, suppress
 
 import qubes.storage
+import qubes.utils
 
 FICLONE = 1074041865        # defined in <linux/fs.h>, assuming sizeof(int)==4
 LOOP_SET_CAPACITY = 0x4C07  # defined in <linux/loop.h>
@@ -297,15 +298,22 @@ class ReflinkVolume(qubes.storage.Volume):
         _import_data_end))
 
     @qubes.storage.Volume.locked
-    @_coroutinized
+    @asyncio.coroutine
     def import_volume(self, src_volume):
         if self.save_on_stop:
             try:
                 success = False
-                _copy_file(src_volume.export(), self._path_import)
+                src_path = yield from qubes.utils.coro_maybe(
+                    src_volume.export())
+                try:
+                    yield from _coroutinized(_copy_file)(
+                        src_path, self._path_import)
+                finally:
+                    yield from qubes.utils.coro_maybe(
+                        src_volume.export_end(src_path))
                 success = True
             finally:
-                self._import_data_end(success)
+                yield from _coroutinized(self._import_data_end)(success)
         return self
 
     def _path_revision(self, number, timestamp=None):

qubes/tests/integ/backup.py

@@ -36,6 +36,7 @@ import qubes.exc
 import qubes.storage.lvm
 import qubes.tests
 import qubes.tests.storage_lvm
+import qubes.utils
 import qubes.vm
 import qubes.vm.appvm
 import qubes.vm.templatevm
@@ -106,6 +107,13 @@ class BackupTestsMixin(object):
 
         f.close()
 
+    def fill_image_vm(self, vm, volume, size=None, sparse=False):
+        path = self.loop.run_until_complete(vm.storage.export(volume))
+        try:
+            self.fill_image(path, size=size, sparse=sparse)
+        finally:
+            self.loop.run_until_complete(vm.storage.export_end(volume, path))
+
     # NOTE: this was create_basic_vms
     def create_backup_vms(self, pool=None):
         template = self.app.default_template
@@ -120,7 +128,7 @@ class BackupTestsMixin(object):
             testnet.create_on_disk(pool=pool))
         testnet.features['service.ntpd'] = True
         vms.append(testnet)
-        self.fill_image(testnet.storage.export('private'), 20*1024*1024)
+        self.fill_image_vm(testnet, 'private', 20*1024*1024)
 
         vmname = self.make_vm_name('test1')
         self.log.debug("Creating %s" % vmname)
@@ -130,7 +138,7 @@ class BackupTestsMixin(object):
         self.loop.run_until_complete(
             testvm1.create_on_disk(pool=pool))
         vms.append(testvm1)
-        self.fill_image(testvm1.storage.export('private'), 100 * 1024 * 1024)
+        self.fill_image_vm(testvm1, 'private', 100 * 1024 * 1024)
 
         vmname = self.make_vm_name('testhvm1')
         self.log.debug("Creating %s" % vmname)
@@ -140,8 +148,7 @@ class BackupTestsMixin(object):
                                       label='red')
         self.loop.run_until_complete(
             testvm2.create_on_disk(pool=pool))
-        self.fill_image(testvm2.storage.export('root'), 1024 * 1024 * 1024, \
-            True)
+        self.fill_image_vm(testvm2, 'root', 1024 * 1024 * 1024, True)
         vms.append(testvm2)
 
         vmname = self.make_vm_name('template')
@@ -150,7 +157,7 @@ class BackupTestsMixin(object):
             name=vmname, label='red')
         self.loop.run_until_complete(
             testvm3.create_on_disk(pool=pool))
-        self.fill_image(testvm3.storage.export('root'), 100 * 1024 * 1024, True)
+        self.fill_image_vm(testvm3, 'root', 100 * 1024 * 1024, True)
         vms.append(testvm3)
 
         vmname = self.make_vm_name('custom')
@@ -262,11 +269,14 @@ class BackupTestsMixin(object):
             for name, volume in vm.volumes.items():
                 if not volume.rw or not volume.save_on_stop:
                     continue
-                vol_path = volume.export()
+                vol_path = self.loop.run_until_complete(
+                    qubes.utils.coro_maybe(volume.export()))
                 hasher = hashlib.sha1()
                 with open(vol_path, 'rb') as afile:
                     for buf in iter(lambda: afile.read(4096000), b''):
                         hasher.update(buf)
+                self.loop.run_until_complete(
+                    qubes.utils.coro_maybe(volume.export_end(vol_path)))
                 hashes[vm.name][name] = hasher.hexdigest()
         return hashes
 
@@ -382,9 +392,9 @@ class TC_00_Backup(BackupTestsMixin, qubes.tests.SystemTestCase):
             self.fill_image(
                 os.path.join(self.hvmtemplate.dir_path, '00file'),
                 195 * 1024 * 1024 - 4096 * 3)
-            self.fill_image(self.hvmtemplate.storage.export('private'),
+            self.fill_image_vm(self.hvmtemplate, 'private',
                             195 * 1024 * 1024 - 4096 * 3)
-            self.fill_image(self.hvmtemplate.storage.export('root'), 1024 * 1024 * 1024,
+            self.fill_image_vm(self.hvmtemplate, 'root', 1024 * 1024 * 1024,
                             sparse=True)
             vms.append(self.hvmtemplate)
             self.app.save()

qubes/tests/integ/basic.py

@@ -518,9 +518,14 @@ class TC_03_QvmRevertTemplateChanges(qubes.tests.SystemTestCase):
         self.app.save()
 
     def get_rootimg_checksum(self):
-        return subprocess.check_output(
-            ['sha1sum', self.test_template.volumes['root'].export()]).\
-            decode().split(' ')[0]
+        path = self.loop.run_until_complete(
+            self.test_template.storage.export('root'))
+        try:
+            return subprocess.check_output(['sha1sum', path]).\
+                decode().split(' ')[0]
+        finally:
+            self.loop.run_until_complete(
+                self.test_template.storage.export_end('root', path))
 
     def _do_test(self):
         checksum_before = self.get_rootimg_checksum()