Make batch processor fork aware and reinit when needed

opentelemetry-sdk/src/opentelemetry/sdk/trace/export/__init__.py

@@ -14,6 +14,7 @@
 
 import collections
 import logging
+import os
 import sys
 import threading
 import typing
@@ -197,6 +198,11 @@ def __init__(
             None
         ] * self.max_export_batch_size  # type: typing.List[typing.Optional[Span]]
         self.worker_thread.start()
+        # Only available in *nix since py37.
+        if hasattr(os, "register_at_fork"):
+            os.register_at_fork(
+                after_in_child=self._at_fork_reinit
+            )  # pylint: disable=broad-except
 
     def on_start(
         self, span: Span, parent_context: typing.Optional[Context] = None
@@ -220,6 +226,24 @@ def on_end(self, span: ReadableSpan) -> None:
             with self.condition:
                 self.condition.notify()
 
+    def _at_fork_reinit(self):
+        # could be in an inconsistent state after fork, reinitialise by calling `_at_fork_reinit`
+        # (creates a new lock internally https://github.com/python/cpython/blob/main/Python/thread_pthread.h#L727)
+        # if exists, otherwise create a new one.
+        if hasattr(self.condition, "_at_fork_reinit"):
+            self.condition._at_fork_reinit()
+        else:
+            self.condition = threading.Condition(threading.Lock())
+
+        self.queue.clear()
+
+        # worker_thread is local to a process, only the thread that issued fork continues
+        # to exist. A new worker thread must be started in child process.
+        self.worker_thread = threading.Thread(
+            name="OtelBatchSpanProcessor", target=self.worker, daemon=True
+        )
+        self.worker_thread.start()
+
     def worker(self):
         timeout = self.schedule_delay_millis / 1e3
         flush_request = None  # type: typing.Optional[_FlushRequest]

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

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import sys
 import threading
 import time
 import unittest
@@ -30,6 +31,10 @@
     OTEL_BSP_SCHEDULE_DELAY,
 )
 from opentelemetry.sdk.trace import export
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+    InMemorySpanExporter,
+)
+from opentelemetry.test.concurrency_test import ConcurrencyTestBase
 
 
 class MySpanExporter(export.SpanExporter):
@@ -157,7 +162,7 @@ def _create_start_and_end_span(name, span_processor):
     span.end()
 
 
-class TestBatchSpanProcessor(unittest.TestCase):
+class TestBatchSpanProcessor(ConcurrencyTestBase):
     @mock.patch.dict(
         "os.environ",
         {
@@ -356,6 +361,63 @@ def test_batch_span_processor_not_sampled(self):
         self.assertEqual(len(spans_names_list), 0)
         span_processor.shutdown()
 
+    def _check_fork_trace(self, exporter, expected):
+        time.sleep(0.5)  # give some time for the exporter to upload spans
+        spans = exporter.get_finished_spans()
+        for span in spans:
+            self.assertIn(span.name, expected)
+
+    @unittest.skipUnless(
+        hasattr(os, "fork") and sys.version_info >= (3, 7),
+        "needs *nix and minor version 7 or later",
+    )
+    def test_batch_span_processor_fork(self):
+        tracer_provider = trace.TracerProvider()
+        tracer = tracer_provider.get_tracer(__name__)
+
+        exporter = InMemorySpanExporter()
+        span_processor = export.BatchSpanProcessor(
+            exporter,
+            max_queue_size=256,
+            max_export_batch_size=64,
+            schedule_delay_millis=10,
+        )
+        tracer_provider.add_span_processor(span_processor)
+        with tracer.start_as_current_span("foo"):
+            pass
+        time.sleep(0.5)  # give some time for the exporter to upload spans
+
+        self.assertTrue(span_processor.force_flush())
+        self.assertEqual(len(exporter.get_finished_spans()), 1)
+        exporter.clear()
+        pid = os.fork()
+        if pid:
+            with tracer.start_as_current_span("parent"):
+                pass
+            self._check_fork_trace(exporter, ["parent"])
+        else:
+            exporter.clear()
+
+            def _target():
+                with tracer.start_as_current_span(f"span") as s:
+                    s.set_attribute("i", "1")
+                    with tracer.start_as_current_span("temp"):
+                        pass
+
+            self.run_with_many_threads(_target, 100)
+
+            time.sleep(0.5)
+
+            spans = exporter.get_finished_spans()
+            self.assertEqual(len(spans), 200)
+            exporter.clear()
+            with tracer.start_as_current_span("child"):
+                with tracer.start_as_current_span("inner"):
+                    pass
+            self._check_fork_trace(exporter, ["child", "inner"])
+
+        span_processor.shutdown()
+
     def test_batch_span_processor_scheduled_delay(self):
         """Test that spans are exported each schedule_delay_millis"""
         spans_names_list = []