Add Fixed-rate sampling logic for Azure Log Exporter (#848)

Author: lzchen

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

@@ -92,6 +92,7 @@ def __init__(self, *args, **kwargs):
         export_interval=15.0,
         grace_period=5.0,
         instrumentation_key=None,
+        logging_sampling_rate=1.0,
         max_batch_size=100,
         minimum_retry_interval=60,  # minimum retry interval in seconds
         proxy=None,

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

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import random
 import threading
 import time
 import traceback
@@ -107,6 +108,16 @@ def stop(self, timeout=None):  # pragma: NO COVER
             return time.time() - start_time  # time taken to stop
 
 
+class SamplingFilter(logging.Filter):
+
+    def __init__(self, probability=1.0):
+        super(SamplingFilter, self).__init__()
+        self.probability = probability
+
+    def filter(self, record):
+        return random.random() < self.probability
+
+
 class AzureLogHandler(TransportMixin, BaseLogHandler):
     """Handler for logging to Microsoft Azure Monitor.
 
@@ -116,6 +127,8 @@ class AzureLogHandler(TransportMixin, BaseLogHandler):
     def __init__(self, **options):
         self.options = Options(**options)
         utils.validate_instrumentation_key(self.options.instrumentation_key)
+        if not 0 <= self.options.logging_sampling_rate <= 1:
+            raise ValueError('Sampling must be in the range: [0,1]')
         self.export_interval = self.options.export_interval
         self.max_batch_size = self.options.max_batch_size
         self.storage = LocalFileStorage(
@@ -125,6 +138,7 @@ def __init__(self, **options):
             retention_period=self.options.storage_retention_period,
         )
         super(AzureLogHandler, self).__init__()
+        self.addFilter(SamplingFilter(self.options.logging_sampling_rate))
 
     def close(self):
         self.storage.close()

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

@@ -78,6 +78,13 @@ def test_ctor(self):
         self.assertRaises(ValueError, lambda: log_exporter.AzureLogHandler())
         Options._default.instrumentation_key = instrumentation_key
 
+    def test_invalid_sampling_rate(self):
+        with self.assertRaises(ValueError):
+            log_exporter.AzureLogHandler(
+                instrumentation_key='12345678-1234-5678-abcd-12345678abcd',
+                logging_sampling_rate=4.0,
+            )
+
     @mock.patch('requests.post', return_value=mock.Mock())
     def test_exception(self, requests_mock):
         logger = logging.getLogger(self.id())
@@ -207,3 +214,37 @@ def test_log_with_invalid_custom_properties(self, requests_mock):
 
         self.assertFalse('not_a_dict' in post_body)
         self.assertFalse('key_1' in post_body)
+
+    @mock.patch('requests.post', return_value=mock.Mock())
+    def test_log_record_sampled(self, requests_mock):
+        logger = logging.getLogger(self.id())
+        handler = log_exporter.AzureLogHandler(
+            instrumentation_key='12345678-1234-5678-abcd-12345678abcd',
+            logging_sampling_rate=1.0,
+        )
+        logger.addHandler(handler)
+        logger.warning('Hello_World')
+        logger.warning('Hello_World2')
+        logger.warning('Hello_World3')
+        logger.warning('Hello_World4')
+        handler.close()
+        post_body = requests_mock.call_args_list[0][1]['data']
+        self.assertTrue('Hello_World' in post_body)
+        self.assertTrue('Hello_World2' in post_body)
+        self.assertTrue('Hello_World3' in post_body)
+        self.assertTrue('Hello_World4' in post_body)
+
+    @mock.patch('requests.post', return_value=mock.Mock())
+    def test_log_record_not_sampled(self, requests_mock):
+        logger = logging.getLogger(self.id())
+        handler = log_exporter.AzureLogHandler(
+            instrumentation_key='12345678-1234-5678-abcd-12345678abcd',
+            logging_sampling_rate=0.0,
+        )
+        logger.addHandler(handler)
+        logger.warning('Hello_World')
+        logger.warning('Hello_World2')
+        logger.warning('Hello_World3')
+        logger.warning('Hello_World4')
+        handler.close()
+        self.assertFalse(requests_mock.called)