jobpickler: also pickle DigestProto

This is now required by some code paths. Also make a generic routine for
pickling / unpickling, as we may be doing more of this.

1 files changed, 21 insertions, 6 deletions

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