Replace use of asyncio.locks.Event with asyncio.locks.Condition in bfx_websocket_bucket.py.

This commit is contained in:
Davide Casale
2023-06-17 22:20:31 +02:00
parent cc5f9f5b0e
commit b12fedb7a3
2 changed files with 49 additions and 47 deletions

View File

@@ -23,18 +23,20 @@ class BfxWebSocketBucket:
def __init__(self, host, event_emitter, events_per_subscription): def __init__(self, host, event_emitter, events_per_subscription):
self.host, self.event_emitter, self.events_per_subscription = host, event_emitter, events_per_subscription self.host, self.event_emitter, self.events_per_subscription = host, event_emitter, events_per_subscription
self.websocket, self.subscriptions, self.pendings = None, {}, [] self.websocket, self.subscriptions, self.pendings = None, {}, []
self.on_open_event = asyncio.locks.Event() self.condition = asyncio.locks.Condition()
self.handler = PublicChannelsHandler(event_emitter=self.event_emitter, \ self.handler = PublicChannelsHandler(event_emitter=self.event_emitter, \
events_per_subscription=self.events_per_subscription) events_per_subscription=self.events_per_subscription)
async def connect(self): async def connect(self):
async def _connection():
async with websockets.connect(self.host) as websocket: async with websockets.connect(self.host) as websocket:
self.websocket = websocket self.websocket = websocket
self.on_open_event.set()
await self.__recover_state() await self.__recover_state()
async with self.condition:
self.condition.notify()
async for message in websocket: async for message in websocket:
message = json.loads(message) message = json.loads(message)
@@ -60,12 +62,6 @@ class BfxWebSocketBucket:
if (chan_id := message[0]) and message[1] != _HEARTBEAT: if (chan_id := message[0]) and message[1] != _HEARTBEAT:
self.handler.handle(self.subscriptions[chan_id], message[1:]) self.handler.handle(self.subscriptions[chan_id], message[1:])
try:
await _connection()
except websockets.exceptions.ConnectionClosedError as error:
if error.code in (1006, 1012):
self.on_open_event.clear()
async def __recover_state(self): async def __recover_state(self):
for pending in self.pendings: for pending in self.pendings:
await self.websocket.send(json.dumps(pending)) await self.websocket.send(json.dumps(pending))
@@ -107,3 +103,9 @@ class BfxWebSocketBucket:
for subscription in self.subscriptions.values(): for subscription in self.subscriptions.values():
if subscription["subId"] == sub_id: if subscription["subId"] == sub_id:
return subscription["chanId"] return subscription["chanId"]
async def wait(self):
async with self.condition:
await self.condition.wait_for(
lambda: self.websocket is not None and \
self.websocket.open)

View File

@@ -138,7 +138,7 @@ class BfxWebSocketClient:
tasks = [ asyncio.create_task(coroutine) for coroutine in coroutines ] tasks = [ asyncio.create_task(coroutine) for coroutine in coroutines ]
if len(self.buckets) == 0 or \ if len(self.buckets) == 0 or \
(await asyncio.gather(*[bucket.on_open_event.wait() for bucket in self.buckets])): (await asyncio.gather(*[bucket.wait() for bucket in self.buckets])):
self.event_emitter.emit("open") self.event_emitter.emit("open")
if self.credentials: if self.credentials:
@@ -184,8 +184,10 @@ class BfxWebSocketClient:
try: try:
await _connection() await _connection()
except (websockets.exceptions.ConnectionClosedError, socket.gaierror) as error: except (websockets.exceptions.ConnectionClosedError, socket.gaierror) as error:
if isinstance(error, websockets.exceptions.ConnectionClosedError): for task in tasks:
if error.code in (1006, 1012): task.cancel()
if isinstance(error, websockets.exceptions.ConnectionClosedError) and error.code in (1006, 1012):
if error.code == 1006: if error.code == 1006:
self.logger.error("Connection lost: no close frame received " \ self.logger.error("Connection lost: no close frame received " \
"or sent (1006). Trying to reconnect...") "or sent (1006). Trying to reconnect...")
@@ -194,9 +196,6 @@ class BfxWebSocketClient:
self.logger.info("WSS server is about to restart, clients need " \ self.logger.info("WSS server is about to restart, clients need " \
"to reconnect (server sent 20051). Reconnection attempt in progress...") "to reconnect (server sent 20051). Reconnection attempt in progress...")
for task in tasks:
task.cancel()
reconnection = Reconnection(status=True, attempts=1, timestamp=datetime.now()) reconnection = Reconnection(status=True, attempts=1, timestamp=datetime.now())
if self.wss_timeout is not None: if self.wss_timeout is not None:
@@ -211,7 +210,8 @@ class BfxWebSocketClient:
f"the client has been offline for {datetime.now() - reconnection.timestamp})") f"the client has been offline for {datetime.now() - reconnection.timestamp})")
reconnection = reconnection._replace(attempts=reconnection.attempts + 1) reconnection = reconnection._replace(attempts=reconnection.attempts + 1)
else: raise error else:
raise error
if not reconnection.status: if not reconnection.status:
self.event_emitter.emit("disconnection", self.event_emitter.emit("disconnection",