gh-102013: Add PyUnstable_GC_VisitObjects (#102014)

5 files changed, 146 insertions, 0 deletions

diff --git a/Doc/c-api/gcsupport.rst b/Doc/c-api/gcsupport.rst
index 8c90d1e..cb5d64a 100644
--- a/Doc/c-api/gcsupport.rst
+++ b/Doc/c-api/gcsupport.rst
@@ -228,3 +228,36 @@ garbage collection runs.
    Returns the current state, 0 for disabled and 1 for enabled.
 
    .. versionadded:: 3.10
+
+
+Querying Garbage Collector State
+--------------------------------
+
+The C-API provides the following interface for querying information about
+the garbage collector.
+
+.. c:function:: void PyUnstable_GC_VisitObjects(gcvisitobjects_t callback, void *arg)
+
+   Run supplied *callback* on all live GC-capable objects. *arg* is passed through to
+   all invocations of *callback*.
+
+   .. warning::
+      If new objects are (de)allocated by the callback it is undefined if they
+      will be visited.
+
+      Garbage collection is disabled during operation. Explicitly running a collection
+      in the callback may lead to undefined behaviour e.g. visiting the same objects
+      multiple times or not at all.
+
+   .. versionadded:: 3.12
+
+.. c:type:: int (*gcvisitobjects_t)(PyObject *object, void *arg)
+
+   Type of the visitor function to be passed to :c:func:`PyUnstable_GC_VisitObjects`.
+   *arg* is the same as the *arg* passed to ``PyUnstable_GC_VisitObjects``.
+   Return ``0`` to continue iteration, return ``1`` to stop iteration. Other return
+   values are reserved for now so behavior on returning anything else is undefined.
+
+   .. versionadded:: 3.12
+
+
diff --git a/Include/objimpl.h b/Include/objimpl.h
index dde8df3..ef871c5 100644
--- a/Include/objimpl.h
+++ b/Include/objimpl.h
@@ -157,6 +157,25 @@ PyAPI_FUNC(int) PyGC_Enable(void);
 PyAPI_FUNC(int) PyGC_Disable(void);
 PyAPI_FUNC(int) PyGC_IsEnabled(void);
 
+
+#if !defined(Py_LIMITED_API)
+/* Visit all live GC-capable objects, similar to gc.get_objects(None). The
+ * supplied callback is called on every such object with the void* arg set
+ * to the supplied arg. Returning 0 from the callback ends iteration, returning
+ * 1 allows iteration to continue. Returning any other value may result in
+ * undefined behaviour.
+ *
+ * If new objects are (de)allocated by the callback it is undefined if they
+ * will be visited.
+
+ * Garbage collection is disabled during operation. Explicitly running a
+ * collection in the callback may lead to undefined behaviour e.g. visiting the
+ * same objects multiple times or not at all.
+ */
+typedef int (*gcvisitobjects_t)(PyObject*, void*);
+PyAPI_FUNC(void) PyUnstable_GC_VisitObjects(gcvisitobjects_t callback, void* arg);
+#endif
+
 /* Test if a type has a GC head */
 #define PyType_IS_GC(t) PyType_HasFeature((t), Py_TPFLAGS_HAVE_GC)
 
diff --git a/Misc/NEWS.d/next/C API/2023-02-18-00-55-14.gh-issue-102013.83mrtI.rst b/Misc/NEWS.d/next/C API/2023-02-18-00-55-14.gh-issue-102013.83mrtI.rst
new file mode 100644
index 0000000..0350237
--- /dev/null
+++ b/Misc/NEWS.d/next/C API/2023-02-18-00-55-14.gh-issue-102013.83mrtI.rst
@@ -0,0 +1 @@
+Add a new (unstable) C-API function for iterating over GC'able objects using a callback: ``PyUnstable_VisitObjects``.
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index ea67017..f45d031 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -3310,6 +3310,73 @@ function_set_kw_defaults(PyObject *self, PyObject *args)
     Py_RETURN_NONE;
 }
 
+struct gc_visit_state_basic {
+    PyObject *target;
+    int found;
+};
+
+static int
+gc_visit_callback_basic(PyObject *obj, void *arg)
+{
+    struct gc_visit_state_basic *state = (struct gc_visit_state_basic *)arg;
+    if (obj == state->target) {
+        state->found = 1;
+        return 0;
+    }
+    return 1;
+}
+
+static PyObject *
+test_gc_visit_objects_basic(PyObject *Py_UNUSED(self),
+                            PyObject *Py_UNUSED(ignored))
+{
+    PyObject *obj;
+    struct gc_visit_state_basic state;
+
+    obj = PyList_New(0);
+    if (obj == NULL) {
+        return NULL;
+    }
+    state.target = obj;
+    state.found = 0;
+    
+    PyUnstable_GC_VisitObjects(gc_visit_callback_basic, &state);
+    Py_DECREF(obj);
+    if (!state.found) {
+        PyErr_SetString(
+             PyExc_AssertionError,
+             "test_gc_visit_objects_basic: Didn't find live list");
+         return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+static int
+gc_visit_callback_exit_early(PyObject *obj, void *arg)
+ {
+    int *visited_i = (int *)arg;
+    (*visited_i)++;
+    if (*visited_i == 2) {
+        return 0;
+    }
+    return 1;
+}
+
+static PyObject *
+test_gc_visit_objects_exit_early(PyObject *Py_UNUSED(self),
+                                 PyObject *Py_UNUSED(ignored))
+{
+    int visited_i = 0;
+    PyUnstable_GC_VisitObjects(gc_visit_callback_exit_early, &visited_i);
+    if (visited_i != 2) {
+        PyErr_SetString(
+            PyExc_AssertionError,
+            "test_gc_visit_objects_exit_early: did not exit when expected");
+    }
+    Py_RETURN_NONE;
+}
+
+
 static PyObject *test_buildvalue_issue38913(PyObject *, PyObject *);
 
 static PyMethodDef TestMethods[] = {
@@ -3452,6 +3519,8 @@ static PyMethodDef TestMethods[] = {
     {"function_set_defaults", function_set_defaults, METH_VARARGS, NULL},
     {"function_get_kw_defaults", function_get_kw_defaults, METH_O, NULL},
     {"function_set_kw_defaults", function_set_kw_defaults, METH_VARARGS, NULL},
+    {"test_gc_visit_objects_basic", test_gc_visit_objects_basic, METH_NOARGS, NULL},
+    {"test_gc_visit_objects_exit_early", test_gc_visit_objects_exit_early, METH_NOARGS, NULL},
     {NULL, NULL} /* sentinel */
 };
 
diff --git a/Modules/gcmodule.c b/Modules/gcmodule.c
index 5879c5e..4eaa549 100644
--- a/Modules/gcmodule.c
+++ b/Modules/gcmodule.c
@@ -2401,3 +2401,27 @@ PyObject_GC_IsFinalized(PyObject *obj)
     }
     return 0;
 }
+
+void
+PyUnstable_GC_VisitObjects(gcvisitobjects_t callback, void *arg)
+{
+    size_t i;
+    GCState *gcstate = get_gc_state();
+    int origenstate = gcstate->enabled;
+    gcstate->enabled = 0;
+    for (i = 0; i < NUM_GENERATIONS; i++) {
+        PyGC_Head *gc_list, *gc;
+        gc_list = GEN_HEAD(gcstate, i);
+        for (gc = GC_NEXT(gc_list); gc != gc_list; gc = GC_NEXT(gc)) {
+            PyObject *op = FROM_GC(gc);
+            Py_INCREF(op);
+            int res = callback(op, arg);
+            Py_DECREF(op);
+            if (!res) {
+                goto done;
+            }
+        }
+    }
+done:
+    gcstate->enabled = origenstate;
+}