FIX: prevent hanging on server shutdown by ianhi · Pull Request #546 · jupyterlab/jupyter-collaboration

ianhi · 2026-02-12T19:24:22Z

fixes #161 supersedes #514

🤖 full disclosure this a PR that I used claude for nearly in its entirety. I have reviewed every line to the best of my ability and confirmed that the fix works for me locally. I did multiple local rounds of review, tightening up exceptions, logic and improving comments. But ultimately I learned more about async through this PR than I knew at the start, so very possible that there is a subtlety i missed.

To debug I ended up designing a loop for claude where it used jupyter-collab-mcp to interact with the notebooks (creating rooms crucial for causing hte bug) and then sent sigterms from a shell with PYTHONFAULTHANDLER=1 and looped through possible fixes before getting it working. I then guiding claude through coming up with a minimal fix.

There were three issues preventing shutdown on receiving ctrl-C

1 client serve loop never finishing

the server was stopped. but client rooms were never stopped. We needed to sent a stop message to the clients to cause their loops to stop

jupyter-collaboration/projects/jupyter-server-ydoc/jupyter_server_ydoc/websocketserver.py

Lines 65 to 75 in c750ddd

    
           # results in corrupted data. 
        
           # room_tasks = list() 
        
           # for name, room in list(self.rooms.items()): 
        
           #     for task in room.background_tasks: 
        
           #         task.cancel()  # FIXME should be upstreamed 
        
           #         room_tasks.append(task) 
        
           # if room_tasks: 
        
           #     _, pending = await asyncio.wait(room_tasks, timeout=3) 
        
           #     if pending: 
        
           #         msg = f"{len(pending)} room task(s) are pending." 
        
           #         self.log.warning(msg)

2 Rooms never being stopped

rooms were persisted after a client disconnect (I think for allowing reconnect). But on shutdown they were never being stopped, so the task was staying alive. now we explcitly stop all rooms.

This option means they weren't being autoclosed on disconnect (which I think is good)

jupyter-collaboration/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py

Lines 119 to 120 in c750ddd

    
           rooms_ready=False, 
        
           auto_clean_rooms=False,

3 the root task was never awaited

because we didn't hold on to a reference to the server start task there was no way to await it and ultimately trigger the worker threads were never cleaned up. now we save a reference and await it at the end of clean

jupyter-collaboration/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py

Lines 82 to 84 in c750ddd

    
           await ensure_async(super().prepare()) 
        
           if not self._websocket_server.started.is_set(): 
        
               self.create_task(self._websocket_server.start())

github-actions · 2026-02-12T19:24:34Z

👈 Launch a Binder on branch ianhi/jupyter-collaboration/fix%2Fclean-shutdown-rooms

ianhi · 2026-02-12T19:42:33Z

+        for _, room in list(self.rooms.items()):
+            for client in list(room.clients):
+                tasks.extend(client._background_tasks)  # type: ignore[attr-defined]
+                client._message_queue.put_nowait(b"")  # type: ignore[attr-defined]


might be nice to have teh client have a disconnect method that just handles this. looks rather mysterious from here.

ianhi · 2026-02-12T19:55:13Z

hmm even though this works more than before it still doesn't work all of the time. :'( draft for now

On SIGTERM, non-daemon anyio WorkerThread instances blocked Python's threading._shutdown() because the root asyncio task was never awaited, so its done-callbacks (which send shutdown signals to worker threads) never fired. Fix: - Track the root asyncio task (server.start()) so clean() can await it - Stop all rooms explicitly during shutdown (auto_clean_rooms=False means pycrdt never stops them automatically) - Await the root task with a timeout so anyio's task group __aexit__ runs and worker threads get their shutdown signal - Cancel timed-out tasks in stop_extension() instead of ignoring them Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

ianhi · 2026-02-16T16:39:23Z

Found the final issue with the help of claude and PYTHONFAULTHANDLER=1

the main serve task was never being awaited so there were some workthreads hanging around preventing shutdown. updated description incoming. but i think this works now!

ianhi · 2026-02-16T16:56:24Z

not sure who to ask for a review here. but I can confirm that this fixes it for me locally. Not really sure why the windows 3.14t test failed :/ it doesn't seem obviously related to these changes?

davidbrochart

This is why structured concurrency is helpful. Jupyverse's RTC implementation is much simpler and doesn't have this bug.

davidbrochart · 2026-02-17T07:49:26Z

+            task = asyncio.create_task(self._websocket_server.start())
+            self._background_tasks.add(task)
+            task.add_done_callback(self._background_tasks.discard)
+            self._websocket_server._start_task = task


You could have self.create_task() return the task:

Suggested change

task = asyncio.create_task(self._websocket_server.start())

self._background_tasks.add(task)

task.add_done_callback(self._background_tasks.discard)

self._websocket_server._start_task = task

self._websocket_server._start_task = self.create_task(self._websocket_server.start())

davidbrochart · 2026-02-17T07:52:22Z

+
+        # Now disconnect existing clients so their serve() tasks complete.
+        tasks: list[asyncio.Task] = []
+        for _, room in list(self.rooms.items()):


Suggested change

for _, room in list(self.rooms.items()):

for room in list(self.rooms.values()):

davidbrochart · 2026-02-17T07:52:51Z

+        # Rooms persist after client disconnect so users can reconnect without
+        # a full re-sync (see auto_clean_rooms=False in app.py). On shutdown
+        # we must stop them explicitly.
+        for _, room in list(self.rooms.items()):


Suggested change

for _, room in list(self.rooms.items()):

for room in list(self.rooms.values()):

ianhi · 2026-02-17T16:20:01Z

Thank you for the review! I cleaned up based on your suggestions. Two follow ups:

I'm down to use systems that work better! But how would jupyverse have fixed this, per the docs it still would have used this library, and I thought the issues were internal? (This may be showing my async noviceness)
Im noticing now that I get clean shutdowns with multiple notebooks, but not when I have a jupyter terminal open. any ideas on that?

krassowski · 2026-02-17T17:15:21Z

jupyverse does not use jupyter_server_ydoc and instead reimplements it. It still depends on other packages from jupyter-collaboration (like the ones providing the UI).

krassowski · 2026-04-30T09:39:35Z

@ianhi do you think this is worth merging and releasing or is (2) a blocker?

ianhi mentioned this pull request Feb 12, 2026

fix: clean up background tasks on close #514

Open

ianhi commented Feb 12, 2026

View reviewed changes

ianhi marked this pull request as draft February 12, 2026 19:55

krassowski added the bug Something isn't working label Feb 13, 2026

ianhi force-pushed the fix/clean-shutdown-rooms branch from eb0b575 to d0ad3a4 Compare February 16, 2026 16:13

ianhi marked this pull request as ready for review February 16, 2026 16:55

davidbrochart reviewed Feb 17, 2026

View reviewed changes

ianhi added 2 commits February 17, 2026 11:07

cleaner loop

27572f9

simplify

a5f0151

Merge branch 'main' into fix/clean-shutdown-rooms

a5eca49

pr-triage-board Bot added this to PR triage (experimental) May 6, 2026

	# results in corrupted data.
	# room_tasks = list()
	# for name, room in list(self.rooms.items()):
	# for task in room.background_tasks:
	# task.cancel() # FIXME should be upstreamed
	# room_tasks.append(task)
	# if room_tasks:
	# _, pending = await asyncio.wait(room_tasks, timeout=3)
	# if pending:
	# msg = f"{len(pending)} room task(s) are pending."
	# self.log.warning(msg)

	await ensure_async(super().prepare())
	if not self._websocket_server.started.is_set():
	self.create_task(self._websocket_server.start())

	for _, room in list(self.rooms.items()):
	for room in list(self.rooms.values()):

Uh oh!

Uh oh!

Conversation

ianhi commented Feb 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

1 client serve loop never finishing

2 Rooms never being stopped

3 the root task was never awaited

Uh oh!

github-actions Bot commented Feb 12, 2026

Uh oh!

ianhi Feb 12, 2026

Choose a reason for hiding this comment

Uh oh!

ianhi commented Feb 12, 2026

Uh oh!

ianhi commented Feb 16, 2026

Uh oh!

ianhi commented Feb 16, 2026

Uh oh!

davidbrochart left a comment

Choose a reason for hiding this comment

Uh oh!

davidbrochart Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

davidbrochart Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

davidbrochart Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

ianhi commented Feb 17, 2026

Uh oh!

krassowski commented Feb 17, 2026

Uh oh!

krassowski commented Apr 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ianhi commented Feb 12, 2026 •

edited

Loading