mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-18 01:31:02 +02:00
Working on processor API
This commit is contained in:
Cyber MacGeddon
parent
e66c06f99a
commit
1d150cdc21
9 changed files with 588 additions and 547 deletions
|
|
@ -1,110 +1,45 @@
|
|||
|
||||
import asyncio
|
||||
from pulsar.schema import JsonSchema
|
||||
import pulsar
|
||||
from prometheus_client import Histogram, Info, Counter, Enum
|
||||
import time
|
||||
|
||||
from . base_processor import BaseProcessor
|
||||
from .. exceptions import TooManyRequests
|
||||
|
||||
default_rate_limit_retry = 10
|
||||
default_rate_limit_timeout = 7200
|
||||
class Consumer:
|
||||
|
||||
class Consumer(BaseProcessor):
|
||||
def __init__(
|
||||
self, taskgroup, client, queue, subscriber, schema,
|
||||
handler
|
||||
):
|
||||
|
||||
def __init__(self, **params):
|
||||
self.taskgroup = taskgroup
|
||||
self.client = client
|
||||
self.queue = queue
|
||||
self.subscriber = subscriber
|
||||
self.schema = schema
|
||||
self.handler = handler
|
||||
|
||||
if not hasattr(__class__, "state_metric"):
|
||||
__class__.state_metric = Enum(
|
||||
'processor_state', 'Processor state',
|
||||
states=['starting', 'running', 'stopped']
|
||||
)
|
||||
__class__.state_metric.state('starting')
|
||||
self.running = True
|
||||
self.task = None
|
||||
|
||||
__class__.state_metric.state('starting')
|
||||
async def start(self):
|
||||
|
||||
super(Consumer, self).__init__(**params)
|
||||
self.running = True
|
||||
|
||||
self.subscriber = params.get("subscriber")
|
||||
self.input_schema = params.get("input_schema")
|
||||
|
||||
self.rate_limit_retry = params.get(
|
||||
"rate_limit_retry", default_rate_limit_retry
|
||||
)
|
||||
self.rate_limit_timeout = params.get(
|
||||
"rate_limit_timeout", default_rate_limit_timeout
|
||||
self.consumer = self.client.subscribe(
|
||||
self.queue, self.subscriber, self.schema
|
||||
)
|
||||
|
||||
if self.input_schema == None:
|
||||
raise RuntimeError("input_schema must be specified")
|
||||
self.task = self.taskgroup.create_task(self.run())
|
||||
|
||||
if not hasattr(__class__, "request_metric"):
|
||||
__class__.request_metric = Histogram(
|
||||
'request_latency', 'Request latency (seconds)',
|
||||
["flow"]
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "pubsub_metric"):
|
||||
__class__.pubsub_metric = Info(
|
||||
'pubsub', 'Pub/sub configuration',
|
||||
["flow"]
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "processing_metric"):
|
||||
__class__.processing_metric = Counter(
|
||||
'processing_count', 'Processing count',
|
||||
["flow", "status"]
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "rate_limit_metric"):
|
||||
__class__.rate_limit_metric = Counter(
|
||||
'rate_limit_count', 'Rate limit event count',
|
||||
["flow"]
|
||||
)
|
||||
|
||||
# __class__.pubsub_metric.info({
|
||||
# "input_queue": self.input_queue,
|
||||
# "subscriber": self.subscriber,
|
||||
# "input_schema": self.input_schema.__name__,
|
||||
# "rate_limit_retry": str(self.rate_limit_retry),
|
||||
# "rate_limit_timeout": str(self.rate_limit_timeout),
|
||||
# })
|
||||
|
||||
# self.consumer = self.client.subscribe(
|
||||
# self.input_queue, self.subscriber,
|
||||
# consumer_type=pulsar.ConsumerType.Shared,
|
||||
# schema=JsonSchema(self.input_schema),
|
||||
# )
|
||||
|
||||
self.config_queues = {}
|
||||
self.config_handlers = [self.on_config_pubsub]
|
||||
|
||||
print("Initialised consumer.", flush=True)
|
||||
|
||||
async def on_config_pubsub(self, version, config):
|
||||
|
||||
print("Configuring using config", version)
|
||||
|
||||
flows = config.get("flows", {}).get(self.ident, {})
|
||||
|
||||
print(flows)
|
||||
|
||||
print("Configuration completed OK")
|
||||
|
||||
async def run(self):
|
||||
|
||||
__class__.state_metric.state('running')
|
||||
|
||||
while True:
|
||||
|
||||
await asyncio.sleep(1)
|
||||
|
||||
continue
|
||||
while self.running:
|
||||
|
||||
msg = await asyncio.to_thread(self.consumer.receive)
|
||||
|
||||
expiry = time.time() + self.rate_limit_timeout
|
||||
# expiry = time.time() + self.rate_limit_timeout
|
||||
expiry = time.time() + 10
|
||||
|
||||
# This loop is for retry on rate-limit / resource limits
|
||||
while True:
|
||||
|
|
@ -117,20 +52,24 @@ class Consumer(BaseProcessor):
|
|||
# be retried
|
||||
self.consumer.negative_acknowledge(msg)
|
||||
|
||||
__class__.processing_metric.labels(status="error").inc()
|
||||
# FIXME
|
||||
# __class__.processing_metric.labels(status="error").inc()
|
||||
|
||||
# Break out of retry loop, processes next message
|
||||
break
|
||||
|
||||
try:
|
||||
|
||||
with __class__.request_metric.time():
|
||||
await self.handle(msg)
|
||||
print("Handle...")
|
||||
# FIXME
|
||||
# with __class__.request_metric.time():
|
||||
await self.handler(msg, self.consumer)
|
||||
print("Handled.")
|
||||
|
||||
# Acknowledge successful processing of the message
|
||||
self.consumer.acknowledge(msg)
|
||||
|
||||
__class__.processing_metric.labels(status="success").inc()
|
||||
# __class__.processing_metric.labels(status="success").inc()
|
||||
|
||||
# Break out of retry loop
|
||||
break
|
||||
|
|
@ -139,14 +78,15 @@ class Consumer(BaseProcessor):
|
|||
|
||||
print("TooManyRequests: will retry...", flush=True)
|
||||
|
||||
__class__.rate_limit_metric.inc()
|
||||
# FIXME
|
||||
# __class__.rate_limit_metric.inc()
|
||||
|
||||
# Sleep
|
||||
time.sleep(self.rate_limit_retry)
|
||||
|
||||
# Contine from retry loop, just causes a reprocessing
|
||||
continue
|
||||
|
||||
|
||||
except Exception as e:
|
||||
|
||||
print("Exception:", e, flush=True)
|
||||
|
|
@ -155,33 +95,8 @@ class Consumer(BaseProcessor):
|
|||
# be retried
|
||||
self.consumer.negative_acknowledge(msg)
|
||||
|
||||
__class__.processing_metric.labels(status="error").inc()
|
||||
# __class__.processing_metric.labels(status="error").inc()
|
||||
|
||||
# Break out of retry loop, processes next message
|
||||
break
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser, default_input_queue, default_subscriber):
|
||||
|
||||
BaseProcessor.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,
|
||||
default=default_rate_limit_retry,
|
||||
help=f'Rate limit retry (default: {default_rate_limit_retry})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--rate-limit-timeout',
|
||||
type=int,
|
||||
default=default_rate_limit_timeout,
|
||||
help=f'Rate limit timeout (default: {default_rate_limit_timeout})'
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue