mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-17 17:21:02 +02:00
Comments, maybe broke something
This commit is contained in:
Cyber MacGeddon
parent
38cea4c26d
commit
4c22f23f3a
7 changed files with 145 additions and 109 deletions
|
|
@ -1,13 +1,15 @@
|
|||
|
||||
# Base class for processors. Implements:
|
||||
# - Pulsar client, subscribe and consume basic
|
||||
# - the async startup logic
|
||||
# - Initialising metrics
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import argparse
|
||||
import pulsar
|
||||
from pulsar.schema import JsonSchema
|
||||
import _pulsar
|
||||
import time
|
||||
import uuid
|
||||
from prometheus_client import start_http_server, Info, Enum
|
||||
from prometheus_client import start_http_server, Info
|
||||
|
||||
from .. schema import ConfigPush, config_push_queue
|
||||
from .. log_level import LogLevel
|
||||
|
|
@ -17,73 +19,104 @@ from . producer import Producer
|
|||
from . consumer import Consumer
|
||||
|
||||
default_config_queue = config_push_queue
|
||||
config_subscriber_id = str(uuid.uuid4())
|
||||
|
||||
# Async processor
|
||||
class AsyncProcessor:
|
||||
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
# Store the identity
|
||||
self.id = params.get("id")
|
||||
|
||||
# Register a pulsar client
|
||||
self.client = PulsarClient(**params)
|
||||
|
||||
# The processor runs all activity in a taskgroup, it's mandatory
|
||||
# that this is provded
|
||||
self.taskgroup = params.get("taskgroup")
|
||||
if self.taskgroup is None:
|
||||
raise RuntimeError("Essential taskgroup missing")
|
||||
|
||||
# Pubsub parameters passed in
|
||||
if not hasattr(__class__, "params_metric"):
|
||||
__class__.params_metric = Info(
|
||||
'params', 'Parameters configuration'
|
||||
)
|
||||
|
||||
# FIXME: Maybe outputs information it should not
|
||||
# Record metrics for the processor
|
||||
__class__.params_metric.info({
|
||||
k: str(params[k])
|
||||
for k in params
|
||||
})
|
||||
|
||||
# Get the configuration topic
|
||||
self.config_push_queue = params.get(
|
||||
"config_push_queue", default_config_queue
|
||||
)
|
||||
|
||||
# This records registered configuration handlers
|
||||
self.config_handlers = []
|
||||
|
||||
# Create a random ID for this subscription to the configuration
|
||||
# service
|
||||
config_subscriber_id = str(uuid.uuid4())
|
||||
|
||||
# Subscribe to config queue
|
||||
self.config_sub_task = self.subscribe(
|
||||
flow = None,
|
||||
queue = self.config_push_queue, subscriber = config_subscriber_id,
|
||||
schema = ConfigPush, handler = self.on_config_change,
|
||||
|
||||
# This causes new subscriptions to view the entire history of
|
||||
# configuration
|
||||
start_of_messages = True
|
||||
)
|
||||
|
||||
self.running = True
|
||||
|
||||
# This is called to start dynamic behaviour. An over-ride point for
|
||||
# extra functionality
|
||||
async def start(self):
|
||||
await self.config_sub_task.start()
|
||||
|
||||
# This is called to stop all threads. An over-ride point for extra
|
||||
# functionality
|
||||
def stop(self):
|
||||
self.client.close()
|
||||
self.running = False
|
||||
|
||||
# Returns the pulsar host
|
||||
@property
|
||||
def pulsar_host(self): return self.client.pulsar_host
|
||||
|
||||
def on_config(self, handler):
|
||||
# Register a new event handler for configuration change
|
||||
def register_config_handler(self, handler):
|
||||
self.config_handlers.append(handler)
|
||||
|
||||
async def start(self):
|
||||
await self.config_sub_task.start()
|
||||
|
||||
# Called when a new configuration message push occurs
|
||||
async def on_config_change(self, message, consumer):
|
||||
|
||||
# Get configuration data and version number
|
||||
config = message.value().config
|
||||
version = message.value().version
|
||||
|
||||
# Acknowledge the message
|
||||
consumer.acknowledge(message)
|
||||
|
||||
# Invoke message handlers
|
||||
print("Config change event", config, version, flush=True)
|
||||
for ch in self.config_handlers:
|
||||
await ch(config, version)
|
||||
|
||||
# This is the 'main' body of the handler. It is a point to override
|
||||
# if needed. By default does nothing. Processors are implemented
|
||||
# by adding consumer/producer functionality so maybe nothing is needed
|
||||
# in the run() body
|
||||
async def run(self):
|
||||
while self.running:
|
||||
await asyncio.sleep(2)
|
||||
|
||||
# Subscribe to a topic, returns a Consumer
|
||||
def subscribe(
|
||||
self, flow, queue, subscriber, schema, handler, metrics=None,
|
||||
start_of_messages=False
|
||||
|
|
@ -101,8 +134,8 @@ class AsyncProcessor:
|
|||
start_of_messages=start_of_messages,
|
||||
)
|
||||
|
||||
# Open a mechanism to publish messages to a topic. Returns a subscriber
|
||||
def publish(self, queue, schema, metrics=None):
|
||||
|
||||
return Producer(
|
||||
client = self.client,
|
||||
queue = queue,
|
||||
|
|
@ -110,51 +143,45 @@ class AsyncProcessor:
|
|||
metrics = metrics,
|
||||
)
|
||||
|
||||
def set_processor_state(self, flow, state):
|
||||
|
||||
if not hasattr(__class__, "state_metric"):
|
||||
__class__.state_metric = Enum(
|
||||
'processor_state', 'Processor state',
|
||||
["flow"],
|
||||
states=['starting', 'running', 'stopped']
|
||||
)
|
||||
|
||||
__class__.state_metric.labels("flow").state(state)
|
||||
|
||||
def set_pubsub_info(self, flow, info) :
|
||||
|
||||
if not hasattr(__class__, "pubsub_metric"):
|
||||
__class__.pubsub_metric = Info(
|
||||
'pubsub', 'Pub/sub configuration',
|
||||
["flow"]
|
||||
)
|
||||
|
||||
__class__.pubsub_metric.labels(flow=flow).info(info)
|
||||
|
||||
# Startup fabric. This runs in 'async' mode, creates a taskgroup and
|
||||
# runs the producer.
|
||||
@classmethod
|
||||
async def launch_async(cls, args, ident):
|
||||
|
||||
async with asyncio.TaskGroup() as tg:
|
||||
try:
|
||||
|
||||
try:
|
||||
# Create a taskgroup. This seems complicated, when an exception
|
||||
# occurs, unhandled it looks like it cancels all threads in the
|
||||
# taskgroup, but I have observed this not working. I think
|
||||
# it's fixed now that exceptions are caught in the right place.
|
||||
async with asyncio.TaskGroup() as tg:
|
||||
|
||||
p = cls(**args | { "taskgroup": tg })
|
||||
|
||||
# FIXME: Two sort of 'ident' things going on here?
|
||||
p.module = ident
|
||||
p.config_ident = args.get("ident", "FIXME")
|
||||
# Create a processor instance, and include the taskgroup
|
||||
# as a paramter. A processor identity ident is used as
|
||||
# - subscriber name
|
||||
# - an identifier for flow configuration
|
||||
p = cls(**args | { "taskgroup": tg, "id", ident })
|
||||
|
||||
await p.start()
|
||||
# Start the processor
|
||||
await p.start()
|
||||
|
||||
task2 = tg.create_task(p.run())
|
||||
# Run the processor
|
||||
task = tg.create_task(p.run())
|
||||
|
||||
except Exception as e:
|
||||
print("Exception, closing taskgroup", flush=True)
|
||||
raise e
|
||||
# The taskgroup causes everything to wait until
|
||||
# all threads have stopped
|
||||
|
||||
# This is here to output a debug message, shouldn't be needed.
|
||||
except Exception as e:
|
||||
print("Exception, closing taskgroup", flush=True)
|
||||
raise e
|
||||
|
||||
# Startup fabric. launch calls launch_async in async mode.
|
||||
@classmethod
|
||||
def launch(cls, ident, doc):
|
||||
|
||||
# Start assembling CLI arguments
|
||||
parser = argparse.ArgumentParser(
|
||||
prog=ident,
|
||||
description=doc
|
||||
|
|
@ -166,22 +193,29 @@ class AsyncProcessor:
|
|||
help=f'Configuration identity (default: {ident})',
|
||||
)
|
||||
|
||||
# Invoke the class-specific add_args, which manages adding all the
|
||||
# command-line arguments
|
||||
cls.add_args(parser)
|
||||
|
||||
# Parse arguments
|
||||
args = parser.parse_args()
|
||||
args = vars(args)
|
||||
|
||||
# Debug
|
||||
print(args, flush=True)
|
||||
|
||||
# Start the Prometheus metrics service if needed
|
||||
if args["metrics"]:
|
||||
start_http_server(args["metrics_port"])
|
||||
|
||||
# Loop forever, exception handler
|
||||
while True:
|
||||
|
||||
print("Starting...", flush=True)
|
||||
|
||||
try:
|
||||
|
||||
# Launch the processor in an asyncio handler
|
||||
asyncio.run(cls.launch_async(
|
||||
args, ident
|
||||
))
|
||||
|
|
@ -194,6 +228,7 @@ class AsyncProcessor:
|
|||
print("Pulsar Interrupted.", flush=True)
|
||||
return
|
||||
|
||||
# Exceptions from a taskgroup come in as an exception group
|
||||
except ExceptionGroup as e:
|
||||
|
||||
print("Exception group:", flush=True)
|
||||
|
|
@ -206,10 +241,13 @@ class AsyncProcessor:
|
|||
print("Type:", type(e), flush=True)
|
||||
print("Exception:", e, flush=True)
|
||||
|
||||
# Retry occurs here
|
||||
print("Will retry...", flush=True)
|
||||
time.sleep(4)
|
||||
print("Retrying...", flush=True)
|
||||
|
||||
# The command-line arguments are built using a stack of add_args
|
||||
# invocations
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,10 @@
|
|||
|
||||
# Base class for processor with management of flows in & out which are managed
|
||||
# by configuration. This is probably all processor types, except for the
|
||||
# configuration service which can't manage itself.
|
||||
|
||||
import json
|
||||
|
||||
from pulsar.schema import JsonSchema
|
||||
|
||||
from .. schema import Error
|
||||
|
|
@ -13,45 +18,51 @@ from .. base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
|
|||
class Flow:
|
||||
pass
|
||||
|
||||
# Parent class for configurable processors, configured with flows by
|
||||
# the config service
|
||||
class FlowProcessor(AsyncProcessor):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
self.id = params.get("id")
|
||||
self.subscriber = params.get("subscriber")
|
||||
|
||||
ProcessorMetrics(id=self.id).info(
|
||||
{
|
||||
"subscriber": self.subscriber,
|
||||
}
|
||||
)
|
||||
# Initialise base class
|
||||
super(FlowProcessor, self).__init__(**params)
|
||||
|
||||
super(FlowProcessor, self).__init__(
|
||||
**params | {
|
||||
"id": self.id,
|
||||
}
|
||||
)
|
||||
# Initialise metrics, records the parameters
|
||||
ProcessorMetrics(id=self.id).info(params)
|
||||
|
||||
self.on_config(self.on_configuration)
|
||||
# Register configuration handler
|
||||
self.register_config_handler(self.on_configuration)
|
||||
|
||||
# Initialise flow information state
|
||||
self.flows = {}
|
||||
|
||||
# These can be overriden by a derived class
|
||||
# These can be overriden by a derived class:
|
||||
|
||||
# Consumer specification, array of ("name", SchemaType, handler)
|
||||
self.consumer_spec = []
|
||||
|
||||
# Producer specification, array of ("name", SchemaType)
|
||||
self.producer_spec = []
|
||||
|
||||
# Configuration specification, collects some flow variables from
|
||||
# config, array of "name"
|
||||
self.config_spec = []
|
||||
|
||||
print("Service initialised.")
|
||||
|
||||
# Register a new consumer name
|
||||
def register_consumer(self, name, schema, handler):
|
||||
self.consumer_spec.append((name, schema, handler))
|
||||
|
||||
# Register a producer name
|
||||
def register_producer(self, name, schema):
|
||||
self.producer_spec.append((name, schema))
|
||||
|
||||
# Register a configuration variable
|
||||
def register_config(self, name):
|
||||
self.config_spec.append((name,))
|
||||
self.config_spec.append(name)
|
||||
|
||||
# Start processing for a new flow
|
||||
async def start_flow(self, flow, defn):
|
||||
|
||||
flow_obj = Flow()
|
||||
|
|
@ -117,6 +128,7 @@ class FlowProcessor(AsyncProcessor):
|
|||
|
||||
print("Started flow: ", flow)
|
||||
|
||||
# Stop processing for a new flow
|
||||
async def stop_flow(self, flow):
|
||||
|
||||
for c in self.flows[flow].consumer:
|
||||
|
|
@ -126,43 +138,50 @@ class FlowProcessor(AsyncProcessor):
|
|||
|
||||
print("Stopped flow: ", flow, flush=True)
|
||||
|
||||
# Event handler - called for a configuration change
|
||||
async def on_configuration(self, config, version):
|
||||
|
||||
print("Got config version", version, flush=True)
|
||||
|
||||
# Skip over invalid data
|
||||
if "flows" not in config: return
|
||||
|
||||
# Check there's configuration information for me
|
||||
if self.id in config["flows"]:
|
||||
|
||||
# Get my flow config
|
||||
flow_config = json.loads(config["flows"][self.id])
|
||||
|
||||
wanted_keys = flow_config.keys()
|
||||
current_keys = self.flows.keys()
|
||||
# Get list of flows which should be running and are currently
|
||||
# running
|
||||
wanted_flows = flow_config.keys()
|
||||
current_flows = self.flows.keys()
|
||||
|
||||
for key in wanted_keys:
|
||||
if key not in current_keys:
|
||||
await self.start_flow(key, flow_config[key])
|
||||
# Start all the flows which arent currently running
|
||||
for flow in wanted_flows:
|
||||
if flow not in current_flows:
|
||||
await self.start_flow(flow, flow_config[flow])
|
||||
|
||||
for key in current_keys:
|
||||
if key not in wanted_keys:
|
||||
await self.stop_flow(key)
|
||||
# Stop all the unwanted flows which are due to be stopped
|
||||
for flow in current_flows:
|
||||
if flow not in wanted_flows:
|
||||
await self.stop_flow(flow)
|
||||
|
||||
print("Handled config update")
|
||||
|
||||
else:
|
||||
|
||||
print("No configuration settings for me!", flush=True)
|
||||
|
||||
# Start threads, just call parent
|
||||
async def start(self):
|
||||
await super(FlowProcessor, self).start()
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser, default_subscriber):
|
||||
def add_args(parser):
|
||||
|
||||
AsyncProcessor.add_args(parser)
|
||||
|
||||
parser.add_argument(
|
||||
'-s', '--subscriber',
|
||||
default=default_subscriber,
|
||||
help=f'Queue subscriber name (default: {default_subscriber})'
|
||||
)
|
||||
|
||||
# parser.add_argument(
|
||||
# '--rate-limit-retry',
|
||||
# type=int,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue