Scan collectors config directory and allow multiple instances of a Co…

…llector

src/collectors/MemcachedCollector/MemcachedCollector.py

@@ -18,6 +18,9 @@ def get_default_config(self):
             # Connection settings
             'host':     'localhost',
             'port':     '11211',
+
+            # Set to true to send stats as <prefix>.<host>.<port>.memcached.(...)
+            'override_hostname':  False,
 
             # Which rows of 'status' you would like to publish.
             # 'telnet host port' and type stats and hit enter to see the list of
@@ -31,6 +34,8 @@ def get_stats(self):
         ignored = ('libevent', 'pid', 'pointer_size', 'time', 'version')
 
         config = self.config
+        if config['override_hostname']:
+            config['hostname'] = "%s.%s" % (config['host'].replace('.', '_'), str(config['port']))
         stats = {}
         # connect
         try:

src/diamond/collector.py

@@ -22,7 +22,7 @@ class Collector(object):
     The Collector class is a base class for all metric collectors.
     """
 
-    def __init__(self, config, handlers):
+    def __init__(self, config, handlers, custom_config=None):
         """
         Create a new instance of the Collector class
         """
@@ -53,6 +53,9 @@ def __init__(self, config, handlers):
         if os.path.exists(configfile):
             # Merge Collector config file
             self.config.merge(configobj.ConfigObj(configfile))
+        if custom_config:
+            # Merge custom config file
+            self.config.merge(custom_config)
 
     def get_default_config(self):
         """

src/diamond/server.py

@@ -10,6 +10,7 @@
 import inspect
 import pwd
 import grp
+import glob
 
 # Path Fix
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),"../")))
@@ -184,14 +185,33 @@ def load_collectors(self, path, filter=None):
         # Return Collector classes
         return collectors
 
-    def init_collector(self, cls):
+    def load_alternative_collectors(self, collectors, path):
+        # Initialize return value
+        alternative_collectors = {}
+
+        # Get a list of files in the directory, if the directory exists
+        if not os.path.exists(path):
+            raise OSError, "Directory does not exist: %s" % (path)
+
+        # Log
+        self.log.debug("Loading alternative Collectors from: %s" % (path))
+
+        for f in glob.glob(os.path.join(path, '*-*.conf')):
+            if os.path.isfile(f):
+                fname = os.path.basename(f)
+                cls_name, alt_name = fname.split('-', 1)
+                if cls_name in collectors:
+                    alternative_collectors[fname[:-5]] = (collectors[cls_name], f)
+        return alternative_collectors
+
+    def init_collector(self, cls, custom_config=None):
         """
         Initialize collector
         """
         collector = None
         try:
             # Initialize Collector
-            collector = cls(self.config, self.handlers)
+            collector = cls(self.config, self.handlers, custom_config=custom_config)
             # Log
             self.log.debug("Initialized Collector: %s" % (cls.__name__))
         except Exception, e:
@@ -201,7 +221,7 @@ def init_collector(self, cls):
         # Return collector
         return collector
 
-    def schedule_collector(self, c, interval_task=True):
+    def schedule_collector(self, c, interval_task=True, custom_name=None):
         """
         Schedule collector
         """
@@ -216,6 +236,8 @@ def schedule_collector(self, c, interval_task=True):
 
         # Get collector schedule
         for name,schedule in c.get_schedule().items():
+            if custom_name is not None:
+                name = custom_name
             # Get scheduler args
             func, args, splay, interval = schedule
 
@@ -261,12 +283,24 @@ def run(self):
         # Load collectors
         collectors = self.load_collectors(self.config['server']['collectors_path'])
 
+        # Load collectors with alternative configurations
+        alternative_collectors = self.load_alternative_collectors(collectors,
+                                                                  self.config['server']['collectors_config_path'])
         # Setup Collectors
         for cls in collectors.values():
             # Initialize Collector
             c = self.init_collector(cls)
             # Schedule Collector
             self.schedule_collector(c)
+
+        # Setup Collectors
+        for alt_name, cls_and_config in alternative_collectors.iteritems():
+            cls, config_file = cls_and_config
+            custom_config = configobj.ConfigObj(config_file)
+            # Initialize Collector
+            c = self.init_collector(cls, custom_config=custom_config)
+            # Schedule Collector
+            self.schedule_collector(c, custom_name=alt_name)
 
         # Start main loop
         self.mainloop()

test.py

@@ -6,6 +6,8 @@
 import unittest
 import inspect
 import traceback
+import tempfile
+import os.path
 
 from StringIO import StringIO
 from contextlib import nested
@@ -97,6 +99,10 @@ def getCollectorTests(path):
 
 
 class BaseCollectorTest(unittest.TestCase):
+
+    def setUp(self):
+        unittest.TestCase.setUp(self)
+        self.tmpfile = None
 
     def test_SetCustomHostname(self):
         config = configobj.ConfigObj()
@@ -109,6 +115,34 @@ def test_SetCustomHostname(self):
         c = Collector(config, [])
         self.assertEquals('custom.localhost', c.get_hostname())
 
+    def test_CustomConfigOverridesSetCustomHostname(self):
+        temp_dir = tempfile.gettempdir()
+        if not temp_dir:
+            self.skipTest("No temporary directory in system")
+
+        config = configobj.ConfigObj()
+        config['server'] = {}
+        config['server']['collectors_config_path'] = temp_dir
+        config['collectors'] = {}
+        config['collectors']['default'] = {}
+
+
+        self.tmpfile = os.path.join(temp_dir, 'Collector.conf')
+        collectorConfig = configobj.ConfigObj(self.tmpfile)
+        collectorConfig['hostname'] = "Collector.hostname"
+        collectorConfig.write()
+
+        customConfig = configobj.ConfigObj()
+        customConfig['hostname'] = 'custom.hostname'
+
+        c = Collector(config, [], custom_config=customConfig)
+
+        self.assertEquals('custom.hostname', c.get_hostname())
+
+    def tearDown(self):
+        if self.tmpfile:
+            os.unlink(self.tmpfile)
+
 
 
 ################################################################################