Support multiple graphite servers/prefixes
authorRail Aliiev <rail@mozilla.com>
Fri, 07 Nov 2014 08:56:22 -0500
changeset 536 66e0a416d40ea2adcf12aa0c94bca6071bf80779
parent 535 b6a57c573d8cbf816da453b075069e69d42d0a21
child 537 38de0f6e950384ce22ffa56efe064651e49a6858
push id530
push userraliiev@mozilla.com
push dateFri, 07 Nov 2014 14:00:39 +0000
Support multiple graphite servers/prefixes
cloudtools/graphite.py
scripts/aws_watch_pending.py
tests/test_cloudtools_graphite.py
--- a/cloudtools/graphite.py
+++ b/cloudtools/graphite.py
@@ -5,24 +5,20 @@ import time
 log = logging.getLogger(__name__)
 
 
 class GraphiteLogger(object):
     # to be used by modules
 
     def __init__(self):
         self._data = {}
+        self._servers = []
 
-    def connect(self, host, port):
-        try:
-            log.debug("Connecting to graphite at %s:%s", host, port)
-            self._sock = socket.create_connection((host, port), timeout=10)
-        except Exception:
-            log.exception("Couldn't connect to graphite at %s:%s", host, port)
-            log.warn("Ignoring all grapite submissions!")
+    def add_destination(self, host, port, prefix):
+        self._servers.append((host, port, prefix))
 
     @staticmethod
     def _generate_line(prefix, name, value, timestamp):
         return "{prefix}.{name} {value} {timestamp}\n".format(
             prefix=prefix, name=name, value=value, timestamp=timestamp)
 
     def add(self, name, value, timestamp=None, collect=False):
         # graphite needs numbers, not strings
@@ -40,27 +36,33 @@ class GraphiteLogger(object):
             self._data[name] = (value, timestamp)
 
     def generate_data(self, prefix):
         data = []
         for name, (value, timestamp) in sorted(self._data.iteritems()):
             data.append(self._generate_line(prefix, name, value, timestamp))
         return "".join(data)
 
-    def sendall(self, prefix):
-        data = self.generate_data(prefix)
-        if data:
+    def sendall(self):
+        if not self._data:
+            log.warning("Nothing to submit to graphite")
+            return
+
+        for host, port, prefix in self._servers:
+            data = self.generate_data(prefix)
             log.debug("Graphite send: \n%s", data)
-            if self._sock:
-                self._sock.sendall(data)
-            else:
-                log.warn("Not sending to graphite")
-            self._data = {}
-        else:
-            log.warning("Nothing to submit to graphite")
+            try:
+                log.debug("Connecting to graphite at %s:%s", host, port)
+                sock = socket.create_connection((host, port), timeout=10)
+                sock.sendall(data)
+            except Exception:
+                log.exception("Couldn't send graphite data to %s:%s", host,
+                              port)
+                log.warn("Ignoring all grapite submissions!")
+        self._data = {}
 
 _graphite_logger = GraphiteLogger()
 
 
 def get_graphite_logger():
     global _graphite_logger
     return _graphite_logger
 
--- a/scripts/aws_watch_pending.py
+++ b/scripts/aws_watch_pending.py
@@ -478,14 +478,24 @@ if __name__ == '__main__':
         dburl=secrets['db'],
         regions=args.regions,
         builder_map=config['buildermap'],
         region_priorities=config['region_priorities'],
         dryrun=args.dryrun,
         spot_config=config.get("spot"),
         ondemand_config=config.get("ondemand"),
     )
-    if config.get("graphite_host") and config.get("graphite_port"):
-        gr_log.connect(host=config.get("graphite_host"),
-                       port=config.get("graphite_port"))
-        gr_log.sendall(prefix=config.get("graphite_prefix",
-                                         "aws_watch_pending"))
+
+    if all([config.get("graphite_host"), config.get("graphite_port"),
+            config.get("graphite_prefix")]):
+        gr_log.add_destination(
+            host=config["graphite_host"], port=config["graphite_port"],
+            prefix=config["graphite_prefix"])
+
+    for entry in secrets.get("graphite_hosts", []):
+        host = entry.get("host")
+        port = entry.get("port")
+        prefix = entry.get("prefix")
+        if all([host, port, prefix]):
+            gr_log.add_destination(host, port, prefix)
+
+    gr_log.sendall()
     log.debug("done")
--- a/tests/test_cloudtools_graphite.py
+++ b/tests/test_cloudtools_graphite.py
@@ -54,8 +54,57 @@ class TestGraphiteLogger(unittest.TestCa
         gl.add("name", 55.66, 2222)
         self.assertDictEqual(gl._data, {"name": (55.66, 2222)})
 
     def test_collect(self):
         gl = get_graphite_logger()
         gl.add("name", 44, collect=True)
         gl.add("name", 55.66, 2222, collect=True)
         self.assertDictEqual(gl._data, {"name": (99.66, 2222)})
+
+    def test_add_destination(self):
+        gl = get_graphite_logger()
+        gl.add_destination("host0", 888, "prefix0")
+        gl.add_destination("host9", 999, "prefix9")
+        self.assertEqual(gl._servers,
+                         [("host0", 888, "prefix0"),
+                          ("host9", 999, "prefix9")])
+
+    def test_single_server(self):
+        gl = get_graphite_logger()
+        gl.add("name", 44)
+        gl.add_destination("host1", 1111, "prefix1")
+        with mock.patch("socket.create_connection") as conn:
+            gl.sendall()
+            expected_calls = [
+                mock.call(("host1", 1111), timeout=10),
+            ]
+            conn.assert_has_calls(expected_calls)
+
+    def test_multiple_servers(self):
+        gl = get_graphite_logger()
+        gl.add_destination("host1", 1111, "prefix1")
+        gl.add_destination("host2", 2222, "prefix2")
+        gl.add("name", 44)
+        with mock.patch("socket.create_connection") as conn:
+            gl.sendall()
+            expected_calls = [
+                mock.call(("host1", 1111), timeout=10),
+                mock.call(("host2", 2222), timeout=10),
+            ]
+            conn.assert_has_calls(expected_calls, any_order=True)
+
+    def test_multiple_servers_sendall(self):
+        gl = get_graphite_logger()
+        gl.add_destination("host1", 1111, "prefix1")
+        gl.add_destination("host2", 2222, "prefix2")
+        with mock.patch("time.time") as m_time:
+            m_time.return_value = 9999
+            gl.add("name", 44)
+        with mock.patch("socket.create_connection") as conn:
+            sock = mock.MagicMock()
+            conn.return_value = sock
+            gl.sendall()
+            expected_calls = [
+                mock.call("prefix1.name 44 9999\n"),
+                mock.call("prefix2.name 44 9999\n"),
+            ]
+            sock.sendall.assert_has_calls(expected_calls)