From e6bc376f75f7e65d897d1e233ae68254a5bda545 Mon Sep 17 00:00:00 2001 From: Steven van Rossum Date: Sun, 25 Oct 2026 11:08:40 +0100 Subject: [PATCH 1/6] Add unit tests from PR #40606 --- .../data/BeamFnDataGrpcMultiplexerTest.java | 263 ++++++++++++++++++ .../InboundUnregisteredConsumerReproTest.java | 176 ++++++++++++ .../ForwardingClientResponseObserverTest.java | 14 + .../stream/OutboundObserverFactoryTest.java | 20 ++ 4 files changed, 473 insertions(+) create mode 100644 sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/InboundUnregisteredConsumerReproTest.java diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java index 6ee06d0a36a4..1579a5197942 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java @@ -21,21 +21,35 @@ import static org.hamcrest.core.StringContains.containsString; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; +import java.util.List; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicBoolean; import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.pipeline.v1.Endpoints; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.fn.stream.OutboundObserverFactory; import org.apache.beam.sdk.fn.test.TestExecutors; import org.apache.beam.sdk.fn.test.TestExecutors.TestExecutorService; import org.apache.beam.sdk.fn.test.TestStreams; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.ByteStringOutputStream; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientCallStreamObserver; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientResponseObserver; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.hamcrest.MatcherAssert; import org.hamcrest.Matchers; @@ -363,4 +377,253 @@ public void accept(BeamFnApi.Elements input) throws Exception { Iterables.getOnlyElement(errorWasReturned).getMessage(), containsString("Multiplexer hanging up")); } + + @Test + public void testEmptyElementsHandled() { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + // Should not throw and should safely exit + multiplexer.getInboundObserver().onNext(BeamFnApi.Elements.getDefaultInstance()); + } + + @Test + public void testPoisonInstructionId() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String instructionId = "poison_instruction"; + AtomicBoolean closed = new AtomicBoolean(); + + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() { + closed.set(true); + } + }); + + multiplexer.poisonInstructionId(instructionId); + assertTrue(closed.get()); + assertFalse(multiplexer.hasConsumer(instructionId)); + + // Late elements for the poisoned instruction are dropped and leave no entry behind. + multiplexer + .getInboundObserver() + .onNext( + BeamFnApi.Elements.newBuilder() + .addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId)) + .build()); + assertFalse(multiplexer.hasConsumer(instructionId)); + + // Registering a consumer for the poisoned instruction still fails after the late data. + assertThrows( + IllegalArgumentException.class, + () -> + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() {} + })); + assertFalse(multiplexer.hasConsumer(instructionId)); + } + + @Test + public void testToString() { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String str = multiplexer.toString(); + assertThat(str, containsString("apiServiceDescriptor")); + assertThat(str, containsString("consumers")); + } + + @Test + public void testConsumerAcceptThrowsExceptionPoisonsInstruction() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + + String instructionId = "throw_instruction"; + AtomicBoolean closed = new AtomicBoolean(); + + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) throws Exception { + throw new RuntimeException("Accept failure"); + } + + @Override + public void flush() {} + + @Override + public void close() { + closed.set(true); + } + }); + + multiplexer + .getInboundObserver() + .onNext( + BeamFnApi.Elements.newBuilder() + .addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(instructionId)) + .build()); + + assertTrue(closed.get()); + } + + /** A multiplexer whose inbound stream is a mock, so a test can count {@code request(1)} calls. */ + @SuppressWarnings("unchecked") + private static ClientCallStreamObserver startWithMockStream( + BeamFnDataGrpcMultiplexer multiplexer) { + ClientCallStreamObserver stream = mock(ClientCallStreamObserver.class); + ((ClientResponseObserver) + multiplexer.getInboundObserver()) + .beforeStart(stream); + return stream; + } + + private static BeamFnApi.Elements dataFor(String instructionId, String transformId) + throws Exception { + Coder> coder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + ByteStringOutputStream out = new ByteStringOutputStream(); + coder.encode(WindowedValues.valueInGlobalWindow("v"), out); + return BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId(transformId) + .setData(out.toByteString())) + .build(); + } + + private static BeamFnDataGrpcMultiplexer newMultiplexer() { + return new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inboundObserver -> TestStreams.withOnNext((BeamFnApi.Elements val) -> {}).build()); + } + + /** + * Each inbound message releases its flow-control permit once, which asks gRPC for one more + * message. A bundle whose consumer throws must release the permit of the message it was + * processing as well as of the messages still queued. + */ + @Test + public void testPermitsReleasedWhenBundleFails() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = newMultiplexer(); + ClientCallStreamObserver stream = startWithMockStream(multiplexer); + String instructionId = "failing_bundle"; + Coder> coder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE); + BeamFnDataInboundObserver observer = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList( + DataEndpoint.create( + "t1", + coder, + value -> { + throw new IllegalStateException("consumer failed"); + })), + Collections.emptyList()); + multiplexer.registerConsumer(instructionId, observer); + + for (int i = 0; i < 3; i++) { + multiplexer.getInboundObserver().onNext(dataFor(instructionId, "t1")); + } + assertThrows(IllegalStateException.class, observer::awaitCompletion); + multiplexer.poisonInstructionId(instructionId); + + verify(stream, times(3)).request(1); + assertFalse(multiplexer.hasConsumer(instructionId)); + } + + @Test + public void testPermitsReleasedForLateDataOfPoisonedInstruction() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = newMultiplexer(); + ClientCallStreamObserver stream = startWithMockStream(multiplexer); + String instructionId = "poisoned_late_data"; + multiplexer.poisonInstructionId(instructionId); + + for (int i = 0; i < 3; i++) { + multiplexer.getInboundObserver().onNext(dataFor(instructionId, "t1")); + } + + verify(stream, times(3)).request(1); + assertFalse(multiplexer.hasConsumer(instructionId)); + } + + @Test + public void testDuplicateRegistrationDoesNotPoison() throws Exception { + BeamFnDataGrpcMultiplexer multiplexer = newMultiplexer(); + String instructionId = "duplicate"; + AtomicBoolean closed = new AtomicBoolean(); + List received = new ArrayList<>(); + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) { + received.add(input); + } + + @Override + public void flush() {} + + @Override + public void close() { + closed.set(true); + } + }); + assertThrows( + IllegalArgumentException.class, + () -> + multiplexer.registerConsumer( + instructionId, + new CloseableFnDataReceiver() { + @Override + public void accept(BeamFnApi.Elements input) {} + + @Override + public void flush() {} + + @Override + public void close() {} + })); + + // The first registration keeps working. + assertTrue(multiplexer.hasConsumer(instructionId)); + assertFalse(closed.get()); + BeamFnApi.Elements elements = dataFor(instructionId, "t1"); + multiplexer.getInboundObserver().onNext(elements); + assertThat(received, Matchers.contains(elements)); + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/InboundUnregisteredConsumerReproTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/InboundUnregisteredConsumerReproTest.java new file mode 100644 index 000000000000..8bfb3e68d4f2 --- /dev/null +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/InboundUnregisteredConsumerReproTest.java @@ -0,0 +1,176 @@ +/* + * 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.sdk.fn.data; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; +import org.apache.beam.model.pipeline.v1.Endpoints; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.fn.stream.OutboundObserverFactory; +import org.apache.beam.sdk.util.ByteStringOutputStream; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientCallStreamObserver; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientResponseObserver; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Test; + +/** + * Data for an instruction whose consumer is not registered yet must not block the gRPC inbound + * thread. gRPC runs inbound onNext and the outbound onReady handler on one serializing executor, so + * a blocked onNext also stops every bundle that waits on onReady. + */ +public class InboundUnregisteredConsumerReproTest { + private static final Endpoints.ApiServiceDescriptor DESCRIPTOR = + Endpoints.ApiServiceDescriptor.newBuilder().setUrl("test").build(); + + @Test + public void testUnregisteredConsumerDoesNotBlockOnReady() throws Exception { + AtomicReference<@Nullable Runnable> onReadyHandler = new AtomicReference<>(); + AtomicBoolean ready = new AtomicBoolean(false); + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inbound -> { + ClientCallStreamObserver call = + new ClientCallStreamObserver() { + @Override + public boolean isReady() { + return ready.get(); + } + + @Override + public void setOnReadyHandler(Runnable handler) { + onReadyHandler.set(handler); + } + + @Override + public void disableAutoInboundFlowControl() {} + + @Override + public void disableAutoRequestWithInitial(int request) {} + + @Override + public void request(int count) {} + + @Override + public void setMessageCompression(boolean enable) {} + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) {} + + @Override + public void onNext(BeamFnApi.Elements value) {} + + @Override + public void onError(Throwable t) {} + + @Override + public void onCompleted() {} + }; + if (inbound instanceof ClientResponseObserver) { + @SuppressWarnings("unchecked") + ClientResponseObserver r = + (ClientResponseObserver) inbound; + r.beforeStart(call); + } + return call; + }); + + // Stands in for gRPC's per-call SerializingExecutor. + ExecutorService grpcExecutor = Executors.newSingleThreadExecutor(); + try { + String instructionId = "not-yet-registered"; + ByteStringOutputStream encoded = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("elem", encoded); + BeamFnApi.Elements data = + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("t") + .setData(encoded.toByteString())) + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(instructionId) + .setTransformId("t") + .setIsLast(true)) + .build(); + + // 1. The runner sends data before the bundle processor registers its consumer. + Future onNextFuture = + grpcExecutor.submit(() -> multiplexer.getInboundObserver().onNext(data)); + + // 2. The outbound stream becomes ready; gRPC queues onReady on the same executor. + ready.set(true); + CountDownLatch onReadyRan = new CountDownLatch(1); + Future onReadyFuture = + grpcExecutor.submit( + () -> { + Runnable handler = onReadyHandler.get(); + if (handler != null) { + handler.run(); + } + onReadyRan.countDown(); + }); + + boolean onReadyRanInTime = onReadyRan.await(2, TimeUnit.SECONDS); + + // 3. The consumer registers; the data must still arrive. + List received = Collections.synchronizedList(new ArrayList<>()); + BeamFnDataInboundObserver consumer = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList( + DataEndpoint.create("t", StringUtf8Coder.of(), received::add)), + Collections.emptyList()); + multiplexer.registerConsumer(instructionId, consumer); + ExecutorService bundle = Executors.newSingleThreadExecutor(); + try { + Future done = + bundle.submit( + () -> { + consumer.awaitCompletion(); + return null; + }); + done.get(5, TimeUnit.SECONDS); + } finally { + bundle.shutdownNow(); + } + onNextFuture.get(5, TimeUnit.SECONDS); + onReadyFuture.get(5, TimeUnit.SECONDS); + assertEquals(Collections.singletonList("elem"), received); + + assertTrue( + "onReady was stuck behind an inbound onNext waiting for an unregistered consumer", + onReadyRanInTime); + } finally { + grpcExecutor.shutdownNow(); + } + } +} diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java index fc0ebf3ca4e4..0694faa2ad30 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserverTest.java @@ -89,4 +89,18 @@ public void testCallsAreForwardedAndOnReadyHandlerBoundError() { verifyNoMoreInteractions(delegateObserver, callStreamObserver); verify(onDoneHandler).run(); } + + @Test + public void testBeforeStartForwardsToClientResponseObserver() { + @SuppressWarnings("unchecked") + ClientResponseObserver delegateObserver = mock(ClientResponseObserver.class); + @SuppressWarnings("unchecked") + ClientCallStreamObserver callStreamObserver = mock(ClientCallStreamObserver.class); + Runnable onReadyHandler = () -> {}; + ClientResponseObserver observer = + ForwardingClientResponseObserver.create(delegateObserver, onReadyHandler); + observer.beforeStart(callStreamObserver); + verify(callStreamObserver).setOnReadyHandler(onReadyHandler); + verify(delegateObserver).beforeStart(callStreamObserver); + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java index ee523dd83412..06f61acb5d44 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/stream/OutboundObserverFactoryTest.java @@ -20,6 +20,9 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.instanceOf; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.verify; import java.util.concurrent.Executors; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.CallStreamObserver; @@ -50,6 +53,15 @@ public void testDefaultInstantiation() { assertThat(observer, instanceOf(DirectStreamObserver.class)); } + @Test + public void testServerDirectInstantiation() { + StreamObserver observer = + OutboundObserverFactory.serverDirect() + .outboundObserverFor(inbound -> mockResponseObserver, mockRequestObserver); + assertThat(observer, instanceOf(DirectStreamObserver.class)); + verify(mockResponseObserver).setOnReadyHandler(any()); + } + @Test public void testBufferedStreamInstantiation() { StreamObserver observer = @@ -67,6 +79,14 @@ public void testBufferedStreamWithLimitInstantiation() { assertEquals(1, ((BufferingStreamObserver) observer).getBufferSize()); } + @Test + public void testTrivial() { + StreamObserver observer = + OutboundObserverFactory.trivial() + .outboundObserverFor(inbound -> mockResponseObserver, mockRequestObserver); + assertSame(mockResponseObserver, observer); + } + private OutboundObserverFactory.BasicFactory fakeFactory() { return inboundObserver -> { assertThat(inboundObserver, instanceOf(ForwardingClientResponseObserver.class)); From 0bc3fb610f287c0dcc9c9afc6c4b8141a09d7bce Mon Sep 17 00:00:00 2001 From: Steven van Rossum Date: Sun, 25 Oct 2026 11:47:57 +0100 Subject: [PATCH 2/6] Add reproducer tests for Fn data plane bidirectional flow-control deadlock --- .../data/BeamFnDataGrpcMultiplexerTest.java | 342 ++++++++++++++++++ 1 file changed, 342 insertions(+) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java index 1579a5197942..5b35784d292f 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexerTest.java @@ -19,6 +19,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.core.StringContains.containsString; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; @@ -32,9 +33,13 @@ import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.pipeline.v1.Endpoints; import org.apache.beam.sdk.coders.Coder; @@ -50,7 +55,9 @@ import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientCallStreamObserver; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientResponseObserver; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.checkerframework.checker.nullness.qual.Nullable; import org.hamcrest.MatcherAssert; import org.hamcrest.Matchers; import org.junit.Rule; @@ -626,4 +633,339 @@ public void close() {} multiplexer.getInboundObserver().onNext(elements); assertThat(received, Matchers.contains(elements)); } + + @Test(timeout = 10_000) + public void testBidirectionalFlowControlDoesNotDeadlock() throws Exception { + final String transformId = "transform-1"; + final int totalInboundElements = 150; + + AtomicBoolean serverAllowOutbound = new AtomicBoolean(false); + AtomicBoolean callStarted = new AtomicBoolean(false); + AtomicBoolean autoFlowControlDisabled = new AtomicBoolean(false); + AtomicInteger requestedCount = new AtomicInteger(0); + AtomicReference<@Nullable Runnable> onReadyHandler = new AtomicReference<>(); + CountDownLatch firstElementBlockedOnOutbound = new CountDownLatch(1); + List outboundElements = Collections.synchronizedList(new ArrayList<>()); + + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inbound -> { + ClientCallStreamObserver callObserver = + new ClientCallStreamObserver() { + @Override + public boolean isReady() { + if (!serverAllowOutbound.get()) { + firstElementBlockedOnOutbound.countDown(); + return false; + } + return true; + } + + @Override + public void setOnReadyHandler(Runnable handler) { + onReadyHandler.set(handler); + } + + @Override + public void disableAutoInboundFlowControl() { + autoFlowControlDisabled.set(true); + } + + @Override + public void request(int count) { + if (!callStarted.get()) { + throw new IllegalStateException("Not started"); + } + requestedCount.addAndGet(count); + } + + @Override + public void setMessageCompression(boolean enable) {} + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) {} + + @Override + public void onNext(BeamFnApi.Elements value) { + outboundElements.add(value); + } + + @Override + public void onError(Throwable t) {} + + @Override + public void onCompleted() {} + }; + if (inbound instanceof ClientResponseObserver) { + @SuppressWarnings("unchecked") + ClientResponseObserver responseObserver = + (ClientResponseObserver) inbound; + responseObserver.beforeStart(callObserver); + } + callStarted.set(true); + return callObserver; + }); + + StreamObserver outboundObserver = multiplexer.getOutboundObserver(); + StreamObserver inboundObserver = multiplexer.getInboundObserver(); + + List receivedValues = Collections.synchronizedList(new ArrayList<>()); + BeamFnDataInboundObserver dataInboundObserver = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList( + DataEndpoint.create( + transformId, + StringUtf8Coder.of(), + value -> { + receivedValues.add(value); + if (receivedValues.size() == 1) { + // DirectStreamObserver sends 100 messages before checking isReady() on the + // 101st message; emit 100 messages first so the next onNext blocks on + // isReady(). + for (int i = 0; i < 100; i++) { + outboundObserver.onNext(ELEMENTS); + } + } + outboundObserver.onNext(ELEMENTS); + })), + Collections.emptyList()); + + multiplexer.registerConsumer(DATA_INSTRUCTION_ID, dataInboundObserver); + + Future consumerFuture = + executor.submit( + () -> { + dataInboundObserver.awaitCompletion(); + return null; + }); + + ByteStringOutputStream encodedStream = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("elem", encodedStream); + ByteString encodedElement = encodedStream.toByteString(); + BeamFnApi.Elements inboundElement = + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(DATA_INSTRUCTION_ID) + .setTransformId(transformId) + .setData(encodedElement)) + .build(); + BeamFnApi.Elements terminalElement = + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(DATA_INSTRUCTION_ID) + .setTransformId(transformId) + .setIsLast(true)) + .build(); + + CountDownLatch allInboundDelivered = new CountDownLatch(1); + Future grpcThreadFuture = + executor.submit( + () -> { + inboundObserver.onNext(inboundElement); + assertTrue(firstElementBlockedOnOutbound.await(5, TimeUnit.SECONDS)); + for (int i = 1; i < totalInboundElements; i++) { + inboundObserver.onNext(inboundElement); + } + inboundObserver.onNext(terminalElement); + allInboundDelivered.countDown(); + // Simulate gRPC firing onReady on the same SerializingExecutor thread after onNext. + serverAllowOutbound.set(true); + Runnable ready = onReadyHandler.get(); + if (ready != null) { + ready.run(); + } + return null; + }); + + assertTrue( + "Inbound onNext blocked the gRPC thread", allInboundDelivered.await(5, TimeUnit.SECONDS)); + grpcThreadFuture.get(5, TimeUnit.SECONDS); + consumerFuture.get(5, TimeUnit.SECONDS); + + assertTrue(autoFlowControlDisabled.get()); + assertEquals(totalInboundElements, receivedValues.size()); + assertEquals(totalInboundElements + 100, outboundElements.size()); + assertEquals(101, requestedCount.get()); + } + + @Test(timeout = 10_000) + public void testBidirectionalFlowControlWithGrpcRequestPermitsDoesNotDeadlock() throws Exception { + final String transformId = "transform-1"; + final int totalInboundElements = 150; + + AtomicBoolean serverAllowOutbound = new AtomicBoolean(false); + AtomicBoolean callStarted = new AtomicBoolean(false); + AtomicBoolean autoFlowControlDisabled = new AtomicBoolean(false); + AtomicInteger requestedCount = new AtomicInteger(0); + java.util.concurrent.Semaphore inboundPermits = new java.util.concurrent.Semaphore(0); + AtomicReference<@Nullable Runnable> onReadyHandler = new AtomicReference<>(); + CountDownLatch firstElementBlockedOnOutbound = new CountDownLatch(1); + List outboundElements = Collections.synchronizedList(new ArrayList<>()); + + BeamFnDataGrpcMultiplexer multiplexer = + new BeamFnDataGrpcMultiplexer( + DESCRIPTOR, + OutboundObserverFactory.clientDirect(), + inbound -> { + ClientCallStreamObserver callObserver = + new ClientCallStreamObserver() { + @Override + public boolean isReady() { + if (!serverAllowOutbound.get()) { + firstElementBlockedOnOutbound.countDown(); + return false; + } + return true; + } + + @Override + public void setOnReadyHandler(Runnable handler) { + onReadyHandler.set(handler); + } + + @Override + public void disableAutoInboundFlowControl() { + autoFlowControlDisabled.set(true); + } + + @Override + public void request(int count) { + if (!callStarted.get()) { + throw new IllegalStateException("Not started"); + } + requestedCount.addAndGet(count); + inboundPermits.release(count); + } + + @Override + public void setMessageCompression(boolean enable) {} + + @Override + public void cancel(@Nullable String message, @Nullable Throwable cause) {} + + @Override + public void onNext(BeamFnApi.Elements value) { + outboundElements.add(value); + } + + @Override + public void onError(Throwable t) {} + + @Override + public void onCompleted() {} + }; + if (inbound instanceof ClientResponseObserver) { + @SuppressWarnings("unchecked") + ClientResponseObserver responseObserver = + (ClientResponseObserver) inbound; + responseObserver.beforeStart(callObserver); + } + callStarted.set(true); + return callObserver; + }); + + StreamObserver outboundObserver = multiplexer.getOutboundObserver(); + StreamObserver inboundObserver = multiplexer.getInboundObserver(); + + List receivedValues = Collections.synchronizedList(new ArrayList<>()); + BeamFnDataInboundObserver dataInboundObserver = + BeamFnDataInboundObserver.forConsumers( + Collections.singletonList( + DataEndpoint.create( + transformId, + StringUtf8Coder.of(), + value -> { + receivedValues.add(value); + if (receivedValues.size() == 1) { + for (int i = 0; i < 100; i++) { + outboundObserver.onNext(ELEMENTS); + } + } + outboundObserver.onNext(ELEMENTS); + })), + Collections.emptyList()); + + multiplexer.registerConsumer(DATA_INSTRUCTION_ID, dataInboundObserver); + + Future consumerFuture = + executor.submit( + () -> { + dataInboundObserver.awaitCompletion(); + return null; + }); + + ByteStringOutputStream encodedStream = new ByteStringOutputStream(); + StringUtf8Coder.of().encode("elem", encodedStream); + ByteString encodedElement = encodedStream.toByteString(); + BeamFnApi.Elements inboundElement = + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(DATA_INSTRUCTION_ID) + .setTransformId(transformId) + .setData(encodedElement)) + .build(); + BeamFnApi.Elements terminalElement = + BeamFnApi.Elements.newBuilder() + .addData( + BeamFnApi.Elements.Data.newBuilder() + .setInstructionId(DATA_INSTRUCTION_ID) + .setTransformId(transformId) + .setIsLast(true)) + .build(); + + CountDownLatch allInboundDelivered = new CountDownLatch(1); + Future grpcThreadFuture = + executor.submit( + () -> { + // Simulate gRPC's initial request(1) when the stream starts. + inboundPermits.release(1); + + inboundPermits.acquire(); + inboundObserver.onNext(inboundElement); + if (!autoFlowControlDisabled.get()) { + inboundPermits.release(1); + } + assertTrue(firstElementBlockedOnOutbound.await(5, TimeUnit.SECONDS)); + + boolean outboundUnblocked = false; + for (int i = 1; i < totalInboundElements; i++) { + if (!inboundPermits.tryAcquire()) { + if (!outboundUnblocked) { + outboundUnblocked = true; + serverAllowOutbound.set(true); + Runnable ready = onReadyHandler.get(); + if (ready != null) { + ready.run(); + } + } + assertTrue(inboundPermits.tryAcquire(5, TimeUnit.SECONDS)); + } + inboundObserver.onNext(inboundElement); + if (!autoFlowControlDisabled.get()) { + inboundPermits.release(1); + } + } + + assertTrue(outboundUnblocked); + assertTrue(inboundPermits.tryAcquire(5, TimeUnit.SECONDS)); + inboundObserver.onNext(terminalElement); + allInboundDelivered.countDown(); + return null; + }); + + assertTrue( + "Inbound onNext blocked the gRPC thread", allInboundDelivered.await(5, TimeUnit.SECONDS)); + grpcThreadFuture.get(5, TimeUnit.SECONDS); + consumerFuture.get(5, TimeUnit.SECONDS); + + assertTrue(autoFlowControlDisabled.get()); + assertEquals(totalInboundElements, receivedValues.size()); + assertEquals(totalInboundElements + 100, outboundElements.size()); + assertEquals(totalInboundElements + 1, requestedCount.get()); + } } From 0ed04f0003fdd283914f552580a8a8f64ce7cc63 Mon Sep 17 00:00:00 2001 From: Steven van Rossum Date: Sun, 25 Oct 2026 07:36:35 +0100 Subject: [PATCH 3/6] Prevent Fn data plane deadlock via non-blocking CancellableQueue.add and VarHandle flow control --- .../apache/beam/sdk/fn/CancellableQueue.java | 70 +++++-- .../fn/data/BeamFnDataGrpcMultiplexer.java | 194 ++++++++++++++---- .../fn/data/BeamFnDataInboundObserver.java | 37 +++- .../ForwardingClientResponseObserver.java | 4 + .../beam/sdk/fn/CancellableQueueTest.java | 18 ++ 5 files changed, 262 insertions(+), 61 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/CancellableQueue.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/CancellableQueue.java index e67423a8a044..5e0104db58e8 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/CancellableQueue.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/CancellableQueue.java @@ -36,7 +36,7 @@ public class CancellableQueue { private final int capacity; - private final @Nullable Object[] elements; + private @Nullable Object[] elements; private final Lock lock; private final Condition notFull; private final Condition notEmpty; @@ -54,6 +54,35 @@ public CancellableQueue(int capacity) { this.notEmpty = lock.newCondition(); } + /** + * Adds an element to this queue without blocking, growing the backing array if necessary. + * + * @throws Exception if the queue is cancelled. + */ + public void add(T t) throws Exception { + lock.lock(); + try { + if (cancellationException != null) { + throw cancellationException; + } + if (count == elements.length) { + @Nullable Object[] newElements = new Object[elements.length * 2]; + int firstPart = elements.length - takeIndex; + System.arraycopy(elements, takeIndex, newElements, 0, firstPart); + System.arraycopy(elements, 0, newElements, firstPart, takeIndex); + elements = newElements; + takeIndex = 0; + addIndex = count; + } + elements[addIndex] = t; + addIndex = (addIndex + 1) % elements.length; + count += 1; + notEmpty.signal(); + } finally { + lock.unlock(); + } + } + /** * Adds an element to this queue. Will block until the queue is not full or is cancelled. * @@ -80,7 +109,7 @@ public void put(T t) throws Exception, InterruptedException { } /** - * Takes an element from this queue. Will block until the queue is not full or is cancelled. + * Takes an element from this queue. Will block until the queue is not empty or is cancelled. * * @throws InterruptedException if this thread was interrupted waiting for an element. The caller * must invoke {@link #cancel} if the interrupt is unrecoverable. @@ -102,7 +131,9 @@ public T take() throws Exception, InterruptedException { elements[takeIndex] = null; takeIndex = (takeIndex + 1) % elements.length; count -= 1; - notFull.signal(); + if (count < capacity) { + notFull.signal(); + } } finally { lock.unlock(); } @@ -114,37 +145,52 @@ public T take() throws Exception, InterruptedException { * *

The first call to {@link #cancel} sets the exception that will be thrown. Resetting the * queue clears the exception. + * + * @return the number of queued elements cleared by this cancellation. */ - public void cancel(Exception exception) { + public int cancel(Exception exception) { checkNotNull(exception); lock.lock(); try { + int cleared = 0; if (cancellationException == null) { cancellationException = exception; - clearElementsLocked(); + cleared = clearElementsLocked(); } notEmpty.signalAll(); notFull.signalAll(); + return cleared; } finally { lock.unlock(); } } - private void clearElementsLocked() { - for (int i = takeIndex; count > 0; i = (i + 1) % elements.length) { - elements[i] = null; - --count; + private int clearElementsLocked() { + int cleared = count; + if (elements.length > capacity) { + elements = new Object[capacity]; + count = 0; + } else { + for (int i = takeIndex; count > 0; i = (i + 1) % elements.length) { + elements[i] = null; + --count; + } } addIndex = 0; takeIndex = 0; + return cleared; } - /** Enables the queue to be re-used after it has been cancelled. */ - public void reset() { + /** + * Enables the queue to be re-used after it has been cancelled. + * + * @return the number of queued elements cleared by this reset. + */ + public int reset() { lock.lock(); try { cancellationException = null; - clearElementsLocked(); + return clearElementsLocked(); } finally { lock.unlock(); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java index e8afd1b7e6cd..9604df1f86a0 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java @@ -17,6 +17,8 @@ */ package org.apache.beam.sdk.fn.data; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; import java.time.Duration; import java.util.HashSet; import java.util.concurrent.CompletableFuture; @@ -25,10 +27,13 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.function.Consumer; +import java.util.function.IntConsumer; import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.pipeline.v1.Endpoints; import org.apache.beam.sdk.fn.stream.OutboundObserverFactory; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Status; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientCallStreamObserver; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ClientResponseObserver; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.StreamObserver; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; @@ -54,8 +59,20 @@ public class BeamFnDataGrpcMultiplexer implements AutoCloseable { private static final Logger LOG = LoggerFactory.getLogger(BeamFnDataGrpcMultiplexer.class); private static final Duration POISONED_INSTRUCTION_ID_CACHE_TIMEOUT = Duration.ofMinutes(20); + private static final VarHandle FLOW_CONTROL_STATE; + + static { + try { + FLOW_CONTROL_STATE = + MethodHandles.lookup() + .findVarHandle(BeamFnDataGrpcMultiplexer.class, "flowControlState", int.class); + } catch (ReflectiveOperationException e) { + throw new ExceptionInInitializerError(e); + } + } + private final Endpoints.@Nullable ApiServiceDescriptor apiServiceDescriptor; - private final StreamObserver inboundObserver; + private final InboundObserver inboundObserver; private final StreamObserver outboundObserver; private final ConcurrentHashMap< /* instructionId= */ String, @@ -63,12 +80,21 @@ public class BeamFnDataGrpcMultiplexer implements AutoCloseable { receivers; private final Cache poisonedInstructionIds; + /** + * Encodes {@code (paused ? Integer.MIN_VALUE : 0) | bufferedElements}. + * + *

The sign bit is {@code 1} ({@code flowControlState < 0}) iff the single gRPC {@code + * request(1)} flow-control token is parked because {@code bufferedElements >= + * BeamFnDataInboundObserver.DEFAULT_QUEUE_CAPACITY}. + */ + @SuppressWarnings("UnusedVariable") + private volatile int flowControlState; + private static class PoisonedException extends RuntimeException { private PoisonedException() { super("Instruction poisoned"); } } - ; public BeamFnDataGrpcMultiplexer( Endpoints.@Nullable ApiServiceDescriptor apiServiceDescriptor, @@ -114,6 +140,9 @@ public StreamObserver getOutboundObserver() { */ public void registerConsumer( String instructionId, CloseableFnDataReceiver receiver) { + if (receiver instanceof BeamFnDataInboundObserver && inboundObserver.requestObserver != null) { + ((BeamFnDataInboundObserver) receiver).setOnElementsConsumed(inboundObserver); + } receivers.compute( instructionId, (unused, existing) -> { @@ -207,59 +236,124 @@ public void close() throws Exception { * harness to initiate transmitting data without needing for the receiving harness to signal that * it is ready to consume that data. */ - private final class InboundObserver implements StreamObserver { - @SuppressWarnings("LabelledBreakTarget") + private final class InboundObserver + implements ClientResponseObserver, IntConsumer { + private volatile @Nullable ClientCallStreamObserver requestObserver; + @Override - public void onNext(BeamFnApi.Elements value) { - // Have a fast path to handle the common case and provide a short circuit to exit if we detect - // multiple instruction ids. - SINGLE_INSTRUCTION_ID: - { - String instructionId = null; - for (BeamFnApi.Elements.Data data : value.getDataList()) { - if (instructionId == null) { - instructionId = data.getInstructionId(); - } else if (!instructionId.equals(data.getInstructionId())) { - // Multiple instruction ids detected, break out of this block - break SINGLE_INSTRUCTION_ID; - } + public void beforeStart(ClientCallStreamObserver requestObserver) { + this.requestObserver = requestObserver; + requestObserver.disableAutoInboundFlowControl(); + } + + @Override + public void accept(int count) { + if (count <= 0) { + throw new IllegalArgumentException("Count must be positive"); + } + int s = (int) FLOW_CONTROL_STATE.getVolatile(BeamFnDataGrpcMultiplexer.this); + while (true) { + int next = s - count; + if ((s ^ next) < 0) { + throw new IllegalStateException("Count exceeds buffered elements"); } - for (BeamFnApi.Elements.Timers timers : value.getTimersList()) { - if (instructionId == null) { - instructionId = timers.getInstructionId(); - } else if (!instructionId.equals(timers.getInstructionId())) { - // Multiple instruction ids detected, break out of this block - break SINGLE_INSTRUCTION_ID; + boolean wakeUp = + next < (Integer.MIN_VALUE | BeamFnDataInboundObserver.DEFAULT_QUEUE_CAPACITY); + int nextState = wakeUp ? (next & Integer.MAX_VALUE) : next; + int prev = + (int) + FLOW_CONTROL_STATE.compareAndExchange(BeamFnDataGrpcMultiplexer.this, s, nextState); + if (prev == s) { + if (wakeUp) { + @Nullable ClientCallStreamObserver observer = this.requestObserver; + if (observer != null) { + observer.request(1); + } } - } - if (instructionId == null) { return; } - forwardToConsumerForInstructionId(instructionId, value); - return; + s = prev; } + } - // Handle the case if there are multiple instruction ids. - HashSet instructionIds = new HashSet<>(); - for (BeamFnApi.Elements.Data data : value.getDataList()) { - instructionIds.add(data.getInstructionId()); + private void releaseOrParkFlowControlToken() { + @Nullable ClientCallStreamObserver observer = this.requestObserver; + if (observer == null) { + return; } - for (BeamFnApi.Elements.Timers timers : value.getTimersList()) { - instructionIds.add(timers.getInstructionId()); + int s = (int) FLOW_CONTROL_STATE.getVolatile(BeamFnDataGrpcMultiplexer.this); + while (true) { + if (Integer.compareUnsigned(s, BeamFnDataInboundObserver.DEFAULT_QUEUE_CAPACITY) < 0) { + observer.request(1); + return; + } + int prev = + (int) + FLOW_CONTROL_STATE.compareAndExchange( + BeamFnDataGrpcMultiplexer.this, s, s | Integer.MIN_VALUE); + if (prev == s) { + return; + } + s = prev; } - for (String instructionId : instructionIds) { - BeamFnApi.Elements.Builder builder = BeamFnApi.Elements.newBuilder(); - for (BeamFnApi.Elements.Data data : value.getDataList()) { - if (instructionId.equals(data.getInstructionId())) { - builder.addData(data); + } + + @SuppressWarnings("LabelledBreakTarget") + @Override + public void onNext(BeamFnApi.Elements value) { + try { + // Have a fast path to handle the common case and provide a short circuit to exit if we + // detect multiple instruction ids. + SINGLE_INSTRUCTION_ID: + { + String instructionId = null; + for (BeamFnApi.Elements.Data data : value.getDataList()) { + if (instructionId == null) { + instructionId = data.getInstructionId(); + } else if (!instructionId.equals(data.getInstructionId())) { + // Multiple instruction ids detected, break out of this block + break SINGLE_INSTRUCTION_ID; + } + } + for (BeamFnApi.Elements.Timers timers : value.getTimersList()) { + if (instructionId == null) { + instructionId = timers.getInstructionId(); + } else if (!instructionId.equals(timers.getInstructionId())) { + // Multiple instruction ids detected, break out of this block + break SINGLE_INSTRUCTION_ID; + } } + if (instructionId == null) { + return; + } + forwardToConsumerForInstructionId(instructionId, value); + return; + } + + // Handle the case if there are multiple instruction ids. + HashSet instructionIds = new HashSet<>(); + for (BeamFnApi.Elements.Data data : value.getDataList()) { + instructionIds.add(data.getInstructionId()); } for (BeamFnApi.Elements.Timers timers : value.getTimersList()) { - if (instructionId.equals(timers.getInstructionId())) { - builder.addTimers(timers); + instructionIds.add(timers.getInstructionId()); + } + for (String instructionId : instructionIds) { + BeamFnApi.Elements.Builder builder = BeamFnApi.Elements.newBuilder(); + for (BeamFnApi.Elements.Data data : value.getDataList()) { + if (instructionId.equals(data.getInstructionId())) { + builder.addData(data); + } + } + for (BeamFnApi.Elements.Timers timers : value.getTimersList()) { + if (instructionId.equals(timers.getInstructionId())) { + builder.addTimers(timers); + } } + forwardToConsumerForInstructionId(instructionId, builder.build()); } - forwardToConsumerForInstructionId(instructionId, builder.build()); + } finally { + releaseOrParkFlowControlToken(); } } @@ -310,10 +404,20 @@ private void forwardToConsumerForInstructionId(String instructionId, BeamFnApi.E return; } - try { - consumer.accept(value); - } catch (Exception e) { - poisonInstructionId(instructionId); + if (consumer instanceof BeamFnDataInboundObserver && this.requestObserver != null) { + FLOW_CONTROL_STATE.getAndAdd(BeamFnDataGrpcMultiplexer.this, 1); + try { + ((BeamFnDataInboundObserver) consumer).add(value); + } catch (Exception e) { + accept(1); + poisonInstructionId(instructionId); + } + } else { + try { + consumer.accept(value); + } catch (Exception e) { + poisonInstructionId(instructionId); + } } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java index 1b23c35c7faf..a92dcda1f05f 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java @@ -27,6 +27,7 @@ import java.util.Map.Entry; import java.util.NoSuchElementException; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.IntConsumer; import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements; import org.apache.beam.sdk.coders.Coder; @@ -44,6 +45,8 @@ */ public class BeamFnDataInboundObserver implements CloseableFnDataReceiver { + static final int DEFAULT_QUEUE_CAPACITY = 100; + /** * Creates a receiver that is able to consume elements multiplexing on to the provided set of * endpoints. @@ -86,6 +89,7 @@ private CloseException() { private int numEndpointsThatAreIncomplete; private final AtomicBoolean consumingReceivedData; + private @Nullable IntConsumer onElementsConsumed; private BeamFnDataInboundObserver( List> dataEndpoints, List> timerEndpoints) { @@ -99,12 +103,20 @@ private BeamFnDataInboundObserver( .computeIfAbsent(endpoint.getTransformId(), unused -> new HashMap<>()) .put(endpoint.getTimerFamilyId(), new EndpointStatus<>(endpoint)); } - this.queue = new CancellableQueue<>(100); + this.queue = new CancellableQueue<>(DEFAULT_QUEUE_CAPACITY); this.totalNumEndpoints = dataEndpoints.size() + timerEndpoints.size(); this.numEndpointsThatAreIncomplete = totalNumEndpoints; this.consumingReceivedData = new AtomicBoolean(false); } + void setOnElementsConsumed(@Nullable IntConsumer onElementsConsumed) { + this.onElementsConsumed = onElementsConsumed; + } + + void add(BeamFnApi.Elements elements) throws Exception { + queue.add(elements); + } + @Override public void accept(BeamFnApi.Elements elements) throws Exception { queue.put(elements); @@ -117,7 +129,11 @@ public void flush() throws Exception { @Override public void close() throws Exception { - queue.cancel(CloseException.INSTANCE); + int cleared = queue.cancel(CloseException.INSTANCE); + IntConsumer callback = onElementsConsumed; + if (cleared > 0 && callback != null) { + callback.accept(cleared); + } } public boolean isConsumingReceivedData() { @@ -179,6 +195,10 @@ public void awaitCompletion() throws Exception { Iterator timersIterator; { BeamFnApi.Elements elements = queue.take(); + IntConsumer callback = onElementsConsumed; + if (callback != null) { + callback.accept(1); + } dataIterator = createDiscardingIterator(elements.getDataList()); timersIterator = createDiscardingIterator(elements.getTimersList()); } @@ -189,7 +209,11 @@ public void awaitCompletion() throws Exception { } } } catch (Exception e) { - queue.cancel(e); + int cleared = queue.cancel(e); + IntConsumer callback = onElementsConsumed; + if (cleared > 0 && callback != null) { + callback.accept(cleared); + } throw e; } finally { consumingReceivedData.set(false); @@ -287,7 +311,12 @@ public void reset() { status.isDone = false; } } - queue.reset(); + int cleared = queue.reset(); + IntConsumer callback = onElementsConsumed; + if (cleared > 0 && callback != null) { + callback.accept(cleared); + } + onElementsConsumed = null; } /** diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java index d690108f8a2e..2464725b1b31 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/stream/ForwardingClientResponseObserver.java @@ -70,8 +70,12 @@ public void onCompleted() { inboundObserver.onCompleted(); } + @SuppressWarnings("unchecked") @Override public void beforeStart(ClientCallStreamObserver stream) { + if (inboundObserver instanceof ClientResponseObserver) { + ((ClientResponseObserver) inboundObserver).beforeStart(stream); + } stream.setOnReadyHandler(onReadyHandler); } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/CancellableQueueTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/CancellableQueueTest.java index aa85e7a2cec6..ba2357a102fa 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/CancellableQueueTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/fn/CancellableQueueTest.java @@ -247,4 +247,22 @@ public void testMemoryReferenceOnTake() throws Exception { assertTrue(weakReference1.get() == null); assertTrue(weakReference2.get() == null); } + + @Test + public void testAddGrowsQueueWithoutBlocking() throws Exception { + CancellableQueue queue = new CancellableQueue<>(4); + queue.add(0); + queue.add(1); + assertEquals(Integer.valueOf(0), queue.take()); + // Wrap around ring buffer before growing + for (int i = 2; i < 15; i++) { + queue.add(i); + } + for (int i = 1; i < 10; i++) { + assertEquals(Integer.valueOf(i), queue.take()); + } + assertEquals(5, queue.cancel(new RuntimeException("cancelled"))); + assertEquals(0, queue.cancel(new RuntimeException("cancelled again"))); + assertEquals(0, queue.reset()); + } } From f69285f818da2c797432b37116f24a44026e90d6 Mon Sep 17 00:00:00 2001 From: Steven van Rossum Date: Sun, 25 Oct 2026 20:01:44 +0100 Subject: [PATCH 4/6] Guard BeamFnDataInboundObserver.onElementsConsumed reads on cleared > 0 --- .../fn/data/BeamFnDataInboundObserver.java | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java index a92dcda1f05f..800cc683400d 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java @@ -130,9 +130,11 @@ public void flush() throws Exception { @Override public void close() throws Exception { int cleared = queue.cancel(CloseException.INSTANCE); - IntConsumer callback = onElementsConsumed; - if (cleared > 0 && callback != null) { - callback.accept(cleared); + if (cleared > 0) { + IntConsumer callback = onElementsConsumed; + if (callback != null) { + callback.accept(cleared); + } } } @@ -210,9 +212,11 @@ public void awaitCompletion() throws Exception { } } catch (Exception e) { int cleared = queue.cancel(e); - IntConsumer callback = onElementsConsumed; - if (cleared > 0 && callback != null) { - callback.accept(cleared); + if (cleared > 0) { + IntConsumer callback = onElementsConsumed; + if (callback != null) { + callback.accept(cleared); + } } throw e; } finally { @@ -312,9 +316,11 @@ public void reset() { } } int cleared = queue.reset(); - IntConsumer callback = onElementsConsumed; - if (cleared > 0 && callback != null) { - callback.accept(cleared); + if (cleared > 0) { + IntConsumer callback = onElementsConsumed; + if (callback != null) { + callback.accept(cleared); + } } onElementsConsumed = null; } From b768043f2ac294ae39a9a7f7f235a21a4dcbb4ce Mon Sep 17 00:00:00 2001 From: Steven van Rossum Date: Sun, 25 Oct 2026 20:28:10 +0100 Subject: [PATCH 5/6] Document flow-control and happens-before invariants --- .../apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java | 6 ++++++ .../apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java | 4 ++++ 2 files changed, 10 insertions(+) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java index 9604df1f86a0..3b91fcd1dc84 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java @@ -140,6 +140,12 @@ public StreamObserver getOutboundObserver() { */ public void registerConsumer( String instructionId, CloseableFnDataReceiver receiver) { + // beforeStart() is invoked synchronously during BeamFnDataGrpcMultiplexer construction on + // client-side gRPC streams (setting requestObserver) and is never invoked on server-side + // streams or plain StreamObserver test doubles (where forwardToConsumerForInstructionId falls + // back to blocking consumer.accept()). We must also set onElementsConsumed before completing + // the receiver future in receivers.compute() below to establish a happens-before edge to the + // inbound gRPC thread calling consumerFuture.get() -> consumer.add(). if (receiver instanceof BeamFnDataInboundObserver && inboundObserver.requestObserver != null) { ((BeamFnDataInboundObserver) receiver).setOnElementsConsumed(inboundObserver); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java index 800cc683400d..68d87ea51771 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataInboundObserver.java @@ -89,6 +89,10 @@ private CloseException() { private int numEndpointsThatAreIncomplete; private final AtomicBoolean consumingReceivedData; + // Plain field: visibility is established by the happens-before chain from + // setOnElementsConsumed() -> CompletableFuture completion in registerConsumer() -> + // consumerFuture.get() -> queue.add() (unlocking queue.lock) -> queue.take()/cancel()/reset() + // (locking queue.lock). Must only be read after queue.take() returns or when cleared > 0. private @Nullable IntConsumer onElementsConsumed; private BeamFnDataInboundObserver( From fcf1a52b4f79102d573ee550c9a03cecef3aed84 Mon Sep 17 00:00:00 2001 From: Steven van Rossum Date: Mon, 26 Oct 2026 00:58:00 +0100 Subject: [PATCH 6/6] Buffer pre-registration inbound elements via lock-free VarHandle state machine Replace the per-instruction CompletableFuture in BeamFnDataGrpcMultiplexer with a lock-free ReceiverEntry state machine backed by VarHandle and an intrusive BufferedNode Treiber stack: - When inbound Elements arrive before registerConsumer, onNext pushes a BufferedNode onto ReceiverEntry.state via compareAndExchange (tracking flow-control permits in FLOW_CONTROL_STATE) and returns immediately without blocking the gRPC SerializingExecutor thread on CompletableFuture.get(10, SECONDS). - When registerConsumer runs, it atomically detaches any buffered chain, reverses it in place for zero-allocation FIFO delivery, and publishes the registered receiver once drained. - When registerConsumer runs before any onNext (the common fast path), ReceiverEntry is inserted with state = receiver directly, avoiding CompletableFuture and node allocations altogether. --- .../fn/data/BeamFnDataGrpcMultiplexer.java | 413 +++++++++++++----- 1 file changed, 315 insertions(+), 98 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java index 3b91fcd1dc84..ec5b2179764a 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java @@ -21,11 +21,7 @@ import java.lang.invoke.VarHandle; import java.time.Duration; import java.util.HashSet; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; import java.util.function.Consumer; import java.util.function.IntConsumer; import org.apache.beam.model.fnexecution.v1.BeamFnApi; @@ -61,6 +57,10 @@ public class BeamFnDataGrpcMultiplexer implements AutoCloseable { private static final Duration POISONED_INSTRUCTION_ID_CACHE_TIMEOUT = Duration.ofMinutes(20); private static final VarHandle FLOW_CONTROL_STATE; + private static final Object UNREGISTERED = new Object(); + private static final Object DRAINING = new Object(); + private static final Object POISONED = new Object(); + static { try { FLOW_CONTROL_STATE = @@ -74,10 +74,7 @@ public class BeamFnDataGrpcMultiplexer implements AutoCloseable { private final Endpoints.@Nullable ApiServiceDescriptor apiServiceDescriptor; private final InboundObserver inboundObserver; private final StreamObserver outboundObserver; - private final ConcurrentHashMap< - /* instructionId= */ String, - CompletableFuture>> - receivers; + private final ConcurrentHashMap receivers; private final Cache poisonedInstructionIds; /** @@ -90,9 +87,88 @@ public class BeamFnDataGrpcMultiplexer implements AutoCloseable { @SuppressWarnings("UnusedVariable") private volatile int flowControlState; - private static class PoisonedException extends RuntimeException { - private PoisonedException() { - super("Instruction poisoned"); + /** + * A Treiber-stack node for {@link BeamFnApi.Elements} that arrive before a consumer has finished + * registering and draining. Once detached via CAS/getAndSet on {@link ReceiverEntry#STATE}, the + * chain is exclusively owned by the detaching thread and reversed in-place via {@link #reverse} + * for FIFO delivery. + */ + private static final class BufferedNode { + final BeamFnApi.Elements value; + final boolean hasPermit; + final boolean draining; + @Nullable BufferedNode nextNode; + + BufferedNode(BeamFnApi.Elements value, boolean hasPermit, Object nextState) { + this.value = value; + this.hasPermit = hasPermit; + if (nextState instanceof BufferedNode) { + BufferedNode next = (BufferedNode) nextState; + this.draining = next.draining; + this.nextNode = next; + } else { + this.draining = (nextState == DRAINING); + this.nextNode = null; + } + } + + static @Nullable BufferedNode reverse(@Nullable BufferedNode head) { + @Nullable BufferedNode prev = null; + @Nullable BufferedNode curr = head; + while (curr != null) { + @Nullable BufferedNode next = curr.nextNode; + curr.nextNode = prev; + prev = curr; + curr = next; + } + return prev; + } + + static int countPermits(@Nullable BufferedNode head) { + int permits = 0; + for (BufferedNode curr = head; curr != null; curr = curr.nextNode) { + if (curr.hasPermit) { + permits++; + } + } + return permits; + } + } + + /** + * Lock-free state holder for an {@code instructionId} in {@link #receivers}. + * + *

{@link #state} transitions monotonically through: + * + *

    + *
  1. {@link #UNREGISTERED} or a {@link BufferedNode} stack with {@code draining == false} + *
  2. {@link #DRAINING} or a {@link BufferedNode} stack with {@code draining == true} while + * {@link #registerConsumer} drains pre-registration elements in FIFO order + *
  3. The registered {@link CloseableFnDataReceiver} once drained (or immediately on + * construction if registered before any inbound data arrives) + *
  4. {@link #POISONED} when unregistered, poisoned, or closed + *
+ */ + private static final class ReceiverEntry { + private static final VarHandle STATE; + + static { + try { + STATE = MethodHandles.lookup().findVarHandle(ReceiverEntry.class, "state", Object.class); + } catch (ReflectiveOperationException e) { + throw new ExceptionInInitializerError(e); + } + } + + @SuppressWarnings("UnusedVariable") + private volatile Object state; + + ReceiverEntry() { + this.state = UNREGISTERED; + } + + ReceiverEntry(CloseableFnDataReceiver receiver) { + this.state = receiver; } } @@ -140,39 +216,94 @@ public StreamObserver getOutboundObserver() { */ public void registerConsumer( String instructionId, CloseableFnDataReceiver receiver) { + if (poisonedInstructionIds.getIfPresent(instructionId) != null) { + throw new IllegalArgumentException("Instruction id was poisoned"); + } // beforeStart() is invoked synchronously during BeamFnDataGrpcMultiplexer construction on // client-side gRPC streams (setting requestObserver) and is never invoked on server-side // streams or plain StreamObserver test doubles (where forwardToConsumerForInstructionId falls - // back to blocking consumer.accept()). We must also set onElementsConsumed before completing - // the receiver future in receivers.compute() below to establish a happens-before edge to the - // inbound gRPC thread calling consumerFuture.get() -> consumer.add(). + // back to blocking consumer.accept()). We must also set onElementsConsumed before publishing + // the receiver in receivers below to establish a happens-before edge to the inbound gRPC + // thread reading ReceiverEntry.STATE -> consumer.add(). if (receiver instanceof BeamFnDataInboundObserver && inboundObserver.requestObserver != null) { ((BeamFnDataInboundObserver) receiver).setOnElementsConsumed(inboundObserver); } - receivers.compute( - instructionId, - (unused, existing) -> { - if (existing != null) { - if (!existing.complete(receiver)) { - throw new IllegalArgumentException("Instruction id was registered twice"); - } - return existing; - } - if (poisonedInstructionIds.getIfPresent(instructionId) != null) { - throw new IllegalArgumentException("Instruction id was poisoned"); - } - return CompletableFuture.completedFuture(receiver); - }); + ReceiverEntry newEntry = new ReceiverEntry(receiver); + @Nullable ReceiverEntry existing = receivers.putIfAbsent(instructionId, newEntry); + if (existing == null) { + if (poisonedInstructionIds.getIfPresent(instructionId) != null + && receivers.remove(instructionId, newEntry)) { + ReceiverEntry.STATE.setVolatile(newEntry, POISONED); + throw new IllegalArgumentException("Instruction id was poisoned"); + } + return; + } + + Object s = ReceiverEntry.STATE.getVolatile(existing); + while (true) { + if (s instanceof CloseableFnDataReceiver + || s == DRAINING + || (s instanceof BufferedNode && ((BufferedNode) s).draining)) { + throw new IllegalArgumentException("Instruction id was registered twice"); + } + if (s == POISONED || poisonedInstructionIds.getIfPresent(instructionId) != null) { + throw new IllegalArgumentException("Instruction id was poisoned"); + } + Object target = (s == UNREGISTERED) ? receiver : DRAINING; + Object prev = ReceiverEntry.STATE.compareAndExchange(existing, s, target); + if (prev == s) { + if (s != UNREGISTERED) { + inboundObserver.drain(instructionId, existing, (BufferedNode) s, receiver); + } + return; + } + s = prev; + } } /** Unregisters a previously registered consumer. */ public void unregisterConsumer(String instructionId) { - @Nullable CompletableFuture> receiverFuture = - receivers.remove(instructionId); - if (receiverFuture != null && !receiverFuture.isDone()) { - // The future must have been inserted by the inbound observer since registerConsumer completes - // the future. - throw new IllegalArgumentException("Unregistering consumer which was not registered."); + @Nullable ReceiverEntry entry = receivers.remove(instructionId); + if (entry != null) { + Object prev = ReceiverEntry.STATE.getAndSet(entry, POISONED); + if (prev instanceof BufferedNode) { + BufferedNode node = (BufferedNode) prev; + int permits = BufferedNode.countPermits(node); + if (permits > 0) { + inboundObserver.accept(permits); + } + if (!node.draining) { + throw new IllegalArgumentException("Unregistering consumer which was not registered."); + } + } else if (prev == UNREGISTERED) { + // The entry must have been inserted by the inbound observer since registerConsumer + // transitions the state out of UNREGISTERED. + throw new IllegalArgumentException("Unregistering consumer which was not registered."); + } + } + } + + @SuppressWarnings("unchecked") + private @Nullable CloseableFnDataReceiver releaseEntry(ReceiverEntry entry) { + Object prev = ReceiverEntry.STATE.getAndSet(entry, POISONED); + if (prev instanceof BufferedNode) { + int permits = BufferedNode.countPermits((BufferedNode) prev); + if (permits > 0) { + inboundObserver.accept(permits); + } + return null; + } + if (prev instanceof CloseableFnDataReceiver) { + return (CloseableFnDataReceiver) prev; + } + return null; + } + + private static void closeReceiverQuietly(CloseableFnDataReceiver receiver) { + try { + receiver.close(); + } catch (Exception e) { + LOG.warn("Unexpected error closing existing observer", e); } } @@ -184,18 +315,11 @@ public void unregisterConsumer(String instructionId) { */ public void poisonInstructionId(String instructionId) { poisonedInstructionIds.put(instructionId, Boolean.TRUE); - @Nullable CompletableFuture> receiverFuture = - receivers.remove(instructionId); - if (receiverFuture != null) { - // Completing exceptionally has no effect if the future was already notified. In that case - // whatever registered the receiver needs to handle cancelling it. - receiverFuture.completeExceptionally(new PoisonedException()); - if (!receiverFuture.isCompletedExceptionally()) { - try { - receiverFuture.get().close(); - } catch (Exception e) { - LOG.warn("Unexpected error closing existing observer"); - } + @Nullable ReceiverEntry entry = receivers.remove(instructionId); + if (entry != null) { + @Nullable CloseableFnDataReceiver receiver = releaseEntry(entry); + if (receiver != null) { + closeReceiverQuietly(receiver); } } } @@ -208,14 +332,11 @@ boolean hasConsumer(String instructionId) { @Override public void close() throws Exception { Exception exception = null; - for (CompletableFuture> receiver : - ImmutableList.copyOf(receivers.values())) { - // Cancel any observer waiting for the client to complete. If the receiver has already been - // completed or cancelled, this call will be ignored. - receiver.cancel(true); - if (!receiver.isCompletedExceptionally()) { + for (ReceiverEntry entry : ImmutableList.copyOf(receivers.values())) { + @Nullable CloseableFnDataReceiver receiver = releaseEntry(entry); + if (receiver != null) { try { - receiver.get().close(); + receiver.close(); } catch (Exception e) { if (exception == null) { exception = e; @@ -238,9 +359,10 @@ public void close() throws Exception { * A multiplexing {@link StreamObserver} that selects the inbound {@link Consumer} to pass the * elements to. * - *

The inbound observer blocks until the {@link Consumer} is bound allowing for the sending - * harness to initiate transmitting data without needing for the receiving harness to signal that - * it is ready to consume that data. + *

If data arrives for an instruction before its {@link Consumer} is registered, the elements + * are buffered in {@link ReceiverEntry} without blocking the inbound gRPC thread, allowing the + * sending harness to initiate transmitting data without needing the receiving harness to signal + * that it is ready to consume that data. */ private final class InboundObserver implements ClientResponseObserver, IntConsumer { @@ -363,53 +485,71 @@ public void onNext(BeamFnApi.Elements value) { } } + private @Nullable ReceiverEntry getOrCreateEntry(String instructionId) { + @Nullable ReceiverEntry entry = receivers.get(instructionId); + if (entry != null) { + return entry; + } + if (poisonedInstructionIds.getIfPresent(instructionId) != null) { + return null; + } + ReceiverEntry newEntry = new ReceiverEntry(); + @Nullable ReceiverEntry existing = receivers.putIfAbsent(instructionId, newEntry); + if (existing != null) { + return existing; + } + if (poisonedInstructionIds.getIfPresent(instructionId) != null) { + if (receivers.remove(instructionId, newEntry)) { + releaseEntry(newEntry); + } + return null; + } + LOG.debug( + "Received data for instruction {} without consumer ready. " + + "Buffering until consumer is registered.", + instructionId); + return newEntry; + } + + @SuppressWarnings("unchecked") private void forwardToConsumerForInstructionId(String instructionId, BeamFnApi.Elements value) { - CloseableFnDataReceiver consumer; - try { - CompletableFuture> consumerFuture = - receivers.computeIfAbsent( - instructionId, - (unused) -> { - if (poisonedInstructionIds.getIfPresent(instructionId) != null) { - throw new PoisonedException(); - } - LOG.debug( - "Received data for instruction {} without consumer ready. " - + "Waiting for consumer to be registered.", - instructionId); - return new CompletableFuture<>(); - }); - // The consumer may not be registered until the bundle processor is fully constructed so we - // conservatively set - // a high timeout. Poisoning will prevent this for occurring for consumers that will not be - // registered. - consumer = consumerFuture.get(3, TimeUnit.HOURS); - /* - * TODO: On failure we should fail any bundles that were impacted eagerly - * instead of relying on the Runner harness to do all the failure handling. - */ - } catch (TimeoutException e) { - LOG.error( - "Timed out waiting to observe consumer data stream for instruction {}", - instructionId, - e); - outboundObserver.onError(e); + @Nullable ReceiverEntry entry = getOrCreateEntry(instructionId); + if (entry == null) { + LOG.debug("Received data for poisoned instruction {}. Dropping input.", instructionId); return; - } catch (ExecutionException | InterruptedException | PoisonedException e) { - if (e instanceof PoisonedException || e.getCause() instanceof PoisonedException) { - LOG.debug("Received data for poisoned instruction {}. Dropping input.", instructionId); - return; + } + + Object s = ReceiverEntry.STATE.getVolatile(entry); + if (!(s instanceof CloseableFnDataReceiver)) { + boolean hasPermit = this.requestObserver != null; + if (hasPermit) { + FLOW_CONTROL_STATE.getAndAdd(BeamFnDataGrpcMultiplexer.this, 1); + } + while (true) { + if (s == POISONED) { + if (hasPermit) { + accept(1); + } + LOG.debug("Received data for poisoned instruction {}. Dropping input.", instructionId); + return; + } + if (s instanceof CloseableFnDataReceiver) { + if (hasPermit) { + accept(1); + } + break; + } + BufferedNode node = new BufferedNode(value, hasPermit, s); + Object prev = ReceiverEntry.STATE.compareAndExchange(entry, s, node); + if (prev == s) { + return; + } + s = prev; } - LOG.error( - "Client interrupted during handling of data for instruction {}", instructionId, e); - outboundObserver.onError(e); - return; - } catch (RuntimeException e) { - LOG.error("Client failed to handle data for instruction {}", instructionId, e); - outboundObserver.onError(e); - return; } + CloseableFnDataReceiver consumer = + (CloseableFnDataReceiver) s; if (consumer instanceof BeamFnDataInboundObserver && this.requestObserver != null) { FLOW_CONTROL_STATE.getAndAdd(BeamFnDataGrpcMultiplexer.this, 1); try { @@ -427,6 +567,83 @@ private void forwardToConsumerForInstructionId(String instructionId, BeamFnApi.E } } + private void drain( + String instructionId, + ReceiverEntry entry, + BufferedNode initialBatch, + CloseableFnDataReceiver receiver) { + boolean isFlowControlledInboundObserver = + receiver instanceof BeamFnDataInboundObserver && this.requestObserver != null; + @Nullable BufferedNode batch = initialBatch; + while (true) { + for (BufferedNode curr = BufferedNode.reverse(batch); curr != null; curr = curr.nextNode) { + if (ReceiverEntry.STATE.getVolatile(entry) == POISONED) { + int remainingPermits = BufferedNode.countPermits(curr); + if (remainingPermits > 0) { + accept(remainingPermits); + } + closeReceiverQuietly(receiver); + return; + } + boolean hasPermit = curr.hasPermit; + if (isFlowControlledInboundObserver) { + if (!hasPermit) { + FLOW_CONTROL_STATE.getAndAdd(BeamFnDataGrpcMultiplexer.this, 1); + } + try { + ((BeamFnDataInboundObserver) receiver).add(curr.value); + } catch (Exception e) { + int remainingPermits = 1 + BufferedNode.countPermits(curr.nextNode); + accept(remainingPermits); + poisonInstructionId(instructionId); + closeReceiverQuietly(receiver); + return; + } + } else { + try { + receiver.accept(curr.value); + } catch (Exception e) { + int remainingPermits = (hasPermit ? 1 : 0) + BufferedNode.countPermits(curr.nextNode); + if (remainingPermits > 0) { + accept(remainingPermits); + } + poisonInstructionId(instructionId); + closeReceiverQuietly(receiver); + if (e instanceof RuntimeException) { + throw (RuntimeException) e; + } + throw new RuntimeException(e); + } + if (hasPermit) { + accept(1); + } + } + } + + Object s = ReceiverEntry.STATE.getVolatile(entry); + while (true) { + if (s == POISONED) { + closeReceiverQuietly(receiver); + return; + } + if (s == DRAINING) { + Object prev = ReceiverEntry.STATE.compareAndExchange(entry, DRAINING, receiver); + if (prev == DRAINING) { + return; + } + s = prev; + } else { + Object prev = ReceiverEntry.STATE.compareAndExchange(entry, s, DRAINING); + if (prev == s) { + batch = (BufferedNode) s; + break; + } + s = prev; + } + } + } + } + @Override public void onError(Throwable t) { LOG.error(