1 // C Extension module to test pycore_lock.h API
2
3 #include "parts.h"
4 #include "pycore_lock.h"
5 #include "pycore_pythread.h" // PyThread_get_thread_ident_ex()
6
7 #include "clinic/test_lock.c.h"
8
9 #ifdef MS_WINDOWS
10 #define WIN32_LEAN_AND_MEAN
11 #include <windows.h>
12 #else
13 #include <unistd.h> // usleep()
14 #endif
15
16 /*[clinic input]
17 module _testinternalcapi
18 [clinic start generated code]*/
19 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=7bb583d8c9eb9a78]*/
20
21
22 static void
pysleep(int ms)23 pysleep(int ms)
24 {
25 #ifdef MS_WINDOWS
26 Sleep(ms);
27 #else
28 usleep(ms * 1000);
29 #endif
30 }
31
32 static PyObject *
test_lock_basic(PyObject * self,PyObject * obj)33 test_lock_basic(PyObject *self, PyObject *obj)
34 {
35 PyMutex m = (PyMutex){0};
36
37 // uncontended lock and unlock
38 PyMutex_Lock(&m);
39 assert(m._bits == 1);
40 PyMutex_Unlock(&m);
41 assert(m._bits == 0);
42
43 Py_RETURN_NONE;
44 }
45
46 struct test_lock2_data {
47 PyMutex m;
48 PyEvent done;
49 int started;
50 };
51
52 static void
lock_thread(void * arg)53 lock_thread(void *arg)
54 {
55 struct test_lock2_data *test_data = arg;
56 PyMutex *m = &test_data->m;
57 _Py_atomic_store_int(&test_data->started, 1);
58
59 PyMutex_Lock(m);
60 assert(m->_bits == 1);
61
62 PyMutex_Unlock(m);
63 assert(m->_bits == 0);
64
65 _PyEvent_Notify(&test_data->done);
66 }
67
68 static PyObject *
test_lock_two_threads(PyObject * self,PyObject * obj)69 test_lock_two_threads(PyObject *self, PyObject *obj)
70 {
71 // lock attempt by two threads
72 struct test_lock2_data test_data;
73 memset(&test_data, 0, sizeof(test_data));
74
75 PyMutex_Lock(&test_data.m);
76 assert(test_data.m._bits == 1);
77
78 PyThread_start_new_thread(lock_thread, &test_data);
79
80 // wait up to two seconds for the lock_thread to attempt to lock "m"
81 int iters = 0;
82 uint8_t v;
83 do {
84 pysleep(10); // allow some time for the other thread to try to lock
85 v = _Py_atomic_load_uint8_relaxed(&test_data.m._bits);
86 assert(v == 1 || v == 3);
87 iters++;
88 } while (v != 3 && iters < 200);
89
90 // both the "locked" and the "has parked" bits should be set
91 assert(test_data.m._bits == 3);
92
93 PyMutex_Unlock(&test_data.m);
94 PyEvent_Wait(&test_data.done);
95 assert(test_data.m._bits == 0);
96
97 Py_RETURN_NONE;
98 }
99
100 #define COUNTER_THREADS 5
101 #define COUNTER_ITERS 10000
102
103 struct test_data_counter {
104 PyMutex m;
105 Py_ssize_t counter;
106 };
107
108 struct thread_data_counter {
109 struct test_data_counter *test_data;
110 PyEvent done_event;
111 };
112
113 static void
counter_thread(void * arg)114 counter_thread(void *arg)
115 {
116 struct thread_data_counter *thread_data = arg;
117 struct test_data_counter *test_data = thread_data->test_data;
118
119 for (Py_ssize_t i = 0; i < COUNTER_ITERS; i++) {
120 PyMutex_Lock(&test_data->m);
121 test_data->counter++;
122 PyMutex_Unlock(&test_data->m);
123 }
124 _PyEvent_Notify(&thread_data->done_event);
125 }
126
127 static PyObject *
test_lock_counter(PyObject * self,PyObject * obj)128 test_lock_counter(PyObject *self, PyObject *obj)
129 {
130 // Test with rapidly locking and unlocking mutex
131 struct test_data_counter test_data;
132 memset(&test_data, 0, sizeof(test_data));
133
134 struct thread_data_counter thread_data[COUNTER_THREADS];
135 memset(&thread_data, 0, sizeof(thread_data));
136
137 for (Py_ssize_t i = 0; i < COUNTER_THREADS; i++) {
138 thread_data[i].test_data = &test_data;
139 PyThread_start_new_thread(counter_thread, &thread_data[i]);
140 }
141
142 for (Py_ssize_t i = 0; i < COUNTER_THREADS; i++) {
143 PyEvent_Wait(&thread_data[i].done_event);
144 }
145
146 assert(test_data.counter == COUNTER_THREADS * COUNTER_ITERS);
147 Py_RETURN_NONE;
148 }
149
150 #define SLOW_COUNTER_ITERS 100
151
152 static void
slow_counter_thread(void * arg)153 slow_counter_thread(void *arg)
154 {
155 struct thread_data_counter *thread_data = arg;
156 struct test_data_counter *test_data = thread_data->test_data;
157
158 for (Py_ssize_t i = 0; i < SLOW_COUNTER_ITERS; i++) {
159 PyMutex_Lock(&test_data->m);
160 if (i % 7 == 0) {
161 pysleep(2);
162 }
163 test_data->counter++;
164 PyMutex_Unlock(&test_data->m);
165 }
166 _PyEvent_Notify(&thread_data->done_event);
167 }
168
169 static PyObject *
test_lock_counter_slow(PyObject * self,PyObject * obj)170 test_lock_counter_slow(PyObject *self, PyObject *obj)
171 {
172 // Test lock/unlock with occasional "long" critical section, which will
173 // trigger handoff of the lock.
174 struct test_data_counter test_data;
175 memset(&test_data, 0, sizeof(test_data));
176
177 struct thread_data_counter thread_data[COUNTER_THREADS];
178 memset(&thread_data, 0, sizeof(thread_data));
179
180 for (Py_ssize_t i = 0; i < COUNTER_THREADS; i++) {
181 thread_data[i].test_data = &test_data;
182 PyThread_start_new_thread(slow_counter_thread, &thread_data[i]);
183 }
184
185 for (Py_ssize_t i = 0; i < COUNTER_THREADS; i++) {
186 PyEvent_Wait(&thread_data[i].done_event);
187 }
188
189 assert(test_data.counter == COUNTER_THREADS * SLOW_COUNTER_ITERS);
190 Py_RETURN_NONE;
191 }
192
193 struct bench_data_locks {
194 int stop;
195 int use_pymutex;
196 int critical_section_length;
197 char padding[200];
198 PyThread_type_lock lock;
199 PyMutex m;
200 double value;
201 Py_ssize_t total_iters;
202 };
203
204 struct bench_thread_data {
205 struct bench_data_locks *bench_data;
206 Py_ssize_t iters;
207 PyEvent done;
208 };
209
210 static void
thread_benchmark_locks(void * arg)211 thread_benchmark_locks(void *arg)
212 {
213 struct bench_thread_data *thread_data = arg;
214 struct bench_data_locks *bench_data = thread_data->bench_data;
215 int use_pymutex = bench_data->use_pymutex;
216 int critical_section_length = bench_data->critical_section_length;
217
218 double my_value = 1.0;
219 Py_ssize_t iters = 0;
220 while (!_Py_atomic_load_int_relaxed(&bench_data->stop)) {
221 if (use_pymutex) {
222 PyMutex_Lock(&bench_data->m);
223 for (int i = 0; i < critical_section_length; i++) {
224 bench_data->value += my_value;
225 my_value = bench_data->value;
226 }
227 PyMutex_Unlock(&bench_data->m);
228 }
229 else {
230 PyThread_acquire_lock(bench_data->lock, 1);
231 for (int i = 0; i < critical_section_length; i++) {
232 bench_data->value += my_value;
233 my_value = bench_data->value;
234 }
235 PyThread_release_lock(bench_data->lock);
236 }
237 iters++;
238 }
239
240 thread_data->iters = iters;
241 _Py_atomic_add_ssize(&bench_data->total_iters, iters);
242 _PyEvent_Notify(&thread_data->done);
243 }
244
245 /*[clinic input]
246 _testinternalcapi.benchmark_locks
247
248 num_threads: Py_ssize_t
249 use_pymutex: bool = True
250 critical_section_length: int = 1
251 time_ms: int = 1000
252 /
253
254 [clinic start generated code]*/
255
256 static PyObject *
_testinternalcapi_benchmark_locks_impl(PyObject * module,Py_ssize_t num_threads,int use_pymutex,int critical_section_length,int time_ms)257 _testinternalcapi_benchmark_locks_impl(PyObject *module,
258 Py_ssize_t num_threads,
259 int use_pymutex,
260 int critical_section_length,
261 int time_ms)
262 /*[clinic end generated code: output=381df8d7e9a74f18 input=f3aeaf688738c121]*/
263 {
264 // Run from Tools/lockbench/lockbench.py
265 // Based on the WebKit lock benchmarks:
266 // https://github.com/WebKit/WebKit/blob/main/Source/WTF/benchmarks/LockSpeedTest.cpp
267 // See also https://webkit.org/blog/6161/locking-in-webkit/
268 PyObject *thread_iters = NULL;
269 PyObject *res = NULL;
270
271 struct bench_data_locks bench_data;
272 memset(&bench_data, 0, sizeof(bench_data));
273 bench_data.use_pymutex = use_pymutex;
274 bench_data.critical_section_length = critical_section_length;
275
276 bench_data.lock = PyThread_allocate_lock();
277 if (bench_data.lock == NULL) {
278 return PyErr_NoMemory();
279 }
280
281 struct bench_thread_data *thread_data = NULL;
282 thread_data = PyMem_Calloc(num_threads, sizeof(*thread_data));
283 if (thread_data == NULL) {
284 PyErr_NoMemory();
285 goto exit;
286 }
287
288 thread_iters = PyList_New(num_threads);
289 if (thread_iters == NULL) {
290 goto exit;
291 }
292
293 PyTime_t start, end;
294 if (PyTime_PerfCounter(&start) < 0) {
295 goto exit;
296 }
297
298 for (Py_ssize_t i = 0; i < num_threads; i++) {
299 thread_data[i].bench_data = &bench_data;
300 PyThread_start_new_thread(thread_benchmark_locks, &thread_data[i]);
301 }
302
303 // Let the threads run for `time_ms` milliseconds
304 pysleep(time_ms);
305 _Py_atomic_store_int(&bench_data.stop, 1);
306
307 // Wait for the threads to finish
308 for (Py_ssize_t i = 0; i < num_threads; i++) {
309 PyEvent_Wait(&thread_data[i].done);
310 }
311
312 Py_ssize_t total_iters = bench_data.total_iters;
313 if (PyTime_PerfCounter(&end) < 0) {
314 goto exit;
315 }
316
317 // Return the total number of acquisitions and the number of acquisitions
318 // for each thread.
319 for (Py_ssize_t i = 0; i < num_threads; i++) {
320 PyObject *iter = PyLong_FromSsize_t(thread_data[i].iters);
321 if (iter == NULL) {
322 goto exit;
323 }
324 PyList_SET_ITEM(thread_iters, i, iter);
325 }
326
327 assert(end != start);
328 double rate = total_iters * 1e9 / (end - start);
329 res = Py_BuildValue("(dO)", rate, thread_iters);
330
331 exit:
332 PyThread_free_lock(bench_data.lock);
333 PyMem_Free(thread_data);
334 Py_XDECREF(thread_iters);
335 return res;
336 }
337
338 static PyObject *
test_lock_benchmark(PyObject * module,PyObject * obj)339 test_lock_benchmark(PyObject *module, PyObject *obj)
340 {
341 // Just make sure the benchmark runs without crashing
342 PyObject *res = _testinternalcapi_benchmark_locks_impl(
343 module, 1, 1, 1, 100);
344 if (res == NULL) {
345 return NULL;
346 }
347 Py_DECREF(res);
348 Py_RETURN_NONE;
349 }
350
351 static int
init_maybe_fail(void * arg)352 init_maybe_fail(void *arg)
353 {
354 int *counter = (int *)arg;
355 (*counter)++;
356 if (*counter < 5) {
357 // failure
358 return -1;
359 }
360 assert(*counter == 5);
361 return 0;
362 }
363
364 static PyObject *
test_lock_once(PyObject * self,PyObject * obj)365 test_lock_once(PyObject *self, PyObject *obj)
366 {
367 _PyOnceFlag once = {0};
368 int counter = 0;
369 for (int i = 0; i < 10; i++) {
370 int res = _PyOnceFlag_CallOnce(&once, init_maybe_fail, &counter);
371 if (i < 4) {
372 assert(res == -1);
373 }
374 else {
375 assert(res == 0);
376 assert(counter == 5);
377 }
378 }
379 Py_RETURN_NONE;
380 }
381
382 struct test_rwlock_data {
383 Py_ssize_t nthreads;
384 _PyRWMutex rw;
385 PyEvent step1;
386 PyEvent step2;
387 PyEvent step3;
388 PyEvent done;
389 };
390
391 static void
rdlock_thread(void * arg)392 rdlock_thread(void *arg)
393 {
394 struct test_rwlock_data *test_data = arg;
395
396 // Acquire the lock in read mode
397 _PyRWMutex_RLock(&test_data->rw);
398 PyEvent_Wait(&test_data->step1);
399 _PyRWMutex_RUnlock(&test_data->rw);
400
401 _PyRWMutex_RLock(&test_data->rw);
402 PyEvent_Wait(&test_data->step3);
403 _PyRWMutex_RUnlock(&test_data->rw);
404
405 if (_Py_atomic_add_ssize(&test_data->nthreads, -1) == 1) {
406 _PyEvent_Notify(&test_data->done);
407 }
408 }
409 static void
wrlock_thread(void * arg)410 wrlock_thread(void *arg)
411 {
412 struct test_rwlock_data *test_data = arg;
413
414 // First acquire the lock in write mode
415 _PyRWMutex_Lock(&test_data->rw);
416 PyEvent_Wait(&test_data->step2);
417 _PyRWMutex_Unlock(&test_data->rw);
418
419 if (_Py_atomic_add_ssize(&test_data->nthreads, -1) == 1) {
420 _PyEvent_Notify(&test_data->done);
421 }
422 }
423
424 static void
wait_until(uintptr_t * ptr,uintptr_t value)425 wait_until(uintptr_t *ptr, uintptr_t value)
426 {
427 // wait up to two seconds for *ptr == value
428 int iters = 0;
429 uintptr_t bits;
430 do {
431 pysleep(10);
432 bits = _Py_atomic_load_uintptr(ptr);
433 iters++;
434 } while (bits != value && iters < 200);
435 }
436
437 static PyObject *
test_lock_rwlock(PyObject * self,PyObject * obj)438 test_lock_rwlock(PyObject *self, PyObject *obj)
439 {
440 struct test_rwlock_data test_data = {.nthreads = 3};
441
442 _PyRWMutex_Lock(&test_data.rw);
443 assert(test_data.rw.bits == 1);
444
445 _PyRWMutex_Unlock(&test_data.rw);
446 assert(test_data.rw.bits == 0);
447
448 // Start two readers
449 PyThread_start_new_thread(rdlock_thread, &test_data);
450 PyThread_start_new_thread(rdlock_thread, &test_data);
451
452 // wait up to two seconds for the threads to attempt to read-lock "rw"
453 wait_until(&test_data.rw.bits, 8);
454 assert(test_data.rw.bits == 8);
455
456 // start writer (while readers hold lock)
457 PyThread_start_new_thread(wrlock_thread, &test_data);
458 wait_until(&test_data.rw.bits, 10);
459 assert(test_data.rw.bits == 10);
460
461 // readers release lock, writer should acquire it
462 _PyEvent_Notify(&test_data.step1);
463 wait_until(&test_data.rw.bits, 3);
464 assert(test_data.rw.bits == 3);
465
466 // writer releases lock, readers acquire it
467 _PyEvent_Notify(&test_data.step2);
468 wait_until(&test_data.rw.bits, 8);
469 assert(test_data.rw.bits == 8);
470
471 // readers release lock again
472 _PyEvent_Notify(&test_data.step3);
473 wait_until(&test_data.rw.bits, 0);
474 assert(test_data.rw.bits == 0);
475
476 PyEvent_Wait(&test_data.done);
477 Py_RETURN_NONE;
478 }
479
480 static PyObject *
test_lock_recursive(PyObject * self,PyObject * obj)481 test_lock_recursive(PyObject *self, PyObject *obj)
482 {
483 _PyRecursiveMutex m = (_PyRecursiveMutex){0};
484 assert(!_PyRecursiveMutex_IsLockedByCurrentThread(&m));
485
486 _PyRecursiveMutex_Lock(&m);
487 assert(m.thread == PyThread_get_thread_ident_ex());
488 assert(PyMutex_IsLocked(&m.mutex));
489 assert(m.level == 0);
490
491 _PyRecursiveMutex_Lock(&m);
492 assert(m.level == 1);
493 _PyRecursiveMutex_Unlock(&m);
494
495 _PyRecursiveMutex_Unlock(&m);
496 assert(m.thread == 0);
497 assert(!PyMutex_IsLocked(&m.mutex));
498 assert(m.level == 0);
499
500 Py_RETURN_NONE;
501 }
502
503 static PyMethodDef test_methods[] = {
504 {"test_lock_basic", test_lock_basic, METH_NOARGS},
505 {"test_lock_two_threads", test_lock_two_threads, METH_NOARGS},
506 {"test_lock_counter", test_lock_counter, METH_NOARGS},
507 {"test_lock_counter_slow", test_lock_counter_slow, METH_NOARGS},
508 _TESTINTERNALCAPI_BENCHMARK_LOCKS_METHODDEF
509 {"test_lock_benchmark", test_lock_benchmark, METH_NOARGS},
510 {"test_lock_once", test_lock_once, METH_NOARGS},
511 {"test_lock_rwlock", test_lock_rwlock, METH_NOARGS},
512 {"test_lock_recursive", test_lock_recursive, METH_NOARGS},
513 {NULL, NULL} /* sentinel */
514 };
515
516 int
_PyTestInternalCapi_Init_Lock(PyObject * mod)517 _PyTestInternalCapi_Init_Lock(PyObject *mod)
518 {
519 if (PyModule_AddFunctions(mod, test_methods) < 0) {
520 return -1;
521 }
522 return 0;
523 }
524