Azure Monitor metrics exporter atexit fixes (#943)

Author: lzchen

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## Unreleased
 
+- PeriodicMetricTask flush on exit
+([#943](https://github.com/census-instrumentation/opencensus-python/pull/943))
+
 ## 0.7.10
 Released 2020-06-29
 

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

@@ -55,6 +55,7 @@ def __init__(self, **options):
             source=self.__class__.__name__,
         )
         self._atexit_handler = atexit.register(self.shutdown)
+        self.exporter_thread = None
         super(MetricsExporter, self).__init__()
 
     def export_metrics(self, metrics):
@@ -136,21 +137,22 @@ def _create_envelope(self, data_point, timestamp, properties):
         return envelope
 
     def shutdown(self):
-        # flush metrics on exit
-        self.export_metrics(stats_module.stats.get_metrics())
-        if self._atexit_handler is not None:
-            atexit.unregister(self._atexit_handler)
-            self._atexit_handler = None
+        # Flush the exporter thread
+        if self.exporter_thread:
+            self.exporter_thread.close()
+        # Shutsdown storage worker
+        self.storage.close()
 
 
 def new_metrics_exporter(**options):
     exporter = MetricsExporter(**options)
     producers = [stats_module.stats]
     if exporter.options.enable_standard_metrics:
         producers.append(standard_metrics.producer)
-    transport.get_exporter_thread(producers,
-                                  exporter,
-                                  interval=exporter.options.export_interval)
+    exporter.exporter_thread = transport.get_exporter_thread(
+                                    producers,
+                                    exporter,
+                                    interval=exporter.options.export_interval)
     from opencensus.ext.azure.metrics_exporter import heartbeat_metrics
     heartbeat_metrics.enable_heartbeat_metrics(
         exporter.options.connection_string,

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

@@ -37,9 +37,10 @@ def enable_heartbeat_metrics(connection_string, ikey):
             )
             producer = AzureHeartbeatMetricsProducer()
             _HEARTBEAT_METRICS = producer
-            transport.get_exporter_thread([_HEARTBEAT_METRICS],
-                                          exporter,
-                                          exporter.options.export_interval)
+            exporter.exporter_thread = \
+                transport.get_exporter_thread([_HEARTBEAT_METRICS],
+                                              exporter,
+                                              exporter.options.export_interval)
 
 
 class AzureHeartbeatMetricsProducer(MetricProducer):

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

@@ -186,32 +186,65 @@ def test_create_envelope(self):
         self.assertTrue('properties' in envelope.data.baseData)
         self.assertEqual(envelope.data.baseData.properties, properties)
 
+    def test_shutdown(self):
+        mock_thread = mock.Mock()
+        mock_storage = mock.Mock()
+        exporter = metrics_exporter.MetricsExporter(
+            instrumentation_key='12345678-1234-5678-abcd-12345678abcd'
+        )
+        exporter.exporter_thread = mock_thread
+        exporter.storage = mock_storage
+        exporter.shutdown()
+        mock_thread.close.assert_called_once()
+        mock_storage.close.assert_called_once()
+
     @mock.patch('opencensus.ext.azure.metrics_exporter'
                 '.transport.get_exporter_thread')
     def test_new_metrics_exporter(self, exporter_mock):
-        iKey = '12345678-1234-5678-abcd-12345678abcd'
-        exporter = metrics_exporter.new_metrics_exporter(
-            instrumentation_key=iKey)
-
-        self.assertEqual(exporter.options.instrumentation_key, iKey)
-        self.assertEqual(len(exporter_mock.call_args_list), 1)
-        self.assertEqual(len(exporter_mock.call_args[0][0]), 2)
-        producer_class = standard_metrics.AzureStandardMetricsProducer
-        self.assertFalse(isinstance(exporter_mock.call_args[0][0][0],
-                                    producer_class))
-        self.assertTrue(isinstance(exporter_mock.call_args[0][0][1],
-                                   producer_class))
+        with mock.patch('opencensus.ext.azure.metrics_exporter'
+                        '.heartbeat_metrics.enable_heartbeat_metrics') as hb:
+            hb.return_value = None
+            iKey = '12345678-1234-5678-abcd-12345678abcd'
+            exporter = metrics_exporter.new_metrics_exporter(
+                instrumentation_key=iKey)
+
+            self.assertEqual(exporter.options.instrumentation_key, iKey)
+            self.assertEqual(len(exporter_mock.call_args_list), 1)
+            self.assertEqual(len(exporter_mock.call_args[0][0]), 2)
+            producer_class = standard_metrics.AzureStandardMetricsProducer
+            self.assertFalse(isinstance(exporter_mock.call_args[0][0][0],
+                                        producer_class))
+            self.assertTrue(isinstance(exporter_mock.call_args[0][0][1],
+                                       producer_class))
 
     @mock.patch('opencensus.ext.azure.metrics_exporter'
                 '.transport.get_exporter_thread')
     def test_new_metrics_exporter_no_standard_metrics(self, exporter_mock):
-        iKey = '12345678-1234-5678-abcd-12345678abcd'
-        exporter = metrics_exporter.new_metrics_exporter(
-            instrumentation_key=iKey, enable_standard_metrics=False)
-
-        self.assertEqual(exporter.options.instrumentation_key, iKey)
-        self.assertEqual(len(exporter_mock.call_args_list), 1)
-        self.assertEqual(len(exporter_mock.call_args[0][0]), 1)
-        producer_class = standard_metrics.AzureStandardMetricsProducer
-        self.assertFalse(isinstance(exporter_mock.call_args[0][0][0],
-                                    producer_class))
+        with mock.patch('opencensus.ext.azure.metrics_exporter'
+                        '.heartbeat_metrics.enable_heartbeat_metrics') as hb:
+            hb.return_value = None
+            iKey = '12345678-1234-5678-abcd-12345678abcd'
+            exporter = metrics_exporter.new_metrics_exporter(
+                instrumentation_key=iKey, enable_standard_metrics=False)
+
+            self.assertEqual(exporter.options.instrumentation_key, iKey)
+            self.assertEqual(len(exporter_mock.call_args_list), 1)
+            self.assertEqual(len(exporter_mock.call_args[0][0]), 1)
+            producer_class = standard_metrics.AzureStandardMetricsProducer
+            self.assertFalse(isinstance(exporter_mock.call_args[0][0][0],
+                                        producer_class))
+
+    @mock.patch('opencensus.ext.azure.metrics_exporter'
+                '.transport.get_exporter_thread')
+    def test_new_metrics_exporter_heartbeat(self, exporter_mock):
+        with mock.patch('opencensus.ext.azure.metrics_exporter'
+                        '.heartbeat_metrics.enable_heartbeat_metrics') as hb:
+            iKey = '12345678-1234-5678-abcd-12345678abcd'
+            exporter = metrics_exporter.new_metrics_exporter(
+                instrumentation_key=iKey)
+
+            self.assertEqual(exporter.options.instrumentation_key, iKey)
+            self.assertEqual(len(hb.call_args_list), 1)
+            self.assertEqual(len(hb.call_args[0]), 2)
+            self.assertEqual(hb.call_args[0][0], None)
+            self.assertEqual(hb.call_args[0][1], iKey)

opencensus/metrics/transport.py

@@ -62,6 +62,8 @@ def __init__(
             interval = DEFAULT_INTERVAL
 
         self.func = function
+        self.args = args
+        self.kwargs = kwargs
 
         def func(*aa, **kw):
             try:
@@ -82,6 +84,13 @@ def run(self):
         execution_context.set_is_exporter(True)
         super(PeriodicMetricTask, self).run()
 
+    def close(self):
+        try:
+            return self.func(*self.args, **self.kwargs)
+        except Exception as ex:
+            logger.exception("Error handling metric flush: {}".format(ex))
+        self.cancel()
+
 
 def get_exporter_thread(metric_producers, exporter, interval=None):
     """Get a running task that periodically exports metrics.

tests/unit/metrics/test_transport.py

@@ -58,6 +58,7 @@ def test_periodic_task(self):
         self.assertEqual(mock_func.call_count, 2)
         time.sleep(INTERVAL)
         self.assertEqual(mock_func.call_count, 3)
+        task.cancel()
 
     def test_periodic_task_cancel(self):
         mock_func = mock.Mock()
@@ -69,6 +70,14 @@ def test_periodic_task_cancel(self):
         time.sleep(INTERVAL)
         self.assertEqual(mock_func.call_count, 1)
 
+    def test_periodic_task_close(self):
+        mock_func = mock.Mock()
+        task = transport.PeriodicMetricTask(100, mock_func)
+        task.start()
+        mock_func.assert_not_called()
+        task.close()
+        self.assertEqual(mock_func.call_count, 1)
+
 
 @mock.patch('opencensus.metrics.transport.DEFAULT_INTERVAL', INTERVAL)
 @mock.patch('opencensus.metrics.transport.logger')