1 /*
2 * C Extension module to test pycore_critical_section.h API.
3 */
4
5 #include "parts.h"
6
7 #include "pycore_critical_section.h"
8
9 #ifdef Py_GIL_DISABLED
10 #define assert_nogil assert
11 #define assert_gil(x)
12 #else
13 #define assert_gil assert
14 #define assert_nogil(x)
15 #endif
16
17
18 static PyObject *
test_critical_sections(PyObject * self,PyObject * Py_UNUSED (args))19 test_critical_sections(PyObject *self, PyObject *Py_UNUSED(args))
20 {
21 PyObject *d1 = PyDict_New();
22 assert(d1 != NULL);
23
24 PyObject *d2 = PyDict_New();
25 assert(d2 != NULL);
26
27 // Beginning a critical section should lock the associated object and
28 // push the critical section onto the thread's stack (in Py_GIL_DISABLED builds).
29 Py_BEGIN_CRITICAL_SECTION(d1);
30 assert_nogil(PyMutex_IsLocked(&d1->ob_mutex));
31 assert_nogil(_PyCriticalSection_IsActive(PyThreadState_GET()->critical_section));
32 assert_gil(PyThreadState_GET()->critical_section == 0);
33 Py_END_CRITICAL_SECTION();
34 assert_nogil(!PyMutex_IsLocked(&d1->ob_mutex));
35
36 assert_nogil(!PyMutex_IsLocked(&d1->ob_mutex));
37 assert_nogil(!PyMutex_IsLocked(&d2->ob_mutex));
38 Py_BEGIN_CRITICAL_SECTION2(d1, d2);
39 assert_nogil(PyMutex_IsLocked(&d1->ob_mutex));
40 assert_nogil(PyMutex_IsLocked(&d2->ob_mutex));
41 Py_END_CRITICAL_SECTION2();
42 assert_nogil(!PyMutex_IsLocked(&d1->ob_mutex));
43 assert_nogil(!PyMutex_IsLocked(&d2->ob_mutex));
44
45 // Passing the same object twice should work (and not deadlock).
46 assert_nogil(!PyMutex_IsLocked(&d2->ob_mutex));
47 Py_BEGIN_CRITICAL_SECTION2(d2, d2);
48 assert_nogil(PyMutex_IsLocked(&d2->ob_mutex));
49 Py_END_CRITICAL_SECTION2();
50 assert_nogil(!PyMutex_IsLocked(&d2->ob_mutex));
51
52 Py_DECREF(d2);
53 Py_DECREF(d1);
54 Py_RETURN_NONE;
55 }
56
57 static void
lock_unlock_object(PyObject * obj,int recurse_depth)58 lock_unlock_object(PyObject *obj, int recurse_depth)
59 {
60 Py_BEGIN_CRITICAL_SECTION(obj);
61 if (recurse_depth > 0) {
62 lock_unlock_object(obj, recurse_depth - 1);
63 }
64 Py_END_CRITICAL_SECTION();
65 }
66
67 static void
lock_unlock_two_objects(PyObject * a,PyObject * b,int recurse_depth)68 lock_unlock_two_objects(PyObject *a, PyObject *b, int recurse_depth)
69 {
70 Py_BEGIN_CRITICAL_SECTION2(a, b);
71 if (recurse_depth > 0) {
72 lock_unlock_two_objects(a, b, recurse_depth - 1);
73 }
74 Py_END_CRITICAL_SECTION2();
75 }
76
77
78 // Test that nested critical sections do not deadlock if they attempt to lock
79 // the same object.
80 static PyObject *
test_critical_sections_nest(PyObject * self,PyObject * Py_UNUSED (args))81 test_critical_sections_nest(PyObject *self, PyObject *Py_UNUSED(args))
82 {
83 PyObject *a = PyDict_New();
84 assert(a != NULL);
85 PyObject *b = PyDict_New();
86 assert(b != NULL);
87
88 // Locking an object recursively with this API should not deadlock.
89 assert_nogil(!PyMutex_IsLocked(&a->ob_mutex));
90 Py_BEGIN_CRITICAL_SECTION(a);
91 assert_nogil(PyMutex_IsLocked(&a->ob_mutex));
92 lock_unlock_object(a, 10);
93 assert_nogil(PyMutex_IsLocked(&a->ob_mutex));
94 Py_END_CRITICAL_SECTION();
95 assert_nogil(!PyMutex_IsLocked(&a->ob_mutex));
96
97 // Same test but with two objects.
98 Py_BEGIN_CRITICAL_SECTION2(b, a);
99 lock_unlock_two_objects(a, b, 10);
100 assert_nogil(PyMutex_IsLocked(&a->ob_mutex));
101 assert_nogil(PyMutex_IsLocked(&b->ob_mutex));
102 Py_END_CRITICAL_SECTION2();
103
104 Py_DECREF(b);
105 Py_DECREF(a);
106 Py_RETURN_NONE;
107 }
108
109 // Test that a critical section is suspended by a Py_BEGIN_ALLOW_THREADS and
110 // resumed by a Py_END_ALLOW_THREADS.
111 static PyObject *
test_critical_sections_suspend(PyObject * self,PyObject * Py_UNUSED (args))112 test_critical_sections_suspend(PyObject *self, PyObject *Py_UNUSED(args))
113 {
114 PyObject *a = PyDict_New();
115 assert(a != NULL);
116
117 Py_BEGIN_CRITICAL_SECTION(a);
118 assert_nogil(PyMutex_IsLocked(&a->ob_mutex));
119
120 // Py_BEGIN_ALLOW_THREADS should suspend the active critical section
121 Py_BEGIN_ALLOW_THREADS
122 assert_nogil(!PyMutex_IsLocked(&a->ob_mutex));
123 Py_END_ALLOW_THREADS;
124
125 // After Py_END_ALLOW_THREADS the critical section should be resumed.
126 assert_nogil(PyMutex_IsLocked(&a->ob_mutex));
127 Py_END_CRITICAL_SECTION();
128
129 Py_DECREF(a);
130 Py_RETURN_NONE;
131 }
132
133 #ifdef Py_CAN_START_THREADS
134 struct test_data {
135 PyObject *obj1;
136 PyObject *obj2;
137 PyObject *obj3;
138 Py_ssize_t countdown;
139 PyEvent done_event;
140 };
141
142 static void
thread_critical_sections(void * arg)143 thread_critical_sections(void *arg)
144 {
145 const Py_ssize_t NUM_ITERS = 200;
146 struct test_data *test_data = arg;
147 PyGILState_STATE gil = PyGILState_Ensure();
148
149 for (Py_ssize_t i = 0; i < NUM_ITERS; i++) {
150 Py_BEGIN_CRITICAL_SECTION(test_data->obj1);
151 Py_END_CRITICAL_SECTION();
152
153 Py_BEGIN_CRITICAL_SECTION(test_data->obj2);
154 lock_unlock_object(test_data->obj1, 1);
155 Py_END_CRITICAL_SECTION();
156
157 Py_BEGIN_CRITICAL_SECTION2(test_data->obj3, test_data->obj1);
158 lock_unlock_object(test_data->obj2, 2);
159 Py_END_CRITICAL_SECTION2();
160
161 Py_BEGIN_CRITICAL_SECTION(test_data->obj3);
162 Py_BEGIN_ALLOW_THREADS
163 Py_END_ALLOW_THREADS
164 Py_END_CRITICAL_SECTION();
165 }
166
167 PyGILState_Release(gil);
168 if (_Py_atomic_add_ssize(&test_data->countdown, -1) == 1) {
169 // last thread to finish sets done_event
170 _PyEvent_Notify(&test_data->done_event);
171 }
172 }
173
174 static PyObject *
test_critical_sections_threads(PyObject * self,PyObject * Py_UNUSED (args))175 test_critical_sections_threads(PyObject *self, PyObject *Py_UNUSED(args))
176 {
177 const Py_ssize_t NUM_THREADS = 4;
178 struct test_data test_data = {
179 .obj1 = PyDict_New(),
180 .obj2 = PyDict_New(),
181 .obj3 = PyDict_New(),
182 .countdown = NUM_THREADS,
183 };
184 assert(test_data.obj1 != NULL);
185 assert(test_data.obj2 != NULL);
186 assert(test_data.obj3 != NULL);
187
188 for (int i = 0; i < NUM_THREADS; i++) {
189 PyThread_start_new_thread(&thread_critical_sections, &test_data);
190 }
191 PyEvent_Wait(&test_data.done_event);
192
193 Py_DECREF(test_data.obj3);
194 Py_DECREF(test_data.obj2);
195 Py_DECREF(test_data.obj1);
196 Py_RETURN_NONE;
197 }
198
199 static void
pysleep(int ms)200 pysleep(int ms)
201 {
202 #ifdef MS_WINDOWS
203 Sleep(ms);
204 #else
205 usleep(ms * 1000);
206 #endif
207 }
208
209 struct test_data_gc {
210 PyObject *obj;
211 Py_ssize_t num_threads;
212 Py_ssize_t id;
213 Py_ssize_t countdown;
214 PyEvent done_event;
215 PyEvent ready;
216 };
217
218 static void
thread_gc(void * arg)219 thread_gc(void *arg)
220 {
221 struct test_data_gc *test_data = arg;
222 PyGILState_STATE gil = PyGILState_Ensure();
223
224 Py_ssize_t id = _Py_atomic_add_ssize(&test_data->id, 1);
225 if (id == test_data->num_threads - 1) {
226 _PyEvent_Notify(&test_data->ready);
227 }
228 else {
229 // wait for all test threads to more reliably reproduce the issue.
230 PyEvent_Wait(&test_data->ready);
231 }
232
233 if (id == 0) {
234 Py_BEGIN_CRITICAL_SECTION(test_data->obj);
235 // pause long enough that the lock would be handed off directly to
236 // a waiting thread.
237 pysleep(5);
238 PyGC_Collect();
239 Py_END_CRITICAL_SECTION();
240 }
241 else if (id == 1) {
242 pysleep(1);
243 Py_BEGIN_CRITICAL_SECTION(test_data->obj);
244 pysleep(1);
245 Py_END_CRITICAL_SECTION();
246 }
247 else if (id == 2) {
248 // sleep long enough so that thread 0 is waiting to stop the world
249 pysleep(6);
250 Py_BEGIN_CRITICAL_SECTION(test_data->obj);
251 pysleep(1);
252 Py_END_CRITICAL_SECTION();
253 }
254
255 PyGILState_Release(gil);
256 if (_Py_atomic_add_ssize(&test_data->countdown, -1) == 1) {
257 // last thread to finish sets done_event
258 _PyEvent_Notify(&test_data->done_event);
259 }
260 }
261
262 static PyObject *
test_critical_sections_gc(PyObject * self,PyObject * Py_UNUSED (args))263 test_critical_sections_gc(PyObject *self, PyObject *Py_UNUSED(args))
264 {
265 // gh-118332: Contended critical sections should not deadlock with GC
266 const Py_ssize_t NUM_THREADS = 3;
267 struct test_data_gc test_data = {
268 .obj = PyDict_New(),
269 .countdown = NUM_THREADS,
270 .num_threads = NUM_THREADS,
271 };
272 assert(test_data.obj != NULL);
273
274 for (int i = 0; i < NUM_THREADS; i++) {
275 PyThread_start_new_thread(&thread_gc, &test_data);
276 }
277 PyEvent_Wait(&test_data.done_event);
278 Py_DECREF(test_data.obj);
279 Py_RETURN_NONE;
280 }
281
282 #endif
283
284 static PyMethodDef test_methods[] = {
285 {"test_critical_sections", test_critical_sections, METH_NOARGS},
286 {"test_critical_sections_nest", test_critical_sections_nest, METH_NOARGS},
287 {"test_critical_sections_suspend", test_critical_sections_suspend, METH_NOARGS},
288 #ifdef Py_CAN_START_THREADS
289 {"test_critical_sections_threads", test_critical_sections_threads, METH_NOARGS},
290 {"test_critical_sections_gc", test_critical_sections_gc, METH_NOARGS},
291 #endif
292 {NULL, NULL} /* sentinel */
293 };
294
295 int
_PyTestInternalCapi_Init_CriticalSection(PyObject * mod)296 _PyTestInternalCapi_Init_CriticalSection(PyObject *mod)
297 {
298 if (PyModule_AddFunctions(mod, test_methods) < 0) {
299 return -1;
300 }
301 return 0;
302 }
303