[client] Change the implementation of LookupQueue from ArrayBlockingQueue to an unbounded LinkedBlockingDeque to avoid deadlocks during re-enqueue

Author: swuferhong

fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQueue.java

@@ -18,14 +18,16 @@
 package org.apache.fluss.client.lookup;
 
 import org.apache.fluss.annotation.Internal;
+import org.apache.fluss.annotation.VisibleForTesting;
 import org.apache.fluss.config.ConfigOptions;
 import org.apache.fluss.config.Configuration;
 
 import javax.annotation.concurrent.ThreadSafe;
-
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -38,13 +40,18 @@ class LookupQueue {
 
     private volatile boolean closed;
     // buffering both the Lookup and PrefixLookup.
+    // TODO This queue could be refactored into a memory-managed queue similar to
+    // RecordAccumulator, which would significantly improve the efficiency of lookup batching. Trace
+    // by https://github.com/apache/fluss/issues/2124
     private final ArrayBlockingQueue<AbstractLookupQuery<?>> lookupQueue;
+    private final BlockingQueue<AbstractLookupQuery<?>> reEnqueuedLookupQueue;
     private final int maxBatchSize;
     private final long batchTimeoutNanos;
 
     LookupQueue(Configuration conf) {
         this.lookupQueue =
                 new ArrayBlockingQueue<>(conf.get(ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE));
+        this.reEnqueuedLookupQueue = new LinkedBlockingQueue<>();
         this.maxBatchSize = conf.get(ConfigOptions.CLIENT_LOOKUP_MAX_BATCH_SIZE);
         this.batchTimeoutNanos = conf.get(ConfigOptions.CLIENT_LOOKUP_BATCH_TIMEOUT).toNanos();
         this.closed = false;
@@ -63,8 +70,21 @@ void appendLookup(AbstractLookupQuery<?> lookup) {
         }
     }
 
+    void reEnqueue(AbstractLookupQuery<?> lookup) {
+        if (closed) {
+            throw new IllegalStateException(
+                    "Can not re-enqueue lookup operation since the LookupQueue is closed.");
+        }
+
+        try {
+            reEnqueuedLookupQueue.put(lookup);
+        } catch (InterruptedException e) {
+            lookup.future().completeExceptionally(e);
+        }
+    }
+
     boolean hasUnDrained() {
-        return !lookupQueue.isEmpty();
+        return !lookupQueue.isEmpty() || !reEnqueuedLookupQueue.isEmpty();
     }
 
     /** Drain a batch of {@link LookupQuery}s from the lookup queue. */
@@ -78,6 +98,16 @@ List<AbstractLookupQuery<?>> drain() throws Exception {
                 break;
             }
 
+            while (!lookupQueue.isEmpty() && count < maxBatchSize) {
+                AbstractLookupQuery<?> lookup =
+                        reEnqueuedLookupQueue.poll(waitNanos, TimeUnit.NANOSECONDS);
+                if (lookup == null) {
+                    break;
+                }
+                lookupOperations.add(lookup);
+                count++;
+            }
+
             AbstractLookupQuery<?> lookup = lookupQueue.poll(waitNanos, TimeUnit.NANOSECONDS);
             if (lookup == null) {
                 break;
@@ -97,10 +127,21 @@ List<AbstractLookupQuery<?>> drain() throws Exception {
     List<AbstractLookupQuery<?>> drainAll() {
         List<AbstractLookupQuery<?>> lookupOperations = new ArrayList<>(lookupQueue.size());
         lookupQueue.drainTo(lookupOperations);
+        reEnqueuedLookupQueue.drainTo(lookupOperations);
         return lookupOperations;
     }
 
     public void close() {
         closed = true;
     }
+
+    @VisibleForTesting
+    ArrayBlockingQueue<AbstractLookupQuery<?>> getLookupQueue() {
+        return lookupQueue;
+    }
+
+    @VisibleForTesting
+    BlockingQueue<AbstractLookupQuery<?>> getReEnqueuedLookupQueue() {
+        return reEnqueuedLookupQueue;
+    }
 }

fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java

@@ -38,12 +38,10 @@
 import org.apache.fluss.rpc.messages.PrefixLookupResponse;
 import org.apache.fluss.rpc.protocol.ApiError;
 import org.apache.fluss.utils.types.Tuple2;
-
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import javax.annotation.Nullable;
-
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -157,6 +155,9 @@ private Map<Tuple2<Integer, LookupType>, List<AbstractLookupQuery<?>>> groupByLe
             // lookup the leader node
             TableBucket tb = lookup.tableBucket();
             try {
+                // TODO Metadata requests are being sent too frequently here. consider first
+                // collecting the tables that need to be updated and then sending them together in
+                // one request.
                 leader = metadataUpdater.leaderFor(lookup.tablePath(), tb);
             } catch (Exception e) {
                 // if leader is not found, re-enqueue the lookup to send again.
@@ -423,7 +424,7 @@ private void handlePrefixLookupException(
     }
 
     private void reEnqueueLookup(AbstractLookupQuery<?> lookup) {
-        lookupQueue.appendLookup(lookup);
+        lookupQueue.reEnqueue(lookup);
     }
 
     private boolean canRetry(AbstractLookupQuery<?> lookup, Exception exception) {

fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupQueueTest.java

@@ -19,13 +19,18 @@
 
 import org.apache.fluss.config.Configuration;
 import org.apache.fluss.metadata.TableBucket;
-
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 
 import static org.apache.fluss.config.ConfigOptions.CLIENT_LOOKUP_BATCH_TIMEOUT;
 import static org.apache.fluss.config.ConfigOptions.CLIENT_LOOKUP_MAX_BATCH_SIZE;
+import static org.apache.fluss.config.ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE;
 import static org.apache.fluss.record.TestData.DATA1_TABLE_PATH_PK;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 
 /** Tests for {@link LookupQueue}. */
 class LookupQueueTest {
@@ -59,6 +64,59 @@ void testDrainMaxBatchSize() throws Exception {
         assertThat(queue.hasUnDrained()).isFalse();
     }
 
+    @Test
+    @Timeout(value = 10, unit = TimeUnit.SECONDS)
+    void testAppendLookupBlocksWhenQueueIsFull() throws Exception {
+        Configuration conf = new Configuration();
+        conf.set(CLIENT_LOOKUP_QUEUE_SIZE, 5);
+        LookupQueue queue = new LookupQueue(conf);
+
+        appendLookups(queue, 5);
+        assertThat(queue.getLookupQueue()).hasSize(5);
+
+        CompletableFuture<Void> future =
+                CompletableFuture.runAsync(
+                        () -> {
+                            appendLookups(queue, 1); // will be blocked.
+                        });
+
+        assertFalse(future.isDone(), "appendLookup should block and not complete immediately");
+
+        Thread.sleep(100);
+        assertFalse(future.isDone(), "Still blocked after 100ms");
+
+        queue.drain();
+        future.get(1, TimeUnit.SECONDS);
+    }
+
+    @Test
+    void testReEnqueueNotBlock() throws Exception {
+        Configuration conf = new Configuration();
+        conf.set(CLIENT_LOOKUP_QUEUE_SIZE, 5);
+        conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 5);
+        LookupQueue queue = new LookupQueue(conf);
+
+        appendLookups(queue, 5);
+        assertThat(queue.getLookupQueue()).hasSize(5);
+        assertThat(queue.getReEnqueuedLookupQueue()).hasSize(0);
+
+        queue.reEnqueue(
+                new LookupQuery(DATA1_TABLE_PATH_PK, new TableBucket(1, 1), new byte[] {0}));
+        assertThat(queue.getLookupQueue()).hasSize(5);
+        // This batch will be put into re-enqueued lookup queue.
+        assertThat(queue.getReEnqueuedLookupQueue()).hasSize(1);
+        assertThat(queue.hasUnDrained()).isTrue();
+
+        assertThat(queue.drain()).hasSize(5);
+        // drain re-enqueued lookup first.
+        assertThat(queue.getReEnqueuedLookupQueue().isEmpty()).isTrue();
+        assertThat(queue.getLookupQueue()).hasSize(1);
+        assertThat(queue.hasUnDrained()).isTrue();
+
+        assertThat(queue.drain()).hasSize(1);
+        assertThat(queue.hasUnDrained()).isFalse();
+    }
+
     private static void appendLookups(LookupQueue queue, int count) {
         for (int i = 0; i < count; i++) {
             queue.appendLookup(

website/docs/engine-flink/options.md

@@ -117,7 +117,7 @@ See more details about [ALTER TABLE ... SET](engine-flink/ddl.md#set-properties)
 | lookup.partial-cache.max-rows            | Long       | (None)            | The maximum number of rows to store in the cache.                                                                                   |
 | client.lookup.queue-size                 | Integer    | 25600             | The maximum number of pending lookup operations.                                                                                    |
 | client.lookup.max-batch-size             | Integer    | 128               | The maximum batch size of merging lookup operations to one lookup request.                                                          |
-| client.lookup.max-inflight-requests      | Integer    | 128               | The maximum number of unacknowledged lookup requests for lookup operations.                                                         |
+| client.lookup.max-inflight-requests      | Integer    | 128               | The maximum number of unacknowledged lookup requests for lookup operations.                                                         |:
 | client.lookup.batch-timeout              | Duration   | 100ms             | The maximum time to wait for the lookup batch to full, if this timeout is reached, the lookup batch will be closed to send.         | 
 | client.lookup.max-retries                | Integer    | Integer.MAX_VALUE | Setting a value greater than zero will cause the client to resend any lookup request that fails with a potentially transient error. |