Added logic for race conditions and global instance for customer SDKStats (#42655)

* Added logic for race conditions and updated the implementation to use a global instance for customer SDKStats

* Updated CHANGELOG

* Minor adjustments

* Refactored code after statsbeat PR

* Fix test

* Clean up test file

Author: rads-1996

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

@@ -15,7 +15,8 @@
   ([#42573](https://github.com/Azure/azure-sdk-for-python/pull/42573))
 - Customer Facing SDKStats: Exception categorization as per [Spec] - https://github.com/aep-health-and-standards/Telemetry-Collection-Spec/blob/main/ApplicationInsights/sdkstats/customer_facing_sdk_stats.md
   ([#42695](https://github.com/Azure/azure-sdk-for-python/pull/42695))
-
+- Customer Facing SDKStats: Added logic for race conditions and updated the implementation to use a global instance for customer SDKStats metrics
+  ([#42655](https://github.com/Azure/azure-sdk-for-python/pull/42655))
 ### Breaking Changes
 
 ### Bugs Fixed

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

@@ -182,17 +182,17 @@ def __init__(self, language: str, version: str, compute_type: str):
         self.version = version
         self.compute_type = compute_type
 
-## Map from Azure Monitor envelope names to TelemetryType
+## Map from Azure Monitor envelope base_types to TelemetryType
 _TYPE_MAP = {
-                _EVENT_ENVELOPE_NAME: _CUSTOM_EVENT,
-                _METRIC_ENVELOPE_NAME: _CUSTOM_METRIC,
-                _REMOTE_DEPENDENCY_ENVELOPE_NAME: _DEPENDENCY,
-                _EXCEPTION_ENVELOPE_NAME: _EXCEPTION,
-                _PAGE_VIEW_ENVELOPE_NAME: _PAGE_VIEW,
-                _MESSAGE_ENVELOPE_NAME: _TRACE,
-                _REQUEST_ENVELOPE_NAME: _REQUEST,
-                _PERFORMANCE_COUNTER_ENVELOPE_NAME: _PERFORMANCE_COUNTER,
-                _AVAILABILITY_ENVELOPE_NAME: _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

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

@@ -56,13 +56,14 @@
     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,
-    _track_dropped_items_from_storage,
 )
 
 
@@ -181,8 +182,7 @@ def __init__(self, **kwargs: Any) -> None:
             except Exception as e:  # pylint: disable=broad-except
                 logger.warning("Failed to initialize statsbeat metrics: %s", e)
 
-        # Initialize customer sdkstats if enabled
-        self._customer_sdkstats_metrics = None
+        # customer sdkstats initialization
         if self._should_collect_customer_sdkstats():
 
             from azure.monitor.opentelemetry.exporter.statsbeat._customer_sdkstats import collect_customer_sdkstats
@@ -209,16 +209,16 @@ def _handle_transmit_from_storage(self, envelopes: List[TelemetryItem], result:
             if result == ExportResult.FAILED_RETRYABLE:
                 envelopes_to_store = [x.as_dict() for x in envelopes]
                 result_from_storage_put = self.storage.put(envelopes_to_store)
-                if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                    _track_dropped_items_from_storage(self._customer_sdkstats_metrics, result_from_storage_put, envelopes)
+                if self._should_collect_customer_sdkstats():
+                    _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._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                _track_dropped_items(self._customer_sdkstats_metrics, envelopes, DropCode.CLIENT_STORAGE_DISABLED)
+            if self._should_collect_customer_sdkstats():
+                _track_dropped_items(envelopes, DropCode.CLIENT_STORAGE_DISABLED)
 
     # pylint: disable=too-many-branches
     # pylint: disable=too-many-nested-blocks
@@ -256,23 +256,23 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                     result = ExportResult.SUCCESS
 
                     # Track successful items in customer sdkstats
-                    if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                        _track_successful_items(self._customer_sdkstats_metrics, envelopes)
+                    if self._should_collect_customer_sdkstats():
+                        _track_successful_items(envelopes)
                 else:  # 206
                     reach_ingestion = True
                     resend_envelopes = []
                     for error in track_response.errors:
                         if _is_retryable_code(error.status_code):
                             resend_envelopes.append(envelopes[error.index])  # type: ignore
                             # Track retried items in customer sdkstats
-                            if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                                _track_retry_items(self._customer_sdkstats_metrics, resend_envelopes, error)
+                            if self._should_collect_customer_sdkstats():
+                                _track_retry_items(resend_envelopes, error)
                         else:
                             if not self._is_stats_exporter():
                                 # Track dropped items in customer sdkstats, non-retryable scenario
-                                if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                                    if error is not None and hasattr(error, "index") and error.index is not None:
-                                        _track_dropped_items(self._customer_sdkstats_metrics, [envelopes[error.index]], error.status_code)
+                                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)
                                 logger.error(
                                     "Data drop %s: %s %s.",
                                     error.status_code,
@@ -282,13 +282,13 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                     if self.storage and resend_envelopes:
                         envelopes_to_store = [x.as_dict() for x in resend_envelopes]
                         result_from_storage = self.storage.put(envelopes_to_store, 0)
-                        if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                            _track_dropped_items_from_storage(self._customer_sdkstats_metrics, result_from_storage, resend_envelopes)
+                        if self._should_collect_customer_sdkstats():
+                            _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._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                            _track_dropped_items(self._customer_sdkstats_metrics, resend_envelopes, DropCode.CLIENT_STORAGE_DISABLED)
+                        if self._should_collect_customer_sdkstats():
+                            _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:
@@ -301,8 +301,8 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                     result = ExportResult.FAILED_RETRYABLE
                     # Log error for 401: Unauthorized, 403: Forbidden to assist with customer troubleshooting
                     if not self._is_stats_exporter():
-                        if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                            _track_retry_items(self._customer_sdkstats_metrics, envelopes, response_error)
+                        if self._should_collect_customer_sdkstats():
+                            _track_retry_items(envelopes, response_error)
                         if response_error.status_code == 401:
                             logger.error(
                                 "Retryable server side error: %s. " \
@@ -311,8 +311,8 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                                 response_error.message,
                             )
                         elif response_error.status_code == 403:
-                            if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                                _track_retry_items(self._customer_sdkstats_metrics, envelopes, response_error)
+                            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 " \
@@ -327,8 +327,8 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                     result = ExportResult.FAILED_NOT_RETRYABLE
 
                     if not self._is_stats_exporter():
-                        if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                            _track_dropped_items(self._customer_sdkstats_metrics, envelopes, response_error.status_code)
+                        if self._should_collect_customer_sdkstats() and isinstance(response_error.status_code, int):
+                            _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
@@ -346,18 +346,17 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                             result = self._transmit(envelopes)
                         else:
                             if not self._is_stats_exporter():
-                                if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                                    _track_dropped_items(self._customer_sdkstats_metrics, envelopes, DropCode.CLIENT_EXCEPTION, _exception_categories.CLIENT_EXCEPTION.value)
+                                if self._should_collect_customer_sdkstats():
+                                    _track_dropped_items(envelopes, DropCode.CLIENT_EXCEPTION, _exception_categories.CLIENT_EXCEPTION.value)
                                 logger.error(
                                     "Error parsing redirect information.",
                                 )
                             result = ExportResult.FAILED_NOT_RETRYABLE
                     else:
                         if not self._is_stats_exporter():
                             # Track dropped items in customer sdkstats, non-retryable scenario
-                            if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
+                            if self._should_collect_customer_sdkstats():
                                 _track_dropped_items(
-                                    self._customer_sdkstats_metrics,
                                     envelopes,
                                     DropCode.CLIENT_EXCEPTION,
                                     _exception_categories.CLIENT_EXCEPTION.value
@@ -382,8 +381,8 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                             response_error.message,
                         )
                         # Track dropped items in customer sdkstats, non-retryable scenario
-                        if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                            _track_dropped_items(self._customer_sdkstats_metrics, envelopes, response_error.status_code)
+                        if self._should_collect_customer_sdkstats() and isinstance(response_error.status_code, int):
+                            _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
@@ -405,8 +404,8 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                 logger.warning("Retrying due to server request error: %s.", request_error.message)
 
                 # Track retry items in customer sdkstats for client-side exceptions
-                if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                    _track_retry_items(self._customer_sdkstats_metrics, envelopes, request_error)
+                if self._should_collect_customer_sdkstats():
+                    _track_retry_items(envelopes, request_error)
 
                 if self._should_collect_stats():
                     exc_type = request_error.exc_type
@@ -418,8 +417,8 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                 logger.exception("Envelopes could not be exported and are not retryable: %s.", ex)  # pylint: disable=C4769
 
                 # Track dropped items in customer sdkstats for general exceptions
-                if self._customer_sdkstats_metrics and self._should_collect_customer_sdkstats():
-                    _track_dropped_items(self._customer_sdkstats_metrics, envelopes, DropCode.CLIENT_EXCEPTION, _exception_categories.CLIENT_EXCEPTION.value)
+                if self._should_collect_customer_sdkstats():
+                    _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__)
@@ -466,11 +465,9 @@ def _should_collect_stats(self):
 
     # check to see whether its the case of customer sdkstats collection
     def _should_collect_customer_sdkstats(self):
-        env_value = os.environ.get("APPLICATIONINSIGHTS_SDKSTATS_ENABLED_PREVIEW", "")
-        is_customer_sdkstats_enabled = env_value.lower() == "true"
         # Don't collect customer sdkstats for instrumentation collection, sdkstats exporter or customer sdkstats exporter
         return (
-            is_customer_sdkstats_enabled
+            is_customer_sdkstats_enabled()
             and not get_customer_sdkstats_shutdown()
             and not self._is_stats_exporter()
             and not self._is_customer_sdkstats_exporter()