[Dataflow Streaming] [Multi Key] Flush streaming sinks at key boundaries and bundle completion

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

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

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

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java

@@ -135,24 +135,27 @@ public PubsubSink<?> create(
 
   @Override
   public SinkWriter<WindowedValue<T>> writer() {
-    return new PubsubWriter(topic);
+    return new PubsubWriter();
   }
 
   /** The SinkWriter for a PubsubSink. */
   class PubsubWriter implements SinkWriter<WindowedValue<T>> {
     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<T> data) throws IOException {
       if (!stream.isEmpty()) {
@@ -187,18 +190,41 @@ public long add(WindowedValue<T> 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();
     }
   }
 

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<T> 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)) {

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

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java

@@ -196,6 +196,8 @@ public interface KeyTransitionListener {
   private @Nullable FailedWorkHandler onFailedWorkHandler;
 
   private List<Windmill.WorkItemCommitRequest.Builder> outputBuilders = Collections.emptyList();
+  private List<Windmill.OutputMessageBundle> bundleOutputMessages = Collections.emptyList();
+  private List<Windmill.PubSubMessageBundle> bundlePubsubMessages = Collections.emptyList();
 
   // Map<finalizerId, Pair<callbackExpiration, callback>>
   private Map<Long, Pair<Instant, Runnable>> finalizationCallbacks = Collections.emptyMap();
@@ -321,6 +323,8 @@ public void reset() {
     // these lists and maps are returned to callers after processing
     // don't clear and reuse, instead reset the reference.
     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.
@@ -354,6 +358,8 @@ public void start(
       throws CoderException {
     reset();
     this.outputBuilders = new ArrayList<>();
+    this.bundleOutputMessages = new ArrayList<>();
+    this.bundlePubsubMessages = new ArrayList<>();
     this.finalizationCallbacks = new HashMap<>();
     this.keyCoder = keyCoder;
     this.workExecutor = workExecutor;
@@ -704,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();
@@ -874,6 +881,26 @@ public List<Windmill.WorkItemCommitRequest> getWorkItemCommits() {
     return commits;
   }
 
+  public void addBundleOutputMessages(Windmill.OutputMessageBundle outputBundle) {
+    this.bundleOutputMessages.add(outputBundle);
+  }
+
+  public List<Windmill.OutputMessageBundle> getBundleOutputMessages() {
+    return bundleOutputMessages;
+  }
+
+  public void addBundlePubsubMessages(Windmill.PubSubMessageBundle pubsubBundle) {
+    this.bundlePubsubMessages.add(pubsubBundle);
+  }
+
+  public List<Windmill.PubSubMessageBundle> getBundlePubsubMessages() {
+    return bundlePubsubMessages;
+  }
+
+  public boolean multiKeyBundleEnabled() {
+    return multiKeyBundleOptions.multiKeyBundleEnabled();
+  }
+
   // Returns finalization callbacks recorded during the bundle execution
   public Map<Long, Pair<Instant, Runnable>> getFinalizationCallbacks() {
     return finalizationCallbacks;

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java

@@ -350,26 +350,47 @@ public long add(WindowedValue<T> 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);
 
         for (Windmill.KeyedMessageBundle.Builder keyedOutput : productionMap.values()) {
           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();
     }
   }
 

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<ElemT> 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}
      *

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 {