diff --git a/Include/internal/pycore_lock.h b/Include/internal/pycore_lock.h index 03ad1c9fd58..18a8896d97a 100644 --- a/Include/internal/pycore_lock.h +++ b/Include/internal/pycore_lock.h @@ -213,6 +213,45 @@ _PyOnceFlag_CallOnce(_PyOnceFlag *flag, _Py_once_fn_t *fn, void *arg) return _PyOnceFlag_CallOnceSlow(flag, fn, arg); } +// A readers-writer (RW) lock. The lock supports multiple concurrent readers or +// a single writer. The lock is write-preferring: if a writer is waiting while +// the lock is read-locked then, new readers will be blocked. This avoids +// starvation of writers. +// +// In C++, the equivalent synchronization primitive is std::shared_mutex +// with shared ("read") and exclusive ("write") locking. +// +// The two least significant bits are used to indicate if the lock is +// write-locked and if there are parked threads (either readers or writers) +// waiting to acquire the lock. The remaining bits are used to indicate the +// number of readers holding the lock. +// +// 0b000..00000: unlocked +// 0bnnn..nnn00: nnn..nnn readers holding the lock +// 0bnnn..nnn10: nnn..nnn readers holding the lock and a writer is waiting +// 0b00000..010: unlocked with awoken writer about to acquire lock +// 0b00000..001: write-locked +// 0b00000..011: write-locked and readers or other writers are waiting +// +// Note that reader_count must be zero if the lock is held by a writer, and +// vice versa. The lock can only be held by readers or a writer, but not both. +// +// The design is optimized for simplicity of the implementation. The lock is +// not fair: if fairness is desired, use an additional PyMutex to serialize +// writers. The lock is also not reentrant. +typedef struct { + uintptr_t bits; +} _PyRWMutex; + +// Read lock (i.e., shared lock) +PyAPI_FUNC(void) _PyRWMutex_RLock(_PyRWMutex *rwmutex); +PyAPI_FUNC(void) _PyRWMutex_RUnlock(_PyRWMutex *rwmutex); + +// Write lock (i.e., exclusive lock) +PyAPI_FUNC(void) _PyRWMutex_Lock(_PyRWMutex *rwmutex); +PyAPI_FUNC(void) _PyRWMutex_Unlock(_PyRWMutex *rwmutex); + + #ifdef __cplusplus } #endif diff --git a/Modules/_testinternalcapi/test_lock.c b/Modules/_testinternalcapi/test_lock.c index 418f71c1441..83081f73a72 100644 --- a/Modules/_testinternalcapi/test_lock.c +++ b/Modules/_testinternalcapi/test_lock.c @@ -372,6 +372,104 @@ test_lock_once(PyObject *self, PyObject *obj) Py_RETURN_NONE; } +struct test_rwlock_data { + Py_ssize_t nthreads; + _PyRWMutex rw; + PyEvent step1; + PyEvent step2; + PyEvent step3; + PyEvent done; +}; + +static void +rdlock_thread(void *arg) +{ + struct test_rwlock_data *test_data = arg; + + // Acquire the lock in read mode + _PyRWMutex_RLock(&test_data->rw); + PyEvent_Wait(&test_data->step1); + _PyRWMutex_RUnlock(&test_data->rw); + + _PyRWMutex_RLock(&test_data->rw); + PyEvent_Wait(&test_data->step3); + _PyRWMutex_RUnlock(&test_data->rw); + + if (_Py_atomic_add_ssize(&test_data->nthreads, -1) == 1) { + _PyEvent_Notify(&test_data->done); + } +} +static void +wrlock_thread(void *arg) +{ + struct test_rwlock_data *test_data = arg; + + // First acquire the lock in write mode + _PyRWMutex_Lock(&test_data->rw); + PyEvent_Wait(&test_data->step2); + _PyRWMutex_Unlock(&test_data->rw); + + if (_Py_atomic_add_ssize(&test_data->nthreads, -1) == 1) { + _PyEvent_Notify(&test_data->done); + } +} + +static void +wait_until(uintptr_t *ptr, uintptr_t value) +{ + // wait up to two seconds for *ptr == value + int iters = 0; + uintptr_t bits; + do { + pysleep(10); + bits = _Py_atomic_load_uintptr(ptr); + iters++; + } while (bits != value && iters < 200); +} + +static PyObject * +test_lock_rwlock(PyObject *self, PyObject *obj) +{ + struct test_rwlock_data test_data = {.nthreads = 3}; + + _PyRWMutex_Lock(&test_data.rw); + assert(test_data.rw.bits == 1); + + _PyRWMutex_Unlock(&test_data.rw); + assert(test_data.rw.bits == 0); + + // Start two readers + PyThread_start_new_thread(rdlock_thread, &test_data); + PyThread_start_new_thread(rdlock_thread, &test_data); + + // wait up to two seconds for the threads to attempt to read-lock "rw" + wait_until(&test_data.rw.bits, 8); + assert(test_data.rw.bits == 8); + + // start writer (while readers hold lock) + PyThread_start_new_thread(wrlock_thread, &test_data); + wait_until(&test_data.rw.bits, 10); + assert(test_data.rw.bits == 10); + + // readers release lock, writer should acquire it + _PyEvent_Notify(&test_data.step1); + wait_until(&test_data.rw.bits, 3); + assert(test_data.rw.bits == 3); + + // writer releases lock, readers acquire it + _PyEvent_Notify(&test_data.step2); + wait_until(&test_data.rw.bits, 8); + assert(test_data.rw.bits == 8); + + // readers release lock again + _PyEvent_Notify(&test_data.step3); + wait_until(&test_data.rw.bits, 0); + assert(test_data.rw.bits == 0); + + PyEvent_Wait(&test_data.done); + Py_RETURN_NONE; +} + static PyMethodDef test_methods[] = { {"test_lock_basic", test_lock_basic, METH_NOARGS}, {"test_lock_two_threads", test_lock_two_threads, METH_NOARGS}, @@ -380,6 +478,7 @@ static PyMethodDef test_methods[] = { _TESTINTERNALCAPI_BENCHMARK_LOCKS_METHODDEF {"test_lock_benchmark", test_lock_benchmark, METH_NOARGS}, {"test_lock_once", test_lock_once, METH_NOARGS}, + {"test_lock_rwlock", test_lock_rwlock, METH_NOARGS}, {NULL, NULL} /* sentinel */ }; diff --git a/Python/lock.c b/Python/lock.c index e9279f0b92a..f0ff1176941 100644 --- a/Python/lock.c +++ b/Python/lock.c @@ -353,3 +353,109 @@ _PyOnceFlag_CallOnceSlow(_PyOnceFlag *flag, _Py_once_fn_t *fn, void *arg) v = _Py_atomic_load_uint8(&flag->v); } } + +#define _Py_WRITE_LOCKED 1 +#define _PyRWMutex_READER_SHIFT 2 +#define _Py_RWMUTEX_MAX_READERS (UINTPTR_MAX >> _PyRWMutex_READER_SHIFT) + +static uintptr_t +rwmutex_set_parked_and_wait(_PyRWMutex *rwmutex, uintptr_t bits) +{ + // Set _Py_HAS_PARKED and wait until we are woken up. + if ((bits & _Py_HAS_PARKED) == 0) { + uintptr_t newval = bits | _Py_HAS_PARKED; + if (!_Py_atomic_compare_exchange_uintptr(&rwmutex->bits, + &bits, newval)) { + return bits; + } + bits = newval; + } + + _PyParkingLot_Park(&rwmutex->bits, &bits, sizeof(bits), -1, NULL, 1); + return _Py_atomic_load_uintptr_relaxed(&rwmutex->bits); +} + +// The number of readers holding the lock +static uintptr_t +rwmutex_reader_count(uintptr_t bits) +{ + return bits >> _PyRWMutex_READER_SHIFT; +} + +void +_PyRWMutex_RLock(_PyRWMutex *rwmutex) +{ + uintptr_t bits = _Py_atomic_load_uintptr_relaxed(&rwmutex->bits); + for (;;) { + if ((bits & _Py_WRITE_LOCKED)) { + // A writer already holds the lock. + bits = rwmutex_set_parked_and_wait(rwmutex, bits); + continue; + } + else if ((bits & _Py_HAS_PARKED)) { + // Reader(s) hold the lock (or just gave up the lock), but there is + // at least one waiting writer. We can't grab the lock because we + // don't want to starve the writer. Instead, we park ourselves and + // wait for the writer to eventually wake us up. + bits = rwmutex_set_parked_and_wait(rwmutex, bits); + continue; + } + else { + // The lock is unlocked or read-locked. Try to grab it. + assert(rwmutex_reader_count(bits) < _Py_RWMUTEX_MAX_READERS); + uintptr_t newval = bits + (1 << _PyRWMutex_READER_SHIFT); + if (!_Py_atomic_compare_exchange_uintptr(&rwmutex->bits, + &bits, newval)) { + continue; + } + return; + } + } +} + +void +_PyRWMutex_RUnlock(_PyRWMutex *rwmutex) +{ + uintptr_t bits = _Py_atomic_add_uintptr(&rwmutex->bits, -(1 << _PyRWMutex_READER_SHIFT)); + assert(rwmutex_reader_count(bits) > 0 && "lock was not read-locked"); + bits -= (1 << _PyRWMutex_READER_SHIFT); + + if (rwmutex_reader_count(bits) == 0 && (bits & _Py_HAS_PARKED)) { + _PyParkingLot_UnparkAll(&rwmutex->bits); + return; + } +} + +void +_PyRWMutex_Lock(_PyRWMutex *rwmutex) +{ + uintptr_t bits = _Py_atomic_load_uintptr_relaxed(&rwmutex->bits); + for (;;) { + // If there are no active readers and it's not already write-locked, + // then we can grab the lock. + if ((bits & ~_Py_HAS_PARKED) == 0) { + if (!_Py_atomic_compare_exchange_uintptr(&rwmutex->bits, + &bits, + bits | _Py_WRITE_LOCKED)) { + continue; + } + return; + } + + // Otherwise, we have to wait. + bits = rwmutex_set_parked_and_wait(rwmutex, bits); + } +} + +void +_PyRWMutex_Unlock(_PyRWMutex *rwmutex) +{ + uintptr_t old_bits = _Py_atomic_exchange_uintptr(&rwmutex->bits, 0); + + assert((old_bits & _Py_WRITE_LOCKED) && "lock was not write-locked"); + assert(rwmutex_reader_count(old_bits) == 0 && "lock was read-locked"); + + if ((old_bits & _Py_HAS_PARKED) != 0) { + _PyParkingLot_UnparkAll(&rwmutex->bits); + } +}