fix: clean up thread correctly to avoid memory leak

widgetti · Dec 8, 2023 · a08ee06 · a08ee06
1 parent 76b4e26
commit a08ee06
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/solara/server/kernel_context.py b/solara/server/kernel_context.py
@@ -283,6 +283,11 @@ def set_context_for_thread(context: VirtualKernelContext, thread: threading.Thre
     current_context[key] = context
 
 
+def clear_context_for_thread(thread: threading.Thread):
+    key = get_thread_key(thread)
+    current_context.pop(key)
+
+
 def has_current_context() -> bool:
     thread_key = get_current_thread_key()
     return (thread_key in current_context) and (current_context[thread_key] is not None)

diff --git a/solara/server/patch.py b/solara/server/patch.py
@@ -260,15 +260,19 @@ def WidgetContextAwareThread__init__(self, *args, **kwargs):
 def Thread_debug_run(self):
     if self.current_context:
         kernel_context.set_context_for_thread(self.current_context, self)
-        shell = self.current_context.kernel.shell
-        shell.display_pub.register_hook(shell.display_in_reacton_hook)
+        display_pub = self.current_context.kernel.shell.display_pub
+        display_in_reacton_hook = self.current_context.kernel.shell.display_in_reacton_hook
+        display_pub.register_hook(display_in_reacton_hook)
     try:
         context = self.current_context or solara.util.nullcontext()
         with pdb_guard(), context:
             Thread__run(self)
     finally:
-        if self.current_context:
-            shell.display_pub.unregister_hook(shell.display_in_reacton_hook)
+        current_context = self.current_context
+        self.current_context = None
+        kernel_context.clear_context_for_thread(self)
+        if current_context:
+            display_pub.unregister_hook(display_in_reacton_hook)
 
 
 _patched = False