fix: test and avoid memory leaks by checking references to virtual kernel

.github/workflows/test.yaml

@@ -242,6 +242,8 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
+      - uses: ts-graphviz/setup-graphviz@v1
+
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v5
         with:

packages/solara-meta/pyproject.toml

@@ -71,6 +71,7 @@ dev = [
     "types-requests",
     "types-markdown",
     "types-PyYAML",
+    "objgraph",
     "pytest",
     "pytest-mock",
     "pytest-cov",

solara/components/markdown.py

@@ -4,10 +4,9 @@
 import textwrap
 import traceback
 import warnings
-from typing import Any, Dict, List, Union, cast
+from typing import Any, Callable, Dict, List, Union, cast
 
 import ipyvuetify as v
-
 try:
     import pymdownx.emoji
     import pymdownx.highlight
@@ -16,6 +15,7 @@
     has_pymdownx = True
 except ModuleNotFoundError:
     has_pymdownx = False
+import reacton.core
 
 import solara
 import solara.components.applayout
@@ -50,7 +50,7 @@ def ExceptionGuard(children=[]):
             solara.Column(children=children)
 
 
-def _run_solara(code):
+def _run_solara(code, cleanups):
     ast = compile(code, "markdown", "exec")
     local_scope: Dict[Any, Any] = {}
     exec(ast, local_scope)
@@ -63,6 +63,13 @@ def _run_solara(code):
     else:
         raise NameError("No Page of app defined")
     box = v.Html(tag="div")
+
+    rc: reacton.core.RenderContext
+
+    def cleanup():
+        rc.close()
+
+    cleanups.append(cleanup)
     box, rc = solara.render(cast(solara.Element, app), container=box)  # type: ignore
     widget_id = box._model_id
     return (
@@ -224,7 +231,7 @@ def _markdown_template(
     return template
 
 
-def _highlight(src, language, unsafe_solara_execute, extra, *args, **kwargs):
+def _highlight(cleanups, src, language, unsafe_solara_execute, extra, *args, **kwargs):
     """Highlight a block of code"""
 
     if not has_pygments:
@@ -243,7 +250,7 @@ def _highlight(src, language, unsafe_solara_execute, extra, *args, **kwargs):
 
     if run_src_with_solara:
         if unsafe_solara_execute:
-            html_widget = _run_solara(src)
+            html_widget = _run_solara(src, cleanups)
             return src_html + html_widget
         else:
             return src_html + html_no_execute_enabled
@@ -260,8 +267,10 @@ def MarkdownIt(md_text: str, highlight: List[int] = [], unsafe_solara_execute: b
     from mdit_py_plugins.footnote import footnote_plugin  # noqa: F401
     from mdit_py_plugins.front_matter import front_matter_plugin  # noqa: F401
 
+    cleanups = solara.use_ref(cast(List[Callable[[], None]], []))
+
     def highlight_code(code, name, attrs):
-        return _highlight(code, name, unsafe_solara_execute, attrs)
+        return _highlight(cleanups.current, code, name, unsafe_solara_execute, attrs)
 
     md = MarkdownItMod(
         "js-default",
@@ -274,6 +283,15 @@ def highlight_code(code, name, attrs):
     md = md.use(container.container_plugin, name="note")
     html = md.render(md_text)
     hash = hashlib.sha256((html + str(unsafe_solara_execute) + repr(highlight)).encode("utf-8")).hexdigest()
+
+    def cleanup_wrapper():
+        def cleanup():
+            for cleanup in cleanups.current:
+                cleanup()
+
+        return cleanup
+
+    solara.use_effect(cleanup_wrapper)
     return v.VuetifyTemplate.element(template=_markdown_template(html)).key(hash)
 
 
@@ -332,11 +350,12 @@ def Page():
 
     md_text = textwrap.dedent(md_text)
     style = solara.util._flatten_style(style)
+    cleanups = solara.use_ref(cast(List[Callable[[], None]], []))
 
     def make_markdown_object():
         def highlight(src, language, *args, **kwargs):
             try:
-                return _highlight(src, language, unsafe_solara_execute, *args, **kwargs)
+                return _highlight(cleanups.current, src, language, unsafe_solara_execute, *args, **kwargs)
             except Exception as e:
                 logger.exception("Error highlighting code: %s", src)
                 return repr(e)
@@ -372,6 +391,16 @@ def highlight(src, language, *args, **kwargs):
 
     md = solara.use_memo(make_markdown_object, dependencies=[unsafe_solara_execute])
     html = md.convert(md_text)
+
+    def cleanup_wrapper():
+        def cleanup():
+            for cleanup in cleanups.current:
+                cleanup()
+
+        return cleanup
+
+    solara.use_effect(cleanup_wrapper)
+
     # if we update the template value, the whole vue tree will rerender (ipvue/ipyvuetify issue)
     # however, using the hash we simply generate a new widget each time
     hash = hashlib.sha256((html + str(unsafe_solara_execute)).encode("utf-8")).hexdigest()

solara/server/app.py

@@ -7,6 +7,7 @@
 import threading
 import traceback
 import warnings
+import weakref
 from enum import Enum
 from pathlib import Path
 from typing import Any, Dict, List, Optional, cast
@@ -132,7 +133,7 @@ def add_path():
         else:
             # the module itself will be added by reloader
             # automatically
-            with reload.reloader.watch():
+            with kernel_context.without_context(), reload.reloader.watch():
                 self.type = AppType.MODULE
                 try:
                     spec = importlib.util.find_spec(self.name)
@@ -420,6 +421,9 @@ def solara_comm_target(comm, msg_first):
 
     def on_msg(msg):
         nonlocal app
+        comm = comm_ref()
+        assert comm is not None
+        context = kernel_context.get_current_context()
         data = msg["content"]["data"]
         method = data["method"]
         if method == "run":
@@ -435,7 +439,12 @@ def on_msg(msg):
             themes = args.get("themes")
             dark = args.get("dark")
             load_themes(themes, dark)
-            load_app_widget(None, app, path)
+            try:
+                load_app_widget(None, app, path)
+            except Exception as e:
+                msg = f"Error loading app: from path {path} and app {app_name}"
+                logger.exception(msg)
+                raise RuntimeError(msg) from e
             comm.send({"method": "finished", "widget_id": context.container._model_id})
         elif method == "app-status":
             context = kernel_context.get_current_context()
@@ -464,18 +473,22 @@ def on_msg(msg):
         else:
             logger.error("Unknown comm method called on solara.control comm: %s", method)
 
-    comm.on_msg(on_msg)
-
     def reload():
+        comm = comm_ref()
+        assert comm is not None
+        context = kernel_context.get_current_context()
         # we don't reload the app ourself, we send a message to the client
         # this ensures that we don't run code of any client that for some reason is connected
         # but not working anymore. And it indirectly passes a message from the current thread
         # (which is that of the Reloader/watchdog), to the thread of the client
         logger.debug(f"Send reload to client: {context.id}")
         comm.send({"method": "reload"})
 
-    context = kernel_context.get_current_context()
-    context.reload = reload
+    comm.on_msg(on_msg)
+    comm_ref = weakref.ref(comm)
+    del comm
+
+    kernel_context.get_current_context().reload = reload
 
 
 def register_solara_comm_target(kernel: Kernel):

solara/server/kernel.py

@@ -247,6 +247,8 @@ def send(
         header=None,
         metadata=None,
     ):
+        if stream is None:
+            return  # can happen when the kernel is closed but someone was still trying to send a message
         try:
             if isinstance(msg_or_type, dict):
                 msg = msg_or_type
@@ -313,6 +315,39 @@ def __init__(self):
         self.shell.display_pub.session = self.session
         self.shell.display_pub.pub_socket = self.iopub_socket
 
+    def close(self):
+        if self.comm_manager is None:
+            raise RuntimeError("Kernel already closed")
+        self.session.close()
+        self._cleanup_references()
+
+    def _cleanup_references(self):
+        try:
+            # all of these reduce the circular references
+            # making it easier for the garbage collector to clean up
+            self.shell_handlers.clear()
+            self.control_handlers.clear()
+            for comm_object in list(self.comm_manager.comms.values()):  # type: ignore
+                comm_object.close()
+            self.comm_manager.targets.clear()  # type: ignore
+            # self.comm_manager.kernel points to us, but we cannot set it to None
+            # so we remove the circular reference by setting the comm_manager to None
+            self.comm_manager = None  # type: ignore
+            self.session.parent = None  # type: ignore
+
+            self.shell.display_pub.session = None  # type: ignore
+            self.shell.display_pub.pub_socket = None  # type: ignore
+            del self.shell.__dict__
+            self.shell = None  # type: ignore
+            self.session.websockets.clear()
+            self.session.stream = None  # type: ignore
+            self.session = None  # type: ignore
+            self.stream.session = None  # type: ignore
+            self.stream = None  # type: ignore
+            self.iopub_socket = None  # type: ignore
+        except Exception:
+            logger.exception("Error cleaning up references from kernel, not fatal")
+
     async def _flush_control_queue(self):
         pass