websockets: improve closed websocket handling

When a closed websocket is detected set a flag to prevent further status updates until the websocket is removed.

Also make sure that the associated subscriptions are removed if a closed websocket is detected during a notification.

Signed-off-by:  Eric Callahan <arksine.code@gmail.com>
This commit is contained in:
Arksine 2021-02-15 11:26:54 -05:00
parent baf97f8ea4
commit f2135d7483
1 changed files with 4 additions and 1 deletions

View File

@ -276,6 +276,7 @@ class WebsocketManager:
ws.write_message(msg) ws.write_message(msg)
except WebSocketClosedError: except WebSocketClosedError:
self.websockets.pop(ws.uid, None) self.websockets.pop(ws.uid, None)
self.server.remove_subscription(ws)
logging.info(f"Websocket Removed: {ws.uid}") logging.info(f"Websocket Removed: {ws.uid}")
except Exception: except Exception:
logging.exception( logging.exception(
@ -294,6 +295,7 @@ class WebSocket(WebSocketHandler):
self.wsm = app.get_websocket_manager() self.wsm = app.get_websocket_manager()
self.rpc = self.wsm.rpc self.rpc = self.wsm.rpc
self.uid = id(self) self.uid = id(self)
self.is_closed = False
async def open(self): async def open(self):
await self.wsm.add_websocket(self) await self.wsm.add_websocket(self)
@ -311,7 +313,7 @@ class WebSocket(WebSocketHandler):
logging.exception("Websocket Command Error") logging.exception("Websocket Command Error")
def send_status(self, status): def send_status(self, status):
if not status: if not status or self.is_closed:
return return
try: try:
self.write_message({ self.write_message({
@ -319,6 +321,7 @@ class WebSocket(WebSocketHandler):
'method': "notify_status_update", 'method': "notify_status_update",
'params': [status]}) 'params': [status]})
except WebSocketClosedError: except WebSocketClosedError:
self.is_closed = True
logging.info( logging.info(
f"Websocket Closed During Status Update: {self.uid}") f"Websocket Closed During Status Update: {self.uid}")
except Exception: except Exception: