Refactor the shutdown logic to be threadsafe

Occasionally CI would hang during the `run_kernel()` test. This is most likely
because Julia will sometimes hang if there are tasks still running upon exit,
and `shutdown_handler()` previously did not bother closing all the eventloop
tasks before calling `exit()`. Now `shutdown_handler()` will begin the shutdown
process and allow `run_kernel()` to close the kernel properly.

Author: JamesWrigley

src/IJulia.jl

@@ -120,15 +120,15 @@ REPL.REPLDisplay(repl::MiniREPL) = repl.display
     # from the front-end.
     comms::Dict{String, Comm} = Dict{String, Comm}()
 
-    shutdown::Function = exit
+    shutdown::Function = start_shutdown
 
     # the following constants need to be initialized in init().
     publish::RefValue{Socket} = Ref{Socket}()
     raw_input::RefValue{Socket} = Ref{Socket}()
     requests::RefValue{Socket} = Ref{Socket}()
     control::RefValue{Socket} = Ref{Socket}()
     heartbeat::RefValue{Socket} = Ref{Socket}()
-    heartbeat_context::RefValue{Context} = Ref{Context}()
+    zmq_context::RefValue{Context} = Ref{Context}()
     profile::Dict{String, Any} = Dict{String, Any}()
     connection_file::Union{String, Nothing} = nothing
     read_stdout::RefValue{Base.PipeEndpoint} = Ref{Base.PipeEndpoint}()
@@ -180,6 +180,24 @@ function Base.wait(kernel::Kernel)
     end
 end
 
+function start_shutdown(kernel::Kernel)
+    IJulia._shutting_down[] = true
+    kernel.inited = false
+
+    # First we call zmq_ctx_shutdown() to close the context and stop all sockets
+    # from working. We don't call ZMQ.close(::Context) directly because that
+    # currently isn't threadsafe:
+    # https://github.com/JuliaInterop/ZMQ.jl/issues/256
+    ZMQ.lib.zmq_ctx_shutdown(kernel.zmq_context[])
+
+    # Wait for the heartbeat thread to stop
+    @ccall uv_thread_join(kernel.heartbeat_threadid::Ptr{Int})::Cint
+
+    # Now all the sockets should have been cancelled and the eventloop tasks
+    # should be ready to shutdown.
+    notify(kernel.stop_event)
+end
+
 function Base.close(kernel::Kernel)
     # Reset the IO streams first so that any later errors get printed
     if kernel.capture_stdout
@@ -205,17 +223,16 @@ function Base.close(kernel::Kernel)
     end
     popdisplay()
 
+    start_shutdown(kernel)
+    wait(kernel)
+
     # Close all sockets
     close(kernel.publish[])
     close(kernel.raw_input[])
     close(kernel.requests[])
     close(kernel.control[])
-    stop_heartbeat(kernel)
-
-    # The waitloop should now be ready to exit
-    kernel.inited = false
-    notify(kernel.stop_event)
-    wait(kernel)
+    close(kernel.heartbeat[])
+    close(kernel.zmq_context[])
 
     # Reset global variables
     IJulia.n = 0

src/handlers.jl

@@ -181,26 +181,26 @@ function kernel_info_request(socket, kernel, msg)
     send_ipython(socket, kernel,
                  msg_reply(msg, "kernel_info_reply",
                            Dict("protocol_version" => "5.4",
-                                        "implementation" => "ijulia",
-                                        "implementation_version" => string(pkgversion(@__MODULE__)),
-                                        "language_info" =>
-                                        Dict("name" => "julia",
-                                             "version" =>
+                                "implementation" => "ijulia",
+                                "implementation_version" => string(pkgversion(@__MODULE__)),
+                                "language_info" =>
+                                    Dict("name" => "julia",
+                                         "version" =>
                                              string(VERSION.major, '.',
                                                     VERSION.minor, '.',
                                                     VERSION.patch),
-                                             "mimetype" => "application/julia",
-                                             "file_extension" => ".jl"),
-                                        "banner" => "Julia: A fresh approach to technical computing.",
-                                        "help_links" => [
-                                                         Dict("text"=>"Julia Home Page",
-                                                              "url"=>"http://julialang.org/"),
-                                                         Dict("text"=>"Julia Documentation",
-                                                              "url"=>"http://docs.julialang.org/"),
-                                                         Dict("text"=>"Julia Packages",
-                                                              "url"=>"https://juliahub.com/ui/Packages")
-                                                        ],
-                                        "status" => "ok")))
+                                         "mimetype" => "application/julia",
+                                         "file_extension" => ".jl"),
+                                "banner" => "Julia: A fresh approach to technical computing.",
+                                "help_links" => [
+                                    Dict("text"=>"Julia Home Page",
+                                         "url"=>"http://julialang.org/"),
+                                    Dict("text"=>"Julia Documentation",
+                                         "url"=>"http://docs.julialang.org/"),
+                                    Dict("text"=>"Julia Packages",
+                                         "url"=>"https://juliahub.com/ui/Packages")
+                                ],
+                                "status" => "ok")))
 end
 
 """
@@ -226,23 +226,11 @@ request](https://jupyter-client.readthedocs.io/en/latest/messaging.html#kernel-s
 sending the reply this will exit the process.
 """
 function shutdown_request(socket, kernel, msg)
-    # stop heartbeat thread
-    stop_heartbeat(kernel)
-
-    # Shutdown the `requests` socket handler before sending any messages. This
-    # is necessary because otherwise the event loop will be calling
-    # `recv_ipython()` and holding a lock on `requests`, which will cause a
-    # deadlock when we try to send a message to it from the `control` socket
-    # handler.
-    IJulia._shutting_down[] = true
-    @async Base.throwto(kernel.requests_task[], InterruptException())
-
-    # In protocol 5.4 the shutdown reply moved to the control socket
-    shutdown_socket = VersionNumber(msg) >= v"5.4" ? kernel.control[] : kernel.requests[]
-    send_ipython(shutdown_socket, kernel,
+    send_ipython(kernel.control[], kernel,
                  msg_reply(msg, "shutdown_reply", msg.content))
     sleep(0.1) # short delay (like in ipykernel), to hopefully ensure shutdown_reply is sent
-    kernel.shutdown()
+
+    kernel.shutdown(kernel)
 
     nothing
 end

src/heartbeat.jl

@@ -25,23 +25,3 @@ function start_heartbeat(kernel)
     ccall(:uv_thread_create, Cint, (Ptr{Int}, Ptr{Cvoid}, Ptr{Cvoid}),
           kernel.heartbeat_threadid, heartbeat_c, heartbeat)
 end
-
-function stop_heartbeat(kernel)
-    if !isopen(kernel.heartbeat_context[])
-        # Do nothing if it has already been stopped (which can happen in the tests)
-        return
-    end
-
-    # First we call zmq_ctx_shutdown() to ensure that the zmq_proxy() call
-    # returns. We don't call ZMQ.close(::Context) directly because that
-    # currently isn't threadsafe:
-    # https://github.com/JuliaInterop/ZMQ.jl/issues/256
-    ZMQ.lib.zmq_ctx_shutdown(kernel.heartbeat_context[])
-    @ccall uv_thread_join(kernel.heartbeat_threadid::Ptr{Int})::Cint
-
-    # Now that the heartbeat thread has joined and its guaranteed to no longer
-    # be working on the heartbeat socket, we can safely close it and then the
-    # context.
-    close(kernel.heartbeat[])
-    close(kernel.heartbeat_context[])
-end

src/init.jl

@@ -108,12 +108,12 @@ function init(args, kernel, profile=nothing)
         kernel.hmac_key = collect(UInt8, profile["key"])
     end
 
-    kernel.publish[] = Socket(PUB)
-    kernel.raw_input[] = Socket(ROUTER)
-    kernel.requests[] = Socket(ROUTER)
-    kernel.control[] = Socket(ROUTER)
-    kernel.heartbeat_context[] = Context()
-    kernel.heartbeat[] = Socket(kernel.heartbeat_context[], ROUTER)
+    kernel.zmq_context[] = Context()
+    kernel.publish[] = Socket(kernel.zmq_context[], PUB)
+    kernel.raw_input[] = Socket(kernel.zmq_context[], ROUTER)
+    kernel.requests[] = Socket(kernel.zmq_context[], ROUTER)
+    kernel.control[] = Socket(kernel.zmq_context[], ROUTER)
+    kernel.heartbeat[] = Socket(kernel.zmq_context[], ROUTER)
     sep = profile["transport"]=="ipc" ? "-" : ":"
     bind(kernel.publish[], "$(profile["transport"])://$(profile["ip"])$(sep)$(profile["iopub_port"])")
     bind(kernel.requests[], "$(profile["transport"])://$(profile["ip"])$(sep)$(profile["shell_port"])")

src/kernel.jl

@@ -25,4 +25,5 @@ function run_kernel()
     end
 
     wait(IJulia._default_kernel::Kernel)
+    close(IJulia._default_kernel::Kernel)
 end

test/kernel.jl

@@ -217,7 +217,7 @@ end
         end
 
         shutdown_called = false
-        Kernel(profile; capture_stdout=false, capture_stderr=false, shutdown=() -> shutdown_called = true) do kernel
+        Kernel(profile; capture_stdout=false, capture_stderr=false, shutdown=(_) -> shutdown_called = true) do kernel
             jupyter_client(profile) do client
                 @testset "Comms" begin
                     # Try opening a Comm without a target_name, which should