gh-151518: Exercise STW fairness directly

Author: tpn

Lib/test/test_free_threading/test_gc.py

@@ -1,14 +1,12 @@
 import unittest
 
-import textwrap
 import threading
 from threading import Thread
 import time
 from unittest import TestCase
 import gc
 
-from test import support
-from test.support import script_helper, threading_helper
+from test.support import threading_helper
 
 
 class MyObj:
@@ -176,56 +174,6 @@ def reader():
         with threading_helper.start_threads(threads):
             pass
 
-    @support.requires_subprocess()
-    def test_tight_gc_loop_does_not_starve_attach(self):
-        script = textwrap.dedent(f"""
-            import faulthandler
-
-            faulthandler.dump_traceback_later({support.SHORT_TIMEOUT}, exit=True)
-
-            import gc
-            import threading
-            import time
-
-            # Add GC-tracked objects to lengthen the stop-the-world pauses.
-            import abc
-            import argparse
-            import collections
-            import contextlib
-            import decimal
-            import enum
-            import functools
-            import heapq
-            import importlib
-            import inspect
-            import itertools
-            import json
-            import math
-            import operator
-            import random
-            import re
-
-            started = threading.Event()
-            stop = threading.Event()
-
-            def collect():
-                gc.collect()
-                started.set()
-                while not stop.is_set():
-                    gc.collect()
-
-            thread = threading.Thread(target=collect)
-            thread.start()
-            started.wait()
-            # Each reattachment must make progress between consecutive pauses.
-            for _ in range(50):
-                time.sleep(0.02)
-            stop.set()
-            thread.join()
-            faulthandler.cancel_dump_traceback_later()
-        """)
-        script_helper.assert_python_ok("-X", "gil=0", "-c", script)
-
 
 if __name__ == "__main__":
     unittest.main()

Lib/test/test_free_threading/test_threading.py

@@ -1,5 +1,8 @@
 import unittest
-from test.support import threading_helper
+import textwrap
+
+from test import support
+from test.support import script_helper, threading_helper
 
 threading_helper.requires_working_threading(module=True)
 
@@ -22,5 +25,39 @@ def mutate_thread():
         threading_helper.run_concurrently([repr_thread, mutate_thread])
 
 
+class TestThreadState(unittest.TestCase):
+    @support.requires_subprocess()
+    def test_tight_stw_loop_does_not_starve_attach(self):
+        script = textwrap.dedent(f"""
+            import faulthandler
+
+            faulthandler.dump_traceback_later({support.SHORT_TIMEOUT}, exit=True)
+
+            import _testinternalcapi
+            import threading
+            import time
+
+            started = threading.Event()
+            stop = threading.Event()
+
+            def stop_the_world():
+                _testinternalcapi.test_stop_the_world()
+                started.set()
+                while not stop.is_set():
+                    _testinternalcapi.test_stop_the_world()
+
+            thread = threading.Thread(target=stop_the_world)
+            thread.start()
+            started.wait()
+            # Each reattachment must make progress between consecutive pauses.
+            for _ in range(50):
+                time.sleep(0.02)
+            stop.set()
+            thread.join()
+            faulthandler.cancel_dump_traceback_later()
+        """)
+        script_helper.assert_python_ok("-X", "gil=0", "-c", script)
+
+
 if __name__ == "__main__":
     unittest.main()

Modules/_testinternalcapi.c

@@ -208,6 +208,15 @@ get_stack_margin(PyObject *self, PyObject *Py_UNUSED(args))
     return PyLong_FromSize_t(_PyOS_STACK_MARGIN_BYTES);
 }
 
+static PyObject *
+test_stop_the_world(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    _PyEval_StopTheWorld(interp);
+    _PyEval_StartTheWorld(interp);
+    Py_RETURN_NONE;
+}
+
 #ifdef MS_WINDOWS
 static const char *
 classify_address(uintptr_t addr, int jit_enabled, PyInterpreterState *interp)
@@ -3295,6 +3304,7 @@ static PyMethodDef module_functions[] = {
     {"get_c_recursion_remaining", get_c_recursion_remaining, METH_NOARGS},
     {"get_stack_pointer", get_stack_pointer, METH_NOARGS},
     {"get_stack_margin", get_stack_margin, METH_NOARGS},
+    {"test_stop_the_world", test_stop_the_world, METH_NOARGS},
     {"classify_stack_addresses", classify_stack_addresses, METH_VARARGS},
     {"get_jit_code_ranges", get_jit_code_ranges, METH_NOARGS},
     {"get_jit_backend", get_jit_backend, METH_NOARGS},

Python/pystate.c

@@ -2226,6 +2226,22 @@ tstate_try_attach(PyThreadState *tstate)
 #endif
 }
 
+static int
+tstate_try_attach_detached(PyThreadState *tstate, int *state)
+{
+#ifdef Py_GIL_DISABLED
+    assert(*state == _Py_THREAD_DETACHED ||
+           *state == _Py_THREAD_DETACHED_WAITING);
+    return _Py_atomic_compare_exchange_int(&tstate->state,
+                                           state,
+                                           _Py_THREAD_ATTACHED);
+#else
+    assert(tstate->state == _Py_THREAD_DETACHED);
+    tstate->state = _Py_THREAD_ATTACHED;
+    return 1;
+#endif
+}
+
 static void
 tstate_set_detached(PyThreadState *tstate, int detached_state)
 {
@@ -2264,9 +2280,7 @@ tstate_wait_attach(PyThreadState *tstate)
         else {
             assert(state == _Py_THREAD_DETACHED ||
                    state == _Py_THREAD_DETACHED_WAITING);
-            if (_Py_atomic_compare_exchange_int(
-                    &tstate->state, &state, _Py_THREAD_ATTACHED))
-            {
+            if (tstate_try_attach_detached(tstate, &state)) {
                 return;
             }
         }
@@ -2534,9 +2548,12 @@ start_the_world(struct _stoptheworld_state *stw)
                 do {
                     assert(state == _Py_THREAD_SUSPENDED ||
                            state == _Py_THREAD_SUSPENDED_WAITING);
-                    next_state = (state == _Py_THREAD_SUSPENDED_WAITING
-                                  ? _Py_THREAD_DETACHED_WAITING
-                                  : _Py_THREAD_DETACHED);
+                    if (state == _Py_THREAD_SUSPENDED_WAITING) {
+                        next_state = _Py_THREAD_DETACHED_WAITING;
+                    }
+                    else {
+                        next_state = _Py_THREAD_DETACHED;
+                    }
                     // Retry if an attach waiter registered concurrently.
                 } while (!_Py_atomic_compare_exchange_int(
                             &t->state, &state, next_state));