[improve][client] Pull v5 blocking receives straight from the receive buffer

The v5 QueueConsumer topped out at ~150K msg/s while the v4 consumer drains
1M msg/s from the same broker. Profiling showed no CPU shortage: three
threads (the user's receive thread, the V5ReceiveQueue executor and the v4
pinned executor) spent most of their time in park/unpark, because every
message went around a serialized ring of cross-thread round trips:

- receive(timeout) always posted a task to the queue executor and parked,
  even when messages were already buffered;
- offer() only took its no-hop fast path below the low watermark, so a
  consumer-bound run (buffer sitting between low and high) took the slow
  path for every message and re-armed the segment loop from the executor
  thread, hopping back to the v4 pinned executor each time.

Mirror the v4 ConsumerBase model instead: the buffer is a thread-safe
GrowableArrayBlockingQueue that the segment loops append to directly and
that blocking take()/poll() drain on the caller's thread; only the pending
async receives and the paused producers stay confined to the executor. The
fast path now holds up to the high watermark. Async waiters are served
through a volatile flag with a re-check after registration, a receive
cancelled before its task runs no longer polls a message into a cancelled
future, and close() terminates the buffer to wake blocked takers.

pulsar-perf consume against a 1M msg/s producer of 10-byte batched
messages: 90-117K msg/s with runaway latency before, 990K-1004K msg/s with
a flat ~70 ms mean latency after.

Author: merlimat

pulsar-client-v5/src/main/java/org/apache/pulsar/client/impl/v5/V5ReceiveQueue.java

@@ -30,28 +30,27 @@
 import java.util.concurrent.TimeUnit;
 import org.apache.pulsar.client.api.v5.Message;
 import org.apache.pulsar.client.api.v5.PulsarClientException;
+import org.apache.pulsar.common.util.collections.GrowableArrayBlockingQueue;
 
 /**
- * Async-native, single-consumer receive queue shared by the v5 scalable consumers.
+ * Async-native receive queue shared by the v5 scalable consumers.
  *
- * <p>Mirrors the v4 {@code ConsumerBase} delivery model: a buffer of ready messages
- * plus a queue of pending receive futures. Both are confined to one pinned executor
- * (obtained from the client's external executor provider, so one thread per consumer),
- * which means a message and a waiter can never cross — no locks, and no lost wakeups.
- * Every receive future is completed on that executor, so user continuations chained on
- * the returned {@link CompletableFuture} never run on a netty IO thread.
+ * <p>Mirrors the v4 {@code ConsumerBase} delivery model. The buffer of ready messages is a
+ * thread-safe blocking queue: producers (the per-segment receive loops, on a v4 client executor)
+ * append to it directly, and the blocking {@link #take()}/{@link #poll(Duration)} pull from it on
+ * the caller's own thread, so in steady state a message never has to cross to another thread just
+ * to be handed over. Only the bookkeeping that needs ordering — the pending {@link #receiveAsync()}
+ * futures and the paused producers — is confined to one pinned executor (obtained from the
+ * client's external executor provider, so one thread per consumer), and every async receive
+ * future is completed there, so user continuations chained on it never run on a netty IO thread.
  *
- * <p>This replaces the previous {@code LinkedTransferQueue} + {@code supplyAsync(take())}
- * approach: {@link #receiveAsync()} parks no thread, is cancelable, and honours timeouts
- * via the client timer instead of blocking a {@code ForkJoinPool.commonPool()} worker.
- *
- * <p>Backpressure: {@link #offer} returns a future that completes only when the buffer has
- * room. Producers gate re-arming on it, so the buffer is bounded by {@code receiverQueueSize}
- * plus one in-flight message per producer (each producer offers at most one message before
- * observing the pause) — modelled on v4 {@code MultiTopicsConsumerImpl}'s pause/resume of
- * sub-consumers, including its fairness rule: once any producer is paused, further offers
- * pause at the half-way mark too, so active producers can't hold the buffer above the resume
- * threshold and starve the paused ones.
+ * <p>Backpressure: {@link #offer} returns a future that completes only when the buffer has room.
+ * Producers gate re-arming on it, so the buffer is bounded by {@code receiverQueueSize} plus one
+ * in-flight message per producer (each producer offers at most one message before observing the
+ * pause) — modelled on v4 {@code MultiTopicsConsumerImpl}'s pause/resume of sub-consumers,
+ * including its fairness rule: once any producer is paused, further offers pause at the half-way
+ * mark too, so active producers can't hold the buffer above the resume threshold and starve the
+ * paused ones.
  */
 final class V5ReceiveQueue<T> {
 
@@ -64,18 +63,23 @@ final class V5ReceiveQueue<T> {
     private final int highWatermark;
     private final int lowWatermark;
 
-    // All three touched only on `executor`, so plain (non-concurrent) collections are safe.
-    private final ArrayDeque<Message<T>> buffer = new ArrayDeque<>();
+    /**
+     * Ready messages. Thread-safe: producers append from their own threads, blocking receives
+     * poll from the caller thread, and the executor drains it for pending async receives.
+     */
+    private final GrowableArrayBlockingQueue<Message<T>> buffer = new GrowableArrayBlockingQueue<>();
+
+    // Both touched only on `executor`, so plain (non-concurrent) collections are safe.
     private final ArrayDeque<CompletableFuture<Message<T>>> pendingReceives = new ArrayDeque<>();
     // Capacity futures handed back to producers that were paused because the buffer was full.
     private final ArrayDeque<CompletableFuture<Void>> capacityWaiters = new ArrayDeque<>();
-    private boolean closed = false;
 
-    // Snapshots of executor-confined state, readable from producer threads so the offer()
-    // fast path can decide without a hop. Written only on `executor`; may lag by the offers
-    // still queued, so each producer (one in-flight offer at a time) can overshoot by one.
-    private volatile int approxBufferSize = 0;
+    // Snapshots of the executor-confined state, readable from producer and receiver threads so
+    // the hot paths can decide without a hop. Written only on `executor` (except `closed`).
+    /** True while {@link #pendingReceives} may hold a future that the next message belongs to. */
+    private volatile boolean hasPendingReceives = false;
     private volatile boolean producersPaused = false;
+    private volatile boolean closed = false;
 
     V5ReceiveQueue(ExecutorService executor, Timer timer, int receiverQueueSize) {
         this.executor = executor;
@@ -86,54 +90,44 @@ final class V5ReceiveQueue<T> {
 
     /**
      * Deposit a freshly-arrived message. Called from the per-segment receive loops (which
-     * run on a v4 client executor). Hands the message straight to a waiting receive future
-     * if there is one, otherwise buffers it.
+     * run on a v4 client executor). Appends straight to the buffer — waking a blocked
+     * {@link #take()}/{@link #poll(Duration)} if it was empty — and hands it over on the
+     * executor only if an async receive is waiting.
      *
      * @return a future that completes when the sink is ready for the next message — right
-     *     away unless the buffer is filling up, in which case it defers until the consumer
-     *     drains it below the low watermark (backpressure).
+     *     away unless the buffer is full, in which case it defers until the consumer drains it
+     *     below the low watermark (backpressure).
      */
     CompletableFuture<Void> offer(Message<T> msg) {
-        // Fast path, decided on the caller thread: while the buffer is comfortably below the
-        // watermarks and nobody is paused, grant capacity with a shared completed future so
-        // the (fast-consumer) hot path pays no allocation and no serialized hop through our
-        // executor before the segment loop re-arms.
-        if (!producersPaused && approxBufferSize < lowWatermark) {
-            executor.execute(() -> doOffer(msg, null));
+        if (closed) {
+            // Dropped, like a message arriving on an already-closed v4 consumer.
+            return READY;
+        }
+        buffer.put(msg);
+        if (hasPendingReceives) {
+            // A receiver registering concurrently re-checks the buffer after publishing the
+            // flag, and we read the flag after appending, so one side always sees the other.
+            executor.execute(this::drainToPendingReceives);
+        }
+        // Fast path, decided on the caller thread: while the buffer is below the high watermark
+        // and nobody is paused, grant capacity with a shared completed future so the hot path
+        // pays no allocation and no serialized hop through our executor before the segment
+        // loop re-arms.
+        if (!producersPaused && buffer.size() < highWatermark) {
             return READY;
         }
         CompletableFuture<Void> capacity = new CompletableFuture<>();
-        executor.execute(() -> doOffer(msg, capacity));
+        executor.execute(() -> grantCapacity(capacity));
         return capacity;
     }
 
-    /** Runs on {@code executor}. {@code capacity} is null when the fast path already granted it. */
-    private void doOffer(Message<T> msg, CompletableFuture<Void> capacity) {
-        if (closed) {
-            if (capacity != null) {
-                capacity.complete(null);
-            }
-            return;
-        }
-        CompletableFuture<Message<T>> waiter = pollWaiter();
-        if (waiter != null) {
-            // Handed straight to a waiting receiver; the buffer didn't grow.
-            waiter.complete(msg);
-            if (capacity != null) {
-                capacity.complete(null);
-            }
-            return;
-        }
-        buffer.add(msg);
-        approxBufferSize = buffer.size();
-        if (capacity == null) {
-            return;
-        }
+    /** Runs on {@code executor}: park the producer while the buffer is full, else grant now. */
+    private void grantCapacity(CompletableFuture<Void> capacity) {
         // Pause when full — or, once any producer is paused, already at the half-way mark, so
         // active producers can't keep the buffer hovering above the resume threshold while the
         // paused ones starve (v4 MultiTopicsConsumerImpl's fairness clause).
-        if (buffer.size() >= highWatermark
-                || (!capacityWaiters.isEmpty() && buffer.size() > lowWatermark)) {
+        if (!closed && (buffer.size() >= highWatermark
+                || (!capacityWaiters.isEmpty() && buffer.size() > lowWatermark))) {
             capacityWaiters.add(capacity);
             producersPaused = true;
         } else {
@@ -156,6 +150,13 @@ private void maybeResumeProducers() {
         }
     }
 
+    /** After a receive pulled straight from the buffer: resume producers once it drained enough. */
+    private void afterDirectPoll() {
+        if (producersPaused && buffer.size() <= lowWatermark) {
+            executor.execute(this::maybeResumeProducers);
+        }
+    }
+
     /** Receive a message, completing as soon as one is available. Never blocks a thread. */
     CompletableFuture<Message<T>> receiveAsync() {
         CompletableFuture<Message<T>> result = new CompletableFuture<>();
@@ -164,13 +165,17 @@ CompletableFuture<Message<T>> receiveAsync() {
                 result.completeExceptionally(alreadyClosed());
                 return;
             }
+            if (result.isDone()) {
+                // Cancelled before we got here: it must not consume a message that has since
+                // been appended straight to the buffer.
+                return;
+            }
             Message<T> msg = buffer.poll();
             if (msg != null) {
-                approxBufferSize = buffer.size();
                 result.complete(msg);
                 maybeResumeProducers();
             } else {
-                pendingReceives.add(result);
+                addPendingReceive(result);
             }
         });
         return result;
@@ -187,9 +192,13 @@ CompletableFuture<Message<T>> receiveAsync(Duration timeout) {
                 result.completeExceptionally(alreadyClosed());
                 return;
             }
+            if (result.isDone()) {
+                // Cancelled before we got here: it must not consume a message that has since
+                // been appended straight to the buffer.
+                return;
+            }
             Message<T> msg = buffer.poll();
             if (msg != null) {
-                approxBufferSize = buffer.size();
                 result.complete(msg);
                 maybeResumeProducers();
                 return;
@@ -199,19 +208,44 @@ CompletableFuture<Message<T>> receiveAsync(Duration timeout) {
                 result.complete(null);
                 return;
             }
-            pendingReceives.add(result);
             Timeout t = timer.newTimeout(ignored -> executor.execute(() -> {
                 if (!result.isDone()) {
                     pendingReceives.remove(result);
+                    hasPendingReceives = !pendingReceives.isEmpty();
                     result.complete(null);
                 }
             }), millis, TimeUnit.MILLISECONDS);
             // Cancel the timer when the message is handed off (or on close) so it doesn't linger.
             result.whenComplete((r, e) -> t.cancel());
+            addPendingReceive(result);
         });
         return result;
     }
 
+    /** Runs on {@code executor}: park an async receive until the next message arrives. */
+    private void addPendingReceive(CompletableFuture<Message<T>> result) {
+        pendingReceives.add(result);
+        hasPendingReceives = true;
+        // A producer that appended between our poll and this flag store saw the flag clear and
+        // posted no drain, so look again now that the flag is published (both are volatile).
+        drainToPendingReceives();
+    }
+
+    /** Runs on {@code executor}: hand buffered messages to pending async receives, in order. */
+    private void drainToPendingReceives() {
+        CompletableFuture<Message<T>> waiter;
+        while ((waiter = pollWaiter()) != null) {
+            Message<T> msg = buffer.poll();
+            if (msg == null) {
+                pendingReceives.addFirst(waiter);
+                break;
+            }
+            waiter.complete(msg);
+        }
+        hasPendingReceives = !pendingReceives.isEmpty();
+        maybeResumeProducers();
+    }
+
     /**
      * Receive up to {@code maxMessages}, blocking (asynchronously) up to {@code timeout} for
      * the batch. Waits for the first message, then opportunistically drains whatever else is
@@ -255,35 +289,53 @@ private CompletableFuture<Void> drainReady(List<Message<T>> batch, int max) {
             while (batch.size() < max && (m = buffer.poll()) != null) {
                 batch.add(m);
             }
-            approxBufferSize = buffer.size();
             maybeResumeProducers();
             done.complete(null);
         });
         return done;
     }
 
-    // --- Blocking views, for the synchronous receive() API. Block only the caller's thread. ---
+    // --- Blocking views, for the synchronous receive() API. Pull straight from the buffer on the
+    // caller's thread; it parks only while the buffer is empty. ---
 
     Message<T> take() throws PulsarClientException {
+        if (closed) {
+            throw alreadyClosed();
+        }
+        Message<T> msg;
         try {
-            return receiveAsync().get();
+            msg = buffer.take();
         } catch (InterruptedException e) {
+            if (closed) {
+                // close() terminates the buffer, which wakes blocked takers this way.
+                throw alreadyClosed();
+            }
             Thread.currentThread().interrupt();
             throw new PulsarClientException("Receive interrupted", e);
-        } catch (ExecutionException e) {
-            throw unwrap(e);
         }
+        afterDirectPoll();
+        return msg;
     }
 
     Message<T> poll(Duration timeout) throws PulsarClientException {
+        if (closed) {
+            throw alreadyClosed();
+        }
+        Message<T> msg;
         try {
-            return receiveAsync(timeout).get();
+            msg = buffer.poll(timeout.toNanos(), TimeUnit.NANOSECONDS);
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
             throw new PulsarClientException("Receive interrupted", e);
-        } catch (ExecutionException e) {
-            throw unwrap(e);
         }
+        if (msg == null) {
+            if (closed) {
+                throw alreadyClosed();
+            }
+            return null;
+        }
+        afterDirectPoll();
+        return msg;
     }
 
     List<Message<T>> receiveMulti(int maxMessages, Duration timeout) throws PulsarClientException {
@@ -299,22 +351,24 @@ List<Message<T>> receiveMulti(int maxMessages, Duration timeout) throws PulsarCl
 
     /** Fail any outstanding receives so blocked/awaiting callers wake instead of hanging forever. */
     void close() {
+        closed = true;
+        // Wake blocking receivers parked on the buffer; anything offered from now on is dropped.
+        buffer.terminate(null);
         executor.execute(() -> {
-            closed = true;
             CompletableFuture<Message<T>> waiter;
             while ((waiter = pendingReceives.poll()) != null) {
                 if (!waiter.isDone()) {
                     waiter.completeExceptionally(alreadyClosed());
                 }
             }
+            hasPendingReceives = false;
             // Release any paused producers so their receive loops re-arm and observe the close.
             CompletableFuture<Void> capacity;
             while ((capacity = capacityWaiters.poll()) != null) {
                 capacity.complete(null);
             }
             producersPaused = false;
             buffer.clear();
-            approxBufferSize = 0;
         });
     }