Refactor recovery to support Stores belonging to multiple topics (#774)

Author: daniil-quix

quixstreams/app.py

@@ -5,6 +5,7 @@
 import signal
 import time
 import warnings
+from collections import defaultdict
 from pathlib import Path
 from typing import Callable, List, Literal, Optional, Protocol, Tuple, Type, Union
 
@@ -923,40 +924,26 @@ def _on_assign(self, _, topic_partitions: List[TopicPartition]):
             non_changelog_tps = [
                 tp for tp in topic_partitions if tp.topic in non_changelog_topics
             ]
+            committed_tps = self._consumer.committed(
+                partitions=non_changelog_tps, timeout=30
+            )
+            committed_offsets: dict[int, dict[str, int]] = defaultdict(dict)
+            for tp in committed_tps:
+                if tp.error:
+                    raise RuntimeError(
+                        f"Failed to get committed offsets for "
+                        f'"{tp.topic}[{tp.partition}]" from the broker: {tp.error}'
+                    )
+                committed_offsets[tp.partition][tp.topic] = tp.offset
+
             for tp in non_changelog_tps:
-                # Get the latest committed offset for the assigned topic partition
-                tp_committed = self._consumer.committed([tp], timeout=30)[0]
                 # Assign store partitions
-                store_partitions = self._state_manager.on_partition_assign(
+                self._state_manager.on_partition_assign(
                     topic=tp.topic,
                     partition=tp.partition,
-                    committed_offset=tp_committed.offset,
+                    committed_offsets=committed_offsets[tp.partition],
                 )
 
-                # Check if the latest committed offset >= stored offset
-                # Otherwise, the re-processed messages might use already updated
-                # state, which can lead to inconsistent outputs
-                stored_offsets = [
-                    offset
-                    for offset in (
-                        store_tp.get_processed_offset()
-                        for store_tp in store_partitions.values()
-                    )
-                    if offset is not None
-                ]
-                min_stored_offset = min(stored_offsets) + 1 if stored_offsets else None
-                if (
-                    min_stored_offset is not None
-                    and min_stored_offset > tp_committed.offset
-                ):
-                    logger.warning(
-                        f'Warning: offset "{tp_committed.offset}" '
-                        f"for topic partition "
-                        f'"{tp_committed.topic}[{tp_committed.partition}]" '
-                        f'is behind the stored offset "{min_stored_offset}". '
-                        f"It may lead to distortions in produced data."
-                    )
-
     def _on_revoke(self, _, topic_partitions: List[TopicPartition]):
         """
         Revoke partitions from consumer and state

quixstreams/checkpointing/checkpoint.py

@@ -199,7 +199,7 @@ def commit(self):
                     f'Detected a failed transaction for store "{store_name}", '
                     f"the checkpoint is aborted"
                 )
-            transaction.prepare(processed_offset=offset)
+            transaction.prepare(processed_offsets={topic: offset})
 
         # Step 2. Flush producer to trigger all delivery callbacks and ensure that
         # all messages are produced
@@ -299,6 +299,4 @@ def commit(self):
             changelog_offset = (
                 produced_offsets.get(changelog_tp) if changelog_tp is not None else None
             )
-            transaction.flush(
-                processed_offset=offset, changelog_offset=changelog_offset
-            )
+            transaction.flush(changelog_offset=changelog_offset)

quixstreams/sources/base/manager.py

@@ -160,7 +160,7 @@ def _recover_state(self, source: StatefulSource) -> StorePartition:
         store_partitions = state_manager.on_partition_assign(
             topic=None,
             partition=source.assigned_store_partition,
-            committed_offset=OFFSET_BEGINNING,
+            committed_offsets={},
         )
 
         if state_manager.recovery_required:

quixstreams/state/base/partition.py

@@ -2,15 +2,10 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Literal, Optional, Union
 
-from quixstreams.models import SuccessfulConfluentKafkaMessageProto
-from quixstreams.state.exceptions import ColumnFamilyHeaderMissing
 from quixstreams.state.metadata import (
-    CHANGELOG_CF_MESSAGE_HEADER,
-    CHANGELOG_PROCESSED_OFFSET_MESSAGE_HEADER,
     Marker,
 )
 from quixstreams.state.serialization import DumpsFunc, LoadsFunc
-from quixstreams.utils.json import loads as json_loads
 
 from .transaction import PartitionTransaction, PartitionTransactionCache
 
@@ -45,42 +40,34 @@ def __init__(
     def close(self): ...
 
     @abstractmethod
-    def _recover_from_changelog_message(
-        self,
-        changelog_message: SuccessfulConfluentKafkaMessageProto,
-        cf_name: str,
-        processed_offset: Optional[int],
-        committed_offset: int,
-    ): ...
-
-    @abstractmethod
-    def get_processed_offset(self) -> Optional[int]:
+    def get_changelog_offset(self) -> Optional[int]:
         """
-        Get last processed offset for the given partition
+        Get the changelog offset that the state is up-to-date with.
         :return: offset or `None` if there's no processed offset yet
         """
         ...
 
     @abstractmethod
-    def get_changelog_offset(self) -> Optional[int]:
+    def write_changelog_offset(self, offset: int):
         """
-        Get offset that the changelog is up-to-date with.
-        :return: offset or `None` if there's no processed offset yet
+        Write a new changelog offset to the db.
+
+        To be used when we simply need to update the changelog offset without touching
+        the actual data.
+
+        :param offset: new changelog offset
         """
-        ...
 
     @abstractmethod
     def write(
         self,
         cache: PartitionTransactionCache,
-        processed_offset: Optional[int],
         changelog_offset: Optional[int],
     ):
         """
         Update the state with data from the update cache
 
         :param cache: The modified data
-        :param processed_offset: The offset processed to generate the data.
         :param changelog_offset: The changelog message offset of the data.
         """
 
@@ -92,7 +79,6 @@ def get(
         Get a key from the store
 
         :param key: a key encoded to `bytes`
-        :param default: a default value to return if the key is not found.
         :param cf_name: rocksdb column family name. Default - "default"
         :return: a value if the key is present in the store. Otherwise, `default`
         """
@@ -107,6 +93,19 @@ def exists(self, key: bytes, cf_name: str = "default") -> bool:
         :return: `True` if the key is present, `False` otherwise.
         """
 
+    @abstractmethod
+    def recover_from_changelog_message(
+        self, key: bytes, value: Optional[bytes], cf_name: str, offset: int
+    ):
+        """
+        Updates state from a given changelog message.
+
+        :param key: changelog message key
+        :param value: changelog message value
+        :param cf_name: column family name
+        :param offset: changelog message offset
+        """
+
     def begin(self) -> PartitionTransaction:
         """
         Start a new `PartitionTransaction`
@@ -120,58 +119,6 @@ def begin(self) -> PartitionTransaction:
             changelog_producer=self._changelog_producer,
         )
 
-    def recover_from_changelog_message(
-        self,
-        changelog_message: SuccessfulConfluentKafkaMessageProto,
-        committed_offset: int,
-    ) -> None:
-        """
-        Updates state from a given changelog message.
-
-        :param changelog_message: A raw Confluent message read from a changelog topic.
-        :param committed_offset: latest committed offset for the partition
-        """
-        headers = dict(changelog_message.headers() or ())
-        # Parse the column family name from message headers
-        cf_name = headers.get(CHANGELOG_CF_MESSAGE_HEADER, b"").decode()
-        if not cf_name:
-            raise ColumnFamilyHeaderMissing(
-                f"Header '{CHANGELOG_CF_MESSAGE_HEADER}' missing from changelog message"
-            )
-
-        # Parse the processed topic-partition-offset info from the changelog message
-        # headers to determine whether the update should be applied or skipped.
-        # It can be empty if the message was produced by the older version of the lib.
-        processed_offset = json_loads(
-            headers.get(CHANGELOG_PROCESSED_OFFSET_MESSAGE_HEADER, b"null")
-        )
-
-        self._recover_from_changelog_message(
-            changelog_message,
-            cf_name,
-            processed_offset,
-            committed_offset,
-        )
-
-    def _should_apply_changelog(
-        self, processed_offset: Optional[int], committed_offset: int
-    ) -> bool:
-        """
-        Determine whether the changelog update should be skipped.
-
-        :param processed_offset: changelog message processed offset.
-        :param committed_offset: latest committed offset of the source topic partition
-        :return: True if update should be applied, else False.
-        """
-        if processed_offset is not None:
-            # Skip recovering from the message if its processed offset is ahead of the
-            # current committed offset.
-            # This way it will recover to a consistent state if the checkpointing code
-            # produced the changelog messages but failed to commit
-            # the source topic offset.
-            return processed_offset < committed_offset
-        return True
-
     def __enter__(self):
         return self
 

quixstreams/state/base/transaction.py

@@ -20,7 +20,7 @@
 from quixstreams.state.exceptions import InvalidChangelogOffset, StateTransactionError
 from quixstreams.state.metadata import (
     CHANGELOG_CF_MESSAGE_HEADER,
-    CHANGELOG_PROCESSED_OFFSET_MESSAGE_HEADER,
+    CHANGELOG_PROCESSED_OFFSETS_MESSAGE_HEADER,
     DEFAULT_PREFIX,
     SEPARATOR,
     Marker,
@@ -398,7 +398,7 @@ def exists(self, key: K, prefix: bytes, cf_name: str = "default") -> bool:
             return self._partition.exists(key_serialized, cf_name=cf_name)
 
     @validate_transaction_status(PartitionTransactionStatus.STARTED)
-    def prepare(self, processed_offset: Optional[int]):
+    def prepare(self, processed_offsets: Optional[dict[str, int]]):
         """
         Produce changelog messages to the changelog topic for all changes accumulated
         in this transaction and prepare transaction to flush its state to the state
@@ -410,33 +410,32 @@ def prepare(self, processed_offset: Optional[int]):
         If changelog is disabled for this application, no updates will be produced
         to the changelog topic.
 
-        :param processed_offset: the offset of the latest processed message
+        :param processed_offsets: the dict with <topic: offset> of the latest processed message
         """
 
         try:
-            self._prepare(processed_offset=processed_offset)
+            self._prepare(processed_offsets=processed_offsets)
             self._status = PartitionTransactionStatus.PREPARED
         except Exception:
             self._status = PartitionTransactionStatus.FAILED
             raise
 
-    def _prepare(self, processed_offset: Optional[int]):
+    def _prepare(self, processed_offsets: Optional[dict[str, int]]):
         if self._changelog_producer is None:
             return
 
         logger.debug(
             f"Flushing state changes to the changelog topic "
             f'topic_name="{self._changelog_producer.changelog_name}" '
-            f"partition={self._changelog_producer.partition} "
-            f"processed_offset={processed_offset}"
+            f"partition={self._changelog_producer.partition}"
         )
-        source_tp_offset_header = json_dumps(processed_offset)
+        source_tp_offset_header = json_dumps(processed_offsets)
         column_families = self._update_cache.get_column_families()
 
         for cf_name in column_families:
             headers: Headers = {
                 CHANGELOG_CF_MESSAGE_HEADER: cf_name,
-                CHANGELOG_PROCESSED_OFFSET_MESSAGE_HEADER: source_tp_offset_header,
+                CHANGELOG_PROCESSED_OFFSETS_MESSAGE_HEADER: source_tp_offset_header,
             }
 
             updates = self._update_cache.get_updates(cf_name=cf_name)
@@ -461,7 +460,6 @@ def _prepare(self, processed_offset: Optional[int]):
     )
     def flush(
         self,
-        processed_offset: Optional[int] = None,
         changelog_offset: Optional[int] = None,
     ):
         """
@@ -476,18 +474,17 @@ def flush(
             not flush ANY data to the database including the offset to optimize
             I/O.
 
-        :param processed_offset: offset of the last processed message, optional.
         :param changelog_offset: offset of the last produced changelog message,
             optional.
         """
         try:
-            self._flush(processed_offset, changelog_offset)
+            self._flush(changelog_offset)
             self._status = PartitionTransactionStatus.COMPLETE
         except Exception:
             self._status = PartitionTransactionStatus.FAILED
             raise
 
-    def _flush(self, processed_offset: Optional[int], changelog_offset: Optional[int]):
+    def _flush(self, changelog_offset: Optional[int]):
         if self._update_cache.is_empty():
             return
 
@@ -503,7 +500,6 @@ def _flush(self, processed_offset: Optional[int], changelog_offset: Optional[int
 
         self._partition.write(
             cache=self._update_cache,
-            processed_offset=processed_offset,
             changelog_offset=changelog_offset,
         )
 

quixstreams/state/manager.py

@@ -248,15 +248,19 @@ def clear_stores(self) -> None:
             shutil.rmtree(self._state_dir, ignore_errors=True)
 
     def on_partition_assign(
-        self, topic: Optional[str], partition: int, committed_offset: int
+        self,
+        topic: Optional[str],
+        partition: int,
+        committed_offsets: dict[str, int],
     ) -> Dict[str, StorePartition]:
         """
         Assign store partitions for each registered store for the given `TopicPartition`
         and return a list of assigned `StorePartition` objects.
 
         :param topic: Kafka topic name
         :param partition: Kafka topic partition
-        :param committed_offset: latest committed offset for the partition
+        :param committed_offsets: a dict with latest committed offsets
+            of all assigned topics for this partition number.
         :return: list of assigned `StorePartition`
         """
 
@@ -268,7 +272,7 @@ def on_partition_assign(
             self._recovery_manager.assign_partition(
                 topic=topic,
                 partition=partition,
-                committed_offset=committed_offset,
+                committed_offsets=committed_offsets,
                 store_partitions=store_partitions,
             )
         return store_partitions