Lee-W commented on code in PR #68625: URL: https://github.com/apache/airflow/pull/68625#discussion_r3459058180
########## providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py: ########## @@ -0,0 +1,327 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Shared-stream Kafka trigger and producer for event-driven scheduling. + +Triggers that declare the same ``topics`` + ``kafka_config_id`` share a +single Kafka consumer in the triggerer (one poll loop broadcast to every +subscriber) instead of opening one consumer each. +The :class:`KafkaSharedStreamProducer` owns that consumer and commits +offsets only after the derived :class:`~airflow.triggers.base.TriggerEvent` +instances have been persisted, via the shared-stream ack channel. + +Requires an Airflow version whose ``airflow.triggers.shared_stream`` module +provides the producer-side ack channel. +""" + +from __future__ import annotations + +import logging +from collections.abc import AsyncIterator, Hashable, Sequence +from typing import TYPE_CHECKING, Any, NamedTuple, cast + +from asgiref.sync import sync_to_async +from confluent_kafka import TopicPartition + +from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook +from airflow.providers.apache.kafka.hooks.produce import KafkaProducerHook +from airflow.triggers.base import BaseEventTrigger, TriggerEvent +from airflow.triggers.shared_stream import AdvanceItem, SharedStreamProducer + +if TYPE_CHECKING: + from confluent_kafka import Consumer, Producer + +log = logging.getLogger(__name__) + + +class KafkaBrokerPayload(NamedTuple): + """ + The ``broker_payload`` carried alongside each raw event from ``open_stream``. + + ``(topic, partition, offset)`` let :meth:`KafkaSharedStreamProducer.advance` + commit the right ``TopicPartition`` and + :meth:`KafkaSharedStreamProducer.get_advance_lane` key the advance by + partition. ``value`` / ``key`` are populated only when a ``dlq_topic`` is + configured -- they are what a dead-letter send re-publishes, so retaining + them through the whole outstanding window is pointless without a DLQ. + """ + + topic: str + partition: int + offset: int + value: Any = None + key: Any = None + + +class KafkaSharedStreamProducer(SharedStreamProducer): + """ + Broker-side half of a shared Kafka stream running in ack mode. + + Drives one confluent-kafka ``Consumer`` for a shared-stream group and + commits offsets only after every subscriber that derived a + ``TriggerEvent`` from a message has had it persisted -- the ack channel + gates the commit, so a triggerer crash cannot drop a message the broker + already considers delivered. + + .. warning:: + + The Kafka connection used by ``kafka_config_id`` **must** set + ``enable.auto.commit=false``. With auto-commit on, the consumer + commits offsets on its own schedule regardless of the ack channel, + which defeats the persistence gate and loses messages on a triggerer + crash. + + Events a subscriber **rejects** (via ``reject_shared_stream_event``) are + terminal: by default they are dropped (committed past). Set ``dlq_topic`` + to instead re-publish each rejected message to a dead-letter topic before + committing past it. Involuntary failures (ack timeout / overflow) and + broadcasts no subscriber was online for are never dropped -- the offset is + held so Kafka redelivers them. + + :param topics: Topics the shared consumer subscribes to. + :param kafka_config_id: Kafka connection id, defaults to ``kafka_default``. + :param poll_timeout: Seconds the consumer waits on each ``poll`` call. + :param dlq_topic: Optional dead-letter topic for rejected messages, produced + through ``kafka_config_id``. When unset, rejected messages are dropped. + """ + + def __init__( + self, + *, + topics: Sequence[str], + kafka_config_id: str = "kafka_default", + poll_timeout: float = 1.0, + dlq_topic: str | None = None, + ) -> None: + self.topics = list(topics) + self.kafka_config_id = kafka_config_id + self.poll_timeout = poll_timeout + self.dlq_topic = dlq_topic + self._consumer: Consumer | None = None + self._dlq_producer: Producer | None = None + + async def open_stream(self) -> AsyncIterator[tuple[Any, KafkaBrokerPayload]]: + """Open the consumer lazily and yield (value, KafkaBrokerPayload) per message.""" + hook = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id) + await sync_to_async(self._ensure_manual_commit)(hook) + consumer = await sync_to_async(hook.get_consumer)() + self._consumer = consumer + poll = sync_to_async(consumer.poll) + keep_for_dlq = self.dlq_topic is not None + while True: + message = await poll(self.poll_timeout) + if message is None: + continue + if message.error(): + raise RuntimeError(f"Kafka consumer error: {message.error()}") + value = message.value() + # Retain value/key only with a DLQ: the broker payload outlives the raw + # event (it is held until the offset is advanced), so keeping the message + # body when it can never be dead-lettered just wastes memory. + payload = KafkaBrokerPayload( + topic=cast("str", message.topic()), + partition=cast("int", message.partition()), + offset=cast("int", message.offset()), + value=value if keep_for_dlq else None, + key=message.key() if keep_for_dlq else None, + ) + yield value, payload + + def _ensure_manual_commit(self, hook: KafkaConsumerHook) -> None: + """ + Refuse to start unless the connection disables Kafka auto-commit. + + The shared-stream ack channel owns offset commits -- it commits only + after the derived trigger events are persisted. ``enable.auto.commit`` + defaults to ``true`` in Kafka; left on, the consumer commits on its own + schedule regardless of the ack channel and silently drops messages on a + triggerer crash. Fail fast rather than run in a lossy configuration. + """ + auto_commit = hook.get_connection(self.kafka_config_id).extra_dejson.get("enable.auto.commit", True) + # Accept only the idiomatic ``false`` -- both the string "false" and the JSON + # bool ``false`` normalise to "false"; anything else (true, 1, 0, unset, ...) is + # rejected so an ambiguous value never silently runs in a lossy configuration. + if str(auto_commit).strip().lower() != "false": + raise ValueError( + f"KafkaSharedStreamProducer requires enable.auto.commit=false in the " + f"{self.kafka_config_id!r} Kafka connection; the shared-stream ack channel " + f"commits offsets only after trigger events are persisted " + f"(got enable.auto.commit={auto_commit!r})." + ) + + def get_advance_lane(self, broker_payload: KafkaBrokerPayload) -> Hashable: + """Order commits per (topic, partition).""" + return broker_payload.topic, broker_payload.partition + + async def advance(self, batch: Sequence[AdvanceItem]) -> None: + """ + Commit the longest terminally-resolved prefix of one partition's batch. + + Every item shares one ``(topic, partition)`` (the lane). A Kafka commit + is cumulative, so committing offset ``N + 1`` marks everything up to + ``N`` as consumed. We commit only through the last item that was + *terminally handled* and stop at the first that should come back: + + * ``acked`` -- accepted; safe to commit past. + * ``rejected`` -- terminally refused. If ``dlq_topic`` is set the message + is re-published there (and flushed) before being committed past; + otherwise it is dropped. Either way it is not redelivered. + * ``failed`` -- an involuntary failure (ack timeout / overflow). Stop + before it so Kafka redelivers it and everything after; never drop. + * all-zero (a broadcast no subscriber was online for) -- nothing was + accepted, so stop before it and let Kafka redeliver when a subscriber + returns. + + The batch's rejected messages are produced to the DLQ together and + flushed once, before the offset is committed, so a crash cannot commit + past a message that never reached the DLQ. If the commit later fails the + batch is redelivered, which may re-send an already-dead-lettered message + -- DLQ consumers should tolerate duplicates. + """ + first = batch[0].broker_payload + commit_offset: int | None = None + to_dlq: list[KafkaBrokerPayload] = [] + for item in batch: + payload = item.broker_payload + outcome = item.outcome + if outcome.failed or (outcome.acked == 0 and outcome.rejected == 0): + # Involuntary failure, or nothing accepted -- redeliver from here on. + break + if outcome.rejected and self.dlq_topic is not None: + to_dlq.append(payload) + commit_offset = payload.offset + if to_dlq and self.dlq_topic is not None: + await sync_to_async(self._dead_letter)(self.dlq_topic, to_dlq) + if commit_offset is not None: + await sync_to_async(self._commit)(first.topic, first.partition, commit_offset + 1) Review Comment: I'm creating https://github.com/apache/airflow/pull/68888 to mitigate this issue. let me know what you think. thanks! -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
