Skip to content

Commit 95a6cb7

Browse files
P403n1x87P403n1x87
committed
implement both locks in native
1 parent 4baebd8 commit 95a6cb7

File tree

4 files changed

+216
-71
lines changed

4 files changed

+216
-71
lines changed

‎ddtrace/internal/_threads.cpp

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -540,6 +540,9 @@ PyInit__threads(void)
540540
if (PyType_Ready(&LockType) < 0)
541541
return NULL;
542542

543+
if (PyType_Ready(&RLockType) < 0)
544+
return NULL;
545+
543546
_periodic_threads = PyDict_New();
544547
if (_periodic_threads == NULL)
545548
return NULL;
@@ -565,6 +568,13 @@ PyInit__threads(void)
565568
goto error;
566569
}
567570

571+
// RLock
572+
Py_INCREF(&RLockType);
573+
if (PyModule_AddObject(m, "RLock", (PyObject*)&RLockType) < 0) {
574+
Py_DECREF(&RLockType);
575+
goto error;
576+
}
577+
568578
return m;
569579

570580
error:

‎ddtrace/internal/_threads.pyi

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,15 @@
11
import typing as t
22

3-
class Lock:
3+
class _BaseLock:
44
def __init__(self, reentrant: bool = False) -> None: ...
55
def acquire(self, timeout: t.Optional[float] = None) -> bool: ...
66
def release(self) -> None: ...
77
def locked(self) -> bool: ...
8+
def __enter__(self) -> None: ...
9+
def __exit__(self, exc_type, exc_value, traceback) -> t.Literal[False]: ...
10+
11+
class Lock(_BaseLock): ...
12+
class RLock(_BaseLock): ...
813

914
class PeriodicThread:
1015
name: str

‎ddtrace/internal/_threads/lock.hpp

Lines changed: 196 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@
77
#include <mutex>
88

99
// ----------------------------------------------------------------------------
10+
// Lock class
11+
// ----------------------------------------------------------------------------
12+
1013
typedef struct lock
1114
{
1215
PyObject_HEAD
@@ -15,26 +18,13 @@ typedef struct lock
1518
_locked = 0;
1619

1720
std::unique_ptr<std::timed_mutex> _mutex = nullptr;
18-
std::unique_ptr<std::recursive_timed_mutex> _rmutex = nullptr;
1921
} Lock;
2022

2123
// ----------------------------------------------------------------------------
2224
static int
2325
Lock_init(Lock* self, PyObject* args, PyObject* kwargs)
2426
{
25-
// Get the reentrant argument
26-
static const char* kwlist[] = { "reentrant", NULL };
27-
PyObject* reentrant = NULL;
28-
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O", (char**)kwlist, &reentrant)) {
29-
return -1;
30-
}
31-
32-
// if reentrant was requested use a recursive mutex
33-
if (PyObject_IsTrue(reentrant)) {
34-
self->_rmutex = std::make_unique<std::recursive_timed_mutex>();
35-
} else {
36-
self->_mutex = std::make_unique<std::timed_mutex>();
37-
}
27+
self->_mutex = std::make_unique<std::timed_mutex>();
3828

3929
return 0;
4030
}
@@ -44,7 +34,6 @@ static void
4434
Lock_dealloc(Lock* self)
4535
{
4636
self->_mutex = nullptr;
47-
self->_rmutex = nullptr;
4837

4938
Py_TYPE(self)->tp_free((PyObject*)self);
5039
}
@@ -60,8 +49,12 @@ Lock_acquire(Lock* self, PyObject* args, PyObject* kwargs)
6049
return NULL;
6150
}
6251

63-
double timeout_value = 0.0;
64-
if (timeout != Py_None) {
52+
if (timeout == Py_None) {
53+
AllowThreads _;
54+
55+
self->_mutex->lock();
56+
} else {
57+
double timeout_value = 0.0;
6558
if (PyFloat_Check(timeout)) {
6659
timeout_value = PyFloat_AsDouble(timeout);
6760
} else if (PyLong_Check(timeout)) {
@@ -70,32 +63,16 @@ Lock_acquire(Lock* self, PyObject* args, PyObject* kwargs)
7063
PyErr_SetString(PyExc_TypeError, "timeout must be a float or an int");
7164
return NULL;
7265
}
73-
}
7466

75-
if (self->_mutex != nullptr) {
7667
AllowThreads _;
7768

78-
if (timeout == Py_None) {
79-
80-
self->_mutex->lock();
81-
} else if (!self->_mutex->try_lock_for(std::chrono::milliseconds((long long)(timeout_value * 1000)))) {
82-
Py_RETURN_FALSE;
83-
}
84-
self->_locked = 1;
85-
} else if (self->_rmutex != nullptr) {
86-
AllowThreads _;
87-
88-
if (timeout == Py_None) {
89-
self->_rmutex->lock();
90-
} else if (!self->_rmutex->try_lock_for(std::chrono::milliseconds((long long)(timeout_value * 1000)))) {
69+
if (!self->_mutex->try_lock_for(std::chrono::milliseconds((long long)(timeout_value * 1000)))) {
9170
Py_RETURN_FALSE;
9271
}
93-
self->_locked++;
94-
} else {
95-
PyErr_SetString(PyExc_RuntimeError, "Lock not initialized");
96-
return NULL;
9772
}
9873

74+
self->_locked = 1;
75+
9976
Py_RETURN_TRUE;
10077
}
10178

@@ -108,16 +85,8 @@ Lock_release(Lock* self)
10885
return NULL;
10986
}
11087

111-
if (self->_mutex != nullptr) {
112-
self->_mutex->unlock();
113-
self->_locked = 0; // Reset the lock state
114-
} else if (self->_rmutex != nullptr) {
115-
self->_rmutex->unlock();
116-
self->_locked--; // Decrement the lock count for reentrant locks
117-
} else {
118-
PyErr_SetString(PyExc_RuntimeError, "Lock not initialized");
119-
return NULL;
120-
}
88+
self->_mutex->unlock();
89+
self->_locked = 0; // Reset the lock state
12190

12291
Py_RETURN_NONE;
12392
}
@@ -133,11 +102,39 @@ Lock_locked(Lock* self)
133102
Py_RETURN_FALSE;
134103
}
135104

105+
// ----------------------------------------------------------------------------
106+
static PyObject*
107+
Lock_enter(Lock* self, PyObject* args, PyObject* kwargs)
108+
{
109+
110+
AllowThreads _;
111+
112+
self->_mutex->lock();
113+
114+
self->_locked = 1;
115+
116+
Py_RETURN_NONE;
117+
}
118+
119+
// ----------------------------------------------------------------------------
120+
static PyObject*
121+
Lock_exit(Lock* self, PyObject* args, PyObject* kwargs)
122+
{
123+
// This method is called when the lock is used in a "with" statement
124+
if (Lock_release(self) == NULL) {
125+
return NULL; // Propagate any error from release
126+
}
127+
128+
return Py_False;
129+
}
130+
136131
// ----------------------------------------------------------------------------
137132
static PyMethodDef Lock_methods[] = {
138133
{ "acquire", (PyCFunction)Lock_acquire, METH_VARARGS | METH_KEYWORDS, "Acquire the lock with an optional timeout" },
139134
{ "release", (PyCFunction)Lock_release, METH_NOARGS, "Release the lock" },
140135
{ "locked", (PyCFunction)Lock_locked, METH_NOARGS, "Return whether the lock is acquired" },
136+
{ "__enter__", (PyCFunction)Lock_enter, METH_NOARGS, "Enter the lock context" },
137+
{ "__exit__", (PyCFunction)Lock_exit, METH_VARARGS | METH_KEYWORDS, "Exit the lock context" },
141138
{ NULL } /* Sentinel */
142139
};
143140

@@ -159,3 +156,156 @@ static PyTypeObject LockType = {
159156
.tp_init = (initproc)Lock_init,
160157
.tp_new = PyType_GenericNew,
161158
};
159+
160+
// ----------------------------------------------------------------------------
161+
// RLock class
162+
// ----------------------------------------------------------------------------
163+
164+
typedef struct rlock
165+
{
166+
PyObject_HEAD
167+
168+
std::atomic<int>
169+
_locked = 0;
170+
171+
std::unique_ptr<std::recursive_timed_mutex> _mutex = nullptr;
172+
} RLock;
173+
174+
// ----------------------------------------------------------------------------
175+
static int
176+
RLock_init(RLock* self, PyObject* args, PyObject* kwargs)
177+
{
178+
self->_mutex = std::make_unique<std::recursive_timed_mutex>();
179+
180+
return 0;
181+
}
182+
183+
// ----------------------------------------------------------------------------
184+
static void
185+
RLock_dealloc(RLock* self)
186+
{
187+
self->_mutex = nullptr;
188+
189+
Py_TYPE(self)->tp_free((PyObject*)self);
190+
}
191+
192+
// ----------------------------------------------------------------------------
193+
static PyObject*
194+
RLock_acquire(RLock* self, PyObject* args, PyObject* kwargs)
195+
{
196+
// Get timeout argument
197+
static const char* kwlist[] = { "timeout", NULL };
198+
PyObject* timeout = Py_None;
199+
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O", (char**)kwlist, &timeout)) {
200+
return NULL;
201+
}
202+
203+
if (timeout == Py_None) {
204+
AllowThreads _;
205+
206+
self->_mutex->lock();
207+
} else {
208+
double timeout_value = 0.0;
209+
if (PyFloat_Check(timeout)) {
210+
timeout_value = PyFloat_AsDouble(timeout);
211+
} else if (PyLong_Check(timeout)) {
212+
timeout_value = PyLong_AsDouble(timeout);
213+
} else {
214+
PyErr_SetString(PyExc_TypeError, "timeout must be a float or an int");
215+
return NULL;
216+
}
217+
218+
AllowThreads _;
219+
220+
if (!self->_mutex->try_lock_for(std::chrono::milliseconds((long long)(timeout_value * 1000)))) {
221+
Py_RETURN_FALSE;
222+
}
223+
}
224+
225+
self->_locked++;
226+
227+
Py_RETURN_TRUE;
228+
}
229+
230+
// ----------------------------------------------------------------------------
231+
static PyObject*
232+
RLock_release(RLock* self)
233+
{
234+
if (self->_locked <= 0) {
235+
PyErr_SetString(PyExc_RuntimeError, "Lock is not acquired");
236+
return NULL;
237+
}
238+
239+
self->_mutex->unlock();
240+
self->_locked--;
241+
242+
Py_RETURN_NONE;
243+
}
244+
245+
// ----------------------------------------------------------------------------
246+
static PyObject*
247+
RLock_locked(RLock* self)
248+
{
249+
if (self->_locked > 0) {
250+
Py_RETURN_TRUE;
251+
}
252+
253+
Py_RETURN_FALSE;
254+
}
255+
256+
// ----------------------------------------------------------------------------
257+
static PyObject*
258+
RLock_enter(RLock* self, PyObject* args, PyObject* kwargs)
259+
{
260+
AllowThreads _;
261+
262+
self->_mutex->lock();
263+
264+
self->_locked++;
265+
266+
Py_RETURN_NONE;
267+
}
268+
269+
// ----------------------------------------------------------------------------
270+
static PyObject*
271+
RLock_exit(RLock* self, PyObject* args, PyObject* kwargs)
272+
{
273+
// This method is called when the lock is used in a "with" statement
274+
if (RLock_release(self) == NULL) {
275+
return NULL; // Propagate any error from release
276+
}
277+
278+
return Py_False;
279+
}
280+
281+
// ----------------------------------------------------------------------------
282+
static PyMethodDef RLock_methods[] = {
283+
{ "acquire",
284+
(PyCFunction)RLock_acquire,
285+
METH_VARARGS | METH_KEYWORDS,
286+
"Acquire the lock with an optional timeout" },
287+
{ "release", (PyCFunction)RLock_release, METH_NOARGS, "Release the lock" },
288+
{ "locked", (PyCFunction)RLock_locked, METH_NOARGS, "Return whether the lock is acquired at least once" },
289+
{ "__enter__", (PyCFunction)RLock_enter, METH_NOARGS, "Enter the lock context" },
290+
{ "__exit__", (PyCFunction)RLock_exit, METH_VARARGS | METH_KEYWORDS, "Exit the lock context" },
291+
{ NULL } /* Sentinel */
292+
};
293+
294+
// ----------------------------------------------------------------------------
295+
static PyMemberDef RLock_members[] = {
296+
{ NULL } /* Sentinel */
297+
};
298+
299+
// ----------------------------------------------------------------------------
300+
static PyTypeObject RLockType = {
301+
.ob_base = PyVarObject_HEAD_INIT(NULL, 0).tp_name = "ddtrace.internal._threads.RLock",
302+
.tp_basicsize = sizeof(RLock),
303+
.tp_itemsize = 0,
304+
.tp_dealloc = (destructor)RLock_dealloc,
305+
.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
306+
.tp_doc = PyDoc_STR("Native re-entrant lock implementation"),
307+
.tp_methods = RLock_methods,
308+
.tp_members = RLock_members,
309+
.tp_init = (initproc)RLock_init,
310+
.tp_new = PyType_GenericNew,
311+
};

0 commit comments

Comments
 (0)