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

9 files changed, 32 insertions, 48 deletions

diff --git a/Include/internal/pycore_context.h b/Include/internal/pycore_context.h
index ec884e9..3284efb 100644
--- a/Include/internal/pycore_context.h
+++ b/Include/internal/pycore_context.h
@@ -5,6 +5,7 @@
 #  error "this header requires Py_BUILD_CORE define"
 #endif
 
+#include "pycore_freelist.h"      // _PyFreeListState
 #include "pycore_hamt.h"          // PyHamtObject
 
 
@@ -13,7 +14,7 @@ extern PyTypeObject _PyContextTokenMissing_Type;
 /* runtime lifecycle */
 
 PyStatus _PyContext_Init(PyInterpreterState *);
-void _PyContext_Fini(PyInterpreterState *);
+void _PyContext_Fini(_PyFreeListState *);
 
 
 /* other API */
@@ -22,23 +23,6 @@ typedef struct {
     PyObject_HEAD
 } _PyContextTokenMissing;
 
-#ifndef WITH_FREELISTS
-// without freelists
-#  define PyContext_MAXFREELIST 0
-#endif
-
-#ifndef PyContext_MAXFREELIST
-#  define PyContext_MAXFREELIST 255
-#endif
-
-struct _Py_context_state {
-#if PyContext_MAXFREELIST > 0
-    // List of free PyContext objects
-    PyContext *freelist;
-    int numfree;
-#endif
-};
-
 struct _pycontextobject {
     PyObject_HEAD
     PyContext *ctx_prev;
diff --git a/Include/internal/pycore_freelist.h b/Include/internal/pycore_freelist.h
index faa0c11..566d47d 100644
--- a/Include/internal/pycore_freelist.h
+++ b/Include/internal/pycore_freelist.h
@@ -18,11 +18,13 @@ extern "C" {
 #  define PyTuple_MAXFREELIST 2000
 #  define PyList_MAXFREELIST 80
 #  define PyFloat_MAXFREELIST 100
+#  define PyContext_MAXFREELIST 255
 #else
 #  define PyTuple_NFREELISTS 0
 #  define PyTuple_MAXFREELIST 0
 #  define PyList_MAXFREELIST 0
 #  define PyFloat_MAXFREELIST 0
+#  define PyContext_MAXFREELIST 0
 #endif
 
 struct _Py_list_state {
@@ -67,11 +69,20 @@ struct _Py_slice_state {
 #endif
 };
 
+struct _Py_context_state {
+#ifdef WITH_FREELISTS
+    // List of free PyContext objects
+    PyContext *freelist;
+    int numfree;
+#endif
+};
+
 typedef struct _Py_freelist_state {
     struct _Py_float_state float_state;
     struct _Py_tuple_state tuple_state;
     struct _Py_list_state list_state;
     struct _Py_slice_state slice_state;
+    struct _Py_context_state context_state;
 } _PyFreeListState;
 
 #ifdef __cplusplus
diff --git a/Include/internal/pycore_gc.h b/Include/internal/pycore_gc.h
index f8e86a2..52e8b39 100644
--- a/Include/internal/pycore_gc.h
+++ b/Include/internal/pycore_gc.h
@@ -252,7 +252,7 @@ extern void _PyList_ClearFreeList(_PyFreeListState *state, int is_finalization);
 extern void _PySlice_ClearCache(_PyFreeListState *state);
 extern void _PyDict_ClearFreeList(PyInterpreterState *interp);
 extern void _PyAsyncGen_ClearFreeLists(PyInterpreterState *interp);
-extern void _PyContext_ClearFreeList(PyInterpreterState *interp);
+extern void _PyContext_ClearFreeList(_PyFreeListState *state, int is_finalization);
 extern void _Py_ScheduleGC(PyInterpreterState *interp);
 extern void _Py_RunGC(PyThreadState *tstate);
 
diff --git a/Include/internal/pycore_interp.h b/Include/internal/pycore_interp.h
index 134a882..7ec9630 100644
--- a/Include/internal/pycore_interp.h
+++ b/Include/internal/pycore_interp.h
@@ -191,7 +191,6 @@ struct _is {
     struct _Py_tuple_state tuple;
     struct _Py_dict_state dict_state;
     struct _Py_async_gen_state async_gen;
-    struct _Py_context_state context;
     struct _Py_exc_state exc_state;
 
     struct ast_state ast;
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