Merge branch 'aevri/enable_spawn_ci_2' into 'master'

jobpickler: also pickle Digest protos

See merge request BuildStream/buildstream!1652

2 files changed, 22 insertions, 7 deletions

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 96cc4b2b1..413592f41 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -199,7 +199,7 @@ tests-spawn-multiprocessing-start-method:
     - mkdir -p "${INTEGRATION_CACHE}"
     - useradd -Um buildstream
     - chown -R buildstream:buildstream .
-    - su buildstream -c "tox --  ${PYTEST_ARGS} tests/{cachekey,plugins}"
+    - su buildstream -c "tox --  ${PYTEST_ARGS} tests/{cachekey,plugins,internals,sourcecache}"
 
 # Run type checkers
 mypy:
diff --git a/src/buildstream/_scheduler/jobs/jobpickler.py b/src/buildstream/_scheduler/jobs/jobpickler.py
index 0edf88c10..6a4a8213b 100644
--- a/src/buildstream/_scheduler/jobs/jobpickler.py
+++ b/src/buildstream/_scheduler/jobs/jobpickler.py
@@ -23,12 +23,25 @@ import io
 import pickle
 
 from ..._protos.buildstream.v2.artifact_pb2 import Artifact as ArtifactProto
+from ..._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import Digest as DigestProto
 
 # BuildStream toplevel imports
 from ..._loader import Loader
 from ..._messenger import Messenger
 
 
+# Note that `str(type(proto_class))` results in `GeneratedProtocolMessageType`
+# instead of the concrete type, so we come up with our own names here.
+_NAME_TO_PROTO_CLASS = {
+    "artifact": ArtifactProto,
+    "digest": DigestProto,
+}
+
+_PROTO_CLASS_TO_NAME = {
+    cls: name for name, cls in _NAME_TO_PROTO_CLASS.items()
+}
+
+
 # pickle_child_job()
 #
 # Perform the special case pickling required to pickle a child job for
@@ -87,7 +100,8 @@ def pickle_child_job(child_job, projects):
         pickler.dispatch_table[cls] = _reduce_plugin
     for cls in source_classes:
         pickler.dispatch_table[cls] = _reduce_plugin
-    pickler.dispatch_table[ArtifactProto] = _reduce_artifact_proto
+    pickler.dispatch_table[ArtifactProto] = _reduce_proto
+    pickler.dispatch_table[DigestProto] = _reduce_proto
     pickler.dispatch_table[Loader] = _reduce_object
     pickler.dispatch_table[Messenger] = _reduce_object
 
@@ -103,14 +117,15 @@ def _reduce_object(instance):
     return (cls.__new__, (cls,), state)
 
 
-def _reduce_artifact_proto(instance):
-    assert isinstance(instance, ArtifactProto)
+def _reduce_proto(instance):
+    name = _PROTO_CLASS_TO_NAME[type(instance)]
     data = instance.SerializeToString()
-    return (_new_artifact_proto_from_reduction_args, (data,))
+    return (_new_proto_from_reduction_args, (name, data))
 
 
-def _new_artifact_proto_from_reduction_args(data):
-    instance = ArtifactProto()
+def _new_proto_from_reduction_args(name, data):
+    cls = _NAME_TO_PROTO_CLASS[name]
+    instance = cls()
     instance.ParseFromString(data)
     return instance