@@ -2642,6 +2642,108 @@ test_soft_deprecated_macros(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(args)
26422642 Py_RETURN_NONE ;
26432643}
26442644
2645+ /**
2646+ * A spinning barrier is a multithreading barrier similar to threading.Barrier,
2647+ * except that it never parks threads that are waiting on the barrier.
2648+ *
2649+ * This is useful in scenarios where it is desirable to increase contention on
2650+ * the code that follows the barrier. For instance, consider this test:
2651+ *
2652+ * def test_my_method_is_atomic():
2653+ * x = MyClass()
2654+ * b = _testcapi.SpinningBarrier()
2655+ *
2656+ * def thread():
2657+ * b.wait()
2658+ * x.my_method()
2659+ *
2660+ * for _ in range(1_000):
2661+ * threads = [Thread(target=thread), Thread(target=thread)]
2662+ * for t in threads: t.start()
2663+ * for t in threads: t.join()
2664+ *
2665+ * It can be desirable (and sometimes necessary) to increase the contention
2666+ * on x's internal data structure by avoiding threads parking.
2667+ * Here, not parking may become necessary when the code in my_method() is so
2668+ * short that contention-related code paths are never exercised otherwise.
2669+ *
2670+ * It is roughly equivalent to this Python class:
2671+ *
2672+ * class SpinningBarrier:
2673+ * def __init__(self, parties: int):
2674+ * self.parties = AtomicInt(parties) # if only we had atomic integers
2675+ *
2676+ * def wait(self):
2677+ * v = self.parties.decrement_and_get()
2678+ * while True:
2679+ * if v < 0:
2680+ * raise ValueError("wait was called too many times")
2681+ * if v == 0:
2682+ * return
2683+ * v = self.parties.get()
2684+ *
2685+ */
2686+
2687+ typedef struct {
2688+ PyObject_HEAD
2689+ int parties ;
2690+ } SpinningBarrier ;
2691+
2692+ int
2693+ SpinningBarrier_init (PyObject * self , PyObject * args , PyObject * kwargs )
2694+ {
2695+ int parties = 0 ;
2696+ const char * kwlist [] = {"parties" , NULL };
2697+ if (!PyArg_ParseTupleAndKeywords (args , kwargs , "i" , (char * * )kwlist , & parties )) {
2698+ return -1 ;
2699+ }
2700+ if (parties <= 0 ) {
2701+ PyErr_SetString (PyExc_ValueError , "parties must be greater than zero" );
2702+ return -1 ;
2703+ }
2704+
2705+ SpinningBarrier * self_b = (SpinningBarrier * ) self ;
2706+ self_b -> parties = parties ;
2707+
2708+ return 0 ;
2709+ }
2710+
2711+ PyObject *
2712+ SpinningBarrier_wait (PyObject * self , PyObject * Py_UNUSED (args ))
2713+ {
2714+ SpinningBarrier * self_b = (SpinningBarrier * ) self ;
2715+ const long decremented = _Py_atomic_add_int (& self_b -> parties , -1 ) - 1 ;
2716+ long v = decremented ;
2717+ while (1 ) {
2718+ if (v < 0 ) {
2719+ PyErr_SetString (PyExc_ValueError , "wait was called too many times" );
2720+ return NULL ;
2721+ }
2722+ if (v == 0 ) {
2723+ return PyLong_FromLong (decremented );
2724+ }
2725+ v = _Py_atomic_load_int_relaxed (& self_b -> parties );
2726+ if (PyErr_CheckSignals ()) {
2727+ return NULL ;
2728+ }
2729+ }
2730+ }
2731+
2732+ static PyMethodDef SpinningBarrier_methods [] = {
2733+ {"wait" , SpinningBarrier_wait , METH_NOARGS },
2734+ {NULL , NULL },
2735+ };
2736+
2737+ static PyTypeObject SpinningBarrier_Type = {
2738+ PyVarObject_HEAD_INIT (NULL , 0 )
2739+ .tp_name = "SpinningBarrier" ,
2740+ .tp_basicsize = sizeof (SpinningBarrier ),
2741+ .tp_new = PyType_GenericNew ,
2742+ .tp_free = PyObject_Free ,
2743+ .tp_init = & SpinningBarrier_init ,
2744+ .tp_methods = SpinningBarrier_methods ,
2745+ };
2746+
26452747static PyMethodDef TestMethods [] = {
26462748 {"set_errno" , set_errno , METH_VARARGS },
26472749 {"test_config" , test_config , METH_NOARGS },
@@ -3369,6 +3471,12 @@ _testcapi_exec(PyObject *m)
33693471 Py_INCREF (& MethStatic_Type );
33703472 PyModule_AddObject (m , "MethStatic" , (PyObject * )& MethStatic_Type );
33713473
3474+ if (PyType_Ready (& SpinningBarrier_Type ) < 0 ) {
3475+ return -1 ;
3476+ }
3477+ Py_INCREF (& SpinningBarrier_Type );
3478+ PyModule_AddObject (m , "SpinningBarrier" , (PyObject * )& SpinningBarrier_Type );
3479+
33723480 PyModule_AddObject (m , "CHAR_MAX" , PyLong_FromLong (CHAR_MAX ));
33733481 PyModule_AddObject (m , "CHAR_MIN" , PyLong_FromLong (CHAR_MIN ));
33743482 PyModule_AddObject (m , "UCHAR_MAX" , PyLong_FromLong (UCHAR_MAX ));
0 commit comments