refactor(internal): use fork-safe locks and threads consistently

.sg/rules/threading-lock-usage.yml

@@ -0,0 +1,50 @@
+id: threading-lock-usage
+message: Use ddtrace.internal.threads.Lock/RLock instead of threading.Lock/RLock
+severity: error
+language: python
+files:
+  - "ddtrace/**"
+ignores:
+  - "ddtrace/internal/threads.py"
+  - "ddtrace/internal/_unpatched.py"
+  - "ddtrace/vendor/**"
+  - "ddtrace/appsec/_iast/_taint_tracking/_vendor/**"
+  - "ddtrace/testing/**"
+  - "ddtrace/internal/openfeature/**"
+  - "ddtrace/contrib/internal/pytorch/**"
+  - "ddtrace/llmobs/_integrations/_bedrock_inference_profiles.py"
+  - "ddtrace/appsec/_iast/_overhead_control_engine.py"
+  - "ddtrace/profiling/_faulthandler.py"
+rule:
+  any:
+    - pattern: threading.Lock()
+    - pattern: threading.RLock()
+    - pattern: from threading import Lock
+    - pattern: from threading import RLock
+    - pattern: from threading import Lock as $ALIAS
+    - pattern: from threading import RLock as $ALIAS
+    - pattern: from threading import $$$IMPORTS, Lock
+    - pattern: from threading import $$$IMPORTS, Lock as $ALIAS
+    - pattern: from threading import Lock, $$$IMPORTS
+    - pattern: from threading import Lock as $ALIAS, $$$IMPORTS
+    - pattern: from threading import $$$IMPORTS, RLock
+    - pattern: from threading import $$$IMPORTS, RLock as $ALIAS
+    - pattern: from threading import RLock, $$$IMPORTS
+    - pattern: from threading import RLock as $ALIAS, $$$IMPORTS
+note: |
+  gevent's monkey-patching replaces `threading.Lock`/`threading.RLock` with
+  green (cooperative) equivalents. Mixing those with dd-trace-py's real OS
+  threads causes deadlocks, since a green lock held by a real thread can never
+  be released by a greenlet waiting on the event loop.
+
+  `ddtrace.internal.threads.Lock`/`RLock` grab a reference to the real,
+  unpatched primitives before any monkey-patching library can run, and are
+  fork-safe. Use them instead:
+
+  Before:
+    import threading
+    lock = threading.Lock()
+
+  After:
+    from ddtrace.internal.threads import Lock
+    lock = Lock()

.sg/tests/__snapshots__/threading-lock-usage-snapshot.yml

@@ -0,0 +1,50 @@
+id: threading-lock-usage
+snapshots:
+  from threading import Lock as TLock:
+    labels:
+    - source: from threading import Lock as TLock
+      style: primary
+      start: 0
+      end: 35
+  from threading import Lock, Thread:
+    labels:
+    - source: from threading import Lock, Thread
+      style: primary
+      start: 0
+      end: 34
+  from threading import RLock as TRLock:
+    labels:
+    - source: from threading import RLock as TRLock
+      style: primary
+      start: 0
+      end: 37
+  from threading import RLock, Thread:
+    labels:
+    - source: from threading import RLock, Thread
+      style: primary
+      start: 0
+      end: 35
+  from threading import Thread, Lock:
+    labels:
+    - source: from threading import Thread, Lock
+      style: primary
+      start: 0
+      end: 34
+  from threading import Thread, RLock:
+    labels:
+    - source: from threading import Thread, RLock
+      style: primary
+      start: 0
+      end: 35
+  import threading; lock = threading.Lock():
+    labels:
+    - source: threading.Lock()
+      style: primary
+      start: 25
+      end: 41
+  import threading; lock = threading.RLock():
+    labels:
+    - source: threading.RLock()
+      style: primary
+      start: 25
+      end: 42

.sg/tests/threading-lock-usage-test.yml

@@ -0,0 +1,21 @@
+id: threading-lock-usage
+valid:
+  # These should NOT trigger the rule (valid code)
+  - from ddtrace.internal.threads import Lock; lock = Lock()
+  - from ddtrace.internal.threads import RLock; lock = RLock()
+  - import threading; threading.Event()
+  - import threading; threading.current_thread()
+  - from threading import Thread
+
+invalid:
+  # Direct attribute calls
+  - import threading; lock = threading.Lock()
+  - import threading; lock = threading.RLock()
+  # Aliased imports
+  - from threading import Lock as TLock
+  - from threading import RLock as TRLock
+  # Comma-separated imports
+  - from threading import Thread, Lock
+  - from threading import Lock, Thread
+  - from threading import Thread, RLock
+  - from threading import RLock, Thread

ddtrace/_trace/processor/__init__.py

@@ -2,7 +2,6 @@
 from collections import defaultdict
 from itertools import chain
 import logging
-from threading import RLock
 from typing import Optional
 
 from ddtrace._trace.sampler import DatadogSampler
@@ -29,6 +28,7 @@
 from ddtrace.internal.telemetry.constants import TELEMETRY_NAMESPACE
 from ddtrace.internal.telemetry.metrics import MetricRecorder
 from ddtrace.internal.telemetry.metrics import get_metric_recorder
+from ddtrace.internal.threads import RLock
 from ddtrace.internal.writer import AgentlessTraceWriter
 from ddtrace.internal.writer import AgentResponse
 from ddtrace.internal.writer import LogWriter

ddtrace/_trace/tracer.py

@@ -7,7 +7,6 @@
 import logging
 import os
 from os import getpid
-from threading import Lock
 from typing import Any
 from typing import AsyncGenerator
 from typing import Callable
@@ -57,6 +56,7 @@
 from ddtrace.internal.settings._config import config
 from ddtrace.internal.settings.asm import config as asm_config
 from ddtrace.internal.settings.peer_service import _ps_config
+from ddtrace.internal.threads import Lock
 from ddtrace.internal.utils import _get_metas_to_propagate
 from ddtrace.internal.utils.deprecations import DDTraceDeprecationWarning
 from ddtrace.internal.utils.formats import format_trace_id

ddtrace/appsec/sca/_instrumenter.py

@@ -7,7 +7,6 @@
 from __future__ import annotations
 
 import sys
-from threading import Lock
 import types
 from types import FunctionType
 from typing import TYPE_CHECKING
@@ -22,9 +21,12 @@
 from ddtrace.internal.module import ModuleHookType
 from ddtrace.internal.module import ModuleWatchdog
 from ddtrace.internal.telemetry import telemetry_writer
+from ddtrace.internal.threads import Lock
 
 
 if TYPE_CHECKING:
+    from _thread import LockType
+
     from ddtrace.appsec.sca._registry import InstrumentationRegistry
 
 
@@ -172,11 +174,11 @@ class Instrumenter:
 
     def __init__(self, registry: InstrumentationRegistry) -> None:
         self.registry = registry
-        self._instrumentation_locks: dict[str, Lock] = {}
+        self._instrumentation_locks: dict[str, LockType] = {}
         self._locks_lock = Lock()
         set_registry(registry)
 
-    def _get_lock(self, qualified_name: str) -> Lock:
+    def _get_lock(self, qualified_name: str) -> LockType:
         with self._locks_lock:
             if qualified_name not in self._instrumentation_locks:
                 self._instrumentation_locks[qualified_name] = Lock()

ddtrace/appsec/sca/_registry.py

@@ -6,12 +6,12 @@
 
 from dataclasses import dataclass
 from dataclasses import field
-from threading import Lock
 from types import CodeType
 from typing import NamedTuple
 from typing import Optional
 
 from ddtrace.internal.logger import get_logger
+from ddtrace.internal.threads import Lock
 
 
 log = get_logger(__name__)

ddtrace/internal/_threads.pyi

@@ -17,7 +17,7 @@ class PeriodicThread:
     def __init__(
         self,
         interval: float,
-        target: t.Callable[[], None],
+        target: t.Callable[[], object],
         name: t.Optional[str] = None,
         on_shutdown: t.Optional[t.Callable[[], None]] = None,
         no_wait_at_start: bool = False,

ddtrace/internal/ci_visibility/encoder.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import json
-import threading
 from typing import TYPE_CHECKING  # noqa:F401
 from typing import Any  # noqa:F401
 from typing import Optional  # noqa:F401
@@ -25,6 +24,7 @@
 from ddtrace.internal.encoding import JSONEncoderV2
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.settings import env
+from ddtrace.internal.threads import RLock
 from ddtrace.internal.utils.time import StopWatch
 from ddtrace.internal.writer.writer import NoEncodableSpansError
 
@@ -49,7 +49,7 @@ def __init__(self, *args: Any) -> None:
         #      which is called implicitly by Cython.
         super(CIVisibilityEncoderV01, self).__init__()  # type: ignore[call-arg]
         self._metadata: dict[str, dict[str, str]] = {}
-        self._lock = threading.RLock()
+        self._lock = RLock()
         self._is_xdist_worker = env.get("PYTEST_XDIST_WORKER") is not None
         self._init_buffer()
 

ddtrace/internal/datastreams/schemas/schema_sampler.py

@@ -1,4 +1,4 @@
-import threading
+from ddtrace.internal.threads import Lock
 
 
 class SchemaSampler:
@@ -7,7 +7,7 @@ class SchemaSampler:
     def __init__(self):
         self.weight = 0
         self.last_sample_millis = 0
-        self.lock = threading.Lock()
+        self.lock = Lock()
 
     def try_sample(self, current_time_millis):
         if current_time_millis >= self.last_sample_millis + self.SAMPLE_INTERVAL_MILLIS:

ddtrace/internal/excepthook.py

@@ -17,20 +17,20 @@
 """
 
 import sys
-import threading
 from types import TracebackType
 from typing import Any
 from typing import Callable
 from typing import Optional
 
 from ddtrace.internal.logger import get_logger
+from ddtrace.internal.threads import Lock
 
 
 log = get_logger(__name__)
 
 ExceptHookType = Callable[[type[BaseException], BaseException, Optional[TracebackType]], Any]
 
-_lock = threading.Lock()
+_lock = Lock()
 _hooks: list[ExceptHookType] = []
 # The hook that was installed before ddtrace took over sys.excepthook. Captured
 # once, the first time a callback is registered. Always run last.

ddtrace/internal/rate_limiter.py

@@ -1,12 +1,20 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
 from dataclasses import field
 import random
-import threading
 import time
+from typing import TYPE_CHECKING
 from typing import Any  # noqa:F401
 from typing import Callable  # noqa:F401
 from typing import Optional  # noqa:F401
 
+from ddtrace.internal.threads import Lock
+
+
+if TYPE_CHECKING:
+    from _thread import LockType
+
 
 class RateLimiter(object):
     """
@@ -50,7 +58,7 @@ def __init__(self, rate_limit: int, time_window: float = 1e9):
         self.tokens_total = 0
         self.prev_window_rate = None  # type: Optional[float]
 
-        self._lock = threading.Lock()
+        self._lock = Lock()
 
     def is_allowed(self) -> bool:
         """
@@ -198,7 +206,7 @@ class BudgetRateLimiterWithJitter:
     budget: float = field(init=False)
     max_budget: float = field(init=False)
     last_time: float = field(init=False, default_factory=time.monotonic)
-    _lock: threading.Lock = field(init=False, default_factory=threading.Lock)
+    _lock: LockType = field(init=False, default_factory=Lock)
 
     def __post_init__(self):
         if self.limit_rate == float("inf"):

ddtrace/internal/telemetry/dependency_tracker.py

@@ -11,7 +11,6 @@
 
 from importlib.metadata import PackageNotFoundError
 import re
-from threading import Lock
 from typing import Any
 from typing import Iterable
 from typing import Optional
@@ -20,6 +19,7 @@
 from ddtrace.internal.packages import get_module_distribution_versions
 from ddtrace.internal.settings._telemetry import config as telemetry_config
 from ddtrace.internal.settings.appsec_telemetry import config as appsec_telemetry_config
+from ddtrace.internal.threads import Lock
 
 from . import modules
 from .dependency import DependencyEntry

ddtrace/internal/threads.py

@@ -2,6 +2,7 @@
 import typing as t
 
 from ddtrace.internal import forksafe
+from ddtrace.internal._threads import PERIODIC_STOP
 from ddtrace.internal._threads import PeriodicThread as _PeriodicThread
 from ddtrace.internal._threads import periodic_threads
 from ddtrace.internal.logger import get_logger
@@ -28,6 +29,7 @@
 
 __all__ = [
     "Lock",
+    "Thread",
     "PeriodicThread",
     "RLock",
 ]
@@ -82,6 +84,26 @@ def start(self) -> None:
                 _threads_to_start_after_fork.append(t.cast(BoundMethod, super().start))
 
 
+class Thread(PeriodicThread):
+    """A fork-safe thread."""
+
+    # A one-shot thread runs its target exactly once, ever. _before_fork()
+    # joins it to completion before the fork happens, so by the time restart
+    # decisions are made it has already done its job. __autorestart__ = False
+    # stops the child from restarting it; the parent path also skips it
+    # explicitly (see ThreadRestartTimer._restart_threads) since force=True
+    # there bypasses __autorestart__ by design for genuinely periodic threads.
+    __autorestart__ = False
+
+    def __init__(self, target: t.Callable[[], None], name: t.Optional[str] = None) -> None:
+        super().__init__(0.0, self._run_once, name=name, no_wait_at_start=True)
+        self._target = target
+
+    def _run_once(self) -> object:
+        self._target()
+        return PERIODIC_STOP
+
+
 # Set of running periodic threads that need to be restarted after a fork.
 _threads_to_restart_after_fork: set[_PeriodicThread] = set()
 
@@ -117,6 +139,12 @@ def _restart_threads(self) -> None:
                         # to avoid restarting orphaned timer instances that were
                         # caught in periodic_threads during a fork.
                         continue
+                    if isinstance(thread, Thread):
+                        # One-shot threads already ran to completion during
+                        # _before_fork()'s join. force=True below bypasses
+                        # __autorestart__, so they must be excluded explicitly
+                        # to avoid re-running their target in the parent.
+                        continue
                     log.debug("Restarting thread %s after fork", thread.name)
                     try:
                         thread._after_fork(force=True)