[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/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,32 +130,44 @@ 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)
               .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(data.getTimestamp()))
               .build());
+
       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) {
+              // If/when we add support for ordering keys, the flush needs to happen at the key
+              // level
+              context.addBundlePubsubMessages(pubsubMessages);
+            } else {
+              context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+            }
+          }
+        }
+      } finally {
+        outputBuilders.clear();
+      }
+    }
+
     @Override
     public void close() throws IOException {
-      outputBuilders.clear();
+      flush(/* bundleLevel= */ context.multiKeyBundleEnabled());
     }
 
     @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,31 @@ 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) {
+            // If/when we add support for ordering keys, the flush needs to happen at the key level
+            context.addBundlePubsubMessages(pubsubMessages);
+          } else {
+            context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+          }
+        }
+      } finally {
+        outputBuilder = createOutputBuilder();
+      }
+    }
+
     @Override
     public void close() throws IOException {
-      Windmill.PubSubMessageBundle pubsubMessages = outputBuilder.build();
-      if (pubsubMessages.getMessagesCount() > 0) {
-        context.getOutputBuilder().addPubsubMessages(pubsubMessages);
-      }
-      outputBuilder.clear();
+      flush(/* bundleLevel= */ context.multiKeyBundleEnabled());
     }
 
     @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/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();