Initital commit to imporve shutdown behavior.

Author: DylanRussell

opentelemetry-sdk/src/opentelemetry/sdk/_shared_internal/__init__.py

@@ -16,9 +16,11 @@
 
 import collections
 import enum
+import inspect
 import logging
 import os
 import threading
+import time
 import weakref
 from abc import abstractmethod
 from typing import (
@@ -130,17 +132,27 @@ def worker(self):
             # https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/sdk.md#batching-processor.
             # Shutdown will interrupt this sleep. Emit will interrupt this sleep only if the queue is bigger then threshold.
             sleep_interrupted = self._worker_awaken.wait(self._schedule_delay)
+            print(
+                "In worker loop:{}, {}, {}".format(
+                    sleep_interrupted,
+                    self._schedule_delay,
+                    self._schedule_delay_millis,
+                )
+            )
             if self._shutdown:
+                print("Shutdown is set...")
                 break
             self._export(
                 BatchExportStrategy.EXPORT_WHILE_BATCH_EXCEEDS_THRESHOLD
                 if sleep_interrupted
                 else BatchExportStrategy.EXPORT_AT_LEAST_ONE_BATCH
             )
             self._worker_awaken.clear()
+        print("last export bach...")
         self._export(BatchExportStrategy.EXPORT_ALL)
 
     def _export(self, batch_strategy: BatchExportStrategy) -> None:
+        print("export started...:{}".format(batch_strategy))
         with self._export_lock:
             iteration = 0
             # We could see concurrent export calls from worker and force_flush. We call _should_export_batch
@@ -149,6 +161,7 @@ def _export(self, batch_strategy: BatchExportStrategy) -> None:
                 iteration += 1
                 token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True))
                 try:
+                    print("SIZE: {}".format(len(self._queue)))
                     self._exporter.export(
                         [
                             # Oldest records are at the back, so pop from there.
@@ -161,6 +174,7 @@ def _export(self, batch_strategy: BatchExportStrategy) -> None:
                             )
                         ]
                     )
+                    print("export succeded??")
                 except Exception:  # pylint: disable=broad-exception-caught
                     self._logger.exception(
                         "Exception while exporting %s.", self._exporting
@@ -180,16 +194,32 @@ def emit(self, data: Telemetry) -> None:
         if len(self._queue) >= self._max_export_batch_size:
             self._worker_awaken.set()
 
-    def shutdown(self):
+    # LoggerProvider calls shutdown without arguments currently, so the default is used.
+    def shutdown(self, timeout_millis=30000):
         if self._shutdown:
             return
         # Prevents emit and force_flush from further calling export.
         self._shutdown = True
-        # Interrupts sleep in the worker, if it's sleeping.
+        # Interrupts sleep in the worker if it's sleeping.
         self._worker_awaken.set()
-        # Main worker loop should exit after one final export call with flush all strategy.
+        # Wait a tiny bit for the worker thread to wake and call export for a final time.
+        time.sleep(0.1)
+        # We will force shutdown after 30 seconds.
+        for _ in range(10):
+            # If export is not being called, we can shutdown.
+            if not self._export_lock.locked():
+                break
+            time.sleep(timeout_millis / 1000 / 10)
+        # We want to shutdown immediately because we already waited 30 seconds. Some exporter's shutdown support a timeout param.
+        if (
+            "timeout_millis"
+            in inspect.getfullargspec(self._exporter.shutdown).args
+        ):
+            self._exporter.shutdown(timeout_millis=0)  # type: ignore
+        else:
+            self._exporter.shutdown()
+        # Worker thread should be finished at this point and return instantly.
         self._worker_thread.join()
-        self._exporter.shutdown()
 
     # TODO: Fix force flush so the timeout is used https://github.com/open-telemetry/opentelemetry-python/issues/4568.
     def force_flush(self, timeout_millis: Optional[int] = None) -> bool:

opentelemetry-sdk/test-requirements.txt

@@ -6,6 +6,7 @@ packaging==24.0
 pluggy==1.5.0
 psutil==5.9.6; sys_platform != 'win32'
 py-cpuinfo==9.0.0
+requests==2.32.3
 pytest==7.4.4
 tomli==2.0.1
 typing_extensions==4.10.0
@@ -15,3 +16,6 @@ zipp==3.19.2
 -e opentelemetry-api
 -e opentelemetry-semantic-conventions
 -e opentelemetry-sdk
+-e exporter/opentelemetry-exporter-otlp-proto-http
+-e exporter/opentelemetry-exporter-otlp-proto-common
+-e opentelemetry-proto

opentelemetry-sdk/tests/logs/test_export.py

@@ -17,6 +17,7 @@
 import os
 import time
 import unittest
+from concurrent.futures import ThreadPoolExecutor
 from sys import version_info
 from unittest.mock import Mock, patch
 
@@ -331,9 +332,12 @@ def test_simple_log_record_processor_different_msg_types_with_formatter(
         self.assertEqual(expected, emitted)
 
 
+# Many more test cases for the BatchLogRecordProcessor exist under
+# opentelemetry-sdk/tests/shared_internal/test_batch_processor.py.
+# Important: make sure to call .shutdown() on the BatchLogRecordProcessor
+# before the end of the test, otherwise the worker thread will continue
+# to run after the end of the test.
 class TestBatchLogRecordProcessor(unittest.TestCase):
-    # Many more test cases for the BatchLogRecordProcessor exist under
-    # opentelemetry-sdk/tests/shared_internal/test_batch_processor.py.
     def test_emit_call_log_record(self):
         exporter = InMemoryLogExporter()
         log_record_processor = Mock(wraps=BatchLogRecordProcessor(exporter))
@@ -346,6 +350,34 @@ def test_emit_call_log_record(self):
 
         logger.error("error")
         self.assertEqual(log_record_processor.emit.call_count, 1)
+        log_record_processor.shutdown()
+
+    def test_with_multiple_threads(self):
+        exporter = InMemoryLogExporter()
+        batch_processor = BatchLogRecordProcessor(
+            exporter,
+            max_queue_size=3000,
+            max_export_batch_size=50,
+            schedule_delay_millis=30000,
+            export_timeout_millis=500,
+        )
+
+        def bulk_emit(num_emit):
+            for _ in range(num_emit):
+                batch_processor.emit(EMPTY_LOG)
+
+        total_expected_logs = 0
+        with ThreadPoolExecutor(max_workers=69) as executor:
+            for num_logs_to_emit in range(1, 70):
+                executor.submit(bulk_emit, num_logs_to_emit)
+                total_expected_logs += num_logs_to_emit
+
+            executor.shutdown()
+
+        batch_processor.shutdown()
+        # Wait a bit for logs to flush.
+        time.sleep(2)
+        assert len(exporter.get_finished_logs()) == total_expected_logs
 
     @mark.skipif(
         version_info < (3, 10),
@@ -404,6 +436,7 @@ def test_args(self):
         self.assertEqual(
             log_record_processor._batch_processor._export_timeout_millis, 15000
         )
+        log_record_processor.shutdown()
 
     @patch.dict(
         "os.environ",
@@ -432,6 +465,7 @@ def test_env_vars(self):
         self.assertEqual(
             log_record_processor._batch_processor._export_timeout_millis, 15000
         )
+        log_record_processor.shutdown()
 
     def test_args_defaults(self):
         exporter = InMemoryLogExporter()
@@ -451,6 +485,7 @@ def test_args_defaults(self):
         self.assertEqual(
             log_record_processor._batch_processor._export_timeout_millis, 30000
         )
+        log_record_processor.shutdown()
 
     @patch.dict(
         "os.environ",
@@ -481,6 +516,7 @@ def test_args_env_var_value_error(self):
         self.assertEqual(
             log_record_processor._batch_processor._export_timeout_millis, 30000
         )
+        log_record_processor.shutdown()
 
     def test_args_none_defaults(self):
         exporter = InMemoryLogExporter()
@@ -506,6 +542,7 @@ def test_args_none_defaults(self):
         self.assertEqual(
             log_record_processor._batch_processor._export_timeout_millis, 30000
         )
+        log_record_processor.shutdown()
 
     def test_validation_negative_max_queue_size(self):
         exporter = InMemoryLogExporter()

opentelemetry-sdk/tests/shared_internal/test_batch_processor.py

@@ -19,14 +19,19 @@
 import time
 import unittest
 import weakref
-from concurrent.futures import ThreadPoolExecutor
 from platform import system
-from sys import version_info
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
-from pytest import mark
+from requests import Session
+from requests.models import Response
 
+from opentelemetry.exporter.otlp.proto.http._log_exporter import (
+    OTLPLogExporter,
+)
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+    OTLPSpanExporter,
+)
 from opentelemetry.sdk._logs import (
     LogData,
     LogRecord,
@@ -46,13 +51,22 @@
 BASIC_SPAN = ReadableSpan(
     "MySpan",
     instrumentation_scope=InstrumentationScope("example", "example"),
+    context=Mock(
+        **{
+            "trace_state": {"a": "b", "c": "d"},
+            "span_id": 10217189687419569865,
+            "trace_id": 67545097771067222548457157018666467027,
+        }
+    ),
 )
 
 if system() != "Windows":
     multiprocessing.set_start_method("fork")
 
 
 # BatchLogRecodProcessor/BatchSpanProcessor initialize and use BatchProcessor.
+# Important: make sure to call .shutdown() before the end of the test,
+# otherwise the worker thread will continue to run after the end of the test.
 @pytest.mark.parametrize(
     "batch_processor_class,telemetry",
     [(BatchLogRecordProcessor, EMPTY_LOG), (BatchSpanProcessor, BASIC_SPAN)],
@@ -80,6 +94,7 @@ def test_telemetry_exported_once_batch_size_reached(
         after_export = time.time_ns()
         # Shows the worker's 30 second sleep was interrupted within a second.
         assert after_export - before_export < 1e9
+        batch_processor.shutdown()
 
     # pylint: disable=no-self-use
     def test_telemetry_exported_once_schedule_delay_reached(
@@ -96,6 +111,7 @@ def test_telemetry_exported_once_schedule_delay_reached(
         batch_processor._batch_processor.emit(telemetry)
         time.sleep(0.2)
         exporter.export.assert_called_once_with([telemetry])
+        batch_processor.shutdown()
 
     def test_telemetry_flushed_before_shutdown_and_dropped_after_shutdown(
         self, batch_processor_class, telemetry
@@ -136,33 +152,7 @@ def test_force_flush_flushes_telemetry(
             batch_processor._batch_processor.emit(telemetry)
         batch_processor.force_flush()
         exporter.export.assert_called_once_with([telemetry for _ in range(10)])
-
-    @mark.skipif(
-        system() == "Windows" or version_info < (3, 9),
-        reason="This test randomly fails on windows and python 3.8.",
-    )
-    def test_with_multiple_threads(self, batch_processor_class, telemetry):
-        exporter = Mock()
-        batch_processor = batch_processor_class(
-            exporter,
-            max_queue_size=3000,
-            max_export_batch_size=1000,
-            schedule_delay_millis=30000,
-            export_timeout_millis=500,
-        )
-
-        def bulk_emit_and_flush(num_emit):
-            for _ in range(num_emit):
-                batch_processor._batch_processor.emit(telemetry)
-            batch_processor.force_flush()
-
-        with ThreadPoolExecutor(max_workers=69) as executor:
-            for idx in range(69):
-                executor.submit(bulk_emit_and_flush, idx + 1)
-
-            executor.shutdown()
-        # 69 calls to force flush.
-        assert exporter.export.call_count == 69
+        batch_processor.shutdown()
 
     @unittest.skipUnless(
         hasattr(os, "fork"),
@@ -202,6 +192,7 @@ def child(conn):
         batch_processor.force_flush()
         # Single export for the telemetry we emitted at the start of the test.
         assert exporter.export.call_count == 1
+        batch_processor.shutdown()
 
     def test_record_processor_is_garbage_collected(
         self, batch_processor_class, telemetry
@@ -217,3 +208,33 @@ def test_record_processor_is_garbage_collected(
 
         # Then the reference to the processor should no longer exist
         assert weak_ref() is None
+
+    def test_shutdown_waits_30sec_before_cancelling_export(
+        self, batch_processor_class, telemetry, caplog
+    ):
+        resp = Response()
+        resp.status_code = 200
+
+        def export_side_effect(*args, **kwargs):
+            time.sleep(5)
+            return resp
+
+        if type(BASIC_SPAN) is type(telemetry):
+            exporter = OTLPSpanExporter()
+        else:
+            exporter = OTLPLogExporter()
+
+        with patch.object(Session, "post") as mock_post:
+            mock_post.side_effect = export_side_effect
+            processor = batch_processor_class(
+                exporter,
+                max_queue_size=200,
+                max_export_batch_size=10,
+                schedule_delay_millis=30000,
+            )
+            print("emitting..")
+            processor._batch_processor.emit(telemetry)
+            print("shutting down..")
+            processor.shutdown(timeout_millis=4000)
+            print("finished shutting down..")
+        print(caplog.record_tuples)

opentelemetry-sdk/tests/trace/export/test_export.py

@@ -144,9 +144,12 @@ def test_simple_span_processor_not_sampled(self):
         self.assertListEqual([], spans_names_list)
 
 
+# Many more test cases for the BatchSpanProcessor exist under
+# opentelemetry-sdk/tests/shared_internal/test_batch_processor.py.
+# Important: make sure to call .shutdown() on the BatchSpanProcessor
+# before the end of the test, otherwise the worker thread will continue
+# to run after the end of the test.
 class TestBatchSpanProcessor(unittest.TestCase):
-    # Many more test cases for the BatchSpanProcessor exist under
-    # opentelemetry-sdk/tests/shared_internal/test_batch_processor.py.
     @mock.patch.dict(
         "os.environ",
         {
@@ -173,6 +176,7 @@ def test_args_env_var(self):
         self.assertEqual(
             batch_span_processor._batch_processor._export_timeout_millis, 4
         )
+        batch_span_processor.shutdown()
 
     def test_args_env_var_defaults(self):
         batch_span_processor = export.BatchSpanProcessor(
@@ -191,6 +195,7 @@ def test_args_env_var_defaults(self):
         self.assertEqual(
             batch_span_processor._batch_processor._export_timeout_millis, 30000
         )
+        batch_span_processor.shutdown()
 
     @mock.patch.dict(
         "os.environ",
@@ -220,6 +225,7 @@ def test_args_env_var_value_error(self):
         self.assertEqual(
             batch_span_processor._batch_processor._export_timeout_millis, 30000
         )
+        batch_span_processor.shutdown()
 
     def test_on_start_accepts_parent_context(self):
         # pylint: disable=no-self-use