search

permitio / fastapi_websocket_rpc

⚡ FASTAPI Websocket RPC- A fast and durable bidirectional JSON RPC channel over Websockets.
https://permit.io
MIT License
216 stars 25 forks source link

Fix memory leak in `RpcChannel`'s `wait_for_response` #16

Closed jsdanielh closed 2 years ago

jsdanielh commented 2 years ago

Fix memory leak in the method wait_for_response of RpcChannel. As documented in #15, when calling asyncio.as_completed it actually creates a task to wait for the self._close event. In a long lived RpcChannel that event never happens and thus that future continues to run until the close method is called. Since the task is created with each RpcChannel's call method, this could end up leaking objects such as _asyncio.Task, _asyncio.Future and coroutine.

This change changes a the wait_for_response method to use asyncio.wait with a timeout and FIRST_COMPLETED for the returned_when argument such that we can obtain all pending futures after the first completed and cancel them to avoid the leak. This warrants that the wait for the self._close event future will be cancelled and then no future/task/coroutine is leaked.

Also add some styling changes and fix some typos.

This fixes #15.

jsdanielh commented 2 years ago

Per my comments here: #15 (comment)

It would be great to have a few more adjustments

I have applied your suggestion and also went and use asyncio.wait as you suggested in #15 (comment) which I also think it's a better solution. Let me know if you have more comments.

orweis commented 2 years ago

Looks great! @jsdanielh

Being a sensitive spot I would love another eye on this from one of you - @roekatz , @asafc

jsdanielh commented 2 years ago

@jsdanielh @orweis The fix in wait_for_response looks correct to me (canceling the self._closed.wait() task which was left uncanceled in the previous implementation).

But I don't really understand why it's not enough and we also need to discourage the use of None timeout. Do you mean that it leaks if the RPC doesn't return? (because to me that doesn't really count as a leak)

Yeah, you're right. It leaked when using as_completed but I don't think it still leaks now with the change to wait and returning when the first completes. Also I agree that if RPC doesn't return it isn't a leak. I'll amend the commit to remove the warnings I added on using None as a timeout