From a16171a495e6fb4ba3cc7908a9d2753ffbe88a4b Mon Sep 17 00:00:00 2001
From: Arun Pandian <pandiana@google.com>
Date: Wed, 12 Aug 2026 15:02:41 +0000
Subject: [PATCH] [Dataflow Streaming] Reduce contention on CommitQueue

Replaces LinkedBlockingQueue with a ConcurrentLinkedQueue and a
Semaphore. The harness threads don't block on the queue anymore.

Improves the internal reshuffle benchmark run time by ~10%.
---
 .../streaming/WeightedBoundedQueue.java       | 42 ++++++++++++-------
 1 file changed, 27 insertions(+), 15 deletions(-)

diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WeightedBoundedQueue.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WeightedBoundedQueue.java
index 5f039be7b00f..497ac18bfc7d 100644
--- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WeightedBoundedQueue.java
+++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WeightedBoundedQueue.java
@@ -17,25 +17,32 @@
  */
 package org.apache.beam.runners.dataflow.worker.streaming;
 
-import java.util.concurrent.LinkedBlockingQueue;
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** Queue bounded by a {@link WeightedSemaphore}. */
-public final class WeightedBoundedQueue<V> {
+public final class WeightedBoundedQueue<V extends @NonNull Object> {
 
-  private final LinkedBlockingQueue<V> queue;
+  private final ConcurrentLinkedQueue<V> queue;
   private final WeightedSemaphore<V> weightedSemaphore;
+  private final Semaphore availableItems;
 
   private WeightedBoundedQueue(
-      LinkedBlockingQueue<V> linkedBlockingQueue, WeightedSemaphore<V> weightedSemaphore) {
-    this.queue = linkedBlockingQueue;
+      ConcurrentLinkedQueue<V> concurrentLinkedQueue, WeightedSemaphore<V> weightedSemaphore) {
+    this.queue = concurrentLinkedQueue;
     this.weightedSemaphore = weightedSemaphore;
+    this.availableItems = new Semaphore(0);
   }
 
-  public static <V> WeightedBoundedQueue<V> create(WeightedSemaphore<V> weightedSemaphore) {
-    return new WeightedBoundedQueue<>(new LinkedBlockingQueue<>(), weightedSemaphore);
+  public static <V extends @NonNull Object> WeightedBoundedQueue<V> create(
+      WeightedSemaphore<V> weightedSemaphore) {
+    return new WeightedBoundedQueue<>(new ConcurrentLinkedQueue<>(), weightedSemaphore);
   }
 
   /**
@@ -43,16 +50,19 @@ public static <V> WeightedBoundedQueue<V> create(WeightedSemaphore<V> weightedSe
    * limit.
    */
   public void put(V value) {
+    checkStateNotNull(value);
     weightedSemaphore.acquireUninterruptibly(value);
     queue.add(value);
+    availableItems.release();
   }
 
   /** Returns and removes the next value, or null if there is no such value. */
   public @Nullable V poll() {
-    @Nullable V result = queue.poll();
-    if (result != null) {
-      weightedSemaphore.release(result);
+    if (!availableItems.tryAcquire()) {
+      return null;
     }
+    V result = checkStateNotNull(queue.poll());
+    weightedSemaphore.release(result);
     return result;
   }
 
@@ -67,22 +77,24 @@ public void put(V value) {
    * @throws InterruptedException if interrupted while waiting
    */
   public @Nullable V poll(long timeout, TimeUnit unit) throws InterruptedException {
-    @Nullable V result = queue.poll(timeout, unit);
-    if (result != null) {
-      weightedSemaphore.release(result);
+    if (!availableItems.tryAcquire(timeout, unit)) {
+      return null;
     }
+    V result = checkStateNotNull(queue.poll());
+    weightedSemaphore.release(result);
     return result;
   }
 
   /** Returns and removes the next value, or blocks until one is available. */
   public V take() throws InterruptedException {
-    V result = queue.take();
+    availableItems.acquire();
+    V result = checkStateNotNull(queue.poll());
     weightedSemaphore.release(result);
     return result;
   }
 
   @VisibleForTesting
   int size() {
-    return queue.size();
+    return availableItems.availablePermits();
   }
 }