Skip to content

Commit 155c55b

Browse files
authored
feat!(gateway): Introduce proper reconnection without close. (#567)
It may introduce a runtime error however.
1 parent bb7d6fb commit 155c55b

File tree

2 files changed

+22
-4
lines changed

2 files changed

+22
-4
lines changed

interactions/api/gateway.py

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,7 @@ async def __restart(self):
165165
if self.__task:
166166
self.__task: Task
167167
self.__task.cancel()
168-
await self._client.close()
168+
self._client = None # clear pending waits
169169
self.__heartbeater.event.clear()
170170
await self._establish_connection()
171171

@@ -199,6 +199,10 @@ async def _establish_connection(
199199

200200
if stream is None:
201201
continue
202+
if self._client is None:
203+
await self._establish_connection()
204+
break
205+
202206
if self._client.close_code in range(4010, 4014) or self._client.close_code == 4004:
203207
raise GatewayException(self._client.close_code)
204208

@@ -230,6 +234,10 @@ async def _handle_connection(
230234
if op == OpCodeType.HELLO:
231235
self.__heartbeater.delay = data["heartbeat_interval"]
232236
self.__heartbeater.event.set()
237+
238+
if self.__task:
239+
self.__task.cancel() # so we can reduce redundant heartbeat bg tasks.
240+
233241
self.__task = ensure_future(self._manage_heartbeat())
234242

235243
if not self.session_id:
@@ -246,12 +254,17 @@ async def _handle_connection(
246254
if op in (OpCodeType.INVALIDATE_SESSION, OpCodeType.RECONNECT):
247255
log.debug("INVALID_SESSION/RECONNECT")
248256

249-
if data and op != OpCodeType.RECONNECT:
257+
# if data and op != OpCodeType.RECONNECT:
258+
# self.session_id = None
259+
# self.sequence = None
260+
# self._closed = True
261+
262+
if bool(data) is False and op == OpCodeType.INVALIDATE_SESSION:
250263
self.session_id = None
251-
self.sequence = None
252-
self._closed = True
253264

254265
await self.__restart()
266+
elif event == "RESUMED":
267+
log.debug(f"RESUMED (session_id: {self.session_id}, seq: {self.sequence})")
255268
elif event == "READY":
256269
self._ready = data
257270
self.session_id = data["session_id"]
@@ -515,6 +528,9 @@ def __option_type_context(self, context: object, type: int) -> dict:
515528
}
516529
return _resolved
517530

531+
async def restart(self):
532+
await self.__restart()
533+
518534
@property
519535
async def __receive_packet_stream(self) -> Optional[Dict[str, Any]]:
520536
"""
@@ -523,6 +539,7 @@ async def __receive_packet_stream(self) -> Optional[Dict[str, Any]]:
523539
:return: The packet stream.
524540
:rtype: Optional[Dict[str, Any]]
525541
"""
542+
526543
packet: WSMessage = await self._client.receive()
527544
return loads(packet.data) if packet and isinstance(packet.data, str) else None
528545

interactions/api/gateway.pyi

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,3 +78,4 @@ class WebSocketClient:
7878
def shard(self) -> None: ...
7979
@property
8080
def presence(self) -> None: ...
81+
async def restart(self): ...

0 commit comments

Comments
 (0)