Backport PR #22002: Fix TkAgg memory leaks and test for memory growth regressions

richardsheridan · tacaswell · commit ec6fbfb3f574 · 2022-04-28T20:06:34.000-04:00
FIX: TkAgg memory leaks and test for memory growth regressions (#22002) tkinter variables get cleaned up with normal `destroy` and `gc` semantics but tkinter's implementation of trace is effectively global and keeps the callback object alive until the trace is removed. Additionally extend and clean up the tests. Closes #20490 Co-authored-by: Elliott Sales de Andrade <quantum.analyst@gmail.com> (cherry picked from commit 1a016f0)
diff --git a/environment.yml b/environment.yml
@@ -48,6 +48,7 @@ dependencies:
   - nbformat!=5.0.0,!=5.0.1
   - pandas!=0.25.0
   - pikepdf
+  - psutil
   - pydocstyle>=5.1.0
   - pytest!=4.6.0,!=5.4.0
   - pytest-cov
diff --git a/lib/matplotlib/_pylab_helpers.py b/lib/matplotlib/_pylab_helpers.py
@@ -65,6 +65,10 @@ def destroy(cls, num):
         if hasattr(manager, "_cidgcf"):
             manager.canvas.mpl_disconnect(manager._cidgcf)
         manager.destroy()
+        del manager, num
+        # Full cyclic garbage collection may be too expensive to do on every
+        # figure destruction, so we collect only the youngest two generations.
+        # see: https://github.com/matplotlib/matplotlib/pull/3045
         gc.collect(1)
 
     @classmethod
diff --git a/lib/matplotlib/backends/_backend_tk.py b/lib/matplotlib/backends/_backend_tk.py
@@ -431,11 +431,12 @@ def __init__(self, canvas, num, window):
         # to store the DPI, which will be updated by the C code, and the trace
         # will handle it on the Python side.
         window_frame = int(window.wm_frame(), 16)
-        window_dpi = tk.IntVar(master=window, value=96,
-                               name=f'window_dpi{window_frame}')
+        self._window_dpi = tk.IntVar(master=window, value=96,
+                                     name=f'window_dpi{window_frame}')
+        self._window_dpi_cbname = ''
         if _tkagg.enable_dpi_awareness(window_frame, window.tk.interpaddr()):
-            self._window_dpi = window_dpi  # Prevent garbage collection.
-            window_dpi.trace_add('write', self._update_window_dpi)
+            self._window_dpi_cbname = self._window_dpi.trace_add(
+                'write', self._update_window_dpi)
 
         self._shown = False
 
@@ -489,20 +490,26 @@ def destroy(self, *args):
             self.canvas._tkcanvas.after_cancel(self.canvas._idle_draw_id)
         if self.canvas._event_loop_id:
             self.canvas._tkcanvas.after_cancel(self.canvas._event_loop_id)
+        if self._window_dpi_cbname:
+            self._window_dpi.trace_remove('write', self._window_dpi_cbname)
 
         # NOTE: events need to be flushed before issuing destroy (GH #9956),
-        # however, self.window.update() can break user code. This is the
-        # safest way to achieve a complete draining of the event queue,
-        # but it may require users to update() on their own to execute the
-        # completion in obscure corner cases.
+        # however, self.window.update() can break user code. An async callback
+        # is the safest way to achieve a complete draining of the event queue,
+        # but it leaks if no tk event loop is running. Therefore we explicitly
+        # check for an event loop and choose our best guess.
         def delayed_destroy():
             self.window.destroy()
 
             if self._owns_mainloop and not Gcf.get_num_fig_managers():
                 self.window.quit()
 
-        # "after idle after 0" avoids Tcl error/race (GH #19940)
-        self.window.after_idle(self.window.after, 0, delayed_destroy)
+        if cbook._get_running_interactive_framework() == "tk":
+            # "after idle after 0" avoids Tcl error/race (GH #19940)
+            self.window.after_idle(self.window.after, 0, delayed_destroy)
+        else:
+            self.window.update()
+            delayed_destroy()
 
     def get_window_title(self):
         return self.window.wm_title()
diff --git a/lib/matplotlib/tests/test_backend_tk.py b/lib/matplotlib/tests/test_backend_tk.py
@@ -45,8 +45,9 @@ def test_func():
             )
         except subprocess.TimeoutExpired:
             pytest.fail("Subprocess timed out")
-        except subprocess.CalledProcessError:
-            pytest.fail("Subprocess failed to test intended behavior")
+        except subprocess.CalledProcessError as e:
+            pytest.fail("Subprocess failed to test intended behavior\n"
+                        + str(e.stderr))
         else:
             # macOS may actually emit irrelevant errors about Accelerated
             # OpenGL vs. software OpenGL, so suppress them.
@@ -158,14 +159,12 @@ def test_never_update():  # pragma: no cover
 
     plt.draw()  # Test FigureCanvasTkAgg.
     fig.canvas.toolbar.configure_subplots()  # Test NavigationToolbar2Tk.
+    # Test FigureCanvasTk filter_destroy callback
+    fig.canvas.get_tk_widget().after(100, plt.close, fig)
 
     # Check for update() or update_idletasks() in the event queue, functionally
     # equivalent to tkinter.Misc.update.
-    # Must pause >= 1 ms to process tcl idle events plus extra time to avoid
-    # flaky tests on slow systems.
-    plt.pause(0.1)
-
-    plt.close(fig)  # Test FigureCanvasTk filter_destroy callback
+    plt.show(block=True)
 
     # Note that exceptions would be printed to stderr; _isolated_tk_test
     # checks them.
diff --git a/lib/matplotlib/tests/test_backends_interactive.py b/lib/matplotlib/tests/test_backends_interactive.py
@@ -195,8 +195,10 @@ def _test_thread_impl():
     future = ThreadPoolExecutor().submit(fig.canvas.draw)
     plt.pause(0.5)  # flush_events fails here on at least Tkagg (bpo-41176)
     future.result()  # Joins the thread; rethrows any exception.
-    plt.close()
-    fig.canvas.flush_events()  # pause doesn't process events after close
+    plt.close()  # backend is responsible for flushing any events here
+    if plt.rcParams["backend"].startswith("WX"):
+        # TODO: debug why WX needs this only on py3.8
+        fig.canvas.flush_events()
 
 
 _thread_safe_backends = _get_testable_interactive_backends()
@@ -413,3 +415,55 @@ def _lazy_headless():
 @pytest.mark.backend('QtAgg', skip_on_importerror=True)
 def test_lazy_linux_headless():
     proc = _run_helper(_lazy_headless, timeout=_test_timeout, MPLBACKEND="")
+
+
+# The source of this function gets extracted and run in another process, so it
+# must be fully self-contained.
+def _test_figure_leak():
+    import gc
+    import sys
+
+    import psutil
+    from matplotlib import pyplot as plt
+    # Second argument is pause length, but if zero we should skip pausing
+    t = float(sys.argv[1])
+    p = psutil.Process()
+
+    # Warmup cycle, this reasonably allocates a lot
+    for _ in range(2):
+        fig = plt.figure()
+        if t:
+            plt.pause(t)
+        plt.close(fig)
+    mem = p.memory_info().rss
+    gc.collect()
+
+    for _ in range(5):
+        fig = plt.figure()
+        if t:
+            plt.pause(t)
+        plt.close(fig)
+        gc.collect()
+    growth = p.memory_info().rss - mem
+
+    print(growth)
+
+
+# TODO: "0.1" memory threshold could be reduced 10x by fixing tkagg
+@pytest.mark.parametrize("env", _get_testable_interactive_backends())
+@pytest.mark.parametrize("time_mem", [(0.0, 2_000_000), (0.1, 30_000_000)])
+def test_figure_leak_20490(env, time_mem):
+    pytest.importorskip("psutil", reason="psutil needed to run this test")
+
+    # We haven't yet directly identified the leaks so test with a memory growth
+    # threshold.
+    pause_time, acceptable_memory_leakage = time_mem
+    if env["MPLBACKEND"] == "macosx":
+        acceptable_memory_leakage += 10_000_000
+
+    result = _run_helper(
+        _test_figure_leak, str(pause_time), timeout=_test_timeout, **env
+    )
+
+    growth = int(result.stdout)
+    assert growth <= acceptable_memory_leakage
diff --git a/requirements/testing/all.txt b/requirements/testing/all.txt
@@ -2,6 +2,7 @@
 
 certifi
 coverage<6.3
+psutil
 pytest!=4.6.0,!=5.4.0
 pytest-cov
 pytest-rerunfailures
