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 4ff618047c3 [GSoC 2026] Kafka Streams runner: terminate a bounded
pipeline when it is drained (#39700)
4ff618047c3 is described below
commit 4ff618047c33b3bf9593adb66ad9ef918333ac9f
Author: M Junaid Shaukat <[email protected]>
AuthorDate: Tue Aug 11 13:14:08 2026 +0500
[GSoC 2026] Kafka Streams runner: terminate a bounded pipeline when it is
drained (#39700)
* [GSoC 2026] Kafka Streams runner: terminate a bounded pipeline when it is
drained
Kafka Streams runs a topology until something closes the client, so a
bounded
pipeline produced its output and then ran for ever. Every processor already
emits TIMESTAMP_MAX_VALUE once its input is exhausted, so each one now
schedules its own termination when it emits that watermark, and the client
is
closed once they have all reported.
Termination is scheduled rather than reported inline so that the work which
follows the final watermark still runs. The callback waits for every
processor
rather than the first, because one instance can own both sides of a
repartition
topic, and it waits until the topology has finished starting, because
processors register as their tasks are initialized.
---
.../kafka/streams/KafkaStreamsPipelineRunner.java | 50 +++++-
.../KafkaStreamsPortablePipelineResult.java | 7 +-
.../translation/ExecutableStageProcessor.java | 13 +-
.../translation/ExecutableStageTranslator.java | 7 +-
.../streams/translation/FlattenProcessor.java | 14 +-
.../streams/translation/FlattenTranslator.java | 4 +-
.../streams/translation/GroupByKeyTranslator.java | 27 ++-
.../streams/translation/ImpulseProcessor.java | 13 +-
.../streams/translation/ImpulseTranslator.java | 4 +-
.../KafkaStreamsTranslationContext.java | 14 ++
.../kafka/streams/translation/ReadProcessor.java | 13 +-
.../kafka/streams/translation/ReadTranslator.java | 11 +-
.../streams/translation/ShuffleByKeyProcessor.java | 19 +-
.../streams/translation/StageOutputProcessor.java | 12 +-
.../streams/translation/TerminationReporter.java | 111 ++++++++++++
.../streams/translation/TerminationTracker.java | 191 +++++++++++++++++++++
.../translation/UnboundedReadProcessor.java | 12 +-
.../translation/WindowedGroupByKeyProcessor.java | 15 +-
.../kafka/streams/KafkaStreamsRunnerBrokerIT.java | 22 +++
.../ExecutableStageProcessorWatermarkTest.java | 3 +-
.../translation/ShuffleByKeyProcessorTest.java | 4 +-
.../translation/StageOutputProcessorTest.java | 6 +-
.../translation/TerminationTrackerTest.java | 172 +++++++++++++++++++
23 files changed, 717 insertions(+), 27 deletions(-)
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java
index 96b4b2cd7f9..93e1058121e 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java
@@ -80,11 +80,59 @@ public class KafkaStreamsPipelineRunner implements
PortablePipelineRunner {
// Build the result before starting: it registers a state listener, and
Kafka Streams only
// accepts one while the application is still in the CREATED state.
KafkaStreamsPortablePipelineResult result =
- new KafkaStreamsPortablePipelineResult(kafkaStreams,
context.getMetricsContainerStepMap());
+ new KafkaStreamsPortablePipelineResult(
+ kafkaStreams,
+ context.getMetricsContainerStepMap(),
+ // Only once every task is initialized are the processors that
have registered the whole
+ // set, and only then can "all of them are finished" mean the
pipeline is finished.
+ context.getTerminationTracker()::started);
+ // A bounded pipeline finishes; Kafka Streams has no notion of that, so
the runner stops the
+ // client itself once every processor has reached the terminal watermark.
Registered before
+ // start(), so a pipeline that drains quickly cannot finish before
anything is listening.
+ context
+ .getTerminationTracker()
+ .onAllTerminated(
+ () -> closeInBackground(kafkaStreams, jobInfo.jobId(), "the
pipeline is drained"));
kafkaStreams.start();
+ // The job service reads the result's state once, when this method
returns, so returning while
+ // the pipeline is still running would leave the job reported as RUNNING
for good. Blocking here
+ // is what FlinkPipelineRunner does too, by blocking in executor.execute().
+ //
+ // A bounded pipeline unblocks this by draining: the processors report
themselves terminated,
+ // the callback above stops the client, and the result's latch is
released. A streaming pipeline
+ // never reaches the terminal watermark, so this blocks until the job is
cancelled, which is the
+ // intended behaviour for a job that has no end.
+ result.waitUntilFinish();
+ if (Thread.currentThread().isInterrupted()) {
+ // Cancelled: the job service interrupts this thread, and the invocation
future it would
+ // otherwise have used to cancel the result has already been cancelled
with it. Stop the
+ // client so it does not outlive the job — from another thread, since
close() waits on the
+ // stream threads and the joins it does would throw straight back out of
an interrupted one.
+ closeInBackground(kafkaStreams, jobInfo.jobId(), "the job was
cancelled");
+ }
return result;
}
+ /**
+ * Stops the Kafka Streams client from a thread of its own.
+ *
+ * <p>Never called from a thread that {@code close()} itself waits for. When
the pipeline drains,
+ * that is the task thread which reported the last termination; when the job
is cancelled, it is
+ * the interrupted invocation thread. In both cases closing inline would
either wait on the thread
+ * doing the closing or abandon the shutdown part-way.
+ */
+ private static void closeInBackground(KafkaStreams kafkaStreams, String
jobId, String reason) {
+ Thread closer =
+ new Thread(
+ () -> {
+ LOG.info("Stopping the Kafka Streams client for job {}: {}",
jobId, reason);
+ kafkaStreams.close();
+ },
+ "kafka-streams-runner-shutdown-" + jobId);
+ closer.setDaemon(true);
+ closer.start();
+ }
+
private static void checkRequiredOption(String name, @Nullable String value)
{
if (value == null || value.isEmpty()) {
throw new IllegalArgumentException(
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java
index 0d508b8189f..81e80b66e30 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java
@@ -53,11 +53,16 @@ class KafkaStreamsPortablePipelineResult implements
PortablePipelineResult {
* listener, and Kafka Streams rejects one once the application has left the
CREATED state.
*/
KafkaStreamsPortablePipelineResult(
- KafkaStreams kafkaStreams, MetricsContainerStepMap
metricsContainerStepMap) {
+ KafkaStreams kafkaStreams,
+ MetricsContainerStepMap metricsContainerStepMap,
+ Runnable onRunning) {
this.kafkaStreams = kafkaStreams;
this.metricsContainerStepMap = metricsContainerStepMap;
kafkaStreams.setStateListener(
(newState, oldState) -> {
+ if (newState == KafkaStreams.State.RUNNING) {
+ onRunning.run();
+ }
if (newState == KafkaStreams.State.NOT_RUNNING || newState ==
KafkaStreams.State.ERROR) {
terminated.countDown();
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
index ef376606114..fcb8b2ff27e 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
@@ -111,6 +111,9 @@ class ExecutableStageProcessor
// Computes this stage's input watermark from its upstream transform's
reports, holding until
// every partition of the upstream transform has reported (see
WatermarkAggregator).
private final WatermarkAggregator watermarkAggregator;
+ // Reports this stage instance as finished once it emits the terminal
watermark, so a bounded
+ // pipeline can stop itself.
+ private final TerminationReporter terminationReporter;
// The last watermark actually forwarded downstream, so we only forward when
it advances.
private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
@@ -139,7 +142,8 @@ class ExecutableStageProcessor
Set<String> upstreamTransformIds,
MetricsContainerImpl metricsContainer,
Map<String, String> outputChildByPCollectionId,
- int maxBundleSize) {
+ int maxBundleSize,
+ TerminationTracker terminationTracker) {
this.stagePayload = stagePayload;
this.jobInfo = jobInfo;
this.transformId = transformId;
@@ -147,6 +151,7 @@ class ExecutableStageProcessor
this.metricsContainer = metricsContainer;
this.outputChildByPCollectionId =
ImmutableMap.copyOf(outputChildByPCollectionId);
this.maxBundleSize = maxBundleSize;
+ this.terminationReporter = new TerminationReporter(terminationTracker,
transformId);
}
/** A harness output element together with the id of the output PCollection
it belongs to. */
@@ -163,6 +168,7 @@ class ExecutableStageProcessor
@Override
public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
this.context = context;
+ terminationReporter.init(context);
// The SDK harness (stage context + bundle factory) is created lazily on
the first data
// element, so a stage that only forwards watermarks never spins one up.
This mirrors Spark's
// SparkExecutableStageFunction, which likewise does not build a bundle
factory when there are
@@ -336,6 +342,7 @@ class ExecutableStageProcessor
record.key(),
KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1),
record.timestamp()));
+ terminationReporter.watermarkEmitted(ctx, watermarkMillis);
}
@Override
@@ -364,6 +371,10 @@ class ExecutableStageProcessor
} catch (Exception e) {
LOG.warn("Error closing executable stage context", e);
}
+ // Last: this is what stops the pipeline waiting on this stage, and
closing the bundle above can
+ // still forward records downstream. Releasing it first would let the
pipeline be declared
+ // finished while this stage was flushing.
+ terminationReporter.close();
}
private static <T> T checkInitialized(@Nullable T value) {
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
index c59e5b919ab..11f3192befd 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
@@ -109,7 +109,8 @@ class ExecutableStageTranslator implements
PTransformTranslator {
ImmutableSet.of(parentProcessor),
context.getMetricsContainerStepMap().getContainer(transformId),
outputChildByPCollectionId,
- context.getPipelineOptions().getMaxBundleSize()),
+ context.getPipelineOptions().getMaxBundleSize(),
+ context.getTerminationTracker()),
parentProcessor);
if (multiOutput) {
@@ -118,7 +119,9 @@ class ExecutableStageTranslator implements
PTransformTranslator {
outputChildByPCollectionId.forEach(
(outputPCollectionId, relayName) -> {
topology.addProcessor(
- relayName, () -> new StageOutputProcessor(relayName),
transformId);
+ relayName,
+ () -> new StageOutputProcessor(relayName,
context.getTerminationTracker()),
+ transformId);
context.registerPCollectionProducer(outputPCollectionId,
relayName);
context.registerPCollectionPartitionCount(outputPCollectionId,
partitionCount);
});
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java
index e37da677448..5b36a53607e 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java
@@ -61,6 +61,9 @@ class FlattenProcessor
// The last watermark actually forwarded downstream, so we only forward when
it advances.
private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
+ // Reports this Flatten as finished once every branch it merges has gone
terminal.
+ private final TerminationReporter terminationReporter;
+
private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
/**
@@ -68,14 +71,22 @@ class FlattenProcessor
* @param upstreamTransformIds the producers of this Flatten's input
PCollections (known from the
* pipeline graph), whose reports the {@link WatermarkAggregator} waits
for
*/
- FlattenProcessor(String transformId, Set<String> upstreamTransformIds) {
+ FlattenProcessor(
+ String transformId, Set<String> upstreamTransformIds, TerminationTracker
terminationTracker) {
this.transformId = transformId;
this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds);
+ this.terminationReporter = new TerminationReporter(terminationTracker,
transformId);
}
@Override
public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
this.context = context;
+ terminationReporter.init(context);
+ }
+
+ @Override
+ public void close() {
+ terminationReporter.close();
}
@Override
@@ -105,6 +116,7 @@ class FlattenProcessor
record.key(),
KStreamsPayload.watermark(advanced.getMillis(), transformId, 0,
1),
record.timestamp()));
+ terminationReporter.watermarkEmitted(ctx, advanced.getMillis());
}
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
index 793c1e1dc53..998553d514d 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
@@ -84,7 +84,9 @@ class FlattenTranslator implements PTransformTranslator {
topology.addProcessor(
transformId,
- () -> new FlattenProcessor(transformId, upstreamTransformIds),
+ () ->
+ new FlattenProcessor(
+ transformId, upstreamTransformIds,
context.getTerminationTracker()),
parentProcessors.toArray(new String[0]));
context.registerPCollectionProducer(outputPCollectionId, transformId);
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java
index c460436eedf..c0e69302386 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java
@@ -106,7 +106,8 @@ class GroupByKeyTranslator implements PTransformTranslator {
String holdsIndexStoreName = transformId + HOLDS_INDEX_STORE_SUFFIX;
String timerStoreName = transformId + TIMER_STORE_SUFFIX;
String timerIndexStoreName = transformId + TIMER_INDEX_STORE_SUFFIX;
- String repartitionTopic = repartitionTopic(transformId);
+ String repartitionTopic =
+ repartitionTopic(transformId,
context.getPipelineOptions().getApplicationId());
KStreamsPayloadSerde<KV<Object, Object>> payloadSerde = new
KStreamsPayloadSerde<>(inputCoder);
@@ -118,7 +119,9 @@ class GroupByKeyTranslator implements PTransformTranslator {
int upstreamPartitionCount = context.getPartitionCount(inputPCollectionId);
topology.addProcessor(
shuffleName,
- () -> new ShuffleByKeyProcessor(keyCoder, upstreamPartitionCount),
+ () ->
+ new ShuffleByKeyProcessor(
+ keyCoder, upstreamPartitionCount, shuffleName,
context.getTerminationTracker()),
parentProcessor);
// Shuffle through the repartition topic: data partitioned by key,
watermark broadcast.
@@ -151,7 +154,8 @@ class GroupByKeyTranslator implements PTransformTranslator {
keyCoder,
valueCoder,
windowingStrategy,
- context.getPipelineOptions()),
+ context.getPipelineOptions(),
+ context.getTerminationTracker()),
sourceName);
topology.addStateStore(
Stores.keyValueStoreBuilder(
@@ -200,8 +204,19 @@ class GroupByKeyTranslator implements PTransformTranslator
{
}
}
- /** The internal repartition topic name for a GroupByKey transform. */
- static String repartitionTopic(String transformId) {
- return REPARTITION_TOPIC_PREFIX +
transformId.replaceAll("[^a-zA-Z0-9._-]", "_");
+ /**
+ * The internal repartition topic name for a GroupByKey transform.
+ *
+ * <p>Namespaced by application id, as the Impulse and Read bootstrap topics
already are.
+ * Transform ids come from the pipeline's structure, so two jobs running the
same pipeline would
+ * otherwise shuffle through the same topic and read each other's data —
and, because the topic is
+ * created only if it does not already exist, the second job would silently
inherit the first
+ * job's partition count rather than the one it asked for.
+ */
+ static String repartitionTopic(String transformId, String applicationId) {
+ return REPARTITION_TOPIC_PREFIX
+ + applicationId.replaceAll("[^a-zA-Z0-9._-]", "_")
+ + "_"
+ + transformId.replaceAll("[^a-zA-Z0-9._-]", "_");
}
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
index bac91978a29..7c4590a0e5a 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
@@ -69,24 +69,34 @@ class ImpulseProcessor implements Processor<byte[], byte[],
byte[], KStreamsPayl
private final String stateStoreName;
private final String transformId;
+ // Reports this source as finished once it emits the terminal watermark.
+ private final TerminationReporter terminationReporter;
private @Nullable ProcessorContext<byte[], KStreamsPayload<byte[]>> context;
private @Nullable KeyValueStore<String, Boolean> firedStore;
private @Nullable Cancellable scheduledPunctuator;
- ImpulseProcessor(String stateStoreName, String transformId) {
+ ImpulseProcessor(
+ String stateStoreName, String transformId, TerminationTracker
terminationTracker) {
this.stateStoreName = stateStoreName;
this.transformId = transformId;
+ this.terminationReporter = new TerminationReporter(terminationTracker,
transformId);
}
@Override
public void init(ProcessorContext<byte[], KStreamsPayload<byte[]>> context) {
this.context = context;
this.firedStore = context.getStateStore(stateStoreName);
+ terminationReporter.init(context);
this.scheduledPunctuator =
context.schedule(PUNCTUATION_DELAY, PunctuationType.WALL_CLOCK_TIME,
ts -> maybeFire());
}
+ @Override
+ public void close() {
+ terminationReporter.close();
+ }
+
@Override
public void process(Record<byte[], byte[]> record) {
// Records that happen to land on the bootstrap topic are not actual data;
they just provide an
@@ -132,6 +142,7 @@ class ImpulseProcessor implements Processor<byte[], byte[],
byte[], KStreamsPayl
ctx.forward(
new Record<byte[], KStreamsPayload<byte[]>>(
new byte[0], KStreamsPayload.watermark(maxMillis, transformId, 0,
1), 0L));
+ terminationReporter.watermarkEmitted(ctx, maxMillis);
}
/** Cancels the wall-clock punctuator after the impulse has fired to stop
periodic wakeups. */
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java
index 79d8c3cf577..3daf7782362 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java
@@ -79,7 +79,9 @@ class ImpulseTranslator implements PTransformTranslator {
Serdes.ByteArray().deserializer(),
bootstrapTopic);
topology.addProcessor(
- transformId, () -> new ImpulseProcessor(stateStoreName, transformId),
sourceNodeName);
+ transformId,
+ () -> new ImpulseProcessor(stateStoreName, transformId,
context.getTerminationTracker()),
+ sourceNodeName);
topology.addStateStore(
Stores.keyValueStoreBuilder(
Stores.persistentKeyValueStore(stateStoreName), Serdes.String(),
Serdes.Boolean()),
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java
index d0316961566..904bd8a71e3 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java
@@ -61,6 +61,11 @@ public class KafkaStreamsTranslationContext {
// work.
private final MetricsContainerStepMap metricsContainerStepMap = new
MetricsContainerStepMap();
+ // Decides when a bounded pipeline has finished. Owned by the context, so it
is scoped to this one
+ // pipeline: the job server runs several jobs in a single process, and a
tracker shared between
+ // them would let one pipeline finishing stop another.
+ private final TerminationTracker terminationTracker = new
TerminationTracker();
+
public static KafkaStreamsTranslationContext create(
JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) {
return new KafkaStreamsTranslationContext(jobInfo, pipelineOptions, new
Topology());
@@ -102,6 +107,15 @@ public class KafkaStreamsTranslationContext {
return metricsContainerStepMap;
}
+ /**
+ * Returns the tracker that decides when this pipeline has finished.
Processors report themselves
+ * to it as they reach the terminal watermark; the runner asks it to stop
the Kafka Streams client
+ * once they all have.
+ */
+ public TerminationTracker getTerminationTracker() {
+ return terminationTracker;
+ }
+
/**
* Registers the processor node that produces the given Beam PCollection.
Downstream translators
* resolve their parent processor names by looking up the input PCollection
id.
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java
index eb20a8b8358..a6ef768ae2a 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java
@@ -95,6 +95,8 @@ class ReadProcessor<T> implements Processor<byte[], byte[],
byte[], KStreamsPayl
private final Coder<WindowedValue<?>> runnerWireCoder;
private final String stateStoreName;
private final String transformId;
+ // Reports this source as finished once it emits the terminal watermark.
+ private final TerminationReporter terminationReporter;
private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
private @Nullable KeyValueStore<String, Boolean> firedStore;
@@ -106,19 +108,27 @@ class ReadProcessor<T> implements Processor<byte[],
byte[], byte[], KStreamsPayl
Coder<WindowedValue<T>> sdkWireCoder,
Coder<WindowedValue<?>> runnerWireCoder,
String stateStoreName,
- String transformId) {
+ String transformId,
+ TerminationTracker terminationTracker) {
this.source = source;
this.options = options;
this.sdkWireCoder = sdkWireCoder;
this.runnerWireCoder = runnerWireCoder;
this.stateStoreName = stateStoreName;
this.transformId = transformId;
+ this.terminationReporter = new TerminationReporter(terminationTracker,
transformId);
+ }
+
+ @Override
+ public void close() {
+ terminationReporter.close();
}
@Override
public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
this.context = context;
this.firedStore = context.getStateStore(stateStoreName);
+ terminationReporter.init(context);
this.scheduledPunctuator =
context.schedule(PUNCTUATION_DELAY, PunctuationType.WALL_CLOCK_TIME,
ts -> maybeFire());
}
@@ -193,6 +203,7 @@ class ReadProcessor<T> implements Processor<byte[], byte[],
byte[], KStreamsPayl
ctx.forward(
new Record<byte[], KStreamsPayload<?>>(
new byte[0], KStreamsPayload.<Object>watermark(maxMillis,
transformId, 0, 1), 0L));
+ terminationReporter.watermarkEmitted(ctx, maxMillis);
}
/** Cancels the wall-clock punctuator after the read has fired to stop
periodic wakeups. */
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 f83442f9781..049f29651eb 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
@@ -155,7 +155,8 @@ class ReadTranslator implements PTransformTranslator {
stateStoreName,
transformId,
maxElementsPerPoll,
- checkpointEveryNPolls),
+ checkpointEveryNPolls,
+ context.getTerminationTracker()),
sourceNodeName);
topology.addStateStore(
Stores.keyValueStoreBuilder(
@@ -198,7 +199,13 @@ class ReadTranslator implements PTransformTranslator {
transformId,
() ->
new ReadProcessor<>(
- source, options, sdkWireCoder, runnerWireCoder,
stateStoreName, transformId),
+ source,
+ options,
+ sdkWireCoder,
+ runnerWireCoder,
+ stateStoreName,
+ transformId,
+ context.getTerminationTracker()),
sourceNodeName);
KeyValueBytesStoreSupplier storeSupplier =
Stores.persistentKeyValueStore(stateStoreName);
topology.addStateStore(
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java
index 79595cba7c9..638239406ed 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java
@@ -57,11 +57,21 @@ class ShuffleByKeyProcessor
private int upstreamPartition;
+ // Reports this shuffle as finished once it has written the terminal
watermark to the repartition
+ // topic. The downstream side reading that topic reports separately, which
is why the pipeline
+ // waits for every processor rather than the first.
+ private final TerminationReporter terminationReporter;
+
private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
- ShuffleByKeyProcessor(Coder<Object> keyCoder, int upstreamPartitionCount) {
+ ShuffleByKeyProcessor(
+ Coder<Object> keyCoder,
+ int upstreamPartitionCount,
+ String nodeName,
+ TerminationTracker terminationTracker) {
this.keyCoder = keyCoder;
this.upstreamPartitionCount = upstreamPartitionCount;
+ this.terminationReporter = new TerminationReporter(terminationTracker,
nodeName);
}
@Override
@@ -70,6 +80,12 @@ class ShuffleByKeyProcessor
// This processor runs in the upstream transform's task, so the task's
partition is the
// identity of the instance whose reports it is forwarding.
this.upstreamPartition = context.taskId().partition();
+ terminationReporter.init(context);
+ }
+
+ @Override
+ public void close() {
+ terminationReporter.close();
}
@Override
@@ -109,6 +125,7 @@ class ShuffleByKeyProcessor
upstreamPartition,
upstreamPartitionCount),
record.timestamp()));
+ terminationReporter.watermarkEmitted(ctx, report.getWatermarkMillis());
}
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java
index eec3f2bae08..6b8d6276345 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java
@@ -48,15 +48,24 @@ class StageOutputProcessor
private static final Logger LOG =
LoggerFactory.getLogger(StageOutputProcessor.class);
private final String transformId;
+ // Reports this output port as finished once the stage's terminal watermark
reaches it.
+ private final TerminationReporter terminationReporter;
private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
- StageOutputProcessor(String transformId) {
+ StageOutputProcessor(String transformId, TerminationTracker
terminationTracker) {
this.transformId = transformId;
+ this.terminationReporter = new TerminationReporter(terminationTracker,
transformId);
}
@Override
public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
this.context = context;
+ terminationReporter.init(context);
+ }
+
+ @Override
+ public void close() {
+ terminationReporter.close();
}
@Override
@@ -89,5 +98,6 @@ class StageOutputProcessor
report.getSourcePartition(),
report.getTotalSourcePartitions()),
record.timestamp()));
+ terminationReporter.watermarkEmitted(ctx, report.getWatermarkMillis());
}
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java
new file mode 100644
index 00000000000..abc10500d21
--- /dev/null
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java
@@ -0,0 +1,111 @@
+/*
+ * 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.time.Duration;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.kafka.streams.processor.Cancellable;
+import org.apache.kafka.streams.processor.PunctuationType;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * The bit of every watermark-emitting processor that reports it has finished,
so a bounded pipeline
+ * can stop itself. See {@link TerminationTracker} for why the runner has to
work this out at all.
+ *
+ * <p>A processor creates one of these, calls {@link #init} from {@code
Processor#init}, passes
+ * every watermark it emits to {@link #watermarkEmitted}, and calls {@link
#close} from {@code
+ * Processor#close}.
+ *
+ * <h3>Why termination is scheduled rather than reported inline</h3>
+ *
+ * <p>Reporting from inside {@code process()} would announce the processor as
finished while it is
+ * still in the middle of handling the record that carried the terminal
watermark. Scheduling a
+ * punctuator instead defers the report until the current processing has
completed, so anything that
+ * has to happen after the final watermark — flushing a bundle, forwarding
downstream, committing —
+ * still runs first.
+ *
+ * <p>The punctuator is {@link PunctuationType#WALL_CLOCK_TIME} rather than
stream time: no further
+ * records arrive after the terminal watermark, so stream time would never
advance and a stream-time
+ * punctuator would never fire. The interval is the smallest Kafka Streams
accepts — it rejects
+ * anything below a millisecond with "The minimum supported scheduling
interval is 1 millisecond."
+ */
+class TerminationReporter {
+
+ /** Kafka Streams rejects any scheduling interval below this. */
+ private static final Duration IMMEDIATELY = Duration.ofMillis(1);
+
+ private final TerminationTracker tracker;
+ private final String transformId;
+
+ private @Nullable String instanceId;
+ private @Nullable Cancellable scheduled;
+ private boolean reported;
+
+ TerminationReporter(TerminationTracker tracker, String transformId) {
+ this.tracker = tracker;
+ this.transformId = transformId;
+ }
+
+ /** Registers this processor instance as something the pipeline is waiting
on. */
+ void init(ProcessorContext<?, ?> context) {
+ // The task is what makes the id unique: one processor node runs as one
instance per task.
+ this.instanceId = transformId + "#" + context.taskId();
+ tracker.register(instanceId);
+ }
+
+ /**
+ * Called with every watermark the processor emits. Once that watermark is
terminal, schedules the
+ * report that this processor has no further work.
+ */
+ void watermarkEmitted(ProcessorContext<?, ?> context, long watermarkMillis) {
+ if (reported || watermarkMillis <
BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()) {
+ return;
+ }
+ reported = true;
+ scheduled =
+ context.schedule(
+ IMMEDIATELY,
+ PunctuationType.WALL_CLOCK_TIME,
+ timestamp -> {
+ cancelSchedule();
+ String id = instanceId;
+ if (id != null) {
+ tracker.terminate(id);
+ }
+ });
+ }
+
+ /** Stops the pipeline waiting on this processor, e.g. when its task
migrates on a rebalance. */
+ void close() {
+ cancelSchedule();
+ String id = instanceId;
+ if (id != null) {
+ tracker.unregister(id);
+ instanceId = null;
+ }
+ }
+
+ private void cancelSchedule() {
+ Cancellable handle = scheduled;
+ if (handle != null) {
+ handle.cancel();
+ scheduled = null;
+ }
+ }
+}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java
new file mode 100644
index 00000000000..6c570db1aff
--- /dev/null
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java
@@ -0,0 +1,191 @@
+/*
+ * 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.util.HashSet;
+import java.util.Set;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Decides when a bounded pipeline has finished, so the Kafka Streams client
can be stopped.
+ *
+ * <p>Kafka Streams has no notion of a processor being finished: a topology
runs until something
+ * closes the client. A bounded Beam pipeline does finish, though, and the
runner already knows
+ * when: every processor emits a watermark of {@link
+ * org.apache.beam.sdk.transforms.windowing.BoundedWindow#TIMESTAMP_MAX_VALUE}
once its input is
+ * exhausted. This class collects those reports and fires a callback when
there is nothing left to
+ * do.
+ *
+ * <h3>Why no coordination between instances is needed</h3>
+ *
+ * <p>A watermark that crosses a repartition topic is broadcast to
<em>every</em> partition (see
+ * {@link GroupByKeyBroadcastPartitioner}), so every task of every downstream
transform observes the
+ * terminal watermark on its own, whichever instance it happens to run on.
Each instance can
+ * therefore decide to stop from what it sees locally, and they all reach the
same conclusion
+ * without talking to each other.
+ *
+ * <h3>Why every local processor has to be counted, not just the first</h3>
+ *
+ * <p>One instance can own tasks from both sides of a repartition topic. The
upstream side goes
+ * terminal as soon as it has written its data to the topic, while the
downstream side still has to
+ * consume it. Stopping the client when the first processor finishes would cut
that downstream work
+ * off and report the pipeline as done having silently dropped it. So the
callback only fires once
+ * every processor instance registered here has terminated.
+ *
+ * <p>An instance that happens to own only upstream tasks still terminates on
its own, which is
+ * correct: what it wrote is durable in the topic for whichever instance reads
it.
+ *
+ * <h3>Scope</h3>
+ *
+ * <p>One tracker belongs to one pipeline, not to the JVM. The job server runs
many jobs in a single
+ * process, so a shared static tracker would let one pipeline finishing tear
down another.
+ *
+ * <p>A pipeline with an unbounded source never produces a terminal watermark,
so the callback never
+ * fires and the client keeps running — which is the intended behaviour for a
streaming job.
+ */
+public class TerminationTracker {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(TerminationTracker.class);
+
+ /** Processor instances currently running here, by {@code
transformId#taskId}. */
+ private final Set<String> live = new HashSet<>();
+
+ /** Those of {@link #live} that have emitted the terminal watermark. */
+ private final Set<String> terminated = new HashSet<>();
+
+ /**
+ * What to do when the pipeline is finished, cleared as it is taken.
+ *
+ * <p>Clearing it is what stops it running twice: a pipeline only finishes
once, but processors go
+ * on reporting afterwards — the callback stops the client, and closing it
makes every remaining
+ * task close its processors, each of which unregisters and asks again.
+ */
+ private @Nullable Runnable onAllTerminated;
+
+ /**
+ * Whether the topology is fully up, and so whether the registered
processors are the whole set.
+ *
+ * <p>Processors register as their task is initialized, which happens
gradually while the client
+ * starts. Deciding before that is finished reads "every processor is done"
off a set that is
+ * merely incomplete: on a short pipeline the source can drain before the
task downstream of the
+ * repartition topic exists, and stopping there discards the rest of the
pipeline and reports a
+ * successful run that produced nothing.
+ */
+ private boolean started;
+
+ /**
+ * Sets what to do when the pipeline is finished. Must be called before the
topology starts, so
+ * that no processor can terminate before there is anything to call.
+ *
+ * <p>The callback runs on whichever thread completes the picture: usually
the Kafka Streams task
+ * thread reporting the last termination, but the thread reporting startup
when the pipeline
+ * drained before it finished starting. Both are threads {@code
KafkaStreams.close()} waits for,
+ * so stopping the client is the callback's job to hand off to a thread of
its own.
+ */
+ public synchronized void onAllTerminated(Runnable callback) {
+ this.onAllTerminated = callback;
+ }
+
+ /**
+ * Marks the topology as fully started, after which the registered
processors are taken to be the
+ * whole set. Called when Kafka Streams reports {@code RUNNING}, which it
does once every assigned
+ * task has been initialized.
+ *
+ * <p>A pipeline short enough to drain during startup will already have
reported terminations by
+ * then, so this re-checks rather than only gating what comes later.
+ */
+ public void started() {
+ Runnable callback;
+ synchronized (this) {
+ started = true;
+ callback = takeCallbackIfDone();
+ }
+ run(callback);
+ }
+
+ /** Registers a processor instance, called from {@code Processor#init}. */
+ synchronized void register(String instanceId) {
+ live.add(instanceId);
+ }
+
+ /**
+ * Removes a processor instance, called from {@code Processor#close}, so
that a task migrating
+ * away during a rebalance is not waited on forever.
+ */
+ void unregister(String instanceId) {
+ Runnable callback;
+ synchronized (this) {
+ live.remove(instanceId);
+ terminated.remove(instanceId);
+ callback = takeCallbackIfDone();
+ }
+ run(callback);
+ }
+
+ /**
+ * Records that a processor instance has emitted the terminal watermark and
has no further work.
+ */
+ void terminate(String instanceId) {
+ Runnable callback;
+ synchronized (this) {
+ if (!live.contains(instanceId)) {
+ // Terminated after being unregistered, or never registered: nothing
is waiting on it.
+ return;
+ }
+ if (terminated.add(instanceId)) {
+ LOG.debug(
+ "Processor {} reached the terminal watermark ({}/{})",
+ instanceId,
+ terminated.size(),
+ live.size());
+ }
+ callback = takeCallbackIfDone();
+ }
+ run(callback);
+ }
+
+ private static void run(@Nullable Runnable callback) {
+ // Deliberately outside the lock: the callback shuts the pipeline down,
and holding the monitor
+ // while calling into shutdown makes this class part of that path for
anyone who changes what
+ // the callback does later.
+ if (callback != null) {
+ callback.run();
+ }
+ }
+
+ /**
+ * Returns the callback to run if the pipeline is finished, having claimed
the right to run it.
+ */
+ private @Nullable Runnable takeCallbackIfDone() {
+ if (!started || live.isEmpty() || !terminated.containsAll(live)) {
+ return null;
+ }
+ Runnable callback = onAllTerminated;
+ if (callback == null) {
+ // Never set, or already taken — either way there is nothing left to do.
+ return null;
+ }
+ onAllTerminated = null;
+ LOG.info(
+ "All {} processor instances reached the terminal watermark; stopping
the pipeline",
+ live.size());
+ return callback;
+ }
+}
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
index b616e0fa853..96aa1bf0c69 100644
---
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
@@ -99,6 +99,10 @@ class UnboundedReadProcessor<T, CheckpointT extends
CheckpointMark>
/** Set once the source's watermark reaches the end of time; it will produce
nothing more. */
private boolean exhausted;
+ // An unbounded source normally never reaches the terminal watermark, so
this normally never
+ // reports anything. It does matter for a source that is drained or is
bounded in practice.
+ private final TerminationReporter terminationReporter;
+
private @Nullable Cancellable scheduledPunctuator;
UnboundedReadProcessor(
@@ -110,7 +114,9 @@ class UnboundedReadProcessor<T, CheckpointT extends
CheckpointMark>
String stateStoreName,
String transformId,
int maxElementsPerPoll,
- int checkpointEveryNPolls) {
+ int checkpointEveryNPolls,
+ TerminationTracker terminationTracker) {
+ this.terminationReporter = new TerminationReporter(terminationTracker,
transformId);
this.source = source;
this.options = options;
this.sdkWireCoder = sdkWireCoder;
@@ -126,6 +132,7 @@ class UnboundedReadProcessor<T, CheckpointT extends
CheckpointMark>
public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
this.context = context;
this.checkpointStore = context.getStateStore(stateStoreName);
+ terminationReporter.init(context);
this.scheduledPunctuator =
context.schedule(POLL_INTERVAL, PunctuationType.WALL_CLOCK_TIME,
timestamp -> poll());
}
@@ -224,6 +231,7 @@ class UnboundedReadProcessor<T, CheckpointT extends
CheckpointMark>
ctx.forward(
new Record<byte[], KStreamsPayload<?>>(
new byte[0], KStreamsPayload.watermark(watermark.getMillis(),
transformId, 0, 1), 0L));
+ terminationReporter.watermarkEmitted(ctx, watermark.getMillis());
}
/** Creates the reader on first use, resuming from the stored checkpoint
mark if there is one. */
@@ -293,6 +301,8 @@ class UnboundedReadProcessor<T, CheckpointT extends
CheckpointMark>
}
reader = null;
}
+ // Last, so the pipeline is not declared finished while this source is
still closing down.
+ terminationReporter.close();
}
private static <V> V checkInitialized(@Nullable V value) {
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java
index d00b01a3753..4c328542a32 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java
@@ -92,6 +92,10 @@ class WindowedGroupByKeyProcessor<K, V, W extends
BoundedWindow>
private Instant inputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
+ // Reports this GroupByKey instance as finished once it emits the terminal
watermark, which it
+ // only does after firing every pane it was holding.
+ private final TerminationReporter terminationReporter;
+
private @Nullable KeyValueStore<byte[], byte[]> stateStore;
private @Nullable KeyValueStore<byte[], byte[]> holdsIndexStore;
private @Nullable KeyValueStore<byte[], byte[]> timerStore;
@@ -107,7 +111,9 @@ class WindowedGroupByKeyProcessor<K, V, W extends
BoundedWindow>
Coder<K> keyCoder,
Coder<V> valueCoder,
WindowingStrategy<?, W> windowingStrategy,
- PipelineOptions options) {
+ PipelineOptions options,
+ TerminationTracker terminationTracker) {
+ this.terminationReporter = new TerminationReporter(terminationTracker,
transformId);
this.stateStoreName = stateStoreName;
this.holdsIndexStoreName = holdsIndexStoreName;
this.timerStoreName = timerStoreName;
@@ -129,6 +135,12 @@ class WindowedGroupByKeyProcessor<K, V, W extends
BoundedWindow>
this.holdsIndexStore = context.getStateStore(holdsIndexStoreName);
this.timerStore = context.getStateStore(timerStoreName);
this.timerIndexStore = context.getStateStore(timerIndexStoreName);
+ terminationReporter.init(context);
+ }
+
+ @Override
+ public void close() {
+ terminationReporter.close();
}
@Override
@@ -291,6 +303,7 @@ class WindowedGroupByKeyProcessor<K, V, W extends
BoundedWindow>
trigger.key(),
KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1),
trigger.timestamp()));
+ terminationReporter.watermarkEmitted(ctx, watermarkMillis);
}
private @NonNull K decodeKey(byte[] bytes) {
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java
index bb82a403e55..b060561eaab 100644
---
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java
@@ -253,6 +253,28 @@ public class KafkaStreamsRunnerBrokerIT {
}
}
+ @Test
+ public void aBoundedPipelineTerminatesOnItsOwn() throws Exception {
+ // Kafka Streams runs a topology until something stops the client, so a
bounded pipeline used to
+ // run for ever against a real broker: it produced the right answer and
then sat there. The
+ // other tests here cannot see that, because they cancel rather than wait,
and the
+ // ValidatesRunner
+ // suite cannot either, because TopologyTestDriver is synchronous and
always reports DONE.
+ //
+ // Nothing cancels this one. Returning from run() at all is the assertion.
+ KafkaStreamsPipelineOptions options = options(4);
+ Pipeline pipeline = Pipeline.create(options);
+ buildChainedPipeline(pipeline);
+
+ PipelineResult result = runPipeline(pipeline, options);
+
+ assertThat(result.getState(), is(PipelineResult.State.DONE));
+ // And it stopped for the right reason — having produced its output
exactly once. Termination is
+ // driven from a wall-clock punctuator, which is the same mechanism that
duplicates output when
+ // it is used to close bundles on time (#39633), so the count matters as
much as the state.
+ assertThat(counterValue(result), is(1L));
+ }
+
/**
* Polls the pipeline's metrics until the counter reaches {@code expected}
or the timeout hits.
*/
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
index 290e109796a..22728500bf0 100644
---
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
@@ -64,7 +64,8 @@ public class ExecutableStageProcessorWatermarkTest {
// Single-output: no per-output routing (this test drives the
watermark path directly).
ImmutableMap.of(),
// The bundle size bound is irrelevant to the watermark path this test
drives.
- 1000);
+ 1000,
+ new TerminationTracker());
}
/** A report from the upstream transform's given partition. */
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java
index 38669138075..b6716764da7 100644
---
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java
@@ -48,7 +48,9 @@ public class ShuffleByKeyProcessorTest {
new ShuffleByKeyProcessor(
(org.apache.beam.sdk.coders.Coder<Object>)
(org.apache.beam.sdk.coders.Coder<?>) StringUtf8Coder.of(),
- upstreamPartitions);
+ upstreamPartitions,
+ "shuffle-node",
+ new TerminationTracker());
MockProcessorContext<byte[], KStreamsPayload<?>> ctx =
new MockProcessorContext<>(new Properties(), new TaskId(0,
taskPartition), null);
processor.init(ctx);
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java
index 83653f13c89..9d17ae73fcb 100644
---
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java
@@ -56,7 +56,7 @@ public class StageOutputProcessorTest {
@Test
public void watermarkKeepsPartitionIdentityAndRelabelsTransformId() {
MockProcessorContext<byte[], KStreamsPayload<?>> ctx = new
MockProcessorContext<>();
- StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID);
+ StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID, new
TerminationTracker());
processor.init(ctx);
// A report from partition 1 of a 3-instance stage.
@@ -76,7 +76,7 @@ public class StageOutputProcessorTest {
@Test
public void distinctStagePartitionsStayDistinctDownstream() {
MockProcessorContext<byte[], KStreamsPayload<?>> ctx = new
MockProcessorContext<>();
- StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID);
+ StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID, new
TerminationTracker());
processor.init(ctx);
processor.process(watermark(100L, 0, 3));
@@ -93,7 +93,7 @@ public class StageOutputProcessorTest {
@Test
public void dataIsForwardedUnchanged() {
MockProcessorContext<byte[], KStreamsPayload<?>> ctx = new
MockProcessorContext<>();
- StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID);
+ StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID, new
TerminationTracker());
processor.init(ctx);
WindowedValue<byte[]> element = WindowedValues.valueInGlobalWindow(new
byte[] {7});
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/TerminationTrackerTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/TerminationTrackerTest.java
new file mode 100644
index 00000000000..72ef28d8bef
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/TerminationTrackerTest.java
@@ -0,0 +1,172 @@
+/*
+ * 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.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.is;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link TerminationTracker}. */
+@RunWith(JUnit4.class)
+public class TerminationTrackerTest {
+
+ private final AtomicInteger calls = new AtomicInteger();
+
+ /** A tracker for a topology that has finished starting up. */
+ private TerminationTracker tracker() {
+ TerminationTracker tracker = new TerminationTracker();
+ tracker.onAllTerminated(calls::incrementAndGet);
+ tracker.started();
+ return tracker;
+ }
+
+ @Test
+ public void doesNotFireBeforeTheTopologyHasFinishedStarting() {
+ // Processors register as their task is initialized, so mid-startup the
registered set is only
+ // part of the pipeline. A source that drains that quickly would otherwise
look like a finished
+ // pipeline, and stopping there discards every stage that had not started
yet — which shows up
+ // as a run that reports success and produces no output.
+ TerminationTracker tracker = new TerminationTracker();
+ tracker.onAllTerminated(calls::incrementAndGet);
+
+ tracker.register("source#0_0");
+ tracker.terminate("source#0_0");
+ assertThat("the rest of the topology may not exist yet", calls.get(),
is(0));
+
+ // The stage downstream of the repartition topic comes up late and has
real work to do.
+ tracker.register("downstream#1_0");
+ tracker.started();
+ assertThat(calls.get(), is(0));
+
+ tracker.terminate("downstream#1_0");
+ assertThat(calls.get(), is(1));
+ }
+
+ @Test
+ public void firesOnStartupIfEverythingAlreadyTerminated() {
+ // A pipeline short enough to drain entirely during startup still has to
be noticed.
+ TerminationTracker tracker = new TerminationTracker();
+ tracker.onAllTerminated(calls::incrementAndGet);
+ tracker.register("source#0_0");
+ tracker.terminate("source#0_0");
+
+ tracker.started();
+
+ assertThat(calls.get(), is(1));
+ }
+
+ @Test
+ public void firesOnceEveryRegisteredProcessorHasTerminated() {
+ TerminationTracker tracker = tracker();
+ tracker.register("stage#0_0");
+ tracker.register("stage#0_1");
+
+ tracker.terminate("stage#0_0");
+ assertThat("one of two done is not the whole pipeline", calls.get(),
is(0));
+
+ tracker.terminate("stage#0_1");
+ assertThat(calls.get(), is(1));
+ }
+
+ @Test
+ public void doesNotFireWhileAProcessorIsStillRunning() {
+ TerminationTracker tracker = tracker();
+ // The shape that makes counting every processor necessary: one instance
owning both sides of a
+ // repartition topic. The upstream goes terminal as soon as it has written
to the topic, while
+ // the downstream still has to consume it.
+ tracker.register("upstream#0_0");
+ tracker.register("downstream#1_0");
+
+ tracker.terminate("upstream#0_0");
+
+ assertThat("stopping here would cut the downstream off", calls.get(),
is(0));
+ }
+
+ @Test
+ public void doesNotFireWithNothingRegistered() {
+ TerminationTracker tracker = tracker();
+ tracker.terminate("never-registered#0_0");
+ assertThat(calls.get(), is(0));
+ }
+
+ @Test
+ public void firesOnlyOnce() {
+ TerminationTracker tracker = tracker();
+ tracker.register("stage#0_0");
+
+ tracker.terminate("stage#0_0");
+ tracker.terminate("stage#0_0");
+
+ assertThat(calls.get(), is(1));
+ }
+
+ @Test
+ public void shuttingDownDoesNotFireAgain() {
+ // What the callback does is stop the client, which closes every task's
processors, and each of
+ // those unregisters on the way out. So the tracker is asked again several
times after the
+ // pipeline has already been declared finished.
+ TerminationTracker tracker = tracker();
+ tracker.register("source#0_0");
+ tracker.register("stage#1_0");
+ tracker.terminate("source#0_0");
+ tracker.terminate("stage#1_0");
+ assertThat(calls.get(), is(1));
+
+ tracker.unregister("source#0_0");
+ tracker.unregister("stage#1_0");
+
+ assertThat("stopping the pipeline must not stop it a second time",
calls.get(), is(1));
+ }
+
+ @Test
+ public void aProcessorThatMigratesAwayIsNoLongerWaitedOn() {
+ TerminationTracker tracker = tracker();
+ tracker.register("stage#0_0");
+ tracker.register("stage#0_1");
+ tracker.terminate("stage#0_0");
+
+ // Task 0_1 is reassigned to another instance during a rebalance. What is
left here is done, so
+ // this instance has nothing to keep it alive.
+ tracker.unregister("stage#0_1");
+
+ assertThat(calls.get(), is(1));
+ }
+
+ @Test
+ public void unregisteringTheLastProcessorDoesNotFire() {
+ TerminationTracker tracker = tracker();
+ tracker.register("stage#0_0");
+
+ tracker.unregister("stage#0_0");
+
+ assertThat("nothing registered means nothing finished", calls.get(),
is(0));
+ }
+
+ @Test
+ public void withoutACallbackNothingHappens() {
+ TerminationTracker tracker = new TerminationTracker();
+ tracker.register("stage#0_0");
+ tracker.terminate("stage#0_0");
+ // No callback set: the point is that this does not throw.
+ assertThat(calls.get(), is(0));
+ }
+}
