gh-111968: Use per-thread freelists for PyContext in free-threading (gh-114122)

5 files changed, 18 insertions, 28 deletions

diff --git a/Python/context.c b/Python/context.c
index c94c014..1e90811 100644
--- a/Python/context.c
+++ b/Python/context.c
@@ -64,12 +64,12 @@ static int
 contextvar_del(PyContextVar *var);
 
 
-#if PyContext_MAXFREELIST > 0
+#ifdef WITH_FREELISTS
 static struct _Py_context_state *
 get_context_state(void)
 {
-    PyInterpreterState *interp = _PyInterpreterState_GET();
-    return &interp->context;
+    _PyFreeListState *state = _PyFreeListState_GET();
+    return &state->context_state;
 }
 #endif
 
@@ -340,13 +340,9 @@ static inline PyContext *
 _context_alloc(void)
 {
     PyContext *ctx;
-#if PyContext_MAXFREELIST > 0
+#ifdef WITH_FREELISTS
     struct _Py_context_state *state = get_context_state();
-#ifdef Py_DEBUG
-    // _context_alloc() must not be called after _PyContext_Fini()
-    assert(state->numfree != -1);
-#endif
-    if (state->numfree) {
+    if (state->numfree > 0) {
         state->numfree--;
         ctx = state->freelist;
         state->freelist = (PyContext *)ctx->ctx_weakreflist;
@@ -471,13 +467,9 @@ context_tp_dealloc(PyContext *self)
     }
     (void)context_tp_clear(self);
 
-#if PyContext_MAXFREELIST > 0
+#ifdef WITH_FREELISTS
     struct _Py_context_state *state = get_context_state();
-#ifdef Py_DEBUG
-    // _context_alloc() must not be called after _PyContext_Fini()
-    assert(state->numfree != -1);
-#endif
-    if (state->numfree < PyContext_MAXFREELIST) {
+    if (state->numfree >= 0 && state->numfree < PyContext_MAXFREELIST) {
         state->numfree++;
         self->ctx_weakreflist = (PyObject *)state->freelist;
         state->freelist = self;
@@ -1275,28 +1267,27 @@ get_token_missing(void)
 
 
 void
-_PyContext_ClearFreeList(PyInterpreterState *interp)
+_PyContext_ClearFreeList(_PyFreeListState *freelist_state, int is_finalization)
 {
-#if PyContext_MAXFREELIST > 0
-    struct _Py_context_state *state = &interp->context;
-    for (; state->numfree; state->numfree--) {
+#ifdef WITH_FREELISTS
+    struct _Py_context_state *state = &freelist_state->context_state;
+    for (; state->numfree > 0; state->numfree--) {
         PyContext *ctx = state->freelist;
         state->freelist = (PyContext *)ctx->ctx_weakreflist;
         ctx->ctx_weakreflist = NULL;
         PyObject_GC_Del(ctx);
     }
+    if (is_finalization) {
+        state->numfree = -1;
+    }
 #endif
 }
 
 
 void
-_PyContext_Fini(PyInterpreterState *interp)
+_PyContext_Fini(_PyFreeListState *state)
 {
-    _PyContext_ClearFreeList(interp);
-#if defined(Py_DEBUG) && PyContext_MAXFREELIST > 0
-    struct _Py_context_state *state = &interp->context;
-    state->numfree = -1;
-#endif
+    _PyContext_ClearFreeList(state, 1);
 }
 
 
diff --git a/Python/gc_free_threading.c b/Python/gc_free_threading.c
index b1d88ff..b1511eb 100644
--- a/Python/gc_free_threading.c
+++ b/Python/gc_free_threading.c
@@ -16,7 +16,6 @@ _PyGC_ClearAllFreeLists(PyInterpreterState *interp)
 {
     _PyDict_ClearFreeList(interp);
     _PyAsyncGen_ClearFreeLists(interp);
-    _PyContext_ClearFreeList(interp);
 
     HEAD_LOCK(&_PyRuntime);
     _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)interp->threads.head;
diff --git a/Python/gc_gil.c b/Python/gc_gil.c
index 873fad8..edf8417 100644
--- a/Python/gc_gil.c
+++ b/Python/gc_gil.c
@@ -13,7 +13,6 @@ _PyGC_ClearAllFreeLists(PyInterpreterState *interp)
 {
     _PyDict_ClearFreeList(interp);
     _PyAsyncGen_ClearFreeLists(interp);
-    _PyContext_ClearFreeList(interp);
 
     _Py_ClearFreeLists(&interp->freelist_state, 0);
 }
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index c33892a..598cd68 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -1736,7 +1736,6 @@ finalize_interp_types(PyInterpreterState *interp)
     _PyXI_FiniTypes(interp);
     _PyExc_Fini(interp);
     _PyAsyncGen_Fini(interp);
-    _PyContext_Fini(interp);
     _PyFloat_FiniType(interp);
     _PyLong_FiniTypes(interp);
     _PyThread_FiniType(interp);
@@ -1759,6 +1758,7 @@ finalize_interp_types(PyInterpreterState *interp)
     _PyList_Fini(state);
     _PyFloat_Fini(state);
     _PySlice_Fini(state);
+    _PyContext_Fini(state);
 
 #ifdef Py_DEBUG
     _PyStaticObjects_CheckRefcnt(interp);
diff --git a/Python/pystate.c b/Python/pystate.c
index 01dc86f..8374223 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -1461,6 +1461,7 @@ _Py_ClearFreeLists(_PyFreeListState *state, int is_finalization)
     _PyFloat_ClearFreeList(state, is_finalization);
     _PyTuple_ClearFreeList(state, is_finalization);
     _PyList_ClearFreeList(state, is_finalization);
+    _PyContext_ClearFreeList(state, is_finalization);
 }
 
 void