squash: fix SystemExit handling

tmt/queue.py

@@ -30,6 +30,10 @@ class Task(Generic[TaskResultT]):
     #: is saved in this field.
     exc: Optional[Exception]
 
+    #: If set, the task raised :py:class:`SystemExit` exception, and wants to
+    #: terminate the run completely.
+    requested_exit: Optional[SystemExit]
+
     @property
     def name(self) -> str:
         """
@@ -171,11 +175,14 @@ def go(self) -> Iterator['Self']:
                 try:
                     result = future.result()
 
+                except SystemExit as exc:
+                    task = dataclasses.replace(self, result=None, exc=None, requested_exit=exc)
+
                 except Exception as exc:
-                    task = dataclasses.replace(self, result=None, exc=exc)
+                    task = dataclasses.replace(self, result=None, exc=exc, requested_exit=None)
 
                 else:
-                    task = dataclasses.replace(self, result=result, exc=None)
+                    task = dataclasses.replace(self, result=result, exc=None, requested_exit=None)
 
                 task.guest = guest
 

tmt/steps/__init__.py

@@ -1990,6 +1990,7 @@ def enqueue_action(
             result=None,
             guest=None,
             exc=None,
+            requested_exit=None,
             phase=phase
             ))
 
@@ -2007,6 +2008,7 @@ def enqueue_plugin(
             result=None,
             guest=None,
             exc=None,
+            requested_exit=None,
             guests=guests,
             phase=phase
             ))

tmt/steps/finish/__init__.py

@@ -182,6 +182,7 @@ def go(self) -> None:
                     result=None,
                     guest=None,
                     exc=None,
+                    requested_exit=None,
                     guests=guest_copies,
                     source=self.plan.data_directory
                     ),

tmt/steps/prepare/__init__.py

@@ -259,6 +259,7 @@ def go(self) -> None:
                     result=None,
                     guest=None,
                     exc=None,
+                    requested_exit=None,
                     guests=guest_copies
                     ),
                 self._logger)
@@ -314,6 +315,7 @@ def go(self) -> None:
                     result=None,
                     guest=None,
                     exc=None,
+                    requested_exit=None,
                     guests=guest_copies,
                     source=self.plan.data_directory),
                 self._logger)

tmt/steps/provision/__init__.py

@@ -1845,12 +1845,23 @@ def go(self) -> Iterator['ProvisionTask']:
                 try:
                     future.result()
 
+                except SystemExit as exc:
+                    yield ProvisionTask(
+                        logger=new_logger,
+                        result=None,
+                        guest=None,
+                        exc=None,
+                        requested_exit=exc,
+                        phases=[]
+                        )
+
                 except Exception as exc:
                     yield ProvisionTask(
                         logger=new_logger,
                         result=None,
                         guest=None,
                         exc=exc,
+                        requested_exit=None,
                         phases=[]
                         )
 
@@ -1860,6 +1871,7 @@ def go(self) -> Iterator['ProvisionTask']:
                         result=None,
                         guest=phase.guest(),
                         exc=None,
+                        requested_exit=None,
                         phases=[],
                         phase=phase
                         )
@@ -1881,6 +1893,7 @@ def enqueue(
             result=None,
             guest=None,
             exc=None,
+            requested_exit=None,
             phases=phases
             ))
 
@@ -1988,24 +2001,21 @@ def go(self) -> None:
         # Provision guests
         self._guests = []
 
-        # A plugin will only raise SystemExit if the exit is really desired
-        # and no other actions should be done. An example of this is
-        # listing available images. In such case, the workdir is deleted
-        # as it's redundant and save() would throw an error.
-        # nonsaveable_exceptions = (SystemExit, tmt.utils.SpecificationError)
-        save = True
-
         def _run_provision_phases(
-                phases: list[ProvisionPlugin[ProvisionStepData]]) -> list[ProvisionTask]:
+                phases: list[ProvisionPlugin[ProvisionStepData]]
+                ) -> tuple[list[ProvisionTask], list[ProvisionTask]]:
             queue: ProvisionQueue = ProvisionQueue(
                 'provision',
                 self._logger.descend(logger_name=f'{self}.queue'))
 
             queue.enqueue(phases=phases, logger=queue._logger)
 
+            all_tasks: list[ProvisionTask] = []
             failed_tasks: list[ProvisionTask] = []
 
             for outcome in queue.run():
+                all_tasks.append(outcome)
+
                 if outcome.exc:
                     outcome.logger.fail(str(outcome.exc))
 
@@ -2021,29 +2031,32 @@ def _run_provision_phases(
                     if guest.is_ready or self.is_dry_run:
                         self._guests.append(guest)
 
-            return failed_tasks
+            return all_tasks, failed_tasks
 
-        def _run_action_phases(phases: list[Action]) -> list[ActionTask]:
+        def _run_action_phases(phases: list[Action]) -> tuple[list[ActionTask], list[ActionTask]]:
             queue: PhaseQueue[ProvisionStepData] = PhaseQueue(
                 'finish',
                 self._logger.descend(logger_name=f'{self}.queue'))
 
             for action in phases:
                 queue.enqueue_action(phase=action)
 
+            all_tasks: list[ActionTask] = []
             failed_tasks: list[ActionTask] = []
 
             for outcome in queue.run():
                 assert isinstance(outcome, ActionTask)
 
+                all_tasks.append(outcome)
+
                 if outcome.exc:
                     outcome.logger.fail(str(outcome.exc))
 
                     failed_tasks.append(outcome)
 
                     continue
 
-            return failed_tasks
+            return all_tasks, failed_tasks
 
         # Provisioning phases may be intermixed with actions. To perform all
         # phases and action in a consistent manner, we will process them in
@@ -2053,7 +2066,8 @@ def _run_action_phases(phases: list[Action]) -> list[ActionTask]:
         all_phases = self.phases(classes=(Action, ProvisionPlugin))
         all_phases.sort(key=lambda x: x.order)
 
-        failed_tasks: list[Union[ActionTask, ProvisionTask]] = []
+        all_outcomes: list[Union[ActionTask, ProvisionTask]] = []
+        failed_outcomes: list[Union[ActionTask, ProvisionTask]] = []
 
         while all_phases:
             # Start looking for sequences of phases of the same kind. Collect
@@ -2066,7 +2080,10 @@ def _run_action_phases(phases: list[Action]) -> list[ActionTask]:
                 while all_phases and isinstance(all_phases[0], Action):
                     action_phases.append(cast(Action, all_phases.pop(0)))
 
-                failed_tasks += _run_action_phases(action_phases)
+                _all_action_outcomes, _failed_action_outcomes = _run_action_phases(action_phases)
+
+                all_outcomes += _all_action_outcomes
+                failed_outcomes += _failed_action_outcomes
 
             else:
                 plugin_phases: list[ProvisionPlugin[ProvisionStepData]] = [
@@ -2078,26 +2095,44 @@ def _run_action_phases(phases: list[Action]) -> list[ActionTask]:
                             ProvisionPlugin[ProvisionStepData],
                             all_phases.pop(0)))
 
-                failed_tasks += _run_provision_phases(plugin_phases)
+                _all_plugin_outcomes, _failed_plugin_outcomes = _run_provision_phases(
+                    plugin_phases)
+
+                all_outcomes += _all_plugin_outcomes
+                failed_outcomes += _failed_plugin_outcomes
+
+        # A plugin will only raise SystemExit if the exit is really desired
+        # and no other actions should be done. An example of this is
+        # listing available images. In such case, the workdir is deleted
+        # as it's redundant and save() would throw an error.
+        #
+        # TODO: in theory, there may be many, many plugins raising `SystemExit`
+        # but we can re-raise just a single one. It would be better to not use
+        # an exception to signal this, but rather set/return a special object,
+        # leaving the materialization into `SystemExit` to the step and/or tmt.
+        # Or not do any one-shot actions under the disguise of provisioning...
+        exiting_tasks = [
+            outcome for outcome in all_outcomes if outcome.requested_exit is not None
+            ]
+
+        if exiting_tasks:
+            assert exiting_tasks[0].requested_exit is not None
+
+            raise exiting_tasks[0].requested_exit
 
-        if failed_tasks:
+        if failed_outcomes:
             raise tmt.utils.GeneralError(
                 'provision step failed',
-                causes=[outcome.exc for outcome in failed_tasks if outcome.exc is not None]
+                causes=[outcome.exc for outcome in failed_outcomes if outcome.exc is not None]
                 )
 
-        # if isinstance(provison_phase_outcome.exc, nonsaveable_exceptions):
-        #     save = False
-
         # To separate "provision" from the follow-up logging visually
         self.info('')
 
         # Give a summary, update status and save
         self.summary()
         self.status('done')
-
-        if save:
-            self.save()
+        self.save()
 
     def guests(self) -> list[Guest]:
         """ Return the list of all provisioned guests """