fix(worker): commit confirmed partition offsets and drain safely

app/worker/consumer.py

@@ -6,6 +6,7 @@ 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
@@ -47,8 +48,12 @@ class ProcessingWorker:
 
         # Micro-batch buffer: list of (NormalizedRecord, message)
         self.buffer: List[Tuple[NormalizedRecord, Any]] = []
-        self.last_flush_time = time.time()
+        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:
@@ -59,7 +64,8 @@ class ProcessingWorker:
                     "bootstrap.servers": config.kafka.bootstrap_servers,
                     "group.id": config.kafka.consumer_group,
                     "auto.offset.reset": config.kafka.auto_offset_reset,
-                    "enable.auto.commit": False,  # Manual offset commit post ClickHouse ACK
+                    "enable.auto.commit": False,
+                    "enable.auto.offset.store": False,
                 }
                 self._consumer = confluent_kafka.Consumer(conf)
             except ImportError:
@@ -74,120 +80,103 @@ class ProcessingWorker:
             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 flush_buffer(self) -> None:
-        """
-        Flushes buffered normalized records to ClickHouse and commits Kafka offsets.
-        If ClickHouse write fails all retries, routes batch to DLQ.
-        """
-        if not self.buffer:
-            self.last_flush_time = time.time()
-            return
+    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)
 
-        batch_records = [item[0] for item in self.buffer]
-        batch_messages = [item[1] for item in self.buffer]
-        batch_count = len(batch_records)
-
-        start_time = time.time()
-        try:
-            self.clickhouse.insert_batch(batch_records)
-            flush_duration = time.time() - start_time
-
-            # Commit offsets only after ClickHouse ACK
-            if self._consumer and hasattr(self._consumer, "commit"):
-                try:
-                    self._consumer.commit(asynchronous=False)
-                except Exception as e:
-                    logger.error(
-                        f"Kafka offset commit failed after ClickHouse ACK: {e}",
-                        extra={"event": "kafka_commit_failed"},
-                    )
-
-            # Record metrics
-            now_iso = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
-            worker_metrics.set_last_flush_at(now_iso)
-            worker_metrics.batch_flush_duration.observe(flush_duration)
-            worker_metrics.batch_size_at_flush.observe(batch_count)
-
-            for rec in batch_records:
-                worker_metrics.record_processed(
-                    rec.service_type, rec.radio_access_type, 1
-                )
-
-            logger.info(
-                f"Successfully flushed {batch_count} records to ClickHouse in {flush_duration:.3f}s",
-                extra={"event": "clickhouse_batch_flushed"},
-            )
-
-        except Exception as e:
-            logger.error(
-                f"Batch insert to ClickHouse failed after retries: {e}. Routing batch to DLQ.",
-                extra={"event": "ch_batch_fatal_error"},
-            )
-            # Route all records in batch to DLQ to prevent pipeline deadlock
-            for rec, msg in self.buffer:
-                route_to_dlq(
-                    producer=self.dlq_producer,
-                    topic=self.dlq_topic,
-                    original_payload=vars(rec),
-                    failure_reason="ch_insert_failed",
-                    failed_field=None,
-                    raw_id=rec.raw_id,
-                )
-
-            # Commit offsets to prevent infinite replay stall
-            if self._consumer and hasattr(self._consumer, "commit"):
-                try:
-                    self._consumer.commit(asynchronous=False)
-                except Exception as ce:
-                    logger.error(f"Commit after DLQ fallback failed: {ce}")
-
-        finally:
-            self.buffer.clear()
-            worker_metrics.set_in_flight_buffer(0)
-            self.last_flush_time = time.time()
-
-    def process_single_message(self, message: Any) -> None:
-        """
-        Parses, validates, and either buffers or routes message to DLQ.
-        """
-        value_bytes = message.value() if hasattr(message, "value") else message
-        if value_bytes is None:
-            return
-
-        record, failure_reason, failed_field, raw_id = normalize_cdr_payload(
-            value_bytes
+    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())
+            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}")
 
-        if failure_reason:
-            payload_str = (
-                value_bytes.decode("utf-8", errors="replace")
-                if isinstance(value_bytes, bytes)
-                else str(value_bytes)
-            )
-            route_to_dlq(
-                producer=self.dlq_producer,
-                topic=self.dlq_topic,
-                original_payload=payload_str,
-                failure_reason=failure_reason,
-                failed_field=failed_field,
-                raw_id=raw_id,
-            )
-            # Commit DLQ record offset so it is not reprocessed
-            if self._consumer and hasattr(self._consumer, "commit"):
-                try:
-                    self._consumer.commit(message=message, asynchronous=False)
-                except Exception:
-                    pass
+    @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:
+            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),
+                                 "ch_insert_failed", 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))
+            self._commit([message for _, message in self.buffer])
+        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()
 
-        if record:
-            self.buffer.append((record, message))
-            worker_metrics.set_in_flight_buffer(len(self.buffer))
-
-            # Flush if batch size threshold reached
-            if len(self.buffer) >= self.batch_size:
+    def process_single_message(self, message):
+        try:
+            record, reason, field, raw_id = normalize_cdr_payload(message.value() or b"")
+            if reason:
                 self.flush_buffer()
+                route_to_dlq(self.dlq_producer, self.dlq_topic, self._original(message),
+                             reason, field, raw_id)
+                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:
         """
@@ -199,6 +188,9 @@ class ProcessingWorker:
         # 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
@@ -210,6 +202,12 @@ class ProcessingWorker:
 
         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)
 
@@ -217,7 +215,7 @@ class ProcessingWorker:
                     # Check flush timer on idle / interval
                     if (
                         self.buffer
-                        and (time.time() - self.last_flush_time) >= self.flush_interval
+                        and (time.monotonic() - self.last_flush_time) >= self.flush_interval
                     ):
                         self.flush_buffer()
                     continue
@@ -234,10 +232,13 @@ class ProcessingWorker:
                 # Check flush interval during message stream
                 if (
                     self.buffer
-                    and (time.time() - self.last_flush_time) >= self.flush_interval
+                    and (time.monotonic() - self.last_flush_time) >= self.flush_interval
                 ):
                     self.flush_buffer()
 
+        except Exception:
+            self.failed = True
+            raise
         finally:
             self.stop()
 
@@ -249,18 +250,23 @@ class ProcessingWorker:
             f"Stopping worker: flushing {len(self.buffer)} remaining in-memory records...",
             extra={"event": "graceful_shutdown_flushing"},
         )
-        if self.buffer:
-            self.flush_buffer()
-
-        if self._consumer and hasattr(self._consumer, "close"):
-            self._consumer.close()
-
-        self.dlq_producer.close()
-        self.clickhouse.close()
-        self.internal_server.stop()
-        logger.info("Processing Worker stopped cleanly", extra={"event": "worker_stopped"})
+        try:
+            if self.buffer 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__":
-    worker = ProcessingWorker()
-    worker.start()
+    ProcessingWorker().start()