WIP: pickle: use PicklablePluginProxy instead

5 files changed, 55 insertions, 80 deletions

diff --git a/src/buildstream/_elementfactory.py b/src/buildstream/_elementfactory.py
index b2a7f73dc..89ec03b62 100644
--- a/src/buildstream/_elementfactory.py
+++ b/src/buildstream/_elementfactory.py
@@ -18,7 +18,7 @@
 #        Tristan Van Berkom <tristan.vanberkom@codethink.co.uk>
 
 from . import _site
-from ._plugincontext import PluginContext
+from ._plugincontext import PluginContext, PicklablePluginProxy
 from .element import Element
 
 
@@ -63,4 +63,6 @@ class ElementFactory(PluginContext):
         element = element_type(context, project, meta, default_config)
         version = self._format_versions.get(meta.kind, 0)
         self._assert_plugin_format(element, version)
-        return element
+        proxy = PicklablePluginProxy(element, self, meta.kind)
+        element._setup_artifact(proxy, context)
+        return proxy
diff --git a/src/buildstream/_plugincontext.py b/src/buildstream/_plugincontext.py
index 0d322c156..3fee59b4f 100644
--- a/src/buildstream/_plugincontext.py
+++ b/src/buildstream/_plugincontext.py
@@ -25,6 +25,51 @@ from . import utils
 from . import _yaml
 
 
+class PicklablePluginProxy():
+
+    def __init__(self, plugin_to_proxy, factory, kind):
+        object.__setattr__(self, '_PicklablePluginProxy__plugin', plugin_to_proxy)
+        object.__setattr__(self, '_PicklablePluginProxy__factory', factory)
+        object.__setattr__(self, '_PicklablePluginProxy__kind', kind)
+
+    def __getattr__(self, name):
+        return getattr(self.__plugin, name)
+
+    def __setattr__(self, name, value):
+        return setattr(self.__plugin, name, value)
+
+    def __getstate__(self):
+        # print("Pickling plugin", self.__plugin, repr(self.__plugin))
+        plugin_dict = self.__plugin.__dict__.copy()
+        try:
+            # We end up with excessive (infinite?) recursion if we keep these.
+            # TODO: figure out why, pickle is meant to handle cycles.
+            del plugin_dict["_Element__reverse_dependencies"]
+        except KeyError as k:
+            # print(k)
+            pass
+        return {
+            'factory': self.__factory,
+            'kind': self.__kind,
+            'plugin_dict': plugin_dict,
+        }
+
+    def __setstate__(self, state):
+        factory = state['factory']
+        kind = state['kind']
+        plugin_dict = state['plugin_dict']
+
+        cls, _ = factory.lookup(kind)
+        plugin = cls.__new__(cls)
+        plugin.__dict__ = plugin_dict
+
+        object.__setattr__(self, '_PicklablePluginProxy__plugin', plugin)
+        object.__setattr__(self, '_PicklablePluginProxy__factory', factory)
+        object.__setattr__(self, '_PicklablePluginProxy__kind', kind)
+
+        # print("Unpickled plugin", plugin, repr(plugin))
+
+
 # A Context for loading plugin types
 #
 # Args:
diff --git a/src/buildstream/_scheduler/jobs/elementjob.py b/src/buildstream/_scheduler/jobs/elementjob.py
index ca1927336..858d27b5d 100644
--- a/src/buildstream/_scheduler/jobs/elementjob.py
+++ b/src/buildstream/_scheduler/jobs/elementjob.py
@@ -25,72 +25,6 @@ from ..._artifact import Artifact
 from .job import Job, ChildJob
 
 
-# TODO: Make sure we only unpickle each element and source once, probably using
-# some global mechanism and context manager.
-
-
-def make_picklable_source_state(source):
-    print("Pickling source", source, repr(source))
-    # breakpoint()
-    meta_kind = source._meta_kind
-    project = source._get_project()
-    factory = project.config.source_factory
-    source_dict = source.__dict__.copy()
-    return {
-        'factory': factory,
-        'meta_kind': meta_kind,
-        'source_dict': source_dict,
-    }
-
-
-def make_source_from_picklable_source_state(state):
-    factory = state['factory']
-    meta_kind = state['meta_kind']
-    source_dict = state['source_dict']
-    # breakpoint()
-    source_cls, default_config = factory.lookup(meta_kind)
-    source = source_cls.__new__(source_cls)
-    source.__dict__ = source_dict
-    context = source._get_project()._context
-    print("Unpickled source", source, repr(source))
-    return source
-
-
-def make_picklable_element_state(element):
-    print("Pickling element", element, repr(element))
-    meta_kind = element._meta_kind
-    project = element._get_project()
-    factory = project.config.element_factory
-    element_dict = element.__dict__.copy()
-    # del element_dict['_Element__sources']
-    element_dict['_Element__sources'] = [
-        make_picklable_source_state(s)
-        for s in element.sources()
-    ]
-    del element_dict['_Element__artifact']
-    return {
-        'factory': factory,
-        'meta_kind': meta_kind,
-        'element_dict': element_dict,
-    }
-
-
-def make_element_from_picklable_element_state(state):
-    factory = state['factory']
-    meta_kind = state['meta_kind']
-    element_dict = state['element_dict']
-    element_cls, default_config = factory.lookup(meta_kind)
-    element = element_cls.__new__(element_cls)
-    element.__dict__ = element_dict
-    context = element._get_project()._context
-    element_dict['_Element__sources'] = [
-        make_source_from_picklable_source_state(s)
-        for s in element_dict['_Element__sources']
-    ]
-    element_dict['_Element__artifact'] = Artifact(element, context)
-    print("Unpickled element", element, repr(element))
-    return element
-
 import sys
 import pdb
 
@@ -156,16 +90,6 @@ class ElementJob(Job):
         self.set_message_unique_id(element._unique_id)
         self.set_task_id(element._unique_id)
 
-    def __getstate__(self):
-        state = super().__getstate__()
-        state['_element'] = make_picklable_element_state(self._element)
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        self._element = make_element_from_picklable_element_state(
-            self._element)
-
     @property
     def element(self):
         return self._element
diff --git a/src/buildstream/_sourcefactory.py b/src/buildstream/_sourcefactory.py
index eca4b50a9..46dc24fa3 100644
--- a/src/buildstream/_sourcefactory.py
+++ b/src/buildstream/_sourcefactory.py
@@ -18,7 +18,7 @@
 #        Tristan Van Berkom <tristan.vanberkom@codethink.co.uk>
 
 from . import _site
-from ._plugincontext import PluginContext
+from ._plugincontext import PluginContext, PicklablePluginProxy
 from .source import Source
 
 
@@ -62,4 +62,5 @@ class SourceFactory(PluginContext):
         source = source_type(context, project, meta)
         version = self._format_versions.get(meta.kind, 0)
         self._assert_plugin_format(source, version)
-        return source
+        proxy = PicklablePluginProxy(source, self, meta.kind)
+        return proxy
diff --git a/src/buildstream/element.py b/src/buildstream/element.py
index bc82f5141..01eb6c95c 100644
--- a/src/buildstream/element.py
+++ b/src/buildstream/element.py
@@ -283,6 +283,9 @@ class Element(Plugin):
                 # This will taint the artifact, disable pushing.
                 self.__sandbox_config_supported = False
 
+    def _setup_artifact(self, proxy, context):
+        self.__artifact = Artifact(proxy, context)  # Artifact class for direct artifact composite interaction
+
     def __lt__(self, other):
         return self.name < other.name