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 e95268ac1308..44abc8e5b34d 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 403499ca616c..f83442f97813 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 @@ public void translate( // 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. + * + *

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 void addUnboundedReadNodes( + String transformId, + UnboundedSource source, + RunnerApi.Components components, + String outputPCollectionId, + KafkaStreamsTranslationContext context) { + PCollectionNode outputNode = + PipelineNode.pCollection( + outputPCollectionId, components.getPcollectionsOrThrow(outputPCollectionId)); + Coder> sdkWireCoder = sdkWireCoder(outputNode, components); + Coder> 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 readableSource = singleSplitOf(source, context); + Coder 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 @@ private void addReadNodes( context.registerPCollectionProducer(outputPCollectionId, transformId); } + /** + * Splits an unbounded source into the single part this runner reads. + * + *

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 + UnboundedSource singleSplitOf( + UnboundedSource source, KafkaStreamsTranslationContext context) { + List> 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 Coder> 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 000000000000..b616e0fa8534 --- /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. + * + *

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. + * + *

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. + * + *

Restart 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. + * + *

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 + implements Processor> { + + 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 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> sdkWireCoder; + private final Coder> runnerWireCoder; + private final Coder checkpointCoder; + private final String stateStoreName; + private final String transformId; + private final int maxElementsPerPoll; + private final int checkpointEveryNPolls; + + private @Nullable ProcessorContext> context; + private @Nullable KeyValueStore checkpointStore; + private @Nullable UnboundedReader 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 source, + SerializablePipelineOptions options, + Coder> sdkWireCoder, + Coder> runnerWireCoder, + Coder 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> 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 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. + * + *

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. + * + *

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> ctx = checkInitialized(context); + UnboundedReader 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> ctx, UnboundedReader 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 element = + WindowedValues.timestampedValueInGlobalWindow( + currentReader.getCurrent(), currentReader.getCurrentTimestamp()); + ctx.forward( + new Record>( + 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> ctx, Instant watermark) { + if (!watermark.isAfter(lastForwardedWatermark)) { + return; + } + lastForwardedWatermark = watermark; + ctx.forward( + new Record>( + 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 ensureReader() { + UnboundedReader existing = reader; + if (existing != null) { + return existing; + } + try { + UnboundedReader 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 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 currentReader) { + KeyValueStore 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 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 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 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 000000000000..b668ee972d53 --- /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. + * + *

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 RECEIVED = Collections.synchronizedList(new ArrayList<>()); + + @Before + public void reset() { + RECEIVED.clear(); + } + + private static class RecordFn extends DoFn { + @ProcessElement + public void processElement(@Element Long element, OutputReceiver 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 { + private final UnboundedSource delegate = CountingSource.unbounded(); + + @Override + public List> split( + int desiredNumSplits, PipelineOptions options) throws Exception { + return delegate.split(2, options); + } + + @Override + public UnboundedReader createReader( + PipelineOptions options, @Nullable CounterMark checkpointMark) throws IOException { + return delegate.createReader(options, checkpointMark); + } + + @Override + public Coder getCheckpointMarkCoder() { + return delegate.getCheckpointMarkCoder(); + } + + @Override + public Coder 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)); + } +}