bpo-31558: Add gc.freeze() (#3705)

Freeze all the objects tracked by gc - move them to a permanent generation
and ignore all the future collections. This can be used before a POSIX
fork() call to make the gc copy-on-write friendly or to speed up collection.

5 files changed, 169 insertions, 2 deletions

diff --git a/Doc/library/gc.rst b/Doc/library/gc.rst
index 87d6824..153d8fb 100644
--- a/Doc/library/gc.rst
+++ b/Doc/library/gc.rst
@@ -174,6 +174,33 @@ The :mod:`gc` module provides the following functions:
    .. versionadded:: 3.1
 
 
+.. function:: freeze()
+
+   Freeze all the objects tracked by gc - move them to a permanent generation
+   and ignore all the future collections. This can be used before a POSIX
+   fork() call to make the gc copy-on-write friendly or to speed up collection.
+   Also collection before a POSIX fork() call may free pages for future
+   allocation which can cause copy-on-write too so it's advised to disable gc
+   in master process and freeze before fork and enable gc in child process.
+
+   .. versionadded:: 3.7
+
+
+.. function:: unfreeze()
+
+   Unfreeze the objects in the permanent generation, put them back into the
+   oldest generation.
+
+   .. versionadded:: 3.7
+
+
+.. function:: get_freeze_count()
+
+   Return the number of objects in the permanent generation.
+
+   .. versionadded:: 3.7
+
+
 The following variables are provided for read-only access (you can mutate the
 values but should not rebind them):
 
diff --git a/Include/internal/mem.h b/Include/internal/mem.h
index f3a8f56..471cdf4 100644
--- a/Include/internal/mem.h
+++ b/Include/internal/mem.h
@@ -167,6 +167,8 @@ struct _gc_runtime_state {
     /* linked lists of container objects */
     struct gc_generation generations[NUM_GENERATIONS];
     PyGC_Head *generation0;
+    /* a permanent generation which won't be collected */
+    struct gc_generation permanent_generation;
     struct gc_generation_stats generation_stats[NUM_GENERATIONS];
     /* true if we are currently running the collector */
     int collecting;
diff --git a/Lib/test/test_gc.py b/Lib/test/test_gc.py
index 914efec..904fc7d 100644
--- a/Lib/test/test_gc.py
+++ b/Lib/test/test_gc.py
@@ -734,6 +734,12 @@ class GCTests(unittest.TestCase):
         self.assertEqual(new[1]["collections"], old[1]["collections"])
         self.assertEqual(new[2]["collections"], old[2]["collections"] + 1)
 
+    def test_freeze(self):
+        gc.freeze()
+        self.assertGreater(gc.get_freeze_count(), 0)
+        gc.unfreeze()
+        self.assertEqual(gc.get_freeze_count(), 0)
+
 
 class GCCallbackTests(unittest.TestCase):
     def setUp(self):
diff --git a/Modules/clinic/gcmodule.c.h b/Modules/clinic/gcmodule.c.h
index 771c573..c31ef65 100644
--- a/Modules/clinic/gcmodule.c.h
+++ b/Modules/clinic/gcmodule.c.h
@@ -255,4 +255,74 @@ PyDoc_STRVAR(gc_is_tracked__doc__,
 
 #define GC_IS_TRACKED_METHODDEF    \
     {"is_tracked", (PyCFunction)gc_is_tracked, METH_O, gc_is_tracked__doc__},
-/*[clinic end generated code: output=5a58583f00ab018e input=a9049054013a1b77]*/
+
+PyDoc_STRVAR(gc_freeze__doc__,
+"freeze($module, /)\n"
+"--\n"
+"\n"
+"Freeze all current tracked objects and ignore them for future collections.\n"
+"\n"
+"This can be used before a POSIX fork() call to make the gc copy-on-write friendly.\n"
+"Note: collection before a POSIX fork() call may free pages for future allocation\n"
+"which can cause copy-on-write.");
+
+#define GC_FREEZE_METHODDEF    \
+    {"freeze", (PyCFunction)gc_freeze, METH_NOARGS, gc_freeze__doc__},
+
+static PyObject *
+gc_freeze_impl(PyObject *module);
+
+static PyObject *
+gc_freeze(PyObject *module, PyObject *Py_UNUSED(ignored))
+{
+    return gc_freeze_impl(module);
+}
+
+PyDoc_STRVAR(gc_unfreeze__doc__,
+"unfreeze($module, /)\n"
+"--\n"
+"\n"
+"Unfreeze all objects in the permanent generation.\n"
+"\n"
+"Put all objects in the permanent generation back into oldest generation.");
+
+#define GC_UNFREEZE_METHODDEF    \
+    {"unfreeze", (PyCFunction)gc_unfreeze, METH_NOARGS, gc_unfreeze__doc__},
+
+static PyObject *
+gc_unfreeze_impl(PyObject *module);
+
+static PyObject *
+gc_unfreeze(PyObject *module, PyObject *Py_UNUSED(ignored))
+{
+    return gc_unfreeze_impl(module);
+}
+
+PyDoc_STRVAR(gc_get_freeze_count__doc__,
+"get_freeze_count($module, /)\n"
+"--\n"
+"\n"
+"Return the number of objects in the permanent generation.");
+
+#define GC_GET_FREEZE_COUNT_METHODDEF    \
+    {"get_freeze_count", (PyCFunction)gc_get_freeze_count, METH_NOARGS, gc_get_freeze_count__doc__},
+
+static int
+gc_get_freeze_count_impl(PyObject *module);
+
+static PyObject *
+gc_get_freeze_count(PyObject *module, PyObject *Py_UNUSED(ignored))
+{
+    PyObject *return_value = NULL;
+    int _return_value;
+
+    _return_value = gc_get_freeze_count_impl(module);
+    if ((_return_value == -1) && PyErr_Occurred()) {
+        goto exit;
+    }
+    return_value = PyLong_FromLong((long)_return_value);
+
+exit:
+    return return_value;
+}
+/*[clinic end generated code: output=4f41ec4588154f2b input=a9049054013a1b77]*/
diff --git a/Modules/gcmodule.c b/Modules/gcmodule.c
index 832e27e..6e26c7a 100644
--- a/Modules/gcmodule.c
+++ b/Modules/gcmodule.c
@@ -71,6 +71,10 @@ _PyGC_Initialize(struct _gc_runtime_state *state)
         state->generations[i] = generations[i];
     };
     state->generation0 = GEN_HEAD(0);
+    struct gc_generation permanent_generation = {
+          {{&state->permanent_generation.head, &state->permanent_generation.head, 0}}, 0, 0
+    };
+    state->permanent_generation = permanent_generation;
 }
 
 /*--------------------------------------------------------------------------
@@ -813,6 +817,8 @@ collect(int generation, Py_ssize_t *n_collected, Py_ssize_t *n_uncollectable,
         for (i = 0; i < NUM_GENERATIONS; i++)
             PySys_FormatStderr(" %zd",
                               gc_list_size(GEN_HEAD(i)));
+        PySys_WriteStderr("\ngc: objects in permanent generation: %zd",
+                         gc_list_size(&_PyRuntime.gc.permanent_generation.head));
         t1 = _PyTime_GetMonotonicClock();
 
         PySys_WriteStderr("\n");
@@ -1405,6 +1411,56 @@ gc_is_tracked(PyObject *module, PyObject *obj)
     return result;
 }
 
+/*[clinic input]
+gc.freeze
+
+Freeze all current tracked objects and ignore them for future collections.
+
+This can be used before a POSIX fork() call to make the gc copy-on-write friendly.
+Note: collection before a POSIX fork() call may free pages for future allocation
+which can cause copy-on-write.
+[clinic start generated code]*/
+
+static PyObject *
+gc_freeze_impl(PyObject *module)
+/*[clinic end generated code: output=502159d9cdc4c139 input=b602b16ac5febbe5]*/
+{
+    for (int i = 0; i < NUM_GENERATIONS; ++i) {
+        gc_list_merge(GEN_HEAD(i), &_PyRuntime.gc.permanent_generation.head);
+        _PyRuntime.gc.generations[i].count = 0;
+    }
+    Py_RETURN_NONE;
+}
+
+/*[clinic input]
+gc.unfreeze
+
+Unfreeze all objects in the permanent generation.
+
+Put all objects in the permanent generation back into oldest generation.
+[clinic start generated code]*/
+
+static PyObject *
+gc_unfreeze_impl(PyObject *module)
+/*[clinic end generated code: output=1c15f2043b25e169 input=2dd52b170f4cef6c]*/
+{
+    gc_list_merge(&_PyRuntime.gc.permanent_generation.head, GEN_HEAD(NUM_GENERATIONS-1));
+    Py_RETURN_NONE;
+}
+
+/*[clinic input]
+gc.get_freeze_count -> int
+
+Return the number of objects in the permanent generation.
+[clinic start generated code]*/
+
+static int
+gc_get_freeze_count_impl(PyObject *module)
+/*[clinic end generated code: output=e4e2ebcc77e5cbf3 input=4b759db880a3c6e4]*/
+{
+    return gc_list_size(&_PyRuntime.gc.permanent_generation.head);
+}
+
 
 PyDoc_STRVAR(gc__doc__,
 "This module provides access to the garbage collector for reference cycles.\n"
@@ -1422,7 +1478,10 @@ PyDoc_STRVAR(gc__doc__,
 "get_objects() -- Return a list of all objects tracked by the collector.\n"
 "is_tracked() -- Returns true if a given object is tracked.\n"
 "get_referrers() -- Return the list of objects that refer to an object.\n"
-"get_referents() -- Return the list of objects that an object refers to.\n");
+"get_referents() -- Return the list of objects that an object refers to.\n"
+"freeze() -- Freeze all tracked objects and ignore them for future collections.\n"
+"unfreeze() -- Unfreeze all objects in the permanent generation.\n"
+"get_freeze_count() -- Return the number of objects in the permanent generation.\n");
 
 static PyMethodDef GcMethods[] = {
     GC_ENABLE_METHODDEF
@@ -1441,6 +1500,9 @@ static PyMethodDef GcMethods[] = {
         gc_get_referrers__doc__},
     {"get_referents",  gc_get_referents, METH_VARARGS,
         gc_get_referents__doc__},
+    GC_FREEZE_METHODDEF
+    GC_UNFREEZE_METHODDEF
+    GC_GET_FREEZE_COUNT_METHODDEF
     {NULL,      NULL}           /* Sentinel */
 };