Address feedback

Author: rads-1996

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

@@ -6,6 +6,7 @@
 - Add `StatsbeatManager.add_metric_callback` to let SDKs/distros add their own metric
   observations to built-in statsbeat metrics
   ([#47363](https://github.com/Azure/azure-sdk-for-python/pull/47363))
+
 ### Breaking Changes
 - Customer Facing SDKStats: Renamed metric dimension attributes from snake_case/dotted to camelCase
   (`compute_type` -> `computeType`, `telemetry_type` -> `telemetryType`, `telemetry_success` -> `telemetrySuccess`,

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

@@ -2,7 +2,7 @@
 # Licensed under the MIT License.
 import logging
 import threading
-from typing import Callable, Iterable, Optional, Any, Dict
+from typing import Callable, Iterable, List, Optional, Any, Dict
 
 from opentelemetry.metrics import CallbackOptions, Observation
 from opentelemetry.sdk.metrics import MeterProvider
@@ -20,8 +20,6 @@
     _get_stats_long_export_interval,
     _get_stats_short_export_interval,
     _get_connection_string_for_region_from_config,
-    _ADDITIONAL_CALLBACKS,
-    _ADDITIONAL_CALLBACKS_LOCK,
 )
 from azure.monitor.opentelemetry.exporter._utils import Singleton
 
@@ -164,6 +162,11 @@ def __init__(self) -> None:
         # Set during first initialization, preserved in shutdown for potential re-initialization
         self._config: Optional[StatsbeatConfig] = None  # type: ignore
 
+        # Extra observation callbacks contributed by SDKs/distros. Keyed by built-in
+        # statsbeat metric name. Registered directly on the singleton instance, e.g.
+        # ``StatsbeatManager()._additional_callbacks.setdefault(name, []).append(cb)``.
+        self._additional_callbacks: Dict[str, List[Callable[[CallbackOptions], Iterable[Observation]]]] = {}
+
     @staticmethod
     def _validate_config(config: Optional[StatsbeatConfig]) -> bool:
         """Validate that a configuration has all required fields.
@@ -381,24 +384,4 @@ def is_initialized(self) -> bool:
         with self._lock:
             return self._initialized
 
-    def add_metric_callback(
-        self,
-        metric_name: str,
-        callback: Callable[[CallbackOptions], Iterable[Observation]],
-    ) -> bool:
-        """Register an extra observation callback that an SDK/Distro with its own network sdkstats metric can use to
-        contribute rows to a built-in statsbeat metric.
-
-        :param metric_name: Name of the built-in statsbeat metric to extend.
-        :type metric_name: str
-        :param callback: OpenTelemetry observable-gauge callback ``(CallbackOptions) -> Iterable[Observation]``.
-        :type callback: Callable[[CallbackOptions], Iterable[Observation]]
-        :returns: ``True`` if newly registered, ``False`` if already registered.
-        :rtype: bool
-        """
-        with _ADDITIONAL_CALLBACKS_LOCK:
-            callbacks = _ADDITIONAL_CALLBACKS.setdefault(metric_name, [])
-            if callback in callbacks:
-                return False
-            callbacks.append(callback)
-            return True
+

sdk/monitor/azure-monitor-opentelemetry-exporter/api.metadata.yml

@@ -41,7 +41,7 @@
     get_statsbeat_browser_sdk_loader_feature_set,
 )
 from azure.monitor.opentelemetry.exporter.statsbeat._utils import (
-    _iter_additional_observations,
+    _get_additional_observations,
 )
 from azure.monitor.opentelemetry.exporter import _utils
 
@@ -382,7 +382,7 @@ def _get_success_count(self, options: CallbackOptions) -> Iterable[Observation]:
             if count != 0:
                 observations.append(Observation(int(count), dict(attributes)))
                 _REQUESTS_MAP[_REQ_SUCCESS_NAME[1]] = 0
-        observations.extend(_iter_additional_observations(_REQ_SUCCESS_NAME[0], options))
+        observations.extend(_get_additional_observations(_REQ_SUCCESS_NAME[0], options))
         return observations
 
     # pylint: disable=unused-argument
@@ -397,7 +397,7 @@ def _get_failure_count(self, options: CallbackOptions) -> Iterable[Observation]:
                     attributes["statusCode"] = code
                     observations.append(Observation(int(count), dict(attributes)))
                     _REQUESTS_MAP[_REQ_FAILURE_NAME[1]][code] = 0  # type: ignore
-        observations.extend(_iter_additional_observations(_REQ_FAILURE_NAME[0], options))
+        observations.extend(_get_additional_observations(_REQ_FAILURE_NAME[0], options))
         return observations
 
     # pylint: disable=unused-argument
@@ -414,7 +414,7 @@ def _get_average_duration(self, options: CallbackOptions) -> Iterable[Observatio
                 observations.append(Observation(result * 1000, dict(attributes)))
                 _REQUESTS_MAP[_REQ_DURATION_NAME[1]] = 0
                 _REQUESTS_MAP["count"] = 0
-        observations.extend(_iter_additional_observations(_REQ_DURATION_NAME[0], options))
+        observations.extend(_get_additional_observations(_REQ_DURATION_NAME[0], options))
         return observations
 
     # pylint: disable=unused-argument
@@ -429,7 +429,7 @@ def _get_retry_count(self, options: CallbackOptions) -> Iterable[Observation]:
                     attributes["statusCode"] = code
                     observations.append(Observation(int(count), dict(attributes)))
                     _REQUESTS_MAP[_REQ_RETRY_NAME[1]][code] = 0  # type: ignore
-        observations.extend(_iter_additional_observations(_REQ_RETRY_NAME[0], options))
+        observations.extend(_get_additional_observations(_REQ_RETRY_NAME[0], options))
         return observations
 
     # pylint: disable=unused-argument
@@ -444,7 +444,7 @@ def _get_throttle_count(self, options: CallbackOptions) -> Iterable[Observation]
                     attributes["statusCode"] = code
                     observations.append(Observation(int(count), dict(attributes)))
                     _REQUESTS_MAP[_REQ_THROTTLE_NAME[1]][code] = 0  # type: ignore
-        observations.extend(_iter_additional_observations(_REQ_THROTTLE_NAME[0], options))
+        observations.extend(_get_additional_observations(_REQ_THROTTLE_NAME[0], options))
         return observations
 
     # pylint: disable=unused-argument
@@ -459,7 +459,7 @@ def _get_exception_count(self, options: CallbackOptions) -> Iterable[Observation
                     attributes["exceptionType"] = code
                     observations.append(Observation(int(count), dict(attributes)))
                     _REQUESTS_MAP[_REQ_EXCEPTION_NAME[1]][code] = 0  # type: ignore
-        observations.extend(_iter_additional_observations(_REQ_EXCEPTION_NAME[0], options))
+        observations.extend(_get_additional_observations(_REQ_EXCEPTION_NAME[0], options))
         return observations
 
 

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

@@ -3,8 +3,7 @@
 import os
 import logging
 import json
-import threading
-from collections.abc import Iterable, Callable, Iterator  # pylint: disable=import-error
+from collections.abc import Iterable  # pylint: disable=import-error
 from typing import Optional, Dict, List
 from opentelemetry.metrics import CallbackOptions, Observation
 
@@ -28,9 +27,6 @@
     _REQUESTS_MAP_LOCK,
 )
 
-_ADDITIONAL_CALLBACKS: Dict[str, List[Callable[[CallbackOptions], Iterable[Observation]]]] = {}
-_ADDITIONAL_CALLBACKS_LOCK = threading.Lock()
-
 
 def _get_stats_connection_string(endpoint: str) -> str:
     cs_env = os.environ.get(_APPLICATIONINSIGHTS_STATS_CONNECTION_STRING_ENV_NAME)
@@ -172,34 +168,36 @@ def _get_connection_string_for_region_from_config(target_region: str, settings:
         return None
 
 
-def _iter_additional_observations(metric_name: str, options: CallbackOptions) -> Iterator[Observation]:
-    """Yield observations contributed via :func:`add_metric_callback`.
+def _get_additional_observations(metric_name: str, options: CallbackOptions) -> List[Observation]:
+    """Return observations contributed by extra callbacks registered on :class:`StatsbeatManager`.
 
     Invoked by the built-in ``_StatsbeatMetrics`` callbacks at collection time.
-    Snapshots the registered callbacks under the registry lock to avoid
-    mutation during iteration, then releases the lock before invoking user
-    callbacks (so they cannot deadlock against the registry). Exceptions raised
-    by individual callbacks are caught, logged, and skipped.
+    Reads ``StatsbeatManager()._additional_callbacks`` (a live mutable dict on the
+    singleton instance), which SDKs/distros populate directly. Exceptions raised by
+    individual callbacks are caught, logged, and skipped.
 
     :param metric_name: Name of the built-in statsbeat metric being collected.
     :type metric_name: str
     :param options: OpenTelemetry callback options forwarded to each registered callback.
     :type options: ~opentelemetry.metrics.CallbackOptions
-    :returns: Iterator over observations contributed by registered callbacks.
-    :rtype: Iterator[~opentelemetry.metrics.Observation]
+    :returns: List of observations contributed by registered callbacks.
+    :rtype: list[~opentelemetry.metrics.Observation]
     """
+    # Lazy import to avoid a circular import between _manager and _utils.
+    from azure.monitor.opentelemetry.exporter.statsbeat._manager import StatsbeatManager  # pylint: disable=import-outside-toplevel
 
-    with _ADDITIONAL_CALLBACKS_LOCK:
-        callbacks = tuple(_ADDITIONAL_CALLBACKS.get(metric_name, ()))
+    callbacks = StatsbeatManager()._additional_callbacks.get(metric_name, ())  # pylint: disable=protected-access
 
+    observations: List[Observation] = []
     iter_logger = logging.getLogger(__name__)
     for cb in callbacks:
         try:
-            yield from cb(options)
+            observations.extend(cb(options))
         except Exception:  # pylint: disable=broad-except
             iter_logger.debug(
                 "Extra statsbeat callback %r for %r raised; skipping.",
                 cb,
                 metric_name,
                 exc_info=True,
             )
+    return observations