Skip to content

Commit 1613ce3

Browse files
committed
Add shared-stream producer and trigger to Apache Kafka provider
Event-driven triggers and asset watchers that watch the same Kafka topics each open their own consumer in the triggerer, which does not scale as event-driven schedules grow. Backing them with the shared-stream producer-side ack channel (#67523) lets siblings on the same topics and connection share one consumer, while offset commits remain gated on trigger-event persistence so a triggerer crash cannot drop a message the broker already considers delivered (at-least-once). Signed-off-by: PoAn Yang <payang@apache.org>
1 parent 7e7e4a9 commit 1613ce3

6 files changed

Lines changed: 867 additions & 0 deletions

File tree

‎providers/apache/kafka/provider.yaml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,7 @@ triggers:
110110
python-modules:
111111
- airflow.providers.apache.kafka.triggers.await_message
112112
- airflow.providers.apache.kafka.triggers.msg_queue
113+
- airflow.providers.apache.kafka.triggers.shared_stream
113114

114115
connection-types:
115116
- hook-class-name: airflow.providers.apache.kafka.hooks.base.KafkaBaseHook

‎providers/apache/kafka/src/airflow/providers/apache/kafka/get_provider_info.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ def get_provider_info():
8383
"python-modules": [
8484
"airflow.providers.apache.kafka.triggers.await_message",
8585
"airflow.providers.apache.kafka.triggers.msg_queue",
86+
"airflow.providers.apache.kafka.triggers.shared_stream",
8687
],
8788
}
8889
],
Lines changed: 327 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,327 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""
18+
Shared-stream Kafka trigger and producer for event-driven scheduling.
19+
20+
Triggers that declare the same ``topics`` + ``kafka_config_id`` share a
21+
single Kafka consumer in the triggerer (one poll loop broadcast to every
22+
subscriber) instead of opening one consumer each.
23+
The :class:`KafkaSharedStreamProducer` owns that consumer and commits
24+
offsets only after the derived :class:`~airflow.triggers.base.TriggerEvent`
25+
instances have been persisted, via the shared-stream ack channel.
26+
27+
Requires an Airflow version whose ``airflow.triggers.shared_stream`` module
28+
provides the producer-side ack channel.
29+
"""
30+
31+
from __future__ import annotations
32+
33+
import logging
34+
from collections.abc import AsyncIterator, Hashable, Sequence
35+
from typing import TYPE_CHECKING, Any, NamedTuple, cast
36+
37+
from asgiref.sync import sync_to_async
38+
from confluent_kafka import TopicPartition
39+
40+
from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
41+
from airflow.providers.apache.kafka.hooks.produce import KafkaProducerHook
42+
from airflow.triggers.base import BaseEventTrigger, TriggerEvent
43+
from airflow.triggers.shared_stream import AdvanceItem, SharedStreamProducer
44+
45+
if TYPE_CHECKING:
46+
from confluent_kafka import Consumer, Producer
47+
48+
log = logging.getLogger(__name__)
49+
50+
51+
class KafkaBrokerPayload(NamedTuple):
52+
"""
53+
The ``broker_payload`` carried alongside each raw event from ``open_stream``.
54+
55+
``(topic, partition, offset)`` let :meth:`KafkaSharedStreamProducer.advance`
56+
commit the right ``TopicPartition`` and
57+
:meth:`KafkaSharedStreamProducer.get_advance_lane` key the advance by
58+
partition. ``value`` / ``key`` are populated only when a ``dlq_topic`` is
59+
configured -- they are what a dead-letter send re-publishes, so retaining
60+
them through the whole outstanding window is pointless without a DLQ.
61+
"""
62+
63+
topic: str
64+
partition: int
65+
offset: int
66+
value: Any = None
67+
key: Any = None
68+
69+
70+
class KafkaSharedStreamProducer(SharedStreamProducer):
71+
"""
72+
Broker-side half of a shared Kafka stream running in ack mode.
73+
74+
Drives one confluent-kafka ``Consumer`` for a shared-stream group and
75+
commits offsets only after every subscriber that derived a
76+
``TriggerEvent`` from a message has had it persisted -- the ack channel
77+
gates the commit, so a triggerer crash cannot drop a message the broker
78+
already considers delivered.
79+
80+
.. warning::
81+
82+
The Kafka connection used by ``kafka_config_id`` **must** set
83+
``enable.auto.commit=false``. With auto-commit on, the consumer
84+
commits offsets on its own schedule regardless of the ack channel,
85+
which defeats the persistence gate and loses messages on a triggerer
86+
crash.
87+
88+
Events a subscriber **rejects** (via ``reject_shared_stream_event``) are
89+
terminal: by default they are dropped (committed past). Set ``dlq_topic``
90+
to instead re-publish each rejected message to a dead-letter topic before
91+
committing past it. Involuntary failures (ack timeout / overflow) and
92+
broadcasts no subscriber was online for are never dropped -- the offset is
93+
held so Kafka redelivers them.
94+
95+
:param topics: Topics the shared consumer subscribes to.
96+
:param kafka_config_id: Kafka connection id, defaults to ``kafka_default``.
97+
:param poll_timeout: Seconds the consumer waits on each ``poll`` call.
98+
:param dlq_topic: Optional dead-letter topic for rejected messages, produced
99+
through ``kafka_config_id``. When unset, rejected messages are dropped.
100+
"""
101+
102+
def __init__(
103+
self,
104+
*,
105+
topics: Sequence[str],
106+
kafka_config_id: str = "kafka_default",
107+
poll_timeout: float = 1.0,
108+
dlq_topic: str | None = None,
109+
) -> None:
110+
self.topics = list(topics)
111+
self.kafka_config_id = kafka_config_id
112+
self.poll_timeout = poll_timeout
113+
self.dlq_topic = dlq_topic
114+
self._consumer: Consumer | None = None
115+
self._dlq_producer: Producer | None = None
116+
117+
async def open_stream(self) -> AsyncIterator[tuple[Any, KafkaBrokerPayload]]:
118+
"""Open the consumer lazily and yield (value, KafkaBrokerPayload) per message."""
119+
hook = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id)
120+
await sync_to_async(self._ensure_manual_commit)(hook)
121+
consumer = await sync_to_async(hook.get_consumer)()
122+
self._consumer = consumer
123+
poll = sync_to_async(consumer.poll)
124+
keep_for_dlq = self.dlq_topic is not None
125+
while True:
126+
message = await poll(self.poll_timeout)
127+
if message is None:
128+
continue
129+
if message.error():
130+
raise RuntimeError(f"Kafka consumer error: {message.error()}")
131+
value = message.value()
132+
# Retain value/key only with a DLQ: the broker payload outlives the raw
133+
# event (it is held until the offset is advanced), so keeping the message
134+
# body when it can never be dead-lettered just wastes memory.
135+
payload = KafkaBrokerPayload(
136+
topic=cast("str", message.topic()),
137+
partition=cast("int", message.partition()),
138+
offset=cast("int", message.offset()),
139+
value=value if keep_for_dlq else None,
140+
key=message.key() if keep_for_dlq else None,
141+
)
142+
yield value, payload
143+
144+
def _ensure_manual_commit(self, hook: KafkaConsumerHook) -> None:
145+
"""
146+
Refuse to start unless the connection disables Kafka auto-commit.
147+
148+
The shared-stream ack channel owns offset commits -- it commits only
149+
after the derived trigger events are persisted. ``enable.auto.commit``
150+
defaults to ``true`` in Kafka; left on, the consumer commits on its own
151+
schedule regardless of the ack channel and silently drops messages on a
152+
triggerer crash. Fail fast rather than run in a lossy configuration.
153+
"""
154+
auto_commit = hook.get_connection(self.kafka_config_id).extra_dejson.get("enable.auto.commit", True)
155+
# Accept only the idiomatic ``false`` -- both the string "false" and the JSON
156+
# bool ``false`` normalise to "false"; anything else (true, 1, 0, unset, ...) is
157+
# rejected so an ambiguous value never silently runs in a lossy configuration.
158+
if str(auto_commit).strip().lower() != "false":
159+
raise ValueError(
160+
f"KafkaSharedStreamProducer requires enable.auto.commit=false in the "
161+
f"{self.kafka_config_id!r} Kafka connection; the shared-stream ack channel "
162+
f"commits offsets only after trigger events are persisted "
163+
f"(got enable.auto.commit={auto_commit!r})."
164+
)
165+
166+
def get_advance_lane(self, broker_payload: KafkaBrokerPayload) -> Hashable:
167+
"""Order commits per (topic, partition)."""
168+
return broker_payload.topic, broker_payload.partition
169+
170+
async def advance(self, batch: Sequence[AdvanceItem]) -> None:
171+
"""
172+
Commit the longest terminally-resolved prefix of one partition's batch.
173+
174+
Every item shares one ``(topic, partition)`` (the lane). A Kafka commit
175+
is cumulative, so committing offset ``N + 1`` marks everything up to
176+
``N`` as consumed. We commit only through the last item that was
177+
*terminally handled* and stop at the first that should come back:
178+
179+
* ``acked`` -- accepted; safe to commit past.
180+
* ``rejected`` -- terminally refused. If ``dlq_topic`` is set the message
181+
is re-published there (and flushed) before being committed past;
182+
otherwise it is dropped. Either way it is not redelivered.
183+
* ``failed`` -- an involuntary failure (ack timeout / overflow). Stop
184+
before it so Kafka redelivers it and everything after; never drop.
185+
* all-zero (a broadcast no subscriber was online for) -- nothing was
186+
accepted, so stop before it and let Kafka redeliver when a subscriber
187+
returns.
188+
189+
The batch's rejected messages are produced to the DLQ together and
190+
flushed once, before the offset is committed, so a crash cannot commit
191+
past a message that never reached the DLQ. If the commit later fails the
192+
batch is redelivered, which may re-send an already-dead-lettered message
193+
-- DLQ consumers should tolerate duplicates.
194+
"""
195+
first = batch[0].broker_payload
196+
commit_offset: int | None = None
197+
to_dlq: list[KafkaBrokerPayload] = []
198+
for item in batch:
199+
payload = item.broker_payload
200+
outcome = item.outcome
201+
if outcome.failed or (outcome.acked == 0 and outcome.rejected == 0):
202+
# Involuntary failure, or nothing accepted -- redeliver from here on.
203+
break
204+
if outcome.rejected and self.dlq_topic is not None:
205+
to_dlq.append(payload)
206+
commit_offset = payload.offset
207+
if to_dlq and self.dlq_topic is not None:
208+
await sync_to_async(self._dead_letter)(self.dlq_topic, to_dlq)
209+
if commit_offset is not None:
210+
await sync_to_async(self._commit)(first.topic, first.partition, commit_offset + 1)
211+
212+
def _commit(self, topic: str, partition: int, offset: int) -> None:
213+
if self._consumer is not None:
214+
self._consumer.commit(offsets=[TopicPartition(topic, partition, offset)], asynchronous=False)
215+
216+
def _dead_letter(self, topic: str, payloads: list[KafkaBrokerPayload]) -> None:
217+
if self._dlq_producer is None:
218+
self._dlq_producer = KafkaProducerHook(kafka_config_id=self.kafka_config_id).get_producer()
219+
producer = self._dlq_producer
220+
for payload in payloads:
221+
producer.produce(topic, value=payload.value, key=payload.key)
222+
# Flush once so the whole dead-letter batch is durable before we commit past it.
223+
producer.flush()
224+
225+
async def aclose(self) -> None:
226+
"""Flush the DLQ producer and close the consumer when the poll ends; best-effort."""
227+
producer = self._dlq_producer
228+
if producer is not None:
229+
self._dlq_producer = None
230+
try:
231+
await sync_to_async(producer.flush)()
232+
except Exception:
233+
log.warning("Failed to flush Kafka DLQ producer", exc_info=True)
234+
consumer = self._consumer
235+
if consumer is not None:
236+
self._consumer = None
237+
try:
238+
await sync_to_async(consumer.close)()
239+
except Exception:
240+
log.warning("Failed to close Kafka consumer", exc_info=True)
241+
242+
243+
class KafkaSharedStreamTrigger(BaseEventTrigger):
244+
"""
245+
Event-driven trigger that watches Kafka topics through a shared consumer.
246+
247+
Triggers that declare the same ``topics`` + ``kafka_config_id`` share one
248+
underlying Kafka consumer in the triggerer (a single poll loop broadcast
249+
to every subscriber). Each subscriber fires a ``TriggerEvent`` per
250+
message; override :meth:`filter_shared_stream` to fire only for the
251+
messages this trigger cares about.
252+
253+
Designed to back an :class:`~airflow.sdk.AssetWatcher` for event-driven
254+
scheduling. The offset is committed only after the derived
255+
``TriggerEvent`` is persisted -- see :class:`KafkaSharedStreamProducer`
256+
for the ``enable.auto.commit=false`` requirement.
257+
258+
:param topics: Topics to watch.
259+
:param kafka_config_id: Kafka connection id, defaults to ``kafka_default``.
260+
:param poll_timeout: Seconds the consumer waits on each ``poll`` call.
261+
:param dlq_topic: Optional dead-letter topic for messages a subscriber
262+
rejects; see :class:`KafkaSharedStreamProducer`. When unset, rejected
263+
messages are dropped.
264+
"""
265+
266+
def __init__(
267+
self,
268+
*,
269+
topics: Sequence[str],
270+
kafka_config_id: str = "kafka_default",
271+
poll_timeout: float = 1.0,
272+
dlq_topic: str | None = None,
273+
) -> None:
274+
super().__init__()
275+
self.topics = list(topics)
276+
self.kafka_config_id = kafka_config_id
277+
self.poll_timeout = poll_timeout
278+
self.dlq_topic = dlq_topic
279+
280+
def serialize(self) -> tuple[str, dict[str, Any]]:
281+
return (
282+
"airflow.providers.apache.kafka.triggers.shared_stream.KafkaSharedStreamTrigger",
283+
{
284+
"topics": self.topics,
285+
"kafka_config_id": self.kafka_config_id,
286+
"poll_timeout": self.poll_timeout,
287+
"dlq_topic": self.dlq_topic,
288+
},
289+
)
290+
291+
def shared_stream_key(self) -> Hashable:
292+
"""Triggers on the same topics + connection share one consumer."""
293+
return ("kafka-shared-stream", tuple(sorted(self.topics)), self.kafka_config_id)
294+
295+
@classmethod
296+
def create_shared_stream_producer(cls, kwargs: dict[str, Any]) -> KafkaSharedStreamProducer:
297+
return KafkaSharedStreamProducer(
298+
topics=kwargs["topics"],
299+
kafka_config_id=kwargs["kafka_config_id"],
300+
poll_timeout=kwargs["poll_timeout"],
301+
dlq_topic=kwargs.get("dlq_topic"),
302+
)
303+
304+
async def filter_shared_stream(self, shared_stream: AsyncIterator[Any]) -> AsyncIterator[TriggerEvent]:
305+
"""Fire one ``TriggerEvent`` per message. Override to filter or transform."""
306+
async for value in shared_stream:
307+
yield TriggerEvent(self._decode(value))
308+
309+
async def run(self) -> AsyncIterator[TriggerEvent]:
310+
"""
311+
Not supported -- this trigger runs only through the shared-stream manager.
312+
313+
``shared_stream_key`` always returns non-``None``, so the triggerer drives
314+
this trigger through :meth:`filter_shared_stream`; the ``_SharedStreamGroup``
315+
owns the Kafka consumer and offset commits. There is no standalone path:
316+
committing offsets safely needs the ack channel to gate them on
317+
trigger-event persistence, which only the manager provides.
318+
"""
319+
raise NotImplementedError(
320+
"KafkaSharedStreamTrigger runs only through the triggerer's shared-stream "
321+
"manager (via filter_shared_stream); it has no standalone run() path."
322+
)
323+
yield # pragma: no cover - marks this as an async generator
324+
325+
@staticmethod
326+
def _decode(value: Any) -> Any:
327+
return value.decode("utf-8") if isinstance(value, bytes) else value

0 commit comments

Comments
 (0)