Mux is working with new services

This commit is contained in:
Cyber MacGeddon 2025-05-02 23:49:27 +01:00
parent bc33122e6e
commit 98149f82b4
6 changed files with 80 additions and 72 deletions

View file

@ -6,10 +6,12 @@ from ... schema import config_response_queue
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class ConfigRequestor(ServiceRequestor): class ConfigRequestor(ServiceRequestor):
def __init__(self, pulsar_client, timeout=120): def __init__(self, pulsar_client, consumer, subscriber, timeout=120):
super(ConfigRequestor, self).__init__( super(ConfigRequestor, self).__init__(
pulsar_client=pulsar_client, pulsar_client=pulsar_client,
consumer_name = consumer,
subscription = subscriber,
request_queue=config_request_queue, request_queue=config_request_queue,
response_queue=config_response_queue, response_queue=config_response_queue,
request_schema=ConfigRequest, request_schema=ConfigRequest,

View file

@ -6,10 +6,12 @@ from ... schema import flow_response_queue
from . requestor import ServiceRequestor from . requestor import ServiceRequestor
class FlowRequestor(ServiceRequestor): class FlowRequestor(ServiceRequestor):
def __init__(self, pulsar_client, timeout=120): def __init__(self, pulsar_client, consumer, subscriber, timeout=120):
super(FlowRequestor, self).__init__( super(FlowRequestor, self).__init__(
pulsar_client=pulsar_client, pulsar_client=pulsar_client,
consumer_name = consumer,
subscription = subscriber,
request_queue=flow_request_queue, request_queue=flow_request_queue,
response_queue=flow_response_queue, response_queue=flow_response_queue,
request_schema=FlowRequest, request_schema=FlowRequest,

View file

@ -8,10 +8,12 @@ from . serialize import serialize_document_package, serialize_document_info
from . serialize import to_document_package, to_document_info, to_criteria from . serialize import to_document_package, to_document_info, to_criteria
class LibrarianRequestor(ServiceRequestor): class LibrarianRequestor(ServiceRequestor):
def __init__(self, pulsar_client, timeout=120): def __init__(self, pulsar_client, consumer, subscriber, timeout=120):
super(LibrarianRequestor, self).__init__( super(LibrarianRequestor, self).__init__(
pulsar_client=pulsar_client, pulsar_client=pulsar_client,
consumer_name = consumer,
subscription = subscriber,
request_queue=librarian_request_queue, request_queue=librarian_request_queue,
response_queue=librarian_response_queue, response_queue=librarian_response_queue,
request_schema=LibrarianRequest, request_schema=LibrarianRequest,

View file

@ -40,6 +40,12 @@ request_response_dispatchers = {
"triples": TriplesQueryRequestor, "triples": TriplesQueryRequestor,
} }
global_dispatchers = {
"config": ConfigRequestor,
"flow": FlowRequestor,
"librarian": LibrarianRequestor,
}
sender_dispatchers = { sender_dispatchers = {
"text-load": TextLoad, "text-load": TextLoad,
"document-load": DocumentLoad, "document-load": DocumentLoad,
@ -58,14 +64,10 @@ import_dispatchers = {
} }
class DispatcherWrapper: class DispatcherWrapper:
def __init__(self, mgr, name, impl): def __init__(self, handler):
self.mgr = mgr self.handler = handler
self.name = name async def process(self, *args):
self.impl = impl return await self.handler(*args)
async def process(self, data, responder):
return await self.mgr.process_impl(
data, responder, self.name, self.impl
)
class DispatcherManager: class DispatcherManager:
@ -87,24 +89,26 @@ class DispatcherManager:
del self.flows[id] del self.flows[id]
return return
def dispatch_config(self): def dispatch_global_service(self):
return DispatcherWrapper(self, "config", ConfigRequestor) return DispatcherWrapper(self.process_global_service)
def dispatch_flow(self): async def process_global_service(self, data, responder, params):
return DispatcherWrapper(self, "flow", FlowRequestor)
def dispatch_librarian(self): kind = params.get("kind")
return DispatcherWrapper(self, "librarian", LibrarianRequestor) 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: if key in self.dispatchers:
return await self.dispatchers[key].process(data, responder) return await self.dispatchers[key].process(data, responder)
dispatcher = impl( dispatcher = global_dispatchers[kind](
pulsar_client = self.pulsar_client pulsar_client = self.pulsar_client,
timeout = 120,
consumer = f"api-gateway-{kind}-request",
subscriber = f"api-gateway-{kind}-request",
) )
await dispatcher.start() await dispatcher.start()
@ -113,19 +117,19 @@ class DispatcherManager:
return await dispatcher.process(data, responder) return await dispatcher.process(data, responder)
def dispatch_service(self): def dispatch_flow_import(self):
return self return self.process_flow_import
def dispatch_import(self): def dispatch_flow_export(self):
return self.invoke_import return self.process_flow_export
def dispatch_export(self):
return self.invoke_export
def dispatch_socket(self): 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") flow = params.get("flow")
kind = params.get("kind") kind = params.get("kind")
@ -156,7 +160,7 @@ class DispatcherManager:
return dispatcher return dispatcher
async def invoke_export(self, ws, running, params): async def process_flow_export(self, ws, running, params):
flow = params.get("flow") flow = params.get("flow")
kind = params.get("kind") kind = params.get("kind")
@ -189,25 +193,20 @@ class DispatcherManager:
return dispatcher return dispatcher
async def invoke_socket(self, ws, running, params): async def process_socket(self, ws, running, params):
flow = params.get("flow")
if flow not in self.flows:
raise RuntimeError("Invalid flow")
dispatcher = Mux(self, ws, running) dispatcher = Mux(self, ws, running)
return dispatcher return dispatcher
async def process(self, data, responder, params): async def process_flow_service(self, data, responder, params):
flow = params.get("flow") flow = params.get("flow")
kind = params.get("kind") kind = params.get("kind")
return await self.invoke(data, responder, flow, 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: if flow not in self.flows:
raise RuntimeError("Invalid flow") raise RuntimeError("Invalid flow")

View file

@ -33,18 +33,17 @@ class Mux:
data = msg.json() data = msg.json()
# if data["service"] not in self.services:
# raise RuntimeError("Bad service")
if "request" not in data: if "request" not in data:
raise RuntimeError("Bad message") raise RuntimeError("Bad message")
if "id" not in data: if "id" not in data:
raise RuntimeError("Bad message") raise RuntimeError("Bad message")
await self.q.put( await self.q.put((
(data["id"], data["service"], data["request"]) data["id"], data.get("flow"),
) data["service"],
data["request"]
))
except Exception as e: except Exception as e:
print("receive exception:", str(e), flush=True) print("receive exception:", str(e), flush=True)
@ -75,7 +74,7 @@ class Mux:
# worker[0] still running, move on # worker[0] still running, move on
break 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 # Wait for outstanding requests to go below MAX_OUTSTANDING_REQUESTS
while len(workers) > MAX_OUTSTANDING_REQUESTS: while len(workers) > MAX_OUTSTANDING_REQUESTS:
@ -94,7 +93,7 @@ class Mux:
}) })
worker = asyncio.create_task( worker = asyncio.create_task(
self.request_task(request, responder, "0000", svc) self.request_task(request, responder, flow, svc)
) )
workers.append(worker) workers.append(worker)
@ -102,8 +101,19 @@ class Mux:
async def request_task(self, request, responder, flow, svc): async def request_task(self, request, responder, flow, svc):
try: 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: except Exception as e:
await self.ws.send_json({"error": str(e)}) await self.ws.send_json({"error": str(e)})
@ -120,7 +130,8 @@ class Mux:
await self.maybe_tidy_workers(workers) await self.maybe_tidy_workers(workers)
# Get next request on queue # 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: except TimeoutError:
continue continue
@ -141,7 +152,7 @@ class Mux:
print(id, svc, request) print(id, svc, request)
await self.start_request_task( await self.start_request_task(
self.ws, id, svc, request, workers self.ws, id, flow, svc, request, workers
) )
except Exception as e: except Exception as e:

View file

@ -23,42 +23,34 @@ class EndpointManager:
} }
self.endpoints = [ 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( MetricsEndpoint(
endpoint_path = "/api/v1/metrics", endpoint_path = "/api/metrics",
prometheus_url = prometheus_url, prometheus_url = prometheus_url,
auth = auth, 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( VariableEndpoint(
endpoint_path = "/api/v1/flow/{flow}/service/{kind}", endpoint_path = "/api/v1/flow/{flow}/service/{kind}",
auth = auth, auth = auth,
dispatcher = dispatcher_manager.dispatch_service(), dispatcher = dispatcher_manager.dispatch_flow_service(),
), ),
SocketEndpoint( SocketEndpoint(
endpoint_path = "/api/v1/flow/{flow}/import/{kind}", endpoint_path = "/api/v1/flow/{flow}/import/{kind}",
auth = auth, auth = auth,
dispatcher = dispatcher_manager.dispatch_import() dispatcher = dispatcher_manager.dispatch_flow_import()
), ),
SocketEndpoint( SocketEndpoint(
endpoint_path = "/api/v1/flow/{flow}/export/{kind}", endpoint_path = "/api/v1/flow/{flow}/export/{kind}",
auth = auth, auth = auth,
dispatcher = dispatcher_manager.dispatch_export() dispatcher = dispatcher_manager.dispatch_flow_export()
),
SocketEndpoint(
endpoint_path = "/api/v1/flow/{flow}/socket",
auth = auth,
dispatcher = dispatcher_manager.dispatch_socket()
), ),
] ]