fix(api): await Kafka acknowledgements and bound HTTP requests
This commit is contained in:
@@ -1,185 +1,91 @@
|
||||
import asyncio
|
||||
import logging
|
||||
"""Kafka publications complete only on broker delivery callbacks."""
|
||||
import threading
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
from concurrent.futures import Future
|
||||
|
||||
from app.common.config import config
|
||||
from app.common.logging import setup_logger
|
||||
|
||||
logger = setup_logger("ingestion-api")
|
||||
|
||||
|
||||
class MessageProducer(ABC):
|
||||
@abstractmethod
|
||||
def produce(self, topic: str, key: Optional[str], value: bytes) -> bool:
|
||||
"""Produce a message synchronously or enqueue in broker buffer."""
|
||||
pass
|
||||
def produce(self, topic: str, key: str | None, value: bytes) -> Future:
|
||||
"""Return a future resolved by the destination acknowledgement."""
|
||||
|
||||
@abstractmethod
|
||||
def is_healthy(self) -> bool:
|
||||
"""Check if producer is healthy and connected to broker."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def flush(self, timeout: float = 5.0) -> int:
|
||||
"""Flush any pending messages."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def close(self) -> None:
|
||||
"""Close producer connection."""
|
||||
def close(self, timeout: float = 5.0) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class ConfluentKafkaProducer(MessageProducer):
|
||||
"""
|
||||
Production Kafka producer implementation wrapping confluent-kafka.
|
||||
Uses 'raw_id' as partition key and LZ4 compression.
|
||||
"""
|
||||
|
||||
def __init__(self, bootstrap_servers: Optional[str] = None):
|
||||
try:
|
||||
import confluent_kafka
|
||||
self._ck = confluent_kafka
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"confluent-kafka library is not installed. Run 'pip install confluent-kafka'."
|
||||
)
|
||||
|
||||
self.bootstrap_servers = bootstrap_servers or config.kafka.bootstrap_servers
|
||||
self._producer_config = {
|
||||
"bootstrap.servers": self.bootstrap_servers,
|
||||
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": 100000,
|
||||
"queue.buffering.max.ms": 5,
|
||||
"queue.buffering.max.messages": config.api.queue_max_size,
|
||||
"linger.ms": 5,
|
||||
"acks": "all",
|
||||
"retries": 3,
|
||||
}
|
||||
self._producer = self._ck.Producer(self._producer_config)
|
||||
self._healthy = True
|
||||
"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 _delivery_report(self, err: Any, msg: Any) -> None:
|
||||
if err is not None:
|
||||
self._healthy = False
|
||||
raw_id = msg.key().decode("utf-8") if msg and msg.key() else None
|
||||
logger.error(
|
||||
f"Kafka delivery failed: {err}",
|
||||
extra={"raw_id": raw_id, "event": "kafka_produce_failed"},
|
||||
)
|
||||
else:
|
||||
self._healthy = True
|
||||
def _poll(self):
|
||||
while not self._stop.is_set():
|
||||
self._producer.poll(0.01)
|
||||
|
||||
def produce(self, topic: str, key: Optional[str], value: bytes) -> bool:
|
||||
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:
|
||||
key_bytes = key.encode("utf-8") if key else None
|
||||
self._producer.produce(
|
||||
topic=topic,
|
||||
key=key_bytes,
|
||||
value=value,
|
||||
on_delivery=self._delivery_report,
|
||||
)
|
||||
# Serve delivery callback queue non-blockingly
|
||||
self._producer.poll(0)
|
||||
return True
|
||||
except BufferError:
|
||||
logger.warning(
|
||||
"Kafka local producer buffer full",
|
||||
extra={"raw_id": key, "event": "kafka_buffer_full"},
|
||||
)
|
||||
self._producer.poll(0.1)
|
||||
return False
|
||||
except Exception as e:
|
||||
self._healthy = False
|
||||
logger.error(
|
||||
f"Kafka produce error: {e}",
|
||||
extra={"raw_id": key, "event": "kafka_error"},
|
||||
)
|
||||
return False
|
||||
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 is_healthy(self) -> bool:
|
||||
return self._healthy
|
||||
|
||||
def flush(self, timeout: float = 5.0) -> int:
|
||||
def flush(self, timeout=5.0):
|
||||
return self._producer.flush(timeout)
|
||||
|
||||
def close(self) -> None:
|
||||
self.flush()
|
||||
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):
|
||||
"""
|
||||
In-memory mock producer for offline testing and verification.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.messages = []
|
||||
self._healthy = True
|
||||
|
||||
def produce(self, topic: str, key: Optional[str], value: bytes) -> bool:
|
||||
if not self._healthy:
|
||||
return False
|
||||
self.messages.append({"topic": topic, "key": key, "value": value})
|
||||
return 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) -> bool:
|
||||
def is_healthy(self):
|
||||
return self._healthy
|
||||
|
||||
def set_healthy(self, healthy: bool) -> None:
|
||||
def set_healthy(self, healthy):
|
||||
self._healthy = healthy
|
||||
|
||||
def flush(self, timeout: float = 5.0) -> int:
|
||||
def flush(self, timeout=5.0):
|
||||
return 0
|
||||
|
||||
def close(self) -> None:
|
||||
def close(self, timeout=5.0):
|
||||
pass
|
||||
|
||||
|
||||
class AsyncProducerWorker:
|
||||
"""
|
||||
Background worker consuming from the bounded async queue and writing to the Kafka producer.
|
||||
"""
|
||||
|
||||
def __init__(self, queue: asyncio.Queue, producer: MessageProducer, topic: str):
|
||||
self.queue = queue
|
||||
self.producer = producer
|
||||
self.topic = topic
|
||||
self._running = False
|
||||
self._task: Optional[asyncio.Task] = None
|
||||
|
||||
async def start(self) -> None:
|
||||
self._running = True
|
||||
self._task = asyncio.create_task(self._run_loop())
|
||||
|
||||
async def _run_loop(self) -> None:
|
||||
while self._running:
|
||||
try:
|
||||
item = await self.queue.get()
|
||||
if item is None: # Sentinel to stop
|
||||
self.queue.task_done()
|
||||
break
|
||||
|
||||
key, value = item
|
||||
success = self.producer.produce(self.topic, key, value)
|
||||
if not success:
|
||||
logger.error(
|
||||
"Failed to publish message from async queue to Kafka",
|
||||
extra={"raw_id": key, "event": "queue_dispatch_failed"},
|
||||
)
|
||||
self.queue.task_done()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Unexpected error in async producer worker: {e}",
|
||||
extra={"event": "worker_unexpected_error"},
|
||||
)
|
||||
|
||||
async def stop(self) -> None:
|
||||
self._running = False
|
||||
if self._task:
|
||||
self._task.cancel()
|
||||
try:
|
||||
await self._task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self.producer.flush()
|
||||
|
||||
Reference in New Issue
Block a user