Merge pull request #186 from chaostoolkit/exit-shortcut

Add exit functions

Author: Lawouach

CHANGELOG.md

@@ -4,6 +4,22 @@
 
 [Unreleased]: https://github.com/chaostoolkit/chaostoolkit-lib/compare/1.13.1...HEAD
 
+### Added
+
+- Add two functions to programmatically exit the experiment as soon as feasible
+  by the Python VM [#185][185]:
+  * `chaoslib.exit.exit_gracefully`: termintes but abides to the rollback strategy
+  * `chaoslib.exit.exit_ungracefully`: termintes but bypasses rollbacks entirely
+    and does not wait for background actions/probes still running
+  This should mostly be useful to have a harsh way to interrupt an execution
+  and is therefore an advanced concept with undesirable side effects (though
+  Chaos Toolkit tries to do as right as it can).
+
+  CAVEAT: Only works on Unix/Linux systems implementing SIGUSR1/SIGUSR2 signals
+
+[185]: https://github.com/chaostoolkit/chaostoolkit-lib/issues/185
+
+
 ## [1.13.1][] - 2020-09-07
 
 [1.13.1]: https://github.com/chaostoolkit/chaostoolkit-lib/compare/1.13.0...1.13.1

chaoslib/activity.py

@@ -219,6 +219,7 @@ def run_activity(activity: Activity, configuration: Configuration,
     This is an internal function and should probably avoid being called
     outside this package.
     """
+    result = None
     try:
         provider = activity["provider"]
         activity_type = provider["type"]

chaoslib/exit.py

@@ -0,0 +1,144 @@
+"""
+This module is advanced usage and mostly interesting to users who need to
+be able to terminate an experiment's execution as fast as possible and,
+potentially, without much care for cleaning up afterwards.
+
+If you want to interrup an execution but can affort to wait for graceful
+completion (current activty, rollbacks...) it's probably best to rely on
+the control interface.
+
+If you need this utility here, you can simply do as follows:
+
+
+```python
+from chaoslib.exit import exit_ungracefully
+
+def my_probe():
+    # whatever condition comes up that shows you need to terminate asap
+    exit_ungracefully()
+```
+
+Then in your experiment
+
+```json
+"method": [
+    {
+        "type": "probe",
+        "name": "interrupt-when-system-is-unhappy",
+        "background": True,
+        "provider": {
+            "type": "python",
+            "module": "mymod",
+            "func": "my_probe"
+        }
+    }
+    ...
+]
+```
+
+This will start your probe in the background.
+
+WARNING: Only available on Unix/Linux systems.
+"""
+from contextlib import contextmanager
+import inspect
+import os
+import platform
+import signal
+from types import FrameType
+
+from logzero import logger
+
+__all__ = ["exit_gracefully", "exit_ungracefully", "exit_signals"]
+
+
+@contextmanager
+def exit_signals():
+    """
+    Register the handlers for SIGUSR1 and SIGUSR2 signals. Puts back the
+    original handlers when the call ends.
+
+    SIGUSR1 is used to terminate the experiment now while keeping the
+    rollbacks if they were declared.
+    SIGUSR2 is used to terminate the experiment without ever running the
+    rollbacks.
+
+    Generally speaking using signals this way is a bit of an overkill but
+    the Python VM has no other mechanism to interrupt blocking calls.
+
+    WARNING: Only available on Unix/Linux systems.
+    """
+    if hasattr(signal, "SIGUSR1") and hasattr(signal, "SIGUSR2"):
+        # keep a reference to the original handlers
+        sigusr1_handler = signal.signal(signal.SIGUSR1, _leave_now)
+        sigusr2_handler = signal.signal(signal.SIGUSR2, _leave_now)
+        try:
+            yield
+        finally:
+            signal.signal(signal.SIGUSR1, sigusr1_handler)
+            signal.signal(signal.SIGUSR2, sigusr2_handler)
+    else:
+        # On a system that doesn't support SIGUSR signals
+        # not much we can do...
+        logger.debug(
+            "System '{}' does not expose SIGUSR signals".format(
+                platform.platform()))
+        yield
+
+
+def exit_gracefully():
+    """
+    Sends a user signal to the chaostoolkit process which should terminate
+    the current execution immediatly, but gracefully.
+
+    WARNING: Only available on Unix/Linux systems.
+    """
+    if not hasattr(signal, "SIGUSR1"):
+        frames = inspect.getouterframes(inspect.currentframe())
+        info = frames[1]
+        logger.error(
+            "Cannot call 'chaoslib.exit.exit_ungracefully() [{} - line {}] "
+            "as it relies on the SIGUSR1 signal which is not available on "
+            "your platform".format(info.filename, info.lineno))
+        return
+
+    os.kill(os.getpid(), signal.SIGUSR1)
+
+
+def exit_ungracefully():
+    """
+    Sends a user signal to the chaostoolkit process which should terminate
+    the current execution immediatly, but not gracefully.
+
+    This means the rollbacks will not be executed, although controls
+    will be correctly terminated.
+
+    WARNING: Only available on Unix/Linux systems.
+    """
+    if not hasattr(signal, "SIGUSR2"):
+        frames = inspect.getouterframes(inspect.currentframe())
+        info = frames[1]
+        logger.error(
+            "Cannot call 'chaoslib.exit.exit_ungracefully() [{} - line {}] "
+            "as it relies on the SIGUSR2 signal which is not available on "
+            "your platform".format(info.filename, info.lineno))
+        return
+
+    os.kill(os.getpid(), signal.SIGUSR2)
+
+
+###############################################################################
+# Internals
+###############################################################################
+def _leave_now(signum: int, frame: FrameType = None) -> None:
+    """
+    Signal handler only interested in SIGUSR1 and SIGUSR2 to indicate
+    requested termination of the experiment.
+    """
+    logger.warning("Caught signal num: '{}'".format(signum))
+
+    if signum == signal.SIGUSR1:
+        raise SystemExit(20)
+
+    elif signum == signal.SIGUSR2:
+        raise SystemExit(30)

chaoslib/run.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-from concurrent.futures import Future, ThreadPoolExecutor
+from concurrent.futures import Future, ThreadPoolExecutor, TimeoutError
 from datetime import datetime
 import platform
 import threading
@@ -13,6 +13,7 @@
 from chaoslib.control import initialize_controls, controls, cleanup_controls, \
     Control, initialize_global_controls, cleanup_global_controls
 from chaoslib.exceptions import ChaosException, InterruptExecution
+from chaoslib.exit import exit_signals
 from chaoslib.configuration import load_configuration
 from chaoslib.hypothesis import run_steady_state_hypothesis
 from chaoslib.rollback import run_rollbacks
@@ -296,10 +297,11 @@ def run(self, experiment: Experiment,
             journal: Journal = None) -> Journal:
 
         self.configure(experiment, settings)
-        journal = self._run(
-            self.strategy, self.schedule, experiment, journal,
-            self.config, self.secrets, self.settings, experiment_vars,
-            self.event_registry)
+        with exit_signals():
+            journal = self._run(
+                self.strategy, self.schedule, experiment, journal,
+                self.config, self.secrets, self.settings, experiment_vars,
+                self.event_registry)
         return journal
 
     def _run(self, strategy: Strategy, schedule: Schedule,  # noqa: C901
@@ -334,6 +336,7 @@ def _run(self, strategy: Strategy, schedule: Schedule,  # noqa: C901
             "rollbacks", {}).get("strategy", "default")
         logger.info("Rollbacks strategy: {}".format(rollback_strategy))
 
+        exit_gracefully_with_rollbacks = True
         with_ssh = has_steady_state_hypothesis_with_probes(experiment)
         if not with_ssh:
             logger.info(
@@ -374,17 +377,27 @@ def _run(self, strategy: Strategy, schedule: Schedule,  # noqa: C901
                 journal["status"] = "interrupted"
                 logger.fatal(str(i))
                 event_registry.interrupted(experiment, journal)
-            except (KeyboardInterrupt, SystemExit):
+            except KeyboardInterrupt:
                 journal["status"] = "interrupted"
-                logger.warning("Received an exit signal, "
-                               "leaving without applying rollbacks.")
+                logger.warning("Received a termination signal (Ctrl-C)...")
+                event_registry.signal_exit()
+            except SystemExit as x:
+                journal["status"] = "interrupted"
+                logger.warning("Received the exit signal: {}".format(x.code))
+
+                exit_gracefully_with_rollbacks = x.code != 30
+                if not exit_gracefully_with_rollbacks:
+                    logger.warning("Ignoring rollbacks as per signal")
                 event_registry.signal_exit()
             finally:
                 hypo_pool.shutdown(wait=True)
 
-            run_rollback(
-                rollback_strategy, rollback_pool, experiment, journal,
-                configuration, secrets, event_registry, dry)
+            # just in case a signal overrode everything else to tell us not to
+            # play them anyway (see the exit.py module)
+            if exit_gracefully_with_rollbacks:
+                run_rollback(
+                    rollback_strategy, rollback_pool, experiment, journal,
+                    configuration, secrets, event_registry, dry)
 
             journal["end"] = datetime.utcnow().isoformat()
             journal["duration"] = time.time() - started_at
@@ -533,9 +546,8 @@ def run_method(strategy: Strategy, activity_pool: ThreadPoolExecutor,
     try:
         state = apply_activities(
             experiment, configuration, secrets, activity_pool, journal, dry)
-        journal["run"] = state
         event_registry.method_completed(experiment, state)
-        return journal["run"]
+        return state
     except InterruptExecution:
         event_registry.method_completed(experiment)
         raise
@@ -717,27 +729,66 @@ def apply_activities(experiment: Experiment, configuration: Configuration,
                      journal: Journal, dry: bool = False) -> List[Run]:
     with controls(level="method", experiment=experiment, context=experiment,
                   configuration=configuration, secrets=secrets) as control:
+        result = []
         runs = []
-        for run in run_activities(
-                experiment, configuration, secrets, pool, dry):
-            runs.append(run)
-            if journal["status"] in ["aborted", "failed", "interrupted"]:
-                break
+        method = experiment.get("method", [])
+        wait_for_background_activities = True
 
-        if pool:
-            logger.debug("Waiting for background activities to complete...")
-            pool.shutdown(wait=True)
-
-        result = []
-        for run in runs:
-            if not run:
-                continue
-            if isinstance(run, dict):
-                result.append(run)
-            else:
-                result.append(run.result())
+        try:
+            for run in run_activities(
+                    experiment, configuration, secrets, pool, dry):
+                runs.append(run)
+                if journal["status"] in ["aborted", "failed", "interrupted"]:
+                    break
+        except SystemExit as x:
+            # when we got a signal for an ungraceful exit, we can decide
+            # not to wait for background activities. Their statuses will
+            # remain failed.
+            wait_for_background_activities = x.code != 30  # see exit.py
+            raise
+        finally:
+            background_activity_timeout = None
 
-        control.with_state(result)
+            if wait_for_background_activities and pool:
+                logger.debug("Waiting for background activities to complete")
+                pool.shutdown(wait=True)
+            elif pool:
+                logger.debug(
+                    "Do not wait for the background activities to finish "
+                    "as per signal")
+                background_activity_timeout = 0.1
+                pool.shutdown(wait=False)
+
+            for index, run in enumerate(runs):
+                if not run:
+                    continue
+
+                if isinstance(run, dict):
+                    result.append(run)
+                else:
+                    try:
+                        # background activities
+                        result.append(
+                            run.result(timeout=background_activity_timeout))
+                    except TimeoutError:
+                        # we want an entry for the background activity in our
+                        # results anyway, we won't have anything meaningful
+                        # to say about it
+                        result.append({
+                            "activity": method[index],
+                            "status": "failed",
+                            "output": None,
+                            "duration": None,
+                            "start": None,
+                            "end": None,
+                            "exception": None
+                        })
+
+            # now let's ensure the journal has all activities in their correct
+            # order (background ones included)
+            journal["run"] = result
+
+            control.with_state(result)
 
     return result