Skip to content

Cancel WebSocketTestSession on close#2427

Merged
Kludex merged 7 commits intomasterfrom
cancel-websocket-on-disconnect
Jan 20, 2024
Merged

Cancel WebSocketTestSession on close#2427
Kludex merged 7 commits intomasterfrom
cancel-websocket-on-disconnect

Conversation

@Kludex
Copy link
Copy Markdown
Owner

@Kludex Kludex commented Jan 19, 2024

@Kludex Kludex requested review from agronholm and graingert January 19, 2024 20:38
@Kludex Kludex force-pushed the cancel-websocket-on-disconnect branch from 4dbfbc7 to 339836e Compare January 19, 2024 20:41
Comment on lines +131 to +145
async def run_app(tg: anyio.abc.TaskGroup) -> None:
try:
await self.app(self.scope, self._asgi_receive, self._asgi_send)
except anyio.get_cancelled_exc_class():
...
except BaseException as exc:
self._send_queue.put(exc)
raise
tg.cancel_scope.cancel()

async with anyio.create_task_group() as tg:
tg.start_soon(run_app, tg)
await self.event.wait()
tg.cancel_scope.cancel()
Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the relevant code.

try:
self.close(1000)
finally:
self.portal.start_task_soon(self._notify_close)
Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And this.

Comment on lines +257 to +270
def test_websocket_not_block_on_close(test_client_factory: Callable[..., TestClient]):
def app(scope: Scope):
async def asgi(receive: Receive, send: Send):
websocket = WebSocket(scope, receive=receive, send=send)
await websocket.accept()
while True:
await anyio.sleep(0.1)

return asgi

client = test_client_factory(app)
with client.websocket_connect("/") as websocket:
...
assert websocket.event.is_set()
Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the relevant test.

@@ -1,16 +1,20 @@
from __future__ import annotations
Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If reviewer complains about this, I can undo the changes.

@Kludex
Copy link
Copy Markdown
Owner Author

Kludex commented Jan 20, 2024

@encode/maintainers someone? 🙏

@Kludex Kludex mentioned this pull request Jan 20, 2024
2 tasks
Copy link
Copy Markdown
Contributor

@frankie567 frankie567 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! This approach with the event flag makes total sense.

except BaseException as exc:
self._send_queue.put(exc)
raise
tg.cancel_scope.cancel()
Copy link
Copy Markdown
Contributor

@frankie567 frankie567 Jan 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't we also cancel the task group before raising the exception just above?

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did that now.

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks.

@Kludex Kludex force-pushed the cancel-websocket-on-disconnect branch from e7ae532 to 8f78e7a Compare January 20, 2024 13:53
@Kludex Kludex merged commit 3ae161e into master Jan 20, 2024
@Kludex Kludex deleted the cancel-websocket-on-disconnect branch January 20, 2024 13:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Background Thread created by WebSocketTestSession should exit when close() is called

2 participants