Skip to content

Commit 5ae75e1

Browse files
authored
gh-111964: Add _PyRWMutex a "readers-writer" lock (gh-112859)
This adds `_PyRWMutex`, a "readers-writer" lock, which wil be used to serialize global stop-the-world pauses with per-interpreter pauses.
1 parent 40574da commit 5ae75e1

File tree

3 files changed

+244
-0
lines changed

3 files changed

+244
-0
lines changed

Include/internal/pycore_lock.h

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,45 @@ _PyOnceFlag_CallOnce(_PyOnceFlag *flag, _Py_once_fn_t *fn, void *arg)
213213
return _PyOnceFlag_CallOnceSlow(flag, fn, arg);
214214
}
215215

216+
// A readers-writer (RW) lock. The lock supports multiple concurrent readers or
217+
// a single writer. The lock is write-preferring: if a writer is waiting while
218+
// the lock is read-locked then, new readers will be blocked. This avoids
219+
// starvation of writers.
220+
//
221+
// In C++, the equivalent synchronization primitive is std::shared_mutex
222+
// with shared ("read") and exclusive ("write") locking.
223+
//
224+
// The two least significant bits are used to indicate if the lock is
225+
// write-locked and if there are parked threads (either readers or writers)
226+
// waiting to acquire the lock. The remaining bits are used to indicate the
227+
// number of readers holding the lock.
228+
//
229+
// 0b000..00000: unlocked
230+
// 0bnnn..nnn00: nnn..nnn readers holding the lock
231+
// 0bnnn..nnn10: nnn..nnn readers holding the lock and a writer is waiting
232+
// 0b00000..010: unlocked with awoken writer about to acquire lock
233+
// 0b00000..001: write-locked
234+
// 0b00000..011: write-locked and readers or other writers are waiting
235+
//
236+
// Note that reader_count must be zero if the lock is held by a writer, and
237+
// vice versa. The lock can only be held by readers or a writer, but not both.
238+
//
239+
// The design is optimized for simplicity of the implementation. The lock is
240+
// not fair: if fairness is desired, use an additional PyMutex to serialize
241+
// writers. The lock is also not reentrant.
242+
typedef struct {
243+
uintptr_t bits;
244+
} _PyRWMutex;
245+
246+
// Read lock (i.e., shared lock)
247+
PyAPI_FUNC(void) _PyRWMutex_RLock(_PyRWMutex *rwmutex);
248+
PyAPI_FUNC(void) _PyRWMutex_RUnlock(_PyRWMutex *rwmutex);
249+
250+
// Write lock (i.e., exclusive lock)
251+
PyAPI_FUNC(void) _PyRWMutex_Lock(_PyRWMutex *rwmutex);
252+
PyAPI_FUNC(void) _PyRWMutex_Unlock(_PyRWMutex *rwmutex);
253+
254+
216255
#ifdef __cplusplus
217256
}
218257
#endif

Modules/_testinternalcapi/test_lock.c

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,6 +372,104 @@ test_lock_once(PyObject *self, PyObject *obj)
372372
Py_RETURN_NONE;
373373
}
374374

375+
struct test_rwlock_data {
376+
Py_ssize_t nthreads;
377+
_PyRWMutex rw;
378+
PyEvent step1;
379+
PyEvent step2;
380+
PyEvent step3;
381+
PyEvent done;
382+
};
383+
384+
static void
385+
rdlock_thread(void *arg)
386+
{
387+
struct test_rwlock_data *test_data = arg;
388+
389+
// Acquire the lock in read mode
390+
_PyRWMutex_RLock(&test_data->rw);
391+
PyEvent_Wait(&test_data->step1);
392+
_PyRWMutex_RUnlock(&test_data->rw);
393+
394+
_PyRWMutex_RLock(&test_data->rw);
395+
PyEvent_Wait(&test_data->step3);
396+
_PyRWMutex_RUnlock(&test_data->rw);
397+
398+
if (_Py_atomic_add_ssize(&test_data->nthreads, -1) == 1) {
399+
_PyEvent_Notify(&test_data->done);
400+
}
401+
}
402+
static void
403+
wrlock_thread(void *arg)
404+
{
405+
struct test_rwlock_data *test_data = arg;
406+
407+
// First acquire the lock in write mode
408+
_PyRWMutex_Lock(&test_data->rw);
409+
PyEvent_Wait(&test_data->step2);
410+
_PyRWMutex_Unlock(&test_data->rw);
411+
412+
if (_Py_atomic_add_ssize(&test_data->nthreads, -1) == 1) {
413+
_PyEvent_Notify(&test_data->done);
414+
}
415+
}
416+
417+
static void
418+
wait_until(uintptr_t *ptr, uintptr_t value)
419+
{
420+
// wait up to two seconds for *ptr == value
421+
int iters = 0;
422+
uintptr_t bits;
423+
do {
424+
pysleep(10);
425+
bits = _Py_atomic_load_uintptr(ptr);
426+
iters++;
427+
} while (bits != value && iters < 200);
428+
}
429+
430+
static PyObject *
431+
test_lock_rwlock(PyObject *self, PyObject *obj)
432+
{
433+
struct test_rwlock_data test_data = {.nthreads = 3};
434+
435+
_PyRWMutex_Lock(&test_data.rw);
436+
assert(test_data.rw.bits == 1);
437+
438+
_PyRWMutex_Unlock(&test_data.rw);
439+
assert(test_data.rw.bits == 0);
440+
441+
// Start two readers
442+
PyThread_start_new_thread(rdlock_thread, &test_data);
443+
PyThread_start_new_thread(rdlock_thread, &test_data);
444+
445+
// wait up to two seconds for the threads to attempt to read-lock "rw"
446+
wait_until(&test_data.rw.bits, 8);
447+
assert(test_data.rw.bits == 8);
448+
449+
// start writer (while readers hold lock)
450+
PyThread_start_new_thread(wrlock_thread, &test_data);
451+
wait_until(&test_data.rw.bits, 10);
452+
assert(test_data.rw.bits == 10);
453+
454+
// readers release lock, writer should acquire it
455+
_PyEvent_Notify(&test_data.step1);
456+
wait_until(&test_data.rw.bits, 3);
457+
assert(test_data.rw.bits == 3);
458+
459+
// writer releases lock, readers acquire it
460+
_PyEvent_Notify(&test_data.step2);
461+
wait_until(&test_data.rw.bits, 8);
462+
assert(test_data.rw.bits == 8);
463+
464+
// readers release lock again
465+
_PyEvent_Notify(&test_data.step3);
466+
wait_until(&test_data.rw.bits, 0);
467+
assert(test_data.rw.bits == 0);
468+
469+
PyEvent_Wait(&test_data.done);
470+
Py_RETURN_NONE;
471+
}
472+
375473
static PyMethodDef test_methods[] = {
376474
{"test_lock_basic", test_lock_basic, METH_NOARGS},
377475
{"test_lock_two_threads", test_lock_two_threads, METH_NOARGS},
@@ -380,6 +478,7 @@ static PyMethodDef test_methods[] = {
380478
_TESTINTERNALCAPI_BENCHMARK_LOCKS_METHODDEF
381479
{"test_lock_benchmark", test_lock_benchmark, METH_NOARGS},
382480
{"test_lock_once", test_lock_once, METH_NOARGS},
481+
{"test_lock_rwlock", test_lock_rwlock, METH_NOARGS},
383482
{NULL, NULL} /* sentinel */
384483
};
385484

Python/lock.c

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -353,3 +353,109 @@ _PyOnceFlag_CallOnceSlow(_PyOnceFlag *flag, _Py_once_fn_t *fn, void *arg)
353353
v = _Py_atomic_load_uint8(&flag->v);
354354
}
355355
}
356+
357+
#define _Py_WRITE_LOCKED 1
358+
#define _PyRWMutex_READER_SHIFT 2
359+
#define _Py_RWMUTEX_MAX_READERS (UINTPTR_MAX >> _PyRWMutex_READER_SHIFT)
360+
361+
static uintptr_t
362+
rwmutex_set_parked_and_wait(_PyRWMutex *rwmutex, uintptr_t bits)
363+
{
364+
// Set _Py_HAS_PARKED and wait until we are woken up.
365+
if ((bits & _Py_HAS_PARKED) == 0) {
366+
uintptr_t newval = bits | _Py_HAS_PARKED;
367+
if (!_Py_atomic_compare_exchange_uintptr(&rwmutex->bits,
368+
&bits, newval)) {
369+
return bits;
370+
}
371+
bits = newval;
372+
}
373+
374+
_PyParkingLot_Park(&rwmutex->bits, &bits, sizeof(bits), -1, NULL, 1);
375+
return _Py_atomic_load_uintptr_relaxed(&rwmutex->bits);
376+
}
377+
378+
// The number of readers holding the lock
379+
static uintptr_t
380+
rwmutex_reader_count(uintptr_t bits)
381+
{
382+
return bits >> _PyRWMutex_READER_SHIFT;
383+
}
384+
385+
void
386+
_PyRWMutex_RLock(_PyRWMutex *rwmutex)
387+
{
388+
uintptr_t bits = _Py_atomic_load_uintptr_relaxed(&rwmutex->bits);
389+
for (;;) {
390+
if ((bits & _Py_WRITE_LOCKED)) {
391+
// A writer already holds the lock.
392+
bits = rwmutex_set_parked_and_wait(rwmutex, bits);
393+
continue;
394+
}
395+
else if ((bits & _Py_HAS_PARKED)) {
396+
// Reader(s) hold the lock (or just gave up the lock), but there is
397+
// at least one waiting writer. We can't grab the lock because we
398+
// don't want to starve the writer. Instead, we park ourselves and
399+
// wait for the writer to eventually wake us up.
400+
bits = rwmutex_set_parked_and_wait(rwmutex, bits);
401+
continue;
402+
}
403+
else {
404+
// The lock is unlocked or read-locked. Try to grab it.
405+
assert(rwmutex_reader_count(bits) < _Py_RWMUTEX_MAX_READERS);
406+
uintptr_t newval = bits + (1 << _PyRWMutex_READER_SHIFT);
407+
if (!_Py_atomic_compare_exchange_uintptr(&rwmutex->bits,
408+
&bits, newval)) {
409+
continue;
410+
}
411+
return;
412+
}
413+
}
414+
}
415+
416+
void
417+
_PyRWMutex_RUnlock(_PyRWMutex *rwmutex)
418+
{
419+
uintptr_t bits = _Py_atomic_add_uintptr(&rwmutex->bits, -(1 << _PyRWMutex_READER_SHIFT));
420+
assert(rwmutex_reader_count(bits) > 0 && "lock was not read-locked");
421+
bits -= (1 << _PyRWMutex_READER_SHIFT);
422+
423+
if (rwmutex_reader_count(bits) == 0 && (bits & _Py_HAS_PARKED)) {
424+
_PyParkingLot_UnparkAll(&rwmutex->bits);
425+
return;
426+
}
427+
}
428+
429+
void
430+
_PyRWMutex_Lock(_PyRWMutex *rwmutex)
431+
{
432+
uintptr_t bits = _Py_atomic_load_uintptr_relaxed(&rwmutex->bits);
433+
for (;;) {
434+
// If there are no active readers and it's not already write-locked,
435+
// then we can grab the lock.
436+
if ((bits & ~_Py_HAS_PARKED) == 0) {
437+
if (!_Py_atomic_compare_exchange_uintptr(&rwmutex->bits,
438+
&bits,
439+
bits | _Py_WRITE_LOCKED)) {
440+
continue;
441+
}
442+
return;
443+
}
444+
445+
// Otherwise, we have to wait.
446+
bits = rwmutex_set_parked_and_wait(rwmutex, bits);
447+
}
448+
}
449+
450+
void
451+
_PyRWMutex_Unlock(_PyRWMutex *rwmutex)
452+
{
453+
uintptr_t old_bits = _Py_atomic_exchange_uintptr(&rwmutex->bits, 0);
454+
455+
assert((old_bits & _Py_WRITE_LOCKED) && "lock was not write-locked");
456+
assert(rwmutex_reader_count(old_bits) == 0 && "lock was read-locked");
457+
458+
if ((old_bits & _Py_HAS_PARKED) != 0) {
459+
_PyParkingLot_UnparkAll(&rwmutex->bits);
460+
}
461+
}

0 commit comments

Comments
 (0)