_context.py: Add disable_fork() method

Calling disable_fork() will prevent the scheduler from running but will
allow communication with casd in the main process.

3 files changed, 24 insertions, 1 deletions

diff --git a/src/buildstream/_cas/cascache.py b/src/buildstream/_cas/cascache.py
index 7bfe7f856..5280615b3 100644
--- a/src/buildstream/_cas/cascache.py
+++ b/src/buildstream/_cas/cascache.py
@@ -79,6 +79,7 @@ class CASCache():
 
         self._casd_channel = None
         self._local_cas = None
+        self._fork_disabled = False
 
     def __getstate__(self):
         state = self.__dict__.copy()
@@ -95,7 +96,7 @@ class CASCache():
 
         if not self._local_cas:
             # gRPC doesn't support fork without exec, which is used in the main process.
-            assert not utils._is_main_process()
+            assert self._fork_disabled or not utils._is_main_process()
 
             self._casd_channel = grpc.insecure_channel('unix:' + self._casd_socket_path)
             self._local_cas = local_cas_pb2_grpc.LocalContentAddressableStorageStub(self._casd_channel)
@@ -128,6 +129,14 @@ class CASCache():
         if not (os.path.isdir(headdir) and os.path.isdir(objdir)):
             raise CASCacheError("CAS repository check failed for '{}'".format(self.casdir))
 
+    # notify_fork_disabled():
+    #
+    # Called by Context when fork() is disabled. This will enable communication
+    # with casd via gRPC in the main process.
+    #
+    def notify_fork_disabled(self):
+        self._fork_disabled = True
+
     # release_resources():
     #
     # Release resources used by CASCache.
diff --git a/src/buildstream/_context.py b/src/buildstream/_context.py
index 5aa1f595f..211ee1cae 100644
--- a/src/buildstream/_context.py
+++ b/src/buildstream/_context.py
@@ -150,6 +150,8 @@ class Context():
         # Whether file contents are required for all artifacts in the local cache
         self.require_artifact_files = True
 
+        self.fork_allowed = True
+
         # Whether elements must be rebuilt when their dependencies have changed
         self._strict_build_plan = None
 
@@ -487,3 +489,13 @@ class Context():
         if self._cascache is None:
             self._cascache = CASCache(self.cachedir)
         return self._cascache
+
+    # disable_fork():
+    #
+    # This will prevent the scheduler from running but will allow communication
+    # with casd in the main process.
+    #
+    def disable_fork(self):
+        self.fork_allowed = False
+        cascache = self.get_cascache()
+        cascache.notify_fork_disabled()
diff --git a/src/buildstream/_scheduler/scheduler.py b/src/buildstream/_scheduler/scheduler.py
index b28b26f0b..bd76e00b1 100644
--- a/src/buildstream/_scheduler/scheduler.py
+++ b/src/buildstream/_scheduler/scheduler.py
@@ -115,6 +115,8 @@ class Scheduler():
     #
     def run(self, queues):
 
+        assert self.context.fork_allowed
+
         # Hold on to the queues to process
         self.queues = queues