added short-lived caching of Packages Collection objects

2 files changed, 54 insertions, 19 deletions

diff --git a/src/lib/Bcfg2/Server/Plugins/Packages/Collection.py b/src/lib/Bcfg2/Server/Plugins/Packages/Collection.py
index 959dac03b..3ea14ce75 100644
--- a/src/lib/Bcfg2/Server/Plugins/Packages/Collection.py
+++ b/src/lib/Bcfg2/Server/Plugins/Packages/Collection.py
@@ -1,3 +1,4 @@
+import sys
 import copy
 import logging
 import Bcfg2.Server.Plugin
@@ -18,9 +19,16 @@ except ImportError:
 # startup.  (It would also prevent machines that change groups from
 # working properly; e.g., if you reinstall a machine with a new OS,
 # then returning a cached Collection object would give the wrong
-# sources to that client.)
+# sources to that client.)  These are keyed by the collection
+# cachekey, a unique key identifying the collection by its _config_,
+# which could be shared among multiple clients.
 collections = dict()
 
+# cache mapping of hostname -> collection cachekey.  this _is_ used to
+# return a Collection object when one is requested, so each entry is
+# very short-lived -- it's purged at the end of each client run.
+clients = dict()
+
 class Collection(Bcfg2.Server.Plugin.Debuggable):
     def __init__(self, metadata, sources, basepath, debug=False):
         """ don't call this directly; use the factory function """
@@ -32,11 +40,11 @@ class Collection(Bcfg2.Server.Plugin.Debuggable):
         self.virt_pkgs = dict()
 
         try:
-            self.config = sources[0].config
+            self.setup = sources[0].setup
             self.cachepath = sources[0].basepath
             self.ptype = sources[0].ptype
         except IndexError:
-            self.config = None
+            self.setup = None
             self.cachepath = None
             self.ptype = "unknown"
 
@@ -290,9 +298,33 @@ class Collection(Bcfg2.Server.Plugin.Debuggable):
     def sort(self, cmp=None, key=None, reverse=False):
         self.sources.sort(cmp, key, reverse)
 
+def get_collection_class(source_type):
+    modname = "Bcfg2.Server.Plugins.Packages.%s" % source_type.title()
+    
+    try:
+        module = sys.modules[modname]
+    except KeyError:
+        try:
+            module = __import__(modname).Server.Plugins.Packages
+        except ImportError:
+            msg = "Packages: Unknown source type %s" % source_type
+            logger.error(msg)
+            raise Bcfg2.Server.Plugin.PluginExecutionError(msg)
+    
+    try:
+        cclass = getattr(module, source_type.title() + "Collection")
+    except AttributeError:
+        msg = "Packages: No collection class found for %s sources" % source_type
+        logger.error(msg)
+        raise Bcfg2.Server.Plugin.PluginExecutionError(msg)
+    
+    return cclass
+
 def clear_cache():
     global collections
+    global clients
     collections = dict()
+    clients = dict()
 
 def factory(metadata, sources, basepath, debug=False):
     global collections
@@ -302,6 +334,9 @@ def factory(metadata, sources, basepath, debug=False):
         # instantiate a dummy Collection object
         return Collection(metadata, [], basepath)
 
+    if metadata.hostname in clients:
+        return collections[clients[metadata.hostname]]
+
     sclasses = set()
     relevant = list()
 
@@ -326,24 +361,16 @@ def factory(metadata, sources, basepath, debug=False):
                          metadata.hostname)
         cclass = Collection
     else:
-        stype = sclasses.pop().__name__.replace("Source", "")
-        try:
-            module = \
-                getattr(__import__("Bcfg2.Server.Plugins.Packages.%s" %
-                                   stype.title()).Server.Plugins.Packages,
-                        stype.title())
-            cclass = getattr(module, "%sCollection" % stype.title())
-        except ImportError:
-            logger.error("Packages: Unknown source type %s" % stype)
-        except AttributeError:
-            logger.warning("Packages: No collection class found for %s sources"
-                           % stype)
+        cclass = get_collection_class(sclasses.pop().__name__.replace("Source",
+                                                                      ""))
 
     if debug:
         logger.error("Packages: Using %s for Collection of sources for %s" %
                      (cclass.__name__, metadata.hostname))
 
     collection = cclass(metadata, relevant, basepath, debug=debug)
-    collections[metadata.hostname] = collection
+    ckey = collection.cachekey
+    clients[metadata.hostname] = ckey
+    collections[ckey] = collection
     return collection
 
diff --git a/src/lib/Bcfg2/Server/Plugins/Packages/__init__.py b/src/lib/Bcfg2/Server/Plugins/Packages/__init__.py
index 3f5c46aab..4070b13ca 100644
--- a/src/lib/Bcfg2/Server/Plugins/Packages/__init__.py
+++ b/src/lib/Bcfg2/Server/Plugins/Packages/__init__.py
@@ -15,7 +15,8 @@ from Bcfg2.Server.Plugins.Packages.PackagesConfig import PackagesConfig
 class Packages(Bcfg2.Server.Plugin.Plugin,
                Bcfg2.Server.Plugin.StructureValidator,
                Bcfg2.Server.Plugin.Generator,
-               Bcfg2.Server.Plugin.Connector):
+               Bcfg2.Server.Plugin.Connector,
+               Bcfg2.Server.Plugin.GoalValidator):
     name = 'Packages'
     conflicts = ['Pkgmgr']
     experimental = True
@@ -101,8 +102,7 @@ class Packages(Bcfg2.Server.Plugin.Plugin,
 
     def HandlesEntry(self, entry, metadata):
         if entry.tag == 'Package':
-            if self.config.getboolean("global", "magic_groups",
-                                      default=True) == True:
+            if self.config.getboolean("global", "magic_groups", default=True):
                 collection = self._get_collection(metadata)
                 if collection.magic_groups_match():
                     return True
@@ -267,3 +267,11 @@ class Packages(Bcfg2.Server.Plugin.Plugin,
     def get_additional_data(self, metadata):
         collection = self._get_collection(metadata)
         return dict(sources=collection.get_additional_data())
+
+    def validate_goals(self, metadata, _):
+        """ we abuse the GoalValidator plugin since validate_goals()
+        is the very last thing called during a client config run.  so
+        we use this to clear the collection cache for this client,
+        which must persist only the duration of a client run """
+        if metadata.hostname in Collection.clients:
+            del Collection.clients[metadata.hostname]