fix: Memory leak from streams rooted in the global registry

holoviews/core/dimension.py

@@ -566,6 +566,7 @@ class LabelledData(param.Parameterized):
     )
 
     _deep_indexable = False
+    _streams = None
 
     def __init__(self, data, id=None, plot_id=None, **params):
         """All LabelledData subclasses must supply data to the
@@ -847,6 +848,7 @@ def map(self, map_fn, specs=None, clone=True):
     def __getstate__(self):
         """Ensures pickles save options applied to this objects."""
         obj_dict = self.__dict__.copy()
+        obj_dict.pop("_streams", None)
         try:
             if Store.save_option_state and (obj_dict.get("_id", None) is not None):
                 custom_key = "_custom_option_{}".format(obj_dict["_id"])

holoviews/plotting/bokeh/plot.py

@@ -267,6 +267,7 @@ def cleanup(self):
         """
         plots = self.traverse(lambda x: x, [BokehPlot])
         for plot in plots:
+            plot._unwatch_session()
             if not isinstance(
                 plot, (GenericCompositePlot, GenericElementPlot, GenericOverlayPlot)
             ):
@@ -289,11 +290,7 @@ def cleanup(self):
                 stream._subscribers = [
                     (p, subscriber)
                     for p, subscriber in stream._subscribers
-                    if subscriber
-                    and (
-                        not is_param_method(subscriber)
-                        or get_method_owner(subscriber) not in plots
-                    )
+                    if not is_param_method(subscriber) or get_method_owner(subscriber) not in plots
                 ]
 
     def _fontsize(self, key, label="fontsize", common=True):
@@ -467,7 +464,7 @@ def _link_dimensioned_streams(self):
         """
         streams = [s for s in self.streams if any(k in self.dimensions for k in s.contents)]
         for s in streams:
-            s.add_subscriber(self._stream_update, 1)
+            s.add_subscriber(self._stream_update, 1.05)
 
     def _stream_update(self, **kwargs):
         contents = [k for s in self.streams for k in s.contents]

holoviews/plotting/mpl/plot.py

@@ -356,7 +356,7 @@ def _link_dimensioned_streams(self):
         """
         streams = [s for s in self.streams if any(k in self.dimensions for k in s.contents)]
         for s in streams:
-            s.add_subscriber(self._stream_update, 1)
+            s.add_subscriber(self._stream_update, 1.05)
 
     def _stream_update(self, **kwargs):
         contents = [k for s in self.streams for k in s.contents]

holoviews/plotting/plot.py

@@ -109,29 +109,38 @@ def document(self):
 
     @document.setter
     def document(self, doc):
+        self._unwatch_session()
         if (
             doc
             and hasattr(doc, "on_session_destroyed")
             and self.root is self.handles.get("plot")
             and not isinstance(self, GenericAdjointLayoutPlot)
         ):
             doc.on_session_destroyed(self._session_destroy)
-            if self._document:
-                if isinstance(self._document.callbacks._session_destroyed_callbacks, set):
-                    self._document.callbacks._session_destroyed_callbacks.discard(
-                        self._session_destroy
-                    )
-                else:
-                    self._document.callbacks._session_destroyed_callbacks.pop(
-                        self._session_destroy, None
-                    )
 
         self._document = doc
         if self.subplots:
             for plot in self.subplots.values():
                 if plot is not None:
                     plot.document = doc
 
+    def _unwatch_session(self):
+        """Drop this plot's session destroy hook from its current document.
+
+        A document holds its callbacks, and through them this plot, strongly,
+        so the hook is dropped whenever the plot leaves the document, i.e. on
+        reassignment and on cleanup. Only the hook belonging to this plot is
+        removed.
+        """
+        doc = self._document
+        if doc is None or not hasattr(doc, "callbacks"):
+            return
+        callbacks = doc.callbacks._session_destroyed_callbacks
+        if isinstance(callbacks, set):
+            callbacks.discard(self._session_destroy)
+        else:
+            callbacks.pop(self._session_destroy, None)
+
     @property
     def pane(self):
         return self._pane
@@ -195,17 +204,15 @@ def cleanup(self):
         """
         plots = self.traverse(lambda x: x, [Plot])
         for plot in plots:
+            plot._unwatch_session()
             if not isinstance(plot, (GenericElementPlot, GenericOverlayPlot)):
                 continue
             for stream in set(plot.streams):
                 stream._subscribers = [
                     (p, subscriber)
                     for p, subscriber in stream._subscribers
-                    if subscriber
-                    and (
-                        not util.is_param_method(subscriber)
-                        or util.get_method_owner(subscriber) not in plots
-                    )
+                    if not util.is_param_method(subscriber)
+                    or util.get_method_owner(subscriber) not in plots
                 ]
 
     def _session_destroy(self, session_context):

holoviews/plotting/renderer.py

@@ -7,6 +7,7 @@
 
 import base64
 import os
+import weakref
 from contextlib import contextmanager
 from functools import partial
 from io import BytesIO, StringIO
@@ -234,9 +235,17 @@ class Renderer(Exporter):
     _render_with_panel = False
 
     def __init__(self, **params):
-        self.last_plot = None
+        self._last_plot = None
         super().__init__(**params)
 
+    @property
+    def last_plot(self):
+        return None if self._last_plot is None else self._last_plot()
+
+    @last_plot.setter
+    def last_plot(self, plot):
+        self._last_plot = None if plot is None else weakref.ref(plot)
+
     def __call__(self, obj, fmt="auto", **kwargs):
         plot, fmt = self._validate(obj, fmt)
         info = {"file-ext": fmt, "mime_type": MIME_TYPES[fmt]}

holoviews/plotting/util.py

@@ -1163,7 +1163,7 @@ def attach_streams(plot, obj, precedence=1.1):
 
     def append_refresh(dmap):
         for stream in get_nested_streams(dmap):
-            if plot.refresh not in stream._subscribers:
+            if plot.refresh not in stream.subscribers:
                 stream.add_subscriber(plot.refresh, precedence)
 
     return obj.traverse(append_refresh, [DynamicMap])

holoviews/streams.py

@@ -6,7 +6,6 @@
 
 from __future__ import annotations
 
-import inspect
 import typing as t
 import weakref
 from collections import defaultdict
@@ -54,47 +53,50 @@ class _SkipTrigger:
     pass
 
 
-class _WeakSubscriber:
-    def __init__(self, subscriber):
-        if inspect.ismethod(subscriber):
-            self._ref = weakref.WeakMethod(subscriber)
-        else:
-            self._ref = weakref.ref(subscriber)
+class _StreamRegistry(weakref.WeakKeyDictionary):
+    """Weak index from a source object to the streams sourced from it.
 
-        self._hash = self._generate_hash(subscriber)
-        self._is_param_method = util.is_param_method(subscriber)
+    Both the keys and the values are weak, so the registry never keeps a
+    source or a stream alive. Ownership lives on the source, in
+    ``LabelledData._streams``, and a source and its streams are collected
+    together once nothing else refers to them.
 
-    def __bool__(self) -> bool:
-        method = self._ref()
-        return method is not None and not self._is_param_method
+    Reads resolve the references, i.e. this behaves like a mapping of source
+    to a list of live streams.
+    """
 
-    def __call__(self, *args, **kwargs):
-        method = self._ref()
-        if method is not None:
-            return method(*args, **kwargs)
+    def register(self, source, stream):
+        # super().get to see the stored references rather than live streams
+        refs = super().get(source)
+        if refs is None:
+            self[source] = refs = []
+        if not any(ref() is stream for ref in refs):
+            refs.append(weakref.ref(stream))
+
+    def unregister(self, source, stream):
+        refs = super().get(source)
+        if refs is None:
+            return
+        refs[:] = [ref for ref in refs if ref() is not None and ref() is not stream]
+        if not refs:
+            self.pop(source, None)
 
     @staticmethod
-    def check(subscriber):
-        while isinstance(subscriber, partial):
-            subscriber = subscriber.func
-        return inspect.ismethod(subscriber)
+    def _alive(refs):
+        return [stream for stream in (ref() for ref in refs) if stream is not None]
 
-    @staticmethod
-    def _generate_hash(subscriber) -> int:
-        if inspect.ismethod(subscriber):
-            return hash((id(subscriber.__func__), id(subscriber.__self__)))
-        else:
-            return hash(id(subscriber))
+    def __getitem__(self, key):
+        return self._alive(super().__getitem__(key))
+
+    def get(self, key, default=None):
+        refs = super().get(key)
+        return default if refs is None else self._alive(refs)
 
-    def __eq__(self, other) -> bool:
-        if isinstance(other, _WeakSubscriber):
-            return self._hash == other._hash
-        if self.check(other):
-            return self._hash == self._generate_hash(other)
-        return NotImplemented
+    def items(self):
+        return [(src, self._alive(refs)) for src, refs in super().items()]
 
-    def __hash__(self) -> int:
-        return self._hash
+    def values(self):
+        return [self._alive(refs) for refs in super().values()]
 
 
 @contextmanager
@@ -166,10 +168,10 @@ class Stream(param.Parameterized):
 
     """
 
-    # Mapping from a source to a list of streams
-    # WeakKeyDictionary to allow garbage collection
-    # of unreferenced sources
-    registry = weakref.WeakKeyDictionary()
+    # Weak index from a source to the streams sourced from it, see
+    # _StreamRegistry. The sources own their streams, this only allows them
+    # to be looked up from a clone sharing the same _plot_id.
+    registry = _StreamRegistry()
 
     # Mapping to define callbacks by backend and Stream type.
     # e.g. Stream._callbacks['bokeh'][Stream] = Callback
@@ -370,10 +372,28 @@ def __init__(
         super().__init__(**params)
         self._rename = self._validate_rename(rename)
         if source is not None:
-            if source in self.registry:
-                self.registry[source].append(self)
-            else:
-                self.registry[source] = [self]
+            self._attach(source)
+
+    def _attach(self, source):
+        """Make ``source`` the owner of this stream.
+
+        The source holds the stream strongly, so a stream stays alive for as
+        long as the object it observes and no longer, while the registry only
+        indexes it weakly.
+        """
+        streams = source._streams
+        if streams is None:
+            source._streams = streams = []
+        if self not in streams:
+            streams.append(self)
+        self.registry.register(source, self)
+
+    def _detach(self, source):
+        """Undo ``_attach``, dropping the source's ownership of this stream."""
+        streams = source._streams
+        if streams is not None and self in streams:
+            streams.remove(self)
+        self.registry.unregister(source, self)
 
     def clone(self):
         """Return new stream with identical properties and no subscribers"""
@@ -426,8 +446,6 @@ def add_subscriber(self, subscriber, precedence=0):
         """
         if not callable(subscriber):
             raise TypeError("Subscriber must be a callable.")
-        if _WeakSubscriber.check(subscriber):
-            subscriber = _WeakSubscriber(subscriber)
         self._subscribers.append((precedence, subscriber))
 
     def _validate_rename(self, mapping):
@@ -464,21 +482,14 @@ def source(self):
     @source.setter
     def source(self, source):
         if self.source is not None:
-            source_list = self.registry[self.source]
-            if self in source_list:
-                source_list.remove(self)
-            if not source_list:
-                self.registry.pop(self.source)
+            self._detach(self.source)
 
         if source is None:
             self._source = None
             return
 
         self._source = weakref.ref(source)
-        if source in self.registry:
-            self.registry[source].append(self)
-        else:
-            self.registry[source] = [self]
+        self._attach(source)
 
     def transform(self):
         """Method that can be overwritten by subclasses to process the

holoviews/tests/core/test_dynamic.py

@@ -918,41 +918,23 @@ def test_dynamic_rx(self):
 
 
 class TestStreamSubscribersAddandClear:
-    def setup_method(self):
+    @pytest.mark.parametrize(
+        ("policy", "remaining"),
+        [("all", []), ("user", ["fn3", "fn4"]), ("internal", ["fn1", "fn2"])],
+    )
+    def test_subscriber_clear(self, policy, remaining):
         self.fn1 = lambda x: x
         self.fn2 = lambda x: x**2
         self.fn3 = lambda x: x**3
         self.fn4 = lambda x: x**4
-
-    def test_subscriber_clear_all(self):
-        pointerx = PointerX(x=2)
-        pointerx.add_subscriber(self.fn1, precedence=0)
-        pointerx.add_subscriber(self.fn2, precedence=1)
-        pointerx.add_subscriber(self.fn3, precedence=1.5)
-        pointerx.add_subscriber(self.fn4, precedence=10)
-        assert pointerx.subscribers == [self.fn1, self.fn2, self.fn3, self.fn4]
-        pointerx.clear("all")
-        assert pointerx.subscribers == []
-
-    def test_subscriber_clear_user(self):
-        pointerx = PointerX(x=2)
-        pointerx.add_subscriber(self.fn1, precedence=0)
-        pointerx.add_subscriber(self.fn2, precedence=1)
-        pointerx.add_subscriber(self.fn3, precedence=1.5)
-        pointerx.add_subscriber(self.fn4, precedence=10)
-        assert pointerx.subscribers == [self.fn1, self.fn2, self.fn3, self.fn4]
-        pointerx.clear("user")
-        assert pointerx.subscribers == [self.fn3, self.fn4]
-
-    def test_subscriber_clear_internal(self):
         pointerx = PointerX(x=2)
         pointerx.add_subscriber(self.fn1, precedence=0)
         pointerx.add_subscriber(self.fn2, precedence=1)
         pointerx.add_subscriber(self.fn3, precedence=1.5)
         pointerx.add_subscriber(self.fn4, precedence=10)
         assert pointerx.subscribers == [self.fn1, self.fn2, self.fn3, self.fn4]
-        pointerx.clear("internal")
-        assert pointerx.subscribers == [self.fn1, self.fn2]
+        pointerx.clear(policy)
+        assert pointerx.subscribers == [getattr(self, name) for name in remaining]
 
 
 class TestDynamicStreamReset: