From 2052b322020dd91eb0f187aecd9bddecb58ef41c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?= Date: Wed, 5 Jul 2017 14:12:27 +0200 Subject: [PATCH] events: simplify wait_for_domain_shutdown coroutine 1. Handle timeout externally - using asyncio.wait_for. 2. Add support for waiting for multiple VMs. --- qubesadmin/events/utils.py | 38 ++++++++----------- .../tests/tools/qvm_template_postprocess.py | 17 +++++---- qubesadmin/tools/qvm_template_postprocess.py | 7 ++-- 3 files changed, 29 insertions(+), 33 deletions(-) diff --git a/qubesadmin/events/utils.py b/qubesadmin/events/utils.py index b32c233..f9118b3 100644 --- a/qubesadmin/events/utils.py +++ b/qubesadmin/events/utils.py @@ -32,43 +32,37 @@ class Interrupt(Exception): '''Interrupt events processing''' -def interrupt_on_vm_shutdown(vm, subject, event): +def interrupt_on_vm_shutdown(vms, subject, event): '''Interrupt events processing when given VM was shutdown''' # pylint: disable=unused-argument if event == 'connection-established': - if vm.is_halted(): + if all(vm.is_halted() for vm in vms): + raise Interrupt + elif event == 'domain-shutdown' and subject in vms: + vms.remove(subject) + if not vms: raise Interrupt - elif event == 'domain-shutdown' and vm == subject: - raise Interrupt @asyncio.coroutine -def wait_for_domain_shutdown(vm, timeout, loop=None): +def wait_for_domain_shutdown(vms): ''' Helper function to wait for domain shutdown. This function wait for domain shutdown, but do not initiate the shutdown itself. - :param vm: QubesVM object to wait for shutdown on - :param timeout: Timeout in seconds, use 0 for no timeout - :param loop: asyncio event loop + :param vms: QubesVM object collection to wait for shutdown on ''' - if loop is None: - loop = asyncio.get_event_loop() - events = qubesadmin.events.EventsDispatcher(vm.app) + if not vms: + return + app = list(vms)[0].app + vms = set(vms) + events = qubesadmin.events.EventsDispatcher(app) events.add_handler('domain-shutdown', - functools.partial(interrupt_on_vm_shutdown, vm)) + functools.partial(interrupt_on_vm_shutdown, vms)) events.add_handler('connection-established', - functools.partial(interrupt_on_vm_shutdown, vm)) - events_task = asyncio.ensure_future(events.listen_for_events(), - loop=loop) - if timeout: - # pylint: disable=no-member - loop.call_later(timeout, events_task.cancel) + functools.partial(interrupt_on_vm_shutdown, vms)) try: - yield from events_task - except asyncio.CancelledError: - raise qubesadmin.exc.QubesVMShutdownTimeout( - 'VM %s shutdown timeout expired', vm.name) + yield from events.listen_for_events() except Interrupt: pass diff --git a/qubesadmin/tests/tools/qvm_template_postprocess.py b/qubesadmin/tests/tools/qvm_template_postprocess.py index ceb0775..78d8009 100644 --- a/qubesadmin/tests/tools/qvm_template_postprocess.py +++ b/qubesadmin/tests/tools/qvm_template_postprocess.py @@ -226,6 +226,10 @@ class TC_00_qvm_template_postprocess(qubesadmin.tests.QubesTestCase): self.app.domains.clear_cache() return self.app.domains['test-vm'] + @asyncio.coroutine + def wait_for_shutdown(self, vm): + pass + @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_appmenus') @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_root_img') def test_020_post_install(self, mock_import_root_img, @@ -250,6 +254,7 @@ class TC_00_qvm_template_postprocess(qubesadmin.tests.QubesTestCase): 'qubesadmin.events.utils.wait_for_domain_shutdown') self.addCleanup(patch_domain_shutdown.stop) mock_domain_shutdown = patch_domain_shutdown.start() + mock_domain_shutdown.side_effect = self.wait_for_shutdown else: self.app.expected_calls[ ('test-vm', 'admin.vm.List', None, None)] = \ @@ -267,18 +272,14 @@ class TC_00_qvm_template_postprocess(qubesadmin.tests.QubesTestCase): mock_import_appmenus.assert_called_once_with(self.app.domains[ 'test-vm'], self.source_dir.name) if qubesadmin.tools.qvm_template_postprocess.have_events: - mock_domain_shutdown.assert_called_once_with(self.app.domains[ - 'test-vm'], 60) + mock_domain_shutdown.assert_called_once_with([self.app.domains[ + 'test-vm']]) self.assertEqual(self.app.service_calls, [ ('test-vm', 'qubes.PostInstall', {}), ('test-vm', 'qubes.PostInstall', b''), ]) self.assertAllCalled() - @asyncio.coroutine - def wait_for_shutdown(self, vm, timeout): - pass - @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_appmenus') @mock.patch('qubesadmin.tools.qvm_template_postprocess.import_root_img') def test_021_post_install_reinstall(self, mock_import_root_img, @@ -320,8 +321,8 @@ class TC_00_qvm_template_postprocess(qubesadmin.tests.QubesTestCase): mock_import_appmenus.assert_called_once_with(self.app.domains[ 'test-vm'], self.source_dir.name) if qubesadmin.tools.qvm_template_postprocess.have_events: - mock_domain_shutdown.assert_called_once_with(self.app.domains[ - 'test-vm'], 60) + mock_domain_shutdown.assert_called_once_with([self.app.domains[ + 'test-vm']]) self.assertEqual(self.app.service_calls, [ ('test-vm', 'qubes.PostInstall', {}), ('test-vm', 'qubes.PostInstall', b''), diff --git a/qubesadmin/tools/qvm_template_postprocess.py b/qubesadmin/tools/qvm_template_postprocess.py index 3a380bf..f0af0dc 100644 --- a/qubesadmin/tools/qvm_template_postprocess.py +++ b/qubesadmin/tools/qvm_template_postprocess.py @@ -164,9 +164,10 @@ def call_postinstall_service(vm): if have_events: try: # pylint: disable=no-member - yield from qubesadmin.events.utils.wait_for_domain_shutdown( - vm, qubesadmin.config.defaults['shutdown_timeout']) - except qubesadmin.exc.QubesVMShutdownTimeout: + yield from asyncio.wait_for( + qubesadmin.events.utils.wait_for_domain_shutdown([vm]), + qubesadmin.config.defaults['shutdown_timeout']) + except asyncio.TimeoutError: vm.kill() else: timeout = qubesadmin.config.defaults['shutdown_timeout']