From 9d09241fab868a1b7fd87a63b438ff39b093a77a Mon Sep 17 00:00:00 2001
From: Philipp Rudiger <P.Rudiger@ed.ac.uk>
Date: Fri, 26 Oct 2018 15:20:24 +0100
Subject: [PATCH 1/7] Ensure cleanup on session destroy

---
 holoviews/plotting/bokeh/plot.py    | 9 +++++++++
 holoviews/plotting/bokeh/widgets.py | 3 +++
 2 files changed, 12 insertions(+)

diff --git a/holoviews/plotting/bokeh/plot.py b/holoviews/plotting/bokeh/plot.py
index 91494945d0..554e5ee4d0 100644
--- a/holoviews/plotting/bokeh/plot.py
+++ b/holoviews/plotting/bokeh/plot.py
@@ -103,12 +103,21 @@ def root(self):
 
     @document.setter
     def document(self, doc):
+        if hasattr(doc, 'on_session_destroyed') and self.root is self.handles.get('plot'):
+            doc.on_session_destroyed(self._session_destroy)
+            if self._document:
+                self._document._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 _session_destroy(self, session_context):
+        self.cleanup()
+        self._document = None
+
 
     def __init__(self, *args, **params):
         root = params.pop('root', None)
diff --git a/holoviews/plotting/bokeh/widgets.py b/holoviews/plotting/bokeh/widgets.py
index e79403a36e..b62f3064d2 100644
--- a/holoviews/plotting/bokeh/widgets.py
+++ b/holoviews/plotting/bokeh/widgets.py
@@ -78,6 +78,9 @@ def __init__(self, plot, renderer=None, **params):
         self._queue = []
         self._active = False
 
+        if hasattr(self.plot.document, 'on_session_destroyed'):
+            self.plot.document.on_session_destroyed(self.plot._session_destroy)
+
 
     @classmethod
     def create_widget(self, dim, holomap=None, editable=False):

From 0a797498323b1df5a3f3e7af2af7329b56b3d18b Mon Sep 17 00:00:00 2001
From: Philipp Rudiger <P.Rudiger@ed.ac.uk>
Date: Fri, 26 Oct 2018 22:08:07 +0100
Subject: [PATCH 2/7] Fixed memory leak with weakrefs

---
 holoviews/plotting/bokeh/callbacks.py |  7 +++++++
 holoviews/plotting/bokeh/plot.py      | 25 +++++++++++++------------
 holoviews/streams.py                  | 26 +++++++++++++++++---------
 3 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/holoviews/plotting/bokeh/callbacks.py b/holoviews/plotting/bokeh/callbacks.py
index 639b8bfeff..7862befaa4 100644
--- a/holoviews/plotting/bokeh/callbacks.py
+++ b/holoviews/plotting/bokeh/callbacks.py
@@ -68,8 +68,15 @@ def __init__(self, plot, streams, source, **params):
 
 
     def cleanup(self):
+        self.reset()
+        self.handle_ids = None
+        self.plot = None
+        self.source = None
+        self.streams = []
         if self.comm:
             self.comm.close()
+        Callback._callbacks = {k: cb for k, cb in Callback._callbacks.items()
+                               if cb is not self}
 
 
     def reset(self):
diff --git a/holoviews/plotting/bokeh/plot.py b/holoviews/plotting/bokeh/plot.py
index 554e5ee4d0..a358aecdb1 100644
--- a/holoviews/plotting/bokeh/plot.py
+++ b/holoviews/plotting/bokeh/plot.py
@@ -103,7 +103,8 @@ def root(self):
 
     @document.setter
     def document(self, doc):
-        if hasattr(doc, 'on_session_destroyed') and self.root is self.handles.get('plot'):
+        if (doc and hasattr(doc, 'on_session_destroyed') and
+            self.root is self.handles.get('plot') and not isinstance(self, AdjointLayoutPlot)):
             doc.on_session_destroyed(self._session_destroy)
             if self._document:
                 self._document._session_destroyed_callbacks.pop(self._session_destroy, None)
@@ -116,8 +117,6 @@ def document(self, doc):
 
     def _session_destroy(self, session_context):
         self.cleanup()
-        self._document = None
-
 
     def __init__(self, *args, **params):
         root = params.pop('root', None)
@@ -158,7 +157,7 @@ def _construct_callbacks(self):
 
         cb_classes = set()
         for _, source in sources:
-            streams = Stream.registry.get(id(source), [])
+            streams = Stream.registry.get(source, [])
             registry = Stream._callbacks['bokeh']
             cb_classes |= {(registry[type(stream)], stream) for stream in streams
                            if type(stream) in registry and stream.linked}
@@ -281,19 +280,21 @@ def cleanup(self):
             if not isinstance(plot, (GenericCompositePlot, GenericElementPlot, GenericOverlayPlot)):
                 continue
             streams = list(plot.streams)
+            plot.streams = []
+            if isinstance(plot, GenericElementPlot):
+                for callback in plot.callbacks:
+                    streams += callback.streams
+                    callback.cleanup()
+
             for stream in set(streams):
                 stream._subscribers = [
                     (p, subscriber) for p, subscriber in stream._subscribers
                     if get_method_owner(subscriber) not in plots
                 ]
-            if not isinstance(plot, GenericElementPlot):
-                continue
-            for callback in plot.callbacks:
-                streams += callback.streams
-                callbacks = {k: cb for k, cb in callback._callbacks.items()
-                            if cb is not callback}
-                Callback._callbacks = callbacks
-                callback.cleanup()
+
+        self.subplots.clear()
+        if self.comm:
+            self.comm.close()
 
 
     def _fontsize(self, key, label='fontsize', common=True):
diff --git a/holoviews/streams.py b/holoviews/streams.py
index 28d201011e..5c8f69bb81 100644
--- a/holoviews/streams.py
+++ b/holoviews/streams.py
@@ -5,6 +5,7 @@
 """
 
 import uuid
+import weakref
 from numbers import Number
 from collections import defaultdict
 from contextlib import contextmanager
@@ -72,7 +73,7 @@ class Stream(param.Parameterized):
     """
 
     # Mapping from a source id to a list of streams
-    registry = defaultdict(list)
+    registry = weakref.WeakKeyDictionary()
 
     # Mapping to define callbacks by backend and Stream type.
     # e.g. Stream._callbacks['bokeh'][Stream] = Callback
@@ -199,7 +200,7 @@ def __init__(self, rename={}, source=None, subscribers=[], linked=False,
         Some streams are configured to automatically link to the source
         plot, to disable this set linked=False
         """
-        self._source = source
+        self._source = weakref.ref(source) if source else None
         self._subscribers = []
         for subscriber in subscribers:
             self.add_subscriber(subscriber)
@@ -218,7 +219,10 @@ def __init__(self, rename={}, source=None, subscribers=[], linked=False,
         super(Stream, self).__init__(**params)
         self._rename = self._validate_rename(rename)
         if source is not None:
-            self.registry[id(source)].append(self)
+            if sid in self.registry:
+                self.registry[source].append(self)
+            else:
+                self.registry[source] = [self]
 
 
     @property
@@ -296,22 +300,26 @@ def rename(self, **mapping):
         """
         params = {k: v for k, v in self.get_param_values() if k != 'name'}
         return self.__class__(rename=mapping,
-                              source=self._source,
+                              source=(self._source() if self._source else None),
                               linked=self.linked, **params)
 
     @property
     def source(self):
-        return self._source
+        return self._source() if self._source else None
 
 
     @source.setter
     def source(self, source):
-        if self._source:
-            source_list = self.registry[id(self._source)]
+        old_source = self._source() if self._source else None
+        if old_source:
+            source_list = self.registry[old_source]
             if self in source_list:
                 source_list.remove(self)
-        self._source = source
-        self.registry[id(source)].append(self)
+        self._source = weakref.ref(source)
+        if source in self.registry:
+            self.registry[source].append(self)
+        else:
+            self.registry[source] = [self]
 
 
     def transform(self):

From 49f5725c5e7d6a22685593f3be725733721536a5 Mon Sep 17 00:00:00 2001
From: Philipp Rudiger <P.Rudiger@ed.ac.uk>
Date: Fri, 26 Oct 2018 22:10:03 +0100
Subject: [PATCH 3/7] Added comments

---
 holoviews/streams.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/holoviews/streams.py b/holoviews/streams.py
index 5c8f69bb81..87d6be4d5e 100644
--- a/holoviews/streams.py
+++ b/holoviews/streams.py
@@ -72,7 +72,9 @@ class Stream(param.Parameterized):
     respectively.
     """
 
-    # Mapping from a source id to a list of streams
+    # Mapping from a source to a list of streams
+    # WeakKeyDictionary to allow garbage collection
+    # of unreferenced sources
     registry = weakref.WeakKeyDictionary()
 
     # Mapping to define callbacks by backend and Stream type.
@@ -200,7 +202,10 @@ def __init__(self, rename={}, source=None, subscribers=[], linked=False,
         Some streams are configured to automatically link to the source
         plot, to disable this set linked=False
         """
+
+        # Source is stored as a weakref to allow it to be garbage collected
         self._source = weakref.ref(source) if source else None
+
         self._subscribers = []
         for subscriber in subscribers:
             self.add_subscriber(subscriber)

From 4613030e16659b158fed827143b8fa4093a525c8 Mon Sep 17 00:00:00 2001
From: Philipp Rudiger <P.Rudiger@ed.ac.uk>
Date: Sat, 27 Oct 2018 01:35:33 +0100
Subject: [PATCH 4/7] Store Links in WeakKeyDictionary

---
 holoviews/plotting/bokeh/callbacks.py |  6 ++---
 holoviews/plotting/links.py           | 32 +++++++++++++++++----------
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/holoviews/plotting/bokeh/callbacks.py b/holoviews/plotting/bokeh/callbacks.py
index 7862befaa4..04fbccbdb5 100644
--- a/holoviews/plotting/bokeh/callbacks.py
+++ b/holoviews/plotting/bokeh/callbacks.py
@@ -1150,10 +1150,10 @@ def find_link(cls, plot, link=None):
         sources = plot.hmap.traverse(lambda x: x, [ViewableElement])
         for src in sources:
             if link is None:
-                if id(src) in Link.registry:
-                    return (plot, Link.registry[id(src)])
+                if src in Link.registry:
+                    return (plot, Link.registry[src])
             else:
-                if id(link.target) == id(src):
+                if link.target is src:
                     return (plot, [link])
 
     def validate(self):
diff --git a/holoviews/plotting/links.py b/holoviews/plotting/links.py
index 08f77af9bc..8e5882f000 100644
--- a/holoviews/plotting/links.py
+++ b/holoviews/plotting/links.py
@@ -1,9 +1,8 @@
+import weakref
 from collections import defaultdict
 
 import param
 
-from ..core import ViewableElement
-
 
 class Link(param.Parameterized):
     """
@@ -22,14 +21,8 @@ class Link(param.Parameterized):
     directional links between the source and target object.
     """
 
-    source = param.ClassSelector(class_=ViewableElement, doc="""
-        The source object of the link (required).""")
-
-    target = param.ClassSelector(class_=ViewableElement, doc="""
-        The target object of the link (optional).""")
-
     # Mapping from a source id to a Link instance
-    registry = defaultdict(list)
+    registry = weakref.WeakKeyDictionary()
 
     # Mapping to define callbacks by backend and Link type.
     # e.g. Link._callbacks['bokeh'][Stream] = Callback
@@ -38,7 +31,11 @@ class Link(param.Parameterized):
     def __init__(self, source, target=None, **params):
         if source is None:
             raise ValueError('%s must define a source' % type(self).__name__)
-        super(Link, self).__init__(source=source, target=target, **params)
+
+        # Source is stored as a weakref to allow it to be garbage collected
+        self._source = weakref.ref(source) if source else None
+        self._target = weakref.ref(target) if target else None
+        super(Link, self).__init__(**params)
         self.link()
 
     @classmethod
@@ -49,17 +46,28 @@ def register_callback(cls, backend, callback):
         """
         cls._callbacks[backend][cls] = callback
 
+    @property
+    def source(self):
+        return self._source() if self._source else None
+
+    @property
+    def target(self):
+        return self._target() if self._target else None
+
     def link(self):
         """
         Registers the Link
         """
-        self.registry[id(self.source)].append(self)
+        if self.source in self.registry:
+            self.registry[self.source].append(self)
+        else:
+            self.registry[self.source] = [self]
 
     def unlink(self):
         """
         Unregisters the Link
         """
-        links = self.registry.get(id(self.source))
+        links = self.registry.get(self.source)
         if self in links:
             links.pop(links.index(self))
 

From 1437b356918c386bb2746ca69e0b60236d27b854 Mon Sep 17 00:00:00 2001
From: Philipp Rudiger <P.Rudiger@ed.ac.uk>
Date: Sat, 27 Oct 2018 03:28:08 +0100
Subject: [PATCH 5/7] Fixed flakes

---
 holoviews/plotting/bokeh/plot.py | 3 ++-
 holoviews/streams.py             | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/holoviews/plotting/bokeh/plot.py b/holoviews/plotting/bokeh/plot.py
index a358aecdb1..094cccbbed 100644
--- a/holoviews/plotting/bokeh/plot.py
+++ b/holoviews/plotting/bokeh/plot.py
@@ -17,7 +17,7 @@
 from ..plot import (DimensionedPlot, GenericCompositePlot, GenericLayoutPlot,
                     GenericElementPlot, GenericOverlayPlot)
 from ..util import attach_streams, displayable, collate
-from .callbacks import Callback, LinkCallback
+from .callbacks import LinkCallback
 from .util import (layout_padding, pad_plots, filter_toolboxes, make_axis,
                    update_shared_sources, empty_plot, decode_bytes)
 
@@ -281,6 +281,7 @@ def cleanup(self):
                 continue
             streams = list(plot.streams)
             plot.streams = []
+            plot._document = None
             if isinstance(plot, GenericElementPlot):
                 for callback in plot.callbacks:
                     streams += callback.streams
diff --git a/holoviews/streams.py b/holoviews/streams.py
index 87d6be4d5e..1da3ec5e25 100644
--- a/holoviews/streams.py
+++ b/holoviews/streams.py
@@ -224,7 +224,7 @@ def __init__(self, rename={}, source=None, subscribers=[], linked=False,
         super(Stream, self).__init__(**params)
         self._rename = self._validate_rename(rename)
         if source is not None:
-            if sid in self.registry:
+            if source in self.registry:
                 self.registry[source].append(self)
             else:
                 self.registry[source] = [self]

From 69674876d8f1a17213fc4dde51b8ebb4cc4535a0 Mon Sep 17 00:00:00 2001
From: Philipp Rudiger <P.Rudiger@ed.ac.uk>
Date: Sat, 27 Oct 2018 12:37:48 +0100
Subject: [PATCH 6/7] Fixed minor issues

---
 holoviews/plotting/bokeh/plot.py | 5 ++++-
 holoviews/tests/teststreams.py   | 4 ++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/holoviews/plotting/bokeh/plot.py b/holoviews/plotting/bokeh/plot.py
index 094cccbbed..e0cc133d9d 100644
--- a/holoviews/plotting/bokeh/plot.py
+++ b/holoviews/plotting/bokeh/plot.py
@@ -282,6 +282,10 @@ def cleanup(self):
             streams = list(plot.streams)
             plot.streams = []
             plot._document = None
+
+            if plot.subplots:
+                plot.subplots.clear()
+
             if isinstance(plot, GenericElementPlot):
                 for callback in plot.callbacks:
                     streams += callback.streams
@@ -293,7 +297,6 @@ def cleanup(self):
                     if get_method_owner(subscriber) not in plots
                 ]
 
-        self.subplots.clear()
         if self.comm:
             self.comm.close()
 
diff --git a/holoviews/tests/teststreams.py b/holoviews/tests/teststreams.py
index 15a4548329..ca7012f8e8 100644
--- a/holoviews/tests/teststreams.py
+++ b/holoviews/tests/teststreams.py
@@ -484,12 +484,12 @@ def tearDown(self):
     def test_source_registry(self):
         points = Points([(0, 0)])
         PointerX(source=points)
-        self.assertIn(id(points), Stream.registry)
+        self.assertIn(points, Stream.registry)
 
     def test_source_registry_empty_element(self):
         points = Points([])
         PointerX(source=points)
-        self.assertIn(id(points), Stream.registry)
+        self.assertIn(points, Stream.registry)
 
 
 class TestParameterRenaming(ComparisonTestCase):

From 640284d8eeadf814b9992ad5276697ef46e0d355 Mon Sep 17 00:00:00 2001
From: Philipp Rudiger <P.Rudiger@ed.ac.uk>
Date: Sun, 28 Oct 2018 13:21:30 +0000
Subject: [PATCH 7/7] Fixed flake

---
 holoviews/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/holoviews/__init__.py b/holoviews/__init__.py
index f8a6090466..70109b2013 100644
--- a/holoviews/__init__.py
+++ b/holoviews/__init__.py
@@ -38,7 +38,7 @@
     import IPython                 # noqa (API import)
     from .ipython import notebook_extension
     extension = notebook_extension # noqa (name remapping)
-except ImportError as e:
+except ImportError:
     class notebook_extension(param.ParameterizedFunction):
         def __call__(self, *args, **opts): # noqa (dummy signature)
             raise Exception("IPython notebook not available: use hv.extension instead.")