sandbox/sandbox.py: Add command batching API

This adds the batch() context manager.

3 files changed, 251 insertions, 6 deletions

diff --git a/buildstream/__init__.py b/buildstream/__init__.py
index af2122ef7..fc09f2812 100644
--- a/buildstream/__init__.py
+++ b/buildstream/__init__.py
@@ -27,7 +27,7 @@ if "_BST_COMPLETION" not in os.environ:
     del get_versions
 
     from .utils import UtilError, ProgramNotFoundError
-    from .sandbox import Sandbox, SandboxFlags
+    from .sandbox import Sandbox, SandboxFlags, SandboxCommandError
     from .types import Scope, Consistency
     from .plugin import Plugin
     from .source import Source, SourceError, SourceFetcher
diff --git a/buildstream/sandbox/__init__.py b/buildstream/sandbox/__init__.py
index 5999aba7a..5966d194f 100644
--- a/buildstream/sandbox/__init__.py
+++ b/buildstream/sandbox/__init__.py
@@ -17,6 +17,6 @@
 #  Authors:
 #        Tristan Maat <tristan.maat@codethink.co.uk>
 
-from .sandbox import Sandbox, SandboxFlags
+from .sandbox import Sandbox, SandboxFlags, SandboxCommandError
 from ._sandboxremote import SandboxRemote
 from ._sandboxdummy import SandboxDummy
diff --git a/buildstream/sandbox/sandbox.py b/buildstream/sandbox/sandbox.py
index 7bca70e87..b21b350ff 100644
--- a/buildstream/sandbox/sandbox.py
+++ b/buildstream/sandbox/sandbox.py
@@ -1,5 +1,6 @@
 #
 #  Copyright (C) 2017 Codethink Limited
+#  Copyright (C) 2018 Bloomberg Finance LP
 #
 #  This program is free software; you can redistribute it and/or
 #  modify it under the terms of the GNU Lesser General Public
@@ -29,7 +30,12 @@ See also: :ref:`sandboxing`.
 """
 
 import os
-from .._exceptions import ImplError, BstError
+import shlex
+import contextlib
+from contextlib import contextmanager
+
+from .._exceptions import ImplError, BstError, SandboxError
+from .._message import Message, MessageType
 from ..storage._filebaseddirectory import FileBasedDirectory
 from ..storage._casbaseddirectory import CasBasedDirectory
 
@@ -75,6 +81,19 @@ class SandboxFlags():
     """
 
 
+class SandboxCommandError(SandboxError):
+    """Raised by :class:`.Sandbox` implementations when a command fails.
+
+    Args:
+       message (str): The error message to report to the user
+       collect (str): An optional directory containing partial install contents
+    """
+    def __init__(self, message, *, collect=None):
+        super().__init__(message, reason='command-failed')
+
+        self.collect = collect
+
+
 class Sandbox():
     """Sandbox()
 
@@ -98,6 +117,13 @@ class Sandbox():
         self.__mount_sources = {}
         self.__allow_real_directory = kwargs['allow_real_directory']
 
+        # Plugin ID for logging
+        plugin = kwargs.get('plugin', None)
+        if plugin:
+            self.__plugin_id = plugin._get_unique_id()
+        else:
+            self.__plugin_id = None
+
         # Configuration from kwargs common to all subclasses
         self.__config = kwargs['config']
         self.__stdout = kwargs['stdout']
@@ -125,6 +151,9 @@ class Sandbox():
         # directory via get_directory.
         self._never_cache_vdirs = False
 
+        # Pending command batch
+        self.__batch = None
+
     def get_directory(self):
         """Fetches the sandbox root directory
 
@@ -213,9 +242,16 @@ class Sandbox():
             'artifact': artifact
         })
 
-    def run(self, command, flags, *, cwd=None, env=None):
+    def run(self, command, flags, *, cwd=None, env=None, label=None):
         """Run a command in the sandbox.
 
+        If this is called outside a batch context, the command is immediately
+        executed.
+
+        If this is called in a batch context, the command is added to the batch
+        for later execution. If the command fails, later commands will not be
+        executed. Command flags must match batch flags.
+
         Args:
             command (list): The command to run in the sandboxed environment, as a list
                             of strings starting with the binary to run.
@@ -223,9 +259,10 @@ class Sandbox():
             cwd (str): The sandbox relative working directory in which to run the command.
             env (dict): A dictionary of string key, value pairs to set as environment
                         variables inside the sandbox environment.
+            label (str): An optional label for the command, used for logging. (*Since: 1.4*)
 
         Returns:
-            (int): The program exit code.
+            (int|None): The program exit code, or None if running in batch context.
 
         Raises:
             (:class:`.ProgramNotFoundError`): If a host tool which the given sandbox
@@ -249,7 +286,66 @@ class Sandbox():
         if isinstance(command, str):
             command = [command]
 
-        return self._run(command, flags, cwd=cwd, env=env)
+        if self.__batch:
+            if flags != self.__batch.flags:
+                raise SandboxError("Inconsistent sandbox flags in single command batch")
+
+            batch_command = _SandboxBatchCommand(command, cwd=cwd, env=env, label=label)
+
+            current_group = self.__batch.current_group
+            current_group.append(batch_command)
+            return None
+        else:
+            return self._run(command, flags, cwd=cwd, env=env)
+
+    @contextmanager
+    def batch(self, flags, *, label=None, collect=None):
+        """Context manager for command batching
+
+        This provides a batch context that defers execution of commands until
+        the end of the context. If a command fails, the batch will be aborted
+        and subsequent commands will not be executed.
+
+        Command batches may be nested. Execution will start only when the top
+        level batch context ends.
+
+        Args:
+            flags (:class:`.SandboxFlags`): The flags for this command batch.
+            label (str): An optional label for the batch group, used for logging.
+            collect (str): An optional directory containing partial install contents
+                           on command failure.
+
+        Raises:
+            (:class:`.SandboxCommandError`): If a command fails.
+
+        *Since: 1.4*
+        """
+
+        group = _SandboxBatchGroup(label=label)
+
+        if self.__batch:
+            # Nested batch
+            if flags != self.__batch.flags:
+                raise SandboxError("Inconsistent sandbox flags in single command batch")
+
+            parent_group = self.__batch.current_group
+            parent_group.append(group)
+            self.__batch.current_group = group
+            try:
+                yield
+            finally:
+                self.__batch.current_group = parent_group
+        else:
+            # Top-level batch
+            batch = self._create_batch(group, flags, collect=collect)
+
+            self.__batch = batch
+            try:
+                yield
+            finally:
+                self.__batch = None
+
+            batch.execute()
 
     #####################################################
     #    Abstract Methods for Sandbox implementations   #
@@ -274,6 +370,20 @@ class Sandbox():
         raise ImplError("Sandbox of type '{}' does not implement _run()"
                         .format(type(self).__name__))
 
+    # _create_batch()
+    #
+    # Abstract method for creating a batch object. Subclasses can override
+    # this method to instantiate a subclass of _SandboxBatch.
+    #
+    # Args:
+    #    main_group (:class:`_SandboxBatchGroup`): The top level batch group.
+    #    flags (:class:`.SandboxFlags`): The flags for commands in this batch.
+    #    collect (str): An optional directory containing partial install contents
+    #                   on command failure.
+    #
+    def _create_batch(self, main_group, flags, *, collect=None):
+        return _SandboxBatch(self, main_group, flags, collect=collect)
+
     ################################################
     #               Private methods                #
     ################################################
@@ -422,3 +532,138 @@ class Sandbox():
                 return True
 
         return False
+
+    # _get_plugin_id()
+    #
+    # Get the plugin's unique identifier
+    #
+    def _get_plugin_id(self):
+        return self.__plugin_id
+
+    # _callback()
+    #
+    # If this is called outside a batch context, the specified function is
+    # invoked immediately.
+    #
+    # If this is called in a batch context, the function is added to the batch
+    # for later invocation.
+    #
+    # Args:
+    #    callback (callable): The function to invoke
+    #
+    def _callback(self, callback):
+        if self.__batch:
+            batch_call = _SandboxBatchCall(callback)
+
+            current_group = self.__batch.current_group
+            current_group.append(batch_call)
+        else:
+            callback()
+
+
+# _SandboxBatch()
+#
+# A batch of sandbox commands.
+#
+class _SandboxBatch():
+
+    def __init__(self, sandbox, main_group, flags, *, collect=None):
+        self.sandbox = sandbox
+        self.main_group = main_group
+        self.current_group = main_group
+        self.flags = flags
+        self.collect = collect
+
+    def execute(self):
+        self.main_group.execute(self)
+
+    def execute_group(self, group):
+        if group.label:
+            context = self.sandbox._get_context()
+            cm = context.timed_activity(group.label, unique_id=self.sandbox._get_plugin_id())
+        else:
+            cm = contextlib.suppress()
+
+        with cm:
+            group.execute_children(self)
+
+    def execute_command(self, command):
+        if command.label:
+            context = self.sandbox._get_context()
+            message = Message(self.sandbox._get_plugin_id(), MessageType.STATUS,
+                              'Running {}'.format(command.label))
+            context.message(message)
+
+        exitcode = self.sandbox._run(command.command, self.flags, cwd=command.cwd, env=command.env)
+        if exitcode != 0:
+            cmdline = ' '.join(shlex.quote(cmd) for cmd in command.command)
+            label = command.label or cmdline
+            raise SandboxCommandError("Command '{}' failed with exitcode {}".format(label, exitcode),
+                                      collect=self.collect)
+
+    def execute_call(self, call):
+        call.callback()
+
+
+# _SandboxBatchItem()
+#
+# An item in a command batch.
+#
+class _SandboxBatchItem():
+
+    def __init__(self, *, label=None):
+        self.label = label
+
+
+# _SandboxBatchCommand()
+#
+# A command item in a command batch.
+#
+class _SandboxBatchCommand(_SandboxBatchItem):
+
+    def __init__(self, command, *, cwd, env, label=None):
+        super().__init__(label=label)
+
+        self.command = command
+        self.cwd = cwd
+        self.env = env
+
+    def execute(self, batch):
+        batch.execute_command(self)
+
+
+# _SandboxBatchGroup()
+#
+# A group in a command batch.
+#
+class _SandboxBatchGroup(_SandboxBatchItem):
+
+    def __init__(self, *, label=None):
+        super().__init__(label=label)
+
+        self.children = []
+
+    def append(self, item):
+        self.children.append(item)
+
+    def execute(self, batch):
+        batch.execute_group(self)
+
+    def execute_children(self, batch):
+        for item in self.children:
+            item.execute(batch)
+
+
+# _SandboxBatchCall()
+#
+# A call item in a command batch.
+#
+class _SandboxBatchCall(_SandboxBatchItem):
+
+    def __init__(self, callback):
+        super().__init__()
+
+        self.callback = callback
+
+    def execute(self, batch):
+        batch.execute_call(self)