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

Update WebSocketProxy.__call__ to use asyncio.TaskGroup #16

Open
DontPanicO opened this issue Mar 10, 2023 · 0 comments
Open

Update WebSocketProxy.__call__ to use asyncio.TaskGroup #16

DontPanicO opened this issue Mar 10, 2023 · 0 comments
Labels
enhancement New feature or request websocket-proxy Issues ore PRs related to WebSocketProxy

Comments

@DontPanicO
Copy link
Owner

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.

@DontPanicO DontPanicO added enhancement New feature or request websocket-proxy Issues ore PRs related to WebSocketProxy labels Mar 10, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request websocket-proxy Issues ore PRs related to WebSocketProxy
Projects
None yet
Development

No branches or pull requests

1 participant