WIP: pickle: PluginContext, ElementFactory, SourceFactory

Make PluginContext picklable by passing down unique identifiers that
can used to reconstruct them afterwards.

Clear out the 'types' cache, we'll have to build it up again on-demand.

4 files changed, 40 insertions, 10 deletions

diff --git a/src/buildstream/_elementfactory.py b/src/buildstream/_elementfactory.py
index d6591bf4c..b2a7f73dc 100644
--- a/src/buildstream/_elementfactory.py
+++ b/src/buildstream/_elementfactory.py
@@ -33,9 +33,12 @@ class ElementFactory(PluginContext):
 
     def __init__(self, plugin_base, *,
                  format_versions={},
-                 plugin_origins=None):
+                 plugin_origins=None,
+                 pass_=None):
 
-        super().__init__(plugin_base, Element, [_site.element_plugins],
+        assert pass_ is not None
+
+        super().__init__(plugin_base, Element, [_site.element_plugins], 'element' + str(pass_),
                          plugin_origins=plugin_origins,
                          format_versions=format_versions)
 
diff --git a/src/buildstream/_plugincontext.py b/src/buildstream/_plugincontext.py
index 7a5407cf6..0d322c156 100644
--- a/src/buildstream/_plugincontext.py
+++ b/src/buildstream/_plugincontext.py
@@ -42,10 +42,12 @@ from . import _yaml
 #
 class PluginContext():
 
-    def __init__(self, plugin_base, base_type, site_plugin_path, *,
+    def __init__(self, plugin_base, base_type, site_plugin_path, identifier, *,
                  plugin_origins=None, dependencies=None,
                  format_versions={}):
 
+        self._identifier = identifier
+
         # The plugin kinds which were loaded
         self.loaded_dependencies = []
 
@@ -59,10 +61,26 @@ class PluginContext():
 
         # The PluginSource object
         self._plugin_base = plugin_base
-        self._site_source = plugin_base.make_plugin_source(searchpath=site_plugin_path)
+        self._site_plugin_path = site_plugin_path
+        self._site_source = plugin_base.make_plugin_source(
+            searchpath=self._site_plugin_path,
+            identifier='site_plugin-' + self._identifier)
         self._alternate_sources = {}
         self._format_versions = format_versions
 
+    def __getstate__(self):
+        import copy
+        state = copy.copy(self.__dict__)
+        del state['_site_source']
+        state['_types'] = {}
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__ = state
+        self._site_source = self._plugin_base.make_plugin_source(
+            searchpath=self._site_plugin_path,
+            identifier='site_plugin-' + self._identifier)
+
     # lookup():
     #
     # Fetches a type loaded from a plugin in this plugin context
@@ -80,7 +98,7 @@ class PluginContext():
     def _get_local_plugin_source(self, path):
         if ('local', path) not in self._alternate_sources:
             # key by a tuple to avoid collision
-            source = self._plugin_base.make_plugin_source(searchpath=[path])
+            source = self._plugin_base.make_plugin_source(searchpath=[path], identifier='local_plugin-' + path + '-' + self._identifier)
             # Ensure that sources never get garbage collected,
             # as they'll take the plugins with them.
             self._alternate_sources[('local', path)] = source
@@ -121,7 +139,7 @@ class PluginContext():
                 # The plugin didn't have an accompanying YAML file
                 defaults = None
 
-            source = self._plugin_base.make_plugin_source(searchpath=[os.path.dirname(location)])
+            source = self._plugin_base.make_plugin_source(searchpath=[os.path.dirname(location)], identifier='pip_plugin-' + self._identifier)
             self._alternate_sources[('pip', package_name)] = source
 
         else:
diff --git a/src/buildstream/_project.py b/src/buildstream/_project.py
index c40321c66..526bb8d7d 100644
--- a/src/buildstream/_project.py
+++ b/src/buildstream/_project.py
@@ -98,6 +98,8 @@ class Project():
                  default_mirror=None, parent_loader=None,
                  search_for_project=True):
 
+        self._pass = None
+
         # The project name
         self.name = None
 
@@ -623,6 +625,8 @@ class Project():
         config_no_include = _yaml.node_copy(self._default_config_node)
         _yaml.composite(config_no_include, project_conf_first_pass)
 
+        assert self._pass is None
+        self._pass = 1
         self._load_pass(config_no_include, self.first_pass_config,
                         ignore_unknown=True)
 
@@ -647,6 +651,8 @@ class Project():
         config = _yaml.node_copy(self._default_config_node)
         _yaml.composite(config, project_conf_second_pass)
 
+        assert self._pass == 1
+        self._pass = 2
         self._load_pass(config, self.config)
 
         self._validate_node(config)
@@ -920,10 +926,12 @@ class Project():
         pluginbase = PluginBase(package='buildstream.plugins')
         output.element_factory = ElementFactory(pluginbase,
                                                 plugin_origins=plugin_element_origins,
-                                                format_versions=element_format_versions)
+                                                format_versions=element_format_versions,
+                                                pass_=self._pass)
         output.source_factory = SourceFactory(pluginbase,
                                               plugin_origins=plugin_source_origins,
-                                              format_versions=source_format_versions)
+                                              format_versions=source_format_versions,
+                                              pass_=self._pass)
 
     # _store_origin()
     #
diff --git a/src/buildstream/_sourcefactory.py b/src/buildstream/_sourcefactory.py
index 1d959a140..eca4b50a9 100644
--- a/src/buildstream/_sourcefactory.py
+++ b/src/buildstream/_sourcefactory.py
@@ -33,9 +33,10 @@ class SourceFactory(PluginContext):
 
     def __init__(self, plugin_base, *,
                  format_versions={},
-                 plugin_origins=None):
+                 plugin_origins=None,
+                 pass_=None):
 
-        super().__init__(plugin_base, Source, [_site.source_plugins],
+        super().__init__(plugin_base, Source, [_site.source_plugins], 'source' + str(pass_),
                          format_versions=format_versions,
                          plugin_origins=plugin_origins)