Skip to content

Commit

Permalink
_threadmodule: thread-safety fixes
Browse files Browse the repository at this point in the history
 - Make rlock thread-safe
 - Use atomics to modify interp->num_threads
  • Loading branch information
colesbury committed Apr 23, 2023
1 parent 3cfbc49 commit 5722416
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 28 deletions.
2 changes: 1 addition & 1 deletion Include/internal/pycore_interp.h
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ struct _is {
/* The linked list of threads, newest first. */
PyThreadState *head;
/* Used in Modules/_threadmodule.c. */
long count;
Py_ssize_t count;
/* Support for runtime thread stack size tuning.
A value of 0 means using the platform's default stack size
or the size specified by the THREAD_STACK_SIZE macro. */
Expand Down
57 changes: 30 additions & 27 deletions Modules/_threadmodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -319,7 +319,7 @@ static PyType_Spec lock_type_spec = {
typedef struct {
PyObject_HEAD
PyThread_type_lock rlock_lock;
unsigned long rlock_owner;
uintptr_t rlock_owner;
unsigned long rlock_count;
PyObject *in_weakreflist;
} rlockobject;
Expand All @@ -331,7 +331,6 @@ rlock_traverse(rlockobject *self, visitproc visit, void *arg)
return 0;
}


static void
rlock_dealloc(rlockobject *self)
{
Expand All @@ -352,18 +351,24 @@ rlock_dealloc(rlockobject *self)
Py_DECREF(tp);
}

static int
rlock_is_owned(rlockobject *self)
{
uintptr_t tid = _Py_ThreadId();
uintptr_t owner_tid = _Py_atomic_load_uintptr_relaxed(&self->rlock_owner);
return owner_tid == tid && self->rlock_count > 0;
}

static PyObject *
rlock_acquire(rlockobject *self, PyObject *args, PyObject *kwds)
{
_PyTime_t timeout;
unsigned long tid;
PyLockStatus r = PY_LOCK_ACQUIRED;

if (lock_acquire_parse_args(args, kwds, &timeout) < 0)
return NULL;

tid = PyThread_get_thread_ident();
if (self->rlock_count > 0 && tid == self->rlock_owner) {
if (rlock_is_owned(self)) {
unsigned long count = self->rlock_count + 1;
if (count <= self->rlock_count) {
PyErr_SetString(PyExc_OverflowError,
Expand All @@ -376,7 +381,7 @@ rlock_acquire(rlockobject *self, PyObject *args, PyObject *kwds)
r = acquire_timed(self->rlock_lock, timeout);
if (r == PY_LOCK_ACQUIRED) {
assert(self->rlock_count == 0);
self->rlock_owner = tid;
_Py_atomic_store_uintptr_relaxed(&self->rlock_owner, _Py_ThreadId());
self->rlock_count = 1;
}
else if (r == PY_LOCK_INTR) {
Expand Down Expand Up @@ -405,15 +410,13 @@ the lock is taken and its internal counter initialized to 1.");
static PyObject *
rlock_release(rlockobject *self, PyObject *Py_UNUSED(ignored))
{
unsigned long tid = PyThread_get_thread_ident();

if (self->rlock_count == 0 || self->rlock_owner != tid) {
if (!rlock_is_owned(self)) {
PyErr_SetString(PyExc_RuntimeError,
"cannot release un-acquired lock");
"cannot release un-acquired lock1");
return NULL;
}
if (--self->rlock_count == 0) {
self->rlock_owner = 0;
_Py_atomic_store_uintptr_relaxed(&self->rlock_owner, 0);
PyThread_release_lock(self->rlock_lock);
}
Py_RETURN_NONE;
Expand All @@ -434,11 +437,11 @@ to be available for other threads.");
static PyObject *
rlock_acquire_restore(rlockobject *self, PyObject *args)
{
unsigned long owner;
unsigned long long owner;
unsigned long count;
int r = 1;

if (!PyArg_ParseTuple(args, "(kk):_acquire_restore", &count, &owner))
if (!PyArg_ParseTuple(args, "(kK):_acquire_restore", &count, &owner))
return NULL;

if (!PyThread_acquire_lock(self->rlock_lock, 0)) {
Expand All @@ -451,7 +454,7 @@ rlock_acquire_restore(rlockobject *self, PyObject *args)
return NULL;
}
assert(self->rlock_count == 0);
self->rlock_owner = owner;
_Py_atomic_store_uintptr_relaxed(&self->rlock_owner, (uintptr_t)owner);
self->rlock_count = count;
Py_RETURN_NONE;
}
Expand All @@ -464,7 +467,7 @@ For internal use by `threading.Condition`.");
static PyObject *
rlock_release_save(rlockobject *self, PyObject *Py_UNUSED(ignored))
{
unsigned long owner;
uintptr_t owner;
unsigned long count;

if (self->rlock_count == 0) {
Expand All @@ -476,9 +479,9 @@ rlock_release_save(rlockobject *self, PyObject *Py_UNUSED(ignored))
owner = self->rlock_owner;
count = self->rlock_count;
self->rlock_count = 0;
self->rlock_owner = 0;
_Py_atomic_store_uintptr_relaxed(&self->rlock_owner, 0);
PyThread_release_lock(self->rlock_lock);
return Py_BuildValue("kk", count, owner);
return Py_BuildValue("kK", count, (unsigned long long)owner);
}

PyDoc_STRVAR(rlock_release_save_doc,
Expand All @@ -488,11 +491,9 @@ For internal use by `threading.Condition`.");


static PyObject *
rlock_is_owned(rlockobject *self, PyObject *Py_UNUSED(ignored))
rlock__is_owned(rlockobject *self, PyObject *Py_UNUSED(ignored))
{
unsigned long tid = PyThread_get_thread_ident();

if (self->rlock_count > 0 && self->rlock_owner == tid) {
if (rlock_is_owned(self)) {
Py_RETURN_TRUE;
}
Py_RETURN_FALSE;
Expand Down Expand Up @@ -526,9 +527,11 @@ rlock_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
static PyObject *
rlock_repr(rlockobject *self)
{
return PyUnicode_FromFormat("<%s %s object owner=%ld count=%lu at %p>",
uintptr_t owner = _Py_atomic_load_uintptr_relaxed(&self->rlock_owner);

return PyUnicode_FromFormat("<%s %s object owner=%p count=%llu at %p>",
self->rlock_count ? "locked" : "unlocked",
Py_TYPE(self)->tp_name, self->rlock_owner,
Py_TYPE(self)->tp_name, (const void *)owner,
self->rlock_count, self);
}

Expand All @@ -555,7 +558,7 @@ static PyMethodDef rlock_methods[] = {
METH_VARARGS | METH_KEYWORDS, rlock_acquire_doc},
{"release", (PyCFunction)rlock_release,
METH_NOARGS, rlock_release_doc},
{"_is_owned", (PyCFunction)rlock_is_owned,
{"_is_owned", (PyCFunction)rlock__is_owned,
METH_NOARGS, rlock_is_owned_doc},
{"_acquire_restore", (PyCFunction)rlock_acquire_restore,
METH_VARARGS, rlock_acquire_restore_doc},
Expand Down Expand Up @@ -1272,7 +1275,7 @@ thread_run(void *boot_raw)
#endif
_PyThreadState_SetCurrent(tstate);
PyEval_AcquireThread(tstate);
tstate->interp->threads.count++;
_Py_atomic_add_ssize(&tstate->interp->threads.count, 1);

PyObject *res = PyObject_Call(boot->func, boot->args, boot->kwargs);
if (res == NULL) {
Expand All @@ -1288,7 +1291,7 @@ thread_run(void *boot_raw)
}

thread_bootstate_free(boot);
tstate->interp->threads.count--;
_Py_atomic_add_ssize(&tstate->interp->threads.count, -1);
PyThreadState_Clear(tstate);
_PyThreadState_DeleteCurrent(tstate);

Expand Down Expand Up @@ -1531,7 +1534,7 @@ static PyObject *
thread__count(PyObject *self, PyObject *Py_UNUSED(ignored))
{
PyInterpreterState *interp = _PyInterpreterState_GET();
return PyLong_FromLong(interp->threads.count);
return PyLong_FromSize_t(_Py_atomic_load_ssize(&interp->threads.count));
}

PyDoc_STRVAR(_count_doc,
Expand Down

0 comments on commit 5722416

Please sign in to comment.