Python 3.11: Fix a memory leak switching to greenlets.

This leak was present in the original implementation of Python 3.11 support (#306) Fixes #328 and fixes gevent/gevent#1924 ; I have run gevent's test suite locally on 3.11 with this fix without seeing any regressions.
python-greenlet · Nov 7, 2022 · aa6f251 · aa6f251
1 parent 7389976
commit aa6f251
Show file tree

Hide file tree

Showing 7 changed files with 268 additions and 8 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -5,7 +5,9 @@
 2.0.1 (unreleased)
 ==================
 
-- Nothing changed yet.
+- Python 3.11: Fix a memory leak. See `issue 328
+  <https://github.com/python-greenlet/greenlet/issues/328>`_ and
+  `gevent issue 1924 <https://github.com/gevent/gevent/issues/1924>`_.
 
 
 2.0.0.post0 (2022-11-03)
@@ -154,13 +156,18 @@ Changes
 
 - Add musllinux (Alpine) binary wheels.
 
+.. important:: This preliminary support for Python 3.11 leaks memory.
+               Please upgrade to greenlet 2 if you're using Python 3.11.
 
 1.1.3 (2022-08-25)
 ==================
 
 - Add support for Python 3.11. Please note that Windows binary wheels
   are not available at this time.
 
+.. important:: This preliminary support for Python 3.11 leaks memory.
+               Please upgrade to greenlet 2 if you're using Python 3.11.
+
 1.1.2 (2021-09-29)
 ==================
 

diff --git a/setup.py b/setup.py
@@ -226,7 +226,8 @@ def get_greenlet_version():
         ],
         'test': [
             'objgraph',
-             'faulthandler; python_version == "2.7" and platform_python_implementation == "CPython"',
+            'faulthandler; python_version == "2.7" and platform_python_implementation == "CPython"',
+            'psutil',
         ],
     },
     python_requires=">=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*",

diff --git a/src/greenlet/greenlet.cpp b/src/greenlet/greenlet.cpp
@@ -1438,12 +1438,15 @@ UserGreenlet::inner_bootstrap(OwnedGreenlet& origin_greenlet, OwnedObject& _run)
         result = single_result(result);
     }
     this->release_args();
+    this->python_state.did_finish(PyThreadState_GET());
 
     result = g_handle_exit(result);
     assert(this->thread_state()->borrow_current() == this->_self);
+
     /* jump back to parent */
     this->stack_state.set_inactive(); /* dead */
 
+
     // TODO: Can we decref some things here? Release our main greenlet
     // and maybe parent?
     for (Greenlet* parent = this->_parent;
@@ -2029,7 +2032,6 @@ UserGreenlet::tp_clear()
     this->_parent.CLEAR();
     this->_main_greenlet.CLEAR();
     this->_run_callable.CLEAR();
-
     return 0;
 }
 
@@ -2140,6 +2142,10 @@ Greenlet::~Greenlet()
 
 UserGreenlet::~UserGreenlet()
 {
+    // Python 3.11: If we don't clear out the raw frame datastack
+    // when deleting an unfinished greenlet,
+    // TestLeaks.test_untracked_memory_doesnt_increase_unfinished_thread_dealloc_in_main fails.
+    this->python_state.did_finish(nullptr);
     this->tp_clear();
 }
 

diff --git a/src/greenlet/greenlet_greenlet.hpp b/src/greenlet/greenlet_greenlet.hpp
@@ -146,10 +146,10 @@ namespace greenlet
         int recursion_depth;
         int trash_delete_nesting;
 #if GREENLET_PY311
-        _PyInterpreterFrame *current_frame;
-        _PyStackChunk *datastack_chunk;
-        PyObject **datastack_top;
-        PyObject **datastack_limit;
+        _PyInterpreterFrame* current_frame;
+        _PyStackChunk* datastack_chunk;
+        PyObject** datastack_top;
+        PyObject** datastack_limit;
 #endif
 
     public:
@@ -170,6 +170,7 @@ namespace greenlet
         void set_new_cframe(_PyCFrame& frame) G_NOEXCEPT;
 #endif
         void will_switch_from(PyThreadState *const origin_tstate) G_NOEXCEPT;
+        void did_finish(PyThreadState* tstate) G_NOEXCEPT;
     };
 
     class StackState
@@ -213,9 +214,13 @@ namespace greenlet
         inline intptr_t stack_saved() const G_NOEXCEPT;
         inline char* stack_start() const G_NOEXCEPT;
         static inline StackState make_main() G_NOEXCEPT;
+#ifdef GREENLET_USE_STDIO
         friend std::ostream& operator<<(std::ostream& os, const StackState& s);
+#endif
     };
+#ifdef GREENLET_USE_STDIO
     std::ostream& operator<<(std::ostream& os, const StackState& s);
+#endif
 
     class SwitchingArgs
     {
@@ -933,14 +938,97 @@ void PythonState::set_new_cframe(_PyCFrame& frame) G_NOEXCEPT
 }
 #endif
 
-
 const PythonState::OwnedFrame& PythonState::top_frame() const G_NOEXCEPT
 {
     return this->_top_frame;
 }
 
+void PythonState::did_finish(PyThreadState* tstate) G_NOEXCEPT
+{
+#if GREENLET_PY311
+    // See https://github.com/gevent/gevent/issues/1924 and
+    // https://github.com/python-greenlet/greenlet/issues/328. In
+    // short, Python 3.11 allocates memory for frames as a sort of
+    // linked list that's kept as part of PyThreadState in the
+    // ``datastack_chunk`` member and friends. These are saved and
+    // restored as part of switching greenlets.
+    //
+    // When we initially switch to a greenlet, we set those to NULL.
+    // That causes the frame management code to treat this like a
+    // brand new thread and start a fresh list of chunks, beginning
+    // with a new "root" chunk. As we make calls in this greenlet,
+    // those chunks get added, and as calls return, they get popped.
+    // But the frame code (pystate.c) is careful to make sure that the
+    // root chunk never gets popped.
+    //
+    // Thus, when a greenlet exits for the last time, there will be at
+    // least a single root chunk that we must be responsible for
+    // deallocating.
+    //
+    // The complex part is that these chunks are allocated and freed
+    // using ``_PyObject_VirtualAlloc``/``Free``. Those aren't public
+    // functions, and they aren't exported for linking. It so happens
+    // that we know they are just thin wrappers around the Arena
+    // allocator, so we can use that directly to deallocate in a
+    // compatible way.
+    //
+    // CAUTION: Check this implementation detail on every major version.
+    //
+    // It might be nice to be able to do this in our destructor, but
+    // can we be sure that no one else is using that memory? Plus, as
+    // described below, our pointers may not even be valid anymore. As
+    // a special case, there is one time that we know we can do this,
+    // and that's from the destructor of the associated UserGreenlet
+    // (NOT main greenlet)
+    PyObjectArenaAllocator alloc;
+    _PyStackChunk* chunk = nullptr;
+    if (tstate) {
+        // We really did finish, we can never be switched to again.
+        chunk = tstate->datastack_chunk;
+        // Unfortunately, we can't do much sanity checking. Our
+        // this->datastack_chunk pointer is out of date (evaluation may
+        // have popped down through it already) so we can't verify that
+        // we deallocate it. I don't think we can even check datastack_top
+        // for the same reason.
+
+        PyObject_GetArenaAllocator(&alloc);
+        tstate->datastack_chunk = nullptr;
+        tstate->datastack_limit = nullptr;
+        tstate->datastack_top = nullptr;
+
+    }
+    else if (this->datastack_chunk) {
+        // The UserGreenlet (NOT the main greenlet!) is being deallocated. If we're
+        // still holding a stack chunk, it's garbage because we know
+        // we can never switch back to let cPython clean it up.
+        // Because the last time we got switched away from, and we
+        // haven't run since then, we know our chain is valid and can
+        // be dealloced.
+        chunk = this->datastack_chunk;
+        PyObject_GetArenaAllocator(&alloc);
+    }
+
+    if (alloc.free && chunk) {
+        // In case the arena mechanism has been torn down already.
+        while (chunk) {
+            _PyStackChunk *prev = chunk->previous;
+            chunk->previous = nullptr;
+            alloc.free(alloc.ctx, chunk, chunk->size);
+            chunk = prev;
+        }
+    }
+
+    this->datastack_chunk = nullptr;
+    this->datastack_limit = nullptr;
+    this->datastack_top = nullptr;
+#endif
+}
+
+
+
 
 using greenlet::StackState;
+
 #ifdef GREENLET_USE_STDIO
 #include <iostream>
 using std::cerr;

diff --git a/src/greenlet/tests/__init__.py b/src/greenlet/tests/__init__.py
@@ -45,6 +45,7 @@ def __new__(cls, classname, bases, classDict):
                 classDict[key] = value
         return type.__new__(cls, classname, bases, classDict)
 
+
 class TestCase(TestCaseMetaClass(
         "NewBase",
         (unittest.TestCase,),

diff --git a/src/greenlet/tests/test_cpp.py b/src/greenlet/tests/test_cpp.py
@@ -11,6 +11,8 @@
 def run_unhandled_exception_in_greenlet_aborts():
     # This is used in multiprocessing.Process and must be picklable
     # so it needs to be global.
+
+
     def _():
         _test_extension_cpp.test_exception_switch_and_do_in_g2(
             _test_extension_cpp.test_exception_throw
@@ -63,3 +65,7 @@ def test_unhandled_exception_aborts(self):
     def test_unhandled_exception_in_greenlet_aborts(self):
         # verify that unhandled throw called in greenlet aborts too
         self._do_test_unhandled_exception(run_unhandled_exception_in_greenlet_aborts)
+
+
+if __name__ == '__main__':
+    __import__('unittest').main()