Adding customer sdkstats to feature statsbeat (#43066)

Author: rads-1996

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

@@ -3,9 +3,10 @@
 ## 1.0.0b43 (Unreleased)
 
 ### Features Added
-
 - OneSettings control plane: Add live metrics feature control
   ([#43032](https://github.com/Azure/azure-sdk-for-python/pull/43032))
+- Adding customer sdkstats to feature statsbeat
+  ([#43066](https://github.com/Azure/azure-sdk-for-python/pull/43066))
 
 ### Breaking Changes
 

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

@@ -17,6 +17,7 @@
     "SHUTDOWN": False,
     "CUSTOM_EVENTS_FEATURE_SET": False,
     "LIVE_METRICS_FEATURE_SET": False,
+    "CUSTOMER_SDKSTATS_FEATURE_SET": False,
 }
 _STATSBEAT_STATE_LOCK = threading.Lock()
 _STATSBEAT_FAILURE_COUNT_THRESHOLD = 3
@@ -74,3 +75,10 @@ def set_statsbeat_live_metrics_feature_set():
 def set_statsbeat_shutdown(shutdown: bool):
     with _STATSBEAT_STATE_LOCK:
         _STATSBEAT_STATE["SHUTDOWN"] = shutdown
+
+def get_statsbeat_customer_sdkstats_feature_set(): # pylint: disable=name-too-long
+    return _STATSBEAT_STATE["CUSTOMER_SDKSTATS_FEATURE_SET"]
+
+def set_statsbeat_customer_sdkstats_feature_set(): # pylint: disable=name-too-long
+    with _STATSBEAT_STATE_LOCK:
+        _STATSBEAT_STATE["CUSTOMER_SDKSTATS_FEATURE_SET"] = True

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

@@ -35,6 +35,7 @@
     _REQUESTS_MAP,
     get_statsbeat_live_metrics_feature_set,
     get_statsbeat_custom_events_feature_set,
+    get_statsbeat_customer_sdkstats_feature_set,
 )
 from azure.monitor.opentelemetry.exporter import _utils
 
@@ -76,6 +77,7 @@ class _StatsbeatFeature:
     CUSTOM_EVENTS_EXTENSION = 4
     DISTRO = 8
     LIVE_METRICS = 16
+    CUSTOMER_SDKSTATS = 32
 
 
 class _AttachTypes:
@@ -138,6 +140,8 @@ def __init__(
             self._feature |= _StatsbeatFeature.CUSTOM_EVENTS_EXTENSION
         if get_statsbeat_live_metrics_feature_set():
             self._feature |= _StatsbeatFeature.LIVE_METRICS
+        if get_statsbeat_customer_sdkstats_feature_set():
+            self._feature |= _StatsbeatFeature.CUSTOMER_SDKSTATS
         self._ikey = instrumentation_key
         self._meter_provider = meter_provider
         self._meter = self._meter_provider.get_meter(__name__)
@@ -260,6 +264,9 @@ def _get_feature_metric(self, options: CallbackOptions) -> Iterable[Observation]
         if get_statsbeat_live_metrics_feature_set():
             self._feature |= _StatsbeatFeature.LIVE_METRICS
             _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = self._feature
+        if get_statsbeat_customer_sdkstats_feature_set():
+            self._feature |= _StatsbeatFeature.CUSTOMER_SDKSTATS
+            _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = self._feature
 
         # Don't send observation if no features enabled
         if self._feature is not _StatsbeatFeature.NONE:

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/statsbeat/customer/_customer_sdkstats.py

@@ -3,6 +3,8 @@
 
 from typing import TYPE_CHECKING
 
+from azure.monitor.opentelemetry.exporter.statsbeat._state import set_statsbeat_customer_sdkstats_feature_set
+
 from ._state import get_customer_stats_manager
 
 if TYPE_CHECKING:
@@ -19,6 +21,7 @@ def collect_customer_sdkstats(exporter: "BaseExporter") -> None:  # type: ignore
     if not customer_stats.is_initialized:
         # The initialize method is thread-safe and handles double-initialization
         customer_stats.initialize(connection_string=exporter._connection_string)  # type: ignore
+    set_statsbeat_customer_sdkstats_feature_set()
 
 def shutdown_customer_sdkstats_metrics() -> None:
     # Shutdown customer SDKStats metrics collection.

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

@@ -85,6 +85,7 @@ def setUp(self):
             _STATSBEAT_STATE["SHUTDOWN"] = False
             _STATSBEAT_STATE["CUSTOM_EVENTS_FEATURE_SET"] = False
             _STATSBEAT_STATE["LIVE_METRICS_FEATURE_SET"] = False
+            _STATSBEAT_STATE["CUSTOMER_SDKSTATS_FEATURE_SET"] = False
 
         _StatsbeatMetrics._COMMON_ATTRIBUTES = dict(_StatsbeatMetrics_COMMON_ATTRS)
         _StatsbeatMetrics._NETWORK_ATTRIBUTES = dict(_StatsbeatMetrics_NETWORK_ATTRS)
@@ -635,6 +636,60 @@ def test_get_feature_metric_live_metrics_runtime(self):
         self.assertTrue((_StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] >> 4) & 1 == 1)
         _STATSBEAT_STATE["LIVE_METRICS_FEATURE_SET"] = False
 
+    # pylint: disable=protected-access
+    @mock.patch(
+        "azure.monitor.opentelemetry.exporter.statsbeat._statsbeat_metrics.get_statsbeat_customer_sdkstats_feature_set"
+    )
+    def test_get_feature_metric_customer_sdkstats(self, feature_mock):
+        feature_mock.return_value = True
+        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"], 32)
+        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_customer_sdkstats_runtime(self):
+        mp = MeterProvider()
+        ikey = "1aa11111-bbbb-1ccc-8ddd-eeeeffff3334"
+        endpoint = "https://westus-1.in.applicationinsights.azure.com/"
+        _STATSBEAT_STATE["CUSTOMER_SDKSTATS_FEATURE_SET"] = False
+        metric = _StatsbeatMetrics(
+            mp,
+            ikey,
+            endpoint,
+            True,
+            0,
+            False,
+        )
+        self.assertTrue((metric._feature >> 5) & 1 == 0)
+        attributes = dict(_StatsbeatMetrics._COMMON_ATTRIBUTES)
+        attributes.update(_StatsbeatMetrics._FEATURE_ATTRIBUTES)
+        self.assertEqual(attributes["feature"], 0)
+        self.assertEqual(attributes["type"], _FEATURE_TYPES.FEATURE)
+        _STATSBEAT_STATE["CUSTOMER_SDKSTATS_FEATURE_SET"] = True
+        observations = metric._get_feature_metric(options=None)
+        attributes["feature"] = _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"]
+        for obs in observations:
+            self.assertEqual(obs.value, 1)
+            self.assertEqual(obs.attributes, attributes)
+        self.assertTrue((_StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] >> 5) & 1 == 1)
+        _STATSBEAT_STATE["CUSTOMER_SDKSTATS_FEATURE_SET"] = False
+
     # pylint: disable=protected-access
     def test_get_feature_metric_distro(self):
         mp = MeterProvider()

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

@@ -39,6 +39,7 @@ def setUp(self):
             _STATSBEAT_STATE["SHUTDOWN"] = False
             _STATSBEAT_STATE["CUSTOM_EVENTS_FEATURE_SET"] = False
             _STATSBEAT_STATE["LIVE_METRICS_FEATURE_SET"] = False
+            _STATSBEAT_STATE["CUSTOMER_SDKSTATS_FEATURE_SET"] = False
 
     def tearDown(self):
         """Clean up after tests."""