feat(NativeWriter): drop threads before worker

Author: VianneyRuhlmann

ddtrace/_trace/tracer.py

@@ -278,6 +278,7 @@ def __init__(self) -> None:
         atexit.register(self._atexit)
         forksafe.register_before_fork(self._sample_before_fork)
         forksafe.register(self._child_after_fork)
+        forksafe.register_after_parent(self._parent_after_fork)
 
         self._shutdown_lock = RLock()
 
@@ -325,9 +326,10 @@ def deregister_on_start_span(self, func: Callable) -> Callable:
         return func
 
     def _sample_before_fork(self) -> None:
-        span = self.current_root_span()
         if isinstance(self._writer, NativeWriter):
-            self._writer._exporter.drop()
+            self._writer._exporter.stop_worker()
+            self._writer._native_worker.join()
+        span = self.current_root_span()
         if span is not None and span.context.sampling_priority is None:
             self.sample(span)
 
@@ -494,6 +496,11 @@ def _child_after_fork(self):
         self._recreate()
         self._new_process = True
 
+    def _parent_after_fork(self):
+        if isinstance(self._writer, NativeWriter):
+            pass
+            # self._writer.start_worker_thread()
+
     def _recreate(self):
         """Re-initialize the tracer's processors and trace writer. This method should only be used in tests."""
         # Stop the writer.
@@ -950,6 +957,7 @@ def shutdown(self, timeout: Optional[float] = None) -> None:
 
             atexit.unregister(self._atexit)
             forksafe.unregister(self._child_after_fork)
+            forksafe.unregister(self._parent_after_fork)
             forksafe.unregister_before_fork(self._sample_before_fork)
 
         self.start_span = self._start_span_after_shutdown  # type: ignore[assignment]

ddtrace/internal/forksafe.py

@@ -16,6 +16,7 @@
 
 _registry = []  # type: typing.List[typing.Callable[[], None]]
 _registry_before_fork = []  # type: typing.List[typing.Callable[[], None]]
+_registry_after_parent = []  # type: typing.List[typing.Callable[[], None]]
 
 # Some integrations might require after-fork hooks to be executed after the
 # actual call to os.fork with earlier versions of Python (<= 3.6), else issues
@@ -28,16 +29,6 @@
 _forked = False
 
 
-def set_forked():
-    global _forked
-
-    _forked = True
-
-
-def has_forked():
-    return _forked
-
-
 def run_hooks(registry):
     # type: (typing.List[typing.Callable[[], None]]) -> None
     for hook in list(registry):
@@ -50,6 +41,17 @@ def run_hooks(registry):
 
 ddtrace_before_fork = functools.partial(run_hooks, _registry_before_fork)
 ddtrace_after_in_child = functools.partial(run_hooks, _registry)
+ddtrace_after_in_parent = functools.partial(run_hooks, _registry_after_parent)
+
+
+def set_forked():
+    global _forked
+
+    _forked = True
+
+
+def has_forked():
+    return _forked
 
 
 def register_hook(registry, hook):
@@ -59,6 +61,9 @@ def register_hook(registry, hook):
 
 register_before_fork = functools.partial(register_hook, _registry_before_fork)
 register = functools.partial(register_hook, _registry)
+register_after_parent = functools.partial(register_hook, _registry_after_parent)
+
+register_after_parent(set_forked)
 
 
 def unregister(after_in_child):
@@ -69,6 +74,14 @@ def unregister(after_in_child):
         log.info("after_in_child hook %s was unregistered without first being registered", after_in_child.__name__)
 
 
+def unregister_parent(after_in_parent):
+    # type: (typing.Callable[[], None]) -> None
+    try:
+        _registry.remove(after_in_parent)
+    except ValueError:
+        log.info("after_in_parent hook %s was unregistered without first being registered", after_in_parent.__name__)
+
+
 def unregister_before_fork(before_fork):
     # type: (typing.Callable[[], None]) -> None
     try:
@@ -78,7 +91,9 @@ def unregister_before_fork(before_fork):
 
 
 if hasattr(os, "register_at_fork"):
-    os.register_at_fork(before=ddtrace_before_fork, after_in_child=ddtrace_after_in_child, after_in_parent=set_forked)
+    os.register_at_fork(
+        before=ddtrace_before_fork, after_in_child=ddtrace_after_in_child, after_in_parent=ddtrace_after_in_parent
+    )
 
 _resetable_objects = weakref.WeakSet()  # type: weakref.WeakSet[ResetObject]
 

ddtrace/internal/native/_native.pyi

@@ -86,7 +86,9 @@ class TraceExporter:
         """
         ...
 
-    def drop(self) -> None:...
+    def drop(self) -> None: ...
+    def run_worker(self) -> None: ...
+    def stop_worker(self) -> None: ...
 
 class TraceExporterBuilder:
     """

ddtrace/internal/writer/writer.py

@@ -31,6 +31,7 @@
 from ..agent import get_connection
 from ..constants import _HTTPLIB_NO_TRACE_REQUEST
 from ..encoding import JSONEncoderV2
+from ..gitmetadata import get_git_tags
 from ..logger import get_logger
 from ..serverless import in_azure_function
 from ..serverless import in_gcp_function
@@ -671,18 +672,20 @@ def __init__(
         self._compute_stats_enabled = compute_stats_enabled
         self._response_cb = response_callback
         self._exporter = self._create_exporter()
+        self.start_worker_thread()
+
+    def start_worker_thread(self):
+        self._native_worker = threading.Thread(target=self._exporter.run_worker, daemon=True)
+        self._native_worker.start()
 
     def _create_exporter(self) -> native.TraceExporter:
         """
         Create a new TraceExporter with the current configuration.
         :return: A configured TraceExporter instance.
         """
-        # Shutdown the existing exporter if it exists
-        if hasattr(self, "_exporter") and self._exporter is not None:
-            self._exporter.shutdown(3_000_000_000)  # 3 seconds timeout
-
         stats_interval = float(os.getenv("_DD_TRACE_STATS_WRITER_INTERVAL") or 10.0)
         bucket_size_ns = int(stats_interval * 1e9)  # type: int
+        _, commit_sha, _ = get_git_tags()
 
         builder = (
             native.TraceExporterBuilder()
@@ -691,6 +694,7 @@ def _create_exporter(self) -> native.TraceExporter:
             .set_language_version(compat.PYTHON_VERSION)
             .set_language_interpreter(compat.PYTHON_INTERPRETER)
             .set_tracer_version(ddtrace.__version__)
+            .set_git_commit_sha(commit_sha)
             .set_client_computed_top_level()
             .set_input_format(self._api_version)
             .set_output_format(self._api_version)
@@ -707,10 +711,14 @@ def set_test_session_token(self, token: Optional[str]) -> None:
         :param token: The test session token to use for authentication.
         """
         self._test_session_token = token
+        self._exporter.stop_worker()
+        self._native_worker.join()
         self._exporter = self._create_exporter()
+        self.start_worker_thread()
 
     def recreate(self):
-        self._exporter.drop()
+        if self._native_worker.is_alive():
+            self._exporter.stop_worker()
         return self.__class__(
             agent_url=self.agent_url,
             processing_interval=self._interval,
@@ -793,6 +801,8 @@ def _send_payload(self, payload: bytes, count: int, client: WriterClientBase):
 
         try:
             # TODO: Return agent response from send
+            if not self._native_worker.is_alive:
+                self.start_worker_thread()
             response_body = self._exporter.send(payload, count)
             if self._response_cb:
                 response = Response(body=response_body)
@@ -931,4 +941,7 @@ def on_shutdown(self):
         try:
             self.periodic()
         finally:
+            if self._native_worker.is_alive():
+                self._exporter.stop_worker()
+                self._native_worker.join()
             self._exporter.shutdown(3_000_000_000)  # 3 seconds timeout