Address feedback

Author: rads-1996

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

@@ -156,6 +156,7 @@ def __init__(self) -> None:
         self._initialized: bool = False  # type: ignore
         self._metrics: Optional[_StatsbeatMetrics] = None  # type: ignore
         self._meter_provider: Optional[MeterProvider] = None  # type: ignore
+        self._warmup_timer: Optional[threading.Timer] = None
 
         # Set during first initialization, preserved in shutdown for potential re-initialization
         self._config: Optional[StatsbeatConfig] = None  # type: ignore
@@ -274,10 +275,14 @@ def _flush() -> None:
 
         timer = threading.Timer(_STATSBEAT_INITIAL_EXPORT_WARMUP_SECONDS, _flush)
         timer.daemon = True
+        self._warmup_timer = timer
         timer.start()
 
     def _cleanup(self, shutdown_meter_provider: bool = True) -> None:
         # Clean up resources with optional meter provider shutdown
+        if hasattr(self, '_warmup_timer') and self._warmup_timer:
+            self._warmup_timer.cancel()
+            self._warmup_timer = None
         if shutdown_meter_provider and self._meter_provider:
             try:
                 self._meter_provider.shutdown()

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

@@ -20,6 +20,7 @@
     "LIVE_METRICS_FEATURE_SET": False,
     "CUSTOMER_SDKSTATS_FEATURE_SET": False,
     "BROWSER_SDK_LOADER_FEATURE_SET": False,
+    "FEATURE_ATTRIBUTE_BITS": 0,
 }
 _STATSBEAT_STATE_LOCK = threading.Lock()
 _STATSBEAT_FAILURE_COUNT_THRESHOLD = 3
@@ -115,3 +116,12 @@ def get_statsbeat_browser_sdk_loader_feature_set():  # pylint: disable=name-too-
 def set_statsbeat_browser_sdk_loader_feature_set():  # pylint: disable=name-too-long
     with _STATSBEAT_STATE_LOCK:
         _STATSBEAT_STATE["BROWSER_SDK_LOADER_FEATURE_SET"] = True
+
+
+def get_statsbeat_feature_attribute_bits() -> int:
+    return int(_STATSBEAT_STATE["FEATURE_ATTRIBUTE_BITS"])
+
+
+def set_statsbeat_feature_attribute_bits(feature_bits: int) -> None:
+    with _STATSBEAT_STATE_LOCK:
+        _STATSBEAT_STATE["FEATURE_ATTRIBUTE_BITS"] = int(feature_bits)

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

@@ -33,6 +33,8 @@
 from azure.monitor.opentelemetry.exporter.statsbeat._state import (
     _REQUESTS_MAP_LOCK,
     _REQUESTS_MAP,
+    get_statsbeat_feature_attribute_bits,
+    set_statsbeat_feature_attribute_bits,
     get_statsbeat_live_metrics_feature_set,
     get_statsbeat_custom_events_feature_set,
     get_statsbeat_customer_sdkstats_feature_set,
@@ -133,7 +135,7 @@ def __init__(
             _StatsbeatMetrics._COMMON_ATTRIBUTES["version"] = _get_version()
 
         self._ikey = instrumentation_key
-        self._feature = _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] or _StatsbeatFeature.NONE
+        self._feature = get_statsbeat_feature_attribute_bits()
         if not disable_offline_storage:
             self._feature |= _StatsbeatFeature.DISK_RETRY
         if has_credential:
@@ -166,6 +168,7 @@ def __init__(
 
         _StatsbeatMetrics._NETWORK_ATTRIBUTES["host"] = _shorten_host(endpoint)
         _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = self._feature
+        set_statsbeat_feature_attribute_bits(self._feature)
         _StatsbeatMetrics._INSTRUMENTATION_ATTRIBUTES["feature"] = _utils.get_instrumentations()
 
         self._vm_retry = True  # True if we want to attempt to find if in VM
@@ -264,22 +267,27 @@ def _get_feature_metric(self, options: CallbackOptions) -> Iterable[Observation]
             return observations
         # Feature metric
         # Check if any features were enabled during runtime
-        feature_bits = int(_StatsbeatMetrics._FEATURE_ATTRIBUTES.get("feature") or 0)
+        feature_bits = get_statsbeat_feature_attribute_bits()
         if feature_bits:
             self._feature |= feature_bits
             _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = self._feature
+            set_statsbeat_feature_attribute_bits(self._feature)
         if get_statsbeat_custom_events_feature_set():
             self._feature |= _StatsbeatFeature.CUSTOM_EVENTS_EXTENSION
             _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = self._feature
+            set_statsbeat_feature_attribute_bits(self._feature)
         if get_statsbeat_live_metrics_feature_set():
             self._feature |= _StatsbeatFeature.LIVE_METRICS
             _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = self._feature
+            set_statsbeat_feature_attribute_bits(self._feature)
         if get_statsbeat_customer_sdkstats_feature_set():
             self._feature |= _StatsbeatFeature.CUSTOMER_SDKSTATS
             _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = self._feature
+            set_statsbeat_feature_attribute_bits(self._feature)
         if get_statsbeat_browser_sdk_loader_feature_set():
             self._feature |= _StatsbeatFeature.BROWSER_SDK_LOADER
             _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = self._feature
+            set_statsbeat_feature_attribute_bits(self._feature)
 
         # Don't send observation if no features enabled
         if self._feature is not _StatsbeatFeature.NONE:

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

@@ -270,6 +270,8 @@ def setUp(self):
             self.manager._metrics = None
         if hasattr(self.manager, "_meter_provider"):
             self.manager._meter_provider = None
+        if hasattr(self.manager, "_warmup_timer"):
+            self.manager._warmup_timer = None
 
     def tearDown(self):
         """Clean up after tests."""
@@ -673,6 +675,8 @@ def test_cleanup_with_shutdown(self):
         self.manager._initialized = True
         mock_meter_provider = Mock()
         self.manager._meter_provider = mock_meter_provider
+        mock_timer = Mock()
+        self.manager._warmup_timer = mock_timer
         self.manager._metrics = Mock()
         config_mock = Mock()
         self.manager._config = config_mock
@@ -682,8 +686,10 @@ def test_cleanup_with_shutdown(self):
         self.assertFalse(self.manager._initialized)
         self.assertIsNone(self.manager._meter_provider)
         self.assertIsNone(self.manager._metrics)
+        self.assertIsNone(self.manager._warmup_timer)
         # Config is intact for potential re-initialization
         self.assertEqual(self.manager._config, config_mock)
+        mock_timer.cancel.assert_called_once()
         mock_meter_provider.shutdown.assert_called_once()
 
     def test_cleanup_without_shutdown(self):
@@ -692,6 +698,8 @@ def test_cleanup_without_shutdown(self):
         self.manager._initialized = True
         mock_meter_provider = Mock()
         self.manager._meter_provider = mock_meter_provider
+        mock_timer = Mock()
+        self.manager._warmup_timer = mock_timer
         self.manager._metrics = Mock()
         config_mock = Mock()
         self.manager._config = config_mock
@@ -701,8 +709,10 @@ def test_cleanup_without_shutdown(self):
         self.assertFalse(self.manager._initialized)
         self.assertIsNone(self.manager._meter_provider)
         self.assertIsNone(self.manager._metrics)
+        self.assertIsNone(self.manager._warmup_timer)
         # Config is intact for potential re-initialization
         self.assertEqual(self.manager._config, config_mock)
+        mock_timer.cancel.assert_called_once()
         mock_meter_provider.shutdown.assert_not_called()
 
     def test_cleanup_meter_provider_exception(self):

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

@@ -24,6 +24,8 @@
     _REQUESTS_MAP,
     _STATSBEAT_STATE,
     _STATSBEAT_STATE_LOCK,
+    get_statsbeat_feature_attribute_bits,
+    set_statsbeat_feature_attribute_bits,
 )
 from azure.monitor.opentelemetry.exporter.statsbeat._statsbeat_metrics import (
     _shorten_host,
@@ -87,6 +89,7 @@ def setUp(self):
             _STATSBEAT_STATE["CUSTOM_EVENTS_FEATURE_SET"] = False
             _STATSBEAT_STATE["LIVE_METRICS_FEATURE_SET"] = False
             _STATSBEAT_STATE["CUSTOMER_SDKSTATS_FEATURE_SET"] = False
+            _STATSBEAT_STATE["FEATURE_ATTRIBUTE_BITS"] = 0
 
         _StatsbeatMetrics._COMMON_ATTRIBUTES = dict(_StatsbeatMetrics_COMMON_ATTRS)
         _StatsbeatMetrics._NETWORK_ATTRIBUTES = dict(_StatsbeatMetrics_NETWORK_ATTRS)
@@ -126,6 +129,14 @@ def test_statsbeat_metric_init(self):
         self.assertEqual(metric._attach_metric.name, _ATTACH_METRIC_NAME[0])
         self.assertEqual(metric._feature_metric.name, _FEATURE_METRIC_NAME[0])
 
+    def test_statsbeat_feature_attribute_bits_getter_default(self):
+        self.assertEqual(get_statsbeat_feature_attribute_bits(), 0)
+
+    def test_statsbeat_feature_attribute_bits_setter_and_getter(self):
+        feature_bits = _StatsbeatFeature.DISK_RETRY | _StatsbeatFeature.LIVE_METRICS
+        set_statsbeat_feature_attribute_bits(feature_bits)
+        self.assertEqual(get_statsbeat_feature_attribute_bits(), feature_bits)
+
     @mock.patch("azure.monitor.opentelemetry.exporter._utils._is_attach_enabled")
     def test_statsbeat_metric_init_attach_enabled(self, attach_mock):
         mp = MeterProvider()
@@ -717,6 +728,7 @@ def test_get_feature_metric_instrumentation(self):
         )
         metric._feature = _StatsbeatFeature.NONE
         _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = _StatsbeatFeature.NONE
+        set_statsbeat_feature_attribute_bits(_StatsbeatFeature.NONE)
         attributes = dict(_StatsbeatMetrics._COMMON_ATTRIBUTES)
         attributes.update(_StatsbeatMetrics._INSTRUMENTATION_ATTRIBUTES)
         self.assertEqual(attributes["type"], _FEATURE_TYPES.INSTRUMENTATION)
@@ -749,6 +761,7 @@ def test_get_feature_metric_instrumentation_none(self):
         )
         metric._feature = _StatsbeatFeature.NONE
         _StatsbeatMetrics._FEATURE_ATTRIBUTES["feature"] = _StatsbeatFeature.NONE
+        set_statsbeat_feature_attribute_bits(_StatsbeatFeature.NONE)
         self.assertEqual(_StatsbeatMetrics._INSTRUMENTATION_ATTRIBUTES["feature"], 0)
         with mock.patch("azure.monitor.opentelemetry.exporter._utils.get_instrumentations") as instrumentations:
             instrumentations.return_value = 0

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

@@ -43,6 +43,7 @@ def setUp(self):
             _STATSBEAT_STATE["CUSTOM_EVENTS_FEATURE_SET"] = False
             _STATSBEAT_STATE["LIVE_METRICS_FEATURE_SET"] = False
             _STATSBEAT_STATE["CUSTOMER_SDKSTATS_FEATURE_SET"] = False
+            _STATSBEAT_STATE["FEATURE_ATTRIBUTE_BITS"] = 0
         _StatsbeatMetrics._FEATURE_ATTRIBUTES = dict(_StatsbeatMetrics_FEATURE_ATTRIBUTES)
 
     def tearDown(self):