sandbox/_sandboxremote.py: Implement command batching

Execute batched commands in a single shell script.

1 files changed, 69 insertions, 1 deletions

diff --git a/buildstream/sandbox/_sandboxremote.py b/buildstream/sandbox/_sandboxremote.py
index caf76d163..758160219 100644
--- a/buildstream/sandbox/_sandboxremote.py
+++ b/buildstream/sandbox/_sandboxremote.py
@@ -19,12 +19,14 @@
 #        Jim MacArthur <jim.macarthur@codethink.co.uk>
 
 import os
+import shlex
 from urllib.parse import urlparse
 from functools import partial
 
 import grpc
 
-from . import Sandbox
+from . import Sandbox, SandboxCommandError
+from .sandbox import _SandboxBatch
 from ..storage._filebaseddirectory import FileBasedDirectory
 from ..storage._casbaseddirectory import CasBasedDirectory
 from .. import _signals
@@ -265,3 +267,69 @@ class SandboxRemote(Sandbox):
         self.process_job_output(action_result.output_directories, action_result.output_files)
 
         return 0
+
+    def _create_batch(self, main_group, flags, *, collect=None):
+        return _SandboxRemoteBatch(self, main_group, flags, collect=collect)
+
+
+# _SandboxRemoteBatch()
+#
+# Command batching by shell script generation.
+#
+class _SandboxRemoteBatch(_SandboxBatch):
+
+    def __init__(self, sandbox, main_group, flags, *, collect=None):
+        super().__init__(sandbox, main_group, flags, collect=collect)
+
+        self.script = None
+        self.first_command = None
+        self.cwd = None
+        self.env = None
+
+    def execute(self):
+        self.script = ""
+
+        self.main_group.execute(self)
+
+        first = self.first_command
+        if first and self.sandbox.run(['sh', '-c', '-e', self.script], self.flags, cwd=first.cwd, env=first.env) != 0:
+            raise SandboxCommandError("Command execution failed", collect=self.collect)
+
+    def execute_group(self, group):
+        group.execute_children(self)
+
+    def execute_command(self, command):
+        if self.first_command is None:
+            # First command in batch
+            # Initial working directory and environment of script already matches
+            # the command configuration.
+            self.first_command = command
+        else:
+            # Change working directory for this command
+            if command.cwd != self.cwd:
+                self.script += "mkdir -p {}\n".format(command.cwd)
+                self.script += "cd {}\n".format(command.cwd)
+
+            # Update environment for this command
+            for key in self.env.keys():
+                if key not in command.env:
+                    self.script += "unset {}\n".format(key)
+            for key, value in command.env.items():
+                if key not in self.env or self.env[key] != value:
+                    self.script += "export {}={}\n".format(key, shlex.quote(value))
+
+        # Keep track of current working directory and environment
+        self.cwd = command.cwd
+        self.env = command.env
+
+        # Actual command execution
+        cmdline = ' '.join(shlex.quote(cmd) for cmd in command.command)
+        self.script += "(set -ex; {})".format(cmdline)
+
+        # Error handling
+        label = command.label or cmdline
+        quoted_label = shlex.quote("'{}'".format(label))
+        self.script += " || (echo Command {} failed with exitcode $? >&2 ; exit 1)\n".format(quoted_label)
+
+    def execute_call(self, call):
+        raise SandboxError("SandboxRemote does not support callbacks in command batches")