From 1ca46f43b01dafe34f9ea6198cf04a3f91eb0bd2 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Wed, 25 Dec 2024 13:23:52 +0000
Subject: [PATCH] Mainly there, put responses not sent out properly

---
 trustgraph-flow/trustgraph/gateway/mux.py    | 87 ++++++++++++++++----
 trustgraph-flow/trustgraph/gateway/socket.py | 11 ++-
 2 files changed, 77 insertions(+), 21 deletions(-)

diff --git a/trustgraph-flow/trustgraph/gateway/mux.py b/trustgraph-flow/trustgraph/gateway/mux.py
index 74797069..e1d3161e 100644
--- a/trustgraph-flow/trustgraph/gateway/mux.py
+++ b/trustgraph-flow/trustgraph/gateway/mux.py
@@ -20,41 +20,89 @@ class MuxEndpoint(SocketEndpoint):
             endpoint_path=path, auth=auth,
         )
 
+        # The outstanding request queue, max size is 10
         self.q = asyncio.Queue(maxsize=10)
 
+        # Worker threads, servicing
+        self.workers = []
+
         self.services = services
 
     async def start(self):
         pass
 
+    async def maybe_tidy_workers(self):
+
+        while True:
+            try:
+
+                await asyncio.wait_for(
+                    asyncio.shield(self.workers[0]),
+                    0.05
+                )
+
+                # worker[0] now stopped
+                self.workers = self.workers[1:]
+
+                if len(self.workers) == 0:
+                    break
+
+                # Loop iterates to try the next worker
+
+            except TimeoutError:
+                # worker[0] still running, move on
+                break
+
+    async def start_request_task(self, ws, id, svc, request):
+
+        requestor = self.services[svc]
+
+        async def responder(resp, fin):
+            print("respnd", id)
+            print(svc, request)
+            await ws.send_json({
+                "id": id,
+                "response": resp,
+                "complete": fin,
+            })
+
+        # Wait for outstanding requests to go below 15
+        while len(self.workers) > 15:
+            await asyncio.sleep(0.1)
+
+        worker = asyncio.create_task(
+            requestor.process(request, responder)
+        )
+
+        self.workers.append(worker)
+
     async def async_thread(self, ws, running):
 
         while running.get():
 
             try:
+
+                if len(self.workers) > 0:
+                    await self.maybe_tidy_workers()
+
+                # Get next request on queue
                 id, svc, request = await asyncio.wait_for(self.q.get(), 1)
+
             except TimeoutError:
                 continue
+
             except Exception as e:
+                # This is an internal working error, may not be recoverable
+                print("Exception:", e)
                 await ws.send_json({"id": id, "error": str(e)})
+                break
 
             try:
-
-                print(svc, request)
-
-                requestor = self.services[svc]
-
-                async def responder(resp, fin):
-                    await ws.send_json({
-                        "id": id,
-                        "response": resp,
-                        "complete": fin,
-                    })
-
-                resp = await requestor.process(request, responder)
+                print(id, svc, request)
+                await self.start_request_task(ws, id, svc, request)
 
             except Exception as e:
-
+                print("Exception2:", e)
                 await ws.send_json({"error": str(e)})
 
         running.stop()
@@ -64,9 +112,7 @@ class MuxEndpoint(SocketEndpoint):
         async for msg in ws:
 
             # On error, finish
-            if msg.type == WSMsgType.ERROR:
-                break
-            else:
+            if msg.type == WSMsgType.TEXT:
 
                 try:
 
@@ -90,5 +136,12 @@ class MuxEndpoint(SocketEndpoint):
                     await ws.send_json({"error": str(e)})
                     continue
 
+            elif msg.type == WSMsgType.ERROR:
+                break
+            elif msg.type == WSMsgType.CLOSE:
+                break
+            else:
+                break
+
         running.stop()
 
diff --git a/trustgraph-flow/trustgraph/gateway/socket.py b/trustgraph-flow/trustgraph/gateway/socket.py
index 869792b7..71f07e69 100644
--- a/trustgraph-flow/trustgraph/gateway/socket.py
+++ b/trustgraph-flow/trustgraph/gateway/socket.py
@@ -22,11 +22,14 @@ class SocketEndpoint:
         
         async for msg in ws:
             # On error, finish
-            if msg.type == WSMsgType.ERROR:
-                break
+            if msg.type == WSMsgType.TEXT:
+                # Ignore incoming message
+                continue
+            elif msg.type == WSMsgType.BINARY:
+                # Ignore incoming message
+                continue
             else:
-                # Ignore incoming messages
-                pass
+                break
 
         running.stop()