feat(ISV-7020): add chunking capabilities

* chunking enable only in retry topics
* tests assisted with Claude

Signed-off-by: Marek Szymutko <mszymutk@redhat.com>
Assisted-by: Claude-4.6-Opus

Author: BorekZnovustvoritel

src/retriable_kafka_client/chunking.py

@@ -0,0 +1,96 @@
+"""Module for tracking chunked messages"""
+
+import logging
+import sys
+import uuid
+from collections import defaultdict
+
+from confluent_kafka import Message
+
+from retriable_kafka_client.headers import (
+    CHUNK_GROUP_HEADER,
+    NUMBER_OF_CHUNKS_HEADER,
+    CHUNK_ID_HEADER,
+    deserialize_number_from_bytes,
+    get_header_value,
+)
+from retriable_kafka_client.kafka_utils import MessageGroup
+
+LOGGER = logging.getLogger(__name__)
+
+
+def generate_group_id() -> bytes:
+    """Generate a random group id."""
+    return uuid.uuid4().bytes
+
+
+def calculate_header_size(
+    headers: dict[str, str | bytes] | list[tuple[str, str | bytes]] | None,
+) -> int:
+    """Approximate the space needed for headers within a message."""
+    if not headers:
+        return 0
+    result = 0
+    if not isinstance(headers, dict):
+        headers = dict(headers)
+    for header_key, header_value in headers.items():
+        result += len(header_key) + sys.getsizeof(header_value) + 8
+        # Two 32-bit numbers specifying lengths of fields are also
+        # present, therefore adding 8
+    return result
+
+
+class ChunkingCache:
+    # pylint: disable=too-few-public-methods
+    """Class for storing information about received message fragments."""
+
+    def __init__(self):
+        self._message_chunks: dict[tuple[bytes, str, int], dict[int, Message]] = (
+            defaultdict(dict)
+        )
+
+    def receive(self, message: Message) -> MessageGroup | None:
+        """
+        Receive a message. If the message is whole, or it is
+        the last fragment, returns the whole message group
+        and flushes cache of this group. Otherwise, returns None
+        and the message fragment shall not be processed.
+        """
+        topic: str = message.topic()  # type: ignore[assignment]
+        partition: int = message.partition()  # type: ignore[assignment]
+        if (
+            (group_id := get_header_value(message, CHUNK_GROUP_HEADER)) is not None
+            and (
+                number_of_chunks_raw := get_header_value(
+                    message, NUMBER_OF_CHUNKS_HEADER
+                )
+            )
+            is not None
+            and (chunk_id_raw := get_header_value(message, CHUNK_ID_HEADER)) is not None
+        ):
+            number_of_chunks = deserialize_number_from_bytes(number_of_chunks_raw)
+            chunk_id = deserialize_number_from_bytes(chunk_id_raw)
+            identifier = (group_id, topic, partition)
+            stored_message_ids_from_group = set(
+                self._message_chunks.get(identifier, {}).keys()
+            )
+            stored_message_ids_from_group.add(chunk_id)
+            if not all(
+                i in stored_message_ids_from_group for i in range(number_of_chunks)
+            ):
+                LOGGER.debug(
+                    "Received a message chunk, waiting for the other chunks..."
+                )
+                self._message_chunks[identifier][chunk_id] = message
+                return None
+            LOGGER.debug(
+                "Received all message chunks, assembling group %s composed of %s messages.",
+                group_id.hex(),
+                number_of_chunks,
+            )
+            # Clear cache and reassemble, cache can be empty if
+            # this is just one-message-sized value
+            messages = self._message_chunks.pop(identifier, {})
+            messages[chunk_id] = message
+            return MessageGroup(topic, partition, messages, group_id)
+        return MessageGroup(topic, partition, {0: message}, None)

src/retriable_kafka_client/config.py

@@ -42,6 +42,8 @@ class ProducerConfig(CommonConfig):
     retries: int = field(default=3)
     fallback_factor: float = field(default=2.0)
     fallback_base: float = field(default=5.0)
+    max_message_size: int = field(default=1000000)
+    split_messages: bool = field(default=False)
 
 
 @dataclass

src/retriable_kafka_client/consumer.py

@@ -1,6 +1,5 @@
 """Base Kafka Consumer module"""
 
-import json
 import logging
 import sys
 from concurrent.futures import Executor, Future
@@ -9,8 +8,9 @@
 
 from confluent_kafka import Consumer, Message, KafkaException, TopicPartition
 
+from .chunking import ChunkingCache
 from .health import perform_healthcheck_using_client
-from .kafka_utils import message_to_partition
+from .kafka_utils import message_to_partition, MessageGroup
 from .kafka_settings import KafkaOptions, DEFAULT_CONSUMER_SETTINGS
 from .consumer_tracking import TrackingManager
 from .config import ConsumerConfig
@@ -91,6 +91,8 @@ def __init__(
         self.__retry_manager = RetryManager(config)
         # Store information about pending retried messages
         self.__schedule_cache = RetryScheduleCache()
+        # Store chunking information
+        self.__chunk_tracker = ChunkingCache()
 
     @property
     def _consumer(self) -> Consumer:
@@ -140,7 +142,7 @@ def __on_revoke(self, _: Consumer, partitions: list[TopicPartition]) -> None:
         self.__tracking_manager.register_revoke(partitions)
         self.__perform_commits()
 
-    def __ack_message(self, message: Message, finished_future: Future) -> None:
+    def __ack_message(self, message: MessageGroup, finished_future: Future) -> None:
         """
         Private method only ever intended to be used from within
         _process_message(). It commits offsets and releases
@@ -155,7 +157,7 @@ def __ack_message(self, message: Message, finished_future: Future) -> None:
             if problem := finished_future.exception():
                 LOGGER.error(
                     "Message could not be processed! Message: %s.",
-                    message.value(),
+                    message.deserialize(),
                     exc_info=problem,
                 )
                 self.__retry_manager.resend_message(message)
@@ -212,20 +214,17 @@ def _process_message(self, message: Message) -> Future[Any] | None:
         Returns: Future of the target execution if the message can be processed.
             None otherwise.
         """
-        message_value = message.value()
-        if not message_value:
-            # Discard empty messages
+        message_group = self.__chunk_tracker.receive(message)
+        if not message_group:
             return None
-        try:
-            message_data = json.loads(message_value)
-        except json.decoder.JSONDecodeError:
-            # This message cannot be deserialized, just log and discard it
-            LOGGER.exception("Decoding error: not a valid JSON: %s", message.value())
+        message_data = message_group.deserialize()
+        if not message_data:
+            self.__tracking_manager.schedule_commit(message_group)
             return None
         future = self._executor.submit(self._config.target, message_data)
-        self.__tracking_manager.process_message(message, future)
+        self.__tracking_manager.process_message(message_group, future)
         # The semaphore is released within this callback
-        future.add_done_callback(lambda res: self.__ack_message(message, res))
+        future.add_done_callback(lambda res: self.__ack_message(message_group, res))
         return future
 
     ### Public methods ###

src/retriable_kafka_client/consumer_tracking.py

@@ -3,51 +3,19 @@
 import logging
 from collections import defaultdict
 from concurrent.futures import Future
+from itertools import chain
 from threading import Lock, Semaphore
-from typing import NamedTuple, Any
+from typing import Any, Iterable
 
-from confluent_kafka import Message, TopicPartition
+from confluent_kafka import TopicPartition
+
+from retriable_kafka_client.kafka_utils import TrackingInfo, MessageGroup
 
 LOGGER = logging.getLogger(__name__)
 
 
-class _PartitionInfo(NamedTuple):
-    """
-    Consistently hashable dataclass for storing information about a partition,
-    namely offset information. Can be used as keys in a dictionary.
-    """
-
-    topic: str
-    partition: int
-
-    @staticmethod
-    def from_message(message: Message) -> "_PartitionInfo":
-        """
-        Create a PartitionInfo from a Kafka message.
-        Args:
-             message: Kafka message object
-        Returns: hashable info about a partition
-        """
-        message_topic = message.topic()
-        message_partition = message.partition()
-        # This should never happen with polled messages. Polled messages need
-        # the information asserted to be valid Kafka messages. This can
-        # happen only for custom-created messages objects, which this
-        # method is not intended to be used for
-        assert message_topic is not None and message_partition is not None, (
-            "Invalid message cannot be converted to partition info"
-        )
-        return _PartitionInfo(message_topic, message_partition)
-
-    def to_offset_info(self, offset: int) -> TopicPartition:
-        """
-        Create a Kafka-committable object using the provided offset.
-        Args:
-            offset: The offset to be committed. Make sure to commit
-                offset one higher than the latest processed message.
-        Returns: The committable Kafka object
-        """
-        return TopicPartition(topic=self.topic, partition=self.partition, offset=offset)
+def _flatten_offsets(done_offsets: Iterable[tuple[int, ...]]) -> list[int]:
+    return list(chain(*done_offsets))
 
 
 class TrackingManager:
@@ -85,8 +53,10 @@ class TrackingManager:
     """
 
     def __init__(self, concurrency: int, cancel_wait_time: float):
-        self.__to_process: dict[_PartitionInfo, dict[int, Future]] = defaultdict(dict)
-        self.__to_commit: dict[_PartitionInfo, set[int]] = defaultdict(set)
+        self.__to_process: dict[TrackingInfo, dict[tuple[int, ...], Future]] = (
+            defaultdict(dict)
+        )
+        self.__to_commit: dict[TrackingInfo, set[tuple[int, ...]]] = defaultdict(set)
         self.__access_lock = Lock()  # For handling multithreaded access to this object
         self.__semaphore = Semaphore(concurrency)
         self.__cancel_wait_time = cancel_wait_time
@@ -105,15 +75,15 @@ def pop_committable(self) -> list[TopicPartition]:
         """
         to_commit = []
         with self.__access_lock:
-            for partition_info, pending_to_commit in self.__to_commit.items():
-                if not pending_to_commit:
+            for partition_info, tuples_pending_to_commit in self.__to_commit.items():
+                if not tuples_pending_to_commit:
                     # Nothing to commit
                     continue
 
-                pending_to_process = self.__to_process.get(partition_info, None)
-                if not pending_to_process:
+                tuples_pending_to_process = self.__to_process.get(partition_info, None)
+                if not tuples_pending_to_process:
                     # Nothing is blocking the committing
-                    max_to_commit = max(pending_to_commit)
+                    max_to_commit = max(_flatten_offsets(tuples_pending_to_commit))
                     to_commit.append(
                         TopicPartition(
                             topic=partition_info.topic,
@@ -123,17 +93,18 @@ def pop_committable(self) -> list[TopicPartition]:
                     )
                     self.__to_commit[partition_info] = set()
                     continue
-
-                min_pending_to_process = min(pending_to_process)
+                min_pending_to_process = min(
+                    _flatten_offsets(tuples_pending_to_process)
+                )
                 commit_candidates = {
-                    offset
-                    for offset in pending_to_commit
-                    if offset < min_pending_to_process
+                    offset_tuple
+                    for offset_tuple in tuples_pending_to_commit
+                    if all(offset < min_pending_to_process for offset in offset_tuple)
                 }
                 if not commit_candidates:
                     # Nothing to commit
                     continue
-                max_to_commit = max(commit_candidates)
+                max_to_commit = max(_flatten_offsets(commit_candidates))
                 to_commit.append(
                     TopicPartition(
                         topic=partition_info.topic,
@@ -157,11 +128,11 @@ def reschedule_uncommittable(
             failed_committable: list of data that failed to be committed
         """
         for failed in failed_committable:
-            self.__to_commit.setdefault(
-                _PartitionInfo(topic=failed.topic, partition=failed.partition), set()
-            ).add(failed.offset)
+            self.__to_commit[
+                TrackingInfo(topic=failed.topic, partition=failed.partition)
+            ].add((failed.offset,))
 
-    def process_message(self, message: Message, future: Future[Any]) -> None:
+    def process_message(self, message: MessageGroup, future: Future[Any]) -> None:
         """
         Mark message as pending for processing.
         Args:
@@ -170,15 +141,16 @@ def process_message(self, message: Message, future: Future[Any]) -> None:
         """
         # We cannot really use context manager, the semaphore is released in
         # future's callback or when the future is cancelled
+
         self.__semaphore.acquire()  # pylint: disable=consider-using-with
-        message_offset: int = message.offset()  # type: ignore[assignment]
+        message_offsets = message.offsets
         with self.__access_lock:
             # Mark the message as being processed
-            self.__to_process[_PartitionInfo.from_message(message)][
-                message_offset + 1
+            self.__to_process[TrackingInfo.from_message_group(message)][
+                tuple(message_offset + 1 for message_offset in message_offsets)
             ] = future
 
-    def schedule_commit(self, message: Message) -> bool:
+    def schedule_commit(self, message: MessageGroup) -> bool:
         """
         Mark message as pending for committing when its processing is fully done.
         Args:
@@ -188,12 +160,12 @@ def schedule_commit(self, message: Message) -> bool:
             as pending for processing), False otherwise
         """
         self.__semaphore.release()
-        partition_info = _PartitionInfo.from_message(message)
-        message_offset: int = message.offset()  # type: ignore[assignment]
-        stored_offset = message_offset + 1
+        partition_info = TrackingInfo.from_message_group(message)
+        message_offsets = message.offsets
+        stored_offsets = tuple(message_offset + 1 for message_offset in message_offsets)
         with self.__access_lock:
-            self.__to_process[partition_info].pop(stored_offset, None)
-            self.__to_commit.setdefault(partition_info, set()).add(stored_offset)
+            self.__to_process[partition_info].pop(stored_offsets, None)
+            self.__to_commit.setdefault(partition_info, set()).add(stored_offsets)
         self._cleanup()
         return True
 
@@ -212,7 +184,7 @@ def _cleanup(self) -> None:
                     cache_to_clean.pop(key, None)
 
     def _revoke_processing(
-        self, revoked_partitions: set[_PartitionInfo]
+        self, revoked_partitions: set[TrackingInfo]
     ) -> list[Future[Any]]:
         """
         Cancel all pending tracked futures related to the given partitions.
@@ -253,7 +225,7 @@ def register_revoke(self, partitions: list[TopicPartition] | None = None) -> Non
             revoked_partition_keys = set(self.__to_process.keys())
         else:
             revoked_partition_keys = {
-                _PartitionInfo(partition=partition.partition, topic=partition.topic)
+                TrackingInfo(partition=partition.partition, topic=partition.topic)
                 for partition in partitions
             }
         pending_futures = self._revoke_processing(revoked_partition_keys)