Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Cancel WebSocketTestSession on close #2427

Merged
merged 7 commits into from Jan 20, 2024
Merged

Conversation

Kludex
Copy link
Sponsor Member

@Kludex Kludex commented Jan 19, 2024

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
Sponsor Member 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.

def __exit__(self, *args: typing.Any) -> None:
try:
self.close(1000)
finally:
self.portal.start_task_soon(self._notify_close)
Copy link
Sponsor Member 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
Sponsor Member 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
Sponsor Member 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
Sponsor Member Author

Kludex commented Jan 20, 2024

@encode/maintainers someone? 🙏

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

@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
Member

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
Sponsor Member 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
Sponsor Member 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
5 checks passed
@Kludex Kludex deleted the cancel-websocket-on-disconnect branch January 20, 2024 13:59
nixroxursox pushed a commit to nixroxursox/starlette that referenced this pull request Mar 18, 2024
* Cancel `WebSocketTestSession` on close

* Undo some noise

* Fix test

* Undo pyproject

* Undo anyio bump

* Undo changes on test_authentication

* Always call cancel scope
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