Refactored logic for tracking dropped items from storage (#42542)

* Refactored logic for tracking dropped items from storage

* Updated CHANGELOG

Author: rads-1996

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

@@ -7,6 +7,8 @@
   ([#42382](https://github.com/Azure/azure-sdk-for-python/pull/42382))
 - Refactored the put methods in storage.py for LocalFileBlob and LocalFileStorage
   ([#42502](https://github.com/Azure/azure-sdk-for-python/pull/42502))
+- Customer Facing Statsbeat: Refactored logic for tracking dropped items from storage
+  ([#42542](https://github.com/Azure/azure-sdk-for-python/pull/42542))
 
 ### Breaking Changes
 

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

@@ -46,24 +46,21 @@
 )
 # from azure.monitor.opentelemetry.exporter._configuration import _ConfigurationManager
 from azure.monitor.opentelemetry.exporter._connection_string_parser import ConnectionStringParser
-from azure.monitor.opentelemetry.exporter._storage import (
-    LocalFileStorage,
-    StorageExportResult,
-)
+from azure.monitor.opentelemetry.exporter._storage import LocalFileStorage
 from azure.monitor.opentelemetry.exporter._utils import _get_auth_policy
 from azure.monitor.opentelemetry.exporter.statsbeat._state import (
     get_statsbeat_initial_success,
     get_statsbeat_shutdown,
     increment_and_check_statsbeat_failure_count,
     is_statsbeat_enabled,
     set_statsbeat_initial_success,
-    get_local_storage_setup_state_exception,
 )
 from azure.monitor.opentelemetry.exporter.statsbeat._utils import (
     _update_requests_map,
     _track_dropped_items,
     _track_retry_items,
     _track_successful_items,
+    _track_dropped_items_from_storage,
 )
 
 
@@ -207,24 +204,7 @@ 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._customer_statsbeat_metrics and self._should_collect_customer_statsbeat():
-                    if result_from_storage_put == StorageExportResult.CLIENT_STORAGE_DISABLED:
-                        # Track items that would have been retried but are dropped since client has local storage disabled
-                        _track_dropped_items(self._customer_statsbeat_metrics, envelopes, DropCode.CLIENT_STORAGE_DISABLED)
-                    elif result_from_storage_put == StorageExportResult.CLIENT_READONLY:
-                        # If filesystem is readonly, track dropped items in customer statsbeat
-                        _track_dropped_items(self._customer_statsbeat_metrics, envelopes, DropCode.CLIENT_READONLY)
-                    elif result_from_storage_put == StorageExportResult.CLIENT_PERSISTENCE_CAPACITY_REACHED:
-                        # If data has to be dropped due to persistent storage being full, track dropped items
-                        _track_dropped_items(self._customer_statsbeat_metrics, envelopes, DropCode.CLIENT_PERSISTENCE_CAPACITY)
-                    elif get_local_storage_setup_state_exception() != "":
-                        # For exceptions caught in _check_and_set_folder_permissions during storage setup
-                        _track_dropped_items(self._customer_statsbeat_metrics, envelopes, DropCode.CLIENT_EXCEPTION, result_from_storage_put)
-                    elif isinstance(result_from_storage_put, str):
-                        # For any exceptions occurred in put method of either LocalFileStorage or LocalFileBlob, track dropped item with reason
-                        _track_dropped_items(self._customer_statsbeat_metrics, envelopes, DropCode.CLIENT_EXCEPTION, result_from_storage_put)
-                    else:
-                        # LocalFileBlob.put returns StorageExportResult.LOCAL_FILE_BLOB_SUCCESS here. Don't need to track anything in this case.
-                        pass
+                    _track_dropped_items_from_storage(self._customer_statsbeat_metrics, result_from_storage_put, envelopes)
             elif result == ExportResult.SUCCESS:
                 # Try to send any cached events
                 self._transmit_from_storage()
@@ -295,7 +275,9 @@ def _transmit(self, envelopes: List[TelemetryItem]) -> ExportResult:
                                 )
                     if self.storage and resend_envelopes:
                         envelopes_to_store = [x.as_dict() for x in resend_envelopes]
-                        self.storage.put(envelopes_to_store, 0)
+                        result_from_storage = self.storage.put(envelopes_to_store, 0)
+                        if self._customer_statsbeat_metrics and self._should_collect_customer_statsbeat():
+                            _track_dropped_items_from_storage(self._customer_statsbeat_metrics, 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

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

@@ -7,10 +7,13 @@
     RetryCode,
     RetryCodeType,
     DropCodeType,
+    DropCode,
     _UNKNOWN,
 )
 from azure.monitor.opentelemetry.exporter._utils import _get_telemetry_type
 from azure.monitor.opentelemetry.exporter._generated.models import TelemetryItem
+from azure.monitor.opentelemetry.exporter._storage import StorageExportResult
+from azure.monitor.opentelemetry.exporter.statsbeat._state import get_local_storage_setup_state_exception
 
 
 from azure.monitor.opentelemetry.exporter._constants import (
@@ -178,3 +181,24 @@ def _track_retry_items(customer_statsbeat_metrics, envelopes: List[TelemetryItem
                     retry_code,
                     str(message)
                 )
+
+def _track_dropped_items_from_storage(customer_statsbeat_metrics, result_from_storage_put, envelopes):
+    if customer_statsbeat_metrics:
+        if result_from_storage_put == StorageExportResult.CLIENT_STORAGE_DISABLED:
+            # Track items that would have been retried but are dropped since client has local storage disabled
+            _track_dropped_items(customer_statsbeat_metrics, envelopes, DropCode.CLIENT_STORAGE_DISABLED)
+        elif result_from_storage_put == StorageExportResult.CLIENT_READONLY:
+            # If filesystem is readonly, track dropped items in customer statsbeat
+            _track_dropped_items(customer_statsbeat_metrics, envelopes, DropCode.CLIENT_READONLY)
+        elif result_from_storage_put == StorageExportResult.CLIENT_PERSISTENCE_CAPACITY_REACHED:
+            # If data has to be dropped due to persistent storage being full, track dropped items
+            _track_dropped_items(customer_statsbeat_metrics, envelopes, DropCode.CLIENT_PERSISTENCE_CAPACITY)
+        elif get_local_storage_setup_state_exception() != "":
+            # For exceptions caught in _check_and_set_folder_permissions during storage setup
+            _track_dropped_items(customer_statsbeat_metrics, envelopes, DropCode.CLIENT_EXCEPTION, result_from_storage_put) # pylint: disable=line-too-long
+        elif isinstance(result_from_storage_put, str):
+            # For any exceptions occurred in put method of either LocalFileStorage or LocalFileBlob, track dropped item with reason # pylint: disable=line-too-long
+            _track_dropped_items(customer_statsbeat_metrics, envelopes, DropCode.CLIENT_EXCEPTION, result_from_storage_put) # pylint: disable=line-too-long
+        else:
+            # LocalFileBlob.put returns StorageExportResult.LOCAL_FILE_BLOB_SUCCESS here. Don't need to track anything in this case. # pylint: disable=line-too-long
+            pass