Move heartbeat init logic into worker (#952)

Author: lzchen

contrib/opencensus-ext-azure/examples/traces/client.py

@@ -26,6 +26,6 @@
 tracer = Tracer(exporter=AzureExporter(), sampler=ProbabilitySampler(1.0))
 with tracer.span(name='parent'):
     with tracer.span(name='child'):
-        response = requests.get(url='http://localhost:8080/')
+        response = requests.get(url='http://example.com/')
         print(response.status_code)
         print(response.text)

contrib/opencensus-ext-azure/opencensus/ext/azure/metrics_exporter/heartbeat_metrics/heartbeat.py

@@ -37,20 +37,23 @@ class HeartbeatMetric:
     def __init__(self):
         self.vm_data = {}
         self.vm_retry = False
+        self.init = False
         self.properties = OrderedDict()
-        self._init_properties()
-        self.heartbeat = LongGauge(
-            HeartbeatMetric.NAME,
-            'Heartbeat metric with custom dimensions',
-            'count',
-            list(self.properties.keys()),
-        )
-        self.heartbeat.get_or_create_time_series(
-            list(self.properties.values())
-        )
 
     def get_metrics(self):
-        if self.vm_retry:
+        if not self.init:
+            self._init_properties()
+            self.heartbeat = LongGauge(
+                HeartbeatMetric.NAME,
+                'Heartbeat metric with custom dimensions',
+                'count',
+                list(self.properties.keys()),
+            )
+            self.heartbeat.get_or_create_time_series(
+                list(self.properties.values())
+            )
+            self.init = True
+        elif self.vm_retry:
             # Only need to possibly update if vm retry
             if self._get_azure_compute_metadata() and not self.vm_retry:
                 self._populate_vm_data()
@@ -64,7 +67,10 @@ def get_metrics(self):
                 self.heartbeat.get_or_create_time_series(
                     list(self.properties.values())
                 )
-        return [self.heartbeat.get_metric(datetime.datetime.utcnow())]
+        if self.heartbeat:
+            return [self.heartbeat.get_metric(datetime.datetime.utcnow())]
+        else:
+            return []
 
     def _init_properties(self):
         self.properties[LabelKey("sdk", '')] = LabelValue(

contrib/opencensus-ext-azure/opencensus/ext/azure/trace_exporter/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import atexit
 import json
 import logging
 
@@ -58,6 +59,7 @@ def __init__(self, **options):
         )
         self._telemetry_processors = []
         super(AzureExporter, self).__init__(**options)
+        atexit.register(self._stop, self.options.grace_period)
         heartbeat_metrics.enable_heartbeat_metrics(
             self.options.connection_string, self.options.instrumentation_key)
 
@@ -181,4 +183,4 @@ def emit(self, batch, event=None):
 
     def _stop(self, timeout=None):
         self.storage.close()
-        return self._worker.stop(timeout)
+        self._worker.stop(timeout)

contrib/opencensus-ext-azure/tests/test_azure_heartbeat_metrics.py

@@ -84,7 +84,16 @@ def test_enable_heartbeat_metrics_exists(self, transport_mock):
 
     def test_heartbeat_metric_init(self):
         metric = heartbeat_metrics.HeartbeatMetric()
+        self.assertEqual(len(metric.vm_data), 0)
+        self.assertFalse(metric.vm_retry)
+        self.assertFalse(metric.init)
+        self.assertEqual(len(metric.properties), 0)
 
+    def test_heartbeat_metric_get_metric_init(self):
+        metric = heartbeat_metrics.HeartbeatMetric()
+        self.assertFalse(metric.init)
+        metrics = metric.get_metrics()
+        self.assertTrue(metric.init)
         self.assertEqual(metric.NAME, 'Heartbeat')
         keys = list(metric.properties.keys())
         values = list(metric.properties.values())
@@ -115,6 +124,7 @@ def test_heartbeat_metric_init(self):
             gauge._len_label_keys,
             len(metric.properties.keys())
         )
+        self.assertEqual(len(metrics), 1)
 
     @mock.patch.dict(
         os.environ,
@@ -126,7 +136,9 @@ def test_heartbeat_metric_init(self):
     )
     def test_heartbeat_metric_init_webapp(self):
         metric = heartbeat_metrics.HeartbeatMetric()
-
+        self.assertFalse(metric.init)
+        metric.get_metrics()
+        self.assertTrue(metric.init)
         self.assertEqual(metric.NAME, 'Heartbeat')
         keys = list(metric.properties.keys())
         values = list(metric.properties.values())
@@ -156,7 +168,9 @@ def test_heartbeat_metric_init_webapp(self):
     )
     def test_heartbeat_metric_init_functionapp(self):
         metric = heartbeat_metrics.HeartbeatMetric()
-
+        self.assertFalse(metric.init)
+        metric.get_metrics()
+        self.assertTrue(metric.init)
         self.assertEqual(metric.NAME, 'Heartbeat')
         keys = list(metric.properties.keys())
         values = list(metric.properties.values())
@@ -186,6 +200,10 @@ def test_heartbeat_metric_init_vm(self):
                 )
             )
             metric = heartbeat_metrics.HeartbeatMetric()
+            self.assertFalse(metric.init)
+            self.assertFalse(metric.vm_retry)
+            metric.get_metrics()
+            self.assertTrue(metric.init)
             self.assertFalse(metric.vm_retry)
             self.assertEqual(metric.NAME, 'Heartbeat')
             keys = list(metric.properties.keys())
@@ -213,6 +231,10 @@ def test_heartbeat_metric_not_vm(self):
             throw(requests.exceptions.ConnectionError)
         ):
             metric = heartbeat_metrics.HeartbeatMetric()
+            self.assertFalse(metric.init)
+            self.assertFalse(metric.vm_retry)
+            metric.get_metrics()
+            self.assertTrue(metric.init)
             self.assertFalse(metric.vm_retry)
             self.assertEqual(metric.NAME, 'Heartbeat')
             keys = list(metric.properties.keys())
@@ -224,6 +246,10 @@ def test_heartbeat_metric_not_vm_timeout(self):
             throw(requests.Timeout)
         ):
             metric = heartbeat_metrics.HeartbeatMetric()
+            self.assertFalse(metric.init)
+            self.assertFalse(metric.vm_retry)
+            metric.get_metrics()
+            self.assertTrue(metric.init)
             self.assertFalse(metric.vm_retry)
             self.assertEqual(metric.NAME, 'Heartbeat')
             keys = list(metric.properties.keys())
@@ -235,6 +261,10 @@ def test_heartbeat_metric_vm_retry(self):
             throw(requests.exceptions.RequestException)
         ):
             metric = heartbeat_metrics.HeartbeatMetric()
+            self.assertFalse(metric.init)
+            self.assertFalse(metric.vm_retry)
+            metric.get_metrics()
+            self.assertTrue(metric.init)
             self.assertTrue(metric.vm_retry)
             keys = list(metric.properties.keys())
             self.assertEqual(len(keys), 2)

opencensus/metrics/transport.py

@@ -86,7 +86,10 @@ def run(self):
 
     def close(self):
         try:
-            return self.func(*self.args, **self.kwargs)
+            # Suppress request tracking on flush
+            execution_context.set_is_exporter(True)
+            self.func(*self.args, **self.kwargs)
+            execution_context.set_is_exporter(False)
         except Exception as ex:
             logger.exception("Error handling metric flush: {}".format(ex))
         self.cancel()