Skip to content

Update WebSocketProxy.__call__ to use asyncio.TaskGroup #16

@DontPanicO

Description

@DontPanicO

Feature or Enhancement

Move from asyncio.gather to asyncio.TaskGroup in WebSocketProxy.__call__.

Pitch

asyncio.TaskGroup (only available wih Python >= 3.11) has a more structured cancellation logic and (as stated from python docs) it should be preferred over asyncio.gather when there are no reasons to use one of the two over the other.

  • Important: this will require for the next release to drop support for Python < 3.11

Actually, the implementation is:

async def _forward(
    client: WebSocket, target: websockets.WebSocketClientProtocol
) -> None:
    ...

async def _reverse(
    client: WebSocket, target: websockets.WebSocketClientProtocol
) -> None:
    ...


class WebSocketProxy:
    ...

    async def __call__(self) -> None:
        async with websockets.connect(self._server_endpoint) as target:
            self._forward_task = asyncio.create_task(
                _forward(self._client, target)
            )
            self._reverse_task = asyncio.create_task(
                _reverse(self._client, target)
            )
            await asyncio.gather(self._forward_task, self._reverse_task)

With asyncio.TaskGroup it'd be like:

...

    async def __call__(self) -> None:
        async with asyncio.TaskGroup() as tg:
            async with websockets.connect(self._server_endpoint) as target:
                self._forward_task = tg.create_task(
                    _forward(self._client, target)
                )
                self._reverse_task = tg.create_task(
                    _reverse(self._client, target)
                )

Entering websockets.connect in the taskgroup context ensures that if any failure with target occurs, our child tasks (and the partent too) would properly cancel.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or requestwebsocket-proxyIssues ore PRs related to WebSocketProxy

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions