[GSoC 2026] Kafka Streams runner: terminate a bounded pipeline when it is drained

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java

@@ -80,11 +80,59 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo)
     // 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(

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();
           }

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,14 +142,16 @@ 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;
     this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds);
     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 @@ private static final class PendingOutput {
   @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 @@ private void forwardWatermark(Record<byte[], KStreamsPayload<?>> record, long wa
             record.key(),
             KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1),
             record.timestamp()));
+    terminationReporter.watermarkEmitted(ctx, watermarkMillis);
   }
 
   @Override
@@ -364,6 +371,10 @@ public void close() {
     } 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) {

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java

@@ -109,7 +109,8 @@ public void translate(
                 ImmutableSet.of(parentProcessor),
                 context.getMetricsContainerStepMap().getContainer(transformId),
                 outputChildByPCollectionId,
-                context.getPipelineOptions().getMaxBundleSize()),
+                context.getPipelineOptions().getMaxBundleSize(),
+                context.getTerminationTracker()),
         parentProcessor);
 
     if (multiOutput) {
@@ -118,7 +119,9 @@ public void translate(
       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);
           });

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java

@@ -61,21 +61,32 @@ 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;
 
   /**
    * @param transformId this Flatten's own transform id, stamped on the watermarks it emits
    * @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 @@ public void process(Record<byte[], KStreamsPayload<?>> record) {
               record.key(),
               KStreamsPayload.watermark(advanced.getMillis(), transformId, 0, 1),
               record.timestamp()));
+      terminationReporter.watermarkEmitted(ctx, advanced.getMillis());
     }
   }
 

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java

@@ -84,7 +84,9 @@ public void translate(
 
     topology.addProcessor(
         transformId,
-        () -> new FlattenProcessor(transformId, upstreamTransformIds),
+        () ->
+            new FlattenProcessor(
+                transformId, upstreamTransformIds, context.getTerminationTracker()),
         parentProcessors.toArray(new String[0]));
 
     context.registerPCollectionProducer(outputPCollectionId, transformId);

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java

@@ -106,7 +106,8 @@ public void translate(
     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 @@ public void translate(
     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 @@ public void translate(
                 keyCoder,
                 valueCoder,
                 windowingStrategy,
-                context.getPipelineOptions()),
+                context.getPipelineOptions(),
+                context.getTerminationTracker()),
         sourceName);
     topology.addStateStore(
         Stores.keyValueStoreBuilder(
@@ -200,8 +204,19 @@ private static WindowingStrategy<?, BoundedWindow> hydrateWindowingStrategy(
     }
   }
 
-  /** 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._-]", "_");
   }
 }

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 @@ private void forwardWatermarkMax(ProcessorContext<byte[], KStreamsPayload<byte[]
     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. */

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java

@@ -79,7 +79,9 @@ public void translate(
         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()),

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 MetricsContainerStepMap getMetricsContainerStepMap() {
     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.

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 @@ private void forwardWatermarkMax(ProcessorContext<byte[], KStreamsPayload<?>> ct
     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. */