Recovery: remove the "processed_offsets" mechanism

Author: dangusev

quixstreams/app.py

@@ -6,7 +6,6 @@
 import time
 import uuid
 import warnings
-from collections import defaultdict
 from pathlib import Path
 from typing import Callable, List, Literal, Optional, Protocol, Tuple, Type, Union, cast
 
@@ -1036,18 +1035,6 @@ 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
-
             # Match the assigned TP with a stream ID via DataFrameRegistry
             for tp in non_changelog_tps:
                 stream_ids = self._dataframe_registry.get_stream_ids(
@@ -1056,9 +1043,7 @@ def _on_assign(self, _, topic_partitions: List[TopicPartition]):
                 # Assign store partitions for the given stream ids
                 for stream_id in stream_ids:
                     self._state_manager.on_partition_assign(
-                        stream_id=stream_id,
-                        partition=tp.partition,
-                        committed_offsets=committed_offsets[tp.partition],
+                        stream_id=stream_id, partition=tp.partition
                     )
         self._run_tracker.timeout_refresh()
 

quixstreams/checkpointing/checkpoint.py

@@ -228,20 +228,12 @@ def commit(self):
             partition,
             store_name,
         ), transaction in self._store_transactions.items():
-            topics = self._dataframe_registry.get_topics_for_stream_id(
-                stream_id=stream_id
-            )
-            processed_offsets = {
-                topic: offset
-                for (topic, partition_), offset in self._tp_offsets.items()
-                if topic in topics and partition_ == partition
-            }
             if transaction.failed:
                 raise StoreTransactionFailed(
                     f'Detected a failed transaction for store "{store_name}", '
                     f"the checkpoint is aborted"
                 )
-            transaction.prepare(processed_offsets=processed_offsets)
+            transaction.prepare()
 
         # Step 3. Flush producer to trigger all delivery callbacks and ensure that
         # all messages are produced

quixstreams/sources/base/manager.py

@@ -156,9 +156,7 @@ def _recover_state(self, source: StatefulSource) -> StorePartition:
         self._consumer.assign([changelog_tp])
 
         store_partitions = state_manager.on_partition_assign(
-            stream_id=None,
-            partition=source.assigned_store_partition,
-            committed_offsets={},
+            stream_id=None, partition=source.assigned_store_partition
         )
 
         if state_manager.recovery_required:

quixstreams/state/base/transaction.py

@@ -25,13 +25,11 @@
 )
 from quixstreams.state.metadata import (
     CHANGELOG_CF_MESSAGE_HEADER,
-    CHANGELOG_PROCESSED_OFFSETS_MESSAGE_HEADER,
     DEFAULT_PREFIX,
     SEPARATOR,
     Marker,
 )
 from quixstreams.state.serialization import DumpsFunc, LoadsFunc, deserialize, serialize
-from quixstreams.utils.json import dumps as json_dumps
 
 from .state import State, TransactionState
 
@@ -477,7 +475,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_offsets: Optional[dict[str, int]] = None) -> None:
+    def prepare(self) -> None:
         """
         Produce changelog messages to the changelog topic for all changes accumulated
         in this transaction and prepare transaction to flush its state to the state
@@ -488,18 +486,16 @@ def prepare(self, processed_offsets: Optional[dict[str, int]] = None) -> None:
 
         If changelog is disabled for this application, no updates will be produced
         to the changelog topic.
-
-        :param processed_offsets: the dict with <topic: offset> of the latest processed message
         """
 
         try:
-            self._prepare(processed_offsets=processed_offsets)
+            self._prepare()
             self._status = PartitionTransactionStatus.PREPARED
         except Exception:
             self._status = PartitionTransactionStatus.FAILED
             raise
 
-    def _prepare(self, processed_offsets: Optional[dict[str, int]]):
+    def _prepare(self):
         if self._changelog_producer is None:
             return
 
@@ -508,13 +504,11 @@ def _prepare(self, processed_offsets: Optional[dict[str, int]]):
             f'topic_name="{self._changelog_producer.changelog_name}" '
             f"partition={self._changelog_producer.partition}"
         )
-        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_OFFSETS_MESSAGE_HEADER: source_tp_offset_header,
             }
 
             updates = self._update_cache.get_updates(cf_name=cf_name)

quixstreams/state/manager.py

@@ -295,16 +295,13 @@ def on_partition_assign(
         self,
         stream_id: Optional[str],
         partition: int,
-        committed_offsets: dict[str, int],
     ) -> Dict[str, StorePartition]:
         """
         Assign store partitions for each registered store for the given stream_id
          and partition number, and return a list of assigned `StorePartition` objects.
 
         :param stream_id: stream id
         :param partition: Kafka topic partition number
-        :param committed_offsets: a dict with latest committed offsets
-            of all assigned topics for this partition number.
         :return: list of assigned `StorePartition`
         """
         store_partitions = {}
@@ -315,7 +312,6 @@ def on_partition_assign(
             self._recovery_manager.assign_partition(
                 topic=stream_id,
                 partition=partition,
-                committed_offsets=committed_offsets,
                 store_partitions=store_partitions,
             )
         return store_partitions

quixstreams/state/metadata.py

@@ -4,7 +4,6 @@
 SEPARATOR_LENGTH = len(SEPARATOR)
 
 CHANGELOG_CF_MESSAGE_HEADER = "__column_family__"
-CHANGELOG_PROCESSED_OFFSETS_MESSAGE_HEADER = "__processed_tp_offsets__"
 METADATA_CF_NAME = "__metadata__"
 
 DEFAULT_PREFIX = b""

quixstreams/state/recovery.py

@@ -13,17 +13,13 @@
 from quixstreams.models.types import Headers
 from quixstreams.state.base import StorePartition
 from quixstreams.utils.dicts import dict_values
-from quixstreams.utils.json import loads as json_loads
 
 from .exceptions import (
     ChangelogTopicPartitionNotAssigned,
     ColumnFamilyHeaderMissing,
     InvalidStoreChangelogOffset,
 )
-from .metadata import (
-    CHANGELOG_CF_MESSAGE_HEADER,
-    CHANGELOG_PROCESSED_OFFSETS_MESSAGE_HEADER,
-)
+from .metadata import CHANGELOG_CF_MESSAGE_HEADER
 
 logger = logging.getLogger(__name__)
 
@@ -50,7 +46,6 @@ def __init__(
         changelog_name: str,
         partition_num: int,
         store_partition: StorePartition,
-        committed_offsets: dict[str, int],
         lowwater: int,
         highwater: int,
     ):
@@ -59,7 +54,6 @@ def __init__(
         self._store_partition = store_partition
         self._changelog_lowwater = lowwater
         self._changelog_highwater = highwater
-        self._committed_offsets = committed_offsets
         self._recovery_consume_position: Optional[int] = None
         self._initial_offset: Optional[int] = None
 
@@ -154,40 +148,23 @@ def recover_from_changelog_message(
                 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_offsets = json_loads(
-            headers.get(CHANGELOG_PROCESSED_OFFSETS_MESSAGE_HEADER, b"null")
-        )
-        if processed_offsets is None or self._should_apply_changelog(
-            processed_offsets=processed_offsets
-        ):
-            key = changelog_message.key()
-            if not isinstance(key, bytes):
-                raise TypeError(
-                    f'Invalid changelog key type {type(key)}, expected "bytes"'
-                )
-
-            value = changelog_message.value()
-            if not isinstance(value, (bytes, _NoneType)):
-                raise TypeError(
-                    f'Invalid changelog value type {type(value)}, expected "bytes"'
-                )
+        key = changelog_message.key()
+        if not isinstance(key, bytes):
+            raise TypeError(f'Invalid changelog key type {type(key)}, expected "bytes"')
 
-            self._store_partition.recover_from_changelog_message(
-                cf_name=cf_name,
-                key=key,
-                value=value,
-                offset=changelog_message.offset(),
-            )
-        else:
-            # Even if the changelog update is skipped, roll the changelog offset
-            # to move forward within the changelog topic
-            self._store_partition.write_changelog_offset(
-                offset=changelog_message.offset(),
+        value = changelog_message.value()
+        if not isinstance(value, (bytes, _NoneType)):
+            raise TypeError(
+                f'Invalid changelog value type {type(value)}, expected "bytes"'
             )
 
+        self._store_partition.recover_from_changelog_message(
+            cf_name=cf_name,
+            key=key,
+            value=value,
+            offset=changelog_message.offset(),
+        )
+
     def set_recovery_consume_position(self, offset: int):
         """
         Update the recovery partition with the consumer's position (whenever
@@ -199,26 +176,6 @@ def set_recovery_consume_position(self, offset: int):
         """
         self._recovery_consume_position = offset
 
-    def _should_apply_changelog(self, processed_offsets: dict[str, int]) -> bool:
-        """
-        Determine whether the changelog update should be skipped.
-
-        :param processed_offsets: a dict with processed offsets
-            from the changelog message header processed offset.
-
-        :return: True if update should be applied, else False.
-        """
-        committed_offsets = self._committed_offsets
-        for topic, processed_offset in processed_offsets.items():
-            # Skip recovering from the message if its processed offset is ahead of the
-            # current committed offset.
-            # This is a best-effort to recover to a consistent state
-            # if the checkpointing code produced the changelog messages
-            # but failed to commit the source topic offset.
-            if processed_offset >= committed_offsets[topic]:
-                return False
-        return True
-
 
 class ChangelogProducerFactory:
     """
@@ -411,7 +368,6 @@ def _generate_recovery_partitions(
         topic_name: Optional[str],
         partition_num: int,
         store_partitions: Dict[str, StorePartition],
-        committed_offsets: dict[str, int],
     ) -> List[RecoveryPartition]:
         partitions = []
         for store_name, store_partition in store_partitions.items():
@@ -432,7 +388,6 @@ def _generate_recovery_partitions(
                     changelog_name=changelog_topic.name,
                     partition_num=partition_num,
                     store_partition=store_partition,
-                    committed_offsets=committed_offsets,
                     lowwater=lowwater,
                     highwater=highwater,
                 )
@@ -443,7 +398,6 @@ def assign_partition(
         self,
         topic: Optional[str],
         partition: int,
-        committed_offsets: dict[str, int],
         store_partitions: Dict[str, StorePartition],
     ):
         """
@@ -455,7 +409,6 @@ def assign_partition(
             topic_name=topic,
             partition_num=partition,
             store_partitions=store_partitions,
-            committed_offsets=committed_offsets,
         )
 
         assigned_tps = set(

quixstreams/state/rocksdb/timestamped.py

@@ -171,16 +171,14 @@ def set_for_timestamp(self, timestamp: int, value: Any, prefix: Any) -> None:
         self._set_min_eligible_timestamp(prefix, min_eligible_timestamp)
 
     @validate_transaction_status(PartitionTransactionStatus.STARTED)
-    def prepare(self, processed_offsets: Optional[dict[str, int]] = None) -> None:
+    def prepare(self) -> None:
         """
         This method first calls `_expire()` to remove outdated entries based on
         their timestamps and grace periods, then calls the parent class's
         `prepare()` to prepare the transaction for flush.
-
-        :param processed_offsets: the dict with <topic: offset> of the latest processed message
         """
         self._expire()
-        super().prepare(processed_offsets=processed_offsets)
+        super().prepare()
 
     def _expire(self) -> None:
         """

quixstreams/state/rocksdb/transaction.py

@@ -95,15 +95,13 @@ def _get_items(
         return sorted(merged_items.items(), key=lambda kv: kv[0], reverse=backwards)
 
     @validate_transaction_status(PartitionTransactionStatus.STARTED)
-    def prepare(self, processed_offsets: Optional[dict[str, int]] = None) -> None:
+    def prepare(self) -> None:
         """
         This method first persists the counter and then calls the parent class's
         `prepare()` to prepare the transaction for flush.
-
-        :param processed_offsets: the dict with <topic: offset> of the latest processed message
         """
         self._persist_counter()
-        super().prepare(processed_offsets=processed_offsets)
+        super().prepare()
 
     def _increment_counter(self) -> int:
         """

quixstreams/state/types.py

@@ -232,7 +232,7 @@ def prepared(self) -> bool:
         """
         ...
 
-    def prepare(self, processed_offsets: Optional[dict[str, int]]):
+    def prepare(self):
         """
         Produce changelog messages to the changelog topic for all changes accumulated
         in this transaction and prepare transcation to flush its state to the state
@@ -243,9 +243,6 @@ def prepare(self, processed_offsets: Optional[dict[str, int]]):
 
         If changelog is disabled for this application, no updates will be produced
         to the changelog topic.
-
-        :param processed_offsets: the dict with <topic: offset> of
-            the latest processed message in the current partition
         """
 
     def as_state(self, prefix: Any) -> WindowedState[K, V]: ...