diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java index ee94bc202ee2..ad07688cc8cc 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java +++ b/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 extends Sink> { + private WindmillStreamWriter writer; private final Coder valueCoder; private final Coder> windowsCoder; @@ -71,15 +71,29 @@ class WindmillSink extends Sink> { this.context = context; } + private static ByteString encodeMetadata( + ByteStringOutputStream stream, + Coder> windowsCoder, + Collection 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> windowsCoder, Collection 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 factories() { } public static class Factory implements SinkFactory { + @Override public WindmillSink create( CloudObject spec, @@ -133,6 +148,7 @@ public SinkWriter> writer() { } class WindmillStreamWriter implements SinkWriter> { + private Map productionMap; private final String destinationName; private final ByteStringOutputStream stream; // Kept across encodes for buffer reuse. @@ -144,15 +160,15 @@ private WindmillStreamWriter(String destinationName) { } private ByteString encode(Coder 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 ByteString encode(Coder coder, EncodeT object) throws public long add(WindowedValue 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(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateUtil.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateUtil.java index 9ce2d687b3fe..d95bf95db806 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateUtil.java +++ b/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 <namespace>+<address>}. */ @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; } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateUtilTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateUtilTest.java new file mode 100644 index 000000000000..589edeb1e544 --- /dev/null +++ b/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> 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> tag1 = StateTags.set("foo", VarIntCoder.of()); + StateTag> tag2 = + new StateTag>() { + @Override + public void appendTo(Appendable sb) throws IOException { + WindmillStateUtil.encodeKey(namespace1, tag1); + sb.append("tag2"); + } + + @Override + public String getId() { + return ""; + } + + @Override + public StateSpec> getSpec() { + return null; + } + + @Override + public SetState 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()); + } +} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ByteStringOutputStream.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ByteStringOutputStream.java index 76a6b18890ba..ade84f7a6436 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/ByteStringOutputStream.java +++ b/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. diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/ByteStringOutputStreamTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/ByteStringOutputStreamTest.java index 37ce6a385abb..605d341d476f 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/ByteStringOutputStreamTest.java +++ b/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);