You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "FrankYang0529 (via GitHub)" <gi...@apache.org> on 2026/06/16 12:01:00 UTC
[PR] feat(providers/apache/kafka): Add KafkaSharedStreamProducer [airflow]
FrankYang0529 opened a new pull request, #68625:
URL: https://github.com/apache/airflow/pull/68625
Adds `KafkaSharedStreamProducer` and `KafkaSharedStreamTrigger` to the Apache Kafka provider, backing the producer-side ack channel for shared-stream triggers (#67523). Many event-driven triggers / asset watchers on the same Kafka topics can now share a single consumer in the trigger, with offset commits gated on trigger-event persistence (at-least-once).
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
For user-facing UI changes, please attach before/after screenshots (or a short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
related: #67523
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [X] Yes (Claude Code with Opus 4.8)
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3449656245
##########
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})."
+ )
Review Comment:
There's an async method `aget_connection` at base hook level. IIUC, we could avoid the extra `sync_to_async` call for the early validation.
https://github.com/apache/airflow/blob/db9fa4289fe0d01e764b377c5c7f13d247af8218/task-sdk/src/airflow/sdk/bases/hook.py#L65-L77
##########
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:
This was caught by claude and it make sense to me: `advance` can commit past an uncommitted failed offset across batches (breaks at-least-once for the failed message).
Perhaps worthwhile to verify with real setup with the given scenario.
Full context:
---
advance (lines 200-240) is stateless across calls: it only inspects the current batch, takes first = batch[0], and commits last_good_offset + 1. The break-at-first-failure logic is correct within one batch, but a Kafka offset commit is cumulative and per-partition, so a later batch can commit an offset higher than an earlier offset that was deliberately held back.
Repro (the core multi-subscriber case shared streams are built for — e.g. two AssetWatchers A and B on one consumer):
1. Offset 10 acked by both → dispatched alone, commits offset 11.
2. Offsets 11-13: A acks, B hangs. B hits ack_timeout; the manager's _fail_subscriber_in_outstanding(B) marks B failed on 11-13 and excludes B from future
broadcasts. Pump harvests [11,12,13] (all failed), advance breaks at 11 → no commit (correct so far; committed position still 11).
3. B's trigger fails and unsubscribes. A keeps going. Offset 14 acked by A → dispatched in a separate batch [14]. advance sees first=14, acked and commits
offset 15.
Committed position has now jumped 11 → 15, skipping the failed 11-13. On triggerer restart/rebalance the consumer resumes at 15, so messages 11-13 — which B never persisted — are never redelivered. The whole point of the ack channel is defeated for those messages.
The single-subscriber case happens to be safe (the timeout fails the subscriber on all outstanding events and excludes it, and future zero-subscriber broadcasts carry all-zero counts → no commit, so the floor stays at 11). It's the multi-subscriber path — the primary use case — that loses data.
The core contract docstring says "stop before it so Kafka redelivers it and everything after," which is only true on restart/rebalance and only if no later commit overtakes it. To honor it, the producer needs per-lane (per topic-partition) state: once it declines to commit at offset N, it must never commit ≥ N
until N is reprocessed. Two viable fixes:
- Track self._floor: dict[(topic, partition), int]; in advance, clamp the commit so it can never advance to/past a recorded failed offset for that partition
(the partition then makes no further progress until restart — which is the correct at-least-once behavior), or
- seek() the consumer back to the failed offset on an involuntary failure so it's redelivered in-session.
I traced the manager's pump and ack-timeout paths carefully and don't see this handled upstream — advance is the only place that decides commits, and it gets
no memory of prior held-back offsets. Worth confirming with the author / #67523 owner whether the contract intends the producer to hold the partition floor (I believe it must).
##########
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)
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is not None:
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
+
+ def _dead_letter(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
+ if self._dlq_producer is None:
+ self._dlq_producer = KafkaProducerHook(kafka_config_id=self.kafka_config_id).get_producer()
+ producer = self._dlq_producer
+ for payload in payloads:
+ producer.produce(topic, value=payload.value, key=payload.key)
+ # Flush once so the whole dead-letter batch is durable before we commit past it.
+ producer.flush()
+
+ async def aclose(self) -> None:
+ """Flush the DLQ producer and close the consumer when the poll ends; best-effort."""
+ producer = self._dlq_producer
+ if producer is not None:
+ self._dlq_producer = None
+ try:
+ await sync_to_async(producer.flush)()
+ except Exception:
+ log.warning("Failed to flush Kafka DLQ producer", exc_info=True)
+ consumer = self._consumer
+ if consumer is not None:
+ self._consumer = None
+ try:
+ await sync_to_async(consumer.close)()
+ except Exception:
+ log.warning("Failed to close Kafka consumer", exc_info=True)
+
+
+class KafkaSharedStreamTrigger(BaseEventTrigger):
+ """
+ Event-driven trigger that watches Kafka topics through a shared consumer.
+
+ Triggers that declare the same ``topics`` + ``kafka_config_id`` share one
+ underlying Kafka consumer in the triggerer (a single poll loop broadcast
+ to every subscriber). Each subscriber fires a ``TriggerEvent`` per
+ message; override :meth:`filter_shared_stream` to fire only for the
+ messages this trigger cares about.
+
+ Designed to back an :class:`~airflow.sdk.AssetWatcher` for event-driven
+ scheduling. The offset is committed only after the derived
+ ``TriggerEvent`` is persisted -- see :class:`KafkaSharedStreamProducer`
+ for the ``enable.auto.commit=false`` requirement.
+
+ :param topics: Topics to watch.
+ :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 messages a subscriber
+ rejects; see :class:`KafkaSharedStreamProducer`. 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:
+ super().__init__()
+ self.topics = list(topics)
Review Comment:
Would it be better to sort at initialization stage so that we won't sort again for each `shared_stream_key` call?
```suggestion
self.topics = sorted(list(topics))
```
##########
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)
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is not None:
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
Review Comment:
Though it _shouldn't_ happen. Would it be better to add warning log for the `self._consumer` is None case?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3457108262
##########
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:
During initialization (`TriggerRunner#create_triggers`), if it waits for all triggers, a trigger can avoid miss message.
However, for a single trigger re-subscribe (failed -> unsubscribe -> subscribe), although the SharedStreamProducer can redeliver the message, it cannot make sure the trigger is in the group and receive the message again.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3452371607
##########
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:
Good catch! I add the `self._floor` and `seek` mechanism, so we can redelivery the batch from a failed message.
There is another high level design questions: does `SharedStreamManager` make sure all triggers are in the group before it starts the group? Currently, the `SharedStreamProducer` commits message if acked > 0, failed == 0, and rejected == 0. If a trigger requires the message but it's not in the group yet, the `SharedStreamProducer` doesn't know it and `SharedStreamProducer` still commits the message .
cc @Lee-W
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3472976264
##########
providers/apache/kafka/tests/integration/apache/kafka/triggers/test_shared_stream.py:
##########
@@ -0,0 +1,257 @@
+# 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.
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+import uuid
+
+import pytest
+from confluent_kafka import Consumer, Producer, TopicPartition
+
+pytest.importorskip("airflow.triggers.shared_stream")
+
+from airflow.models import Connection
+from airflow.providers.apache.kafka.triggers.shared_stream import KafkaSharedStreamProducer
+from airflow.triggers.shared_stream import AdvanceItem, AdvanceOutcome
+
+BROKER = "broker:29092"
+
+
+def _produce(topic: str, values: list[bytes]) -> None:
+ """Produce ``values`` to partition 0 of ``topic`` so offsets are deterministic (0, 1, ...)."""
+ producer = Producer({"bootstrap.servers": BROKER})
+ for value in values:
+ producer.produce(topic, value=value, partition=0)
+ assert producer.flush(10) == 0
+
+
+def _poll_one(consumer: Consumer, timeout: float):
+ """Poll ``consumer`` until a non-error message arrives or ``timeout`` elapses."""
+ deadline = time.monotonic() + timeout
+ while time.monotonic() < deadline:
+ message = consumer.poll(1.0)
+ if message is not None and message.error() is None:
+ return message
+ return None
+
+
+async def _take(stream, n: int) -> list:
+ out = []
+ async for item in stream:
+ out.append(item)
+ if len(out) >= n:
+ break
+ return out
+
+
+@pytest.mark.integration("kafka")
+class TestKafkaSharedStreamProducerIntegration:
+ @pytest.fixture(autouse=True)
+ def setup_connections(self, create_connection_without_db):
+ self.group = f"shared-stream-it-{uuid.uuid4().hex[:8]}"
+ self.conn_id = "kafka_shared_stream_it"
+ self.bad_conn_id = "kafka_shared_stream_it_autocommit"
+ create_connection_without_db(
+ Connection(
+ conn_id=self.conn_id,
+ conn_type="kafka",
+ extra=json.dumps(
+ {
+ "bootstrap.servers": BROKER,
+ "group.id": self.group,
+ "enable.auto.commit": False,
+ "auto.offset.reset": "beginning",
+ "socket.timeout.ms": 10000,
+ }
+ ),
+ )
+ )
+ create_connection_without_db(
+ Connection(
+ conn_id=self.bad_conn_id,
+ conn_type="kafka",
+ extra=json.dumps(
+ {
+ "bootstrap.servers": BROKER,
+ "group.id": f"{self.group}-bad",
+ "enable.auto.commit": True,
+ "auto.offset.reset": "beginning",
+ }
+ ),
+ )
+ )
Review Comment:
```suggestion
for conn_id, extra in (
(
self.conn_id,
{
"bootstrap.servers": BROKER,
"group.id": self.group,
"enable.auto.commit": False,
"auto.offset.reset": "beginning",
"socket.timeout.ms": 10000,
},
),
(
self.bad_conn_id,
{
"bootstrap.servers": BROKER,
"group.id": f"{self.group}-bad",
"enable.auto.commit": True,
"auto.offset.reset": "beginning",
}
)
):
create_connection_without_db(
Connection(
conn_id=conn_id,
conn_type="kafka",
extra=json.dumps(extras),
)
)
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3452275098
##########
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)
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is not None:
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
+
+ def _dead_letter(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
+ if self._dlq_producer is None:
+ self._dlq_producer = KafkaProducerHook(kafka_config_id=self.kafka_config_id).get_producer()
+ producer = self._dlq_producer
+ for payload in payloads:
+ producer.produce(topic, value=payload.value, key=payload.key)
+ # Flush once so the whole dead-letter batch is durable before we commit past it.
+ producer.flush()
+
+ async def aclose(self) -> None:
+ """Flush the DLQ producer and close the consumer when the poll ends; best-effort."""
+ producer = self._dlq_producer
+ if producer is not None:
+ self._dlq_producer = None
+ try:
+ await sync_to_async(producer.flush)()
+ except Exception:
+ log.warning("Failed to flush Kafka DLQ producer", exc_info=True)
+ consumer = self._consumer
+ if consumer is not None:
+ self._consumer = None
+ try:
+ await sync_to_async(consumer.close)()
+ except Exception:
+ log.warning("Failed to close Kafka consumer", exc_info=True)
+
+
+class KafkaSharedStreamTrigger(BaseEventTrigger):
+ """
+ Event-driven trigger that watches Kafka topics through a shared consumer.
+
+ Triggers that declare the same ``topics`` + ``kafka_config_id`` share one
+ underlying Kafka consumer in the triggerer (a single poll loop broadcast
+ to every subscriber). Each subscriber fires a ``TriggerEvent`` per
+ message; override :meth:`filter_shared_stream` to fire only for the
+ messages this trigger cares about.
+
+ Designed to back an :class:`~airflow.sdk.AssetWatcher` for event-driven
+ scheduling. The offset is committed only after the derived
+ ``TriggerEvent`` is persisted -- see :class:`KafkaSharedStreamProducer`
+ for the ``enable.auto.commit=false`` requirement.
+
+ :param topics: Topics to watch.
+ :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 messages a subscriber
+ rejects; see :class:`KafkaSharedStreamProducer`. 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:
+ super().__init__()
+ self.topics = list(topics)
Review Comment:
Changed to run `sorted` in `KafkaSharedStreamTrigger#init`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3452277333
##########
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)
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is not None:
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
Review Comment:
Added log when `self._consumer` is None.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r4164634980
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,407 @@
+# 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.
+
+This module builds on the shared-stream ack channel, which was added in
+Airflow 3.3. Importing this module on an older version raises
+``AirflowOptionalProviderFeatureException``.
+"""
+
+from __future__ import annotations
+
+from airflow.providers.apache.kafka.version_compat import AIRFLOW_V_3_3_PLUS
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+if not AIRFLOW_V_3_3_PLUS:
+ raise AirflowOptionalProviderFeatureException("Kafka shared-stream triggers need Airflow 3.3+.")
+
+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 that 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 floor
+ is held and the consumer is sought back so the message is redelivered this
+ session: a failure to the remaining healthy subscribers, a zero-subscriber
+ broadcast to whoever subscribes next.
+
+ :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
+ # lane (topic, partition) -> earliest offset held for redelivery. An
+ # involuntary failure / zero-subscriber broadcast must not be committed
+ # past until it is reprocessed; the floor is carried across advance()
+ # calls so a later batch's cumulative commit cannot overtake it.
+ self._floor: dict[Hashable, int] = {}
+
+ 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 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()
+ 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
+
+ async 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.
+ """
+ connection = await hook.aget_connection(self.kafka_config_id)
+ auto_commit = connection.extra_dejson.get("enable.auto.commit", True)
+ 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 one partition's resolved prefix; hold and seek back the first held offset.
+
+ 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. Within a batch 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`` (ack timeout / overflow) or all-zero (a broadcast no
+ subscriber was online for) -- hold the floor here and ``seek`` the
+ consumer back to it, so it and everything after are redelivered this
+ session. A failure lands on the healthy subscribers the manager left
+ online; a zero-subscriber broadcast reaches whoever subscribes next,
+ instead of waiting for a full consumer rebuild. The floor lifts once a
+ re-read batch acks through the held offset.
+
+ The floor is per-lane and carried across calls (``self._floor``): once a
+ lane holds at offset ``H``, no later batch may commit past ``H`` until a
+ re-read batch starting at or before ``H`` acks through it.
+
+ 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
+ lane = self.get_advance_lane(first)
+ commit_through: int | None = None
+ held_offset: int | None = None
+ to_dlq: list[KafkaBrokerPayload] = []
+ for item in batch:
+ payload = item.broker_payload
+ outcome = item.outcome
+ if outcome.failed > 0 or (outcome.acked == 0 and outcome.rejected == 0):
+ # Held: an involuntary failure, or a broadcast no subscriber took.
+ # Seek back so it -- and everything after -- is redelivered this
+ # session, to the subscribers the manager left online or to
+ # whoever subscribes next.
+ held_offset = payload.offset
+ break
+ if outcome.rejected and self.dlq_topic is not None:
+ to_dlq.append(payload)
+ commit_through = payload.offset
+
+ floor = self._resolve_floor(lane, first.offset, commit_through, held_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_through is not None:
+ target = commit_through + 1
+ if floor is not None:
+ # floor may be a previous batch's held offset, so commit_through
+ # can exceed it. Use minimal value to avoid it.
+ target = min(target, floor)
+ await sync_to_async(self._commit)(first.topic, first.partition, target)
+ if held_offset is not None:
+ # Seek to the lane's floor (the earliest unresolved held offset), not
+ # this batch's break point: an earlier batch may hold a smaller offset,
+ # and seeking forward to a later one would skip and lose it.
+ seek_target = floor if floor is not None else held_offset
+ await sync_to_async(self._seek)(first.topic, first.partition, seek_target)
+
+ def _resolve_floor(
+ self, lane: Hashable, batch_start: int, commit_through: int | None, held_offset: int | None
+ ) -> int | None:
+ """
+ Update and return this lane's redelivery floor.
+
+ The floor is the earliest unresolved held offset on the lane: the minimum
+ of any still-unresolved previous floor and this batch's held offset.
+ A previous floor lifts only when this batch starts at or before it and acks
+ consecutively through it.
+ A later arrived ack before the seek operation, or a failure at a later offset can
+ neither clear nor advance a floor whose offset was never reprocessed.
+ """
+ previous = self._floor.get(lane)
+ cleared = (
+ previous is not None
+ and batch_start <= previous
+ and commit_through is not None
+ and commit_through >= previous
+ )
+ candidates: list[int] = []
+ if previous is not None and not cleared:
+ candidates.append(previous)
+ if held_offset is not None:
+ candidates.append(held_offset)
+ floor: int | None = min(candidates) if candidates else None
+ if floor is None:
+ self._floor.pop(lane, None)
+ else:
+ self._floor[lane] = floor
+ return floor
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ # When advance() raises, the triggerer fails every trigger that shares this
+ # producer, and Kafka redelivers from the last committed offset. Returning
+ # here would let the triggerer treat the batch as advanced although nothing
+ # was committed.
+ raise RuntimeError(f"Cannot commit {topic}[{partition}]@{offset}: no open consumer")
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
+
+ def _seek(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ raise RuntimeError(f"Cannot seek {topic}[{partition}] to {offset}: no open consumer")
+ self._consumer.seek(TopicPartition(topic, partition, offset))
+
+ def _dead_letter(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
+ if self._dlq_producer is None:
+ self._dlq_producer = KafkaProducerHook(kafka_config_id=self.kafka_config_id).get_producer()
+ producer = self._dlq_producer
+ for payload in payloads:
+ producer.produce(topic, value=payload.value, key=payload.key)
+ # Flush once so the whole dead-letter batch is durable before we commit past it.
+ producer.flush()
+
+ async def aclose(self) -> None:
+ """Flush the DLQ producer and close the consumer when the poll ends; best-effort."""
+ producer = self._dlq_producer
+ if producer is not None:
+ self._dlq_producer = None
+ try:
+ await sync_to_async(producer.flush)()
+ except Exception:
+ log.warning("Failed to flush Kafka DLQ producer", exc_info=True)
+ consumer = self._consumer
+ if consumer is not None:
+ self._consumer = None
+ try:
+ await sync_to_async(consumer.close)()
+ except Exception:
+ log.warning("Failed to close Kafka consumer", exc_info=True)
+
+
+class KafkaSharedStreamTrigger(BaseEventTrigger):
+ """
+ Event-driven trigger that watches Kafka topics through a shared consumer.
+
+ Triggers that declare the same ``topics`` + ``kafka_config_id`` share one
+ underlying Kafka consumer in the triggerer (a single poll loop broadcast
+ to every subscriber). Each subscriber fires a ``TriggerEvent`` per
+ message; override :meth:`filter_shared_stream` to fire only for the
+ messages this trigger cares about.
+
+ Designed to back an :class:`~airflow.sdk.AssetWatcher` for event-driven
+ scheduling. The offset is committed only after the derived
+ ``TriggerEvent`` is persisted -- see :class:`KafkaSharedStreamProducer`
+ for the ``enable.auto.commit=false`` requirement.
+
+ :param topics: Topics to watch.
+ :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 messages a subscriber
+ rejects; see :class:`KafkaSharedStreamProducer`. 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:
+ super().__init__()
+ # Sort once here so shared_stream_key() is order-independent without
+ # re-sorting on every call: triggers on the same topics in any order
+ # share one consumer.
+ self.topics = sorted(topics)
+ self.kafka_config_id = kafka_config_id
+ self.poll_timeout = poll_timeout
+ self.dlq_topic = dlq_topic
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+ "airflow.providers.apache.kafka.triggers.shared_stream.KafkaSharedStreamTrigger",
+ {
+ "topics": self.topics,
+ "kafka_config_id": self.kafka_config_id,
+ "poll_timeout": self.poll_timeout,
+ "dlq_topic": self.dlq_topic,
+ },
+ )
+
+ def shared_stream_key(self) -> Hashable:
Review Comment:
ah, bad design by me then. i guess we'll need to live with it
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3452266785
##########
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})."
+ )
Review Comment:
Updated it to use `aget_connection`. Thanks for the suggestion.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3457743431
##########
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:
`create_triggers` can't know the full trigger set — it's open and dynamic (multiple Dag processors, `lazy shared_stream_key` after template rendering, triggers spawned at runtime). There's no well-defined "all triggers" to wait for.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r4162035365
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,398 @@
+# 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
Review Comment:
Thanks for the suggestion. It now raises `AirflowOptionalProviderFeatureException("Kafka shared-stream triggers need Airflow 3.3+.")`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r4162076824
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,398 @@
+# 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 that 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 floor
+ is held and the consumer is sought back so the message is redelivered this
+ session: a failure to the remaining healthy subscribers, a zero-subscriber
+ broadcast to whoever subscribes next.
+
+ :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
+ # lane (topic, partition) -> earliest offset held for redelivery. An
+ # involuntary failure / zero-subscriber broadcast must not be committed
+ # past until it is reprocessed; the floor is carried across advance()
+ # calls so a later batch's cumulative commit cannot overtake it.
+ self._floor: dict[Hashable, int] = {}
+
+ 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 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()
+ 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
+
+ async 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.
+ """
+ connection = await hook.aget_connection(self.kafka_config_id)
+ auto_commit = connection.extra_dejson.get("enable.auto.commit", True)
+ 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 one partition's resolved prefix; hold and seek back the first held offset.
+
+ 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. Within a batch 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`` (ack timeout / overflow) or all-zero (a broadcast no
+ subscriber was online for) -- hold the floor here and ``seek`` the
+ consumer back to it, so it and everything after are redelivered this
+ session. A failure lands on the healthy subscribers the manager left
+ online; a zero-subscriber broadcast reaches whoever subscribes next,
+ instead of waiting for a full consumer rebuild. The floor lifts once a
+ re-read batch acks through the held offset.
+
+ The floor is per-lane and carried across calls (``self._floor``): once a
+ lane holds at offset ``H``, no later batch may commit past ``H`` until a
+ re-read batch starting at or before ``H`` acks through it.
+
+ 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
+ lane = self.get_advance_lane(first)
+ commit_through: int | None = None
+ held_offset: int | None = None
+ to_dlq: list[KafkaBrokerPayload] = []
+ for item in batch:
+ payload = item.broker_payload
+ outcome = item.outcome
+ if outcome.failed > 0 or (outcome.acked == 0 and outcome.rejected == 0):
+ # Held: an involuntary failure, or a broadcast no subscriber took.
+ # Seek back so it -- and everything after -- is redelivered this
+ # session, to the subscribers the manager left online or to
+ # whoever subscribes next.
+ held_offset = payload.offset
+ break
+ if outcome.rejected and self.dlq_topic is not None:
+ to_dlq.append(payload)
+ commit_through = payload.offset
+
+ floor = self._resolve_floor(lane, first.offset, commit_through, held_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_through is not None:
+ target = commit_through + 1
+ if floor is not None:
+ # floor may be a previous batch's held offset, so commit_through
+ # can exceed it. Use minimal value to avoid it.
+ target = min(target, floor)
+ await sync_to_async(self._commit)(first.topic, first.partition, target)
+ if held_offset is not None:
+ # Seek to the lane's floor (the earliest unresolved held offset), not
+ # this batch's break point: an earlier batch may hold a smaller offset,
+ # and seeking forward to a later one would skip and lose it.
+ seek_target = floor if floor is not None else held_offset
+ await sync_to_async(self._seek)(first.topic, first.partition, seek_target)
+
+ def _resolve_floor(
+ self, lane: Hashable, batch_start: int, commit_through: int | None, held_offset: int | None
+ ) -> int | None:
+ """
+ Update and return this lane's redelivery floor.
+
+ The floor is the earliest unresolved held offset on the lane: the minimum
+ of any still-unresolved previous floor and this batch's held offset.
+ A previous floor lifts only when this batch starts at or before it and acks
+ consecutively through it.
+ A later arrived ack before the seek operation, or a failure at a later offset can
+ neither clear nor advance a floor whose offset was never reprocessed.
+ """
+ previous = self._floor.get(lane)
+ cleared = (
+ previous is not None
+ and batch_start <= previous
+ and commit_through is not None
+ and commit_through >= previous
+ )
+ candidates: list[int] = []
+ if previous is not None and not cleared:
+ candidates.append(previous)
+ if held_offset is not None:
+ candidates.append(held_offset)
+ floor: int | None = min(candidates) if candidates else None
+ if floor is None:
+ self._floor.pop(lane, None)
+ else:
+ self._floor[lane] = floor
+ return floor
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
Review Comment:
Agree, changed to raise `RuntimeError` here.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3499025364
##########
providers/apache/kafka/tests/integration/apache/kafka/triggers/test_shared_stream.py:
##########
@@ -0,0 +1,257 @@
+# 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.
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+import uuid
+
+import pytest
+from confluent_kafka import Consumer, Producer, TopicPartition
+
+pytest.importorskip("airflow.triggers.shared_stream")
+
+from airflow.models import Connection
+from airflow.providers.apache.kafka.triggers.shared_stream import KafkaSharedStreamProducer
+from airflow.triggers.shared_stream import AdvanceItem, AdvanceOutcome
+
+BROKER = "broker:29092"
+
+
+def _produce(topic: str, values: list[bytes]) -> None:
+ """Produce ``values`` to partition 0 of ``topic`` so offsets are deterministic (0, 1, ...)."""
+ producer = Producer({"bootstrap.servers": BROKER})
+ for value in values:
+ producer.produce(topic, value=value, partition=0)
+ assert producer.flush(10) == 0
+
+
+def _poll_one(consumer: Consumer, timeout: float):
+ """Poll ``consumer`` until a non-error message arrives or ``timeout`` elapses."""
+ deadline = time.monotonic() + timeout
+ while time.monotonic() < deadline:
+ message = consumer.poll(1.0)
+ if message is not None and message.error() is None:
+ return message
+ return None
+
+
+async def _take(stream, n: int) -> list:
+ out = []
+ async for item in stream:
+ out.append(item)
+ if len(out) >= n:
+ break
+ return out
+
+
+@pytest.mark.integration("kafka")
+class TestKafkaSharedStreamProducerIntegration:
+ @pytest.fixture(autouse=True)
+ def setup_connections(self, create_connection_without_db):
+ self.group = f"shared-stream-it-{uuid.uuid4().hex[:8]}"
+ self.conn_id = "kafka_shared_stream_it"
+ self.bad_conn_id = "kafka_shared_stream_it_autocommit"
+ create_connection_without_db(
+ Connection(
+ conn_id=self.conn_id,
+ conn_type="kafka",
+ extra=json.dumps(
+ {
+ "bootstrap.servers": BROKER,
+ "group.id": self.group,
+ "enable.auto.commit": False,
+ "auto.offset.reset": "beginning",
+ "socket.timeout.ms": 10000,
+ }
+ ),
+ )
+ )
+ create_connection_without_db(
+ Connection(
+ conn_id=self.bad_conn_id,
+ conn_type="kafka",
+ extra=json.dumps(
+ {
+ "bootstrap.servers": BROKER,
+ "group.id": f"{self.group}-bad",
+ "enable.auto.commit": True,
+ "auto.offset.reset": "beginning",
+ }
+ ),
+ )
+ )
Review Comment:
Thanks for the suggestion. Updated it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on PR #68625:
URL: https://github.com/apache/airflow/pull/68625#issuecomment-5436808502
@jason810496 I update real integration setup to PR description and verify it can work. Thank you.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #68625:
URL: https://github.com/apache/airflow/pull/68625#issuecomment-5343508906
This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r4163977734
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,407 @@
+# 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.
+
+This module builds on the shared-stream ack channel, which was added in
+Airflow 3.3. Importing this module on an older version raises
+``AirflowOptionalProviderFeatureException``.
+"""
+
+from __future__ import annotations
+
+from airflow.providers.apache.kafka.version_compat import AIRFLOW_V_3_3_PLUS
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+if not AIRFLOW_V_3_3_PLUS:
+ raise AirflowOptionalProviderFeatureException("Kafka shared-stream triggers need Airflow 3.3+.")
+
+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 that 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 floor
+ is held and the consumer is sought back so the message is redelivered this
+ session: a failure to the remaining healthy subscribers, a zero-subscriber
+ broadcast to whoever subscribes next.
+
+ :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
+ # lane (topic, partition) -> earliest offset held for redelivery. An
+ # involuntary failure / zero-subscriber broadcast must not be committed
+ # past until it is reprocessed; the floor is carried across advance()
+ # calls so a later batch's cumulative commit cannot overtake it.
+ self._floor: dict[Hashable, int] = {}
+
+ 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 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()
+ 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
+
+ async 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.
+ """
+ connection = await hook.aget_connection(self.kafka_config_id)
+ auto_commit = connection.extra_dejson.get("enable.auto.commit", True)
+ 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 one partition's resolved prefix; hold and seek back the first held offset.
+
+ 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. Within a batch 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`` (ack timeout / overflow) or all-zero (a broadcast no
+ subscriber was online for) -- hold the floor here and ``seek`` the
+ consumer back to it, so it and everything after are redelivered this
+ session. A failure lands on the healthy subscribers the manager left
+ online; a zero-subscriber broadcast reaches whoever subscribes next,
+ instead of waiting for a full consumer rebuild. The floor lifts once a
+ re-read batch acks through the held offset.
+
+ The floor is per-lane and carried across calls (``self._floor``): once a
+ lane holds at offset ``H``, no later batch may commit past ``H`` until a
+ re-read batch starting at or before ``H`` acks through it.
+
+ 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
+ lane = self.get_advance_lane(first)
+ commit_through: int | None = None
+ held_offset: int | None = None
+ to_dlq: list[KafkaBrokerPayload] = []
+ for item in batch:
+ payload = item.broker_payload
+ outcome = item.outcome
+ if outcome.failed > 0 or (outcome.acked == 0 and outcome.rejected == 0):
+ # Held: an involuntary failure, or a broadcast no subscriber took.
+ # Seek back so it -- and everything after -- is redelivered this
+ # session, to the subscribers the manager left online or to
+ # whoever subscribes next.
+ held_offset = payload.offset
+ break
+ if outcome.rejected and self.dlq_topic is not None:
+ to_dlq.append(payload)
+ commit_through = payload.offset
+
+ floor = self._resolve_floor(lane, first.offset, commit_through, held_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_through is not None:
+ target = commit_through + 1
+ if floor is not None:
+ # floor may be a previous batch's held offset, so commit_through
+ # can exceed it. Use minimal value to avoid it.
+ target = min(target, floor)
+ await sync_to_async(self._commit)(first.topic, first.partition, target)
+ if held_offset is not None:
+ # Seek to the lane's floor (the earliest unresolved held offset), not
+ # this batch's break point: an earlier batch may hold a smaller offset,
+ # and seeking forward to a later one would skip and lose it.
+ seek_target = floor if floor is not None else held_offset
+ await sync_to_async(self._seek)(first.topic, first.partition, seek_target)
+
+ def _resolve_floor(
+ self, lane: Hashable, batch_start: int, commit_through: int | None, held_offset: int | None
+ ) -> int | None:
+ """
+ Update and return this lane's redelivery floor.
+
+ The floor is the earliest unresolved held offset on the lane: the minimum
+ of any still-unresolved previous floor and this batch's held offset.
+ A previous floor lifts only when this batch starts at or before it and acks
+ consecutively through it.
+ A later arrived ack before the seek operation, or a failure at a later offset can
+ neither clear nor advance a floor whose offset was never reprocessed.
+ """
+ previous = self._floor.get(lane)
+ cleared = (
+ previous is not None
+ and batch_start <= previous
+ and commit_through is not None
+ and commit_through >= previous
+ )
+ candidates: list[int] = []
+ if previous is not None and not cleared:
+ candidates.append(previous)
+ if held_offset is not None:
+ candidates.append(held_offset)
+ floor: int | None = min(candidates) if candidates else None
+ if floor is None:
+ self._floor.pop(lane, None)
+ else:
+ self._floor[lane] = floor
+ return floor
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ # When advance() raises, the triggerer fails every trigger that shares this
+ # producer, and Kafka redelivers from the last committed offset. Returning
+ # here would let the triggerer treat the batch as advanced although nothing
+ # was committed.
+ raise RuntimeError(f"Cannot commit {topic}[{partition}]@{offset}: no open consumer")
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
+
+ def _seek(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ raise RuntimeError(f"Cannot seek {topic}[{partition}] to {offset}: no open consumer")
+ self._consumer.seek(TopicPartition(topic, partition, offset))
+
+ def _dead_letter(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
+ if self._dlq_producer is None:
+ self._dlq_producer = KafkaProducerHook(kafka_config_id=self.kafka_config_id).get_producer()
+ producer = self._dlq_producer
+ for payload in payloads:
+ producer.produce(topic, value=payload.value, key=payload.key)
+ # Flush once so the whole dead-letter batch is durable before we commit past it.
+ producer.flush()
+
+ async def aclose(self) -> None:
+ """Flush the DLQ producer and close the consumer when the poll ends; best-effort."""
+ producer = self._dlq_producer
+ if producer is not None:
+ self._dlq_producer = None
+ try:
+ await sync_to_async(producer.flush)()
+ except Exception:
+ log.warning("Failed to flush Kafka DLQ producer", exc_info=True)
+ consumer = self._consumer
+ if consumer is not None:
+ self._consumer = None
+ try:
+ await sync_to_async(consumer.close)()
+ except Exception:
+ log.warning("Failed to close Kafka consumer", exc_info=True)
+
+
+class KafkaSharedStreamTrigger(BaseEventTrigger):
+ """
+ Event-driven trigger that watches Kafka topics through a shared consumer.
+
+ Triggers that declare the same ``topics`` + ``kafka_config_id`` share one
+ underlying Kafka consumer in the triggerer (a single poll loop broadcast
+ to every subscriber). Each subscriber fires a ``TriggerEvent`` per
+ message; override :meth:`filter_shared_stream` to fire only for the
+ messages this trigger cares about.
+
+ Designed to back an :class:`~airflow.sdk.AssetWatcher` for event-driven
+ scheduling. The offset is committed only after the derived
+ ``TriggerEvent`` is persisted -- see :class:`KafkaSharedStreamProducer`
+ for the ``enable.auto.commit=false`` requirement.
+
+ :param topics: Topics to watch.
+ :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 messages a subscriber
+ rejects; see :class:`KafkaSharedStreamProducer`. 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:
+ super().__init__()
+ # Sort once here so shared_stream_key() is order-independent without
+ # re-sorting on every call: triggers on the same topics in any order
+ # share one consumer.
+ self.topics = sorted(topics)
+ self.kafka_config_id = kafka_config_id
+ self.poll_timeout = poll_timeout
+ self.dlq_topic = dlq_topic
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+ "airflow.providers.apache.kafka.triggers.shared_stream.KafkaSharedStreamTrigger",
+ {
+ "topics": self.topics,
+ "kafka_config_id": self.kafka_config_id,
+ "poll_timeout": self.poll_timeout,
+ "dlq_topic": self.dlq_topic,
+ },
+ )
+
+ def shared_stream_key(self) -> Hashable:
Review Comment:
`shared_stream_key()` can't be renamed in the provider alone. It is the `BaseEventTrigger` method that `run_trigger` in the triggerer calls. With only `create_shared_stream_key()` defined, the triggerer would get `None` from the base method and skip the shared stream.
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,407 @@
+# 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.
+
+This module builds on the shared-stream ack channel, which was added in
+Airflow 3.3. Importing this module on an older version raises
+``AirflowOptionalProviderFeatureException``.
+"""
+
+from __future__ import annotations
+
+from airflow.providers.apache.kafka.version_compat import AIRFLOW_V_3_3_PLUS
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+if not AIRFLOW_V_3_3_PLUS:
+ raise AirflowOptionalProviderFeatureException("Kafka shared-stream triggers need Airflow 3.3+.")
+
+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 that 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 floor
+ is held and the consumer is sought back so the message is redelivered this
+ session: a failure to the remaining healthy subscribers, a zero-subscriber
+ broadcast to whoever subscribes next.
+
+ :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
+ # lane (topic, partition) -> earliest offset held for redelivery. An
+ # involuntary failure / zero-subscriber broadcast must not be committed
+ # past until it is reprocessed; the floor is carried across advance()
+ # calls so a later batch's cumulative commit cannot overtake it.
+ self._floor: dict[Hashable, int] = {}
+
+ 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 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()
+ 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
+
+ async 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.
+ """
+ connection = await hook.aget_connection(self.kafka_config_id)
+ auto_commit = connection.extra_dejson.get("enable.auto.commit", True)
+ 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 one partition's resolved prefix; hold and seek back the first held offset.
+
+ 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. Within a batch 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`` (ack timeout / overflow) or all-zero (a broadcast no
+ subscriber was online for) -- hold the floor here and ``seek`` the
+ consumer back to it, so it and everything after are redelivered this
+ session. A failure lands on the healthy subscribers the manager left
+ online; a zero-subscriber broadcast reaches whoever subscribes next,
+ instead of waiting for a full consumer rebuild. The floor lifts once a
+ re-read batch acks through the held offset.
+
+ The floor is per-lane and carried across calls (``self._floor``): once a
+ lane holds at offset ``H``, no later batch may commit past ``H`` until a
+ re-read batch starting at or before ``H`` acks through it.
+
+ 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
+ lane = self.get_advance_lane(first)
+ commit_through: int | None = None
+ held_offset: int | None = None
+ to_dlq: list[KafkaBrokerPayload] = []
+ for item in batch:
+ payload = item.broker_payload
+ outcome = item.outcome
+ if outcome.failed > 0 or (outcome.acked == 0 and outcome.rejected == 0):
+ # Held: an involuntary failure, or a broadcast no subscriber took.
+ # Seek back so it -- and everything after -- is redelivered this
+ # session, to the subscribers the manager left online or to
+ # whoever subscribes next.
+ held_offset = payload.offset
+ break
+ if outcome.rejected and self.dlq_topic is not None:
+ to_dlq.append(payload)
+ commit_through = payload.offset
+
+ floor = self._resolve_floor(lane, first.offset, commit_through, held_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_through is not None:
+ target = commit_through + 1
+ if floor is not None:
+ # floor may be a previous batch's held offset, so commit_through
+ # can exceed it. Use minimal value to avoid it.
+ target = min(target, floor)
+ await sync_to_async(self._commit)(first.topic, first.partition, target)
+ if held_offset is not None:
+ # Seek to the lane's floor (the earliest unresolved held offset), not
+ # this batch's break point: an earlier batch may hold a smaller offset,
+ # and seeking forward to a later one would skip and lose it.
+ seek_target = floor if floor is not None else held_offset
+ await sync_to_async(self._seek)(first.topic, first.partition, seek_target)
+
+ def _resolve_floor(
+ self, lane: Hashable, batch_start: int, commit_through: int | None, held_offset: int | None
+ ) -> int | None:
+ """
+ Update and return this lane's redelivery floor.
+
+ The floor is the earliest unresolved held offset on the lane: the minimum
+ of any still-unresolved previous floor and this batch's held offset.
+ A previous floor lifts only when this batch starts at or before it and acks
+ consecutively through it.
+ A later arrived ack before the seek operation, or a failure at a later offset can
+ neither clear nor advance a floor whose offset was never reprocessed.
+ """
+ previous = self._floor.get(lane)
+ cleared = (
+ previous is not None
+ and batch_start <= previous
+ and commit_through is not None
+ and commit_through >= previous
+ )
+ candidates: list[int] = []
+ if previous is not None and not cleared:
+ candidates.append(previous)
+ if held_offset is not None:
+ candidates.append(held_offset)
+ floor: int | None = min(candidates) if candidates else None
+ if floor is None:
+ self._floor.pop(lane, None)
+ else:
+ self._floor[lane] = floor
+ return floor
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ # When advance() raises, the triggerer fails every trigger that shares this
+ # producer, and Kafka redelivers from the last committed offset. Returning
+ # here would let the triggerer treat the batch as advanced although nothing
+ # was committed.
+ raise RuntimeError(f"Cannot commit {topic}[{partition}]@{offset}: no open consumer")
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
+
+ def _seek(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ raise RuntimeError(f"Cannot seek {topic}[{partition}] to {offset}: no open consumer")
+ self._consumer.seek(TopicPartition(topic, partition, offset))
+
+ def _dead_letter(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
Review Comment:
Updated it.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r4163443506
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,407 @@
+# 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.
+
+This module builds on the shared-stream ack channel, which was added in
+Airflow 3.3. Importing this module on an older version raises
+``AirflowOptionalProviderFeatureException``.
+"""
+
+from __future__ import annotations
+
+from airflow.providers.apache.kafka.version_compat import AIRFLOW_V_3_3_PLUS
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+if not AIRFLOW_V_3_3_PLUS:
+ raise AirflowOptionalProviderFeatureException("Kafka shared-stream triggers need Airflow 3.3+.")
+
+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 that 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 floor
+ is held and the consumer is sought back so the message is redelivered this
+ session: a failure to the remaining healthy subscribers, a zero-subscriber
+ broadcast to whoever subscribes next.
+
+ :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
+ # lane (topic, partition) -> earliest offset held for redelivery. An
+ # involuntary failure / zero-subscriber broadcast must not be committed
+ # past until it is reprocessed; the floor is carried across advance()
+ # calls so a later batch's cumulative commit cannot overtake it.
+ self._floor: dict[Hashable, int] = {}
+
+ 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 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()
+ 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
+
+ async 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.
+ """
+ connection = await hook.aget_connection(self.kafka_config_id)
+ auto_commit = connection.extra_dejson.get("enable.auto.commit", True)
+ 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 one partition's resolved prefix; hold and seek back the first held offset.
+
+ 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. Within a batch 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`` (ack timeout / overflow) or all-zero (a broadcast no
+ subscriber was online for) -- hold the floor here and ``seek`` the
+ consumer back to it, so it and everything after are redelivered this
+ session. A failure lands on the healthy subscribers the manager left
+ online; a zero-subscriber broadcast reaches whoever subscribes next,
+ instead of waiting for a full consumer rebuild. The floor lifts once a
+ re-read batch acks through the held offset.
+
+ The floor is per-lane and carried across calls (``self._floor``): once a
+ lane holds at offset ``H``, no later batch may commit past ``H`` until a
+ re-read batch starting at or before ``H`` acks through it.
+
+ 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
+ lane = self.get_advance_lane(first)
+ commit_through: int | None = None
+ held_offset: int | None = None
+ to_dlq: list[KafkaBrokerPayload] = []
+ for item in batch:
+ payload = item.broker_payload
+ outcome = item.outcome
+ if outcome.failed > 0 or (outcome.acked == 0 and outcome.rejected == 0):
+ # Held: an involuntary failure, or a broadcast no subscriber took.
+ # Seek back so it -- and everything after -- is redelivered this
+ # session, to the subscribers the manager left online or to
+ # whoever subscribes next.
+ held_offset = payload.offset
+ break
+ if outcome.rejected and self.dlq_topic is not None:
+ to_dlq.append(payload)
+ commit_through = payload.offset
+
+ floor = self._resolve_floor(lane, first.offset, commit_through, held_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_through is not None:
+ target = commit_through + 1
+ if floor is not None:
+ # floor may be a previous batch's held offset, so commit_through
+ # can exceed it. Use minimal value to avoid it.
+ target = min(target, floor)
+ await sync_to_async(self._commit)(first.topic, first.partition, target)
+ if held_offset is not None:
+ # Seek to the lane's floor (the earliest unresolved held offset), not
+ # this batch's break point: an earlier batch may hold a smaller offset,
+ # and seeking forward to a later one would skip and lose it.
+ seek_target = floor if floor is not None else held_offset
+ await sync_to_async(self._seek)(first.topic, first.partition, seek_target)
+
+ def _resolve_floor(
+ self, lane: Hashable, batch_start: int, commit_through: int | None, held_offset: int | None
+ ) -> int | None:
+ """
+ Update and return this lane's redelivery floor.
+
+ The floor is the earliest unresolved held offset on the lane: the minimum
+ of any still-unresolved previous floor and this batch's held offset.
+ A previous floor lifts only when this batch starts at or before it and acks
+ consecutively through it.
+ A later arrived ack before the seek operation, or a failure at a later offset can
+ neither clear nor advance a floor whose offset was never reprocessed.
+ """
+ previous = self._floor.get(lane)
+ cleared = (
+ previous is not None
+ and batch_start <= previous
+ and commit_through is not None
+ and commit_through >= previous
+ )
+ candidates: list[int] = []
+ if previous is not None and not cleared:
+ candidates.append(previous)
+ if held_offset is not None:
+ candidates.append(held_offset)
+ floor: int | None = min(candidates) if candidates else None
+ if floor is None:
+ self._floor.pop(lane, None)
+ else:
+ self._floor[lane] = floor
+ return floor
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ # When advance() raises, the triggerer fails every trigger that shares this
+ # producer, and Kafka redelivers from the last committed offset. Returning
+ # here would let the triggerer treat the batch as advanced although nothing
+ # was committed.
+ raise RuntimeError(f"Cannot commit {topic}[{partition}]@{offset}: no open consumer")
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
+
+ def _seek(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ raise RuntimeError(f"Cannot seek {topic}[{partition}] to {offset}: no open consumer")
+ self._consumer.seek(TopicPartition(topic, partition, offset))
+
+ def _dead_letter(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
+ if self._dlq_producer is None:
+ self._dlq_producer = KafkaProducerHook(kafka_config_id=self.kafka_config_id).get_producer()
+ producer = self._dlq_producer
+ for payload in payloads:
+ producer.produce(topic, value=payload.value, key=payload.key)
+ # Flush once so the whole dead-letter batch is durable before we commit past it.
+ producer.flush()
+
+ async def aclose(self) -> None:
+ """Flush the DLQ producer and close the consumer when the poll ends; best-effort."""
+ producer = self._dlq_producer
+ if producer is not None:
+ self._dlq_producer = None
+ try:
+ await sync_to_async(producer.flush)()
+ except Exception:
+ log.warning("Failed to flush Kafka DLQ producer", exc_info=True)
+ consumer = self._consumer
+ if consumer is not None:
+ self._consumer = None
+ try:
+ await sync_to_async(consumer.close)()
+ except Exception:
+ log.warning("Failed to close Kafka consumer", exc_info=True)
+
+
+class KafkaSharedStreamTrigger(BaseEventTrigger):
+ """
+ Event-driven trigger that watches Kafka topics through a shared consumer.
+
+ Triggers that declare the same ``topics`` + ``kafka_config_id`` share one
+ underlying Kafka consumer in the triggerer (a single poll loop broadcast
+ to every subscriber). Each subscriber fires a ``TriggerEvent`` per
+ message; override :meth:`filter_shared_stream` to fire only for the
+ messages this trigger cares about.
+
+ Designed to back an :class:`~airflow.sdk.AssetWatcher` for event-driven
+ scheduling. The offset is committed only after the derived
+ ``TriggerEvent`` is persisted -- see :class:`KafkaSharedStreamProducer`
+ for the ``enable.auto.commit=false`` requirement.
+
+ :param topics: Topics to watch.
+ :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 messages a subscriber
+ rejects; see :class:`KafkaSharedStreamProducer`. 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:
+ super().__init__()
+ # Sort once here so shared_stream_key() is order-independent without
+ # re-sorting on every call: triggers on the same topics in any order
+ # share one consumer.
+ self.topics = sorted(topics)
+ self.kafka_config_id = kafka_config_id
+ self.poll_timeout = poll_timeout
+ self.dlq_topic = dlq_topic
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+ "airflow.providers.apache.kafka.triggers.shared_stream.KafkaSharedStreamTrigger",
+ {
+ "topics": self.topics,
+ "kafka_config_id": self.kafka_config_id,
+ "poll_timeout": self.poll_timeout,
+ "dlq_topic": self.dlq_topic,
+ },
+ )
+
+ def shared_stream_key(self) -> Hashable:
Review Comment:
```suggestion
def create_shared_stream_key(self) -> Hashable:
```
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,407 @@
+# 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.
+
+This module builds on the shared-stream ack channel, which was added in
+Airflow 3.3. Importing this module on an older version raises
+``AirflowOptionalProviderFeatureException``.
+"""
+
+from __future__ import annotations
+
+from airflow.providers.apache.kafka.version_compat import AIRFLOW_V_3_3_PLUS
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+if not AIRFLOW_V_3_3_PLUS:
+ raise AirflowOptionalProviderFeatureException("Kafka shared-stream triggers need Airflow 3.3+.")
+
+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 that 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 floor
+ is held and the consumer is sought back so the message is redelivered this
+ session: a failure to the remaining healthy subscribers, a zero-subscriber
+ broadcast to whoever subscribes next.
+
+ :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
+ # lane (topic, partition) -> earliest offset held for redelivery. An
+ # involuntary failure / zero-subscriber broadcast must not be committed
+ # past until it is reprocessed; the floor is carried across advance()
+ # calls so a later batch's cumulative commit cannot overtake it.
+ self._floor: dict[Hashable, int] = {}
+
+ 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 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()
+ 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
+
+ async 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.
+ """
+ connection = await hook.aget_connection(self.kafka_config_id)
+ auto_commit = connection.extra_dejson.get("enable.auto.commit", True)
+ 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 one partition's resolved prefix; hold and seek back the first held offset.
+
+ 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. Within a batch 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`` (ack timeout / overflow) or all-zero (a broadcast no
+ subscriber was online for) -- hold the floor here and ``seek`` the
+ consumer back to it, so it and everything after are redelivered this
+ session. A failure lands on the healthy subscribers the manager left
+ online; a zero-subscriber broadcast reaches whoever subscribes next,
+ instead of waiting for a full consumer rebuild. The floor lifts once a
+ re-read batch acks through the held offset.
+
+ The floor is per-lane and carried across calls (``self._floor``): once a
+ lane holds at offset ``H``, no later batch may commit past ``H`` until a
+ re-read batch starting at or before ``H`` acks through it.
+
+ 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
+ lane = self.get_advance_lane(first)
+ commit_through: int | None = None
+ held_offset: int | None = None
+ to_dlq: list[KafkaBrokerPayload] = []
+ for item in batch:
+ payload = item.broker_payload
+ outcome = item.outcome
+ if outcome.failed > 0 or (outcome.acked == 0 and outcome.rejected == 0):
+ # Held: an involuntary failure, or a broadcast no subscriber took.
+ # Seek back so it -- and everything after -- is redelivered this
+ # session, to the subscribers the manager left online or to
+ # whoever subscribes next.
+ held_offset = payload.offset
+ break
+ if outcome.rejected and self.dlq_topic is not None:
+ to_dlq.append(payload)
+ commit_through = payload.offset
+
+ floor = self._resolve_floor(lane, first.offset, commit_through, held_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_through is not None:
+ target = commit_through + 1
+ if floor is not None:
+ # floor may be a previous batch's held offset, so commit_through
+ # can exceed it. Use minimal value to avoid it.
+ target = min(target, floor)
+ await sync_to_async(self._commit)(first.topic, first.partition, target)
+ if held_offset is not None:
+ # Seek to the lane's floor (the earliest unresolved held offset), not
+ # this batch's break point: an earlier batch may hold a smaller offset,
+ # and seeking forward to a later one would skip and lose it.
+ seek_target = floor if floor is not None else held_offset
+ await sync_to_async(self._seek)(first.topic, first.partition, seek_target)
+
+ def _resolve_floor(
+ self, lane: Hashable, batch_start: int, commit_through: int | None, held_offset: int | None
+ ) -> int | None:
+ """
+ Update and return this lane's redelivery floor.
+
+ The floor is the earliest unresolved held offset on the lane: the minimum
+ of any still-unresolved previous floor and this batch's held offset.
+ A previous floor lifts only when this batch starts at or before it and acks
+ consecutively through it.
+ A later arrived ack before the seek operation, or a failure at a later offset can
+ neither clear nor advance a floor whose offset was never reprocessed.
+ """
+ previous = self._floor.get(lane)
+ cleared = (
+ previous is not None
+ and batch_start <= previous
+ and commit_through is not None
+ and commit_through >= previous
+ )
+ candidates: list[int] = []
+ if previous is not None and not cleared:
+ candidates.append(previous)
+ if held_offset is not None:
+ candidates.append(held_offset)
+ floor: int | None = min(candidates) if candidates else None
+ if floor is None:
+ self._floor.pop(lane, None)
+ else:
+ self._floor[lane] = floor
+ return floor
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ # When advance() raises, the triggerer fails every trigger that shares this
+ # producer, and Kafka redelivers from the last committed offset. Returning
+ # here would let the triggerer treat the batch as advanced although nothing
+ # was committed.
+ raise RuntimeError(f"Cannot commit {topic}[{partition}]@{offset}: no open consumer")
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
+
+ def _seek(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ raise RuntimeError(f"Cannot seek {topic}[{partition}] to {offset}: no open consumer")
+ self._consumer.seek(TopicPartition(topic, partition, offset))
+
+ def _dead_letter(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
Review Comment:
```suggestion
def _send_to_dlq(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3449656245
##########
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})."
+ )
Review Comment:
There's an async method `aget_connection` at base hook level. IIUC, we could avoid the extra `sync_to_async` call (since it will create a new thread under the hook, it quite _heavy_ to achieve the async purpose) for the early validation.
https://github.com/apache/airflow/blob/db9fa4289fe0d01e764b377c5c7f13d247af8218/task-sdk/src/airflow/sdk/bases/hook.py#L65-L77
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3449703986
##########
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)
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is not None:
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
Review Comment:
Though it _shouldn't_ happen. Would it be better to add warning log for the `self._consumer` is None case?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3453507566
##########
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:
No — the manager doesn't gate on "all triggers present," and it can't.
The subscriber set is open (triggers come and go across scheduler cycles), and `shared_stream_key` is computed lazily per trigger after template rendering.
While you're right, it's by design. The subscriber set is frozen at fan-out, and `is_clean` is correct *relative to the subscribers online when the event was broadcast* — that's the contract. A trigger not yet subscribed was never counted, so committing isn't wrong; at-least-once is scoped to who's listening, same as a Kafka consumer group resuming from its committed offset.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r3453538736
##########
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:
The one real problem is **triggerer restart**: the first watcher re-subscribes, the consumer churns the backlog, and watchers re-subscribing a moment later miss what committed in that window.
Some early thoughts — it can only be addressed in base, and it's not part of this PR.
Rough directions:
1. hold the first commit for a short configurable time after a group starts, so watchers that re-subscribe within that window are counted before anything commits;
2. or rely on idempotency + the documented restart-redelivery contract and add nothing.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r4155541410
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,398 @@
+# 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
Review Comment:
`airflow.triggers.shared_stream` is only available from Airflow 3.3.0, but the provider still declares `apache-airflow>=2.11.0`.
We'll need to make it 3.3+ only
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,398 @@
+# 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 that 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 floor
+ is held and the consumer is sought back so the message is redelivered this
+ session: a failure to the remaining healthy subscribers, a zero-subscriber
+ broadcast to whoever subscribes next.
+
+ :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
+ # lane (topic, partition) -> earliest offset held for redelivery. An
+ # involuntary failure / zero-subscriber broadcast must not be committed
+ # past until it is reprocessed; the floor is carried across advance()
+ # calls so a later batch's cumulative commit cannot overtake it.
+ self._floor: dict[Hashable, int] = {}
+
+ 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 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()
+ 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
+
+ async 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.
+ """
+ connection = await hook.aget_connection(self.kafka_config_id)
+ auto_commit = connection.extra_dejson.get("enable.auto.commit", True)
+ 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 one partition's resolved prefix; hold and seek back the first held offset.
+
+ 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. Within a batch 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`` (ack timeout / overflow) or all-zero (a broadcast no
+ subscriber was online for) -- hold the floor here and ``seek`` the
+ consumer back to it, so it and everything after are redelivered this
+ session. A failure lands on the healthy subscribers the manager left
+ online; a zero-subscriber broadcast reaches whoever subscribes next,
+ instead of waiting for a full consumer rebuild. The floor lifts once a
+ re-read batch acks through the held offset.
+
+ The floor is per-lane and carried across calls (``self._floor``): once a
+ lane holds at offset ``H``, no later batch may commit past ``H`` until a
+ re-read batch starting at or before ``H`` acks through it.
+
+ 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
+ lane = self.get_advance_lane(first)
+ commit_through: int | None = None
+ held_offset: int | None = None
+ to_dlq: list[KafkaBrokerPayload] = []
+ for item in batch:
+ payload = item.broker_payload
+ outcome = item.outcome
+ if outcome.failed > 0 or (outcome.acked == 0 and outcome.rejected == 0):
+ # Held: an involuntary failure, or a broadcast no subscriber took.
+ # Seek back so it -- and everything after -- is redelivered this
+ # session, to the subscribers the manager left online or to
+ # whoever subscribes next.
+ held_offset = payload.offset
+ break
+ if outcome.rejected and self.dlq_topic is not None:
+ to_dlq.append(payload)
+ commit_through = payload.offset
+
+ floor = self._resolve_floor(lane, first.offset, commit_through, held_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_through is not None:
+ target = commit_through + 1
+ if floor is not None:
+ # floor may be a previous batch's held offset, so commit_through
+ # can exceed it. Use minimal value to avoid it.
+ target = min(target, floor)
+ await sync_to_async(self._commit)(first.topic, first.partition, target)
+ if held_offset is not None:
+ # Seek to the lane's floor (the earliest unresolved held offset), not
+ # this batch's break point: an earlier batch may hold a smaller offset,
+ # and seeking forward to a later one would skip and lose it.
+ seek_target = floor if floor is not None else held_offset
+ await sync_to_async(self._seek)(first.topic, first.partition, seek_target)
+
+ def _resolve_floor(
+ self, lane: Hashable, batch_start: int, commit_through: int | None, held_offset: int | None
+ ) -> int | None:
+ """
+ Update and return this lane's redelivery floor.
+
+ The floor is the earliest unresolved held offset on the lane: the minimum
+ of any still-unresolved previous floor and this batch's held offset.
+ A previous floor lifts only when this batch starts at or before it and acks
+ consecutively through it.
+ A later arrived ack before the seek operation, or a failure at a later offset can
+ neither clear nor advance a floor whose offset was never reprocessed.
+ """
+ previous = self._floor.get(lane)
+ cleared = (
+ previous is not None
+ and batch_start <= previous
+ and commit_through is not None
+ and commit_through >= previous
+ )
+ candidates: list[int] = []
+ if previous is not None and not cleared:
+ candidates.append(previous)
+ if held_offset is not None:
+ candidates.append(held_offset)
+ floor: int | None = min(candidates) if candidates else None
+ if floor is None:
+ self._floor.pop(lane, None)
+ else:
+ self._floor[lane] = floor
+ return floor
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
Review Comment:
non-blocking
since this "shouldn't happen", would raising be safer than warning and returning?
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,398 @@
+# 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 that 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 floor
+ is held and the consumer is sought back so the message is redelivered this
+ session: a failure to the remaining healthy subscribers, a zero-subscriber
+ broadcast to whoever subscribes next.
+
+ :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
+ # lane (topic, partition) -> earliest offset held for redelivery. An
+ # involuntary failure / zero-subscriber broadcast must not be committed
+ # past until it is reprocessed; the floor is carried across advance()
+ # calls so a later batch's cumulative commit cannot overtake it.
+ self._floor: dict[Hashable, int] = {}
+
+ 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 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()
+ 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
+
+ async 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.
+ """
+ connection = await hook.aget_connection(self.kafka_config_id)
+ auto_commit = connection.extra_dejson.get("enable.auto.commit", True)
+ 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 one partition's resolved prefix; hold and seek back the first held offset.
+
+ 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. Within a batch 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`` (ack timeout / overflow) or all-zero (a broadcast no
+ subscriber was online for) -- hold the floor here and ``seek`` the
+ consumer back to it, so it and everything after are redelivered this
+ session. A failure lands on the healthy subscribers the manager left
+ online; a zero-subscriber broadcast reaches whoever subscribes next,
+ instead of waiting for a full consumer rebuild. The floor lifts once a
+ re-read batch acks through the held offset.
+
+ The floor is per-lane and carried across calls (``self._floor``): once a
+ lane holds at offset ``H``, no later batch may commit past ``H`` until a
+ re-read batch starting at or before ``H`` acks through it.
+
+ 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
+ lane = self.get_advance_lane(first)
+ commit_through: int | None = None
+ held_offset: int | None = None
+ to_dlq: list[KafkaBrokerPayload] = []
+ for item in batch:
+ payload = item.broker_payload
+ outcome = item.outcome
+ if outcome.failed > 0 or (outcome.acked == 0 and outcome.rejected == 0):
+ # Held: an involuntary failure, or a broadcast no subscriber took.
+ # Seek back so it -- and everything after -- is redelivered this
+ # session, to the subscribers the manager left online or to
+ # whoever subscribes next.
+ held_offset = payload.offset
+ break
+ if outcome.rejected and self.dlq_topic is not None:
+ to_dlq.append(payload)
+ commit_through = payload.offset
+
+ floor = self._resolve_floor(lane, first.offset, commit_through, held_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_through is not None:
+ target = commit_through + 1
+ if floor is not None:
+ # floor may be a previous batch's held offset, so commit_through
+ # can exceed it. Use minimal value to avoid it.
+ target = min(target, floor)
+ await sync_to_async(self._commit)(first.topic, first.partition, target)
+ if held_offset is not None:
+ # Seek to the lane's floor (the earliest unresolved held offset), not
+ # this batch's break point: an earlier batch may hold a smaller offset,
+ # and seeking forward to a later one would skip and lose it.
+ seek_target = floor if floor is not None else held_offset
+ await sync_to_async(self._seek)(first.topic, first.partition, seek_target)
+
+ def _resolve_floor(
+ self, lane: Hashable, batch_start: int, commit_through: int | None, held_offset: int | None
+ ) -> int | None:
+ """
+ Update and return this lane's redelivery floor.
+
+ The floor is the earliest unresolved held offset on the lane: the minimum
+ of any still-unresolved previous floor and this batch's held offset.
+ A previous floor lifts only when this batch starts at or before it and acks
+ consecutively through it.
+ A later arrived ack before the seek operation, or a failure at a later offset can
+ neither clear nor advance a floor whose offset was never reprocessed.
+ """
+ previous = self._floor.get(lane)
+ cleared = (
+ previous is not None
+ and batch_start <= previous
+ and commit_through is not None
+ and commit_through >= previous
+ )
+ candidates: list[int] = []
+ if previous is not None and not cleared:
+ candidates.append(previous)
+ if held_offset is not None:
+ candidates.append(held_offset)
+ floor: int | None = min(candidates) if candidates else None
+ if floor is None:
+ self._floor.pop(lane, None)
+ else:
+ self._floor[lane] = floor
+ return floor
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ log.warning("Cannot commit %s[%d]@%d: no open consumer", topic, partition, offset)
+ return
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
+
+ def _seek(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ log.warning("Cannot seek %s[%d] to %d: no open consumer", topic, partition, offset)
+ return
+ self._consumer.seek(TopicPartition(topic, partition, offset))
+
+ def _dead_letter(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
+ if self._dlq_producer is None:
+ self._dlq_producer = KafkaProducerHook(kafka_config_id=self.kafka_config_id).get_producer()
+ producer = self._dlq_producer
+ for payload in payloads:
+ producer.produce(topic, value=payload.value, key=payload.key)
+ # Flush once so the whole dead-letter batch is durable before we commit past it.
+ producer.flush()
+
+ async def aclose(self) -> None:
+ """Flush the DLQ producer and close the consumer when the poll ends; best-effort."""
+ producer = self._dlq_producer
+ if producer is not None:
+ self._dlq_producer = None
+ try:
+ await sync_to_async(producer.flush)()
+ except Exception:
+ log.warning("Failed to flush Kafka DLQ producer", exc_info=True)
+ consumer = self._consumer
+ if consumer is not None:
+ self._consumer = None
+ try:
+ await sync_to_async(consumer.close)()
+ except Exception:
+ log.warning("Failed to close Kafka consumer", exc_info=True)
+
+
+class KafkaSharedStreamTrigger(BaseEventTrigger):
+ """
+ Event-driven trigger that watches Kafka topics through a shared consumer.
+
+ Triggers that declare the same ``topics`` + ``kafka_config_id`` share one
+ underlying Kafka consumer in the triggerer (a single poll loop broadcast
+ to every subscriber). Each subscriber fires a ``TriggerEvent`` per
+ message; override :meth:`filter_shared_stream` to fire only for the
+ messages this trigger cares about.
+
+ Designed to back an :class:`~airflow.sdk.AssetWatcher` for event-driven
+ scheduling. The offset is committed only after the derived
+ ``TriggerEvent`` is persisted -- see :class:`KafkaSharedStreamProducer`
+ for the ``enable.auto.commit=false`` requirement.
+
+ :param topics: Topics to watch.
+ :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 messages a subscriber
+ rejects; see :class:`KafkaSharedStreamProducer`. 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:
+ super().__init__()
+ # Sort once here so shared_stream_key() is order-independent without
+ # re-sorting on every call: triggers on the same topics in any order
+ # share one consumer.
+ self.topics = sorted(topics)
+ self.kafka_config_id = kafka_config_id
+ self.poll_timeout = poll_timeout
+ self.dlq_topic = dlq_topic
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+ "airflow.providers.apache.kafka.triggers.shared_stream.KafkaSharedStreamTrigger",
+ {
+ "topics": self.topics,
+ "kafka_config_id": self.kafka_config_id,
+ "poll_timeout": self.poll_timeout,
+ "dlq_topic": self.dlq_topic,
+ },
+ )
+
+ def shared_stream_key(self) -> Hashable:
+ """Triggers on the same topics + connection share one consumer."""
+ return "kafka-shared-stream", tuple(self.topics), self.kafka_config_id
Review Comment:
Dropping `self.kafka_config_id` from this key and all the unit tests still pass, since both triggers in the key tests use `kafka_config_id="c"`. Since this decides whether triggers on different connections share a consumer, could we add a case where only `kafka_config_id` differs and the keys must differ?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #68625:
URL: https://github.com/apache/airflow/pull/68625#discussion_r4162075257
##########
providers/apache/kafka/src/airflow/providers/apache/kafka/triggers/shared_stream.py:
##########
@@ -0,0 +1,398 @@
+# 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 that 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 floor
+ is held and the consumer is sought back so the message is redelivered this
+ session: a failure to the remaining healthy subscribers, a zero-subscriber
+ broadcast to whoever subscribes next.
+
+ :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
+ # lane (topic, partition) -> earliest offset held for redelivery. An
+ # involuntary failure / zero-subscriber broadcast must not be committed
+ # past until it is reprocessed; the floor is carried across advance()
+ # calls so a later batch's cumulative commit cannot overtake it.
+ self._floor: dict[Hashable, int] = {}
+
+ 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 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()
+ 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
+
+ async 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.
+ """
+ connection = await hook.aget_connection(self.kafka_config_id)
+ auto_commit = connection.extra_dejson.get("enable.auto.commit", True)
+ 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 one partition's resolved prefix; hold and seek back the first held offset.
+
+ 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. Within a batch 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`` (ack timeout / overflow) or all-zero (a broadcast no
+ subscriber was online for) -- hold the floor here and ``seek`` the
+ consumer back to it, so it and everything after are redelivered this
+ session. A failure lands on the healthy subscribers the manager left
+ online; a zero-subscriber broadcast reaches whoever subscribes next,
+ instead of waiting for a full consumer rebuild. The floor lifts once a
+ re-read batch acks through the held offset.
+
+ The floor is per-lane and carried across calls (``self._floor``): once a
+ lane holds at offset ``H``, no later batch may commit past ``H`` until a
+ re-read batch starting at or before ``H`` acks through it.
+
+ 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
+ lane = self.get_advance_lane(first)
+ commit_through: int | None = None
+ held_offset: int | None = None
+ to_dlq: list[KafkaBrokerPayload] = []
+ for item in batch:
+ payload = item.broker_payload
+ outcome = item.outcome
+ if outcome.failed > 0 or (outcome.acked == 0 and outcome.rejected == 0):
+ # Held: an involuntary failure, or a broadcast no subscriber took.
+ # Seek back so it -- and everything after -- is redelivered this
+ # session, to the subscribers the manager left online or to
+ # whoever subscribes next.
+ held_offset = payload.offset
+ break
+ if outcome.rejected and self.dlq_topic is not None:
+ to_dlq.append(payload)
+ commit_through = payload.offset
+
+ floor = self._resolve_floor(lane, first.offset, commit_through, held_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_through is not None:
+ target = commit_through + 1
+ if floor is not None:
+ # floor may be a previous batch's held offset, so commit_through
+ # can exceed it. Use minimal value to avoid it.
+ target = min(target, floor)
+ await sync_to_async(self._commit)(first.topic, first.partition, target)
+ if held_offset is not None:
+ # Seek to the lane's floor (the earliest unresolved held offset), not
+ # this batch's break point: an earlier batch may hold a smaller offset,
+ # and seeking forward to a later one would skip and lose it.
+ seek_target = floor if floor is not None else held_offset
+ await sync_to_async(self._seek)(first.topic, first.partition, seek_target)
+
+ def _resolve_floor(
+ self, lane: Hashable, batch_start: int, commit_through: int | None, held_offset: int | None
+ ) -> int | None:
+ """
+ Update and return this lane's redelivery floor.
+
+ The floor is the earliest unresolved held offset on the lane: the minimum
+ of any still-unresolved previous floor and this batch's held offset.
+ A previous floor lifts only when this batch starts at or before it and acks
+ consecutively through it.
+ A later arrived ack before the seek operation, or a failure at a later offset can
+ neither clear nor advance a floor whose offset was never reprocessed.
+ """
+ previous = self._floor.get(lane)
+ cleared = (
+ previous is not None
+ and batch_start <= previous
+ and commit_through is not None
+ and commit_through >= previous
+ )
+ candidates: list[int] = []
+ if previous is not None and not cleared:
+ candidates.append(previous)
+ if held_offset is not None:
+ candidates.append(held_offset)
+ floor: int | None = min(candidates) if candidates else None
+ if floor is None:
+ self._floor.pop(lane, None)
+ else:
+ self._floor[lane] = floor
+ return floor
+
+ def _commit(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ log.warning("Cannot commit %s[%d]@%d: no open consumer", topic, partition, offset)
+ return
+ self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
+
+ def _seek(self, topic: str, partition: int, offset: int) -> None:
+ if self._consumer is None:
+ log.warning("Cannot seek %s[%d] to %d: no open consumer", topic, partition, offset)
+ return
+ self._consumer.seek(TopicPartition(topic, partition, offset))
+
+ def _dead_letter(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
+ if self._dlq_producer is None:
+ self._dlq_producer = KafkaProducerHook(kafka_config_id=self.kafka_config_id).get_producer()
+ producer = self._dlq_producer
+ for payload in payloads:
+ producer.produce(topic, value=payload.value, key=payload.key)
+ # Flush once so the whole dead-letter batch is durable before we commit past it.
+ producer.flush()
+
+ async def aclose(self) -> None:
+ """Flush the DLQ producer and close the consumer when the poll ends; best-effort."""
+ producer = self._dlq_producer
+ if producer is not None:
+ self._dlq_producer = None
+ try:
+ await sync_to_async(producer.flush)()
+ except Exception:
+ log.warning("Failed to flush Kafka DLQ producer", exc_info=True)
+ consumer = self._consumer
+ if consumer is not None:
+ self._consumer = None
+ try:
+ await sync_to_async(consumer.close)()
+ except Exception:
+ log.warning("Failed to close Kafka consumer", exc_info=True)
+
+
+class KafkaSharedStreamTrigger(BaseEventTrigger):
+ """
+ Event-driven trigger that watches Kafka topics through a shared consumer.
+
+ Triggers that declare the same ``topics`` + ``kafka_config_id`` share one
+ underlying Kafka consumer in the triggerer (a single poll loop broadcast
+ to every subscriber). Each subscriber fires a ``TriggerEvent`` per
+ message; override :meth:`filter_shared_stream` to fire only for the
+ messages this trigger cares about.
+
+ Designed to back an :class:`~airflow.sdk.AssetWatcher` for event-driven
+ scheduling. The offset is committed only after the derived
+ ``TriggerEvent`` is persisted -- see :class:`KafkaSharedStreamProducer`
+ for the ``enable.auto.commit=false`` requirement.
+
+ :param topics: Topics to watch.
+ :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 messages a subscriber
+ rejects; see :class:`KafkaSharedStreamProducer`. 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:
+ super().__init__()
+ # Sort once here so shared_stream_key() is order-independent without
+ # re-sorting on every call: triggers on the same topics in any order
+ # share one consumer.
+ self.topics = sorted(topics)
+ self.kafka_config_id = kafka_config_id
+ self.poll_timeout = poll_timeout
+ self.dlq_topic = dlq_topic
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+ "airflow.providers.apache.kafka.triggers.shared_stream.KafkaSharedStreamTrigger",
+ {
+ "topics": self.topics,
+ "kafka_config_id": self.kafka_config_id,
+ "poll_timeout": self.poll_timeout,
+ "dlq_topic": self.dlq_topic,
+ },
+ )
+
+ def shared_stream_key(self) -> Hashable:
+ """Triggers on the same topics + connection share one consumer."""
+ return "kafka-shared-stream", tuple(self.topics), self.kafka_config_id
Review Comment:
Add `test_shared_stream_key_separates_topics_and_connections` for it. 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on PR #68625:
URL: https://github.com/apache/airflow/pull/68625#issuecomment-5948018491
Thanks for the review. Add `KafkaSharedStreamTrigger` to triggers.rst.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add KafkaSharedStreamProducer and KafkaSharedStreamTrigger [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W merged PR #68625:
URL: https://github.com/apache/airflow/pull/68625
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org