[Bug]: Go SDK not faithful to Runner's coder spec

[shunping]

What happened?

Two of the GroupIntoBatches tests are excluded in the prism Java VR test suite.

• 'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testInStreamingMode'
• 'org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testBufferingTimerInFixedWindow'

The stacktrace is shown below:

bundle inst009 stage-001 failed:org.apache.beam.sdk.coders.CoderException: java.io.EOFException: reached end of stream after reading 7 bytes; 69 bytes expected
        at org.apache.beam.sdk.coders.StringUtf8Coder.decode(StringUtf8Coder.java:104)
        at org.apache.beam.sdk.coders.StringUtf8Coder.decode(StringUtf8Coder.java:37)
        at org.apache.beam.sdk.coders.KvCoder.decode(KvCoder.java:84)
        at org.apache.beam.sdk.coders.KvCoder.decode(KvCoder.java:37)
        at org.apache.beam.sdk.values.WindowedValues$FullWindowedValueCoder.decode(WindowedValues.java:847)
        at org.apache.beam.sdk.values.WindowedValues$FullWindowedValueCoder.decode(WindowedValues.java:838)
        at org.apache.beam.sdk.values.WindowedValues$FullWindowedValueCoder.decode(WindowedValues.java:784)
        at org.apache.beam.sdk.fn.data.BeamFnDataInboundObserver.multiplexElements(BeamFnDataInboundObserver.java:232)
        at org.apache.beam.sdk.fn.data.BeamFnDataInboundObserver.awaitCompletion(BeamFnDataInboundObserver.java:186)
        at org.apache.beam.fn.harness.control.ProcessBundleHandler.processBundle(ProcessBundleHandler.java:542)
        at org.apache.beam.fn.harness.control.BeamFnControlClient.delegateOnInstructionRequestType(BeamFnControlClient.java:150)
        at org.apache.beam.fn.harness.control.BeamFnControlClient$InboundObserver.lambda$onNext$0(BeamFnControlClient.java:115)
        at java.base/java.util.concurrent.Executors$RunnableAdapter.call(Executors.java:539)
        at java.base/java.util.concurrent.FutureTask.run(FutureTask.java:264)
        at org.apache.beam.sdk.util.UnboundedScheduledExecutorService$ScheduledFutureTask.run(UnboundedScheduledExecutorService.java:163)
        at java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1136)
        at java.base/java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:635)
        at java.base/java.lang.Thread.run(Thread.java:833)
    Caused by: java.io.EOFException: reached end of stream after reading 7 bytes; 69 bytes expected
        at org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams.readFully(ByteStreams.java:802)
        at org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams.readFully(ByteStreams.java:784)
        at org.apache.beam.sdk.coders.StringUtf8Coder.readString(StringUtf8Coder.java:60)
        at org.apache.beam.sdk.coders.StringUtf8Coder.decode(StringUtf8Coder.java:100)
        ... 17 more
        at org.apache.beam.runners.prism.TestPrismRunner.run(TestPrismRunner.java:72)
        at org.apache.beam.sdk.Pipeline.run(Pipeline.java:325)
        at org.apache.beam.sdk.testing.TestPipeline.run(TestPipeline.java:442)
        at org.apache.beam.sdk.testing.TestPipeline.run(TestPipeline.java:381)
        at org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testInStreamingMode(GroupIntoBatchesTest.java:438)

The encoded bytes for the first element in the TestStream is
{Encoded:[3 107 101 121 69 105 110 115 116 101 105 110]}, which is actually "0x3 key Einstein", but somehow it miss the length before the second string.

Issue Priority

Priority: 2 (default / most bugs should be filed as P2)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam YAML
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Infrastructure
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner

[shunping]

After digging in the code, it seems that when we convert the TestStream's element events from/to pipeline proto, we call CoderUtils.encodeToByteArray and CoderUtils.decodeFromByteArray.

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

Lines 124 to 126 in 09aa10c

	.setEncodedElement(
	ByteString.copyFrom(
	CoderUtils.encodeToByteArray(coder, element.getValue()))));

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

Lines 150 to 153 in 09aa10c

	decodedElements.add(
	TimestampedValue.of(
	CoderUtils.decodeFromByteArray(coder, element.getEncodedElement().toByteArray()),
	new Instant(element.getTimestamp())));

.

If called without a context, CoderUtils.encodeToByteArray/CoderUtils.decodeFromByteArray would use context outer.

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

Line 96 in 09aa10c

return decodeFromByteArray(coder, encodedValue, Coder.Context.OUTER);

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

Line 116 in 09aa10c

return decodeFromByteString(coder, encodedValue, Coder.Context.OUTER);

In such context, StringUtf8Coder will not have length prefix when encoding and decoding the data.

beam/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/StringUtf8Coder.java

Lines 76 to 82 in 09aa10c

	if (context.isWholeStream) {
	byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
	if (outStream instanceof ExposedByteArrayOutputStream) {
	((ExposedByteArrayOutputStream) outStream).writeAndOwn(bytes);
	} else {
	outStream.write(bytes);
	}

beam/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/StringUtf8Coder.java

Lines 95 to 98 in 09aa10c

	if (context.isWholeStream) {
	byte[] bytes = StreamUtils.getBytesWithoutClosing(inStream);
	return new String(bytes, StandardCharsets.UTF_8);
	} else {

This works if there is not fusion break between TestStream and its following stages. However, for portable runner like Prism, there could be a fusion break after TestStream. If so, the TestStream values encoded in the pipeline proto is generated with outer context, while the downstream step uses nested context to decode the WindowedValues in the buffer.

beam/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java

Line 838 in 09aa10c

return decode(inStream, Context.NESTED);

This leads to the stack trace we see in the description.

[shunping]

Prism has some hack when handling the encoded bytes from TestStream:

beam/sdks/go/pkg/beam/runners/prism/internal/execute.go

Lines 288 to 297 in 09aa10c

	// Hack for Java Strings in test stream, since it doesn't encode them correctly.
	forceLP := cID == "StringUtf8Coder" || cID != pyld.GetCoderId()
	if forceLP {
	// slog.Warn("recoding TestStreamValue", "cID", cID, "newUrn", coders[cID].GetSpec().GetUrn(), "payloadCoder", pyld.GetCoderId(), "oldUrn", coders[pyld.GetCoderId()].GetSpec().GetUrn())
	// The coder needed length prefixing. For simplicity, add a length prefix to each
	// encoded element, since we will be sending a length prefixed coder to consume
	// this anyway. This is simpler than trying to find all the re-written coders after the fact.
	mayLP = func(v []byte) []byte {
	var buf bytes.Buffer
	if err := coder.EncodeVarInt((int64)(len(v)), &buf); err != nil {

However, it only works for element type of String. If we have a nested type like KV<String, String> (the ones used in org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testInStreamingMode and org.apache.beam.sdk.transforms.GroupIntoBatchesTest.testBufferingTimerInFixedWindow), the hack won't work.

I think the way to resolve this is to modify TestStream in Java to always use nested context instead.

[shunping]

I think the way to resolve this is to modify TestStream in Java to always use nested context instead.

Unfortunately, this approach is causing internal test failure and we have to seek an alternative way: wrap the original coder with a length-prefix coder to enforce Java to use outer context in WindowValues. This is because that the inner coder of length-prefix coder will always encode/decode using outer context no matter what context the length-prefix coder is in.

beam/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/LengthPrefixCoder.java

Line 64 in 303c4a2

return valueCoder.decode(ByteStreams.limit(inStream, size), Context.OUTER);

[shunping]

When I tried to length-prefix coders in prism, I found out a Go SDK problem (as indicated in the new issue title). Basically, even if prism added LP to the coders, the SDK side strip the LP in some scenarios. My fix in #36424 corrects this behavior.

However, it will inevitably introduce a breaking change to streaming pipeline update as the underlying coder could change. We proceed with this change because it makes Go SDK more compliant with Beam Portability.