From 389252f386bb35d42dd7ae22a585c652620e36ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Sun, 25 Jun 2017 21:03:06 +0200
Subject: [PATCH] tools: handle shutdown order in qvm-shutdown

VMs can have runtime dependencies - for example it isn't possible to
shutdown netvm used by some other running VM(s). Since client-side tools
may not have full knowledge about rules enforcing those dependencies
(for example may not have access to 'netvm' property), implement
best-effort approach:
1. Try to shutdown all requested VMs
2. For those where shutdown request succeed, wait for actual shutdown
3. For others - go back to step 1

And loop until all VMs are shutdown, or all shutdown requests fails.
---
 qubesadmin/tools/qvm_shutdown.py | 90 +++++++++++++++++++++++---------
 1 file changed, 64 insertions(+), 26 deletions(-)

diff --git a/qubesadmin/tools/qvm_shutdown.py b/qubesadmin/tools/qvm_shutdown.py
index e8aff02..5d75574 100644
--- a/qubesadmin/tools/qvm_shutdown.py
+++ b/qubesadmin/tools/qvm_shutdown.py
@@ -27,6 +27,13 @@ from __future__ import print_function
 import sys
 import time
 
+import asyncio
+
+try:
+    import qubesadmin.events.utils
+    have_events = True
+except ImportError:
+    have_events = False
 import qubesadmin.tools
 import qubesadmin.exc
 
@@ -52,33 +59,64 @@ parser.add_argument('--timeout',
 def main(args=None, app=None):  # pylint: disable=missing-docstring
     args = parser.parse_args(args, app=app)
 
-    for vm in args.domains:
-        try:
-            vm.shutdown(force=args.force)
-        except qubesadmin.exc.QubesVMNotStartedError:
-            pass
-
-    if not args.wait:
-        return
-
-    timeout = args.timeout
-    current_vms = list(sorted(args.domains))
-    while timeout >= 0:
-        current_vms = [vm for vm in current_vms
-            if vm.get_power_state() != 'Halted']
-        if not current_vms:
-            return 0
-        args.app.log.info('Waiting for shutdown ({}): {}'.format(
-            timeout, ', '.join([str(vm) for vm in current_vms])))
-        time.sleep(1)
-        timeout -= 1
-
-    args.app.log.info(
-        'Killing remaining qubes: {}'
-        .format(', '.join([str(vm) for vm in current_vms])))
-    for vm in current_vms:
-        vm.kill()
+    if have_events:
+        loop = asyncio.get_event_loop()
+    remaining_domains = args.domains
+    for _ in range(len(args.domains)):
+        this_round_domains = set(remaining_domains)
+        if not this_round_domains:
+            break
+        remaining_domains = set()
+        for vm in this_round_domains:
+            try:
+                vm.shutdown()
+            except qubesadmin.exc.QubesVMNotRunningError:
+                pass
+            except qubesadmin.exc.QubesException as e:
+                if not args.wait:
+                    vm.log.error('Shutdown error: {}'.format(e))
+                else:
+                    remaining_domains.add(vm)
+        if not args.wait:
+            return len(remaining_domains)
+        this_round_domains.difference_update(remaining_domains)
+        if not this_round_domains:
+            # no VM shutdown request succeed, no sense to try again
+            break
+        if have_events:
+            try:
+                # pylint: disable=no-member
+                loop.run_until_complete(asyncio.wait_for(
+                    qubesadmin.events.utils.wait_for_domain_shutdown(
+                        this_round_domains),
+                    args.timeout))
+            except asyncio.TimeoutError:
+                for vm in this_round_domains:
+                    vm.kill()
+        else:
+            timeout = args.timeout
+            current_vms = list(sorted(this_round_domains))
+            while timeout >= 0:
+                current_vms = [vm for vm in current_vms
+                    if vm.get_power_state() != 'Halted']
+                if not current_vms:
+                    break
+                args.app.log.info('Waiting for shutdown ({}): {}'.format(
+                    timeout, ', '.join([str(vm) for vm in current_vms])))
+                time.sleep(1)
+                timeout -= 1
+            if current_vms:
+                args.app.log.info(
+                    'Killing remaining qubes: {}'
+                    .format(', '.join([str(vm) for vm in current_vms])))
+            for vm in current_vms:
+                vm.kill()
 
+    if args.wait:
+        if have_events:
+            loop.close()
+        return len([vm for vm in args.domains
+            if vm.get_power_state() != 'Halted'])
 
 if __name__ == '__main__':
     sys.exit(main())