From 0f161b5c9de4a992fe76fadbb4d6e0755c941065 Mon Sep 17 00:00:00 2001 From: Sergey Sayamov Date: Mon, 13 Mar 2023 20:03:24 +0300 Subject: [PATCH] Fixes #77: Pass task_id into background handler of WS endpoint for easy redefinition of on_receive --- docs/receipts/websockets.md | 7 ++++--- starlette_web/common/ws/base_endpoint.py | 6 +++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/receipts/websockets.md b/docs/receipts/websockets.md index f75fbc0..eb8d9d6 100644 --- a/docs/receipts/websockets.md +++ b/docs/receipts/websockets.md @@ -58,9 +58,10 @@ data manager. task group and cancels all tasks within that scope. Simply return inside background handler, if you want to cancel a single task. -**AVOID** stuff like this https://stackoverflow.com/a/60675826 -and, in general, passing event-handling primitives between tasks, -other than class-level anyio.Lock() +**NOT RECOMMENDED** to use stuff like this https://stackoverflow.com/a/60675826. +However, should you need to do so, you may redefine method +`starlette_web.common.ws.base_endpoint.BaseWSEndpoint.on_receive` to use +`anyio.TaskGroup.start` and return `CancelScope` from within background handler. ## Manually calling websocket.receive diff --git a/starlette_web/common/ws/base_endpoint.py b/starlette_web/common/ws/base_endpoint.py index 6e54b49..1ebd82d 100644 --- a/starlette_web/common/ws/base_endpoint.py +++ b/starlette_web/common/ws/base_endpoint.py @@ -67,7 +67,8 @@ async def on_connect(self, websocket: WebSocket) -> None: async def on_receive(self, websocket: WebSocket, data: Any) -> None: cleaned_data = self._validate(data) - self.task_group.start_soon(self._background_handler_wrap, websocket, cleaned_data) + task_id = get_random_string(50) + self.task_group.start_soon(self._background_handler_wrap, task_id, websocket, cleaned_data) async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None: if self.task_group: @@ -79,8 +80,7 @@ async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None: logger.debug("WS connection has been closed.") - async def _background_handler_wrap(self, websocket: WebSocket, data: Dict): - task_id = get_random_string(50) + async def _background_handler_wrap(self, task_id: str, websocket: WebSocket, data: Dict): task_result = None try: