fix client memory leak

Author: Archmonger

src/reactpy_django/components.py

@@ -7,7 +7,6 @@
 from uuid import uuid4
 from warnings import warn
 
-import orjson
 from django.contrib.staticfiles.finders import find
 from django.core.cache import caches
 from django.http import HttpRequest
@@ -19,7 +18,6 @@
 from reactpy_django.exceptions import ViewNotRegisteredError
 from reactpy_django.html import pyscript
 from reactpy_django.utils import (
-    extend_pyscript_config,
     generate_obj_name,
     import_module,
     render_pyscript_template,
@@ -324,13 +322,14 @@ def _python_to_pyscript(
         # FIXME: This is needed to properly re-render PyScript instances such as
         # when a component is re-rendered due to WebSocket disconnection.
         # There may be a better way to do this in the future.
-        # While this solution allows re-creating PyScript components, it also
-        # results in a browser memory leak. It currently unclear how to properly
-        # clean up unused code (user_workspace_UUID) from PyScript.
         set_rendered(True)
         return None
 
     return html.div(
-        html.div((extra_props or {}) | {"id": f"pyscript-{uuid}"}, initial),
+        html.div(
+            (extra_props or {})
+            | {"id": f"pyscript-{uuid}", "className": "pyscript", "data-uuid": uuid},
+            initial,
+        ),
         pyscript({"async": ""}, executor),
     )

src/reactpy_django/pyscript/component_template.py

@@ -1,18 +1,22 @@
-"""Code within this module is designed to be run directly by PyScript, and is not
-intended to be run in a Django environment.
-
-Our template tag performs string substitutions to turn this file into valid PyScript."""
-
+# pylint: disable=used-before-assignment
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    import asyncio
+
     from reactpy_django.pyscript.layout_manager import ReactPyLayoutManager
 
 
 # User component is inserted below by regex replacement
 def user_workspace_UUID():
     """Encapsulate the user's code with a completely unique function (workspace)
-    to prevent overlapping imports and variable names between different components."""
+    to prevent overlapping imports and variable names between different components.
+
+    This code is designed to be run directly by PyScript, and is not intended to be run
+    in a standard Python environment.
+
+    Our template tag performs string substitutions to turn this file into valid PyScript.
+    """
 
     def root(): ...
 
@@ -21,4 +25,4 @@ def root(): ...
 
 # PyScript allows top-level await, which allows us to not throw errors on components
 # that terminate early (such as hook-less components)
-await ReactPyLayoutManager("UUID").run(user_workspace_UUID)  # noqa: F704
+task_UUID = asyncio.create_task(ReactPyLayoutManager("UUID").run(user_workspace_UUID))

src/reactpy_django/pyscript/layout_manager.py

@@ -9,7 +9,11 @@
 
 class ReactPyLayoutManager:
     """Encapsulate the entire layout manager with a class to prevent overlapping
-    variable names between user code."""
+    variable names between user code.
+
+    This code is designed to be run directly by PyScript, and is not intended to be run
+    in a standard Python environment.
+    """
 
     def __init__(self, uuid):
         self.uuid = uuid
@@ -68,8 +72,38 @@ def event_handler(*args):
         event_name = event_name.lstrip("on_").lower().replace("_", "")
         add_event_listener(element, event_name, event_handler)
 
+    @staticmethod
+    def delete_old_workspaces():
+        dom_workspaces = js.document.querySelectorAll(".pyscript")
+        dom_uuids = {element.dataset.uuid for element in dom_workspaces}
+        python_uuids = {
+            value.split("_")[-1]
+            for value in globals()
+            if value.startswith("user_workspace_")
+        }
+
+        # Delete the workspace if it exists at the moment when we check
+        for uuid in python_uuids - dom_uuids:
+            task_name = f"task_{uuid}"
+            if task_name in globals():
+                task: asyncio.Task = globals()[task_name]
+                task.cancel()
+                del globals()[task_name]
+            else:
+                print(f"Warning: Could not auto delete PyScript task {task_name}")
+
+            workspace_name = f"user_workspace_{uuid}"
+            if workspace_name in globals():
+                del globals()[workspace_name]
+            else:
+                print(
+                    f"Warning: Could not auto delete PyScript workspace {workspace_name}"
+                )
+
     async def run(self, workspace_function: Coroutine):
+        self.delete_old_workspaces()
         root_model = {}
+
         async with Layout(workspace_function()) as layout:
             while True:
                 update = await layout.render()

src/reactpy_django/templates/reactpy/pyscript_component.html

@@ -1,4 +1,2 @@
-{% if reactpy_class %}<div id="pyscript-{{pyscript_uuid}}" class="{{pyscript_class}}">{{pyscript_initial_html}}</div>
-{% endif %}
-{% if not reactpy_class %}<div id="pyscript-{{pyscript_uuid}}">{{pyscript_initial_html}}</div>{% endif %}
+<div id="pyscript-{{pyscript_uuid}}" class="pyscript" data-uuid="{{pyscript_uuid}}">{{pyscript_initial_html}}</div>
 <py-script async>{{pyscript_executor}}</py-script>