fix(worker): preserve micro-batches across DLQ records
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user