Introduce tblib to handle subprocess exceptions

4 files changed, 38 insertions, 7 deletions

diff --git a/requirements/requirements.in b/requirements/requirements.in
index ce721da33..5a602dccd 100644
--- a/requirements/requirements.in
+++ b/requirements/requirements.in
@@ -8,3 +8,4 @@ ruamel.yaml >= 0.16
 setuptools
 pyroaring
 ujson
+tblib
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 7d3d5d6c7..946535eb8 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -8,6 +8,7 @@ ruamel.yaml==0.16.5
 setuptools==40.8.0
 pyroaring==0.2.8
 ujson==1.35
+tblib==1.5.0
 ## The following requirements were added by pip freeze:
 MarkupSafe==1.1.1
 ruamel.yaml.clib==0.1.2
diff --git a/src/buildstream/_exceptions.py b/src/buildstream/_exceptions.py
index ca17577f7..85fcf61ee 100644
--- a/src/buildstream/_exceptions.py
+++ b/src/buildstream/_exceptions.py
@@ -20,6 +20,7 @@
 
 from enum import Enum, unique
 import os
+import sys
 
 # Disable pylint warnings for whole file here:
 # pylint: disable=global-statement
@@ -239,10 +240,12 @@ class LoadErrorReason(Enum):
 #    reason (LoadErrorReason): machine readable error reason
 #
 # This exception is raised when loading or parsing YAML, or when
-# interpreting project YAML
+# interpreting project YAML. Although reason has a default value,
+# the arg must be assigned to a LoadErrorReason. This is a workaround
+# for unpickling subclassed Exception() classes.
 #
 class LoadError(BstError):
-    def __init__(self, message, reason, *, detail=None):
+    def __init__(self, message, reason=None, *, detail=None):
         super().__init__(message, detail=detail, domain=ErrorDomain.LOAD, reason=reason)
 
 
@@ -385,3 +388,19 @@ class ArtifactElementError(BstError):
 class ProfileError(BstError):
     def __init__(self, message, detail=None, reason=None):
         super().__init__(message, detail=detail, domain=ErrorDomain.PROFILE, reason=reason)
+
+
+# SubprocessException
+#
+# Used with 'tblib.pickling_suport' to pickle the exception & traceback
+# object thrown from subprocessing a Stream entry point, e.g. build().
+# The install() method of pickling_support must be called before attempting
+# to pickle this object.
+#
+class SubprocessException:
+    def __init__(self, exception):
+        self.exception = exception
+        __, __, self.tb = sys.exc_info()
+
+    def re_raise(self):
+        raise self.exception.with_traceback(self.tb)
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index e0c338361..21e225b21 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -33,9 +33,18 @@ import queue
 from contextlib import contextmanager, suppress
 from fnmatch import fnmatch
 from typing import List, Tuple
+from tblib import pickling_support
 
 from ._artifactelement import verify_artifact_ref, ArtifactElement
-from ._exceptions import StreamError, ImplError, BstError, ArtifactElementError, ArtifactError, set_last_task_error
+from ._exceptions import (
+    StreamError,
+    ImplError,
+    BstError,
+    ArtifactElementError,
+    ArtifactError,
+    set_last_task_error,
+    SubprocessException,
+)
 from ._message import Message, MessageType
 from ._scheduler import (
     Scheduler,
@@ -59,7 +68,6 @@ from .plugin import Plugin
 from . import utils, _yaml, _site
 from . import Scope
 
-
 # Stream()
 #
 # This is the main, toplevel calling interface in BuildStream core.
@@ -124,10 +132,12 @@ class Stream:
         # Set main process
         utils._set_stream_pid()
 
+        # Add traceback pickling support
+        pickling_support.install()
         try:
             func(*args, **kwargs)
-        except Exception as e:
-            notify.put(Notification(NotificationType.EXCEPTION, exception=e))
+        except Exception as e:  # pylint: disable=broad-except
+            notify.put(Notification(NotificationType.EXCEPTION, exception=SubprocessException(e)))
 
     def run_in_subprocess(self, func, *args, **kwargs):
         assert not self._subprocess
@@ -1710,7 +1720,7 @@ class Stream:
         elif notification.notification_type == NotificationType.TASK_ERROR:
             set_last_task_error(*notification.task_error)
         elif notification.notification_type == NotificationType.EXCEPTION:
-            raise notification.exception
+            raise notification.exception.re_raise()
         else:
             raise StreamError("Unrecognised notification type received")