Files
telecom-cdr-pipeline-ai/app/worker/consumer.py

288 lines
12 KiB
Python

import signal
import time
from datetime import datetime, timezone
from typing import Any, List, Optional, Tuple
from app.api.producer import ConfluentKafkaProducer, MessageProducer
from app.common.config import config
from app.common.logging import setup_logger
from app.common.validation import PayloadError, parse_payload
from app.worker.clickhouse import BaseClickHouseClient, ClickHouseService
from app.worker.dlq import route_to_dlq
from app.worker.internal_server import WorkerInternalServer
from app.worker.metrics import worker_metrics
from app.worker.normalizer import NormalizedRecord, normalize_cdr_payload
logger = setup_logger("processing-worker")
class ProcessingWorker:
"""
Consumes raw CDRs from Kafka, normalizes schema, micro-batches into ClickHouse,
routes malformed payloads to DLQ, and manages offsets with at-least-once semantics.
"""
def __init__(
self,
consumer: Optional[Any] = None,
dlq_producer: Optional[MessageProducer] = None,
clickhouse: Optional[BaseClickHouseClient] = None,
batch_size: Optional[int] = None,
flush_interval: Optional[float] = None,
internal_server: Optional[WorkerInternalServer] = None,
):
self.batch_size = batch_size or config.batch.size
self.flush_interval = flush_interval or config.batch.flush_interval_seconds
# Kafka Consumer
self._consumer = consumer
# Kafka DLQ Producer
self.dlq_producer = dlq_producer or ConfluentKafkaProducer()
self.dlq_topic = config.kafka.dlq_topic
# ClickHouse Client
self.clickhouse = clickhouse or ClickHouseService()
# Internal HTTP Server for metrics & /stats
self.internal_server = internal_server or WorkerInternalServer()
# Micro-batch buffer: list of (NormalizedRecord, message)
self.buffer: List[Tuple[NormalizedRecord, Any]] = []
self.pending_dlq_messages: List[Any] = []
self.last_flush_time = time.monotonic()
self.running = False
self.failed = False
self._last_lag_time = 0.0
self._shutdown_started = False
self._deadline_alarm_installed = False
def _init_consumer(self) -> None:
if self._consumer is None:
try:
import confluent_kafka
conf = {
"bootstrap.servers": config.kafka.bootstrap_servers,
"group.id": config.kafka.consumer_group,
"auto.offset.reset": config.kafka.auto_offset_reset,
"enable.auto.commit": False,
"enable.auto.offset.store": False,
}
self._consumer = confluent_kafka.Consumer(conf)
except ImportError:
raise ImportError(
"confluent-kafka is required. Run 'pip install confluent-kafka'."
)
def _handle_signal(self, signum: int, frame: Any) -> None:
sig_name = "SIGTERM" if signum == signal.SIGTERM else "SIGINT"
logger.info(
f"Received {sig_name}, initiating graceful shutdown...",
extra={"event": "shutdown_signal_received"},
)
self.running = False
if not self._shutdown_started:
self._shutdown_started = True
if self._deadline_alarm_installed:
signal.setitimer(signal.ITIMER_REAL, config.api.shutdown_timeout_seconds)
def _shutdown_timeout(self, signum, frame):
# SystemExit bypasses retry/fallback Exception handlers. Uncommitted
# source offsets remain replayable instead of extending shutdown forever.
self.failed = True
self.running = False
logger.error("Shutdown deadline exceeded; uncommitted records will replay",
extra={"event": "shutdown_timeout"})
raise SystemExit(1)
def _commit(self, messages):
from confluent_kafka import TopicPartition
offsets = {}
for message in messages:
topic = message.topic() if hasattr(message, "topic") else config.kafka.raw_topic
key = (topic, message.partition())
part = message.partition() if hasattr(message, "partition") and callable(message.partition) else getattr(message, "partition", 0)
key = (topic, part)
offsets[key] = max(offsets.get(key, 0), message.offset() + 1)
result = self._consumer.commit(
offsets=[TopicPartition(topic, partition, offset) for (topic, partition), offset in offsets.items()],
asynchronous=False,
)
for partition in result or []:
if partition.error:
raise RuntimeError(f"Partition commit failed: {partition.error}")
@staticmethod
def _original(message):
data = message.value()
try:
return parse_payload(data)
except (PayloadError, TypeError, AttributeError):
return data.decode("utf-8", errors="replace") if isinstance(data, bytes) else str(data)
def flush_buffer(self):
if not self.buffer:
if self.pending_dlq_messages:
self._commit(self.pending_dlq_messages)
self.pending_dlq_messages.clear()
self.last_flush_time = time.monotonic()
return
records = [record for record, _ in self.buffer]
started = time.monotonic()
try:
try:
self.clickhouse.insert_batch(records)
except Exception as exc:
logger.error(str(exc), extra={"event": "ch_batch_fatal_error"})
for record, message in self.buffer:
route_to_dlq(self.dlq_producer, self.dlq_topic, self._original(message),
"normalization_error", raw_id=record.raw_id)
else:
worker_metrics.batch_flush_duration.observe(time.monotonic() - started)
worker_metrics.batch_size_at_flush.observe(len(records))
worker_metrics.set_last_flush_at(datetime.now(timezone.utc).isoformat())
for record, message in self.buffer:
worker_metrics.record_processed(record.service_type, record.radio_access_type)
if hasattr(message, "timestamp"):
_, timestamp_ms = message.timestamp()
if timestamp_ms is not None and timestamp_ms >= 0:
worker_metrics.kafka_to_ch_latency.observe(max(0, time.time() - timestamp_ms / 1000))
messages_to_commit = [message for _, message in self.buffer] + self.pending_dlq_messages
self._commit(messages_to_commit)
self.pending_dlq_messages.clear()
except Exception:
self.failed = True
self.running = False
raise
self.buffer.clear()
worker_metrics.set_in_flight_buffer(0)
self.last_flush_time = time.monotonic()
def process_single_message(self, message):
try:
record, reason, field, raw_id = normalize_cdr_payload(message.value() or b"")
if reason:
route_to_dlq(self.dlq_producer, self.dlq_topic, self._original(message),
reason, field, raw_id)
msg_part = message.partition() if hasattr(message, "partition") and callable(message.partition) else getattr(message, "partition", 0)
has_pending_buffer = any(
(m.partition() if hasattr(m, "partition") and callable(m.partition) else getattr(m, "partition", 0)) == msg_part
for _, m in self.buffer
)
if has_pending_buffer:
self.pending_dlq_messages.append(message)
else:
self._commit([message])
else:
self.buffer.append((record, message))
worker_metrics.set_in_flight_buffer(len(self.buffer))
if len(self.buffer) >= self.batch_size:
self.flush_buffer()
except Exception:
self.failed = True
self.running = False
raise
def update_consumer_lag(self):
partitions = self._consumer.assignment()
if not partitions:
return
committed = self._consumer.committed(partitions, timeout=1)
for partition in committed:
low, high = self._consumer.get_watermark_offsets(partition, timeout=1)
offset = partition.offset if partition.offset >= 0 else low
worker_metrics.set_consumer_lag(partition.partition, max(0, high - offset))
def start(self) -> None:
"""
Main worker execution loop with signal handling and micro-batching timer.
"""
self._init_consumer()
self.internal_server.start()
# Register signals for graceful shutdown
signal.signal(signal.SIGTERM, self._handle_signal)
signal.signal(signal.SIGINT, self._handle_signal)
if hasattr(signal, "setitimer"):
signal.signal(signal.SIGALRM, self._shutdown_timeout)
self._deadline_alarm_installed = True
self._consumer.subscribe([config.kafka.raw_topic])
self.running = True
logger.info(
f"Processing Worker started. Subscribed to '{config.kafka.raw_topic}'. "
f"Batch size: {self.batch_size}, Interval: {self.flush_interval}s",
extra={"event": "worker_started"},
)
try:
while self.running:
if time.monotonic() - self._last_lag_time >= 5:
try:
self.update_consumer_lag()
except Exception as exc:
logger.warning(str(exc), extra={"event": "consumer_lag_unavailable"})
self._last_lag_time = time.monotonic()
# Poll with short 100ms timeout
msg = self._consumer.poll(timeout=0.1)
if msg is None:
# Check flush timer on idle / interval
if (
self.buffer
and (time.monotonic() - self.last_flush_time) >= self.flush_interval
):
self.flush_buffer()
continue
if msg.error():
logger.error(
f"Consumer error: {msg.error()}",
extra={"event": "kafka_consumer_error"},
)
continue
self.process_single_message(msg)
# Check flush interval during message stream
if (
self.buffer
and (time.monotonic() - self.last_flush_time) >= self.flush_interval
):
self.flush_buffer()
except Exception:
self.failed = True
raise
finally:
self.stop()
def stop(self) -> None:
"""
Graceful shutdown: flush any in-memory records and commit offsets.
"""
logger.info(
f"Stopping worker: flushing {len(self.buffer)} remaining in-memory records...",
extra={"event": "graceful_shutdown_flushing"},
)
try:
if (self.buffer or self.pending_dlq_messages) and not self.failed:
self.flush_buffer()
finally:
try:
if self._consumer:
self._consumer.close()
finally:
try:
self.dlq_producer.close()
finally:
self.clickhouse.close()
self.internal_server.stop()
if self._deadline_alarm_installed:
signal.setitimer(signal.ITIMER_REAL, 0)
logger.info("Worker stopped", extra={"event": "worker_stopped"})
if __name__ == "__main__":
ProcessingWorker().start()