From 974d2e97441ef74d1dbc500b845b6bc777370c73 Mon Sep 17 00:00:00 2001 From: Amirhossein Khalili Date: Sun, 20 Sep 2026 10:39:26 +0330 Subject: [PATCH] fix(worker): preserve micro-batches across DLQ records --- app/worker/consumer.py | 25 ++++++++++++++++++++----- tests/test_delivery_regressions.py | 22 +++++++++++++++++++++- 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/app/worker/consumer.py b/app/worker/consumer.py index f43f9d2..588dcf0 100644 --- a/app/worker/consumer.py +++ b/app/worker/consumer.py @@ -48,6 +48,7 @@ class ProcessingWorker: # 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 @@ -100,6 +101,8 @@ class ProcessingWorker: 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()], @@ -119,6 +122,9 @@ class ProcessingWorker: 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] @@ -130,7 +136,7 @@ class ProcessingWorker: 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) + "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)) @@ -141,7 +147,9 @@ class ProcessingWorker: _, 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]) + 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 @@ -154,10 +162,17 @@ class ProcessingWorker: 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]) + 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)) @@ -251,7 +266,7 @@ class ProcessingWorker: extra={"event": "graceful_shutdown_flushing"}, ) try: - if self.buffer and not self.failed: + if (self.buffer or self.pending_dlq_messages) and not self.failed: self.flush_buffer() finally: try: diff --git a/tests/test_delivery_regressions.py b/tests/test_delivery_regressions.py index 76b8cdb..b9e49c6 100644 --- a/tests/test_delivery_regressions.py +++ b/tests/test_delivery_regressions.py @@ -96,10 +96,30 @@ class DeliverySafetyTests(unittest.TestCase): def test_invalid_record_cannot_commit_past_buffered_valid_record(self): self.worker.process_single_message(Message(payload(), offset=0)) self.worker.process_single_message(Message(b"{malformed", offset=1)) - self.assertEqual(self.events[0], ("insert", 1)) + # Micro-batching preserved: invalid record does not force premature storage flush + self.assertEqual([e for e in self.events if e[0] == "insert"], []) + self.assertEqual(self.consumer.commits, []) + self.assertEqual(len(self.worker.buffer), 1) + # Flush commits both valid and completed DLQ offsets + self.worker.flush_buffer() + self.assertEqual([e for e in self.events if e[0] == "insert"], [("insert", 1)]) self.assertEqual(self.consumer.commits[-1], {("cdr.raw", 0): 2}) self.assertEqual(len(self.worker.buffer), 0) + def test_alternating_valid_and_invalid_records_preserves_microbatching(self): + self.worker.process_single_message(Message(payload(raw_id="r1"), offset=0)) + self.worker.process_single_message(Message(b"{bad1", offset=1)) + self.worker.process_single_message(Message(payload(raw_id="r2"), offset=2)) + self.worker.process_single_message(Message(b"{bad2", offset=3)) + self.worker.process_single_message(Message(payload(raw_id="r3"), offset=4)) + # Verify no inserts occurred yet: + self.assertEqual([e for e in self.events if e[0] == "insert"], []) + self.assertEqual(len(self.worker.buffer), 3) + # Explicit flush should insert all 3 valid records in a single batch + self.worker.flush_buffer() + self.assertEqual([e for e in self.events if e[0] == "insert"], [("insert", 3)]) + self.assertEqual(self.consumer.commits[-1], {("cdr.raw", 0): 5}) + def test_batch_commits_only_its_explicit_partition_offsets(self): self.worker.process_single_message(Message(payload(), offset=4, partition=0)) self.worker.process_single_message(Message(payload(raw_id="second"), offset=9, partition=1))