From 29499db7f35fb1efd1a89c62557c2beead3657de Mon Sep 17 00:00:00 2001
From: Hao Song <songhao@vmware.com>
Date: Wed, 18 Sep 2019 17:01:10 -0700
Subject: [PATCH] Support addition of custom tags in heartbeater service (#45)

* Support addition of custom tags in heartbeater service

* Fix pylint error

* Bump version
---
 setup.py                                    |  2 +-
 wavefront_sdk/common/heartbeater_service.py | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index cdf0d4c..3a15b37 100755
--- a/setup.py
+++ b/setup.py
@@ -15,7 +15,7 @@
 
 setuptools.setup(
     name='wavefront-sdk-python',
-    version='1.2',
+    version='1.2.1',
     author='Wavefront by VMware',
     author_email='chitimba@wavefront.com',
     url='https://github.com/wavefrontHQ/wavefront-sdk-python',
diff --git a/wavefront_sdk/common/heartbeater_service.py b/wavefront_sdk/common/heartbeater_service.py
index 5e0898c..9d30abc 100644
--- a/wavefront_sdk/common/heartbeater_service.py
+++ b/wavefront_sdk/common/heartbeater_service.py
@@ -21,6 +21,7 @@
 
 # pylint: disable=E0012,R0205
 # pylint: disable=too-few-public-methods
+# pylint: disable=too-many-instance-attributes
 class HeartbeaterService(object):
     """Service that periodically reports component heartbeats to Wavefront."""
 
@@ -38,6 +39,7 @@ def __init__(self, wavefront_client, application_tags, components, source):
         self.source = source
         self.reporting_interval_seconds = 60 * 5
         self.heartbeat_metric_tags_list = []
+        self.custom_tags_set = set()
         if isinstance(components, str):
             components = [components]
         for component in components:
@@ -54,6 +56,14 @@ def __init__(self, wavefront_client, application_tags, components, source):
         self._timer = None
         self._run()
 
+    def report_custom_tags(self, custom_tags):
+        """
+        Append custom tags for heartbeat reporting.
+
+        @param custom_tags: dict of custom tags.
+        """
+        self.custom_tags_set.add(custom_tags)
+
     def _schedule_timer(self):
         self._timer = threading.Timer(self.reporting_interval_seconds,
                                       self._run)
@@ -72,6 +82,11 @@ def _report(self):
                 self.wavefront_client.send_metric(HEART_BEAT_METRIC, 1.0,
                                                   time.time(), self.source,
                                                   heartbeat)
+            while self.custom_tags_set:
+                self.wavefront_client.send_metric(HEART_BEAT_METRIC, 1.0,
+                                                  time.time(), self.source,
+                                                  self.custom_tags_set.pop())
+
         # pylint: disable=broad-except,fixme
         # TODO: Please make sure we catch more specific exception here.
         except Exception: