Skip to content

Commit

Permalink
fix: websocket send task needed cancelling and only start when needed
Browse files Browse the repository at this point in the history
It kept a reference to the websocket which meant we had a memory leak.
  • Loading branch information
maartenbreddels committed Dec 1, 2023
1 parent 8e7a550 commit 96470e2
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion solara/server/starlette.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,8 @@ def __init__(self, ws: starlette.websockets.WebSocket, portal: anyio.from_thread
self.ws = ws
self.portal = portal
self.to_send: List[Union[str, bytes]] = []
self.task = asyncio.ensure_future(self.process_messages_task())
if settings.main.experimental_performance:
self.task = asyncio.ensure_future(self.process_messages_task())

async def process_messages_task(self):
while True:
Expand Down Expand Up @@ -220,6 +221,8 @@ async def run():
thread_return = anyio.to_thread.run_sync(websocket_thread_runner, ws, portal)
await thread_return
finally:
if settings.main.experimental_performance:
ws_wrapper.task.cancel()
try:
await ws.close()
except: # noqa
Expand Down

0 comments on commit 96470e2

Please # to comment.