Refactor customer sdk stats (#42969)

Author: lzchen

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

@@ -19,6 +19,8 @@
   ([#42655](https://github.com/Azure/azure-sdk-for-python/pull/42655))
 - Customer Facing SDKStats: Added telemetry_success field to dropped items as per [Spec] - https://github.com/aep-health-and-standards/Telemetry-Collection-Spec/pull/606
   ([#42846](https://github.com/Azure/azure-sdk-for-python/pull/42846))
+- Customer Facing SDKStats: Refactor to use `Manager` and `Singleton` pattern
+  ([#42969](https://github.com/Azure/azure-sdk-for-python/pull/42969))
 
 ### Breaking Changes
 

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/_constants.py

@@ -154,6 +154,7 @@
 _UNKNOWN = "UNKNOWN"
 
 # Customer Facing SDKStats
+
 _APPLICATIONINSIGHTS_SDKSTATS_ENABLED_PREVIEW = "APPLICATIONINSIGHTS_SDKSTATS_ENABLED_PREVIEW"
 _APPLICATIONINSIGHTS_SDKSTATS_EXPORT_INTERVAL = "APPLICATIONINSIGHTS_SDKSTATS_EXPORT_INTERVAL"
 _CUSTOMER_SDKSTATS_LANGUAGE = "python"
@@ -179,27 +180,18 @@ class CustomerSdkStatsMetricName(str, Enum, metaclass=CaseInsensitiveEnumMeta):
     ITEM_DROP_COUNT = "preview.item.dropped.count"
     ITEM_RETRY_COUNT = "preview.item.retry.count"
 
-class CustomerSdkStatsProperties:
-    language: str
-    version: str
-    compute_type: str
-    def __init__(self, language: str, version: str, compute_type: str):
-        self.language = language
-        self.version = version
-        self.compute_type = compute_type
-
 ## Map from Azure Monitor envelope base_types to TelemetryType
 _TYPE_MAP = {
-                "EventData": _CUSTOM_EVENT,
-                "MetricData": _CUSTOM_METRIC,
-                "RemoteDependencyData": _DEPENDENCY,
-                "ExceptionData": _EXCEPTION,
-                "PageViewData": _PAGE_VIEW,
-                "MessageData": _TRACE,
-                "RequestData": _REQUEST,
-                "PerformanceCounterData": _PERFORMANCE_COUNTER,
-                "AvailabilityData": _AVAILABILITY,
-            }
+    "EventData": _CUSTOM_EVENT,
+    "MetricData": _CUSTOM_METRIC,
+    "RemoteDependencyData": _DEPENDENCY,
+    "ExceptionData": _EXCEPTION,
+    "PageViewData": _PAGE_VIEW,
+    "MessageData": _TRACE,
+    "RequestData": _REQUEST,
+    "PerformanceCounterData": _PERFORMANCE_COUNTER,
+    "AvailabilityData": _AVAILABILITY,
+}
 
 # Exception categories
 class _exception_categories(Enum):

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/_generated/models/_models_py3.py

@@ -1326,7 +1326,7 @@ def __init__(  # type: ignore
         items_received: Optional[int] = None,
         items_accepted: Optional[int] = None,
         errors: Optional[List["TelemetryErrorDetails"]] = None,
-        **kwargs
+        **kwargs  # type: Any
     ):
         """
         :keyword items_received: The number of items received.

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/_storage.py

@@ -13,7 +13,7 @@
 
 from azure.monitor.opentelemetry.exporter._utils import PeriodicTask
 
-from azure.monitor.opentelemetry.exporter.statsbeat._state import (
+from azure.monitor.opentelemetry.exporter.statsbeat.customer._state import (
     get_local_storage_setup_state_exception,
     get_local_storage_setup_state_readonly,
     set_local_storage_setup_state_exception,
@@ -46,7 +46,7 @@ class StorageExportResult(Enum):
 # pylint: disable=broad-except
 class LocalFileBlob:
     def __init__(self, fullpath: str) -> None:
-        self.fullpath = fullpath
+        self.fullpath: str = fullpath
 
     def delete(self) -> None:
         try:
@@ -82,7 +82,7 @@ def put(self, data: List[Any], lease_period: int = 0) -> Union[StorageExportResu
 
     def lease(self, period: int) -> Optional['LocalFileBlob']:
         timestamp = _now() + _seconds(period)
-        fullpath = self.fullpath
+        fullpath: str = self.fullpath
         if fullpath.endswith(".lock"):
             fullpath = fullpath[: fullpath.rindex("@")]
         fullpath += "@{}.lock".format(_fmt(timestamp))
@@ -189,7 +189,7 @@ def gets(self) -> Generator[LocalFileBlob, None, None]:
         else:
             pass
 
-    def get(self) -> Optional['LocalFileBlob']:
+    def get(self) -> Optional[LocalFileBlob]:
         if not self._enabled:
             return None
         cursor = self.gets()

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/export/_base.py

@@ -52,18 +52,21 @@
 from azure.monitor.opentelemetry.exporter.statsbeat._state import (
     get_statsbeat_initial_success,
     get_statsbeat_shutdown,
-    get_customer_sdkstats_shutdown,
     increment_and_check_statsbeat_failure_count,
     is_statsbeat_enabled,
     set_statsbeat_initial_success,
-    is_customer_sdkstats_enabled,
 )
 from azure.monitor.opentelemetry.exporter.statsbeat._utils import (
     _update_requests_map,
-    _track_dropped_items_from_storage,
-    _track_dropped_items,
-    _track_retry_items,
-    _track_successful_items,
+)
+from azure.monitor.opentelemetry.exporter.statsbeat.customer._utils import (
+    track_dropped_items_from_storage,
+    track_dropped_items,
+    track_retry_items,
+    track_successful_items,
+)
+from azure.monitor.opentelemetry.exporter.statsbeat.customer._state import (
+    get_customer_stats_manager,
 )
 
 
@@ -111,6 +114,7 @@ def __init__(self, **kwargs: Any) -> None:
             self._credential = _get_authentication_credential(**kwargs)
         self._consecutive_redirects = 0  # To prevent circular redirects
         self._disable_offline_storage = kwargs.get("disable_offline_storage", False)
+        self._connection_string = parsed_connection_string._connection_string
         self._endpoint = parsed_connection_string.endpoint
         self._region = parsed_connection_string.region
         self._instrumentation_key = parsed_connection_string.instrumentation_key
@@ -159,10 +163,10 @@ def __init__(self, **kwargs: Any) -> None:
             config.http_logging_policy or HttpLoggingPolicy(**kwargs),
         ]
 
-        self.client = AzureMonitorClient(
+        self.client: AzureMonitorClient = AzureMonitorClient(
             host=self._endpoint, connection_timeout=self._timeout, policies=policies, **kwargs
         )
-        self.storage = None
+        self.storage: Optional[LocalFileStorage] = None
         if not self._disable_offline_storage:
             self.storage = LocalFileStorage(  # pyright: ignore
                 path=self._storage_directory,  # type: ignore
@@ -187,7 +191,7 @@ def __init__(self, **kwargs: Any) -> None:
         # customer sdkstats initialization
         if self._should_collect_customer_sdkstats():
 
-            from azure.monitor.opentelemetry.exporter.statsbeat._customer_sdkstats import collect_customer_sdkstats
+            from azure.monitor.opentelemetry.exporter.statsbeat.customer import collect_customer_sdkstats
             # Collect customer sdkstats metrics
             collect_customer_sdkstats(self)
 
@@ -212,15 +216,15 @@ def _handle_transmit_from_storage(self, envelopes: List[TelemetryItem], result:
                 envelopes_to_store = [x.as_dict() for x in envelopes]
                 result_from_storage_put = self.storage.put(envelopes_to_store)
                 if self._should_collect_customer_sdkstats():
-                    _track_dropped_items_from_storage(result_from_storage_put, envelopes)
+                    track_dropped_items_from_storage(result_from_storage_put, envelopes)
             elif result == ExportResult.SUCCESS:
                 # Try to send any cached events
                 self._transmit_from_storage()
 
         else:
             # Track items that would have been retried but are dropped since client has local storage disabled
             if self._should_collect_customer_sdkstats():
-                _track_dropped_items(envelopes, DropCode.CLIENT_STORAGE_DISABLED)
+                track_dropped_items(envelopes, DropCode.CLIENT_STORAGE_DISABLED)
 
     # pylint: disable=too-many-branches
     # pylint: disable=too-many-nested-blocks
@@ -259,7 +263,7 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
 
                     # Track successful items in customer sdkstats
                     if self._should_collect_customer_sdkstats():
-                        _track_successful_items(envelopes)
+                        track_successful_items(envelopes)
                 else:  # 206
                     reach_ingestion = True
                     resend_envelopes = []
@@ -268,13 +272,13 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                             resend_envelopes.append(envelopes[error.index])  # type: ignore
                             # Track retried items in customer sdkstats
                             if self._should_collect_customer_sdkstats():
-                                _track_retry_items(resend_envelopes, error)
+                                track_retry_items(resend_envelopes, error)
                         else:
                             if not self._is_stats_exporter():
                                 # Track dropped items in customer sdkstats, non-retryable scenario
                                 if self._should_collect_customer_sdkstats():
                                     if error is not None and hasattr(error, "index") and error.index is not None and isinstance(error.status_code, int):
-                                        _track_dropped_items([envelopes[error.index]], error.status_code)
+                                        track_dropped_items([envelopes[error.index]], error.status_code)
                                 logger.error(
                                     "Data drop %s: %s %s.",
                                     error.status_code,
@@ -285,12 +289,12 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                         envelopes_to_store = [x.as_dict() for x in resend_envelopes]
                         result_from_storage = self.storage.put(envelopes_to_store, 0)
                         if self._should_collect_customer_sdkstats():
-                            _track_dropped_items_from_storage(result_from_storage, resend_envelopes)
+                            track_dropped_items_from_storage(result_from_storage, resend_envelopes)
                         self._consecutive_redirects = 0
                     elif resend_envelopes:
                         # Track items that would have been retried but are dropped since client has local storage disabled
                         if self._should_collect_customer_sdkstats():
-                            _track_dropped_items(resend_envelopes, DropCode.CLIENT_STORAGE_DISABLED)
+                            track_dropped_items(resend_envelopes, DropCode.CLIENT_STORAGE_DISABLED)
                     # Mark as not retryable because we already write to storage here
                     result = ExportResult.FAILED_NOT_RETRYABLE
             except HttpResponseError as response_error:
@@ -304,7 +308,7 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                     # Log error for 401: Unauthorized, 403: Forbidden to assist with customer troubleshooting
                     if not self._is_stats_exporter():
                         if self._should_collect_customer_sdkstats():
-                            _track_retry_items(envelopes, response_error)
+                            track_retry_items(envelopes, response_error)
                         if response_error.status_code == 401:
                             logger.error(
                                 "Retryable server side error: %s. " \
@@ -313,8 +317,6 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                                 response_error.message,
                             )
                         elif response_error.status_code == 403:
-                            if self._should_collect_customer_sdkstats():
-                                _track_retry_items(envelopes, response_error)
                             logger.error(
                                 "Retryable server side error: %s. " \
                                 "Your application may be configured with a token credential " \
@@ -330,7 +332,7 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
 
                     if not self._is_stats_exporter():
                         if self._should_collect_customer_sdkstats() and isinstance(response_error.status_code, int):
-                            _track_dropped_items(envelopes, response_error.status_code)
+                            track_dropped_items(envelopes, response_error.status_code)
                 elif _is_redirect_code(response_error.status_code):
                     self._consecutive_redirects = self._consecutive_redirects + 1
                     # pylint: disable=W0212
@@ -349,7 +351,7 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                         else:
                             if not self._is_stats_exporter():
                                 if self._should_collect_customer_sdkstats():
-                                    _track_dropped_items(envelopes, DropCode.CLIENT_EXCEPTION, _exception_categories.CLIENT_EXCEPTION.value)
+                                    track_dropped_items(envelopes, DropCode.CLIENT_EXCEPTION, _exception_categories.CLIENT_EXCEPTION.value)
                                 logger.error(
                                     "Error parsing redirect information.",
                                 )
@@ -358,7 +360,7 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                         if not self._is_stats_exporter():
                             # Track dropped items in customer sdkstats, non-retryable scenario
                             if self._should_collect_customer_sdkstats():
-                                _track_dropped_items(
+                                track_dropped_items(
                                     envelopes,
                                     DropCode.CLIENT_EXCEPTION,
                                     _exception_categories.CLIENT_EXCEPTION.value
@@ -384,14 +386,14 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                         )
                         # Track dropped items in customer sdkstats, non-retryable scenario
                         if self._should_collect_customer_sdkstats() and isinstance(response_error.status_code, int):
-                            _track_dropped_items(envelopes, response_error.status_code)
+                            track_dropped_items(envelopes, response_error.status_code)
                         if _is_invalid_code(response_error.status_code):
                             # Shutdown statsbeat on invalid code from customer endpoint
                             # Import here to avoid circular dependencies
                             from azure.monitor.opentelemetry.exporter.statsbeat._statsbeat import (
                                 shutdown_statsbeat_metrics,
                             )
-                            from azure.monitor.opentelemetry.exporter.statsbeat._customer_sdkstats import (
+                            from azure.monitor.opentelemetry.exporter.statsbeat.customer import (
                                 shutdown_customer_sdkstats_metrics,
                             )
 
@@ -407,7 +409,7 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
 
                 # Track retry items in customer sdkstats for client-side exceptions
                 if self._should_collect_customer_sdkstats():
-                    _track_retry_items(envelopes, request_error)
+                    track_retry_items(envelopes, request_error)
 
                 if self._should_collect_stats():
                     exc_type = request_error.exc_type
@@ -420,7 +422,7 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
 
                 # Track dropped items in customer sdkstats for general exceptions
                 if self._should_collect_customer_sdkstats():
-                    _track_dropped_items(envelopes, DropCode.CLIENT_EXCEPTION, _exception_categories.CLIENT_EXCEPTION.value)
+                    track_dropped_items(envelopes, DropCode.CLIENT_EXCEPTION, _exception_categories.CLIENT_EXCEPTION.value)
 
                 if self._should_collect_stats():
                     _update_requests_map(_REQ_EXCEPTION_NAME[1], value=ex.__class__.__name__)
@@ -467,10 +469,10 @@ def _should_collect_stats(self):
 
     # check to see whether its the case of customer sdkstats collection
     def _should_collect_customer_sdkstats(self):
-        # Don't collect customer sdkstats for instrumentation collection, sdkstats exporter or customer sdkstats exporter
+        manager = get_customer_stats_manager()
         return (
-            is_customer_sdkstats_enabled()
-            and not get_customer_sdkstats_shutdown()
+            manager.is_enabled
+            and not manager.is_shutdown
             and not self._is_stats_exporter()
             and not self._is_customer_sdkstats_exporter()
             and not self._instrumentation_collection