92 lines
2.8 KiB
Python
92 lines
2.8 KiB
Python
"""Kafka publications complete only on broker delivery callbacks."""
|
|
import threading
|
|
from abc import ABC, abstractmethod
|
|
from concurrent.futures import Future
|
|
|
|
from app.common.config import config
|
|
|
|
|
|
class MessageProducer(ABC):
|
|
@abstractmethod
|
|
def produce(self, topic: str, key: str | None, value: bytes) -> Future:
|
|
"""Return a future resolved by the destination acknowledgement."""
|
|
|
|
@abstractmethod
|
|
def close(self, timeout: float = 5.0) -> None:
|
|
pass
|
|
|
|
|
|
class ConfluentKafkaProducer(MessageProducer):
|
|
def __init__(self, bootstrap_servers=None):
|
|
from confluent_kafka import Producer
|
|
self._producer = Producer({
|
|
"bootstrap.servers": bootstrap_servers or config.kafka.bootstrap_servers,
|
|
"compression.type": "lz4",
|
|
"queue.buffering.max.messages": config.api.queue_max_size,
|
|
"linger.ms": 5,
|
|
"acks": "all",
|
|
"enable.idempotence": True,
|
|
"delivery.timeout.ms": int(config.kafka.delivery_timeout_seconds * 1000),
|
|
})
|
|
self._stop = threading.Event()
|
|
self._closed = False
|
|
self._thread = threading.Thread(target=self._poll, name="KafkaDelivery", daemon=True)
|
|
self._thread.start()
|
|
|
|
def _poll(self):
|
|
while not self._stop.is_set():
|
|
self._producer.poll(0.01)
|
|
|
|
def produce(self, topic, key, value):
|
|
future = Future()
|
|
def delivered(error, message):
|
|
if error is not None:
|
|
future.set_exception(RuntimeError(str(error)))
|
|
else:
|
|
future.set_result(message)
|
|
try:
|
|
if self._closed:
|
|
raise RuntimeError("Producer is closed")
|
|
self._producer.produce(topic, key=key, value=value, on_delivery=delivered)
|
|
except Exception as exc:
|
|
future.set_exception(exc)
|
|
return future
|
|
|
|
def flush(self, timeout=5.0):
|
|
return self._producer.flush(timeout)
|
|
|
|
def close(self, timeout=5.0):
|
|
self._closed = True
|
|
remaining = self.flush(timeout)
|
|
self._stop.set()
|
|
self._thread.join(timeout=1)
|
|
if remaining:
|
|
raise TimeoutError(f"{remaining} Kafka deliveries remain unconfirmed")
|
|
|
|
|
|
class MockMessageProducer(MessageProducer):
|
|
def __init__(self):
|
|
self.messages = []
|
|
self._healthy = True
|
|
|
|
def produce(self, topic, key, value):
|
|
future = Future()
|
|
if self._healthy:
|
|
self.messages.append({"topic": topic, "key": key, "value": value})
|
|
future.set_result(None)
|
|
else:
|
|
future.set_exception(ConnectionError("Broker unavailable"))
|
|
return future
|
|
|
|
def is_healthy(self):
|
|
return self._healthy
|
|
|
|
def set_healthy(self, healthy):
|
|
self._healthy = healthy
|
|
|
|
def flush(self, timeout=5.0):
|
|
return 0
|
|
|
|
def close(self, timeout=5.0):
|
|
pass
|