[Dataflow Streaming] Reuse ByteStringOutputStream buffers to reduce GC overhead

Author: arunpandianp

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

@@ -18,7 +18,6 @@
 package org.apache.beam.runners.dataflow.worker;
 
 import static org.apache.beam.runners.dataflow.util.Structs.getString;
-import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import com.google.auto.service.AutoService;
 import java.io.IOException;
@@ -54,6 +53,7 @@
   "nullness" // TODO(https://github.com/apache/beam/issues/20497)
 })
 class WindmillSink<T> extends Sink<WindowedValue<T>> {
+
   private WindmillStreamWriter writer;
   private final Coder<T> valueCoder;
   private final Coder<Collection<? extends BoundedWindow>> windowsCoder;
@@ -71,15 +71,29 @@ class WindmillSink<T> extends Sink<WindowedValue<T>> {
     this.context = context;
   }
 
+  private static ByteString encodeMetadata(
+      ByteStringOutputStream stream,
+      Coder<Collection<? extends BoundedWindow>> windowsCoder,
+      Collection<? extends BoundedWindow> windows,
+      PaneInfo paneInfo)
+      throws IOException {
+    try {
+      PaneInfoCoder.INSTANCE.encode(paneInfo, stream);
+      windowsCoder.encode(windows, stream, Coder.Context.OUTER);
+      return stream.toByteStringAndReset();
+    } catch (Exception e) {
+      stream.toByteStringAndReset();
+      throw e;
+    }
+  }
+
   public static ByteString encodeMetadata(
       Coder<Collection<? extends BoundedWindow>> windowsCoder,
       Collection<? extends BoundedWindow> windows,
       PaneInfo paneInfo)
       throws IOException {
     ByteStringOutputStream stream = new ByteStringOutputStream();
-    PaneInfoCoder.INSTANCE.encode(paneInfo, stream);
-    windowsCoder.encode(windows, stream, Coder.Context.OUTER);
-    return stream.toByteString();
+    return encodeMetadata(stream, windowsCoder, windows, paneInfo);
   }
 
   public static PaneInfo decodeMetadataPane(ByteString metadata) throws IOException {
@@ -109,6 +123,7 @@ public Map<String, SinkFactory> factories() {
   }
 
   public static class Factory implements SinkFactory {
+
     @Override
     public WindmillSink<?> create(
         CloudObject spec,
@@ -133,6 +148,7 @@ public SinkWriter<WindowedValue<T>> writer() {
   }
 
   class WindmillStreamWriter implements SinkWriter<WindowedValue<T>> {
+
     private Map<ByteString, Windmill.KeyedMessageBundle.Builder> productionMap;
     private final String destinationName;
     private final ByteStringOutputStream stream; // Kept across encodes for buffer reuse.
@@ -144,10 +160,10 @@ private WindmillStreamWriter(String destinationName) {
     }
 
     private <EncodeT> ByteString encode(Coder<EncodeT> coder, EncodeT object) throws IOException {
-      checkState(
-          stream.size() == 0,
-          "Expected output stream to be empty but had %s",
-          stream.toByteString());
+      if (stream.size() != 0) {
+        throw new IllegalStateException(
+            "Expected output stream to be empty but had " + stream.toByteString());
+      }
       try {
         coder.encode(object, stream, Coder.Context.OUTER);
         return stream.toByteStringAndReset();
@@ -162,7 +178,8 @@ private <EncodeT> ByteString encode(Coder<EncodeT> coder, EncodeT object) throws
     public long add(WindowedValue<T> data) throws IOException {
       ByteString key, value;
       ByteString id = ByteString.EMPTY;
-      ByteString metadata = encodeMetadata(windowsCoder, data.getWindows(), data.getPaneInfo());
+      ByteString metadata =
+          encodeMetadata(stream, windowsCoder, data.getWindows(), data.getPaneInfo());
       if (valueCoder instanceof KvCoder) {
         KvCoder kvCoder = (KvCoder) valueCoder;
         KV kv = (KV) data.getValue();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateUtil.java

@@ -18,30 +18,50 @@
 package org.apache.beam.runners.dataflow.worker.windmill.state;
 
 import java.io.IOException;
+import java.lang.ref.SoftReference;
 import org.apache.beam.runners.core.StateNamespace;
 import org.apache.beam.runners.core.StateTag;
 import org.apache.beam.sdk.util.ByteStringOutputStream;
 import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 class WindmillStateUtil {
 
+  private static final ThreadLocal<@Nullable SoftReference<@Nullable ByteStringOutputStream>>
+      threadLocalOutputStream = new ThreadLocal<>();
+
   /** Encodes the given namespace and address as {@code &lt;namespace&gt;+&lt;address&gt;}. */
   @VisibleForTesting
   static ByteString encodeKey(StateNamespace namespace, StateTag<?> address) {
+    // Use ByteStringOutputStream rather than concatenation and String.format. We build these keys
+    // a lot, and this leads to better performance results. See associated benchmarks.
+    ByteStringOutputStream stream = getByteStringOutputStream();
     try {
-      // Use ByteStringOutputStream rather than concatenation and String.format. We build these keys
-      // a lot, and this leads to better performance results. See associated benchmarks.
-      ByteStringOutputStream stream = new ByteStringOutputStream();
       // stringKey starts and ends with a slash.  We separate it from the
       // StateTag ID by a '+' (which is guaranteed not to be in the stringKey) because the
       // ID comes from the user.
       namespace.appendTo(stream);
       stream.append('+');
       address.appendTo(stream);
-      return stream.toByteString();
+      return stream.toByteStringAndReset();
     } catch (IOException e) {
+      stream.toByteStringAndReset();
       throw new RuntimeException(e);
+    } catch (RuntimeException e) {
+      stream.toByteStringAndReset();
+      throw e;
+    }
+  }
+
+  private static ByteStringOutputStream getByteStringOutputStream() {
+    @Nullable
+    SoftReference<@Nullable ByteStringOutputStream> refStream = threadLocalOutputStream.get();
+    @Nullable ByteStringOutputStream stream = refStream == null ? null : refStream.get();
+    if (stream == null) {
+      stream = new ByteStringOutputStream();
+      threadLocalOutputStream.set(new SoftReference<>(stream));
     }
+    return stream;
   }
 }