maybe leak

Author: P403n1x87

ddtrace/_trace/context.py

@@ -91,9 +91,6 @@ def __init__(
         if lock is not None:
             self._lock = lock
         else:
-            # DEV: A `forksafe.RLock` is not necessary here since Contexts
-            # are recreated by the tracer after fork
-            # https://github.com/DataDog/dd-trace-py/blob/a1932e8ddb704d259ea8a3188d30bf542f59fd8d/ddtrace/tracer.py#L489-L508
             self._lock = RLock()
 
     def __getstate__(self) -> _ContextState:

ddtrace/internal/_encoding.pyx

@@ -4,7 +4,6 @@ from libc cimport stdint
 from libc.string cimport strlen
 
 from json import dumps as json_dumps
-import threading
 from json import dumps as json_dumps
 
 from ._utils cimport PyBytesLike_Check

ddtrace/internal/_threads.cpp

@@ -516,6 +516,8 @@ static PyTypeObject PeriodicThreadType = {
 // ----------------------------------------------------------------------------
 static PyMethodDef _threads_methods[] = {
     { "reset_locks", (PyCFunction)lock_reset_locks, METH_NOARGS, "Reset all locks (generally after a fork)" },
+    { "begin_reset_locks", (PyCFunction)lock_begin_reset_locks, METH_NOARGS, "Begin resetting locks (before a fork)" },
+    { "end_reset_locks", (PyCFunction)lock_end_reset_locks, METH_NOARGS, "End resetting locks (after a fork)" },
     { NULL, NULL, 0, NULL } /* Sentinel */
 };
 

ddtrace/internal/_threads.pyi

@@ -12,6 +12,8 @@ class Lock(_BaseLock): ...
 class RLock(_BaseLock): ...
 
 def reset_locks() -> None: ...
+def begin_reset_locks() -> None: ...
+def end_reset_locks() -> None: ...
 
 class PeriodicThread:
     name: str

ddtrace/internal/_threads/lock.hpp

@@ -7,7 +7,7 @@
 #include <mutex>
 #include <set>
 
-std::unique_ptr<std::mutex> _lock_set_mutex = std::make_unique<std::mutex>();
+std::mutex _lock_set_mutex;
 
 // ----------------------------------------------------------------------------
 // Lock class
@@ -35,14 +35,25 @@ Lock_init(Lock* self, PyObject* args, PyObject* kwargs)
     {
         AllowThreads _;
 
-        std::lock_guard<std::mutex> guard(*_lock_set_mutex);
+        std::lock_guard<std::mutex> guard(_lock_set_mutex);
 
         lock_set.insert(self);
     }
 
     return 0;
 }
 
+// ----------------------------------------------------------------------------
+static inline void
+_Lock_maybe_leak(Lock* self)
+{
+    // This function is used to ensure that the mutex is not leaked if it is
+    // still locked when the lock object is deallocated.
+    if (self->_locked) {
+        self->_mutex.release(); // DEV: This releases the unique_ptr, not the mutex!
+    }
+}
+
 // ----------------------------------------------------------------------------
 static void
 Lock_dealloc(Lock* self)
@@ -51,11 +62,13 @@ Lock_dealloc(Lock* self)
     {
         AllowThreads _;
 
-        std::lock_guard<std::mutex> guard(*_lock_set_mutex);
+        std::lock_guard<std::mutex> guard(_lock_set_mutex);
 
         lock_set.erase(self);
     }
 
+    _Lock_maybe_leak(self);
+
     self->_mutex = nullptr;
 
     Py_TYPE(self)->tp_free((PyObject*)self);
@@ -153,6 +166,7 @@ Lock_exit(Lock* self, PyObject* args, PyObject* kwargs)
 static inline void
 Lock_reset(Lock* self)
 {
+    _Lock_maybe_leak(self);
     self->_mutex = std::make_unique<std::timed_mutex>();
     self->_locked = 0;
 }
@@ -212,26 +226,38 @@ RLock_init(RLock* self, PyObject* args, PyObject* kwargs)
     {
         AllowThreads _;
 
-        std::lock_guard<std::mutex> guard(*_lock_set_mutex);
+        std::lock_guard<std::mutex> guard(_lock_set_mutex);
 
         rlock_set.insert(self);
     }
 
     return 0;
 }
 
+// ----------------------------------------------------------------------------
+static inline void
+_RLock_maybe_leak(RLock* self)
+{
+    // This function is used to ensure that the mutex is not leaked if it is
+    // still locked when the re-entrant lock object is deallocated.
+    if (self->_locked) {
+        self->_mutex.release(); // DEV: This releases the unique_ptr, not the mutex!
+    }
+}
+
 // ----------------------------------------------------------------------------
 static void
 RLock_dealloc(RLock* self)
 {
     {
         AllowThreads _;
 
-        std::lock_guard<std::mutex> guard(*_lock_set_mutex);
+        std::lock_guard<std::mutex> guard(_lock_set_mutex);
 
         rlock_set.erase(self);
     }
 
+    _RLock_maybe_leak(self);
     self->_mutex = nullptr;
 
     Py_TYPE(self)->tp_free((PyObject*)self);
@@ -329,6 +355,7 @@ RLock_exit(RLock* self, PyObject* args, PyObject* kwargs)
 static inline void
 RLock_reset(RLock* self)
 {
+    _RLock_maybe_leak(self);
     self->_mutex = std::make_unique<std::recursive_timed_mutex>();
     self->_locked = 0;
 }
@@ -380,8 +407,33 @@ lock_reset_locks(PyObject* Py_UNUSED(self), PyObject* Py_UNUSED(args))
         RLock_reset(rlock);
     }
 
-    // Reset the lock set mutex too!
-    _lock_set_mutex = std::make_unique<std::mutex>();
+    _lock_set_mutex.unlock();
+
+    Py_RETURN_NONE;
+}
+
+// ----------------------------------------------------------------------------
+static PyObject*
+lock_begin_reset_locks(PyObject* Py_UNUSED(self), PyObject* Py_UNUSED(args))
+{
+    // This function is called before a fork to ensure that the lock set mutex
+    // is not held by any thread.
+    {
+        AllowThreads _;
+
+        _lock_set_mutex.lock();
+    }
+
+    Py_RETURN_NONE;
+}
+
+// ----------------------------------------------------------------------------
+static PyObject*
+lock_end_reset_locks(PyObject* Py_UNUSED(self), PyObject* Py_UNUSED(args))
+{
+    // This function is called after a fork to ensure that the lock set mutex
+    // is released and can be used by the new process.
+    _lock_set_mutex.unlock();
 
     Py_RETURN_NONE;
 }

ddtrace/internal/rate_limiter.py

@@ -3,7 +3,6 @@
 from dataclasses import dataclass
 from dataclasses import field
 import random
-import threading
 import time
 from typing import Any  # noqa:F401
 from typing import Callable  # noqa:F401
@@ -204,7 +203,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: Lock = field(init=False, default_factory=Lock)
 
     def __post_init__(self):
         if self.limit_rate == float("inf"):

ddtrace/internal/threads.py

@@ -4,14 +4,15 @@
 from ddtrace.internal._threads import Lock
 from ddtrace.internal._threads import PeriodicThread
 from ddtrace.internal._threads import RLock
+from ddtrace.internal._threads import begin_reset_locks
+from ddtrace.internal._threads import end_reset_locks
 from ddtrace.internal._threads import periodic_threads
 from ddtrace.internal._threads import reset_locks
 
 
 __all__ = [
     "Lock",
     "PeriodicThread",
-    "periodic_threads",
     "RLock",
 ]
 
@@ -36,4 +37,6 @@ def _() -> None:
     periodic_threads.clear()
 
 
+forksafe.register_before_fork(begin_reset_locks)
 forksafe.register(reset_locks)
+forksafe.register_after_parent(end_reset_locks)