async def _receive(self) -> None:
try:
msg = await asyncio.wait_for(self._socket.recv(),
timeout=self._timeout)
except asyncio.TimeoutError:
await self._socket.ping()
return
# Other exceptions for socket.recv():
# - ConnectionClosed
# - ConnectionClosedOK
# - ConnectionClosedError
# which should be handled by self._connect()
else:
try:
parsed = json.loads(msg)
except ValueError as e:
logger.error(
format_msg(
'stream message "%s" is an invalid JSON: reason: %s',
msg, e))
return
else:
await self._handle_message(parsed)
async def _reconnect(self, fails: int, exception: Exception) -> None:
logger.error(
format_msg('socket error %s, reconnecting %s...',
repr_exception(exception), fails))
if self._connected_task is not None:
self._connected_task.cancel()
try:
await self._connected_task
except Exception:
pass
self._connected_task = None
self._before_connect()
async def close(self, code: int = DEFAULT_STREAM_CLOSE_CODE) -> None:
"""Close the current socket connection
Args:
code (:obj:`int`, optional): socket close code, defaults to 4999
"""
if not self._conn_task:
raise StreamDisconnectedException(self._uri)
# A lot of incomming messages might prevent
# the socket from gracefully shutting down,
# which leads `websockets` to fail connection
# and result in a 1006 close code (ConnectionClosedError).
# In that situation, we can not properly figure out whether the socket
# is closed by socket.close() or network connection error.
# So just set up a flag to do the trick
self._closing = True
tasks = [self._conn_task]
if self._socket:
tasks.append(
# make socket.close run in background
self._socket.close(code))
self._conn_task.cancel()
# Make sure:
# - conn_task is cancelled
# - socket is closed
for coro in asyncio.as_completed(tasks):
try:
await coro
except Exception as e:
logger.error(format_msg('close tasks error: %s', e))
self._socket = None
self._closing = False