fix: respect timeout_millis in BatchProcessor.force_flush (#4568)

Author: chimchim89

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

@@ -169,12 +169,15 @@ def worker(self):
             self._worker_awaken.clear()
         self._export(BatchExportStrategy.EXPORT_ALL)
 
-    def _export(self, batch_strategy: BatchExportStrategy) -> None:
+    def _export(self, batch_strategy: BatchExportStrategy, flush_should_end: Optional[float] = None) -> bool:
+        # Returns True if all batches were exported, False if flush_should_end was reached.
         with self._export_lock:
             iteration = 0
             # We could see concurrent export calls from worker and force_flush. We call _should_export_batch
             # once the lock is obtained to see if we still need to make the requested export.
             while self._should_export_batch(batch_strategy, iteration):
+                if flush_should_end is not None and time.time() >= flush_should_end:
+                    return False
                 iteration += 1
                 token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True))
                 try:
@@ -195,6 +198,7 @@ def _export(self, batch_strategy: BatchExportStrategy) -> None:
                         "Exception while exporting %s.", self._exporting
                     )
                 detach(token)
+        return True
 
     def emit(self, data: Telemetry) -> None:
         if self._shutdown:
@@ -236,10 +240,13 @@ def shutdown(self, timeout_millis: int = 30000):
         # call is ongoing and the thread isn't finished. In this case we will return instead of waiting on
         # the thread to finish.
 
-    # 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:
         if self._shutdown:
             return False
+        flush_should_end = (
+            time.time() + (timeout_millis / 1000)
+            if timeout_millis is not None
+            else None
+        )
         # Blocking call to export.
-        self._export(BatchExportStrategy.EXPORT_ALL)
-        return True
+        return self._export(BatchExportStrategy.EXPORT_ALL, flush_should_end)

opentelemetry-sdk/tests/shared_internal/test_batch_processor.py

@@ -171,6 +171,75 @@ def test_force_flush_flushes_telemetry(
         exporter.export.assert_called_once_with([telemetry for _ in range(10)])
         batch_processor.shutdown()
 
+    # pylint: disable=no-self-use
+    def test_force_flush_returns_true_when_all_exported(
+        self, batch_processor_class, telemetry
+    ):
+        exporter = Mock()
+        batch_processor = batch_processor_class(
+            exporter,
+            max_queue_size=15,
+            max_export_batch_size=15,
+            schedule_delay_millis=30000,
+            export_timeout_millis=500,
+        )
+        for _ in range(10):
+            batch_processor._batch_processor.emit(telemetry)
+        result = batch_processor.force_flush(timeout_millis=5000)
+        assert result is True
+        exporter.export.assert_called_once()
+        batch_processor.shutdown()
+
+    # pylint: disable=no-self-use
+    def test_force_flush_returns_false_when_timeout_exceeded(
+        self, batch_processor_class, telemetry
+    ):
+        call_count = 0
+
+        def slow_export(batch):
+            nonlocal call_count
+            call_count += 1
+            # Sleep long enough that the deadline is exceeded after first batch.
+            time.sleep(0.2)
+
+        exporter = Mock()
+        exporter.export.side_effect = slow_export
+        batch_processor = batch_processor_class(
+            exporter,
+            max_queue_size=50,
+            max_export_batch_size=1,
+            schedule_delay_millis=30000,
+            export_timeout_millis=500,
+        )
+        for _ in range(10):
+            batch_processor._batch_processor.emit(telemetry)
+        # 100ms timeout, each export takes 200ms, so deadline is hit after first batch.
+        result = batch_processor.force_flush(timeout_millis=100)
+        assert result is False
+        # Exporter was called at least once but not for all batches.
+        assert 1 <= call_count < 10
+        batch_processor.shutdown()
+
+    # pylint: disable=no-self-use
+    def test_force_flush_returns_false_when_shutdown(
+        self, batch_processor_class, telemetry
+    ):
+        exporter = Mock()
+        batch_processor = batch_processor_class(
+            exporter,
+            max_queue_size=15,
+            max_export_batch_size=15,
+            schedule_delay_millis=30000,
+            export_timeout_millis=500,
+        )
+        batch_processor.shutdown()
+        for _ in range(10):
+            batch_processor._batch_processor.emit(telemetry)
+        result = batch_processor.force_flush(timeout_millis=5000)
+        assert result is False
+        # Nothing should have been exported after shutdown.
+        exporter.export.assert_not_called()
+
     @unittest.skipUnless(
         hasattr(os, "fork"),
         "needs *nix",