clean up possible open tasks on exit (#257)

SimonDanisch · Oct 2, 2024 · ab70056 · ab70056
1 parent 637afc0
commit ab70056
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/src/Bonito.jl b/src/Bonito.jl
@@ -94,6 +94,16 @@ function __init__()
     if !has_html_display()
         browser_display()
     end
+    atexit() do
+        for (task, (task, close_ref)) in Bonito.SERVER_CLEANUP_TASKS
+            close_ref[] = false
+        end
+        Bonito.CURRENT_SESSION[] = nothing
+        if !isnothing(Bonito.GLOBAL_SERVER[])
+            close(Bonito.GLOBAL_SERVER[])
+        end
+        Bonito.GLOBAL_SERVER[] = nothing
+    end
 end
 
 end # module
diff --git a/src/connection/websocket.jl b/src/connection/websocket.jl
@@ -52,7 +52,7 @@ function (connection::WebSocketConnection)(context, websocket::WebSocket)
 end
 
 
-const SERVER_CLEANUP_TASKS = Dict{Server, Task}()
+const SERVER_CLEANUP_TASKS = Dict{Server, Tuple{Task, Base.RefValue{Bool}}}()
 
 """
     abstract type CleanupPolicy end
@@ -166,7 +166,8 @@ end
 
 function add_cleanup_task!(server::Server)
     get!(SERVER_CLEANUP_TASKS, server) do
-        @async while true
+        close_ref = Base.RefValue(true)
+        task = @async while close_ref[]
             try
                 sleep(1)
                 cleanup_server(server)
@@ -176,6 +177,7 @@ function add_cleanup_task!(server::Server)
                 end
             end
         end
+        return (task, close_ref)
     end
 end