mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 08:26:21 +02:00
e10e9d2295
* Concurrency in consumers * Add concurrency to consumer spec * Add concurrency command-line option to all LLMs (default 1)
179 lines
4.9 KiB
Python
179 lines
4.9 KiB
Python
|
|
# Subscriber is similar to consumer: It provides a service to take stuff
|
|
# off of a queue and make it available using an internal broker system,
|
|
# so suitable for when multiple recipients are reading from the same queue
|
|
|
|
from pulsar.schema import JsonSchema
|
|
import asyncio
|
|
import _pulsar
|
|
import time
|
|
|
|
class Subscriber:
|
|
|
|
def __init__(self, client, topic, subscription, consumer_name,
|
|
schema=None, max_size=100, metrics=None):
|
|
self.client = client
|
|
self.topic = topic
|
|
self.subscription = subscription
|
|
self.consumer_name = consumer_name
|
|
self.schema = schema
|
|
self.q = {}
|
|
self.full = {}
|
|
self.max_size = max_size
|
|
self.lock = asyncio.Lock()
|
|
self.running = True
|
|
self.metrics = metrics
|
|
self.task = None
|
|
|
|
self.consumer = None
|
|
|
|
def __del__(self):
|
|
|
|
self.running = False
|
|
|
|
async def start(self):
|
|
|
|
self.consumer = self.client.subscribe(
|
|
topic = self.topic,
|
|
subscription_name = self.subscription,
|
|
consumer_name = self.consumer_name,
|
|
schema = JsonSchema(self.schema),
|
|
)
|
|
|
|
self.task = asyncio.create_task(self.run())
|
|
|
|
async def stop(self):
|
|
self.running = False
|
|
|
|
if self.task:
|
|
await self.task
|
|
|
|
async def join(self):
|
|
await self.stop()
|
|
|
|
if self.task:
|
|
await self.task
|
|
|
|
async def run(self):
|
|
|
|
while self.running:
|
|
|
|
if self.metrics:
|
|
self.metrics.state("stopped")
|
|
|
|
try:
|
|
|
|
if self.metrics:
|
|
self.metrics.state("running")
|
|
|
|
print("Subscriber running...", flush=True)
|
|
|
|
while self.running:
|
|
|
|
try:
|
|
msg = await asyncio.to_thread(
|
|
self.consumer.receive,
|
|
timeout_millis=250
|
|
)
|
|
except _pulsar.Timeout:
|
|
continue
|
|
except Exception as e:
|
|
print("Exception:", e, flush=True)
|
|
print(type(e))
|
|
raise e
|
|
|
|
if self.metrics:
|
|
self.metrics.received()
|
|
|
|
# Acknowledge successful reception of the message
|
|
self.consumer.acknowledge(msg)
|
|
|
|
try:
|
|
id = msg.properties()["id"]
|
|
except:
|
|
id = None
|
|
|
|
value = msg.value()
|
|
|
|
async with self.lock:
|
|
|
|
# FIXME: Hard-coded timeouts
|
|
|
|
if id in self.q:
|
|
|
|
try:
|
|
# FIXME: Timeout means data goes missing
|
|
await asyncio.wait_for(
|
|
self.q[id].put(value),
|
|
timeout=1
|
|
)
|
|
|
|
except Exception as e:
|
|
self.metrics.dropped()
|
|
print("Q Put:", e, flush=True)
|
|
|
|
for q in self.full.values():
|
|
try:
|
|
# FIXME: Timeout means data goes missing
|
|
await asyncio.wait_for(
|
|
q.put(value),
|
|
timeout=1
|
|
)
|
|
except Exception as e:
|
|
self.metrics.dropped()
|
|
print("Q Put:", e, flush=True)
|
|
|
|
except Exception as e:
|
|
print("Subscriber exception:", e, flush=True)
|
|
|
|
finally:
|
|
|
|
if self.consumer:
|
|
self.consumer.unsubscribe()
|
|
self.consumer.close()
|
|
self.consumer = None
|
|
|
|
|
|
if self.metrics:
|
|
self.metrics.state("stopped")
|
|
|
|
if not self.running:
|
|
return
|
|
|
|
# If handler drops out, sleep a retry
|
|
await asyncio.sleep(1)
|
|
|
|
async def subscribe(self, id):
|
|
|
|
async with self.lock:
|
|
|
|
q = asyncio.Queue(maxsize=self.max_size)
|
|
self.q[id] = q
|
|
|
|
return q
|
|
|
|
async def unsubscribe(self, id):
|
|
|
|
async with self.lock:
|
|
|
|
if id in self.q:
|
|
# self.q[id].shutdown(immediate=True)
|
|
del self.q[id]
|
|
|
|
async def subscribe_all(self, id):
|
|
|
|
async with self.lock:
|
|
|
|
q = asyncio.Queue(maxsize=self.max_size)
|
|
self.full[id] = q
|
|
|
|
return q
|
|
|
|
async def unsubscribe_all(self, id):
|
|
|
|
async with self.lock:
|
|
|
|
if id in self.full:
|
|
# self.full[id].shutdown(immediate=True)
|
|
del self.full[id]
|
|
|