mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-17 09:11:03 +02:00
181 lines
4.8 KiB
Python
181 lines
4.8 KiB
Python
|
|
import _pulsar
|
|
import asyncio
|
|
import time
|
|
|
|
from .. exceptions import TooManyRequests
|
|
|
|
class Consumer:
|
|
|
|
def __init__(
|
|
self, taskgroup, flow, client, queue, subscriber, schema,
|
|
handler, rate_limit_retry_time = 10,
|
|
rate_limit_timeout = 7200, metrics = None,
|
|
start_of_messages=False,
|
|
):
|
|
|
|
self.taskgroup = taskgroup
|
|
self.flow = flow
|
|
self.client = client
|
|
self.queue = queue
|
|
self.subscriber = subscriber
|
|
self.schema = schema
|
|
self.handler = handler
|
|
self.rate_limit_retry_time = rate_limit_retry_time
|
|
self.rate_limit_timeout = rate_limit_timeout
|
|
self.start_of_messages = start_of_messages
|
|
|
|
self.running = True
|
|
self.task = None
|
|
|
|
self.metrics = metrics
|
|
|
|
self.consumer = None
|
|
|
|
def __del__(self):
|
|
self.running = False
|
|
|
|
if hasattr(self, "consumer"):
|
|
self.consumer.close()
|
|
|
|
async def stop(self):
|
|
|
|
self.running = False
|
|
await self.task
|
|
|
|
async def start(self):
|
|
|
|
self.running = True
|
|
|
|
# Puts it in the stopped state, the run thread should set running
|
|
if self.metrics:
|
|
self.metrics.state("stopped")
|
|
|
|
self.task = self.taskgroup.create_task(self.run())
|
|
|
|
async def run(self):
|
|
|
|
if self.metrics:
|
|
self.metrics.state("stopped")
|
|
|
|
while self.running:
|
|
|
|
try:
|
|
|
|
print(self.queue, "subscribing...", flush=True)
|
|
|
|
self.consumer = await asyncio.to_thread(
|
|
self.client.subscribe,
|
|
queue = self.queue, subscriber = self.subscriber,
|
|
schema = self.schema,
|
|
start_of_messages = self.start_of_messages
|
|
)
|
|
|
|
except Exception as e:
|
|
|
|
print("Exception:", e, flush=True)
|
|
await asyncio.sleep(2)
|
|
continue
|
|
|
|
print(self.queue, "subscribed", flush=True)
|
|
|
|
if self.metrics:
|
|
self.metrics.state("running")
|
|
|
|
try:
|
|
|
|
await self.consume()
|
|
|
|
if self.metrics:
|
|
self.metrics.state("stopped")
|
|
|
|
except Exception as e:
|
|
|
|
print("Exception:", e, flush=True)
|
|
self.consumer.close()
|
|
self.consumer = None
|
|
await asyncio.sleep(2)
|
|
continue
|
|
|
|
async def consume(self):
|
|
|
|
while self.running:
|
|
|
|
try:
|
|
msg = await asyncio.to_thread(
|
|
self.consumer.receive,
|
|
timeout_millis=2000
|
|
)
|
|
except _pulsar.Timeout:
|
|
continue
|
|
except Exception as e:
|
|
raise e
|
|
|
|
expiry = time.time() + self.rate_limit_timeout
|
|
|
|
# This loop is for retry on rate-limit / resource limits
|
|
while self.running:
|
|
|
|
if time.time() > expiry:
|
|
|
|
print("Gave up waiting for rate-limit retry", flush=True)
|
|
|
|
# Message failed to be processed, this causes it to
|
|
# be retried
|
|
self.consumer.negative_acknowledge(msg)
|
|
|
|
if self.metrics:
|
|
self.metrics.process("error")
|
|
|
|
# Break out of retry loop, processes next message
|
|
break
|
|
|
|
try:
|
|
|
|
print("Handle...", flush=True)
|
|
|
|
if self.metrics:
|
|
|
|
with self.metrics.record_time():
|
|
await self.handler(msg, self, self.flow)
|
|
|
|
else:
|
|
await self.handler(msg, self.consumer)
|
|
|
|
print("Handled.", flush=True)
|
|
|
|
# Acknowledge successful processing of the message
|
|
self.consumer.acknowledge(msg)
|
|
|
|
if self.metrics:
|
|
self.metrics.process("success")
|
|
|
|
# Break out of retry loop
|
|
break
|
|
|
|
except TooManyRequests:
|
|
|
|
print("TooManyRequests: will retry...", flush=True)
|
|
|
|
if self.metrics:
|
|
self.metrics.rate_limit()
|
|
|
|
# Sleep
|
|
await asyncio.sleep(self.rate_limit_retry)
|
|
|
|
# Contine from retry loop, just causes a reprocessing
|
|
continue
|
|
|
|
except Exception as e:
|
|
|
|
print("Exception:", e, flush=True)
|
|
|
|
# Message failed to be processed, this causes it to
|
|
# be retried
|
|
self.consumer.negative_acknowledge(msg)
|
|
|
|
if self.metrics:
|
|
self.metrics.process("error")
|
|
|
|
# Break out of retry loop, processes next message
|
|
break
|