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.

Sign up for GitHub

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

Ensure WebSocket is always marked as closed after getting ConnectionClosed error #132

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Ensure WebSocket is always marked as closed after getting ConnectionClosed error #132

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 11 additions & 16 deletions replit_river/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,10 +116,15 @@ async def is_websocket_open(self) -> bool:
async with self._ws_lock:
return await self._ws_wrapper.is_open()

async def _begin_close_session_countdown(self) -> None:
"""Begin the countdown to close session, this should be called when
websocket is closed.
async def _handle_connection_closed(self) -> None:
"""
Handle the WebSocket connection being closing.
This will trigger connection retries, if enabled, and starts a
session reconnection timer.
If the timer expires before reconnection, the session will be closed.
"""
# ensure websocket is closed and initiate connection retries if applicable.
await self.close_websocket(self._ws_wrapper, not self._is_server)
# calculate the value now before establishing it so that there are no
# await points between the check and the assignment to avoid a TOCTOU
# race.
Expand All @@ -145,12 +150,7 @@ async def serve(self) -> None:
try:
await self._handle_messages_from_ws(tg)
except ConnectionClosed:
if self._retry_connection_callback:
self._task_manager.create_task(
self._retry_connection_callback()
)

await self._begin_close_session_countdown()
await self._handle_connection_closed()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

seems like we lost the _retry_connection_callback callback here?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is already handled down inside close_websocket:

river-python/replit_river/session.py

Lines 454 to 455 in 235406c

if should_retry and self._retry_connection_callback:
self._task_manager.create_task(self._retry_connection_callback())

logger.debug("ConnectionClosed while serving", exc_info=True)
except FailedSendingMessageException:
# Expected error if the connection is closed.
Expand Down Expand Up @@ -310,10 +310,7 @@ async def _heartbeat(
"%r closing websocket because of heartbeat misses",
self.session_id,
)
await self.close_websocket(
self._ws_wrapper, should_retry=not self._is_server
)
await self._begin_close_session_countdown()
await self._handle_connection_closed()
continue
except FailedSendingMessageException:
# this is expected during websocket closed period
Expand Down Expand Up @@ -344,9 +341,7 @@ async def _send_transport_message(
websocket: websockets.WebSocketCommonProtocol,
) -> None:
try:
await send_transport_message(
msg, websocket, self._begin_close_session_countdown
)
await send_transport_message(msg, websocket, self._handle_connection_closed)
except WebsocketClosedException as e:
raise e
except FailedSendingMessageException as e:
Expand Down
Loading