This is an automated email from the ASF dual-hosted git repository.
je-ik pushed a commit to branch feat/18479-kafka-streams-runner-skeleton
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to
refs/heads/feat/18479-kafka-streams-runner-skeleton by this push:
new 5861f31e8ac [GSoC 2026] Kafka Streams runner: read unbounded sources
(#39611)
5861f31e8ac is described below
commit 5861f31e8ac33749ea67061cd83ed213552c19ea
Author: M Junaid Shaukat <[email protected]>
AuthorDate: Tue Aug 4 21:09:06 2026 +0500
[GSoC 2026] Kafka Streams runner: read unbounded sources (#39611)
[GSoC 2026] Kafka Streams runner: read unbounded sources
---
.../kafka/streams/KafkaStreamsPipelineOptions.java | 10 +
.../kafka/streams/translation/ReadTranslator.java | 123 ++++++++-
.../translation/UnboundedReadProcessor.java | 304 +++++++++++++++++++++
.../streams/translation/UnboundedReadTest.java | 209 ++++++++++++++
4 files changed, 635 insertions(+), 11 deletions(-)
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java
index e95268ac130..44abc8e5b34 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java
@@ -77,6 +77,16 @@ public interface KafkaStreamsPipelineOptions extends
PortablePipelineOptions {
void setTopicReplicationFactor(short topicReplicationFactor);
+ @Description(
+ "How many non-empty polls of an unbounded source to make before storing
its checkpoint mark."
+ + " Taking a mark can be costly for some sources, so it is not worth
doing on every poll;"
+ + " the cost of a larger value is that more elements are replayed
after a restart, since"
+ + " the reader resumes from the last mark that was stored.")
+ @Default.Integer(10)
+ int getReadCheckpointNumBundles();
+
+ void setReadCheckpointNumBundles(int readCheckpointNumBundles);
+
@Description("Directory where Kafka Streams stores local state.")
@Default.InstanceFactory(StateDirDefaultFactory.class)
String getStateDir();
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java
index 403499ca616..f83442f9781 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java
@@ -18,12 +18,14 @@
package org.apache.beam.runners.kafka.streams.translation;
import java.io.IOException;
+import java.util.List;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import
org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.WireCoderSetting;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.fnexecution.wire.WireCoders;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.BoundedSource;
+import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.util.construction.ReadTranslation;
import org.apache.beam.sdk.util.construction.RehydratedComponents;
import org.apache.beam.sdk.util.construction.graph.PipelineNode;
@@ -77,25 +79,92 @@ class ReadTranslator implements PTransformTranslator {
// Read produces exactly one output PCollection; downstream consumers are
separate PTransforms
// whose inputs reference this PCollection id and are wired by their own
translators.
String outputPCollectionId =
Iterables.getOnlyElement(transform.getOutputsMap().values());
- addReadNodes(
- transformId,
- boundedSource(transform),
- pipeline.getComponents(),
- outputPCollectionId,
- context);
- }
-
- private static BoundedSource<?> boundedSource(RunnerApi.PTransform
transform) {
try {
RunnerApi.ReadPayload payload =
RunnerApi.ReadPayload.parseFrom(transform.getSpec().getPayload());
- return ReadTranslation.boundedSourceFromProto(payload);
+ // The same URN carries both kinds of source; the payload says which,
and they need different
+ // processors. A bounded source is drained once and ends time; an
unbounded one is polled
+ // repeatedly and moves the watermark as its reader reports progress.
+ if (payload.getIsBounded() == RunnerApi.IsBounded.Enum.UNBOUNDED) {
+ addUnboundedReadNodes(
+ transformId,
+ ReadTranslation.unboundedSourceFromProto(payload),
+ pipeline.getComponents(),
+ outputPCollectionId,
+ context);
+ } else {
+ addReadNodes(
+ transformId,
+ ReadTranslation.boundedSourceFromProto(payload),
+ pipeline.getComponents(),
+ outputPCollectionId,
+ context);
+ }
} catch (IOException e) {
throw new RuntimeException(
- "Failed to read the BoundedSource from transform " +
transform.getUniqueName(), e);
+ "Failed to read the source from transform " +
transform.getUniqueName(), e);
}
}
+ /**
+ * Adds the source, {@link UnboundedReadProcessor}, and the store holding
its checkpoint mark.
+ *
+ * <p>The store keeps encoded bytes rather than the mark itself, since the
mark's coder comes from
+ * the source and is only known here.
+ */
+ private <T, CheckpointT extends UnboundedSource.CheckpointMark> void
addUnboundedReadNodes(
+ String transformId,
+ UnboundedSource<T, CheckpointT> source,
+ RunnerApi.Components components,
+ String outputPCollectionId,
+ KafkaStreamsTranslationContext context) {
+ PCollectionNode outputNode =
+ PipelineNode.pCollection(
+ outputPCollectionId,
components.getPcollectionsOrThrow(outputPCollectionId));
+ Coder<WindowedValue<T>> sdkWireCoder = sdkWireCoder(outputNode,
components);
+ Coder<WindowedValue<?>> runnerWireCoder = runnerWireCoder(outputNode,
components);
+
+ Topology topology = context.getTopology();
+ String sourceNodeName = transformId + SOURCE_SUFFIX;
+ String stateStoreName = transformId + STATE_STORE_SUFFIX;
+ String bootstrapTopic = context.getReadBootstrapTopic(transformId);
+ SerializablePipelineOptions options =
+ new SerializablePipelineOptions(context.getPipelineOptions());
+ // Split here rather than in the processor: splitting belongs to
translation, where it happens
+ // once for the pipeline instead of once per task instance, and the
contract says nothing about
+ // splitting a source that has already been split.
+ UnboundedSource<T, CheckpointT> readableSource = singleSplitOf(source,
context);
+ Coder<CheckpointT> checkpointCoder =
readableSource.getCheckpointMarkCoder();
+ int maxElementsPerPoll = context.getPipelineOptions().getMaxBundleSize();
+ int checkpointEveryNPolls =
context.getPipelineOptions().getReadCheckpointNumBundles();
+
+ topology.addSource(
+ sourceNodeName,
+ Serdes.ByteArray().deserializer(),
+ Serdes.ByteArray().deserializer(),
+ bootstrapTopic);
+ topology.addProcessor(
+ transformId,
+ () ->
+ new UnboundedReadProcessor<>(
+ readableSource,
+ options,
+ sdkWireCoder,
+ runnerWireCoder,
+ checkpointCoder,
+ stateStoreName,
+ transformId,
+ maxElementsPerPoll,
+ checkpointEveryNPolls),
+ sourceNodeName);
+ topology.addStateStore(
+ Stores.keyValueStoreBuilder(
+ Stores.persistentKeyValueStore(stateStoreName), Serdes.String(),
Serdes.ByteArray()),
+ transformId);
+
+ context.registerPCollectionProducer(outputPCollectionId, transformId);
+ }
+
/**
* Adds the source, {@link ReadProcessor}, and state store for the read. The
type variable {@code
* T} captures the {@link BoundedSource}'s element type so the processor and
its wire coders are
@@ -138,6 +207,38 @@ class ReadTranslator implements PTransformTranslator {
context.registerPCollectionProducer(outputPCollectionId, transformId);
}
+ /**
+ * Splits an unbounded source into the single part this runner reads.
+ *
+ * <p>A source is not obliged to be readable in its unsplit form — {@code
split} is where several
+ * of them do their setup — so it is asked to split even though only one
part is wanted. The count
+ * passed to {@code split} is only a hint, so what comes back has to be
checked: taking the first
+ * of several splits would quietly drop whatever the others would have
produced, which is data
+ * loss rather than a missing feature, so it fails instead.
+ */
+ private static <T, CheckpointT extends UnboundedSource.CheckpointMark>
+ UnboundedSource<T, CheckpointT> singleSplitOf(
+ UnboundedSource<T, CheckpointT> source,
KafkaStreamsTranslationContext context) {
+ List<? extends UnboundedSource<T, CheckpointT>> splits;
+ try {
+ splits = source.split(1, context.getPipelineOptions());
+ } catch (Exception e) {
+ throw new RuntimeException("Failed to split unbounded source " + source,
e);
+ }
+ if (splits.size() != 1) {
+ throw new UnsupportedOperationException(
+ "Unbounded source "
+ + source
+ + " split into "
+ + splits.size()
+ + " parts, but the Kafka Streams runner reads a source with a
single reader and"
+ + " would therefore drop the data of every part but the first.
Reading several"
+ + " splits in parallel is not supported yet; see"
+ + " https://github.com/apache/beam/issues/18479.";);
+ }
+ return splits.get(0);
+ }
+
/** The coder the SDK harness would use on the wire, keeping unknown element
coders intact. */
private static <T> Coder<WindowedValue<T>> sdkWireCoder(
PCollectionNode outputNode, RunnerApi.Components components) {
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java
new file mode 100644
index 00000000000..b616e0fa853
--- /dev/null
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java
@@ -0,0 +1,304 @@
+/*
+ * 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.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+import java.io.IOException;
+import java.time.Duration;
+import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderException;
+import org.apache.beam.sdk.io.UnboundedSource;
+import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark;
+import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.util.CoderUtils;
+import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.kafka.streams.processor.Cancellable;
+import org.apache.kafka.streams.processor.PunctuationType;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.apache.kafka.streams.processor.api.Record;
+import org.apache.kafka.streams.state.KeyValueStore;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Reads an {@link UnboundedSource} and forwards its elements and watermark
downstream.
+ *
+ * <p>Where the bounded {@link ReadProcessor} drains its source once and jumps
the watermark to the
+ * end of time, an unbounded source never finishes: it is polled repeatedly,
and its watermark
+ * advances gradually as the reader reports progress. That difference is what
makes this a streaming
+ * runner rather than a batch one — downstream windows close because the
source says time has moved
+ * on, not because the input ran out.
+ *
+ * <p>Polling happens on a wall-clock punctuator rather than in {@code
process}, because the
+ * processor's bootstrap topic is empty and nothing else would drive it. Each
turn reads at most
+ * {@link #maxElementsPerPoll} elements so a busy source cannot monopolise the
Kafka Streams thread
+ * and starve the rest of the topology, then forwards the reader's watermark
if it advanced.
+ *
+ * <p><b>Restart</b> is what the checkpoint mark is for. {@link
UnboundedReader#getCheckpointMark()}
+ * describes the position the reader has consumed to; it is written to a
persistent state store, and
+ * on {@link #init} the reader is created from the stored mark rather than
from scratch, so a task
+ * that moves or restarts resumes where it left off instead of re-reading from
the beginning. The
+ * store is changelogged and, under exactly-once, its writes commit atomically
with the records the
+ * processor forwarded, so the mark can never be ahead of the data that was
actually emitted.
+ *
+ * <p>The source handed to this processor has already been split by {@link
ReadTranslator}, which is
+ * where splitting belongs: it happens once for the pipeline rather than once
per task instance, and
+ * the contract does not define splitting an already-split source. Reading
several splits in
+ * parallel arrives with the topic-based shuffle work (#18479). As in the
bounded processor, Kafka
+ * Streams disallows negative record timestamps, so each forwarded {@link
Record} carries the Unix
+ * epoch and the Beam event time travels inside the {@link WindowedValue}.
+ */
+class UnboundedReadProcessor<T, CheckpointT extends CheckpointMark>
+ implements Processor<byte[], byte[], byte[], KStreamsPayload<?>> {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(UnboundedReadProcessor.class);
+
+ /** Sole entry in the state store; the value is the encoded checkpoint mark.
*/
+ static final String CHECKPOINT_KEY = "checkpoint";
+
+ /** How often the source is polled. */
+ private static final Duration POLL_INTERVAL = Duration.ofMillis(50);
+
+ private final UnboundedSource<T, CheckpointT> source;
+ private final SerializablePipelineOptions options;
+ // See ReadProcessor: a source produces decoded objects, but the downstream
stage's harness input
+ // expects the runner-side wire form, so each element is transcoded through
these two coders.
+ private final Coder<WindowedValue<T>> sdkWireCoder;
+ private final Coder<WindowedValue<?>> runnerWireCoder;
+ private final Coder<CheckpointT> checkpointCoder;
+ private final String stateStoreName;
+ private final String transformId;
+ private final int maxElementsPerPoll;
+ private final int checkpointEveryNPolls;
+
+ private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
+ private @Nullable KeyValueStore<String, byte[]> checkpointStore;
+ private @Nullable UnboundedReader<T> reader;
+ private boolean readerStarted;
+ private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
+ /** Set once the source's watermark reaches the end of time; it will produce
nothing more. */
+ private boolean exhausted;
+
+ private @Nullable Cancellable scheduledPunctuator;
+
+ UnboundedReadProcessor(
+ UnboundedSource<T, CheckpointT> source,
+ SerializablePipelineOptions options,
+ Coder<WindowedValue<T>> sdkWireCoder,
+ Coder<WindowedValue<?>> runnerWireCoder,
+ Coder<CheckpointT> checkpointCoder,
+ String stateStoreName,
+ String transformId,
+ int maxElementsPerPoll,
+ int checkpointEveryNPolls) {
+ this.source = source;
+ this.options = options;
+ this.sdkWireCoder = sdkWireCoder;
+ this.runnerWireCoder = runnerWireCoder;
+ this.checkpointCoder = checkpointCoder;
+ this.stateStoreName = stateStoreName;
+ this.transformId = transformId;
+ this.maxElementsPerPoll = maxElementsPerPoll;
+ this.checkpointEveryNPolls = checkpointEveryNPolls;
+ }
+
+ @Override
+ public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
+ this.context = context;
+ this.checkpointStore = context.getStateStore(stateStoreName);
+ this.scheduledPunctuator =
+ context.schedule(POLL_INTERVAL, PunctuationType.WALL_CLOCK_TIME,
timestamp -> poll());
+ }
+
+ @Override
+ public void process(Record<byte[], byte[]> record) {
+ // The bootstrap topic carries no real data; a record arriving on it is
just another chance to
+ // poll. The reader's own position decides what is actually emitted.
+ poll();
+ }
+
+ /**
+ * Drains what the source currently has, in batches, then publishes the
watermark.
+ *
+ * <p>A batch is capped at {@link #maxElementsPerPoll} so that the
checkpoint mark and the
+ * watermark are updated as the reader progresses rather than only at the
end. Batches run back to
+ * back while the source keeps filling them, since returning after every
batch would cap
+ * throughput at one batch per punctuation interval.
+ *
+ * <p>The run is bounded all the same. A source that always has data — which
is the normal case
+ * for one that is keeping up — would otherwise never let this method
return, and the Kafka
+ * Streams thread would never get back to committing or to the rest of the
topology. So at most
+ * {@link #checkpointEveryNPolls} batches are taken before yielding, which
is also where the
+ * checkpoint mark is stored, and the next punctuation carries on from there.
+ */
+ private void poll() {
+ if (exhausted) {
+ return;
+ }
+ ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
+ UnboundedReader<T> currentReader = ensureReader();
+ for (int batch = 0; batch < checkpointEveryNPolls; batch++) {
+ int emitted = readBatch(ctx, currentReader);
+ Instant watermark = currentReader.getWatermark();
+ forwardWatermarkIfAdvanced(ctx, watermark);
+ if (!watermark.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) {
+ // The source has declared it will produce nothing further, so stop
polling it. Store the
+ // final position first, since the loop will not come back to it.
+ storeCheckpoint(currentReader);
+ exhausted = true;
+ Cancellable punctuator = scheduledPunctuator;
+ if (punctuator != null) {
+ punctuator.cancel();
+ scheduledPunctuator = null;
+ }
+ return;
+ }
+ if (emitted < maxElementsPerPoll) {
+ // Short batch: the source has nothing more for now, so store what was
read and wait for
+ // the next punctuation rather than spinning on a reader that keeps
returning false.
+ if (emitted > 0) {
+ storeCheckpoint(currentReader);
+ }
+ return;
+ }
+ }
+ // Yielded on the batch bound rather than on an empty source, so record
the position reached.
+ storeCheckpoint(currentReader);
+ }
+
+ /** Forwards up to {@link #maxElementsPerPoll} elements, returning how many
were available. */
+ private int readBatch(
+ ProcessorContext<byte[], KStreamsPayload<?>> ctx, UnboundedReader<T>
currentReader) {
+ int emitted = 0;
+ try {
+ while (emitted < maxElementsPerPoll) {
+ // start() positions the reader on its first element; advance() moves
to the next. Either
+ // returning false means nothing is available right now — not that the
source is finished,
+ // which is the difference from a bounded read.
+ boolean hasElement = readerStarted ? currentReader.advance() :
currentReader.start();
+ readerStarted = true;
+ if (!hasElement) {
+ break;
+ }
+ WindowedValue<T> element =
+ WindowedValues.timestampedValueInGlobalWindow(
+ currentReader.getCurrent(),
currentReader.getCurrentTimestamp());
+ ctx.forward(
+ new Record<byte[], KStreamsPayload<?>>(
+ new byte[0], KStreamsPayload.data(toRunnerWire(element)), 0L));
+ emitted++;
+ }
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to read unbounded source for
transform " + transformId, e);
+ }
+ return emitted;
+ }
+
+ /** Publishes the reader's watermark, which is what lets downstream windows
close. */
+ private void forwardWatermarkIfAdvanced(
+ ProcessorContext<byte[], KStreamsPayload<?>> ctx, Instant watermark) {
+ if (!watermark.isAfter(lastForwardedWatermark)) {
+ return;
+ }
+ lastForwardedWatermark = watermark;
+ ctx.forward(
+ new Record<byte[], KStreamsPayload<?>>(
+ new byte[0], KStreamsPayload.watermark(watermark.getMillis(),
transformId, 0, 1), 0L));
+ }
+
+ /** Creates the reader on first use, resuming from the stored checkpoint
mark if there is one. */
+ private UnboundedReader<T> ensureReader() {
+ UnboundedReader<T> existing = reader;
+ if (existing != null) {
+ return existing;
+ }
+ try {
+ UnboundedReader<T> created = source.createReader(options.get(),
restoreCheckpoint());
+ reader = created;
+ return created;
+ } catch (Exception e) {
+ throw new RuntimeException(
+ "Failed to create a reader for unbounded source in transform " +
transformId, e);
+ }
+ }
+
+ private @Nullable CheckpointT restoreCheckpoint() {
+ KeyValueStore<String, byte[]> store = checkInitialized(checkpointStore);
+ byte[] encoded = store.get(CHECKPOINT_KEY);
+ if (encoded == null) {
+ return null;
+ }
+ try {
+ CheckpointT mark = CoderUtils.decodeFromByteArray(checkpointCoder,
encoded);
+ LOG.info("Unbounded read {} resuming from a stored checkpoint mark",
transformId);
+ return mark;
+ } catch (CoderException e) {
+ throw new RuntimeException(
+ "Failed to decode the checkpoint mark for transform " + transformId,
e);
+ }
+ }
+
+ private void storeCheckpoint(UnboundedReader<T> currentReader) {
+ KeyValueStore<String, byte[]> store = checkInitialized(checkpointStore);
+ @SuppressWarnings("unchecked")
+ CheckpointT mark = (CheckpointT) currentReader.getCheckpointMark();
+ try {
+ store.put(CHECKPOINT_KEY, CoderUtils.encodeToByteArray(checkpointCoder,
mark));
+ } catch (CoderException e) {
+ throw new RuntimeException(
+ "Failed to encode the checkpoint mark for transform " + transformId,
e);
+ }
+ }
+
+ /** Transcodes a raw element into the runner-side wire form the SDK harness
input expects. */
+ private WindowedValue<?> toRunnerWire(WindowedValue<T> element) {
+ try {
+ byte[] wireBytes = CoderUtils.encodeToByteArray(sdkWireCoder, element);
+ return CoderUtils.decodeFromByteArray(runnerWireCoder, wireBytes);
+ } catch (CoderException e) {
+ throw new RuntimeException(
+ "Failed to transcode an unbounded-read element to wire form for
transform " + transformId,
+ e);
+ }
+ }
+
+ @Override
+ public void close() {
+ UnboundedReader<T> currentReader = reader;
+ if (currentReader != null) {
+ try {
+ currentReader.close();
+ } catch (IOException e) {
+ LOG.warn("Error closing the reader for unbounded source {}",
transformId, e);
+ }
+ reader = null;
+ }
+ }
+
+ private static <V> V checkInitialized(@Nullable V value) {
+ if (value == null) {
+ throw new IllegalStateException("UnboundedReadProcessor used before
init()");
+ }
+ return value;
+ }
+}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java
new file mode 100644
index 00000000000..b668ee972d5
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java
@@ -0,0 +1,209 @@
+/*
+ * 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.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.greaterThan;
+
+import java.io.IOException;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions;
+import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.io.CountingSource;
+import org.apache.beam.sdk.io.CountingSource.CounterMark;
+import org.apache.beam.sdk.io.Read;
+import org.apache.beam.sdk.io.UnboundedSource;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.kafka.streams.TopologyTestDriver;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.junit.Before;
+import org.junit.Test;
+
+/**
+ * Runs a pipeline whose source is unbounded, which is the shape the runner
exists for: a Kafka
+ * Streams application is a long-running stream processor, and until now the
runner could only read
+ * sources that finish.
+ *
+ * <p>Two things separate this from the bounded read. The source is polled
repeatedly rather than
+ * drained once, so elements arrive over several turns of the wall clock; and
the watermark comes
+ * from the reader's own progress rather than jumping to the end of time when
the input runs out,
+ * which is what lets downstream windows close on a stream that never ends.
+ */
+public class UnboundedReadTest {
+
+ /** How many elements one poll of the source may take. */
+ private static final int ELEMENTS_PER_POLL = 5;
+
+ /** Elements the finite variant of the source produces before ending time. */
+ private static final int ELEMENTS = 12;
+
+ /** Elements the pipeline has seen, recorded in order. */
+ private static final List<Long> RECEIVED = Collections.synchronizedList(new
ArrayList<>());
+
+ @Before
+ public void reset() {
+ RECEIVED.clear();
+ }
+
+ private static class RecordFn extends DoFn<Long, Long> {
+ @ProcessElement
+ public void processElement(@Element Long element, OutputReceiver<Long>
out) {
+ RECEIVED.add(element);
+ out.output(element);
+ }
+ }
+
+ /**
+ * A genuinely unbounded pipeline. Nothing caps the source — capping it with
{@code
+ * withMaxNumRecords} would turn it back into a bounded read and test the
wrong path — so the work
+ * is bounded instead by how many elements a single poll may take and how
many turns the test
+ * drives.
+ */
+ private static Pipeline unboundedPipeline() {
+ KafkaStreamsPipelineOptions options =
+
KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class);
+ options.setMaxBundleSize(ELEMENTS_PER_POLL);
+ Pipeline pipeline = Pipeline.create(options);
+ pipeline
+ .apply("read", Read.from(CountingSource.unbounded()))
+ .apply("record", ParDo.of(new RecordFn()));
+ return pipeline;
+ }
+
+ @Test
+ public void anUnboundedSourceIsPolledAndItsElementsReachTheHarness() {
+ Pipeline pipeline = unboundedPipeline();
+ KafkaStreamsTranslationContext context =
KafkaStreamsTestRunner.translate(pipeline);
+
+ try (TopologyTestDriver driver =
+ new TopologyTestDriver(
+ context.getTopology(),
KafkaStreamsTestRunner.streamsConfig(pipeline))) {
+ // Several turns, because an unbounded read yields what is available now
rather than
+ // everything at once.
+ for (int turn = 0; turn < 4; turn++) {
+ driver.advanceWallClockTime(Duration.ofMillis(100));
+ }
+ }
+
+ // More than a single poll's worth, which is the point: a bounded read
drains once, whereas
+ // this one has to be asked again on each turn of the clock and keep going
from where it was.
+ assertThat(
+ "expected several polls' worth of elements, got " + RECEIVED.size(),
+ RECEIVED.size(),
+ is(greaterThan(ELEMENTS_PER_POLL)));
+ // The source counts from zero, so what arrived has to start there and be
contiguous — no gap
+ // and no repeat, which is what the checkpoint mark between polls is for.
+ for (int i = 0; i < RECEIVED.size(); i++) {
+ assertThat(RECEIVED.get(i), is((long) i));
+ }
+ }
+
+ @Test
+ public void aSourceThatReachesTheEndOfTimeStopsBeingPolled() {
+ // CountingSource.unbounded() with a limit reports the terminal watermark
once it has produced
+ // its elements, which is a source saying it will yield nothing further.
Polling must stop
+ // there rather than spinning on a reader that can only return false.
+ KafkaStreamsPipelineOptions options =
+
KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class);
+ options.setMaxBundleSize(ELEMENTS_PER_POLL);
+ Pipeline pipeline = Pipeline.create(options);
+ pipeline
+ .apply("read",
Read.from(CountingSource.unbounded()).withMaxNumRecords(ELEMENTS))
+ .apply("record", ParDo.of(new RecordFn()));
+
+ KafkaStreamsTranslationContext context =
KafkaStreamsTestRunner.translate(pipeline);
+ try (TopologyTestDriver driver =
+ new TopologyTestDriver(
+ context.getTopology(),
KafkaStreamsTestRunner.streamsConfig(pipeline))) {
+ for (int turn = 0; turn < 10; turn++) {
+ driver.advanceWallClockTime(Duration.ofMillis(100));
+ }
+ }
+
+ // Every element exactly once: the source finished, and the turns after it
finished added
+ // nothing.
+ assertThat(RECEIVED.size(), is(ELEMENTS));
+ }
+
+ /** A source that ignores the requested split count and always returns two
parts. */
+ private static class TwoSplitSource extends UnboundedSource<Long,
CounterMark> {
+ private final UnboundedSource<Long, CounterMark> delegate =
CountingSource.unbounded();
+
+ @Override
+ public List<? extends UnboundedSource<Long, CounterMark>> split(
+ int desiredNumSplits, PipelineOptions options) throws Exception {
+ return delegate.split(2, options);
+ }
+
+ @Override
+ public UnboundedReader<Long> createReader(
+ PipelineOptions options, @Nullable CounterMark checkpointMark) throws
IOException {
+ return delegate.createReader(options, checkpointMark);
+ }
+
+ @Override
+ public Coder<CounterMark> getCheckpointMarkCoder() {
+ return delegate.getCheckpointMarkCoder();
+ }
+
+ @Override
+ public Coder<Long> getOutputCoder() {
+ return delegate.getOutputCoder();
+ }
+ }
+
+ @Test
+ public void aSourceThatSplitsIntoSeveralPartsIsRejectedRatherThanTruncated()
{
+ // The count passed to split() is only a hint. Reading the first part of
several and ignoring
+ // the rest would silently drop their data, so translation has to fail
instead.
+ Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions());
+ pipeline
+ .apply("read", Read.from(new TwoSplitSource()))
+ .apply("record", ParDo.of(new RecordFn()));
+
+ try {
+ KafkaStreamsTestRunner.translate(pipeline);
+ throw new AssertionError("expected a multi-split source to be rejected");
+ } catch (UnsupportedOperationException e) {
+ assertThat(e.getMessage(), containsString("split into 2 parts"));
+ assertThat(e.getMessage(), containsString("drop the data"));
+ }
+ }
+
+ @Test
+ public void theSourceIsTranslatedAsAnUnboundedRead() {
+ // The bounded and unbounded reads share a URN and are told apart by the
payload, so this pins
+ // down that the pipeline really did take the unbounded path.
+ KafkaStreamsTranslationContext context =
KafkaStreamsTestRunner.translate(unboundedPipeline());
+
+ boolean hasReadProcessor =
+ context.getTopology().describe().subtopologies().stream()
+ .flatMap(subtopology -> subtopology.nodes().stream())
+ .anyMatch(node -> node.name().contains("read"));
+ assertThat(hasReadProcessor, is(true));
+ }
+}
