From 98149f82b43d6a15260a6d16347dca9dd29b3ebf Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 2 May 2025 23:49:27 +0100 Subject: [PATCH] Mux is working with new services --- .../trustgraph/gateway/dispatch/config.py | 4 +- .../trustgraph/gateway/dispatch/flow.py | 4 +- .../trustgraph/gateway/dispatch/librarian.py | 4 +- .../trustgraph/gateway/dispatch/manager.py | 71 +++++++++---------- .../trustgraph/gateway/dispatch/mux.py | 35 +++++---- .../trustgraph/gateway/endpoint/manager.py | 34 ++++----- 6 files changed, 80 insertions(+), 72 deletions(-) diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/config.py b/trustgraph-flow/trustgraph/gateway/dispatch/config.py index 4b6a0439..3aeedb6f 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/config.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/config.py @@ -6,10 +6,12 @@ from ... schema import config_response_queue from . requestor import ServiceRequestor class ConfigRequestor(ServiceRequestor): - def __init__(self, pulsar_client, timeout=120): + def __init__(self, pulsar_client, consumer, subscriber, timeout=120): super(ConfigRequestor, self).__init__( pulsar_client=pulsar_client, + consumer_name = consumer, + subscription = subscriber, request_queue=config_request_queue, response_queue=config_response_queue, request_schema=ConfigRequest, diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/flow.py b/trustgraph-flow/trustgraph/gateway/dispatch/flow.py index 7a2f8a39..0b38e9be 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/flow.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/flow.py @@ -6,10 +6,12 @@ from ... schema import flow_response_queue from . requestor import ServiceRequestor class FlowRequestor(ServiceRequestor): - def __init__(self, pulsar_client, timeout=120): + def __init__(self, pulsar_client, consumer, subscriber, timeout=120): super(FlowRequestor, self).__init__( pulsar_client=pulsar_client, + consumer_name = consumer, + subscription = subscriber, request_queue=flow_request_queue, response_queue=flow_response_queue, request_schema=FlowRequest, diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py b/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py index 8705614a..f280b392 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/librarian.py @@ -8,10 +8,12 @@ from . serialize import serialize_document_package, serialize_document_info from . serialize import to_document_package, to_document_info, to_criteria class LibrarianRequestor(ServiceRequestor): - def __init__(self, pulsar_client, timeout=120): + def __init__(self, pulsar_client, consumer, subscriber, timeout=120): super(LibrarianRequestor, self).__init__( pulsar_client=pulsar_client, + consumer_name = consumer, + subscription = subscriber, request_queue=librarian_request_queue, response_queue=librarian_response_queue, request_schema=LibrarianRequest, diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index 4fecdf56..ff9dc8ad 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -40,6 +40,12 @@ request_response_dispatchers = { "triples": TriplesQueryRequestor, } +global_dispatchers = { + "config": ConfigRequestor, + "flow": FlowRequestor, + "librarian": LibrarianRequestor, +} + sender_dispatchers = { "text-load": TextLoad, "document-load": DocumentLoad, @@ -58,14 +64,10 @@ import_dispatchers = { } class DispatcherWrapper: - def __init__(self, mgr, name, impl): - self.mgr = mgr - self.name = name - self.impl = impl - async def process(self, data, responder): - return await self.mgr.process_impl( - data, responder, self.name, self.impl - ) + def __init__(self, handler): + self.handler = handler + async def process(self, *args): + return await self.handler(*args) class DispatcherManager: @@ -87,24 +89,26 @@ class DispatcherManager: del self.flows[id] return - def dispatch_config(self): - return DispatcherWrapper(self, "config", ConfigRequestor) + def dispatch_global_service(self): + return DispatcherWrapper(self.process_global_service) - def dispatch_flow(self): - return DispatcherWrapper(self, "flow", FlowRequestor) + async def process_global_service(self, data, responder, params): - def dispatch_librarian(self): - return DispatcherWrapper(self, "librarian", LibrarianRequestor) + kind = params.get("kind") + return await self.invoke_global_service(data, responder, kind) - async def process_impl(self, data, responder, name, impl): + async def invoke_global_service(self, data, responder, kind): - key = (None, name) + key = (None, kind) if key in self.dispatchers: return await self.dispatchers[key].process(data, responder) - dispatcher = impl( - pulsar_client = self.pulsar_client + dispatcher = global_dispatchers[kind]( + pulsar_client = self.pulsar_client, + timeout = 120, + consumer = f"api-gateway-{kind}-request", + subscriber = f"api-gateway-{kind}-request", ) await dispatcher.start() @@ -113,19 +117,19 @@ class DispatcherManager: return await dispatcher.process(data, responder) - def dispatch_service(self): - return self + def dispatch_flow_import(self): + return self.process_flow_import - def dispatch_import(self): - return self.invoke_import - - def dispatch_export(self): - return self.invoke_export + def dispatch_flow_export(self): + return self.process_flow_export def dispatch_socket(self): - return self.invoke_socket + return self.process_socket - async def invoke_import(self, ws, running, params): + def dispatch_flow_service(self): + return DispatcherWrapper(self.process_flow_service) + + async def process_flow_import(self, ws, running, params): flow = params.get("flow") kind = params.get("kind") @@ -156,7 +160,7 @@ class DispatcherManager: return dispatcher - async def invoke_export(self, ws, running, params): + async def process_flow_export(self, ws, running, params): flow = params.get("flow") kind = params.get("kind") @@ -189,25 +193,20 @@ class DispatcherManager: return dispatcher - async def invoke_socket(self, ws, running, params): - - flow = params.get("flow") - - if flow not in self.flows: - raise RuntimeError("Invalid flow") + async def process_socket(self, ws, running, params): dispatcher = Mux(self, ws, running) return dispatcher - async def process(self, data, responder, params): + async def process_flow_service(self, data, responder, params): flow = params.get("flow") kind = params.get("kind") return await self.invoke(data, responder, flow, kind) - async def invoke(self, data, responder, flow, kind): + async def invoke_flow_service(self, data, responder, flow, kind): if flow not in self.flows: raise RuntimeError("Invalid flow") diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/mux.py b/trustgraph-flow/trustgraph/gateway/dispatch/mux.py index 786d31ff..e2c5a921 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/mux.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/mux.py @@ -33,18 +33,17 @@ class Mux: data = msg.json() -# if data["service"] not in self.services: -# raise RuntimeError("Bad service") - if "request" not in data: raise RuntimeError("Bad message") if "id" not in data: raise RuntimeError("Bad message") - await self.q.put( - (data["id"], data["service"], data["request"]) - ) + await self.q.put(( + data["id"], data.get("flow"), + data["service"], + data["request"] + )) except Exception as e: print("receive exception:", str(e), flush=True) @@ -75,7 +74,7 @@ class Mux: # worker[0] still running, move on break - async def start_request_task(self, ws, id, svc, request, workers): + async def start_request_task(self, ws, id, flow, svc, request, workers): # Wait for outstanding requests to go below MAX_OUTSTANDING_REQUESTS while len(workers) > MAX_OUTSTANDING_REQUESTS: @@ -94,7 +93,7 @@ class Mux: }) worker = asyncio.create_task( - self.request_task(request, responder, "0000", svc) + self.request_task(request, responder, flow, svc) ) workers.append(worker) @@ -102,8 +101,19 @@ class Mux: async def request_task(self, request, responder, flow, svc): try: - dispatcher = self.dispatcher_manager.dispatch_service() - await dispatcher.invoke(request, responder, "0000", svc) + + if flow: + + await self.dispatcher_manager.invoke_flow_service( + request, responder, flow, svc + ) + + else: + + await self.dispatcher_manager.invoke_global_service( + request, responder, svc + ) + except Exception as e: await self.ws.send_json({"error": str(e)}) @@ -120,7 +130,8 @@ class Mux: await self.maybe_tidy_workers(workers) # Get next request on queue - id, svc, request = await asyncio.wait_for(self.q.get(), 1) + item = await asyncio.wait_for(self.q.get(), 1) + id, flow, svc, request = item except TimeoutError: continue @@ -141,7 +152,7 @@ class Mux: print(id, svc, request) await self.start_request_task( - self.ws, id, svc, request, workers + self.ws, id, flow, svc, request, workers ) except Exception as e: diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/manager.py b/trustgraph-flow/trustgraph/gateway/endpoint/manager.py index 6e59f667..75a39766 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/manager.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/manager.py @@ -23,42 +23,34 @@ class EndpointManager: } self.endpoints = [ - ConstantEndpoint( - endpoint_path = "/api/v1/librarian", auth = auth, - dispatcher = dispatcher_manager.dispatch_librarian(), - ), - ConstantEndpoint( - endpoint_path = "/api/v1/config", auth = auth, - dispatcher = dispatcher_manager.dispatch_config(), - ), - ConstantEndpoint( - endpoint_path = "/api/v1/flow", auth = auth, - dispatcher = dispatcher_manager.dispatch_flow(), - ), MetricsEndpoint( - endpoint_path = "/api/v1/metrics", + endpoint_path = "/api/metrics", prometheus_url = prometheus_url, auth = auth, ), + VariableEndpoint( + endpoint_path = "/api/v1/{kind}", auth = auth, + dispatcher = dispatcher_manager.dispatch_global_service(), + ), + SocketEndpoint( + endpoint_path = "/api/v1/socket", + auth = auth, + dispatcher = dispatcher_manager.dispatch_socket() + ), VariableEndpoint( endpoint_path = "/api/v1/flow/{flow}/service/{kind}", auth = auth, - dispatcher = dispatcher_manager.dispatch_service(), + dispatcher = dispatcher_manager.dispatch_flow_service(), ), SocketEndpoint( endpoint_path = "/api/v1/flow/{flow}/import/{kind}", auth = auth, - dispatcher = dispatcher_manager.dispatch_import() + dispatcher = dispatcher_manager.dispatch_flow_import() ), SocketEndpoint( endpoint_path = "/api/v1/flow/{flow}/export/{kind}", auth = auth, - dispatcher = dispatcher_manager.dispatch_export() - ), - SocketEndpoint( - endpoint_path = "/api/v1/flow/{flow}/socket", - auth = auth, - dispatcher = dispatcher_manager.dispatch_socket() + dispatcher = dispatcher_manager.dispatch_flow_export() ), ]