Add custom events extension tracking in feature statsbeat (#33667)

Author: lzchen

sdk/monitor/azure-monitor-opentelemetry-exporter/CHANGELOG.md

@@ -14,6 +14,8 @@
     ([#33432](https://github.com/Azure/azure-sdk-for-python/pull/33432))
 - Shutdown statsbeat on customer getting 400 error code
     ([#33489](https://github.com/Azure/azure-sdk-for-python/pull/33489))
+- Track custom events extension in feature statsbeat
+    ([#33667](https://github.com/Azure/azure-sdk-for-python/pull/33667))
 - Readme examples are updated with correct imports
     ([#33691](https://github.com/Azure/azure-sdk-for-python/pull/33691))
 

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/export/logs/_exporter.py

@@ -26,6 +26,12 @@
     BaseExporter,
     ExportResult,
 )
+from azure.monitor.opentelemetry.exporter.statsbeat._state import (
+    get_statsbeat_shutdown,
+    get_statsbeat_custom_events_feature_set,
+    is_statsbeat_enabled,
+    set_statsbeat_custom_events_feature_set,
+)
 
 _logger = logging.getLogger(__name__)
 
@@ -127,6 +133,7 @@ def _convert_log_to_envelope(log_data: LogData) -> TelemetryItem:
 
     # Event telemetry
     if _log_data_is_event(log_data):
+        _set_statsbeat_custom_events_feature()
         envelope.name = 'Microsoft.ApplicationInsights.Event'
         data = TelemetryEventData(
             name=str(log_record.body)[:32768],
@@ -196,3 +203,7 @@ def _is_ignored_attribute(key: str) -> bool:
         _APPLICATION_INSIGHTS_EVENT_MARKER_ATTRIBUTE,
     )
 )
+
+def _set_statsbeat_custom_events_feature():
+    if is_statsbeat_enabled() and not get_statsbeat_shutdown() and not get_statsbeat_custom_events_feature_set():
+        set_statsbeat_custom_events_feature_set()

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/statsbeat/_state.py

@@ -13,6 +13,7 @@
     "INITIAL_FAILURE_COUNT": 0,
     "INITIAL_SUCCESS": False,
     "SHUTDOWN": False,
+    "CUSTOM_EVENTS_FEATURE_SET": False,
 }
 _STATSBEAT_STATE_LOCK = threading.Lock()
 _STATSBEAT_FAILURE_COUNT_THRESHOLD = 3
@@ -47,3 +48,12 @@ def get_statsbeat_initial_success():
 
 def get_statsbeat_shutdown():
     return _STATSBEAT_STATE["SHUTDOWN"]
+
+
+def get_statsbeat_custom_events_feature_set():
+    return _STATSBEAT_STATE["CUSTOM_EVENTS_FEATURE_SET"]
+
+
+def set_statsbeat_custom_events_feature_set():
+    with _STATSBEAT_STATE_LOCK:
+        _STATSBEAT_STATE["CUSTOM_EVENTS_FEATURE_SET"] = True

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/statsbeat/_statsbeat_metrics.py

@@ -27,6 +27,7 @@
 from azure.monitor.opentelemetry.exporter.statsbeat._state import (
     _REQUESTS_MAP_LOCK,
     _REQUESTS_MAP,
+    get_statsbeat_custom_events_feature_set,
 )
 from azure.monitor.opentelemetry.exporter import _utils
 
@@ -51,6 +52,7 @@ class _StatsbeatFeature:
     NONE = 0
     DISK_RETRY = 1
     AAD = 2
+    CUSTOM_EVENTS_EXTENSION = 4
 
 
 class _AttachTypes:
@@ -102,6 +104,8 @@ def __init__(
             self._feature |= _StatsbeatFeature.DISK_RETRY
         if has_credential:
             self._feature |= _StatsbeatFeature.AAD
+        if get_statsbeat_custom_events_feature_set():
+            self._feature |= _StatsbeatFeature.CUSTOM_EVENTS_EXTENSION
         self._ikey = instrumentation_key
         self._meter_provider = meter_provider
         self._meter = self._meter_provider.get_meter(__name__)
@@ -215,6 +219,11 @@ def _get_feature_metric(self, options: CallbackOptions) -> Iterable[Observation]
         if not self._meets_long_interval_threshold(_FEATURE_METRIC_NAME[0]):
             return observations
         # Feature metric
+        # Check if any features were enabled during runtime
+        if get_statsbeat_custom_events_feature_set():
+            self._feature |= _StatsbeatFeature.CUSTOM_EVENTS_EXTENSION
+            _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = self._feature
+
         # Don't send observation if no features enabled
         if self._feature is not _StatsbeatFeature.NONE:
             attributes = dict(_StatsbeatMetrics._COMMON_ATTRIBUTES)

sdk/monitor/azure-monitor-opentelemetry-exporter/tests/statsbeat/test_statsbeat.py

@@ -59,6 +59,7 @@ def func(*_args, **_kwargs):
 class TestStatsbeat(unittest.TestCase):
     def setUp(self):
         _statsbeat._STATSBEAT_METRICS = None
+        _STATSBEAT_STATE["SHUTDOWN"] = False
 
     @mock.patch.object(MeterProvider, 'shutdown')
     @mock.patch.object(MeterProvider, 'force_flush')
@@ -217,11 +218,11 @@ def test_collect_statsbeat_metrics_no_aad(
         )
 
     def test_shutdown_statsbeat_metrics(self):
-        _STATSBEAT_STATE["SHUTDOWN"] = False
         metric_mock = mock.Mock()
         mp_mock = mock.Mock()
         metric_mock._meter_provider = mp_mock
         _statsbeat._STATSBEAT_METRICS = metric_mock
+        self.assertFalse(_STATSBEAT_STATE["SHUTDOWN"])
         _statsbeat.shutdown_statsbeat_metrics()
         mp_mock.shutdown.assert_called_once()
         self.assertIsNone(_statsbeat._STATSBEAT_METRICS)
@@ -274,6 +275,11 @@ def setUp(self):
             _StatsbeatMetrics_INSTRUMENTATION_ATTRIBUTES
         )
         _REQUESTS_MAP.clear()
+        _STATSBEAT_STATE["INITIAL_FAILURE_COUNT"] = 0
+        _STATSBEAT_STATE["INITIAL_SUCCESS"] = False
+        _STATSBEAT_STATE["SHUTDOWN"] = False
+        _STATSBEAT_STATE["CUSTOM_EVENTS_FEATURE_SET"] = False
+
 
     def test_statsbeat_metric_init(self):
         mp = MeterProvider()
@@ -685,6 +691,55 @@ def test_get_feature_metric_aad(self):
         self.assertEqual(len(observations), 1)
         self.assertEqual(_StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"], _StatsbeatFeature.AAD)
 
+    # pylint: disable=protected-access
+    def test_get_feature_metric_custom_events_init(self):
+        mp = MeterProvider()
+        ikey = "1aa11111-bbbb-1ccc-8ddd-eeeeffff3334"
+        endpoint = "https://westus-1.in.applicationinsights.azure.com/"
+        _STATSBEAT_STATE["CUSTOM_EVENTS_FEATURE_SET"] = True
+        metric = _StatsbeatMetrics(
+            mp,
+            ikey,
+            endpoint,
+            True,
+            0,
+            False,
+        )
+        attributes = dict(_StatsbeatMetrics._COMMON_ATTRIBUTES)
+        attributes.update(_StatsbeatMetrics._FEATURE_ATTRIBUTES)
+        self.assertEqual(attributes["feature"], 4)
+        self.assertEqual(attributes["type"], _FEATURE_TYPES.FEATURE)
+        observations = metric._get_feature_metric(options=None)
+        for obs in observations:
+            self.assertEqual(obs.value, 1)
+            self.assertEqual(obs.attributes, attributes)
+
+    # pylint: disable=protected-access
+    def test_get_feature_metric_custom_events_runtime(self):
+        mp = MeterProvider()
+        ikey = "1aa11111-bbbb-1ccc-8ddd-eeeeffff3334"
+        endpoint = "https://westus-1.in.applicationinsights.azure.com/"
+        metric = _StatsbeatMetrics(
+            mp,
+            ikey,
+            endpoint,
+            True,
+            0,
+            False,
+        )
+        attributes = dict(_StatsbeatMetrics._COMMON_ATTRIBUTES)
+        attributes.update(_StatsbeatMetrics._FEATURE_ATTRIBUTES)
+        self.assertEqual(attributes["feature"], 0)
+        self.assertEqual(attributes["type"], _FEATURE_TYPES.FEATURE)
+        _STATSBEAT_STATE["CUSTOM_EVENTS_FEATURE_SET"] = True
+        observations = metric._get_feature_metric(options=None)
+        attributes.update(_StatsbeatMetrics._FEATURE_ATTRIBUTES)
+        self.assertEqual(attributes["feature"], 4)
+        self.assertEqual(metric._feature, 4)
+        for obs in observations:
+            self.assertEqual(obs.value, 1)
+            self.assertEqual(obs.attributes, attributes)
+
     # pylint: disable=protected-access
     def test_get_feature_metric_instrumentation(self):
         mp = MeterProvider()