apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-05-25 23:35:12 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update socket to complete the change

Browse source
This commit is contained in:
Cyber MacGeddon 2025-08-28 10:27:47 +01:00
parent 976529de69
commit 5f31b9efd8
1 changed files with 37 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

49
trustgraph-flow/trustgraph/gateway/endpoint/socket.py
View file

@ -25,7 +25,7 @@ class SocketEndpoint:
await dispatcher.run() await dispatcher.run()
async def listener(self, ws, dispatcher, running): async def listener(self, ws, dispatcher, running):
"""Enhanced listener with graceful shutdown"""
async for msg in ws: async for msg in ws:
# On error, finish # On error, finish
@ -36,13 +36,16 @@ class SocketEndpoint:
await dispatcher.receive(msg) await dispatcher.receive(msg)
continue continue
else: else:
# Graceful shutdown on close
logger.info("Websocket closing, initiating graceful shutdown")
running.stop()
# Allow time for dispatcher cleanup
await asyncio.sleep(1.0)
break break
running.stop()
await ws.close()
async def handle(self, request): async def handle(self, request):
"""Enhanced handler with better cleanup"""
try: try:
token = request.query['token'] token = request.query['token']
except: except:
@ -55,7 +58,9 @@ class SocketEndpoint:
ws = web.WebSocketResponse(max_msg_size=52428800) ws = web.WebSocketResponse(max_msg_size=52428800)
await ws.prepare(request) await ws.prepare(request)
dispatcher = None
try: try:
async with asyncio.TaskGroup() as tg: async with asyncio.TaskGroup() as tg:
@ -80,9 +85,6 @@ class SocketEndpoint:
logger.debug("Task group closed") logger.debug("Task group closed")
# Finally?
await dispatcher.destroy()
except ExceptionGroup as e: except ExceptionGroup as e:
logger.error("Exception group occurred:", exc_info=True) logger.error("Exception group occurred:", exc_info=True)
@ -90,11 +92,34 @@ class SocketEndpoint:
for se in e.exceptions: for se in e.exceptions:
logger.error(f" Exception type: {type(se)}") logger.error(f" Exception type: {type(se)}")
logger.error(f" Exception: {se}") logger.error(f" Exception: {se}")
# Attempt graceful dispatcher shutdown
if dispatcher and hasattr(dispatcher, 'destroy'):
try:
await asyncio.wait_for(
dispatcher.destroy(),
timeout=5.0
)
except asyncio.TimeoutError:
logger.warning("Dispatcher shutdown timed out")
except Exception as de:
logger.error(f"Error during dispatcher cleanup: {de}")
except Exception as e: except Exception as e:
logger.error(f"Socket exception: {e}", exc_info=True) logger.error(f"Socket exception: {e}", exc_info=True)
await ws.close() finally:
# Ensure dispatcher cleanup
if dispatcher and hasattr(dispatcher, 'destroy'):
try:
await dispatcher.destroy()
except Exception as de:
logger.error(f"Error in final dispatcher cleanup: {de}")
# Ensure websocket is closed
if ws and not ws.closed:
await ws.close()
return ws return ws
async def start(self): async def start(self):