[buildstream] 03/17: Introduce tblib to handle subprocess exceptions

[no...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

not-in-ldap pushed a commit to branch tpollard/buildsubprocess
in repository https://gitbox.apache.org/repos/asf/buildstream.git

commit 9b38304326b32b93859721b5e0d31e169d97d09e
Author: Tom Pollard <to...@codethink.co.uk>
AuthorDate: Thu Oct 31 12:16:18 2019 +0000

    Introduce tblib to handle subprocess exceptions
---
 requirements/requirements.in   |  1 +
 requirements/requirements.txt  |  1 +
 src/buildstream/_exceptions.py | 23 +++++++++++++++++++++--
 src/buildstream/_stream.py     | 20 +++++++++++++++-----
 4 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/requirements/requirements.in b/requirements/requirements.in
index ce721da..5a602dc 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 9620908..c9c2ea6 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -8,6 +8,7 @@ ruamel.yaml==0.16.5
 setuptools==39.0.1
 pyroaring==0.2.9
 ujson==1.35
+tblib==1.5.0
 ## The following requirements were added by pip freeze:
 MarkupSafe==1.1.1
 ruamel.yaml.clib==0.2.0
diff --git a/src/buildstream/_exceptions.py b/src/buildstream/_exceptions.py
index 51f5427..072be20 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)
 
 
@@ -394,3 +397,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 c2263c3..c02ba6a 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")