[Dataflow Streaming] Optimize WindmillStateUtil.encodeKey by appending directly to ByteStringOutputStream

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/ByteStringAppendable.java

@@ -0,0 +1,61 @@
+/*
+ * 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.util.common;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import javax.annotation.Nullable;
+import javax.annotation.concurrent.NotThreadSafe;
+import org.apache.beam.sdk.util.ByteStringOutputStream;
+import org.apache.beam.sdk.util.Preconditions;
+
+/*
+ * An Appendable implementation that appends CharSequence in UTF-8
+ * to the passed in ByteStringOutputStream. This can be used over
+ * new OutputStreamWriter(stream, StandardCharsets.UTF_8) when thread safety of
+ * OutputStreamWriter is not required.
+ */
+@NotThreadSafe
+public class ByteStringAppendable implements Appendable {
+
+  private final ByteStringOutputStream stream;
+
+  public ByteStringAppendable(ByteStringOutputStream stream) {
+    this.stream = stream;
+  }
+
+  @Override
+  public Appendable append(@Nullable CharSequence csq) throws IOException {
+    Preconditions.checkArgumentNotNull(csq);
+    stream.write(csq.toString().getBytes(StandardCharsets.UTF_8));
+    return this;
+  }
+
+  @Override
+  public Appendable append(@Nullable CharSequence csq, int start, int end) throws IOException {
+    Preconditions.checkArgumentNotNull(csq);
+    stream.write(csq.subSequence(start, end).toString().getBytes(StandardCharsets.UTF_8));
+    return this;
+  }
+
+  @Override
+  public Appendable append(char c) throws IOException {
+    stream.write(String.valueOf(c).getBytes(StandardCharsets.UTF_8));
+    return this;
+  }
+}

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

@@ -18,31 +18,29 @@
 package org.apache.beam.runners.dataflow.worker.windmill.state;
 
 import java.io.IOException;
-import java.io.OutputStreamWriter;
-import java.nio.charset.StandardCharsets;
 import org.apache.beam.runners.core.StateNamespace;
 import org.apache.beam.runners.core.StateTag;
+import org.apache.beam.runners.dataflow.worker.util.common.ByteStringAppendable;
 import org.apache.beam.sdk.util.ByteStringOutputStream;
 import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.ByteString;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 
 class WindmillStateUtil {
+
   /** Encodes the given namespace and address as {@code <namespace>+<address>}. */
   @VisibleForTesting
   static ByteString encodeKey(StateNamespace namespace, StateTag<?> address) {
     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();
-      OutputStreamWriter writer = new OutputStreamWriter(stream, StandardCharsets.UTF_8);
-
+      ByteStringAppendable appendable = new ByteStringAppendable(stream);
       // 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(writer);
-      writer.write('+');
-      address.appendTo(writer);
-      writer.flush();
+      namespace.appendTo(appendable);
+      appendable.append('+');
+      address.appendTo(appendable);
       return stream.toByteString();
     } catch (IOException e) {
       throw new RuntimeException(e);

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/common/ByteStringAppendableTest.java

@@ -0,0 +1,97 @@
+/*
+ * 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.util.common;
+
+import static org.junit.Assert.assertEquals;
+
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.nio.charset.StandardCharsets;
+import org.apache.beam.sdk.util.ByteStringOutputStream;
+import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.ByteString;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ByteStringAppendableTest {
+
+  @Test
+  public void equivalentToOutputStreamWriter() throws IOException {
+    String randomString = "⣏⓫⦎⊺ⱄ\u243B♼⢓␜\u2065✝⡳oⶤⱲ⨻1⅒ↀ◅⡪⋲";
+    ByteString byteString1, byteString2;
+    {
+      ByteStringOutputStream stream = new ByteStringOutputStream();
+      OutputStreamWriter writer = new OutputStreamWriter(stream, StandardCharsets.UTF_8);
+      writer.append(randomString);
+      writer.flush();
+      byteString1 = stream.toByteString();
+    }
+
+    {
+      ByteStringOutputStream stream = new ByteStringOutputStream();
+      ByteStringAppendable appendable = new ByteStringAppendable(stream);
+      appendable.append(randomString);
+      byteString2 = stream.toByteString();
+    }
+    assertEquals(byteString1, byteString2);
+  }
+
+  @Test
+  public void equivalentToOutputStreamWriterSubstr() throws IOException {
+    String randomString = "⣏⓫⦎⊺ⱄ\u243B♼⢓␜\u2065✝⡳oⶤⱲ⨻1⅒ↀ◅⡪⋲";
+    ByteString byteString1, byteString2;
+    {
+      ByteStringOutputStream stream = new ByteStringOutputStream();
+      OutputStreamWriter writer = new OutputStreamWriter(stream, StandardCharsets.UTF_8);
+      writer.append(randomString, 3, 10);
+      writer.flush();
+      byteString1 = stream.toByteString();
+    }
+
+    {
+      ByteStringOutputStream stream = new ByteStringOutputStream();
+      ByteStringAppendable appendable = new ByteStringAppendable(stream);
+      appendable.append(randomString, 3, 10);
+      byteString2 = stream.toByteString();
+    }
+    assertEquals(byteString1, byteString2);
+  }
+
+  @Test
+  public void equivalentToOutputStreamWriterChar() throws IOException {
+    for (char c=0; c<=255; ++c) {
+      ByteString byteString1, byteString2;
+      {
+        ByteStringOutputStream stream = new ByteStringOutputStream();
+        OutputStreamWriter writer = new OutputStreamWriter(stream, StandardCharsets.UTF_8);
+        writer.append(c);
+        writer.flush();
+        byteString1 = stream.toByteString();
+      }
+
+      {
+        ByteStringOutputStream stream = new ByteStringOutputStream();
+        ByteStringAppendable appendable = new ByteStringAppendable(stream);
+        appendable.append(c);
+        byteString2 = stream.toByteString();
+      }
+      assertEquals(byteString1, byteString2);
+    }
+  }
+}