qui-domains: adjust shutdown failure actions

Author: Jayant-kernel

qui/tray/domains.py

@@ -192,7 +192,7 @@ async def perform_action(self):
 class ShutdownItem(VMActionMenuItem):
     """Shutdown menu Item. When activated shutdowns the domain."""
 
-    def __init__(self, vm, icon_cache, force=False):
+    def __init__(self, vm, icon_cache, force=False, follow_shift=True):
         if force:
             super().__init__(
                 vm,
@@ -205,8 +205,11 @@ def __init__(self, vm, icon_cache, force=False):
                 vm, label=_("Shutdown"), icon_cache=icon_cache, icon_name="shutdown"
             )
         self.force = force
+        self.follow_shift = follow_shift
 
     def set_force(self, force):
+        if not self.follow_shift:
+            return
         self.force = force
         if self.force:
             self.label.set_text(_("Force shutdown"))
@@ -332,18 +335,47 @@ async def perform_action(self, *_args, **_kwargs):
                     ).format(self.vm.name, str(ex)),
                 )
                 return
+            if isinstance(ex, exc.QubesVMInUseError):
+                title = _("Qube {0} is in use").format(self.vm.name)
+                markup = _(
+                    "The qube <b>{0}</b> could not be shut down because it "
+                    "is in use by connected qubes.\n\n"
+                    "<b>Warning:</b> force shutdown may cause unexpected "
+                    "issues in connected qubes."
+                ).format(self.vm.name)
+                action = "force"
+            elif isinstance(ex, exc.QubesVMShutdownTimeoutError):
+                title = _("Qube {0} shutdown timed out").format(self.vm.name)
+                markup = _(
+                    "The qube <b>{0}</b> did not shut down within the "
+                    "expected time.\n\nYou can retry the shutdown or, "
+                    "if the problem persists, kill the qube."
+                ).format(self.vm.name)
+                action = "timeout"
+            else:
+                title = _("Error restarting qube {0}").format(self.vm.name)
+                markup = _(
+                    "The qube <b>{0}</b> could not be shut down. "
+                    "The following error occurred:\n"
+                    "<tt>{1}</tt>\n\n"
+                    "Would you like to kill the qube?"
+                ).format(self.vm.name, str(ex))
+                action = "kill"
+
             dialog = Gtk.MessageDialog(
-                None, 0, Gtk.MessageType.ERROR, Gtk.ButtonsType.OK_CANCEL
-            )
-            dialog.set_title("Error restarting qube")
-            dialog.set_markup(
-                f"The qube {self.vm.name} couldn't be shut down "
-                "normally. The following error occurred: \n"
-                f"<tt>{str(ex)}</tt>\n\n"
-                "Do you want to force shutdown? \n\n<b>Warning:</b> "
-                "this may cause unexpected issues in connected qubes."
+                None, 0, Gtk.MessageType.ERROR, Gtk.ButtonsType.NONE
             )
-            dialog.connect("response", self.react_to_question)
+            dialog.set_title(title)
+            dialog.set_markup(markup)
+            dialog.add_button(_("Cancel"), Gtk.ResponseType.CANCEL)
+            if action == "force":
+                dialog.add_button(_("Force shutdown"), Gtk.ResponseType.OK)
+            elif action == "timeout":
+                dialog.add_button(_("Retry shutdown"), Gtk.ResponseType.OK)
+                dialog.add_button(_("Kill"), Gtk.ResponseType.YES)
+            elif action == "kill":
+                dialog.add_button(_("Kill"), Gtk.ResponseType.OK)
+            dialog.connect("response", self.react_to_question, action)
             GLib.idle_add(dialog.show)
 
         try:
@@ -366,20 +398,29 @@ async def perform_action(self, *_args, **_kwargs):
                 ).format(self.vm.name, str(ex)),
             )
 
-    def react_to_question(self, widget, response):
-        if response == Gtk.ResponseType.OK:
-            try:
+    def react_to_question(self, widget, response, action):
+        if response not in (Gtk.ResponseType.OK, Gtk.ResponseType.YES):
+            self.give_up = True
+            widget.destroy()
+            return
+        try:
+            if action == "force":
                 self.vm.shutdown(force=True)
-            except exc.QubesException as ex:
-                show_error(
-                    _("Error shutting down qube"),
-                    _(
-                        "The following error occurred while attempting to "
-                        "shut down qube {0}:\n{1}"
-                    ).format(self.vm.name, str(ex)),
-                )
-                self.give_up = True
-        else:
+            elif action == "timeout":
+                if response == Gtk.ResponseType.YES:
+                    self.vm.kill()
+                elif response == Gtk.ResponseType.OK:
+                    self.vm.shutdown(force=False)
+            elif action == "kill" and response == Gtk.ResponseType.OK:
+                self.vm.kill()
+        except exc.QubesException as ex:
+            show_error(
+                _("Error shutting down qube"),
+                _(
+                    "The following error occurred while attempting to "
+                    "shut down qube {0}:\n{1}"
+                ).format(self.vm.name, str(ex)),
+            )
             self.give_up = True
         widget.destroy()
 
@@ -538,7 +579,7 @@ def __init__(self):
 class StartedMenu(Gtk.Menu):
     """The sub-menu for a started domain"""
 
-    def __init__(self, vm, app, icon_cache):
+    def __init__(self, vm, app, icon_cache, shutdown_failed=False):
         super().__init__()
         self.vm = vm
         self.app = app
@@ -552,7 +593,17 @@ def __init__(self, vm, app, icon_cache):
 
         self.add(PreferencesItem(self.vm, icon_cache))
         self.add(PauseItem(self.vm, icon_cache))
-        self.add(ShutdownItem(self.vm, icon_cache, force=app.shift_pressed))
+        self.add(
+            ShutdownItem(
+                self.vm,
+                icon_cache,
+                force=app.shift_pressed and not shutdown_failed,
+                follow_shift=not shutdown_failed,
+            )
+        )
+        if shutdown_failed:
+            self.add(ShutdownItem(self.vm, icon_cache, force=True, follow_shift=False))
+            self.add(KillItem(self.vm, icon_cache))
         if self.vm.klass != "DispVM" or not self.vm.auto_cleanup:
             self.add(RestartItem(self.vm, icon_cache, force=app.shift_pressed))
 
@@ -700,6 +751,7 @@ def __init__(self, vm, app, icon_cache, state=None):
         self.vm = vm
         self.app = app
         self.icon_cache = icon_cache
+        self.shutdown_failed = False
         self.decorator = qui.decorators.DomainDecorator(vm)
 
         # Main horizontal box
@@ -753,7 +805,12 @@ def _set_submenu(self, state):
                 self.vm, self.icon_cache, working_correctly=(state == "Running")
             )
         elif state == "Running":
-            submenu = StartedMenu(self.vm, self.app, self.icon_cache)
+            submenu = StartedMenu(
+                self.vm,
+                self.app,
+                self.icon_cache,
+                shutdown_failed=self.shutdown_failed,
+            )
         elif state == "Paused":
             submenu = PausedMenu(self.vm, self.icon_cache)
         else:
@@ -1166,6 +1223,16 @@ def update_domain_item(self, vm, event, **kwargs):
                 # it's a fragile DispVM
                 state = "Transient"
 
+        if event == "domain-shutdown-failed":
+            item.shutdown_failed = True
+        elif event in (
+            "domain-start",
+            "domain-pre-start",
+            "domain-shutdown",
+            "domain-pre-shutdown",
+        ):
+            item.shutdown_failed = False
+
         item.update_state(state)
 
         if event == "domain-shutdown":