[Dataflow Streaming] Reuse ByteStringOutputStream buffers to reduce GC overhead

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.reset();
+      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,15 +160,15 @@ 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();
       } catch (Exception e) {
-        stream.toByteStringAndReset();
+        stream.reset();
         throw e;
       }
     }
@@ -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,78 @@
 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 RefHolder> threadLocalRefHolder = new ThreadLocal<>();
+
   /** Encodes the given namespace and address as {@code &lt;namespace&gt;+&lt;address&gt;}. */
   @VisibleForTesting
   static ByteString encodeKey(StateNamespace namespace, StateTag<?> address) {
+    RefHolder refHolder = getRefHolderFromThreadLocal();
+    // 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;
+    boolean releaseThreadLocal;
+    if (refHolder.inUse) {
+      // If the thread local stream is already in use, create a new one
+      stream = new ByteStringOutputStream();
+      releaseThreadLocal = false;
+    } else {
+      stream = getByteStringOutputStream(refHolder);
+      refHolder.inUse = true;
+      releaseThreadLocal = true;
+    }
     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) {
       throw new RuntimeException(e);
+    } finally {
+      stream.reset();
+      if (releaseThreadLocal) {
+        refHolder.inUse = false;
+      }
+    }
+  }
+
+  private static class RefHolder {
+    public SoftReference<@Nullable ByteStringOutputStream> streamRef =
+        new SoftReference<>(new ByteStringOutputStream());
+
+    // Boolean is true when the thread local stream is already in use by the current thread.
+    // Used to avoid reusing the same stream from nested calls if any.
+    public boolean inUse = false;
+  }
+
+  private static RefHolder getRefHolderFromThreadLocal() {
+    @Nullable RefHolder refHolder = threadLocalRefHolder.get();
+    if (refHolder == null) {
+      refHolder = new RefHolder();
+      threadLocalRefHolder.set(refHolder);
+    }
+    return refHolder;
+  }
+
+  private static ByteStringOutputStream getByteStringOutputStream(RefHolder refHolder) {
+    @Nullable
+    ByteStringOutputStream stream = refHolder.streamRef == null ? null : refHolder.streamRef.get();
+    if (stream == null) {
+      stream = new ByteStringOutputStream();
+      refHolder.streamRef = new SoftReference<>(stream);
     }
+    return stream;
   }
 }

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateUtilTest.java

@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.windmill.state;
+
+import static org.junit.Assert.assertEquals;
+
+import java.io.IOException;
+import org.apache.beam.runners.core.StateNamespace;
+import org.apache.beam.runners.core.StateNamespaceForTest;
+import org.apache.beam.runners.core.StateTag;
+import org.apache.beam.runners.core.StateTags;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.state.SetState;
+import org.apache.beam.sdk.state.StateSpec;
+import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class WindmillStateUtilTest {
+
+  @Test
+  public void testEncodeKey() {
+    StateNamespaceForTest namespace = new StateNamespaceForTest("key");
+    StateTag<SetState<Integer>> foo = StateTags.set("foo", VarIntCoder.of());
+    ByteString bytes = WindmillStateUtil.encodeKey(namespace, foo);
+    assertEquals("key+ufoo", bytes.toStringUtf8());
+  }
+
+  @Test
+  public void testEncodeKeyNested() {
+    // Hypothetical case where a namespace/tag encoding depends on a call to encodeKey
+    // This tests if thread locals in WindmillStateUtil are not reused with nesting
+    StateNamespaceForTest namespace1 = new StateNamespaceForTest("key");
+    StateTag<SetState<Integer>> tag1 = StateTags.set("foo", VarIntCoder.of());
+    StateTag<SetState<Integer>> tag2 =
+        new StateTag<SetState<Integer>>() {
+          @Override
+          public void appendTo(Appendable sb) throws IOException {
+            WindmillStateUtil.encodeKey(namespace1, tag1);
+            sb.append("tag2");
+          }
+
+          @Override
+          public String getId() {
+            return "";
+          }
+
+          @Override
+          public StateSpec<SetState<Integer>> getSpec() {
+            return null;
+          }
+
+          @Override
+          public SetState<Integer> bind(StateBinder binder) {
+            return null;
+          }
+        };
+
+    StateNamespace namespace2 =
+        new StateNamespaceForTest("key") {
+          @Override
+          public void appendTo(Appendable sb) throws IOException {
+            WindmillStateUtil.encodeKey(namespace1, tag1);
+            sb.append("namespace2");
+          }
+        };
+    ByteString bytes = WindmillStateUtil.encodeKey(namespace2, tag2);
+    assertEquals("namespace2+tag2", bytes.toStringUtf8());
+  }
+}

sdks/java/core/src/main/java/org/apache/beam/sdk/util/ByteStringOutputStream.java

@@ -158,6 +158,16 @@ public ByteString toByteStringAndReset() {
     return rval;
   }
 
+  /*
+   * Resets the output stream to be re-used possibly re-using any existing buffers.
+   */
+  public void reset() {
+    if (size() == 0) {
+      return;
+    }
+    toByteStringAndReset();
+  }
+
   /**
    * Creates a byte string with the given size containing the prefix of the contents of this output
    * stream.

sdks/java/core/src/test/java/org/apache/beam/sdk/util/ByteStringOutputStreamTest.java

@@ -223,6 +223,19 @@ public void appendEquivalentToOutputStreamWriterChar() throws IOException {
     }
   }
 
+  @Test
+  public void testReset() throws IOException {
+    try (ByteStringOutputStream stream = new ByteStringOutputStream()) {
+      stream.reset();
+      assertEquals(ByteString.EMPTY, stream.toByteString());
+      stream.append("test");
+      stream.reset();
+      assertEquals(ByteString.EMPTY, stream.toByteString());
+      stream.reset();
+      assertEquals(ByteString.EMPTY, stream.toByteString());
+    }
+  }
+
   // Grow the elements based upon an approximation of the fibonacci sequence.
   private static int next(int current) {
     double a = Math.max(1, current * (1 + Math.sqrt(5)) / 2.0);