Exactly Once Processing (#385)

Author: tim-quix

README.md

@@ -18,7 +18,7 @@ Quix Streams has the following benefits:
 - Support for many serialization formats, including JSON (and Quix-specific).
 - Support for stateful operations using RocksDB.
 - Support for aggregations over tumbling and hopping time windows.
-- "At-least-once" Kafka processing guarantees.
+- "At-least-once" and "exactly-once" Kafka processing guarantees.
 - Designed to run and scale resiliently via container orchestration (like Kubernetes).
 - Easily runs locally and in Jupyter Notebook for convenient development and debugging.
 - Seamless integration with the fully managed [Quix Cloud](https://quix.io/product) platform.
@@ -160,9 +160,9 @@ Here are some of the planned improvements:
 - [x] [Windowed aggregations over Tumbling & Hopping windows](https://quix.io/docs/quix-streams/v2-0-latest/windowing.html)
 - [x] [State recovery based on Kafka changelog topics](https://quix.io/docs/quix-streams/advanced/stateful-processing.html#fault-tolerance-recovery)
 - [x] [Group-by operation](https://quix.io/docs/quix-streams/groupby.html)
+- [X] ["Exactly Once" delivery guarantees for Kafka message processing (AKA transactions)](https://quix.io/docs/quix-streams/configuration.html#processing-guarantees)
 - [ ] Joins
 - [ ] Windowed aggregations over Sliding windows
-- [ ] "Exactly Once" delivery guarantees for Kafka message processing (AKA transactions)
 - [ ] Support for Avro and Protobuf formats
 - [ ] Schema Registry support
 

docs/configuration.md

@@ -35,6 +35,10 @@ See more `auto.offset.reset` in this [article](https://www.quix.io/blog/kafka-au
 **Options**: `"latest"`, `"earliest"`.  
 **Default** - `"latest"`.
 
+- **`processing_guarantee`** - Use "at-least-once" or "exactly-once" processing guarantees.  
+See [Processing Guarantees](#processing-guarantees) for more information.  
+**Options**: `"at-least-once"` or `"exactly-once".  
+**Default** - `"at-least-once"`.
 
 ## Authentication
 
@@ -83,6 +87,77 @@ app = Application(
 `ConnectionConfig.from_librdkafka_dict(config, ignore_extras=True)` will additionally 
 ignore irrelevant settings (but you will lose some validation checks).
 
+## Processing Guarantees
+
+This section concerns the `processing_guarantee` setting.
+
+### What are "processing guarantees"?
+Kafka broadly has three guarantee levels/semantics associated with handling messages.
+
+From weakest to strongest: `at-most-once`, `at-least-once`, and `exactly-once`. 
+Stronger guarantees generally have larger overhead, and thus reduced speed. 
+
+These guarantees can be read literally: when consuming Kafka messages, you can 
+guarantee each will be processed `X` times:
+
+- `at-most-once` - a message will be processed not more than once.  
+If the processing fails, the message will not be processed again.
+- `at-least-once` - a message may be processed more than once.  
+It may lead to duplicates in the output topics, but the message won’t be lost and is guaranteed to be processed.
+- `exactly-once` - a message will be processed exactly once.  
+The message is guaranteed to be processed without duplicated outputs but at the cost of higher latency and complexity.
+
+### What options does Quix Streams offer?
+
+Currently, Quix Streams offers `at-least-once` and `exactly-once`.
+
+The default guarantee is `at-least-once`. 
+
+### Performance comparison
+
+It is difficult to offer hard numbers for the speed difference between `exactly-once` 
+and `at-least-once` because many factors are at play, though latency is perhaps 
+the main contributor since `exactly-once` requires more communication with the broker.
+
+With default `Application` settings, `exactly-once` will likely perform at worst 
+10% slower than `at-least-once`, but given the nature of infrastructure/architecture, 
+even this is not a guarantee.
+
+You can minimize the impact of latency somewhat by adjusting the [`commit_interval` of 
+the `Application`](#main-configuration-parameters), but be aware of [the
+performance considerations around `commit_interval`](advanced/checkpointing.md)).
+
+
+### Exactly-Once producer buffering
+
+Because `exactly-once` messages do not fully commit or produce until the 
+[checkpoint](advanced/checkpointing.md) successfully completes, any resulting produced 
+messages are not immediately readable like with `at-least-once`.
+
+Basically, the [`commit_interval` of an `Application`](#main-configuration-parameters) 
+can be interpreted as the maximum amount of time before a produced message can be read 
+(depending on how far into the checkpoint a message is processed). If you adjust it, 
+be aware of other [performance considerations around `commit_interval`](advanced/checkpointing.md)).
+
+> NOTE: `groupby` doubles this effect since it produces to another topic under the hood. 
+
+
+### What processing guarantee is right for me?
+
+To pick the right guarantee, you need to weigh an increase in speed 
+vs. the potential to double-process a result, which may require other infrastructural
+considerations to handle appropriately.
+
+In general, you may consider using `at-least-once` guarantees when:
+- The latency and processing speed are the primary concerns.
+- The downstream consumers can gracefully handle duplicated messages.
+
+You may consider using `exactly-once` instead when:
+- Consistency and correctness of the outputs are critical.
+- Downstream consumers of the output topics cannot handle duplicated data.
+- Note: you may want to pick a smaller value for the `commit_interval` to commit checkpoints more often and reduce the latency.  
+For more information about tuning the `commit_interval`, see the ["Configuring the Checkpointing" page](advanced/checkpointing.md). 
+
 
 ## State
 - **`state_dir`** - path to the application state directory.  

quixstreams/app.py

@@ -4,7 +4,7 @@
 import os
 import signal
 import warnings
-from typing import Optional, List, Callable, Union
+from typing import Optional, List, Callable, Union, Literal, get_args
 
 from confluent_kafka import TopicPartition
 from typing_extensions import Self
@@ -44,6 +44,7 @@
 __all__ = ("Application",)
 
 logger = logging.getLogger(__name__)
+ProcessingGuarantee = Literal["at-least-once", "exactly-once"]
 MessageProcessedCallback = Callable[[str, int, int], None]
 
 # Enforce idempotent producing for the internal RowProducer
@@ -78,7 +79,7 @@ class Application:
     ```python
     from quixstreams import Application
 
-    # Set up an `app = Application` and  `sdf = StreamingDataFrame`;
+    # Set up an `app = Application` and `sdf = StreamingDataFrame`;
     # add some operations to `sdf` and then run everything.
 
     app = Application(broker_address='localhost:9092', consumer_group='group')
@@ -114,6 +115,7 @@ def __init__(
         topic_manager: Optional[TopicManager] = None,
         request_timeout: float = 30,
         topic_create_timeout: float = 60,
+        processing_guarantee: ProcessingGuarantee = "at-least-once",
     ):
         """
         :param broker_address: Connection settings for Kafka.
@@ -162,6 +164,7 @@ def __init__(
         :param topic_manager: A `TopicManager` instance
         :param request_timeout: timeout (seconds) for REST-based requests
         :param topic_create_timeout: timeout (seconds) for topic create finalization
+        :param processing_guarantee: Use "exactly-once" or "at-least-once" processing.
 
         <br><br>***Error Handlers***<br>
         To handle errors, `Application` accepts callbacks triggered when
@@ -180,6 +183,13 @@ def __init__(
             > NOTE: It is recommended to just use `quix_sdk_token` instead.
         """
         configure_logging(loglevel=loglevel)
+
+        if processing_guarantee not in get_args(ProcessingGuarantee):
+            raise ValueError(
+                f'Must provide a valid "processing_guarantee"; expected one of: '
+                f'{get_args(ProcessingGuarantee)}, got "{processing_guarantee}"'
+            )
+
         producer_extra_config = producer_extra_config or {}
         consumer_extra_config = consumer_extra_config or {}
 
@@ -248,6 +258,7 @@ def __init__(
         self._commit_interval = commit_interval
         self._producer_extra_config = producer_extra_config
         self._consumer_extra_config = consumer_extra_config
+        self._processing_guarantee = processing_guarantee
         self._consumer = RowConsumer(
             broker_address=broker_address,
             consumer_group=consumer_group,
@@ -264,6 +275,7 @@ def __init__(
                 "max.poll.interval.ms", _default_max_poll_interval_ms
             )
             / 1000,  # convert to seconds
+            transactional=self._uses_exactly_once,
         )
         self._consumer_poll_timeout = consumer_poll_timeout
         self._producer_poll_timeout = producer_poll_timeout
@@ -303,12 +315,17 @@ def __init__(
             producer=self._producer,
             consumer=self._consumer,
             state_manager=self._state_manager,
+            exactly_once=self._uses_exactly_once,
         )
 
     @property
-    def is_quix_app(self):
+    def is_quix_app(self) -> bool:
         return self._is_quix_app
 
+    @property
+    def _uses_exactly_once(self) -> bool:
+        return self._processing_guarantee == "exactly-once"
+
     @classmethod
     def Quix(
         cls,
@@ -331,6 +348,7 @@ def Quix(
         topic_manager: Optional[QuixTopicManager] = None,
         request_timeout: float = 30,
         topic_create_timeout: float = 60,
+        processing_guarantee: Literal["at-least-once", "exactly-once"] = "exactly-once",
     ) -> Self:
         """
         >***NOTE:*** DEPRECATED: use Application with `quix_sdk_token` argument instead.
@@ -398,6 +416,7 @@ def Quix(
         :param topic_manager: A `QuixTopicManager` instance
         :param request_timeout: timeout (seconds) for REST-based requests
         :param topic_create_timeout: timeout (seconds) for topic create finalization
+        :param processing_guarantee: Use "exactly-once" or "at-least-once" processing.
 
         <br><br>***Error Handlers***<br>
         To handle errors, `Application` accepts callbacks triggered when
@@ -445,6 +464,7 @@ def Quix(
             request_timeout=request_timeout,
             topic_create_timeout=topic_create_timeout,
             quix_config_builder=quix_config_builder,
+            processing_guarantee=processing_guarantee,
         )
         return app
 
@@ -625,14 +645,18 @@ def get_consumer(self, auto_commit_enable: bool = True) -> Consumer:
         Create and return a pre-configured Consumer instance.
         The Consumer is initialized with params passed to Application.
 
-        It's useful for consuming data from Kafka outside the standard Application processing flow.
-        (e.g. to consume test data from a topic).
-        Using it within the StreamingDataFrame functions is not recommended, as it creates a new Consumer instance
+        It's useful for consuming data from Kafka outside the standard
+        Application processing flow.
+        (e.g., to consume test data from a topic).
+        Using it within the StreamingDataFrame functions is not recommended, as it
+        creates a new Consumer instance
         each time, which is not optimized for repeated use in a streaming pipeline.
 
-        Note: By default this consumer does not autocommit consumed offsets to allow exactly-once processing.
+        Note: By default, this consumer does not autocommit the consumed offsets to allow
+        at-least-once processing.
         To store the offset call store_offsets() after processing a message.
-        If autocommit is necessary set `enable.auto.offset.store` to True in the consumer config when creating the app.
+        If autocommit is necessary set `enable.auto.offset.store` to True in
+        the consumer config when creating the app.
 
         Example Snippet:
 
@@ -705,14 +729,16 @@ def run(
             f'broker_address="{self._broker_address}" '
             f'consumer_group="{self._consumer_group}" '
             f'auto_offset_reset="{self._auto_offset_reset}" '
-            f"commit_interval={self._commit_interval}s"
+            f"commit_interval={self._commit_interval}s "
+            f'processing_guarantee="{self._processing_guarantee}"'
         )
         if self.is_quix_app:
             self._quix_runtime_init()
 
         self._setup_topics()
 
         exit_stack = contextlib.ExitStack()
+        exit_stack.enter_context(self._processing_context)
         exit_stack.enter_context(self._state_manager)
         exit_stack.enter_context(self._consumer)
         exit_stack.push(

quixstreams/checkpointing/__init__.py

@@ -1,2 +1,2 @@
-from .checkpoint import Checkpoint as Checkpoint
-from .exceptions import InvalidStoredOffset as InvalidStoredOffset
+from .checkpoint import Checkpoint
+from .exceptions import InvalidStoredOffset

quixstreams/checkpointing/checkpoint.py

@@ -33,6 +33,7 @@ def __init__(
         producer: RowProducer,
         consumer: Consumer,
         state_manager: StateStoreManager,
+        exactly_once: bool = False,
     ):
         self._created_at = time.monotonic()
         # A mapping of <(topic, partition): processed offset>
@@ -45,6 +46,10 @@ def __init__(
         self._state_manager = state_manager
         self._consumer = consumer
         self._producer = producer
+        self._exactly_once = exactly_once
+
+        if self._exactly_once:
+            self._producer.begin_transaction()
 
     def expired(self) -> bool:
         """
@@ -102,6 +107,15 @@ def get_store_transaction(
         self._store_transactions[(topic, partition, store_name)] = transaction
         return transaction
 
+    def close(self):
+        """
+        Perform cleanup (when the checkpoint is empty) instead of committing.
+
+        Needed for exactly-once, as Kafka transactions are timeboxed.
+        """
+        if self._exactly_once:
+            self._producer.abort_transaction()
+
     def commit(self):
         """
         Commit the checkpoint.
@@ -113,10 +127,6 @@ def commit(self):
          4. Flush each state store partition to the disk.
         """
 
-        if not self._tp_offsets:
-            # No messages have been processed during this checkpoint, return
-            return
-
         # Step 1. Produce the changelogs
         for (
             topic,
@@ -148,15 +158,20 @@ def commit(self):
             TopicPartition(topic=topic, partition=partition, offset=offset + 1)
             for (topic, partition), offset in self._tp_offsets.items()
         ]
-        logger.debug("Checkpoint: commiting consumer")
-        try:
-            partitions = self._consumer.commit(offsets=offsets, asynchronous=False)
-        except KafkaException as e:
-            raise CheckpointConsumerCommitError(e.args[0]) from None
+        if self._exactly_once:
+            self._producer.commit_transaction(
+                offsets, self._consumer.consumer_group_metadata()
+            )
+        else:
+            logger.debug("Checkpoint: committing consumer")
+            try:
+                partitions = self._consumer.commit(offsets=offsets, asynchronous=False)
+            except KafkaException as e:
+                raise CheckpointConsumerCommitError(e.args[0]) from None
 
-        for partition in partitions:
-            if partition.error:
-                raise CheckpointConsumerCommitError(partition.error)
+            for partition in partitions:
+                if partition.error:
+                    raise CheckpointConsumerCommitError(partition.error)
 
         # Step 4. Flush state store partitions to the disk together with changelog
         # offsets
@@ -175,10 +190,6 @@ def commit(self):
             changelog_offset = (
                 produced_offsets.get(changelog_tp) if changelog_tp is not None else None
             )
-            if changelog_offset is not None:
-                # Increment the changelog offset by one to match the high watermark
-                # in Kafka
-                changelog_offset += 1
             transaction.flush(
                 processed_offset=offset, changelog_offset=changelog_offset
             )

quixstreams/kafka/consumer.py

@@ -9,7 +9,7 @@
     Consumer as ConfluentConsumer,
     KafkaError,
 )
-from confluent_kafka.admin import ClusterMetadata
+from confluent_kafka.admin import ClusterMetadata, GroupMetadata
 
 from .configuration import ConnectionConfig
 from quixstreams.exceptions import PartitionAssignmentError, KafkaPartitionError
@@ -548,6 +548,12 @@ def close(self):
         self._consumer.close()
         logger.debug("Kafka consumer closed")
 
+    def consumer_group_metadata(self) -> GroupMetadata:
+        """
+        Used by the producer during consumer offset sending for an EOS transaction.
+        """
+        return self._consumer.consumer_group_metadata()
+
     @property
     def _consumer(self) -> ConfluentConsumer:
         """