Merge pull request #19480 from vallsv/fix-callback-registry-mem-leak

QuLogic · web-flow · commit 918478793712 · 2021-02-18T20:03:10.000-05:00
Fix CallbackRegistry memory leak
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -197,16 +197,17 @@ def __setstate__(self, state):
             s: {proxy: cid for cid, proxy in d.items()}
             for s, d in self.callbacks.items()}
 
-    def connect(self, s, func):
-        """Register *func* to be called when signal *s* is generated."""
-        self._func_cid_map.setdefault(s, {})
+    @_api.rename_parameter("3.4", "s", "signal")
+    def connect(self, signal, func):
+        """Register *func* to be called when signal *signal* is generated."""
+        self._func_cid_map.setdefault(signal, {})
         proxy = _weak_or_strong_ref(func, self._remove_proxy)
-        if proxy in self._func_cid_map[s]:
-            return self._func_cid_map[s][proxy]
+        if proxy in self._func_cid_map[signal]:
+            return self._func_cid_map[signal][proxy]
         cid = next(self._cid_gen)
-        self._func_cid_map[s][proxy] = cid
-        self.callbacks.setdefault(s, {})
-        self.callbacks[s][cid] = proxy
+        self._func_cid_map[signal][proxy] = cid
+        self.callbacks.setdefault(signal, {})
+        self.callbacks[signal][cid] = proxy
         return cid
 
     # Keep a reference to sys.is_finalizing, as sys may have been cleared out
@@ -215,32 +216,45 @@ def _remove_proxy(self, proxy, *, _is_finalizing=sys.is_finalizing):
         if _is_finalizing():
             # Weakrefs can't be properly torn down at that point anymore.
             return
-        for signal, proxies in list(self._func_cid_map.items()):
-            try:
-                del self.callbacks[signal][proxies[proxy]]
-            except KeyError:
-                pass
-            if len(self.callbacks[signal]) == 0:
-                del self.callbacks[signal]
-                del self._func_cid_map[signal]
+        for signal, proxy_to_cid in list(self._func_cid_map.items()):
+            cid = proxy_to_cid.pop(proxy, None)
+            if cid is not None:
+                del self.callbacks[signal][cid]
+                self._pickled_cids.discard(cid)
+                break
+        else:
+            # Not found
+            return
+        # Clean up empty dicts
+        if len(self.callbacks[signal]) == 0:
+            del self.callbacks[signal]
+            del self._func_cid_map[signal]
 
     def disconnect(self, cid):
         """
         Disconnect the callback registered with callback id *cid*.
 
         No error is raised if such a callback does not exist.
         """
-        for eventname, callbackd in list(self.callbacks.items()):
-            try:
-                del callbackd[cid]
-            except KeyError:
-                continue
-            else:
-                for signal, functions in list(self._func_cid_map.items()):
-                    for function, value in list(functions.items()):
-                        if value == cid:
-                            del functions[function]
-                return
+        self._pickled_cids.discard(cid)
+        # Clean up callbacks
+        for signal, cid_to_proxy in list(self.callbacks.items()):
+            proxy = cid_to_proxy.pop(cid, None)
+            if proxy is not None:
+                break
+        else:
+            # Not found
+            return
+
+        proxy_to_cid = self._func_cid_map[signal]
+        for current_proxy, current_cid in list(proxy_to_cid.items()):
+            if current_cid == cid:
+                assert proxy is current_proxy
+                del proxy_to_cid[current_proxy]
+        # Clean up empty dicts
+        if len(self.callbacks[signal]) == 0:
+            del self.callbacks[signal]
+            del self._func_cid_map[signal]
 
     def process(self, s, *args, **kwargs):
         """
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -182,31 +182,45 @@ def setup(self):
         self.signal = 'test'
         self.callbacks = cbook.CallbackRegistry()
 
-    def connect(self, s, func):
-        return self.callbacks.connect(s, func)
+    def connect(self, s, func, pickle):
+        cid = self.callbacks.connect(s, func)
+        if pickle:
+            self.callbacks._pickled_cids.add(cid)
+        return cid
+
+    def disconnect(self, cid):
+        return self.callbacks.disconnect(cid)
+
+    def count(self):
+        count1 = len(self.callbacks._func_cid_map.get(self.signal, []))
+        count2 = len(self.callbacks.callbacks.get(self.signal))
+        assert count1 == count2
+        return count1
 
     def is_empty(self):
         assert self.callbacks._func_cid_map == {}
         assert self.callbacks.callbacks == {}
+        assert self.callbacks._pickled_cids == set()
 
     def is_not_empty(self):
         assert self.callbacks._func_cid_map != {}
         assert self.callbacks.callbacks != {}
 
-    def test_callback_complete(self):
+    @pytest.mark.parametrize('pickle', [True, False])
+    def test_callback_complete(self, pickle):
         # ensure we start with an empty registry
         self.is_empty()
 
         # create a class for testing
         mini_me = Test_callback_registry()
 
         # test that we can add a callback
-        cid1 = self.connect(self.signal, mini_me.dummy)
+        cid1 = self.connect(self.signal, mini_me.dummy, pickle)
         assert type(cid1) == int
         self.is_not_empty()
 
         # test that we don't add a second callback
-        cid2 = self.connect(self.signal, mini_me.dummy)
+        cid2 = self.connect(self.signal, mini_me.dummy, pickle)
         assert cid1 == cid2
         self.is_not_empty()
         assert len(self.callbacks._func_cid_map) == 1
@@ -217,6 +231,68 @@ def test_callback_complete(self):
         # check we now have no callbacks registered
         self.is_empty()
 
+    @pytest.mark.parametrize('pickle', [True, False])
+    def test_callback_disconnect(self, pickle):
+        # ensure we start with an empty registry
+        self.is_empty()
+
+        # create a class for testing
+        mini_me = Test_callback_registry()
+
+        # test that we can add a callback
+        cid1 = self.connect(self.signal, mini_me.dummy, pickle)
+        assert type(cid1) == int
+        self.is_not_empty()
+
+        self.disconnect(cid1)
+
+        # check we now have no callbacks registered
+        self.is_empty()
+
+    @pytest.mark.parametrize('pickle', [True, False])
+    def test_callback_wrong_disconnect(self, pickle):
+        # ensure we start with an empty registry
+        self.is_empty()
+
+        # create a class for testing
+        mini_me = Test_callback_registry()
+
+        # test that we can add a callback
+        cid1 = self.connect(self.signal, mini_me.dummy, pickle)
+        assert type(cid1) == int
+        self.is_not_empty()
+
+        self.disconnect("foo")
+
+        # check we still have callbacks registered
+        self.is_not_empty()
+
+    @pytest.mark.parametrize('pickle', [True, False])
+    def test_registration_on_non_empty_registry(self, pickle):
+        # ensure we start with an empty registry
+        self.is_empty()
+
+        # setup the registry with a callback
+        mini_me = Test_callback_registry()
+        self.connect(self.signal, mini_me.dummy, pickle)
+
+        # Add another callback
+        mini_me2 = Test_callback_registry()
+        self.connect(self.signal, mini_me2.dummy, pickle)
+
+        # Remove and add the second callback
+        mini_me2 = Test_callback_registry()
+        self.connect(self.signal, mini_me2.dummy, pickle)
+
+        # We still have 2 references
+        self.is_not_empty()
+        assert self.count() == 2
+
+        # Removing the last 2 references
+        mini_me = None
+        mini_me2 = None
+        self.is_empty()
+
     def dummy(self):
         pass
 
