add interruptible shutdown

Author: liustve

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/logs/otlp_aws_logs_exporter.py

@@ -6,7 +6,8 @@
 import logging
 import random
 from io import BytesIO
-from time import sleep, time
+from threading import Event
+from time import time
 from typing import Dict, Optional, Sequence
 
 from requests import Response
@@ -60,6 +61,7 @@ def __init__(
             compression=Compression.Gzip,
             session=AwsAuthSession(aws_region=self._aws_region, service="logs"),
         )
+        self._shutdown_event = Event()
 
     def export(self, batch: Sequence[LogData]) -> LogExportResult:
         """
@@ -114,9 +116,18 @@ def export(self, batch: Sequence[LogData]) -> LogExportResult:
                 resp.reason,
                 backoff_seconds,
             )
-            sleep(backoff_seconds)
+            # Use interruptible sleep that can be interrupted by shutdown
+            if self._shutdown_event.wait(backoff_seconds):
+                _logger.info("Export interrupted by shutdown")
+                return LogExportResult.FAILURE
+            
             retry_num += 1
 
+    def shutdown(self) -> None:
+        """Shutdown the exporter and interrupt any ongoing waits."""
+        self._shutdown_event.set()
+        return super().shutdown()
+
     def _send(self, serialized_data: bytes, timeout_sec: float):
         try:
             response = self._session.post(

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/exporter/otlp/aws/logs/test_otlp_aws_logs_exporter.py

@@ -73,19 +73,17 @@ def test_should_not_export_again_if_not_retryable(self, mock_request):
 
         self.assertEqual(result, LogExportResult.FAILURE)
 
-    @patch(
-        "amazon.opentelemetry.distro.exporter.otlp.aws.logs.otlp_aws_logs_exporter.sleep", side_effect=lambda x: None
-    )
+    @patch("threading.Event.wait", side_effect=lambda x: False)
     @patch("requests.Session.post", return_value=retryable_response_no_header)
-    def test_should_export_again_with_backoff_if_retryable_and_no_retry_after_header(self, mock_request, mock_sleep):
+    def test_should_export_again_with_backoff_if_retryable_and_no_retry_after_header(self, mock_request, mock_wait):
         """Tests that multiple export requests are made with exponential delay if the response status code is retryable.
         But there is no Retry-After header."""
         self.exporter._timeout = 10000  # Large timeout to avoid early exit
         result = self.exporter.export(self.logs)
 
-        self.assertEqual(mock_sleep.call_count, _MAX_RETRYS - 1)
+        self.assertEqual(mock_wait.call_count, _MAX_RETRYS - 1)
 
-        delays = mock_sleep.call_args_list
+        delays = mock_wait.call_args_list
 
         for index, delay in enumerate(delays):
             expected_base = 2**index
@@ -97,30 +95,26 @@ def test_should_export_again_with_backoff_if_retryable_and_no_retry_after_header
         self.assertEqual(mock_request.call_count, _MAX_RETRYS)
         self.assertEqual(result, LogExportResult.FAILURE)
 
-    @patch(
-        "amazon.opentelemetry.distro.exporter.otlp.aws.logs.otlp_aws_logs_exporter.sleep", side_effect=lambda x: None
-    )
+    @patch("threading.Event.wait", side_effect=lambda x: False)
     @patch(
         "requests.Session.post",
         side_effect=[retryable_response_header, retryable_response_header, retryable_response_header, good_response],
     )
-    def test_should_export_again_with_server_delay_if_retryable_and_retry_after_header(self, mock_request, mock_sleep):
+    def test_should_export_again_with_server_delay_if_retryable_and_retry_after_header(self, mock_request, mock_wait):
         """Tests that multiple export requests are made with the server's suggested
         delay if the response status code is retryable and there is a Retry-After header."""
         self.exporter._timeout = 10000  # Large timeout to avoid early exit
         result = self.exporter.export(self.logs)
-        delays = mock_sleep.call_args_list
+        delays = mock_wait.call_args_list
 
         for delay in delays:
             self.assertEqual(delay[0][0], 10)
 
-        self.assertEqual(mock_sleep.call_count, 3)
+        self.assertEqual(mock_wait.call_count, 3)
         self.assertEqual(mock_request.call_count, 4)
         self.assertEqual(result, LogExportResult.SUCCESS)
 
-    @patch(
-        "amazon.opentelemetry.distro.exporter.otlp.aws.logs.otlp_aws_logs_exporter.sleep", side_effect=lambda x: None
-    )
+    @patch("threading.Event.wait", side_effect=lambda x: False)
     @patch(
         "requests.Session.post",
         side_effect=[
@@ -131,13 +125,13 @@ def test_should_export_again_with_server_delay_if_retryable_and_retry_after_head
         ],
     )
     def test_should_export_again_with_backoff_delay_if_retryable_and_bad_retry_after_header(
-        self, mock_request, mock_sleep
+        self, mock_request, mock_wait
     ):
         """Tests that multiple export requests are made with exponential delay if the response status code is retryable.
         but the Retry-After header is invalid or malformed."""
         self.exporter._timeout = 10000  # Large timeout to avoid early exit
         result = self.exporter.export(self.logs)
-        delays = mock_sleep.call_args_list
+        delays = mock_wait.call_args_list
 
         for index, delay in enumerate(delays):
             expected_base = 2**index
@@ -146,7 +140,7 @@ def test_should_export_again_with_backoff_delay_if_retryable_and_bad_retry_after
             self.assertGreaterEqual(actual_delay, expected_base * 0.8)
             self.assertLessEqual(actual_delay, expected_base * 1.2)
 
-        self.assertEqual(mock_sleep.call_count, 3)
+        self.assertEqual(mock_wait.call_count, 3)
         self.assertEqual(mock_request.call_count, 4)
         self.assertEqual(result, LogExportResult.SUCCESS)
 
@@ -158,29 +152,44 @@ def test_export_connection_error_retry(self, mock_request):
         self.assertEqual(mock_request.call_count, 2)
         self.assertEqual(result, LogExportResult.SUCCESS)
 
-    @patch(
-        "amazon.opentelemetry.distro.exporter.otlp.aws.logs.otlp_aws_logs_exporter.sleep", side_effect=lambda x: None
-    )
+    @patch("threading.Event.wait", side_effect=lambda x: False)
     @patch("requests.Session.post", return_value=retryable_response_no_header)
-    def test_should_stop_retrying_when_deadline_exceeded(self, mock_request, mock_sleep):
+    def test_should_stop_retrying_when_deadline_exceeded(self, mock_request, mock_wait):
         """Tests that the exporter stops retrying when the deadline is exceeded."""
         self.exporter._timeout = 5  # Short timeout to trigger deadline check
 
-        # Mock time to simulate time passing
         with patch("amazon.opentelemetry.distro.exporter.otlp.aws.logs.otlp_aws_logs_exporter.time") as mock_time:
             # First call returns start time, subsequent calls simulate time passing
             mock_time.side_effect = [0, 0, 1, 2, 4, 8]  # Exponential backoff would be 1, 2, 4 seconds
 
             result = self.exporter.export(self.logs)
 
             # Should stop before max retries due to deadline
-            self.assertLess(mock_sleep.call_count, _MAX_RETRYS)
+            self.assertLess(mock_wait.call_count, _MAX_RETRYS)
             self.assertLess(mock_request.call_count, _MAX_RETRYS + 1)
             self.assertEqual(result, LogExportResult.FAILURE)
 
             # Verify total time passed is at the timeout limit
             self.assertGreaterEqual(5, self.exporter._timeout)
 
+    @patch("requests.Session.post", return_value=retryable_response_no_header)
+    def test_export_interrupted_by_shutdown(self, mock_request):
+        """Tests that export can be interrupted by shutdown during retry wait."""
+        self.exporter._timeout = 10000
+        
+        # Mock Event.wait to call shutdown on first call, then return True (interrupted)
+        # We cannot call shutdown() at the beginning since the exporter would just automatically return a FAILURE result without even attempting the export.
+        def mock_wait_with_shutdown(timeout):
+            self.exporter.shutdown()  
+            return True
+            
+        with patch.object(self.exporter._shutdown_event, 'wait', side_effect=mock_wait_with_shutdown):
+            result = self.exporter.export(self.logs)
+            
+            # Should make one request, then get interrupted during retry wait
+            self.assertEqual(mock_request.call_count, 1)
+            self.assertEqual(result, LogExportResult.FAILURE)
+
     @staticmethod
     def generate_test_log_data(count=5):
         logs = []