Skip to content

Commit

Permalink
Fix issue with not correctly waiting for connections (#8294)
Browse files Browse the repository at this point in the history
## Describe your changes

[In this PR](#8068), we
accidentally removed the waiting for `has_connection.` This means that
an app without any connection is likely to just do a lot of necessary
loops, causing higher CPU usage compared to previous versions. This PR
adds the `has_connection` await.

---

**Contribution License Agreement**

By submitting this pull request you agree that all contributions to this
project are made under the Apache 2.0 license.
  • Loading branch information
LukasMasuch committed Mar 13, 2024
1 parent c61e33b commit aa53d4f
Showing 1 changed file with 17 additions and 7 deletions.
24 changes: 17 additions & 7 deletions lib/streamlit/runtime/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -609,7 +609,22 @@ async def _loop_coroutine(self) -> None:
async_objs.started.set_result(None)

while not async_objs.must_stop.is_set():
if self._state == RuntimeState.ONE_OR_MORE_SESSIONS_CONNECTED:
if self._state == RuntimeState.NO_SESSIONS_CONNECTED: # type: ignore[comparison-overlap]
# mypy 1.4 incorrectly thinks this if-clause is unreachable,
# because it thinks self._state must be INITIAL | ONE_OR_MORE_SESSIONS_CONNECTED.

# Wait for new websocket connections (new sessions):
_, pending_tasks = await asyncio.wait( # type: ignore[unreachable]
(
asyncio.create_task(async_objs.must_stop.wait()),
asyncio.create_task(async_objs.has_connection.wait()),
),
return_when=asyncio.FIRST_COMPLETED,
)
# Clean up pending tasks to avoid memory leaks
for task in pending_tasks:
task.cancel()
elif self._state == RuntimeState.ONE_OR_MORE_SESSIONS_CONNECTED:
async_objs.need_send_data.clear()

for active_session_info in self._session_mgr.list_active_sessions():
Expand All @@ -628,16 +643,11 @@ async def _loop_coroutine(self) -> None:
# Yield for a few milliseconds between session message
# flushing.
await asyncio.sleep(0.01)
elif self._state == RuntimeState.NO_SESSIONS_CONNECTED: # type: ignore[comparison-overlap]
# mypy 1.4 incorrectly thinks this if-clause is unreachable,
# because it thinks self._state must be INITIAL | ONE_OR_MORE_SESSIONS_CONNECTED.

# This will jump to the asyncio.wait below.
pass
else:
# Break out of the thread loop if we encounter any other state.
break

# Wait for new proto messages that need to be sent out:
_, pending_tasks = await asyncio.wait(
(
asyncio.create_task(async_objs.must_stop.wait()),
Expand Down

0 comments on commit aa53d4f

Please sign in to comment.