fix: Weakly reference callables for pub-sub model (#1035)

* Fix memory leaks seen when creating many objects that use pub-sub events API
* Store weaker references to bound methods/callables in the events API

Author: kratsg

src/pyhf/events.py

@@ -1,3 +1,5 @@
+import weakref
+
 __events = {}
 __disabled_events = set([])
 
@@ -6,10 +8,16 @@ def noop(*args, **kwargs):
     pass
 
 
-class Callables(list):
+class WeakList(list):
+    def append(self, item):
+        list.append(self, weakref.WeakMethod(item, self.remove))
+
+
+class Callables(WeakList):
     def __call__(self, *args, **kwargs):
-        for f in self:
-            f(*args, **kwargs)
+        for func in self:
+            # weakref: needs to be de-ref'd first before calling
+            func()(*args, **kwargs)
 
     def __repr__(self):
         return "Callables(%s)" % list.__repr__(self)

tests/test_events.py

@@ -6,31 +6,43 @@ def test_subscribe_event():
     ename = 'test'
 
     m = mock.Mock()
-    events.subscribe(ename)(m)
+    events.subscribe(ename)(m.__call__)
 
     assert ename in events.__events
-    assert m in events.__events.get(ename)
+    assert m.__call__ == events.__events.get(ename)[0]()
     del events.__events[ename]
 
 
 def test_event():
     ename = 'test'
 
     m = mock.Mock()
-    events.subscribe(ename)(m)
+    events.subscribe(ename)(m.__call__)
 
     events.trigger(ename)()
     m.assert_called_once()
     del events.__events[ename]
 
 
+def test_event_weakref():
+    ename = 'test'
+
+    m = mock.Mock()
+    events.subscribe(ename)(m.__call__)
+    assert len(events.trigger(ename)) == 1
+    # should be weakly referenced
+    del m
+    assert len(events.trigger(ename)) == 0
+    del events.__events[ename]
+
+
 def test_disable_event():
     ename = 'test'
 
     m = mock.Mock()
     noop, noop_m = events.noop, mock.Mock()
     events.noop = noop_m
-    events.subscribe(ename)(m)
+    events.subscribe(ename)(m.__call__)
 
     events.disable(ename)
     assert m.called is False

tests/test_interpolate.py

@@ -1,7 +1,6 @@
 import pyhf
 import numpy as np
 import pytest
-import mock
 
 
 @pytest.fixture
@@ -66,25 +65,27 @@ def test_interpolator_structure(interpcode, random_histosets_alphasets_pair):
     )
 
 
-def test_interpolator_subscription(interpcode, random_histosets_alphasets_pair):
+def test_interpolator_subscription(mocker, interpcode, random_histosets_alphasets_pair):
     histogramssets, alphasets = random_histosets_alphasets_pair
     ename = 'tensorlib_changed'
 
-    # inject into our interpolator class
     interpolator_cls = pyhf.interpolators.get(interpcode)
-    with mock.patch('{0:s}._precompute'.format(interpolator_cls.__module__)) as m:
-        interpolator_cls(histogramssets.tolist(), subscribe=False)
-        assert m.call_count == 1
-        assert m not in pyhf.events.__events.get(ename, [])
-        pyhf.events.trigger(ename)()
-        assert m.call_count == 1
-
-    with mock.patch('{0:s}._precompute'.format(interpolator_cls.__module__)) as m:
-        interpolator_cls(histogramssets.tolist(), subscribe=True)
-        assert m.call_count == 1
-        assert m in pyhf.events.__events.get(ename, [])
-        pyhf.events.trigger(ename)()
-        assert m.call_count == 2
+    spy = mocker.spy(interpolator_cls, '_precompute')
+
+    interpolator_cls(histogramssets.tolist(), subscribe=False)
+    assert spy.call_count == 1
+    assert len(pyhf.events.__events.get(ename, [])) == 0
+    pyhf.events.trigger(ename)()
+    assert spy.call_count == 1
+
+    # reset and go again
+    spy.reset_mock()
+
+    interpolator_cls(histogramssets.tolist(), subscribe=True)
+    assert spy.call_count == 1
+    assert len(pyhf.events.__events.get(ename, [])) == 1
+    pyhf.events.trigger(ename)()
+    assert spy.call_count == 2
 
 
 def test_interpolator_alphaset_change(

tests/test_tensor.py

@@ -429,7 +429,7 @@ def test_trigger_tensorlib_changed_name(mocker):
     pyhf.set_backend(numpy_64)
 
     func = mocker.Mock()
-    pyhf.events.subscribe('tensorlib_changed')(func)
+    pyhf.events.subscribe('tensorlib_changed')(func.__call__)
 
     assert func.call_count == 0
     pyhf.set_backend(jax_64)
@@ -443,7 +443,7 @@ def test_trigger_tensorlib_changed_precision(mocker):
     pyhf.set_backend(jax_64)
 
     func = mocker.Mock()
-    pyhf.events.subscribe('tensorlib_changed')(func)
+    pyhf.events.subscribe('tensorlib_changed')(func.__call__)
 
     assert func.call_count == 0
     pyhf.set_backend(jax_32)