Make message receive and handling async

[halleysfifthinc]

Motivation

All messages from the front-end/server are received and handled synchronously, including custom comm messages (comm_open, comm_msg, and comm_close). So, any currently executing cell blocks the IJulia kernel from receiving and handling any IOPub/comm messages. For example, in the following WebIO MWE, a JS function updates an "output" Observable, and the JS function is triggered by setting an ("input") observable:

using WebIO, Observables
s = Scope()
s["in"] = Observable{String}("")
s["out"] = Observable{String}("")
onjs(s["in"], js"""
function (val)
    _webIOScope.setObservableValue("out",val);
end""")

you can't observe a new s["out"] value (aka the result of the JS function) during execution of the same cell that set s["in"] (which triggers the JS function).

Example Julia function that fails (hangs) without async comms

function julia_js_julia(_in, out, str)
    ch = Channel{String}()
    obsf = on(out) do val
        put!(ch, val)
    end
    t = @async take!(ch)
    _in[] = str
    out = fetch(t)
    off(obsf)
        
    return out
end

*This example function isn't thread-safe. (The scp["in"] observable isn't locked, so concurrently setting it could lead to interleaved/mismatched updates to the scp["out"] observable.)

One example of an actual use-case/benefit is PlotlyJS.to_image, which uses the same
Julia => JS => Julia observable setup to retrieve the results of a plotly.js function call.
Currently, the PlotlyJS.to_image function soft-fails because the observable that holds the generated
image is only updated after the current cell finishes execution (when IJulia can process the
comm_msg from WebIO in the Jupyter frontend/client).

Testing

I've manually tested that the above WebIO MWE works with this PR, and that interrupting still works. I realize this is a fairly fundamental rearchitecturing of the message receiving/handling, but I'm not sure what else to test and/or if there is a good way to test any of this in CI. I'm open to any hints/pointers if you want more thorough testing/test cases.

Fixes #858.

P.S. Breadcrumb for the future: This new architecture has a lot of parallels (easily adapted) to the new subshells feature that was recently implemented in ipython/ipykernel#1249.

[JamesWrigley]

This sounds like a good idea, but it absolutely needs tests before merging. At some point I'll start writing tests for more of the internals which you should be able to modify for this PR, but feel free to have a go already if you have time :)

[halleysfifthinc]

👍 I will wait until you've added more internals tests before I do anything further. I am/have been running IJulia with this PR to give any bugs the opportunity to surface.

[JamesWrigley]

If you rebase this on master I think we can continue with it 🙂 Couple things:

• We should use Threads.@spawn instead of @async.
• We should run CI with multiple threads by default to try to catch any race conditions.

[halleysfifthinc]

Will do! The use of @async was actually intentional. My goal was to keep IJulia specific activity on the interactive thread. We could potentially go even further and @spawn cell execution on non-interactive threads to more intentionally separate user and IJulia activity. That could theoretically be helpful in some situations, but this PR will already be a(nother) significant rearchitecture of a core part of IJulia. (And I can't think of a specific motivating example.)

[JamesWrigley]

Keeping it on the interactive threads make sense, but for that we should use Threads.@spawn :interactive. @async has the unfortunate side-effect of pinning the parent task to the same thread so it's kinda discouraged now.

[halleysfifthinc]

pinning the parent task to the same thread

Right.. Is that not equivalent to Threads.@spawn :interactive? The rest of the IJulia kernel is synchronous/not using tasks, so it will always be on the first (aka interactive) thread, and we want the rest of the IJulia activity to stay on that thread too, just allowed to be asynchronous/concurrent?

Happy to learn more if I'm wrong, this was my first serious foray into async/concurrent programming!

[JamesWrigley]

That is technically true, but @async is still deprecated so I'd prefer we stick with Threads.@spawn and explicitly specifying the threadpool. One other advantage is that if there's multiple threads in the interactive threadpool then we can use all of them instead of one.

[codecov]

Codecov Report

❌ Patch coverage is 87.02290% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 66.59%. Comparing base (6062a30) to head (00fe58b).
⚠️ Report is 4 commits behind head on master.

Files with missing lines	Patch %	Lines
src/eventloop.jl	86.74%	11 Missing ⚠️
src/handlers.jl	66.66%	4 Missing ⚠️
src/msg.jl	92.85%	1 Missing ⚠️
src/stdio.jl	75.00%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (6062a30) and HEAD (00fe58b). Click for more details.

HEAD has 2 uploads less than BASE

Flag	BASE (6062a30)	HEAD (00fe58b)
	6	4

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1140      +/-   ##
==========================================
- Coverage   72.47%   66.59%   -5.89%     
==========================================
  Files          18       18              
  Lines        1326     1368      +42     
==========================================
- Hits          961      911      -50     
- Misses        365      457      +92     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[halleysfifthinc]

I've left some comments to explain some design decisions and/or about open questions I have.

I'm still unsure how to add tests for this, and I'd welcome any brainstorming.

[halleysfifthinc]

I'm not sure that this needs to be in a while loop vs something like

Suggested change

	try
	waitall([control_task, kernel.requests_task[], kernel.iopub_task[]])
	catch
	# send interrupts (user SIGINT) to the code-execution task
	if isa(e, InterruptException)
	@async Base.throwto(kernel.requests_task[], e)
	@async Base.throwto(kernel.iopub_task[], e)
	else
	rethrow()
	end
	finally
	wait(kernel.close_event)
	end

And maybe not even the finally clause? Basically, with the wait, this task shouldn't be scheduled again unless one of the message handling tasks fails, which we aren't trying to recover from. So if we do get back here, its because we want to/have to stop.

[JamesWrigley]

Yeah I agree, I was looking at this recently and thought the control flow was a bit strange 😅

[halleysfifthinc]

I am now wondering if the async handling should be expanded to most messages besides "execute_request"? In particular, "complete_request" and "inspect_request" are (should be?) side-effect free, and would be really convenient to be able to e.g. see the docs for a functions when writing a new cell while another cell is mid-execution.

[JamesWrigley]

Sorry I missed this 🙈 I'll try to review it this week but feel free to ping me if I forget.

[JamesWrigley]

I'm not quite convinced that what we're doing here is safe. If I understand correctly the reasoning is:

1. ZMQ sockets are not thread-safe.
2. Thus we use @async to ensure that all tasks are running on the same thread.
3. Thus we can safely recv/send in different tasks as long as we lock appropriately to prevent one recv being interleaved with another recv (likewise for send)

But that's making the assumption that ZMQ.jl's recv and send don't do anything to the socket internally that may conflict with each other, and I don't think that's true. Imagine this sequence:

• Task 1 is sending and yields immediately after calling zmq_msg_send(): https://github.com/JuliaInterop/ZMQ.jl/blob/1e1b458180311b19127937e8dd0befa79a93d54f/src/comm.jl#L8
  Let's say that zmq_msg_send() fails because we have to try again (EAGAIN).
• Task 2 is receiving and yields immediately after calling zmq_msg_recv(): https://github.com/JuliaInterop/ZMQ.jl/blob/1e1b458180311b19127937e8dd0befa79a93d54f/src/comm.jl#L80
  Let's say it fails for some non-EAGAIN reason (maybe a corrupted message or something). This overwrites the internal error code from zmq_msg_send().
• Control switches back to Task 1 which calls zmq_errno(), which returns the error code from the call to zmq_msg_recv() and thus incorrectly fails instead of trying again.

Now I'm pretty sure that neither send() or recv() will yield in those places so in practice this particular situation couldn't happen right now, but that's an implementation detail of ZMQ and certainly not something we can rely on. But I also can't think of a good alternative yet 🤔

Also, I fixed some lingering-task issues in #1190 which seems to have caused some merge conflicts, sorry about that 🙈

[JamesWrigley]

Yeah I agree, I was looking at this recently and thought the control flow was a bit strange 😅

[JamesWrigley]

Hmm a nice design would be to use a poller that could poll the iopub socket and an internal inproc socket that we send messages to. But ZMQ doesn't have a poller yet... JuliaInterop/ZMQ.jl#52

[JamesWrigley]

Using timeouts would also work, but myeh 🤷

[halleysfifthinc]

1. ZMQ sockets are not thread-safe.
2. Thus we use @async to ensure that all tasks are running on the same thread.
3. Thus we can safely recv/send in different tasks as long as we lock appropriately to prevent one recv being interleaved with another recv (likewise for send)

So the @async keeping things on the same thread is unrelated to the ZMQ sockets.

The actual motivating factor behind splitting the socket locks into read/write is because the read channel/task yields (waiting to read from the socket) while holding the lock. This caused a deadlock when another task tries to send, even though the socket is otherwise quiet (not actively receiving).

To avoid the split locks, we need a way to (in the receive channel/task) release the lock on a yielding wait (i.e. the socket doesn't have anything to read so the task yields). I couldn't figure out how to do that back when I first made this PR. I'll take another look to see if I can figure it out now.

[JamesWrigley]

Ok, now ZMQ.jl has a poller 😅 Not in a release yet but you can dev ZMQ for now and I'll tag a release before this is merged. I would propose this implementation for all sockets that need to be used by multiple tasks:

• One socket that connects to the Jupyter ZMQ address, and another inproc socket that IJulia tasks can send messages to to be forwarded to the Jupyter socket.
• The Jupyter socket is only touched by a single task that polls both the Jupyter socket and its corresponding inproc socket. Received messages from the Jupyter socket are put on a Channel and received messages from the inproc socket are forwarded to the Jupyter socket. Only the inproc socket needs to be locked (for both sends and recvs).

I believe that's fully threadsafe 🤞 What do you think?

[JamesWrigley]

I realized today that the Poller uses tasks internally, so they also need to be robust against InterruptException's. Will check that later. I think the cleanest way would be to bubble up the exception to wait(::Poller), so here we would just need to wrap the polling loop in a try-catch.

[JamesWrigley]

Gentle bump, did you have any luck with this?

[halleysfifthinc]

Taking another look now.. IIUC, the main/only(?) reason to use an inproc for the second, internal socket (instead of e.g. a Channel) is so the "conductor" task can have a single/unified thing to wait on (e.g. ZMQ.Poller), that is woken by either an incoming message on the Jupyter socket from the client, or an outgoing message to forward on the internal socket?

[JamesWrigley]

Yep exactly. That way we can guarantee that each socket is only ever touched by a single IJulia task at a time and also have a proper event-driven loop.

Also, I take back what I said about the inproc sockets needing to have a lock for sends and recvs. The recv socket will only ever be used from the conductor task so only the send socket needs a lock in case multiple tasks try to send stuff.

[halleysfifthinc]

I'm stumped for the moment, so I pushed the current state of things for a second opinion on any obvious design flaws.

Summary of design constraints as I understand them:

• Sending by the kernel must be synchronous (implicitly assumed by required ordering of e.g. status messages, etc), but receiving can be asynchronous.
• Lack of thread-safety in these ZMQ sockets prevents interleaved send/recv actions, therefore a single (thread-safe) lock must be used to limit send/recv actions to only one thread/task at a time.
• The original (synchronous) eventloop had a default/"resting" state of holding the socket locks while waiting to receive. This is untenable for an asynchronous eventloop because a receiver task would immediately re-acquire the socket lock after a successful read, preventing any other tasks from reliably/fairly acquiring the lock to send.

The solution to avoiding a receiver task's default/"resting" state of holding the lock is obviously polling, so that the lock is only held when we can immediately read without blocking, quickly followed by a release. I had originally looked into polling, but couldn't figure how to get ZMQ polling to interact with Julia's task scheduler properly.

Using the new ZMQ.Poller, the current design of receiving tasks is basically

poller = ZMQ.Poller([kernel.requests[]])
while true
    pr = wait(poller)
    msg::Msg = recv_ipython(pr.socket, kernel)
    # route via Channel to other async eventloop task ...
end

while other tasks handle the messages, where handling may require sending/receiving on at least 3 different sockets (an execute_request that reads from stdin would send and/or receive from the IOPub, stdin, and Shell sockets).

The recv socket will only ever be used from the conductor task so only the send socket needs a lock in case multiple tasks try to send stuff.

Agreed, and unless I'm missing something, I'm not actually sure why a second inproc socket is needed at all with my current design. By keeping a shared send/recv lock per socket, only one task can use the socket at any given time, and no task should be able to be blocked while holding the lock.

Unfortunately, the kernel hangs during the precompile script after line 53 (sending the "complete_request"). Some printlns at the beginning and end of the locked regions suggest that its not deadlocked by the socket locks, and a healthy sprinkling of printlns in the receiver tasks and precompile script resulted in getting just past line 55 (receiving the response). Any ideas?

[halleysfifthinc]

Addendum: Claude noticed a part of the Poller thread-safety warning that I had missed (which seems like the reason you had recommended inproc?):

It is also not threadsafe to use any of the sockets being monitored while the function is executing.

To fix that, I tried adding internal inproc sockets to forward to the actual output sockets, but that didn't fix the hang.

[JamesWrigley]

Reduced to:

using ZMQ
using IJulia: IJulia, Kernel, create_profile

profile = create_profile(8000; key="a0436f6c-1916-498b-8eb9-e81ab9368e84")
Kernel(profile; capture_stdout=false, capture_stderr=false, capture_stdin=false) do kernel
    # Connect as a client to the kernel
    requests_socket = ZMQ.Socket(ZMQ.DEALER)
    ip = profile["ip"]
    port = profile["shell_port"]
    ZMQ.connect(requests_socket, "tcp://$(ip):$(port)")

    # Kernel info
    idents = ["626c4427-479d61edd6b98ccca470f2d6"]
    signature = "306b616a72292e9a736fe42b3c7d6fd51e10653ea2c5bc8f33810a06d33df8b5"
    header = "{\"msg_id\": \"626c4427-479d61edd6b98ccca470f2d6_3346283_0\", \"msg_type\": \"kernel_info_request\", \"username\": \"james\", \"session\": \"626c4427-479d61edd6b98ccca470f2d6\", \"date\": \"2025-11-02T18:59:07.097698Z\", \"version\": \"5.4\"}"
    parent_header = "{}"
    metadata = "{}"
    content = "{}"

    for i in 1:2
        @info "Iteration $i"
        ZMQ.send_multipart(requests_socket, [only(idents), "<IDS|MSG>", signature, header, parent_header, metadata, content])
        ZMQ.recv_multipart(requests_socket, String)
    end

    close(requests_socket)
end

The first iteration succeeds but the second hangs because the requests task hangs at wait(poller) 🐙

[JamesWrigley]

Turned out to be a bug in the Poller, an internal flag was not being reset properly: JuliaInterop/ZMQ.jl#266

Now precompilation works for me with 6b435e4, but the kernel tests are hanging somewhere 🤔 Could you have a look at where? Lemme know if it looks like another poller bug and I'll investigate.

[JamesWrigley]

I rewrote the poller 😅 Now it's all in Julia instead of calling zmq_poll(), which allows for a drastically simpler implementation. Do you wanna try running IJulia again? For me all the tests pass now without hangs. I temporarily changed IJulia to point at JuliaInterop/ZMQ.jl#266 to see if CI is happy.

EDIT: given that CI appears to be hanging when precompiling I assume it's not happy.
EDIT: and somehow it's now hanging for me as well.

[JamesWrigley]

False alarm, I made a mistake in the sources section so CI was pointing to ZMQ.jl master 🙃 Fixed it now and tests pass, apart from 1.10 which doesn't support sources. Should be ready to try now. In the meantime I'll clean up the ZMQ PR for merging.