From 82059f559c99c1dc8dcf9175c05b54871225fa71 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Mon, 17 Aug 2026 23:38:54 +0000 Subject: [PATCH 1/3] [Dataflow Streaming] [Multi Key] Flush streaming sinks at key boundaries and bundle completion When multi-key bundling is enabled in Dataflow Streaming Engine, outputs (productions) emitted during a key's processElement lifecycle must be attributed to that specific key's WorkItemCommitRequest, while outputs emitted during DoFn finishBundle must be attributed to the bundle-level MultiKeyWorkItemCommitRequest. This change: 1. Adds finishKey(@Nullable Object key) to the SinkWriter interface and implements it across streaming sinks (WindmillSink, PubsubSink, PubsubDynamicSink) and non-streaming sinks. 2. In WindmillSink, PubsubSink, and PubsubDynamicSink: - When multi-key bundles are enabled: flushes to the active key's commit builder in finishKey; flushes remaining finishBundle outputs to bundle-level outputs (bundleOutputMessages / bundlePubsubMessages) in close(). - When multi-key bundles are disabled: finishKey does not flush; close() flushes all outputs directly into the single key's commit request. - Updates abort() to cleanly discard buffered messages and reset internal streams without flushing. 3. In StreamingModeExecutionContext and StreamingWorkScheduler, tracks bundleOutputMessages and bundlePubsubMessages and attaches them to MultiKeyWorkItemCommitRequest. --- .../runners/dataflow/worker/AvroByteSink.java | 4 + .../beam/runners/dataflow/worker/IsmSink.java | 4 + .../dataflow/worker/PubsubDynamicSink.java | 53 +- .../runners/dataflow/worker/PubsubSink.java | 52 +- .../runners/dataflow/worker/ShuffleSink.java | 4 + .../worker/SizeReportingSinkWrapper.java | 6 + .../worker/StreamingModeExecutionContext.java | 26 + .../runners/dataflow/worker/WindmillSink.java | 31 +- .../worker/util/common/worker/Sink.java | 4 + .../util/common/worker/WriteOperation.java | 9 +- .../processing/StreamingWorkScheduler.java | 48 +- .../IntrinsicMapTaskExecutorFactoryTest.java | 3 + .../worker/PubsubDynamicSinkTest.java | 150 ++++++ .../dataflow/worker/PubsubSinkTest.java | 130 +++++ .../worker/StreamingDataflowWorkerTest.java | 494 ++++++++++++++++++ .../util/common/worker/ExecutorTestUtils.java | 4 + .../common/worker/WriteOperationTest.java | 59 +++ 17 files changed, 1041 insertions(+), 40 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AvroByteSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AvroByteSink.java index 7d47d85bded9..5565a88b2ef5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AvroByteSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AvroByteSink.java @@ -31,6 +31,7 @@ import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.MimeTypes; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A sink that writes Avro files. Records are written to the Avro file as a series of byte arrays. @@ -72,6 +73,9 @@ public long add(T value) throws IOException { return encodedElem.length; } + @Override + public void finishKey(@Nullable Object key) throws IOException {} + @Override public void close() throws IOException { fileWriter.close(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IsmSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IsmSink.java index 43a3a71c6097..e83c64d34bcb 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IsmSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IsmSink.java @@ -49,6 +49,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Optional; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.CountingOutputStream; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A {@link Sink} that writes Ism files. @@ -295,6 +296,9 @@ private void finish() throws IOException { .encode(Footer.of(startOfIndex, startOfBloomFilter, numberOfKeysWritten), out); } + @Override + public void finishKey(@Nullable Object key) throws IOException {} + @Override public void close() throws IOException { finish(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSink.java index 9098cdc6717e..f2dbf70da2e2 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSink.java @@ -113,9 +113,12 @@ public ByteString getDataFromMessage(PubsubMessage formatted, ByteStringOutputSt return stream.toByteStringAndReset(); } - public void close(Windmill.PubSubMessageBundle.Builder outputBuilder) throws IOException { - context.getOutputBuilder().addPubsubMessages(outputBuilder); - outputBuilder.clear(); + private Windmill.PubSubMessageBundle.Builder createOutputBuilder(String topic) { + return Windmill.PubSubMessageBundle.newBuilder() + .setTopic(topic) + .setTimestampLabel(timestampLabel) + .setIdLabel(idLabel) + .setWithAttributes(true); } @Override @@ -127,16 +130,7 @@ public long add(WindowedValue data) throws IOException { !dataTopic.isEmpty(), "No topic set for message when using dynamic topics."); ByteString byteString = getDataFromMessage(data.getValue(), stream); Windmill.PubSubMessageBundle.Builder builder = - outputBuilders.computeIfAbsent( - dataTopic, - topic -> - context - .getOutputBuilder() - .addPubsubMessagesBuilder() - .setTopic(topic) - .setTimestampLabel(timestampLabel) - .setIdLabel(idLabel) - .setWithAttributes(true)); + outputBuilders.computeIfAbsent(dataTopic, this::createOutputBuilder); builder.addMessages( Windmill.Message.newBuilder() .setData(byteString) @@ -145,14 +139,43 @@ public long add(WindowedValue data) throws IOException { return byteString.size(); } + private void flush(boolean bundleLevel) { + try { + for (Windmill.PubSubMessageBundle.Builder builder : outputBuilders.values()) { + if (builder.getMessagesCount() > 0) { + Windmill.PubSubMessageBundle pubsubMessages = builder.build(); + if (bundleLevel) { + context.addBundlePubsubMessages(pubsubMessages); + } else { + context.getOutputBuilder().addPubsubMessages(pubsubMessages); + } + } + } + } finally { + outputBuilders.clear(); + } + } + + @Override + public void finishKey(@Nullable Object key) throws IOException { + if (context.multiKeyBundleEnabled()) { + flush(/*bundleLevel=*/ false); + } + } + @Override public void close() throws IOException { - outputBuilders.clear(); + if (context.multiKeyBundleEnabled()) { + flush(/*bundleLevel=*/ true); + } else { + flush(/*bundleLevel=*/ false); + } } @Override public void abort() throws IOException { - close(); + outputBuilders.clear(); + stream.reset(); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java index 2f4b26b89ab4..ae4655a44e8e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java @@ -135,7 +135,7 @@ public PubsubSink create( @Override public SinkWriter> writer() { - return new PubsubWriter(topic); + return new PubsubWriter(); } /** The SinkWriter for a PubsubSink. */ @@ -143,16 +143,19 @@ class PubsubWriter implements SinkWriter> { private Windmill.PubSubMessageBundle.Builder outputBuilder; private ByteStringOutputStream stream; // Kept across adds for buffer reuse. - private PubsubWriter(String topic) { - outputBuilder = - Windmill.PubSubMessageBundle.newBuilder() - .setTopic(topic) - .setTimestampLabel(timestampLabel) - .setIdLabel(idLabel) - .setWithAttributes(withAttributes); + private PubsubWriter() { + outputBuilder = createOutputBuilder(); stream = new ByteStringOutputStream(); } + private Windmill.PubSubMessageBundle.Builder createOutputBuilder() { + return Windmill.PubSubMessageBundle.newBuilder() + .setTopic(topic) + .setTimestampLabel(timestampLabel) + .setIdLabel(idLabel) + .setWithAttributes(withAttributes); + } + @Override public long add(WindowedValue data) throws IOException { if (!stream.isEmpty()) { @@ -187,18 +190,41 @@ public long add(WindowedValue data) throws IOException { return byteString.size(); } + private void flush(boolean bundleLevel) { + try { + Windmill.PubSubMessageBundle pubsubMessages = outputBuilder.build(); + if (pubsubMessages.getMessagesCount() > 0) { + if (bundleLevel) { + context.addBundlePubsubMessages(pubsubMessages); + } else { + context.getOutputBuilder().addPubsubMessages(pubsubMessages); + } + } + } finally { + outputBuilder = createOutputBuilder(); + } + } + + @Override + public void finishKey(@Nullable Object key) throws IOException { + if (context.multiKeyBundleEnabled()) { + flush(/*bundleLevel=*/ false); + } + } + @Override public void close() throws IOException { - Windmill.PubSubMessageBundle pubsubMessages = outputBuilder.build(); - if (pubsubMessages.getMessagesCount() > 0) { - context.getOutputBuilder().addPubsubMessages(pubsubMessages); + if (context.multiKeyBundleEnabled()) { + flush(/*bundleLevel=*/ true); + } else { + flush(/*bundleLevel=*/ false); } - outputBuilder.clear(); } @Override public void abort() throws IOException { - close(); + outputBuilder = createOutputBuilder(); + stream.reset(); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ShuffleSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ShuffleSink.java index 49ab341eee1c..a231945201e2 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ShuffleSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/ShuffleSink.java @@ -39,6 +39,7 @@ import org.apache.beam.sdk.values.WindowedValues.WindowedValueCoder; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.Ints; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A sink that writes to a shuffle dataset. @@ -300,6 +301,9 @@ public long add(WindowedValue windowedElem) throws IOException { return bytes; } + @Override + public void finishKey(@Nullable Object key) throws IOException {} + @Override public void close() throws IOException { try (Closeable trackedCloseState = tracker.enterState(writeState)) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SizeReportingSinkWrapper.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SizeReportingSinkWrapper.java index ed6a77f6d9c4..1a9919baf49c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SizeReportingSinkWrapper.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SizeReportingSinkWrapper.java @@ -20,6 +20,7 @@ import java.io.IOException; import org.apache.beam.runners.dataflow.worker.util.common.worker.Sink; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.checkerframework.checker.nullness.qual.Nullable; /** * A wrapper for Sink that reports bytes buffered (or written) to {@link DataflowExecutionContext}. @@ -65,6 +66,11 @@ public long add(T value) throws IOException { return size; } + @Override + public void finishKey(@Nullable Object key) throws IOException { + underlyingWriter.finishKey(key); + } + @Override public void close() throws IOException { underlyingWriter.close(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index 365ebbdc1f9d..a01552c9a5a0 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -195,6 +195,8 @@ public interface KeyTransitionListener { private List executedWorks = Collections.emptyList(); private List outputBuilders = Collections.emptyList(); + private List bundleOutputMessages = Collections.emptyList(); + private List bundlePubsubMessages = Collections.emptyList(); // Map> private Map> finalizationCallbacks = Collections.emptyMap(); @@ -321,6 +323,8 @@ public void reset() { // don't clear and reuse, instead reset the reference. this.executedWorks = Collections.emptyList(); this.outputBuilders = Collections.emptyList(); + this.bundleOutputMessages = Collections.emptyList(); + this.bundlePubsubMessages = Collections.emptyList(); this.finalizationCallbacks = Collections.emptyMap(); // Work from prior bundles might have a reference to the old workBatchFailed. // If the work gets retried it'll get the new workBatchFailed to notify failure. @@ -355,6 +359,8 @@ public void start( reset(); this.executedWorks = new ArrayList<>(); this.outputBuilders = new ArrayList<>(); + this.bundleOutputMessages = new ArrayList<>(); + this.bundlePubsubMessages = new ArrayList<>(); this.finalizationCallbacks = new HashMap<>(); this.keyCoder = keyCoder; this.workExecutor = workExecutor; @@ -862,6 +868,26 @@ public List getWorkItemCommits() { return commits; } + public void addBundleOutputMessages(Windmill.OutputMessageBundle outputBundle) { + this.bundleOutputMessages.add(outputBundle); + } + + public List getBundleOutputMessages() { + return bundleOutputMessages; + } + + public void addBundlePubsubMessages(Windmill.PubSubMessageBundle pubsubBundle) { + this.bundlePubsubMessages.add(pubsubBundle); + } + + public List getBundlePubsubMessages() { + return bundlePubsubMessages; + } + + public boolean multiKeyBundleEnabled() { + return multiKeyBundleOptions.multiKeyBundleEnabled(); + } + // Returns list of Work that was executed in the bundle public List getExecutedWorks() { return executedWorks; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java index 9d8a0f0da309..30131c51713e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java @@ -350,8 +350,7 @@ public long add(WindowedValue data) throws IOException { return (long) key.size() + value.size() + metadata.size() + id.size() + offsetSize; } - @Override - public void close() throws IOException { + private void flush(boolean bundleLevel) { try { outputBuilder.setDestinationStreamId(destinationName); @@ -359,17 +358,39 @@ public void close() throws IOException { outputBuilder.addBundles(keyedOutput.build()); } if (outputBuilder.getBundlesCount() > 0) { - context.getOutputBuilder().addOutputMessages(outputBuilder.build()); + Windmill.OutputMessageBundle bundle = outputBuilder.build(); + if (bundleLevel) { + context.addBundleOutputMessages(bundle); + } else { + context.getOutputBuilder().addOutputMessages(bundle); + } } } finally { outputBuilder.clear(); + productionMap.clear(); + } + } + + @Override + public void finishKey(@Nullable Object key) throws IOException { + if (context.multiKeyBundleEnabled()) { + flush(/*bundleLevel=*/ false); + } + } + + @Override + public void close() throws IOException { + if (context.multiKeyBundleEnabled()) { + flush(/*bundleLevel=*/ true); + } else { + flush(/*bundleLevel=*/ false); } - productionMap.clear(); } @Override public void abort() throws IOException { - close(); + outputBuilder.clear(); + productionMap.clear(); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/Sink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/Sink.java index 7bf9ec99b6ae..8bf0cb8def6f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/Sink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/Sink.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.dataflow.worker.util.common.worker; import java.io.IOException; +import org.checkerframework.checker.nullness.qual.Nullable; /** * Abstract base class for Sinks. @@ -36,6 +37,9 @@ public interface SinkWriter extends AutoCloseable { /** Adds a value to the sink. Returns the size in bytes of the data written. */ public long add(ElemT value) throws IOException; + /** Called when all elements for a specific key have been processed. */ + public void finishKey(@Nullable Object key) throws IOException; + /** * {@inheritDoc} * diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperation.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperation.java index a97c9920b9a3..013913d97c59 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperation.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperation.java @@ -107,7 +107,14 @@ public void finish() throws Exception { } @Override - public void finishKey(@Nullable Object key) throws Exception {} + public void finishKey(@Nullable Object key) throws Exception { + try (Closeable scope = context.enterProcess()) { + checkStarted(); + if (writer != null) { + writer.finishKey(key); + } + } + } @Override public void abort() throws Exception { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java index 63cfad5a9a6f..ff2ff8ac94e8 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java @@ -253,10 +253,15 @@ private void processWork( executeWork(work, stageInfo, computationState, handle, keyTransitionListener); workBatch = executeWorkResult.workBatch(); List workItemCommits = executeWorkResult.workItemCommits(); + List bundleOutputMessages = + executeWorkResult.bundleOutputMessages(); + List bundlePubsubMessages = + executeWorkResult.bundlePubsubMessages(); commitFinalizer.cacheCommitFinalizers(executeWorkResult.finalizationCallbacks()); - commitWorkBatch(computationState, workBatch, workItemCommits); + commitWorkBatch( + computationState, workBatch, workItemCommits, bundleOutputMessages, bundlePubsubMessages); recordProcessingStats(workBatch, workItemCommits, executeWorkResult.stateBytesRead()); LOG.debug("Processing done for work batch size: {}", workBatch.size()); @@ -330,6 +335,8 @@ private ExecuteWorkResult executeWork( List workBatch; List workItemCommits; + List bundleOutputMessages; + List bundlePubsubMessages; Map> finalizationCallbacks; long stateBytesRead; { @@ -342,6 +349,8 @@ private ExecuteWorkResult executeWork( // context workBatch = context.getExecutedWorks(); workItemCommits = context.getWorkItemCommits(); + bundleOutputMessages = context.getBundleOutputMessages(); + bundlePubsubMessages = context.getBundlePubsubMessages(); finalizationCallbacks = context.getFinalizationCallbacks(); stateBytesRead = context.getStateBytesRead(); @@ -352,7 +361,12 @@ private ExecuteWorkResult executeWork( computationWorkExecutor = null; return ExecuteWorkResult.create( - workBatch, workItemCommits, finalizationCallbacks, stateBytesRead); + workBatch, + workItemCommits, + bundleOutputMessages, + bundlePubsubMessages, + finalizationCallbacks, + stateBytesRead); } catch (Throwable t) { if (computationWorkExecutor != null) { // If processing failed due to a thrown exception, close the executionState. Do not @@ -387,12 +401,15 @@ private StageInfo getStageInfo(ComputationState computationState) { private void commitWorkBatch( ComputationState computationState, List workBatch, - List workItemCommits) { + List workItemCommits, + List bundleOutputMessages, + List bundlePubsubMessages) { if (workBatch.isEmpty()) { return; } if (workBatch.size() > 1 || multiKeyBundleOptions.multiKeyBundleEnabled()) { - commitMultiKeyWorkBatch(computationState, workBatch, workItemCommits); + commitMultiKeyWorkBatch( + computationState, workBatch, workItemCommits, bundleOutputMessages, bundlePubsubMessages); } else { commitSingleKeyWork(computationState, workBatch.get(0), workItemCommits.get(0)); } @@ -401,11 +418,19 @@ private void commitWorkBatch( private void commitMultiKeyWorkBatch( ComputationState computationState, List workBatch, - List workItemCommits) { + List workItemCommits, + List bundleOutputMessages, + List bundlePubsubMessages) { checkState(!workBatch.isEmpty()); checkState(workBatch.size() == workItemCommits.size()); Windmill.MultiKeyWorkItemCommitRequest.Builder multiKeyBuilder = Windmill.MultiKeyWorkItemCommitRequest.newBuilder(); + if (!bundleOutputMessages.isEmpty()) { + multiKeyBuilder.addAllOutputMessages(bundleOutputMessages); + } + if (!bundlePubsubMessages.isEmpty()) { + multiKeyBuilder.addAllPubsubMessages(bundlePubsubMessages); + } Work primaryWork = workBatch.get(0); Work.KeyGroup keyGroup = primaryWork.getKeyGroup(); @@ -525,16 +550,27 @@ abstract static class ExecuteWorkResult { static ExecuteWorkResult create( List workBatch, List workItemCommits, + List bundleOutputMessages, + List bundlePubsubMessages, Map> finalizationCallbacks, long stateBytesRead) { return new AutoValue_StreamingWorkScheduler_ExecuteWorkResult( - workBatch, workItemCommits, finalizationCallbacks, stateBytesRead); + workBatch, + workItemCommits, + bundleOutputMessages, + bundlePubsubMessages, + finalizationCallbacks, + stateBytesRead); } abstract List workBatch(); abstract List workItemCommits(); + abstract List bundleOutputMessages(); + + abstract List bundlePubsubMessages(); + // Map> abstract Map> finalizationCallbacks(); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java index d3a424758f66..0adbdade7225 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorFactoryTest.java @@ -451,6 +451,9 @@ public long add(Integer outputElem) { return 4; } + @Override + public void finishKey(@Nullable Object key) throws IOException {} + @Override public void close() {} diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSinkTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSinkTest.java index d8822ce4937b..befa6d7a7482 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSinkTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSinkTest.java @@ -161,4 +161,154 @@ public void testWriteDynamicDestinations() throws Exception { .build(); assertEquals(expectedCommit, outputBuilder.build()); } + + @Test + public void testSingleKey_finishKeyDoesNotFlush_closeAttachesToKey() throws Exception { + when(mockContext.multiKeyBundleEnabled()).thenReturn(false); + + Windmill.WorkItemCommitRequest.Builder outputBuilder = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("key")) + .setWorkToken(0); + when(mockContext.getOutputBuilder()).thenReturn(outputBuilder); + + Map spec = new HashMap<>(); + spec.put(PropertyNames.OBJECT_TYPE_NAME, "PubsubDynamicSink"); + spec.put(PropertyNames.PUBSUB_TIMESTAMP_ATTRIBUTE, "ts"); + spec.put(PropertyNames.PUBSUB_ID_ATTRIBUTE, "id"); + + CloudObject cloudSinkSpec = CloudObject.fromSpec(spec); + PubsubDynamicSink sink = + (PubsubDynamicSink) + SinkRegistry.defaultRegistry() + .create( + cloudSinkSpec, + WindowedValues.getFullCoder(VoidCoder.of(), IntervalWindow.getCoder()), + null, + mockContext, + null) + .getUnderlyingSink(); + + Sink.SinkWriter> writer = sink.writer(); + byte[] payload0 = "msg0".getBytes(StandardCharsets.UTF_8); + byte[] payload1 = "msg1".getBytes(StandardCharsets.UTF_8); + + writer.add( + WindowedValues.timestampedValueInGlobalWindow( + new PubsubMessage(payload0, null).withTopic("topic1"), new Instant(0))); + + // In single-key mode, finishKey does not flush + writer.finishKey("key"); + assertEquals(0, outputBuilder.getPubsubMessagesCount()); + + // close flushes all outputs into the key's outputBuilder + writer.add( + WindowedValues.timestampedValueInGlobalWindow( + new PubsubMessage(payload1, null).withTopic("topic2"), new Instant(1000))); + writer.close(); + + assertEquals(2, outputBuilder.getPubsubMessagesCount()); + Map bundlesByTopic = new HashMap<>(); + for (Windmill.PubSubMessageBundle bundle : outputBuilder.getPubsubMessagesList()) { + bundlesByTopic.put(bundle.getTopic(), bundle); + } + assertEquals(1, bundlesByTopic.get("topic1").getMessagesCount()); + assertEquals(1, bundlesByTopic.get("topic2").getMessagesCount()); + } + + @Test + public void testMultiKey_finishKeyFlushesToKey_closeFlushesToBundleLevel() throws Exception { + when(mockContext.multiKeyBundleEnabled()).thenReturn(true); + + Windmill.WorkItemCommitRequest.Builder outputBuilder = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1); + when(mockContext.getOutputBuilder()).thenReturn(outputBuilder); + + Map spec = new HashMap<>(); + spec.put(PropertyNames.OBJECT_TYPE_NAME, "PubsubDynamicSink"); + spec.put(PropertyNames.PUBSUB_TIMESTAMP_ATTRIBUTE, "ts"); + spec.put(PropertyNames.PUBSUB_ID_ATTRIBUTE, "id"); + + CloudObject cloudSinkSpec = CloudObject.fromSpec(spec); + PubsubDynamicSink sink = + (PubsubDynamicSink) + SinkRegistry.defaultRegistry() + .create( + cloudSinkSpec, + WindowedValues.getFullCoder(VoidCoder.of(), IntervalWindow.getCoder()), + null, + mockContext, + null) + .getUnderlyingSink(); + + Sink.SinkWriter> writer = sink.writer(); + byte[] payload0 = "msg0".getBytes(StandardCharsets.UTF_8); + byte[] payload1 = "msg1".getBytes(StandardCharsets.UTF_8); + + writer.add( + WindowedValues.timestampedValueInGlobalWindow( + new PubsubMessage(payload0, null).withTopic("topic1"), new Instant(0))); + + // In multi-key mode, finishKey flushes to the active key's outputBuilder + writer.finishKey("key1"); + assertEquals(1, outputBuilder.getPubsubMessagesCount()); + assertEquals("topic1", outputBuilder.getPubsubMessages(0).getTopic()); + + // Messages added during finishBundle are flushed to bundle-level outputs in close() + writer.add( + WindowedValues.timestampedValueInGlobalWindow( + new PubsubMessage(payload1, null).withTopic("topic2"), new Instant(1000))); + writer.close(); + + // Verify key's outputBuilder still only has 1 bundle from finishKey + assertEquals(1, outputBuilder.getPubsubMessagesCount()); + + // Verify mockContext.addBundlePubsubMessages was called with the bundle from close() + org.mockito.ArgumentCaptor captor = + org.mockito.ArgumentCaptor.forClass(Windmill.PubSubMessageBundle.class); + org.mockito.Mockito.verify(mockContext).addBundlePubsubMessages(captor.capture()); + Windmill.PubSubMessageBundle bundleLevel = captor.getValue(); + assertEquals("topic2", bundleLevel.getTopic()); + assertEquals(1, bundleLevel.getMessagesCount()); + } + + @Test + public void testAbort_doesNotFlushToContext() throws Exception { + Windmill.WorkItemCommitRequest.Builder outputBuilder = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("key")) + .setWorkToken(0); + when(mockContext.getOutputBuilder()).thenReturn(outputBuilder); + + Map spec = new HashMap<>(); + spec.put(PropertyNames.OBJECT_TYPE_NAME, "PubsubDynamicSink"); + spec.put(PropertyNames.PUBSUB_TIMESTAMP_ATTRIBUTE, "ts"); + spec.put(PropertyNames.PUBSUB_ID_ATTRIBUTE, "id"); + + CloudObject cloudSinkSpec = CloudObject.fromSpec(spec); + PubsubDynamicSink sink = + (PubsubDynamicSink) + SinkRegistry.defaultRegistry() + .create( + cloudSinkSpec, + WindowedValues.getFullCoder(VoidCoder.of(), IntervalWindow.getCoder()), + null, + mockContext, + null) + .getUnderlyingSink(); + + Sink.SinkWriter> writer = sink.writer(); + byte[] payload0 = "msg0".getBytes(StandardCharsets.UTF_8); + writer.add( + WindowedValues.timestampedValueInGlobalWindow( + new PubsubMessage(payload0, null).withTopic("topic1"), new Instant(0))); + + writer.abort(); + + assertEquals(0, outputBuilder.getPubsubMessagesCount()); + org.mockito.Mockito.verify(mockContext, org.mockito.Mockito.never()) + .addBundlePubsubMessages(org.mockito.ArgumentMatchers.any()); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java index 5327cd172410..7fe7e40b596a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java @@ -180,4 +180,134 @@ public void testExceptionAfterEncoding() throws Exception { CoderException.class, () -> writer.add(WindowedValues.timestampedValueInGlobalWindow("e0", new Instant(0)))); } + + @Test + public void testSingleKey_finishKeyDoesNotFlush_closeAttachesToKey() throws Exception { + when(mockContext.multiKeyBundleEnabled()).thenReturn(false); + + Windmill.WorkItemCommitRequest.Builder outputBuilder = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("key")) + .setWorkToken(0); + when(mockContext.getOutputBuilder()).thenReturn(outputBuilder); + + Map spec = new HashMap<>(); + spec.put(PropertyNames.OBJECT_TYPE_NAME, ""); + spec.put(PropertyNames.PUBSUB_TOPIC, "topic"); + spec.put(PropertyNames.PUBSUB_TIMESTAMP_ATTRIBUTE, "ts"); + spec.put(PropertyNames.PUBSUB_ID_ATTRIBUTE, "id"); + CloudObject cloudSinkSpec = CloudObject.fromSpec(spec); + PubsubSink.Factory factory = new PubsubSink.Factory(); + PubsubSink sink = + (PubsubSink) + factory.create( + cloudSinkSpec, + WindowedValues.getFullCoder(StringUtf8Coder.of(), IntervalWindow.getCoder()), + null, + mockContext, + null); + + Sink.SinkWriter> writer = sink.writer(); + writer.add(WindowedValues.timestampedValueInGlobalWindow("e0", new Instant(0))); + + // In single key mode, finishKey should not flush + writer.finishKey("key"); + assertEquals(0, outputBuilder.getPubsubMessagesCount()); + + // close should flush and attach to the key's outputBuilder + writer.add(WindowedValues.timestampedValueInGlobalWindow("e1", new Instant(1000))); + writer.close(); + + assertEquals(1, outputBuilder.getPubsubMessagesCount()); + Windmill.PubSubMessageBundle bundle = outputBuilder.getPubsubMessages(0); + assertEquals("topic", bundle.getTopic()); + assertEquals(2, bundle.getMessagesCount()); + assertEquals("e0", bundle.getMessages(0).getData().toStringUtf8()); + assertEquals("e1", bundle.getMessages(1).getData().toStringUtf8()); + } + + @Test + public void testMultiKey_finishKeyFlushesToKey_closeFlushesToBundleLevel() throws Exception { + when(mockContext.multiKeyBundleEnabled()).thenReturn(true); + + Windmill.WorkItemCommitRequest.Builder outputBuilder = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1); + when(mockContext.getOutputBuilder()).thenReturn(outputBuilder); + + Map spec = new HashMap<>(); + spec.put(PropertyNames.OBJECT_TYPE_NAME, ""); + spec.put(PropertyNames.PUBSUB_TOPIC, "topic"); + spec.put(PropertyNames.PUBSUB_TIMESTAMP_ATTRIBUTE, "ts"); + spec.put(PropertyNames.PUBSUB_ID_ATTRIBUTE, "id"); + CloudObject cloudSinkSpec = CloudObject.fromSpec(spec); + PubsubSink.Factory factory = new PubsubSink.Factory(); + PubsubSink sink = + (PubsubSink) + factory.create( + cloudSinkSpec, + WindowedValues.getFullCoder(StringUtf8Coder.of(), IntervalWindow.getCoder()), + null, + mockContext, + null); + + Sink.SinkWriter> writer = sink.writer(); + writer.add(WindowedValues.timestampedValueInGlobalWindow("e0", new Instant(0))); + + // In multi-key mode, finishKey flushes to the active key's outputBuilder + writer.finishKey("key1"); + assertEquals(1, outputBuilder.getPubsubMessagesCount()); + assertEquals("e0", outputBuilder.getPubsubMessages(0).getMessages(0).getData().toStringUtf8()); + + // Messages added during finishBundle are flushed to bundle-level outputs in close() + writer.add(WindowedValues.timestampedValueInGlobalWindow("e1", new Instant(1000))); + writer.close(); + + // Verify key's outputBuilder still only has 1 bundle from finishKey + assertEquals(1, outputBuilder.getPubsubMessagesCount()); + + // Verify mockContext.addBundlePubsubMessages was called with the bundle from close() + org.mockito.ArgumentCaptor captor = + org.mockito.ArgumentCaptor.forClass(Windmill.PubSubMessageBundle.class); + org.mockito.Mockito.verify(mockContext).addBundlePubsubMessages(captor.capture()); + Windmill.PubSubMessageBundle bundleLevel = captor.getValue(); + assertEquals("topic", bundleLevel.getTopic()); + assertEquals(1, bundleLevel.getMessagesCount()); + assertEquals("e1", bundleLevel.getMessages(0).getData().toStringUtf8()); + } + + @Test + public void testAbort_doesNotFlushToContext() throws Exception { + Windmill.WorkItemCommitRequest.Builder outputBuilder = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("key")) + .setWorkToken(0); + when(mockContext.getOutputBuilder()).thenReturn(outputBuilder); + + Map spec = new HashMap<>(); + spec.put(PropertyNames.OBJECT_TYPE_NAME, ""); + spec.put(PropertyNames.PUBSUB_TOPIC, "topic"); + spec.put(PropertyNames.PUBSUB_TIMESTAMP_ATTRIBUTE, "ts"); + spec.put(PropertyNames.PUBSUB_ID_ATTRIBUTE, "id"); + CloudObject cloudSinkSpec = CloudObject.fromSpec(spec); + PubsubSink.Factory factory = new PubsubSink.Factory(); + PubsubSink sink = + (PubsubSink) + factory.create( + cloudSinkSpec, + WindowedValues.getFullCoder(StringUtf8Coder.of(), IntervalWindow.getCoder()), + null, + mockContext, + null); + + Sink.SinkWriter> writer = sink.writer(); + writer.add(WindowedValues.timestampedValueInGlobalWindow("e0", new Instant(0))); + + writer.abort(); + + assertEquals(0, outputBuilder.getPubsubMessagesCount()); + org.mockito.Mockito.verify(mockContext, org.mockito.Mockito.never()) + .addBundlePubsubMessages(org.mockito.ArgumentMatchers.any()); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java index c48b30ecf640..1f7a5d0e6c82 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java @@ -1650,6 +1650,471 @@ public void testMultiKeyCommit_queuedWorkItemFailsAndSubsequentWorkItemPickedUp( worker.stop(); } + private void runMultiKeyCombinationTest( + Map> processOutputs, List> finishBundleOutputs) + throws Exception { + org.junit.Assume.assumeTrue( + "Multi-key bundling is only supported in Streaming Engine", streamingEngine); + server.clearCommitsReceived(); + StreamingDataflowWorker worker = + makeMultiKeyEnabledWorker( + new ConfigurableMultiKeyDoFn(processOutputs, finishBundleOutputs)); + worker.start(); + + String batchInputText = + "work {" + + " computation_id: \"" + + DEFAULT_COMPUTATION_ID + + "\"" + + " input_data_watermark: 0" + + " work {" + + " key: \"key1\"" + + " sharding_key: 1" + + " work_token: 1" + + " cache_token: 2" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data1\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"key2\"" + + " sharding_key: 2" + + " work_token: 2" + + " cache_token: 3" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data2\"" + + " }" + + " }" + + " }" + + "}"; + Windmill.GetWorkResponse batchInput = + buildInput( + batchInputText, + CoderUtils.encodeToByteArray( + CollectionCoder.of(IntervalWindow.getCoder()), + Collections.singletonList(DEFAULT_WINDOW))); + + server.whenGetDataCalled().answerByDefault(StreamingDataflowWorkerTest::emptyDataResponder); + server.whenGetWorkCalled().thenReturn(batchInput); + + Map result = server.waitForAndGetCommits(2); + assertEquals(2, result.size()); + + // Verify Key 1 commit: should only contain process outputs for key1 + assertTrue(result.containsKey(1L)); + Windmill.WorkItemCommitRequest commit1 = result.get(1L); + assertEquals("key1", commit1.getKey().toStringUtf8()); + List expectedKey1Outputs = processOutputs.getOrDefault("key1", Collections.emptyList()); + if (expectedKey1Outputs.isEmpty()) { + assertEquals(0, commit1.getOutputMessagesCount()); + } else { + assertEquals(1, commit1.getOutputMessagesCount()); + Windmill.OutputMessageBundle outputBundle1 = commit1.getOutputMessages(0); + assertEquals(DEFAULT_DESTINATION_STREAM_ID, outputBundle1.getDestinationStreamId()); + assertEquals(1, outputBundle1.getBundlesCount()); + Windmill.KeyedMessageBundle keyedBundle1 = outputBundle1.getBundles(0); + assertEquals("key1", keyedBundle1.getKey().toStringUtf8()); + assertEquals(expectedKey1Outputs.size(), keyedBundle1.getMessagesCount()); + for (int i = 0; i < expectedKey1Outputs.size(); i++) { + assertEquals( + expectedKey1Outputs.get(i), keyedBundle1.getMessages(i).getData().toStringUtf8()); + } + } + // Verify Key 2 commit: should only contain process outputs for key2 (NOT finishBundle outputs) + assertTrue(result.containsKey(2L)); + Windmill.WorkItemCommitRequest commit2 = result.get(2L); + assertEquals("key2", commit2.getKey().toStringUtf8()); + List expectedKey2Outputs = processOutputs.getOrDefault("key2", Collections.emptyList()); + if (expectedKey2Outputs.isEmpty()) { + assertEquals(0, commit2.getOutputMessagesCount()); + } else { + assertEquals(1, commit2.getOutputMessagesCount()); + Windmill.OutputMessageBundle outputBundle2 = commit2.getOutputMessages(0); + assertEquals(DEFAULT_DESTINATION_STREAM_ID, outputBundle2.getDestinationStreamId()); + assertEquals(1, outputBundle2.getBundlesCount()); + Windmill.KeyedMessageBundle keyedBundle2 = outputBundle2.getBundles(0); + assertEquals("key2", keyedBundle2.getKey().toStringUtf8()); + assertEquals(expectedKey2Outputs.size(), keyedBundle2.getMessagesCount()); + for (int i = 0; i < expectedKey2Outputs.size(); i++) { + assertEquals( + expectedKey2Outputs.get(i), keyedBundle2.getMessages(i).getData().toStringUtf8()); + } + } + + // Verify MultiKey commit: should contain all finishBundle outputs at the bundle level + List multiKeyCommits = + server.getMultiKeyCommitsReceived(); + assertEquals(1, multiKeyCommits.size()); + Windmill.MultiKeyWorkItemCommitRequest multiKeyCommit = multiKeyCommits.get(0); + if (finishBundleOutputs.isEmpty()) { + assertEquals(0, multiKeyCommit.getOutputMessagesCount()); + } else { + assertEquals(1, multiKeyCommit.getOutputMessagesCount()); + Windmill.OutputMessageBundle outputBundle_fb = multiKeyCommit.getOutputMessages(0); + assertEquals(DEFAULT_DESTINATION_STREAM_ID, outputBundle_fb.getDestinationStreamId()); + Map> expectedFbByKey = new HashMap<>(); + for (KV kv : finishBundleOutputs) { + expectedFbByKey.computeIfAbsent(kv.getKey(), k -> new ArrayList<>()).add(kv.getValue()); + } + assertEquals(expectedFbByKey.size(), outputBundle_fb.getBundlesCount()); + for (Windmill.KeyedMessageBundle keyedBundle : outputBundle_fb.getBundlesList()) { + String key = keyedBundle.getKey().toStringUtf8(); + assertTrue(expectedFbByKey.containsKey(key)); + List expectedValues = expectedFbByKey.get(key); + assertEquals(expectedValues.size(), keyedBundle.getMessagesCount()); + for (int i = 0; i < expectedValues.size(); i++) { + assertEquals(expectedValues.get(i), keyedBundle.getMessages(i).getData().toStringUtf8()); + } + } + } + + worker.stop(); + } + + @Test + public void testMultiKey_allCombinationsOfProcessAndFinishBundleOutputs() throws Exception { + if (!streamingEngine) { + return; + } + List> key1Options = + List.of(Collections.emptyList(), List.of("k1_out1"), List.of("k1_out1", "k1_out2")); + + List> key2Options = + List.of(Collections.emptyList(), List.of("k2_out1"), List.of("k2_out1", "k2_out2")); + + List>> finishBundleOptions = + List.of( + Collections.emptyList(), + List.of(KV.of("fb_key", "fb_val1")), + List.of(KV.of("fb_key1", "fb_val1"), KV.of("fb_key2", "fb_val2"))); + + for (List k1Out : key1Options) { + for (List k2Out : key2Options) { + for (List> fbOut : finishBundleOptions) { + Map> processOutputs = new HashMap<>(); + if (!k1Out.isEmpty()) { + processOutputs.put("key1", k1Out); + } + if (!k2Out.isEmpty()) { + processOutputs.put("key2", k2Out); + } + runMultiKeyCombinationTest(processOutputs, fbOut); + } + } + } + } + + @Test + public void testMultiKey_productionsAttachedToEachKey() throws Exception { + runMultiKeyCombinationTest( + ImmutableMap.of( + "key1", ImmutableList.of("data1"), + "key2", ImmutableList.of("data2")), + Collections.emptyList()); + } + + @Test + public void testMultiKey_finishBundleProductionsAttachedToBundleLevel() throws Exception { + runMultiKeyCombinationTest( + Collections.emptyMap(), ImmutableList.of(KV.of("finish_key", "finish_value"))); + } + + @Test + public void testMultiKey_processElementsAndFinishBundle_bufferedMessageAttachedToBundleLevel() + throws Exception { + runMultiKeyCombinationTest( + ImmutableMap.of( + "key1", ImmutableList.of("data1"), + "key2", ImmutableList.of("data2")), + ImmutableList.of(KV.of("finish_key", "buffer_emitted"))); + } + + @Test + public void testSingleKey_processAndFinishBundleOutputsAttachedToSameKey() throws Exception { + KvCoder kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); + List instructions = + Arrays.asList( + makeSourceInstruction(kvCoder), + makeDoFnInstruction( + new ConfigurableMultiKeyDoFn( + ImmutableMap.of("key1", ImmutableList.of("data1")), + ImmutableList.of(KV.of("finish_key", "finish_value"))), + 0, + kvCoder), + makeSinkInstruction(kvCoder, 1)); + + StreamingDataflowWorker worker = + makeWorker(defaultWorkerParams().setInstructions(instructions).build()); + worker.start(); + + String input = + "work {" + + " computation_id: \"" + + DEFAULT_COMPUTATION_ID + + "\"" + + " input_data_watermark: 0" + + " work {" + + " key: \"key1\"" + + " sharding_key: 1" + + " work_token: 1" + + " cache_token: 2" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data1\"" + + " }" + + " }" + + " }" + + "}"; + Windmill.GetWorkResponse workResponse = + buildInput( + input, + CoderUtils.encodeToByteArray( + CollectionCoder.of(IntervalWindow.getCoder()), + Collections.singletonList(DEFAULT_WINDOW))); + + server.clearCommitsReceived(); + server.whenGetDataCalled().answerByDefault(StreamingDataflowWorkerTest::emptyDataResponder); + server.whenGetWorkCalled().thenReturn(workResponse); + + Map result = server.waitForAndGetCommits(1); + assertEquals(1, result.size()); + + assertTrue(result.containsKey(1L)); + Windmill.WorkItemCommitRequest commit = result.get(1L); + assertEquals("key1", commit.getKey().toStringUtf8()); + // In single-key mode, finishKey does not flush; close flushes all outputs into the single key's + // commit + assertEquals(1, commit.getOutputMessagesCount()); + Windmill.OutputMessageBundle outputBundle = commit.getOutputMessages(0); + assertEquals(DEFAULT_DESTINATION_STREAM_ID, outputBundle.getDestinationStreamId()); + assertEquals(2, outputBundle.getBundlesCount()); + Map outputsByKey = new HashMap<>(); + for (Windmill.KeyedMessageBundle bundle : outputBundle.getBundlesList()) { + assertEquals(1, bundle.getMessagesCount()); + outputsByKey.put( + bundle.getKey().toStringUtf8(), bundle.getMessages(0).getData().toStringUtf8()); + } + assertEquals("data1", outputsByKey.get("key1")); + assertEquals("finish_value", outputsByKey.get("finish_key")); + + worker.stop(); + } + + @Test + public void testSingleKey_multiKeyBundleEnabled_finishBundleAttachesToBundleLevel() + throws Exception { + org.junit.Assume.assumeTrue( + "Multi-key bundling is only supported in Streaming Engine", streamingEngine); + + server.clearCommitsReceived(); + StreamingDataflowWorker worker = + makeMultiKeyEnabledWorker( + new ConfigurableMultiKeyDoFn( + ImmutableMap.of("key1", ImmutableList.of("data1")), + ImmutableList.of(KV.of("finish_key", "finish_value")))); + worker.start(); + + String input = + "work {" + + " computation_id: \"" + + DEFAULT_COMPUTATION_ID + + "\"" + + " input_data_watermark: 0" + + " work {" + + " key: \"key1\"" + + " sharding_key: 1" + + " work_token: 1" + + " cache_token: 2" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data1\"" + + " }" + + " }" + + " }" + + "}"; + Windmill.GetWorkResponse workResponse = + buildInput( + input, + CoderUtils.encodeToByteArray( + CollectionCoder.of(IntervalWindow.getCoder()), + Collections.singletonList(DEFAULT_WINDOW))); + + server.whenGetDataCalled().answerByDefault(StreamingDataflowWorkerTest::emptyDataResponder); + server.whenGetWorkCalled().thenReturn(workResponse); + + Map result = server.waitForAndGetCommits(1); + assertEquals(1, result.size()); + + assertTrue(result.containsKey(1L)); + Windmill.WorkItemCommitRequest commit1 = result.get(1L); + assertEquals("key1", commit1.getKey().toStringUtf8()); + // In multi-key mode, key1 only contains its own process outputs + assertEquals(1, commit1.getOutputMessagesCount()); + Windmill.OutputMessageBundle outputBundle1 = commit1.getOutputMessages(0); + assertEquals(DEFAULT_DESTINATION_STREAM_ID, outputBundle1.getDestinationStreamId()); + assertEquals(1, outputBundle1.getBundlesCount()); + assertEquals("key1", outputBundle1.getBundles(0).getKey().toStringUtf8()); + assertEquals("data1", outputBundle1.getBundles(0).getMessages(0).getData().toStringUtf8()); + + // finishBundle outputs are flushed to the bundle level of MultiKeyWorkItemCommitRequest + List multiKeyCommits = + server.getMultiKeyCommitsReceived(); + assertEquals(1, multiKeyCommits.size()); + Windmill.MultiKeyWorkItemCommitRequest multiKeyCommit = multiKeyCommits.get(0); + assertEquals(1, multiKeyCommit.getOutputMessagesCount()); + Windmill.OutputMessageBundle bundleLevel = multiKeyCommit.getOutputMessages(0); + assertEquals(DEFAULT_DESTINATION_STREAM_ID, bundleLevel.getDestinationStreamId()); + assertEquals(1, bundleLevel.getBundlesCount()); + assertEquals("finish_key", bundleLevel.getBundles(0).getKey().toStringUtf8()); + assertEquals("finish_value", bundleLevel.getBundles(0).getMessages(0).getData().toStringUtf8()); + + worker.stop(); + } + + @Test + public void testMultiKey_threeKeys_withIntermediateEmptyKey() throws Exception { + org.junit.Assume.assumeTrue( + "Multi-key bundling is only supported in Streaming Engine", streamingEngine); + + server.clearCommitsReceived(); + StreamingDataflowWorker worker = + makeMultiKeyEnabledWorker( + new ConfigurableMultiKeyDoFn( + ImmutableMap.of( + "key1", ImmutableList.of("data1"), + "key3", ImmutableList.of("data3")), + ImmutableList.of(KV.of("finish_key", "finish_value")))); + worker.start(); + + String batchInputText = + "work {" + + " computation_id: \"" + + DEFAULT_COMPUTATION_ID + + "\"" + + " input_data_watermark: 0" + + " work {" + + " key: \"key1\"" + + " sharding_key: 1" + + " work_token: 1" + + " cache_token: 2" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data1\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"key2\"" + + " sharding_key: 2" + + " work_token: 2" + + " cache_token: 3" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data2\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"key3\"" + + " sharding_key: 3" + + " work_token: 3" + + " cache_token: 4" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data3\"" + + " }" + + " }" + + " }" + + "}"; + Windmill.GetWorkResponse batchInput = + buildInput( + batchInputText, + CoderUtils.encodeToByteArray( + CollectionCoder.of(IntervalWindow.getCoder()), + Collections.singletonList(DEFAULT_WINDOW))); + + server.whenGetDataCalled().answerByDefault(StreamingDataflowWorkerTest::emptyDataResponder); + server.whenGetWorkCalled().thenReturn(batchInput); + + Map result = server.waitForAndGetCommits(3); + assertEquals(3, result.size()); + + // Verify Key 1 commit + assertTrue(result.containsKey(1L)); + Windmill.WorkItemCommitRequest commit1 = result.get(1L); + assertEquals("key1", commit1.getKey().toStringUtf8()); + assertEquals(1, commit1.getOutputMessagesCount()); + assertEquals( + "data1", + commit1.getOutputMessages(0).getBundles(0).getMessages(0).getData().toStringUtf8()); + + // Verify Key 2 commit (empty outputs) + assertTrue(result.containsKey(2L)); + Windmill.WorkItemCommitRequest commit2 = result.get(2L); + assertEquals("key2", commit2.getKey().toStringUtf8()); + assertEquals(0, commit2.getOutputMessagesCount()); + + // Verify Key 3 commit + assertTrue(result.containsKey(3L)); + Windmill.WorkItemCommitRequest commit3 = result.get(3L); + assertEquals("key3", commit3.getKey().toStringUtf8()); + assertEquals(1, commit3.getOutputMessagesCount()); + assertEquals( + "data3", + commit3.getOutputMessages(0).getBundles(0).getMessages(0).getData().toStringUtf8()); + + // Verify MultiKey commit: should contain all 3 requests and finishBundle outputs at bundle + // level + List multiKeyCommits = + server.getMultiKeyCommitsReceived(); + assertEquals(1, multiKeyCommits.size()); + Windmill.MultiKeyWorkItemCommitRequest multiKeyCommit = multiKeyCommits.get(0); + assertEquals(3, multiKeyCommit.getRequestsCount()); + assertEquals(1, multiKeyCommit.getOutputMessagesCount()); + Windmill.OutputMessageBundle bundleLevel = multiKeyCommit.getOutputMessages(0); + assertEquals("finish_key", bundleLevel.getBundles(0).getKey().toStringUtf8()); + assertEquals("finish_value", bundleLevel.getBundles(0).getMessages(0).getData().toStringUtf8()); + + worker.stop(); + } + private StreamingDataflowWorker makeMultiKeyEnabledWorker() { return makeMultiKeyEnabledWorker(new WorkDoFn()); } @@ -5456,6 +5921,35 @@ public void processElement(ProcessContext c) { } } + static class ConfigurableMultiKeyDoFn extends DoFn, KV> { + private final Map> processOutputs; + private final List> finishBundleOutputs; + + ConfigurableMultiKeyDoFn( + Map> processOutputs, List> finishBundleOutputs) { + this.processOutputs = processOutputs; + this.finishBundleOutputs = finishBundleOutputs; + } + + @ProcessElement + public void processElement(ProcessContext c) { + String key = c.element().getKey(); + List outputs = processOutputs.get(key); + if (outputs != null) { + for (String output : outputs) { + c.output(KV.of(key, output)); + } + } + } + + @FinishBundle + public void finishBundle(FinishBundleContext c) { + for (KV output : finishBundleOutputs) { + c.output(output, new Instant(0), DEFAULT_WINDOW); + } + } + } + @AutoValue abstract static class StreamingDataflowWorkerTestParams { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ExecutorTestUtils.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ExecutorTestUtils.java index ac7c787b1d26..1872881bbf92 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ExecutorTestUtils.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/ExecutorTestUtils.java @@ -30,6 +30,7 @@ import org.apache.beam.runners.dataflow.worker.counters.CounterSet; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; /** Utilities for tests. */ @SuppressWarnings({ @@ -169,6 +170,9 @@ public long add(String outputElem) { return outputElem.length(); } + @Override + public void finishKey(@Nullable Object key) throws IOException {} + @Override public void close() { assertFalse(closed); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperationTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperationTest.java index 2a83cc56cc0a..95c7863d79ff 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperationTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WriteOperationTest.java @@ -21,9 +21,11 @@ import static org.hamcrest.CoreMatchers.hasItems; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -182,4 +184,61 @@ public void testWriteOperationContext() throws Exception { inOrder.verify(sinkWriter).close(); inOrder.verify(finishCloseable).close(); } + + @Test + public void testFinishKey() throws Exception { + OperationContext mockContext = mock(OperationContext.class); + when(mockContext.counterFactory()).thenReturn(counterSet); + when(mockContext.nameContext()).thenReturn(NameContextsForTests.nameContextForTest()); + Closeable startCloseable = mock(Closeable.class); + Closeable processCloseable = mock(Closeable.class); + when(mockContext.enterStart()).thenReturn(startCloseable); + when(mockContext.enterProcess()).thenReturn(processCloseable); + + Sink sink = mock(Sink.class); + Sink.SinkWriter sinkWriter = mock(Sink.SinkWriter.class); + when(sink.writer()).thenReturn(sinkWriter); + + WriteOperation operation = WriteOperation.forTest(sink, mockContext); + operation.start(); + operation.finishKey("key1"); + + verify(mockContext).enterProcess(); + verify(sinkWriter).finishKey("key1"); + verify(processCloseable).close(); + } + + @Test + public void testFinishKey_unstarted_throwsException() throws Exception { + Sink sink = mock(Sink.class); + WriteOperation operation = WriteOperation.forTest(sink, context); + + assertThrows(AssertionError.class, () -> operation.finishKey("key1")); + } + + @Test + public void testFinishKey_nullKey() throws Exception { + Sink sink = mock(Sink.class); + Sink.SinkWriter sinkWriter = mock(Sink.SinkWriter.class); + when(sink.writer()).thenReturn(sinkWriter); + + WriteOperation operation = WriteOperation.forTest(sink, context); + operation.start(); + operation.finishKey(null); + + verify(sinkWriter).finishKey(null); + } + + @Test + public void testFinishKey_exceptionPropagates() throws Exception { + Sink sink = mock(Sink.class); + Sink.SinkWriter sinkWriter = mock(Sink.SinkWriter.class); + when(sink.writer()).thenReturn(sinkWriter); + doThrow(new IOException("finishKey error")).when(sinkWriter).finishKey("key1"); + + WriteOperation operation = WriteOperation.forTest(sink, context); + operation.start(); + + assertThrows(IOException.class, () -> operation.finishKey("key1")); + } } From 4cbbf843eeb1dcc8598227671113286f4288ce60 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Thu, 3 Sep 2026 12:06:47 +0000 Subject: [PATCH 2/3] improve tests --- .../worker/PubsubDynamicSinkTest.java | 26 +++++++++++--- .../dataflow/worker/PubsubSinkTest.java | 13 ++++--- .../worker/StreamingDataflowWorkerTest.java | 35 +++---------------- 3 files changed, 33 insertions(+), 41 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSinkTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSinkTest.java index befa6d7a7482..bf548ffeb8b6 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSinkTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSinkTest.java @@ -18,6 +18,9 @@ package org.apache.beam.runners.dataflow.worker; import static org.junit.Assert.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import java.nio.charset.StandardCharsets; @@ -41,6 +44,7 @@ import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.MockitoAnnotations; @@ -214,6 +218,12 @@ public void testSingleKey_finishKeyDoesNotFlush_closeAttachesToKey() throws Exce } assertEquals(1, bundlesByTopic.get("topic1").getMessagesCount()); assertEquals(1, bundlesByTopic.get("topic2").getMessagesCount()); + Pubsub.PubsubMessage pubsubMsg0 = + Pubsub.PubsubMessage.parseFrom(bundlesByTopic.get("topic1").getMessages(0).getData()); + assertEquals(ByteString.copyFrom(payload0), pubsubMsg0.getData()); + Pubsub.PubsubMessage pubsubMsg1 = + Pubsub.PubsubMessage.parseFrom(bundlesByTopic.get("topic2").getMessages(0).getData()); + assertEquals(ByteString.copyFrom(payload1), pubsubMsg1.getData()); } @Test @@ -255,6 +265,10 @@ public void testMultiKey_finishKeyFlushesToKey_closeFlushesToBundleLevel() throw writer.finishKey("key1"); assertEquals(1, outputBuilder.getPubsubMessagesCount()); assertEquals("topic1", outputBuilder.getPubsubMessages(0).getTopic()); + assertEquals(1, outputBuilder.getPubsubMessages(0).getMessagesCount()); + Pubsub.PubsubMessage pubsubMsg0 = + Pubsub.PubsubMessage.parseFrom(outputBuilder.getPubsubMessages(0).getMessages(0).getData()); + assertEquals(ByteString.copyFrom(payload0), pubsubMsg0.getData()); // Messages added during finishBundle are flushed to bundle-level outputs in close() writer.add( @@ -266,12 +280,15 @@ public void testMultiKey_finishKeyFlushesToKey_closeFlushesToBundleLevel() throw assertEquals(1, outputBuilder.getPubsubMessagesCount()); // Verify mockContext.addBundlePubsubMessages was called with the bundle from close() - org.mockito.ArgumentCaptor captor = - org.mockito.ArgumentCaptor.forClass(Windmill.PubSubMessageBundle.class); - org.mockito.Mockito.verify(mockContext).addBundlePubsubMessages(captor.capture()); + ArgumentCaptor captor = + ArgumentCaptor.forClass(Windmill.PubSubMessageBundle.class); + verify(mockContext).addBundlePubsubMessages(captor.capture()); Windmill.PubSubMessageBundle bundleLevel = captor.getValue(); assertEquals("topic2", bundleLevel.getTopic()); assertEquals(1, bundleLevel.getMessagesCount()); + Pubsub.PubsubMessage pubsubMsg1 = + Pubsub.PubsubMessage.parseFrom(bundleLevel.getMessages(0).getData()); + assertEquals(ByteString.copyFrom(payload1), pubsubMsg1.getData()); } @Test @@ -308,7 +325,6 @@ public void testAbort_doesNotFlushToContext() throws Exception { writer.abort(); assertEquals(0, outputBuilder.getPubsubMessagesCount()); - org.mockito.Mockito.verify(mockContext, org.mockito.Mockito.never()) - .addBundlePubsubMessages(org.mockito.ArgumentMatchers.any()); + verify(mockContext, never()).addBundlePubsubMessages(any()); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java index 7fe7e40b596a..80c2c9f8305e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/PubsubSinkTest.java @@ -19,6 +19,9 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import java.io.IOException; @@ -43,6 +46,7 @@ import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.MockitoAnnotations; @@ -268,9 +272,9 @@ public void testMultiKey_finishKeyFlushesToKey_closeFlushesToBundleLevel() throw assertEquals(1, outputBuilder.getPubsubMessagesCount()); // Verify mockContext.addBundlePubsubMessages was called with the bundle from close() - org.mockito.ArgumentCaptor captor = - org.mockito.ArgumentCaptor.forClass(Windmill.PubSubMessageBundle.class); - org.mockito.Mockito.verify(mockContext).addBundlePubsubMessages(captor.capture()); + ArgumentCaptor captor = + ArgumentCaptor.forClass(Windmill.PubSubMessageBundle.class); + verify(mockContext).addBundlePubsubMessages(captor.capture()); Windmill.PubSubMessageBundle bundleLevel = captor.getValue(); assertEquals("topic", bundleLevel.getTopic()); assertEquals(1, bundleLevel.getMessagesCount()); @@ -307,7 +311,6 @@ public void testAbort_doesNotFlushToContext() throws Exception { writer.abort(); assertEquals(0, outputBuilder.getPubsubMessagesCount()); - org.mockito.Mockito.verify(mockContext, org.mockito.Mockito.never()) - .addBundlePubsubMessages(org.mockito.ArgumentMatchers.any()); + verify(mockContext, never()).addBundlePubsubMessages(any()); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java index dad1e1feaf54..6fc19c303240 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java @@ -33,6 +33,7 @@ import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; +import static org.junit.Assume.assumeTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.nullable; @@ -1662,8 +1663,7 @@ public void testMultiKeyCommit_queuedWorkItemFailsAndSubsequentWorkItemPickedUp( private void runMultiKeyCombinationTest( Map> processOutputs, List> finishBundleOutputs) throws Exception { - org.junit.Assume.assumeTrue( - "Multi-key bundling is only supported in Streaming Engine", streamingEngine); + assumeTrue("Multi-key bundling is only supported in Streaming Engine", streamingEngine); server.clearCommitsReceived(); StreamingDataflowWorker worker = makeMultiKeyEnabledWorker( @@ -1826,31 +1826,6 @@ public void testMultiKey_allCombinationsOfProcessAndFinishBundleOutputs() throws } } - @Test - public void testMultiKey_productionsAttachedToEachKey() throws Exception { - runMultiKeyCombinationTest( - ImmutableMap.of( - "key1", ImmutableList.of("data1"), - "key2", ImmutableList.of("data2")), - Collections.emptyList()); - } - - @Test - public void testMultiKey_finishBundleProductionsAttachedToBundleLevel() throws Exception { - runMultiKeyCombinationTest( - Collections.emptyMap(), ImmutableList.of(KV.of("finish_key", "finish_value"))); - } - - @Test - public void testMultiKey_processElementsAndFinishBundle_bufferedMessageAttachedToBundleLevel() - throws Exception { - runMultiKeyCombinationTest( - ImmutableMap.of( - "key1", ImmutableList.of("data1"), - "key2", ImmutableList.of("data2")), - ImmutableList.of(KV.of("finish_key", "buffer_emitted"))); - } - @Test public void testSingleKey_processAndFinishBundleOutputsAttachedToSameKey() throws Exception { KvCoder kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); @@ -1929,8 +1904,7 @@ public void testSingleKey_processAndFinishBundleOutputsAttachedToSameKey() throw @Test public void testSingleKey_multiKeyBundleEnabled_finishBundleAttachesToBundleLevel() throws Exception { - org.junit.Assume.assumeTrue( - "Multi-key bundling is only supported in Streaming Engine", streamingEngine); + assumeTrue("Multi-key bundling is only supported in Streaming Engine", streamingEngine); server.clearCommitsReceived(); StreamingDataflowWorker worker = @@ -2004,8 +1978,7 @@ public void testSingleKey_multiKeyBundleEnabled_finishBundleAttachesToBundleLeve @Test public void testMultiKey_threeKeys_withIntermediateEmptyKey() throws Exception { - org.junit.Assume.assumeTrue( - "Multi-key bundling is only supported in Streaming Engine", streamingEngine); + assumeTrue("Multi-key bundling is only supported in Streaming Engine", streamingEngine); server.clearCommitsReceived(); StreamingDataflowWorker worker = From 9deeefcb85ac38fb9ac8d4c2ea899481a1bbf8f0 Mon Sep 17 00:00:00 2001 From: Arun Pandian Date: Thu, 3 Sep 2026 12:27:09 +0000 Subject: [PATCH 3/3] add checks --- .../dataflow/worker/StreamingModeExecutionContext.java | 1 + .../windmill/work/processing/StreamingWorkScheduler.java | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index 998efbd025b4..f090f414fa66 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -710,6 +710,7 @@ private void flushStateInternal() { } private void validateCommitRequestSize() { + // TODO: Validate size of outputs at MultiKeyWorkItemCommitRequest level. Windmill.WorkItemCommitRequest.Builder currentBuilder = getOutputBuilder(); Work currentWork = getWork(); long byteLimit = operationalLimits.getMaxWorkItemCommitBytes(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java index a891607a57d4..3d08fae0cf50 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java @@ -404,6 +404,12 @@ private void commitWorkBatch( commitMultiKeyWorkBatch( computationState, workBatch, workItemCommits, bundleOutputMessages, bundlePubsubMessages); } else { + checkState( + bundleOutputMessages.isEmpty(), + "bundleOutputMessages should be empty when calling commitSingleKeyWork"); + checkState( + bundlePubsubMessages.isEmpty(), + "bundlePubsubMessages should be empty when calling commitSingleKeyWork"); commitSingleKeyWork(computationState, workBatch.get(0), workItemCommits.get(0)); } }