bpo-31356: Add context manager to temporarily disable GC (GH-4224)

1 files changed, 106 insertions, 0 deletions

diff --git a/Modules/gcmodule.c b/Modules/gcmodule.c
index 8ba1093..c057d25 100644
--- a/Modules/gcmodule.c
+++ b/Modules/gcmodule.c
@@ -1067,6 +1067,10 @@ static PyObject *
 gc_enable_impl(PyObject *module)
 /*[clinic end generated code: output=45a427e9dce9155c input=81ac4940ca579707]*/
 {
+    if(_PyRuntime.gc.disabled_threads){
+        PyErr_WarnEx(PyExc_RuntimeWarning, "Garbage collector enabled while another "
+            "thread is inside gc.ensure_enabled",1);
+    }
     _PyRuntime.gc.enabled = 1;
     Py_RETURN_NONE;
 }
@@ -1508,6 +1512,102 @@ static PyMethodDef GcMethods[] = {
     {NULL,      NULL}           /* Sentinel */
 };
 
+typedef struct {
+    PyObject_HEAD
+    int previous_gc_state;
+} ensure_disabled_object;
+
+
+static void
+ensure_disabled_object_dealloc(ensure_disabled_object *m_obj)
+{
+    Py_TYPE(m_obj)->tp_free((PyObject*)m_obj);
+}
+
+static PyObject *
+ensure_disabled__enter__method(ensure_disabled_object *self, PyObject *args)
+{
+    PyGILState_STATE gstate = PyGILState_Ensure();
+    ++_PyRuntime.gc.disabled_threads;
+    self->previous_gc_state = _PyRuntime.gc.enabled;
+    gc_disable_impl(NULL);
+    PyGILState_Release(gstate);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+ensure_disabled__exit__method(ensure_disabled_object *self, PyObject *args)
+{
+    PyGILState_STATE gstate = PyGILState_Ensure();
+    --_PyRuntime.gc.disabled_threads;
+    if(self->previous_gc_state){
+        gc_enable_impl(NULL);
+    }else{
+        gc_disable_impl(NULL);
+    }
+    PyGILState_Release(gstate);
+    Py_RETURN_NONE;
+}
+
+
+
+static struct PyMethodDef ensure_disabled_object_methods[] = {
+    {"__enter__",       (PyCFunction) ensure_disabled__enter__method,      METH_NOARGS},
+    {"__exit__",        (PyCFunction) ensure_disabled__exit__method,       METH_VARARGS},
+    {NULL,         NULL}       /* sentinel */
+};
+
+static PyObject *
+new_disabled_obj(PyTypeObject *type, PyObject *args, PyObject *kwdict){
+    ensure_disabled_object *self;
+    self = (ensure_disabled_object *)type->tp_alloc(type, 0);
+    return (PyObject *) self;
+};
+
+static PyTypeObject gc_ensure_disabled_type = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "gc.ensure_disabled",                       /* tp_name */
+    sizeof(ensure_disabled_object),             /* tp_size */
+    0,                                          /* tp_itemsize */
+    /* methods */
+    (destructor) ensure_disabled_object_dealloc,/* tp_dealloc */
+    0,                                          /* tp_print */
+    0,                                          /* tp_getattr */
+    0,                                          /* tp_setattr */
+    0,                                          /* tp_reserved */
+    0,                                          /* tp_repr */
+    0,                                          /* tp_as_number */
+    0,                                          /*tp_as_sequence*/
+    0,                                          /*tp_as_mapping*/
+    0,                                          /*tp_hash*/
+    0,                                          /*tp_call*/
+    0,                                          /*tp_str*/
+    PyObject_GenericGetAttr,                    /*tp_getattro*/
+    0,                                          /*tp_setattro*/
+    0,                                          /*tp_as_buffer*/
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,   /*tp_flags*/
+    0,                                          /*tp_doc*/
+    0,                                          /* tp_traverse */
+    0,                                          /* tp_clear */
+    0,                                          /* tp_richcompare */
+    0,                                          /* tp_weaklistoffset */
+    0,                                          /* tp_iter */
+    0,                                          /* tp_iternext */
+    ensure_disabled_object_methods,             /* tp_methods */
+    0,                                          /* tp_members */
+    0,                                          /* tp_getset */
+    0,                                          /* tp_base */
+    0,                                          /* tp_dict */
+    0,                                          /* tp_descr_get */
+    0,                                          /* tp_descr_set */
+    0,                                          /* tp_dictoffset */
+    0,                                          /* tp_init */
+    PyType_GenericAlloc,                        /* tp_alloc */
+    new_disabled_obj,                           /* tp_new */
+    PyObject_Del,                               /* tp_free */
+};
+
+
 static struct PyModuleDef gcmodule = {
     PyModuleDef_HEAD_INIT,
     "gc",              /* m_name */
@@ -1548,6 +1648,12 @@ PyInit_gc(void)
     if (PyModule_AddObject(m, "callbacks", _PyRuntime.gc.callbacks) < 0)
         return NULL;
 
+    if (PyType_Ready(&gc_ensure_disabled_type) < 0)
+        return NULL;
+    if (PyModule_AddObject(m, "ensure_disabled", (PyObject*) &gc_ensure_disabled_type) < 0)
+        return NULL;
+
+
 #define ADD_INT(NAME) if (PyModule_AddIntConstant(m, #NAME, NAME) < 0) return NULL
     ADD_INT(DEBUG_STATS);
     ADD_INT(DEBUG_COLLECTABLE);