Prevent Fn data plane deadlock via non-blocking CancellableQueue.add and gRPC inbound flow control

sdks/java/core/src/main/java/org/apache/beam/sdk/fn/CancellableQueue.java

@@ -36,7 +36,7 @@
 public class CancellableQueue<T extends @NonNull Object> {
 
   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 {
    *
    * <p>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();
     }

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,21 +59,42 @@
 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<BeamFnApi.Elements> inboundObserver;
+  private final InboundObserver inboundObserver;
   private final StreamObserver<BeamFnApi.Elements> outboundObserver;
   private final ConcurrentHashMap<
           /* instructionId= */ String,
           CompletableFuture<CloseableFnDataReceiver<BeamFnApi.Elements>>>
       receivers;
   private final Cache</* instructionId= */ String, /* unused= */ Boolean> poisonedInstructionIds;
 
+  /**
+   * Encodes {@code (paused ? Integer.MIN_VALUE : 0) | bufferedElements}.
+   *
+   * <p>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,15 @@ public StreamObserver<BeamFnApi.Elements> getOutboundObserver() {
    */
   public void registerConsumer(
       String instructionId, CloseableFnDataReceiver<BeamFnApi.Elements> 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);
+    }
     receivers.compute(
         instructionId,
         (unused, existing) -> {
@@ -207,59 +242,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<BeamFnApi.Elements> {
-    @SuppressWarnings("LabelledBreakTarget")
+  private final class InboundObserver
+      implements ClientResponseObserver<BeamFnApi.Elements, BeamFnApi.Elements>, IntConsumer {
+    private volatile @Nullable ClientCallStreamObserver<BeamFnApi.Elements> 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<BeamFnApi.Elements> 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<BeamFnApi.Elements> 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<String> instructionIds = new HashSet<>();
-      for (BeamFnApi.Elements.Data data : value.getDataList()) {
-        instructionIds.add(data.getInstructionId());
+    private void releaseOrParkFlowControlToken() {
+      @Nullable ClientCallStreamObserver<BeamFnApi.Elements> 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<String> 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 +410,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);
+        }
       }
     }