diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupBatch.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupBatch.java index 0d378ab4f21..347a0d2b030 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupBatch.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupBatch.java @@ -18,71 +18,66 @@ package org.apache.fluss.client.lookup; import org.apache.fluss.annotation.Internal; -import org.apache.fluss.exception.FlussRuntimeException; import org.apache.fluss.metadata.TableBucket; -import javax.annotation.Nullable; - import java.util.ArrayList; +import java.util.Collections; import java.util.List; -/** A batch that contains the lookup operations that send to same tablet bucket together. */ +/** A batch of lookup operations accumulated for the same table bucket and lookup kind. */ @Internal -public class LookupBatch { +final class LookupBatch { - private final LookupBatchKey lookupBatchKey; + private final TableBucket tableBucket; + private final LookupType lookupType; + private final boolean historical; + private final long createdNanos; + private final List> lookups; - private final List lookups; + private boolean completed; - LookupBatch(LookupBatchKey lookupBatchKey) { - this.lookupBatchKey = lookupBatchKey; + LookupBatch(AbstractLookupQuery firstLookup, long createdNanos) { + this.tableBucket = firstLookup.tableBucket(); + this.lookupType = firstLookup.lookupType(); + this.historical = firstLookup.originalPartitionName() != null; + this.createdNanos = createdNanos; this.lookups = new ArrayList<>(); + this.lookups.add(firstLookup); } - public void addLookup(LookupQuery lookup) { + void addLookup(AbstractLookupQuery lookup) { lookups.add(lookup); } - public List lookups() { - return lookups; + TableBucket tableBucket() { + return tableBucket; } - public TableBucket tableBucket() { - return lookupBatchKey.tableBucket(); + LookupType lookupType() { + return lookupType; } - public @Nullable String originalPartitionName() { - return lookupBatchKey.originalPartitionName(); + boolean historical() { + return historical; } - LookupBatchKey lookupBatchKey() { - return lookupBatchKey; + List> lookups() { + return Collections.unmodifiableList(lookups); } - /** Complete the lookup operations using given values . */ - public void complete(List values) { - // if the size of return values of lookup operation are not equal to the number of lookups, - // should complete an exception. - if (values.size() != lookups.size()) { - completeExceptionally( - new FlussRuntimeException( - String.format( - "The number of return values of lookup operation is not equal to the number of " - + "lookups. Return %d values, but expected %d.", - values.size(), lookups.size()))); - } else { - for (int i = 0; i < values.size(); i++) { - AbstractLookupQuery lookup = lookups.get(i); - // single value. - lookup.future().complete(values.get(i)); - } - } + int size() { + return lookups.size(); + } + + long waitedNanos(long nowNanos) { + return nowNanos - createdNanos; } - /** Complete the lookup operations with given exception. */ - public void completeExceptionally(Exception exception) { - for (LookupQuery lookup : lookups) { - lookup.future().completeExceptionally(exception); + boolean markCompleted() { + if (completed) { + return false; } + completed = true; + return true; } } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQueue.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQueue.java index 515f0ebb01c..9a7959b2ad3 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQueue.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQueue.java @@ -21,143 +21,472 @@ import org.apache.fluss.annotation.VisibleForTesting; import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; +import org.apache.fluss.metadata.TableBucket; import javax.annotation.concurrent.ThreadSafe; +import java.util.ArrayDeque; import java.util.ArrayList; +import java.util.Deque; +import java.util.HashMap; +import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; -import java.util.concurrent.ArrayBlockingQueue; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.LinkedBlockingQueue; +import java.util.Map; +import java.util.Objects; +import java.util.Set; import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.ReentrantLock; -/** - * A queue that buffers the pending lookup operations and provides a list of {@link LookupQuery} - * when call method {@link #drain()}. - */ +import static org.apache.fluss.utils.Preconditions.checkArgument; +import static org.apache.fluss.utils.Preconditions.checkState; + +/** Accumulates lookup operations into independently drainable batches per table bucket. */ @ThreadSafe @Internal 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> lookupQueue; - private final BlockingQueue> reEnqueuedLookupQueue; + private final ReentrantLock stateLock = new ReentrantLock(); + + // Appenders wait here when their current batch is full or absent and there is not enough queue + // capacity to create another batch. + private final Condition appendCondition = stateLock.newCondition(); + + // The sender waits here when no batch or retry is currently drainable. + private final Condition drainCondition = stateLock.newCondition(); + private final Map> batches; + private final Deque> reEnqueuedLookups; + private final Map inFlightRequestsByBucket; + private final int queueSize; private final int maxBatchSize; + private final int maxInFlightRequestsPerBucket; private final long batchTimeoutNanos; + private boolean closed; + private int remainingQueueSize; + private int waitingAppenders; + LookupQueue(Configuration conf) { - this.lookupQueue = - new ArrayBlockingQueue<>(conf.get(ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE)); - this.reEnqueuedLookupQueue = new LinkedBlockingQueue<>(); + this.queueSize = conf.get(ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE); this.maxBatchSize = conf.get(ConfigOptions.CLIENT_LOOKUP_MAX_BATCH_SIZE); + this.maxInFlightRequestsPerBucket = + conf.get(ConfigOptions.CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET); + checkArgument(queueSize > 0, "Lookup queue size must be greater than 0."); + checkArgument(maxBatchSize > 0, "Lookup batch size must be greater than 0."); + checkArgument( + queueSize >= maxBatchSize, + "Lookup queue size (%s) must not be smaller than batch size (%s).", + queueSize, + maxBatchSize); + checkArgument( + maxInFlightRequestsPerBucket > 0, + "Maximum in-flight lookup requests per bucket must be greater than 0."); + + this.batches = new LinkedHashMap<>(); + this.reEnqueuedLookups = new ArrayDeque<>(); + this.inFlightRequestsByBucket = new HashMap<>(); + this.remainingQueueSize = queueSize; this.batchTimeoutNanos = conf.get(ConfigOptions.CLIENT_LOOKUP_BATCH_TIMEOUT).toNanos(); - this.closed = false; } void appendLookup(AbstractLookupQuery lookup) { - if (closed) { - throw new IllegalStateException( - "Can not append lookup operation since the LookupQueue is closed."); + AccumulatorKey key = AccumulatorKey.of(lookup); + stateLock.lock(); + try { + while (true) { + if (closed) { + throw new IllegalStateException( + "Can not append lookup operation since the LookupQueue is closed."); + } + + LookupBatch appendedBatch = tryAppend(key, lookup); + if (appendedBatch != null) { + if (appendedBatch.size() == maxBatchSize) { + drainCondition.signal(); + } + return; + } + + if (remainingQueueSize >= maxBatchSize) { + Deque bucketBatches = + batches.computeIfAbsent(key, ignored -> new ArrayDeque<>()); + bucketBatches.addLast(new LookupBatch(lookup, System.nanoTime())); + remainingQueueSize -= maxBatchSize; + // Wake the sender to start the batch timeout. + drainCondition.signal(); + return; + } + + waitingAppenders++; + try { + // A waiting appender makes underfilled batches ready, so wake the sender before + // waiting for one completed batch to return its reserved capacity. + if (waitingAppenders == 1) { + drainCondition.signal(); + } + appendCondition.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + lookup.future().completeExceptionally(e); + return; + } finally { + waitingAppenders--; + } + } + } finally { + stateLock.unlock(); } + } + /** Re-enqueues a retry without blocking an RPC callback thread on queue capacity. */ + void reEnqueue(AbstractLookupQuery lookup) { + stateLock.lock(); try { - lookupQueue.put(lookup); - } catch (InterruptedException e) { - lookup.future().completeExceptionally(e); + if (closed) { + throw new IllegalStateException( + "Can not re-enqueue lookup operation since the LookupQueue is closed."); + } + reEnqueuedLookups.addLast(lookup); + drainCondition.signal(); + } finally { + stateLock.unlock(); } } - void reEnqueue(AbstractLookupQuery lookup) { - if (closed) { - throw new IllegalStateException( - "Can not re-enqueue lookup operation since the LookupQueue is closed."); + boolean hasUnDrained() { + stateLock.lock(); + try { + return hasUnDrainedUnsafe(); + } finally { + stateLock.unlock(); } + } + + /** Waits until at least one bucket batch is ready, then drains all currently ready batches. */ + List drain() throws InterruptedException { + return drain(false); + } + + /** Drains all batches that can currently be sent while respecting the per-bucket limit. */ + List drainAll() throws InterruptedException { + return drain(true); + } + void addInFlightRequests(Set tableBuckets) { + stateLock.lock(); try { - reEnqueuedLookupQueue.put(lookup); - } catch (InterruptedException e) { - lookup.future().completeExceptionally(e); + for (TableBucket tableBucket : tableBuckets) { + inFlightRequestsByBucket.merge(tableBucket, 1, Integer::sum); + } + } finally { + stateLock.unlock(); } } - boolean hasUnDrained() { - return !lookupQueue.isEmpty() || !reEnqueuedLookupQueue.isEmpty(); - } - - /** Drain a batch of {@link LookupQuery}s from the lookup queue. */ - List> drain() throws Exception { - final long startNanos = System.nanoTime(); - List> lookupOperations = new ArrayList<>(maxBatchSize); - int count = 0; - while (true) { - long waitNanos = batchTimeoutNanos - (System.nanoTime() - startNanos); - if (waitNanos <= 0) { - break; - } - - long nextRetryDelayNanos = Long.MAX_VALUE; - int reEnqueuedToCheck = reEnqueuedLookupQueue.size(); - while (reEnqueuedToCheck > 0 && count < maxBatchSize) { - AbstractLookupQuery lookup = reEnqueuedLookupQueue.poll(); - if (lookup == null) { - break; + void completeInFlightRequests(Set tableBuckets) { + stateLock.lock(); + try { + boolean bucketDroppedBelowInFlightLimit = false; + for (TableBucket tableBucket : tableBuckets) { + int inFlight = inFlightRequestsByBucket.getOrDefault(tableBucket, 0); + int remaining = inFlight - 1; + checkState( + remaining >= 0, + "No in-flight lookup request exists for table bucket %s.", + tableBucket); + if (inFlight >= maxInFlightRequestsPerBucket + && remaining < maxInFlightRequestsPerBucket) { + bucketDroppedBelowInFlightLimit = true; } - long retryDelayMs = lookup.nextRetryTimeMs() - System.currentTimeMillis(); - if (retryDelayMs <= 0) { - lookupOperations.add(lookup); - count++; + if (remaining == 0) { + inFlightRequestsByBucket.remove(tableBucket); } else { - nextRetryDelayNanos = - Math.min( - nextRetryDelayNanos, - TimeUnit.MILLISECONDS.toNanos(retryDelayMs)); - reEnqueuedLookupQueue.add(lookup); + inFlightRequestsByBucket.put(tableBucket, remaining); + } + } + if (bucketDroppedBelowInFlightLimit) { + drainCondition.signal(); + } + } finally { + stateLock.unlock(); + } + } + + void releaseBatchCapacity(LookupBatch batch) { + stateLock.lock(); + try { + if (!batch.markCompleted()) { + return; + } + remainingQueueSize += maxBatchSize; + checkState( + remainingQueueSize <= queueSize, + "Lookup queue capacity was released more than once."); + + // Retries must reserve returned capacity before regular appenders are woken up. + // Otherwise, under sustained full-queue pressure, new lookups can repeatedly take every + // released slot and starve the retry deque. + if (!reEnqueuedLookups.isEmpty()) { + int reEnqueuedBefore = reEnqueuedLookups.size(); + boolean retryWaitingForCapacity = appendReadyRetries(System.currentTimeMillis()); + if (reEnqueuedLookups.size() < reEnqueuedBefore || retryWaitingForCapacity) { + drainCondition.signal(); + } + } + + if (waitingAppenders > 0 && remainingQueueSize >= maxBatchSize) { + appendCondition.signalAll(); + } + } finally { + stateLock.unlock(); + } + } + + void close() { + stateLock.lock(); + try { + closed = true; + appendCondition.signalAll(); + drainCondition.signalAll(); + } finally { + stateLock.unlock(); + } + } + + private List drain(boolean drainAll) throws InterruptedException { + stateLock.lock(); + try { + while (true) { + long nowNanos = System.nanoTime(); + boolean retryWaitingForCapacity = appendReadyRetries(System.currentTimeMillis()); + List readyBatches = + drainReadyBatches(nowNanos, drainAll, retryWaitingForCapacity); + if (!readyBatches.isEmpty()) { + return readyBatches; + } + if (!hasUnDrainedUnsafe() && (drainAll || closed)) { + return readyBatches; + } + + long waitNanos = nextReadyWaitNanos(nowNanos, System.currentTimeMillis()); + if (waitNanos == Long.MAX_VALUE) { + drainCondition.await(); + } else if (waitNanos > 0) { + drainCondition.awaitNanos(waitNanos); } - reEnqueuedToCheck--; } + } finally { + stateLock.unlock(); + } + } + + private LookupBatch tryAppend(AccumulatorKey key, AbstractLookupQuery lookup) { + Deque bucketBatches = batches.get(key); + if (bucketBatches == null) { + return null; + } + LookupBatch lastBatch = bucketBatches.peekLast(); + if (lastBatch == null || lastBatch.size() >= maxBatchSize) { + return null; + } + lastBatch.addLookup(lookup); + return lastBatch; + } + + private boolean appendReadyRetries(long nowMs) { + boolean waitingForCapacity = false; + Iterator> iterator = reEnqueuedLookups.iterator(); + while (iterator.hasNext()) { + AbstractLookupQuery lookup = iterator.next(); + if (lookup.nextRetryTimeMs() > nowMs) { + continue; + } + + AccumulatorKey key = AccumulatorKey.of(lookup); + if (tryAppend(key, lookup) != null) { + iterator.remove(); + } else if (remainingQueueSize >= maxBatchSize) { + Deque bucketBatches = + batches.computeIfAbsent(key, ignored -> new ArrayDeque<>()); + bucketBatches.addLast(new LookupBatch(lookup, System.nanoTime())); + remainingQueueSize -= maxBatchSize; + iterator.remove(); + } else { + waitingForCapacity = true; + } + } + return waitingForCapacity; + } + + /** + * Drains ready batches fairly across accumulator keys. + * + *

The queue stores batches by {@link AccumulatorKey}, where different lookup kinds for the + * same {@link TableBucket} use different keys. If we keep taking from one key until that key is + * empty or the bucket permit is exhausted, a key with continuous backlog can stay at the head + * of the {@link LinkedHashMap} and starve later keys for the same bucket. To avoid that, each + * pass selects at most one batch from each key. A selected key that still has more batches is + * moved to the tail, so the next pass starts from the other keys first. + */ + private List drainReadyBatches( + long nowNanos, boolean drainAll, boolean retryWaitingForCapacity) { + List readyBatches = new ArrayList<>(); + Map selectedByBucket = new HashMap<>(); + boolean drainedAnyBatch; + do { + drainedAnyBatch = false; + List keys = new ArrayList<>(batches.keySet()); + for (AccumulatorKey key : keys) { + Deque bucketBatches = batches.get(key); + if (bucketBatches == null || bucketBatches.isEmpty()) { + batches.remove(key); + continue; + } - long lookupWaitNanos = waitNanos; - if (count == 0 && nextRetryDelayNanos != Long.MAX_VALUE) { - lookupWaitNanos = Math.min(waitNanos, Math.max(1L, nextRetryDelayNanos)); + LookupBatch batch = bucketBatches.peekFirst(); + TableBucket tableBucket = batch.tableBucket(); + int inFlight = inFlightRequestsByBucket.getOrDefault(tableBucket, 0); + int selected = selectedByBucket.getOrDefault(tableBucket, 0); + if (inFlight + selected >= maxInFlightRequestsPerBucket) { + continue; + } + if (!drainAll + && !closed + && batch.size() < maxBatchSize + && batch.waitedNanos(nowNanos) < batchTimeoutNanos + && waitingAppenders == 0 + && !retryWaitingForCapacity) { + continue; + } + + readyBatches.add(bucketBatches.removeFirst()); + selectedByBucket.put(tableBucket, selected + 1); + drainedAnyBatch = true; + if (bucketBatches.isEmpty()) { + batches.remove(key); + } else { + // This key has used its turn in this pass. Move it behind the other keys so a + // later drain can pick another lookup kind for the same bucket before this + // key's remaining backlog. + batches.remove(key); + batches.put(key, bucketBatches); + } } - AbstractLookupQuery lookup = lookupQueue.poll(lookupWaitNanos, TimeUnit.NANOSECONDS); - if (lookup == null) { - break; + } while (drainedAnyBatch); + return readyBatches; + } + + private long nextReadyWaitNanos(long nowNanos, long nowMs) { + long waitNanos = Long.MAX_VALUE; + for (Deque bucketBatches : batches.values()) { + LookupBatch batch = bucketBatches.peekFirst(); + if (batch == null + || inFlightRequestsByBucket.getOrDefault(batch.tableBucket(), 0) + >= maxInFlightRequestsPerBucket) { + continue; } - lookupOperations.add(lookup); - count++; - int transferred = lookupQueue.drainTo(lookupOperations, maxBatchSize - count); - count += transferred; - if (count >= maxBatchSize) { - break; + waitNanos = + Math.min( + waitNanos, + Math.max(0L, batchTimeoutNanos - batch.waitedNanos(nowNanos))); + } + for (AbstractLookupQuery lookup : reEnqueuedLookups) { + long retryDelayMs = lookup.nextRetryTimeMs() - nowMs; + if (retryDelayMs > 0) { + waitNanos = Math.min(waitNanos, TimeUnit.MILLISECONDS.toNanos(retryDelayMs)); } } - return lookupOperations; + return waitNanos; } - /** Drain all the {@link LookupQuery}s from the lookup queue. */ - List> drainAll() { - List> lookupOperations = new ArrayList<>(lookupQueue.size()); - lookupQueue.drainTo(lookupOperations); - reEnqueuedLookupQueue.drainTo(lookupOperations); - return lookupOperations; + private boolean hasUnDrainedUnsafe() { + return !batches.isEmpty() || !reEnqueuedLookups.isEmpty(); } - public void close() { - closed = true; + @VisibleForTesting + int remainingQueueSize() { + stateLock.lock(); + try { + return remainingQueueSize; + } finally { + stateLock.unlock(); + } + } + + int maxBatchSize() { + return maxBatchSize; } @VisibleForTesting - ArrayBlockingQueue> getLookupQueue() { - return lookupQueue; + int queuedBatchCount() { + stateLock.lock(); + try { + int count = 0; + for (Deque bucketBatches : batches.values()) { + count += bucketBatches.size(); + } + return count; + } finally { + stateLock.unlock(); + } } @VisibleForTesting - BlockingQueue> getReEnqueuedLookupQueue() { - return reEnqueuedLookupQueue; + int reEnqueuedLookupCount() { + stateLock.lock(); + try { + return reEnqueuedLookups.size(); + } finally { + stateLock.unlock(); + } + } + + @VisibleForTesting + int waitingAppenderCount() { + stateLock.lock(); + try { + return waitingAppenders; + } finally { + stateLock.unlock(); + } + } + + private static final class AccumulatorKey { + private final TableBucket tableBucket; + private final LookupType lookupType; + private final boolean historical; + + private AccumulatorKey(TableBucket tableBucket, LookupType lookupType, boolean historical) { + this.tableBucket = tableBucket; + this.lookupType = lookupType; + this.historical = historical; + } + + private static AccumulatorKey of(AbstractLookupQuery lookup) { + return new AccumulatorKey( + lookup.tableBucket(), + lookup.lookupType(), + lookup.originalPartitionName() != null); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + AccumulatorKey that = (AccumulatorKey) o; + return historical == that.historical + && Objects.equals(tableBucket, that.tableBucket) + && lookupType == that.lookupType; + } + + @Override + public int hashCode() { + return Objects.hash(tableBucket, lookupType, historical); + } } } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupRequestBatch.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupRequestBatch.java new file mode 100644 index 00000000000..966283ea0ec --- /dev/null +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupRequestBatch.java @@ -0,0 +1,89 @@ +/* + * 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.fluss.client.lookup; + +import org.apache.fluss.annotation.Internal; +import org.apache.fluss.exception.FlussRuntimeException; +import org.apache.fluss.metadata.TableBucket; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.List; + +/** Lookup operations encoded in one bucket entry of a lookup RPC. */ +@Internal +public class LookupRequestBatch { + + private final LookupBatchKey lookupBatchKey; + + private final List lookups; + + LookupRequestBatch(LookupBatchKey lookupBatchKey) { + this.lookupBatchKey = lookupBatchKey; + this.lookups = new ArrayList<>(); + } + + /** Adds a lookup operation to this RPC bucket entry. */ + public void addLookup(LookupQuery lookup) { + lookups.add(lookup); + } + + /** Returns the lookup operations in this RPC bucket entry. */ + public List lookups() { + return lookups; + } + + /** Returns the table bucket targeted by this RPC bucket entry. */ + public TableBucket tableBucket() { + return lookupBatchKey.tableBucket(); + } + + /** Returns the original partition name for a historical lookup, or null otherwise. */ + public @Nullable String originalPartitionName() { + return lookupBatchKey.originalPartitionName(); + } + + LookupBatchKey lookupBatchKey() { + return lookupBatchKey; + } + + /** Complete the lookup operations using given values. */ + public void complete(List values) { + if (values.size() != lookups.size()) { + completeExceptionally( + new FlussRuntimeException( + String.format( + "The number of return values of lookup operation is not equal to the number of " + + "lookups. Return %d values, but expected %d.", + values.size(), lookups.size()))); + } else { + for (int i = 0; i < values.size(); i++) { + AbstractLookupQuery lookup = lookups.get(i); + lookup.future().complete(values.get(i)); + } + } + } + + /** Complete the lookup operations with given exception. */ + public void completeExceptionally(Exception exception) { + for (LookupQuery lookup : lookups) { + lookup.future().completeExceptionally(exception); + } + } +} diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java index d1b59a75cc2..332222c2628 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupSender.java @@ -78,7 +78,7 @@ class LookupSender implements Runnable { private final LookupQueue lookupQueue; - private final Semaphore maxInFlightReuqestsSemaphore; + private final Semaphore maxInFlightRequestsSemaphore; private final int maxRetries; @@ -97,7 +97,7 @@ class LookupSender implements Runnable { int maxRequestTimeoutMs) { this.metadataUpdater = metadataUpdater; this.lookupQueue = lookupQueue; - this.maxInFlightReuqestsSemaphore = new Semaphore(maxFlightRequests); + this.maxInFlightRequestsSemaphore = new Semaphore(maxFlightRequests); this.maxRetries = maxRetries; this.running = true; this.acks = acks; @@ -123,7 +123,7 @@ public void run() { // okay we stopped accepting requests but there may still be requests in the accumulator or // waiting for acknowledgment, wait until these are completed. // TODO Check the in flight request count in the accumulator. - if (!forceClose && lookupQueue.hasUnDrained()) { + while (!forceClose && lookupQueue.hasUnDrained()) { try { runOnce(true); } catch (Exception e) { @@ -137,83 +137,134 @@ public void run() { /** Run a single iteration of sending. */ private void runOnce(boolean drainAll) throws Exception { - List> lookups = - drainAll ? lookupQueue.drainAll() : lookupQueue.drain(); - sendLookups(lookups); - } - - private void sendLookups(List> lookups) throws Exception { - if (lookups.isEmpty()) { - return; + List lookupBatches = drainAll ? lookupQueue.drainAll() : lookupQueue.drain(); + Map, List> batchesByLeaderAndType = + groupByLeaderAndType(lookupBatches); + for (Map.Entry, List> entry : + batchesByLeaderAndType.entrySet()) { + sendLookupBatches(entry.getKey().f0, entry.getKey().f1, entry.getValue()); } - // group by to lookup batches - Map, List>> lookupBatches = - groupByLeaderAndType(lookups); - - // if no lookup batches, sleep a bit to avoid busy loop. This case will happen when there is - // no leader for all the lookup request in queue. - if (lookupBatches.isEmpty() && !lookupQueue.hasUnDrained()) { - // TODO: may use wait/notify mechanism to avoid active sleep, and use a dynamic sleep - // time based on the request waited time. - Thread.sleep(100); - } - - // now, send the batches - lookupBatches.forEach( - (destAndType, batch) -> sendLookups(destAndType.f0, destAndType.f1, batch)); } - private Map, List>> groupByLeaderAndType( - List> lookups) { - // -> lookup batches - Map, List>> lookupBatchesByLeader = - new HashMap<>(); - for (AbstractLookupQuery lookup : lookups) { - int leader; - // lookup the leader node - TableBucket tb = lookup.tableBucket(); + private Map, List> groupByLeaderAndType( + List lookupBatches) { + Map, List> batchesByLeaderAndType = + new LinkedHashMap<>(); + for (LookupBatch lookupBatch : lookupBatches) { + AbstractLookupQuery firstLookup = lookupBatch.lookups().get(0); + int destination; 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); + destination = + metadataUpdater.leaderFor( + firstLookup.tablePath(), lookupBatch.tableBucket()); } catch (PartitionNotExistException e) { - // Metadata refresh confirmed that the queued lookup carries a deleted partition - // id. Complete it instead of repeatedly enqueueing the stale TableBucket; a - // primary key lookuper can then reroute by partition name. - lookup.future().completeExceptionally(e); + lookupQueue.releaseBatchCapacity(lookupBatch); + for (AbstractLookupQuery lookup : lookupBatch.lookups()) { + lookup.future().completeExceptionally(e); + } continue; } catch (Exception e) { - // if leader is not found, re-enqueue the lookup to send again. - LOG.warn("Failed to lookup the leader for {} when lookup", tb, e); - reEnqueueLookup(lookup); + LOG.warn( + "Failed to lookup the leader for {} when lookup", + lookupBatch.tableBucket(), + e); + for (AbstractLookupQuery lookup : lookupBatch.lookups()) { + reEnqueueLookup(lookup); + } + lookupQueue.releaseBatchCapacity(lookupBatch); continue; } - lookupBatchesByLeader - .computeIfAbsent(Tuple2.of(leader, lookup.lookupType()), k -> new ArrayList<>()) - .add(lookup); + + batchesByLeaderAndType + .computeIfAbsent( + Tuple2.of(destination, lookupBatch.lookupType()), + ignored -> new ArrayList<>()) + .add(lookupBatch); + } + return batchesByLeaderAndType; + } + + private void sendLookupBatches( + int destination, LookupType lookupType, List lookupBatches) { + Map, List> batchesByTableAndKind = new LinkedHashMap<>(); + for (LookupBatch lookupBatch : lookupBatches) { + batchesByTableAndKind + .computeIfAbsent( + Tuple2.of( + lookupBatch.tableBucket().getTableId(), + lookupBatch.historical()), + ignored -> new ArrayList<>()) + .add(lookupBatch); + } + + for (List batchesForTableAndKind : batchesByTableAndKind.values()) { + for (List requestBatches : packLookupBatches(batchesForTableAndKind)) { + List> lookups = new ArrayList<>(); + for (LookupBatch requestBatch : requestBatches) { + lookups.addAll(requestBatch.lookups()); + } + sendLookups( + destination, + lookupType, + lookups, + () -> { + for (LookupBatch requestBatch : requestBatches) { + lookupQueue.releaseBatchCapacity(requestBatch); + } + }); + } + } + } + + @VisibleForTesting + List> packLookupBatches(List lookupBatches) { + List> requestGroups = new ArrayList<>(); + List currentGroup = new ArrayList<>(); + int currentSize = 0; + for (LookupBatch lookupBatch : lookupBatches) { + if (!currentGroup.isEmpty() + && currentSize + lookupBatch.size() > lookupQueue.maxBatchSize()) { + requestGroups.add(currentGroup); + currentGroup = new ArrayList<>(); + currentSize = 0; + } + currentGroup.add(lookupBatch); + currentSize += lookupBatch.size(); + } + if (!currentGroup.isEmpty()) { + requestGroups.add(currentGroup); } - return lookupBatchesByLeader; + return requestGroups; } @VisibleForTesting - void sendLookups( - int destination, LookupType lookupType, List> lookupBatches) { + void sendLookups(int destination, LookupType lookupType, List> lookups) { + sendLookups(destination, lookupType, lookups, null); + } + + private void sendLookups( + int destination, + LookupType lookupType, + List> lookups, + @Nullable Runnable releaseOriginalBatchCapacity) { if (lookupType == LookupType.LOOKUP) { - sendLookupRequest(destination, lookupBatches, false); + sendLookupRequest(destination, lookups, false, releaseOriginalBatchCapacity); } else if (lookupType == LookupType.LOOKUP_WITH_INSERT_IF_NOT_EXISTS) { - sendLookupRequest(destination, lookupBatches, true); + sendLookupRequest(destination, lookups, true, releaseOriginalBatchCapacity); } else if (lookupType == LookupType.PREFIX_LOOKUP) { - sendPrefixLookupRequest(destination, lookupBatches); + sendPrefixLookupRequest(destination, lookups, releaseOriginalBatchCapacity); } else { throw new IllegalArgumentException("Unsupported lookup type: " + lookupType); } } private void sendLookupRequest( - int destination, List> lookups, boolean insertIfNotExists) { + int destination, + List> lookups, + boolean insertIfNotExists, + @Nullable Runnable releaseOriginalBatchCapacity) { // table id -> (bucket and original partition name -> lookups) - Map> lookupByTableId = new LinkedHashMap<>(); + Map> lookupByTableId = new LinkedHashMap<>(); for (AbstractLookupQuery abstractLookupQuery : lookups) { LookupQuery lookup = (LookupQuery) abstractLookupQuery; TableBucket tb = lookup.tableBucket(); @@ -221,12 +272,13 @@ private void sendLookupRequest( LookupBatchKey batchKey = new LookupBatchKey(tb, lookup.originalPartitionName()); lookupByTableId .computeIfAbsent(tableId, k -> new LinkedHashMap<>()) - .computeIfAbsent(batchKey, k -> new LookupBatch(batchKey)) + .computeIfAbsent(batchKey, k -> new LookupRequestBatch(batchKey)) .addLookup(lookup); } TabletServerGateway gateway = metadataUpdater.newTabletServerClientForNode(destination); if (gateway == null) { + List futureCompletions = new ArrayList<>(); lookupByTableId.forEach( (tableId, lookupsByBatchKey) -> handleLookupRequestException( @@ -235,15 +287,18 @@ private void sendLookupRequest( + destination + " is not found in metadata cache."), destination, - lookupsByBatchKey.values())); + lookupsByBatchKey.values(), + futureCompletions)); + runFutureCompletionsAfterReleasingOriginalBatchCapacity( + releaseOriginalBatchCapacity, futureCompletions); return; } lookupByTableId.forEach( (tableId, lookupsByBatchKey) -> { - List> lookupRequestGroups = + List> lookupRequestGroups = packLookupRequestGroups(lookupsByBatchKey.values()); - for (Map lookupsByBatchKeyInRequest : + for (Map lookupsByBatchKeyInRequest : lookupRequestGroups) { sendLookupRequestAndHandleResponse( destination, @@ -255,7 +310,8 @@ private void sendLookupRequest( acks, maxRequestTimeoutMs), tableId, - lookupsByBatchKeyInRequest); + lookupsByBatchKeyInRequest, + releaseOriginalBatchCapacity); } }); } @@ -267,11 +323,11 @@ private void sendLookupRequest( * multiple original partitions may target the same {@link TableBucket}; responses carry the * original partition name so they can be dispatched using {@link LookupBatchKey}. */ - private List> packLookupRequestGroups( - Collection lookupBatches) { - Map normalLookups = new LinkedHashMap<>(); - Map historicalLookups = new LinkedHashMap<>(); - for (LookupBatch lookupBatch : lookupBatches) { + private List> packLookupRequestGroups( + Collection lookupBatches) { + Map normalLookups = new LinkedHashMap<>(); + Map historicalLookups = new LinkedHashMap<>(); + for (LookupRequestBatch lookupBatch : lookupBatches) { if (lookupBatch.originalPartitionName() == null) { normalLookups.put(lookupBatch.lookupBatchKey(), lookupBatch); } else { @@ -279,7 +335,7 @@ private List> packLookupRequestGroups( } } - List> lookupRequestGroups = new ArrayList<>(2); + List> lookupRequestGroups = new ArrayList<>(2); if (!normalLookups.isEmpty()) { lookupRequestGroups.add(normalLookups); } @@ -290,7 +346,9 @@ private List> packLookupRequestGroups( } private void sendPrefixLookupRequest( - int destination, List> prefixLookups) { + int destination, + List> prefixLookups, + @Nullable Runnable releaseOriginalBatchCapacity) { // table id -> (bucket -> lookups) Map> lookupByTableId = new HashMap<>(); for (AbstractLookupQuery abstractLookupQuery : prefixLookups) { @@ -305,6 +363,7 @@ private void sendPrefixLookupRequest( TabletServerGateway gateway = metadataUpdater.newTabletServerClientForNode(destination); if (gateway == null) { + List futureCompletions = new ArrayList<>(); lookupByTableId.forEach( (tableId, lookupsByBucket) -> handlePrefixLookupException( @@ -313,7 +372,10 @@ private void sendPrefixLookupRequest( + destination + " is not found in metadata cache."), destination, - lookupsByBucket)); + lookupsByBucket, + futureCompletions)); + runFutureCompletionsAfterReleasingOriginalBatchCapacity( + releaseOriginalBatchCapacity, futureCompletions); return; } @@ -324,7 +386,8 @@ private void sendPrefixLookupRequest( gateway, makePrefixLookupRequest(tableId, prefixLookupBatch.values()), tableId, - prefixLookupBatch)); + prefixLookupBatch, + releaseOriginalBatchCapacity)); } private void sendLookupRequestAndHandleResponse( @@ -332,35 +395,69 @@ private void sendLookupRequestAndHandleResponse( TabletServerGateway gateway, LookupRequest lookupRequest, long tableId, - Map lookupsByBatchKey) { - // TODO: Normal and historical lookups share the in-flight permits for simplicity. See - // https://github.com/apache/fluss/issues/3766. + Map lookupsByBatchKey, + @Nullable Runnable releaseOriginalBatchCapacity) { + Set tableBuckets = + lookupsByBatchKey.keySet().stream() + .map(LookupBatchKey::tableBucket) + .collect(Collectors.toSet()); + boolean acquired = false; + try { - maxInFlightReuqestsSemaphore.acquire(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new FlussRuntimeException("interrupted:", e); + acquireInFlightRequest(tableBuckets); + acquired = true; + gateway.lookup(lookupRequest) + .whenComplete( + (lookupResponse, e) -> { + List futureCompletions = new ArrayList<>(); + try { + if (e != null) { + handleLookupRequestException( + e, + destination, + lookupsByBatchKey.values(), + futureCompletions); + } else { + try { + handleLookupResponse( + tableId, + destination, + lookupResponse, + lookupsByBatchKey, + futureCompletions); + } catch (Throwable t) { + handleLookupRequestException( + t, + destination, + lookupsByBatchKey.values(), + futureCompletions); + } + } + } finally { + try { + releaseInFlightRequest(tableBuckets); + } finally { + runFutureCompletionsAfterReleasingOriginalBatchCapacity( + releaseOriginalBatchCapacity, futureCompletions); + } + } + }); + } catch (Throwable t) { + List futureCompletions = new ArrayList<>(); + try { + handleLookupRequestException( + t, destination, lookupsByBatchKey.values(), futureCompletions); + } finally { + try { + if (acquired) { + releaseInFlightRequest(tableBuckets); + } + } finally { + runFutureCompletionsAfterReleasingOriginalBatchCapacity( + releaseOriginalBatchCapacity, futureCompletions); + } + } } - gateway.lookup(lookupRequest) - .thenAccept( - lookupResponse -> { - try { - handleLookupResponse( - tableId, destination, lookupResponse, lookupsByBatchKey); - } finally { - maxInFlightReuqestsSemaphore.release(); - } - }) - .exceptionally( - e -> { - try { - handleLookupRequestException( - e, destination, lookupsByBatchKey.values()); - return null; - } finally { - maxInFlightReuqestsSemaphore.release(); - } - }); } private void sendPrefixLookupRequestAndHandleResponse( @@ -368,42 +465,100 @@ private void sendPrefixLookupRequestAndHandleResponse( TabletServerGateway gateway, PrefixLookupRequest prefixLookupRequest, long tableId, - Map lookupsByBucket) { + Map lookupsByBucket, + @Nullable Runnable releaseOriginalBatchCapacity) { + Set tableBuckets = lookupsByBucket.keySet(); + boolean acquired = false; + + try { + acquireInFlightRequest(tableBuckets); + acquired = true; + gateway.prefixLookup(prefixLookupRequest) + .whenComplete( + (prefixLookupResponse, e) -> { + List futureCompletions = new ArrayList<>(); + try { + if (e != null) { + handlePrefixLookupException( + e, destination, lookupsByBucket, futureCompletions); + } else { + try { + handlePrefixLookupResponse( + tableId, + destination, + prefixLookupResponse, + lookupsByBucket, + futureCompletions); + } catch (Throwable t) { + handlePrefixLookupException( + t, + destination, + lookupsByBucket, + futureCompletions); + } + } + } finally { + try { + releaseInFlightRequest(tableBuckets); + } finally { + runFutureCompletionsAfterReleasingOriginalBatchCapacity( + releaseOriginalBatchCapacity, futureCompletions); + } + } + }); + } catch (Throwable t) { + List futureCompletions = new ArrayList<>(); + try { + handlePrefixLookupException(t, destination, lookupsByBucket, futureCompletions); + } finally { + try { + if (acquired) { + releaseInFlightRequest(tableBuckets); + } + } finally { + runFutureCompletionsAfterReleasingOriginalBatchCapacity( + releaseOriginalBatchCapacity, futureCompletions); + } + } + } + } + + private static void runFutureCompletionsAfterReleasingOriginalBatchCapacity( + @Nullable Runnable releaseOriginalBatchCapacity, List futureCompletions) { + try { + if (releaseOriginalBatchCapacity != null) { + releaseOriginalBatchCapacity.run(); + } + } finally { + // Complete user futures only after releasing original batch capacity, because + // CompletableFuture continuations run synchronously and may append new lookups. + for (Runnable futureCompletion : futureCompletions) { + futureCompletion.run(); + } + } + } + + private void acquireInFlightRequest(Set tableBuckets) { try { - maxInFlightReuqestsSemaphore.acquire(); + maxInFlightRequestsSemaphore.acquire(); + lookupQueue.addInFlightRequests(tableBuckets); } catch (InterruptedException e) { Thread.currentThread().interrupt(); - throw new FlussRuntimeException("interrupted:", e); + throw new FlussRuntimeException("Interrupted while sending lookup request.", e); } - gateway.prefixLookup(prefixLookupRequest) - .thenAccept( - prefixLookupResponse -> { - try { - handlePrefixLookupResponse( - tableId, - destination, - prefixLookupResponse, - lookupsByBucket); - } finally { - maxInFlightReuqestsSemaphore.release(); - } - }) - .exceptionally( - e -> { - try { - handlePrefixLookupException(e, destination, lookupsByBucket); - return null; - } finally { - maxInFlightReuqestsSemaphore.release(); - } - }); + } + + private void releaseInFlightRequest(Set tableBuckets) { + lookupQueue.completeInFlightRequests(tableBuckets); + maxInFlightRequestsSemaphore.release(); } private void handleLookupResponse( long tableId, int destination, LookupResponse lookupResponse, - Map lookupsByBatchKey) { + Map lookupsByBatchKey, + List futureCompletions) { for (PbLookupRespForBucket pbLookupRespForBucket : lookupResponse.getBucketsRespsList()) { TableBucket tableBucket = new TableBucket( @@ -418,10 +573,16 @@ private void handleLookupResponse( pbLookupRespForBucket.hasOriginalPartitionName() ? pbLookupRespForBucket.getOriginalPartitionName() : null); - LookupBatch lookupBatch = lookupsByBatchKey.get(lookupBatchKey); + LookupRequestBatch lookupBatch = lookupsByBatchKey.get(lookupBatchKey); if (pbLookupRespForBucket.hasErrorCode()) { ApiError error = ApiError.fromErrorMessage(pbLookupRespForBucket); - handleLookupError(tableBucket, destination, error, lookupBatch.lookups(), "lookup"); + handleLookupError( + tableBucket, + destination, + error, + lookupBatch.lookups(), + "lookup", + futureCompletions); } else { List byteValues = pbLookupRespForBucket.getValuesList().stream() @@ -434,7 +595,7 @@ private void handleLookupResponse( } }) .collect(Collectors.toList()); - lookupBatch.complete(byteValues); + futureCompletions.add(() -> lookupBatch.complete(byteValues)); } } } @@ -443,7 +604,8 @@ private void handlePrefixLookupResponse( long tableId, int destination, PrefixLookupResponse prefixLookupResponse, - Map prefixLookupsByBucket) { + Map prefixLookupsByBucket, + List futureCompletions) { for (PbPrefixLookupRespForBucket pbRespForBucket : prefixLookupResponse.getBucketsRespsList()) { TableBucket tableBucket = @@ -462,7 +624,8 @@ private void handlePrefixLookupResponse( destination, error, prefixLookupBatch.lookups(), - "prefix lookup"); + "prefix lookup", + futureCompletions); } else { List> result = new ArrayList<>(pbRespForBucket.getValueListsCount()); for (int i = 0; i < pbRespForBucket.getValueListsCount(); i++) { @@ -473,22 +636,33 @@ private void handlePrefixLookupResponse( } result.add(keyResult); } - prefixLookupBatch.complete(result); + futureCompletions.add(() -> prefixLookupBatch.complete(result)); } } } private void handleLookupRequestException( - Throwable t, int destination, Collection lookupBatches) { + Throwable t, + int destination, + Collection lookupBatches, + List futureCompletions) { ApiError error = ApiError.fromThrowable(t); - for (LookupBatch lookupBatch : lookupBatches) { + for (LookupRequestBatch lookupBatch : lookupBatches) { handleLookupError( - lookupBatch.tableBucket(), destination, error, lookupBatch.lookups(), "lookup"); + lookupBatch.tableBucket(), + destination, + error, + lookupBatch.lookups(), + "lookup", + futureCompletions); } } private void handlePrefixLookupException( - Throwable t, int destination, Map lookupsByBucket) { + Throwable t, + int destination, + Map lookupsByBucket, + List futureCompletions) { ApiError error = ApiError.fromThrowable(t); for (PrefixLookupBatch lookupBatch : lookupsByBucket.values()) { handleLookupError( @@ -496,7 +670,8 @@ private void handlePrefixLookupException( destination, error, lookupBatch.lookups(), - "prefix lookup"); + "prefix lookup", + futureCompletions); } } @@ -531,14 +706,15 @@ private boolean canRetry(AbstractLookupQuery lookup, Exception exception) { * @param tableBucket the table bucket * @param error the error from server response * @param lookups the list of lookups to handle - * @param lookupType the type of lookup ("" for regular lookup, "prefix " for prefix lookup) + * @param futureCompletions deferred successful or terminal future completions */ private void handleLookupError( TableBucket tableBucket, int destination, ApiError error, List> lookups, - String lookupType) { + String lookupType, + List futureCompletions) { ApiException exception = error.exception(); LOG.error( "Failed to {} from node {} for bucket {}", @@ -582,7 +758,7 @@ private void handleLookupError( lookupType, tableBucket, error.formatErrMsg()); - lookup.future().completeExceptionally(exception); + futureCompletions.add(() -> lookup.future().completeExceptionally(exception)); } } } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java index 3c7512945dc..ee99e73e3b7 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java @@ -21,7 +21,7 @@ import org.apache.fluss.client.admin.ClusterHealthStatus; import org.apache.fluss.client.admin.OffsetSpec; import org.apache.fluss.client.admin.ProducerOffsetsResult; -import org.apache.fluss.client.lookup.LookupBatch; +import org.apache.fluss.client.lookup.LookupRequestBatch; import org.apache.fluss.client.lookup.PrefixLookupBatch; import org.apache.fluss.client.metadata.AcquireKvSnapshotLeaseResult; import org.apache.fluss.client.metadata.ActiveKvSnapshots; @@ -208,7 +208,7 @@ public static PutKvRequest makePutKvRequest( public static LookupRequest makeLookupRequest( long tableId, - Collection lookupBatches, + Collection lookupBatches, boolean insertIfNotExists, short acks, int timeoutMs) { diff --git a/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupQueueTest.java b/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupQueueTest.java index 063d8b28cb5..8f19eee0a7c 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupQueueTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupQueueTest.java @@ -22,106 +22,359 @@ import org.junit.jupiter.api.Test; +import java.time.Duration; +import java.util.Collections; +import java.util.List; 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_MAX_INFLIGHT_REQUESTS_PER_BUCKET; 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.apache.fluss.testutils.common.CommonTestUtils.waitUntil; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Tests for {@link LookupQueue}. */ class LookupQueueTest { @Test - void testDrainMaxBatchSize() throws Exception { - Configuration conf = new Configuration(); - conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 10); - conf.setString(CLIENT_LOOKUP_BATCH_TIMEOUT.key(), "1ms"); - LookupQueue queue = new LookupQueue(conf); + void testFullBatchDrainsImmediatelyAndReleasesCapacity() throws Exception { + LookupQueue queue = createQueue(10, 5, "10s"); - // drain empty - assertThat(queue.drain()).hasSize(0); + // Creating this batch reserves five queue slots. Filling it makes it immediately drainable, + // without waiting for the long batch timeout. + appendLookups(queue, new TableBucket(1, 0), 5); - appendLookups(queue, 1); - assertThat(queue.drain()).hasSize(1); + List batches = queue.drain(); + assertThat(batches).singleElement().extracting(LookupBatch::size).isEqualTo(5); + assertThat(queue.remainingQueueSize()).isEqualTo(5); assertThat(queue.hasUnDrained()).isFalse(); - appendLookups(queue, 9); - assertThat(queue.drain()).hasSize(9); - assertThat(queue.hasUnDrained()).isFalse(); + // Queue capacity belongs to the batch until its RPC attempt finishes. + queue.releaseBatchCapacity(batches.get(0)); + assertThat(queue.remainingQueueSize()).isEqualTo(10); + } - appendLookups(queue, 10); - assertThat(queue.drain()).hasSize(10); - assertThat(queue.hasUnDrained()).isFalse(); + @Test + void testUnderfilledBatchWaitsForItsTimeout() throws Exception { + LookupQueue queue = createQueue(10, 5, "100ms"); + appendLookups(queue, new TableBucket(1, 0), 1); + + // The batch is not full, so drain() should wait for its batch timeout before returning it. + long startNanos = System.nanoTime(); + List batches = queue.drain(); + long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos); + + assertThat(elapsedMs).isGreaterThanOrEqualTo(50); + assertThat(batches).singleElement().extracting(LookupBatch::size).isEqualTo(1); + queue.releaseBatchCapacity(batches.get(0)); + } - appendLookups(queue, 20); - assertThat(queue.drain()).hasSize(10); + @Test + void testReadyBucketDoesNotWaitForAnotherBucket() throws Exception { + LookupQueue queue = createQueue(20, 2, "10s"); + TableBucket slowBucket = new TableBucket(1, 0); + TableBucket readyBucket = new TableBucket(1, 1); + + appendLookups(queue, slowBucket, 1); + appendLookups(queue, readyBucket, 2); + + // The ready bucket is full and must drain immediately; it must not wait for the underfilled + // batch belonging to another bucket. + List readyBatches = queue.drain(); + assertThat(readyBatches).singleElement(); + assertThat(readyBatches.get(0).tableBucket()).isEqualTo(readyBucket); assertThat(queue.hasUnDrained()).isTrue(); - assertThat(queue.drainAll()).hasSize(10); - assertThat(queue.hasUnDrained()).isFalse(); + queue.releaseBatchCapacity(readyBatches.get(0)); + + // Force-drain the remaining underfilled batch to clean up the test. + List remainingBatches = queue.drainAll(); + assertThat(remainingBatches).singleElement(); + assertThat(remainingBatches.get(0).tableBucket()).isEqualTo(slowBucket); + queue.releaseBatchCapacity(remainingBatches.get(0)); } @Test - void testAppendLookupBlocksWhenQueueIsFull() throws Exception { + void testCreatingNewBatchBlocksUntilCapacityIsReleased() throws Exception { + LookupQueue queue = createQueue(10, 5, "10s"); + // Capacity is reserved per batch. These two one-element batches reserve 5 slots each and + // therefore consume all 10 queue slots. + appendLookups(queue, new TableBucket(1, 0), 1); + appendLookups(queue, new TableBucket(1, 1), 1); + assertThat(queue.remainingQueueSize()).isZero(); + + // A third bucket needs a new batch, so its appender must wait for five slots to be + // released. + CompletableFuture appendFuture = + CompletableFuture.runAsync(() -> appendLookups(queue, new TableBucket(1, 2), 1)); + waitUntil( + () -> queue.waitingAppenderCount() == 1, + Duration.ofSeconds(5), + "lookup appender waiting for batch capacity"); + assertThat(appendFuture).isNotDone(); + + List pressuredBatches = queue.drain(); + assertThat(pressuredBatches).hasSize(2); + + // Releasing either drained batch returns five slots and unblocks the third appender. + queue.releaseBatchCapacity(pressuredBatches.get(0)); + appendFuture.get(1, TimeUnit.SECONDS); + assertThat(queue.queuedBatchCount()).isEqualTo(1); + + queue.releaseBatchCapacity(pressuredBatches.get(1)); + List lastBatch = queue.drainAll(); + assertThat(lastBatch).singleElement(); + queue.releaseBatchCapacity(lastBatch.get(0)); + assertThat(queue.remainingQueueSize()).isEqualTo(10); + } + + @Test + void testCompletingBatchWakesAllAppendersThatCanShareNewBatch() throws Exception { + LookupQueue queue = createQueue(2, 2, "10s"); + TableBucket firstBucket = new TableBucket(1, 0); + TableBucket sharedBucket = new TableBucket(1, 1); + + // The first full batch reserves all queue capacity until its RPC attempt completes. + appendLookups(queue, firstBucket, 2); + LookupBatch firstBatch = queue.drain().get(0); + + // Both appenders need a new batch for the same bucket, so both initially wait for capacity. + CompletableFuture firstAppend = + CompletableFuture.runAsync(() -> queue.appendLookup(newLookup(sharedBucket))); + CompletableFuture secondAppend = + CompletableFuture.runAsync(() -> queue.appendLookup(newLookup(sharedBucket))); + waitUntil( + () -> queue.waitingAppenderCount() == 2, + Duration.ofSeconds(5), + "two lookup appenders waiting for batch capacity"); + + // One completion provides capacity for one new batch. Both appenders must wake: one creates + // that batch and the other fills its remaining slot without reserving additional capacity. + queue.releaseBatchCapacity(firstBatch); + CompletableFuture.allOf(firstAppend, secondAppend).get(1, TimeUnit.SECONDS); + + List sharedBatches = queue.drain(); + assertThat(sharedBatches).singleElement(); + assertThat(sharedBatches.get(0).tableBucket()).isEqualTo(sharedBucket); + assertThat(sharedBatches.get(0).size()).isEqualTo(2); + queue.releaseBatchCapacity(sharedBatches.get(0)); + } + + @Test + void testRetryDoesNotBlockCallbackAndReservesCapacityAgain() throws Exception { + LookupQueue queue = createQueue(1, 1, "10s"); + LookupQuery lookup = newLookup(new TableBucket(1, 0)); + queue.appendLookup(lookup); + LookupBatch firstAttempt = queue.drain().get(0); + + // The RPC callback only places the failed lookup in the retry deque. It returns immediately + // even though the first attempt still owns the queue's only reserved slot. + queue.reEnqueue(lookup); + assertThat(queue.reEnqueuedLookupCount()).isEqualTo(1); + assertThat(queue.remainingQueueSize()).isZero(); + + // Once the first attempt releases its slot, the sender thread can reserve it for the retry. + queue.releaseBatchCapacity(firstAttempt); + LookupBatch retryAttempt = queue.drain().get(0); + assertThat(retryAttempt.lookups()).containsExactly(lookup); + assertThat(queue.remainingQueueSize()).isZero(); + + queue.releaseBatchCapacity(retryAttempt); + assertThat(queue.remainingQueueSize()).isEqualTo(1); + } + + @Test + void testSameBucketDifferentLookupKindsAreDrainedRoundRobin() throws Exception { Configuration conf = new Configuration(); - conf.set(CLIENT_LOOKUP_QUEUE_SIZE, 5); + conf.set(CLIENT_LOOKUP_QUEUE_SIZE, 4); + conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 1); + conf.set(CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET, 1); + conf.setString(CLIENT_LOOKUP_BATCH_TIMEOUT.key(), "10s"); LookupQueue queue = new LookupQueue(conf); + TableBucket tableBucket = new TableBucket(1, 0); + LookupQuery normalLookup1 = newLookup(tableBucket); + LookupQuery normalLookup2 = newLookup(tableBucket); + LookupQuery historicalLookup = + new LookupQuery( + DATA1_TABLE_PATH_PK, tableBucket, new byte[] {1}, false, "dt=20200101"); + PrefixLookupQuery prefixLookup = + new PrefixLookupQuery(DATA1_TABLE_PATH_PK, tableBucket, new byte[] {2}); + + queue.appendLookup(normalLookup1); + queue.appendLookup(normalLookup2); + queue.appendLookup(historicalLookup); + queue.appendLookup(prefixLookup); + + LookupBatch firstBatch = queue.drain().get(0); + assertThat(firstBatch.lookups()).containsExactly(normalLookup1); + queue.addInFlightRequests(Collections.singleton(tableBucket)); + queue.completeInFlightRequests(Collections.singleton(tableBucket)); + queue.releaseBatchCapacity(firstBatch); + + // The normal lookup key still has backlog, but it must not monopolize the bucket permit. + // The next drain should move to the historical lookup key for the same table bucket. + LookupBatch secondBatch = queue.drain().get(0); + assertThat(secondBatch.lookups()).containsExactly(historicalLookup); + queue.releaseBatchCapacity(secondBatch); - appendLookups(queue, 5); - assertThat(queue.getLookupQueue()).hasSize(5); + LookupBatch thirdBatch = queue.drain().get(0); + assertThat(thirdBatch.lookups()).containsExactly(prefixLookup); + queue.releaseBatchCapacity(thirdBatch); + + LookupBatch fourthBatch = queue.drain().get(0); + assertThat(fourthBatch.lookups()).containsExactly(normalLookup2); + queue.releaseBatchCapacity(fourthBatch); + } + + @Test + void testReadyRetryReservesReleasedCapacityBeforeWaitingAppenders() throws Exception { + LookupQueue queue = createQueue(1, 1, "10s"); + TableBucket retryBucket = new TableBucket(1, 0); + TableBucket newLookupBucket = new TableBucket(1, 1); + LookupQuery retryLookup = newLookup(retryBucket); + queue.appendLookup(retryLookup); + LookupBatch firstAttempt = queue.drain().get(0); - CompletableFuture future = - CompletableFuture.runAsync( - () -> { - appendLookups(queue, 1); // will be blocked. - }); + // The retry is ready but cannot reserve capacity until the first attempt releases its slot. + queue.reEnqueue(retryLookup); + assertThat(queue.remainingQueueSize()).isZero(); - // appendLookup should block and not complete immediately. - assertThat(future.isDone()).isFalse(); + CompletableFuture appendFuture = + CompletableFuture.runAsync(() -> queue.appendLookup(newLookup(newLookupBucket))); + waitUntil( + () -> queue.waitingAppenderCount() == 1, + Duration.ofSeconds(5), + "lookup appender waiting for batch capacity"); - Thread.sleep(100); - // Still blocked after 100ms. - assertThat(future.isDone()).isFalse(); + // Releasing the first attempt must reserve the slot for the ready retry before waking the + // waiting appender. Otherwise, sustained appends can repeatedly consume every released slot + // and starve retries forever. + queue.releaseBatchCapacity(firstAttempt); + assertThat(appendFuture).isNotDone(); + LookupBatch retryAttempt = queue.drain().get(0); + assertThat(retryAttempt.lookups()).containsExactly(retryLookup); + assertThat(queue.remainingQueueSize()).isZero(); - queue.drain(); - future.get(1, TimeUnit.SECONDS); - assertThat(future.isDone()).isTrue(); + queue.releaseBatchCapacity(retryAttempt); + appendFuture.get(1, TimeUnit.SECONDS); + LookupBatch newLookupAttempt = queue.drain().get(0); + assertThat(newLookupAttempt.tableBucket()).isEqualTo(newLookupBucket); + queue.releaseBatchCapacity(newLookupAttempt); } @Test - void testReEnqueueNotBlock() throws Exception { + void testPerBucketInFlightLimitSkipsOnlyLimitedBucket() throws Exception { Configuration conf = new Configuration(); - conf.set(CLIENT_LOOKUP_QUEUE_SIZE, 5); - conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 5); + conf.set(CLIENT_LOOKUP_QUEUE_SIZE, 10); + conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 1); + conf.set(CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET, 5); + conf.setString(CLIENT_LOOKUP_BATCH_TIMEOUT.key(), "10s"); LookupQueue queue = new LookupQueue(conf); + TableBucket limitedBucket = new TableBucket(1, 0); + TableBucket otherBucket = new TableBucket(1, 1); - appendLookups(queue, 5); - assertThat(queue.getLookupQueue()).hasSize(5); - assertThat(queue.getReEnqueuedLookupQueue()).hasSize(0); + // Pretend the limited bucket already has five outstanding RPCs, exhausting its permit + // count. + for (int i = 0; i < 5; i++) { + queue.addInFlightRequests(Collections.singleton(limitedBucket)); + } + appendLookups(queue, limitedBucket, 2); + appendLookups(queue, otherBucket, 1); - 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(); + // drain() skips only the saturated bucket and still returns work for the other bucket. + List readyBatches = queue.drain(); + assertThat(readyBatches).singleElement(); + assertThat(readyBatches.get(0).tableBucket()).isEqualTo(otherBucket); + queue.releaseBatchCapacity(readyBatches.get(0)); - assertThat(queue.drain()).hasSize(5); - // drain re-enqueued lookup first. - assertThat(queue.getReEnqueuedLookupQueue().isEmpty()).isTrue(); - assertThat(queue.getLookupQueue()).hasSize(1); - assertThat(queue.hasUnDrained()).isTrue(); + // Releasing one permit makes one batch from the previously saturated bucket drainable. + queue.completeInFlightRequests(Collections.singleton(limitedBucket)); + List newlyReady = queue.drain(); + assertThat(newlyReady).singleElement(); + assertThat(newlyReady.get(0).tableBucket()).isEqualTo(limitedBucket); + queue.releaseBatchCapacity(newlyReady.get(0)); + } - assertThat(queue.drain()).hasSize(1); - assertThat(queue.hasUnDrained()).isFalse(); + @Test + void testLookupKindsAndHistoricalLookupsUseSeparateBatches() throws Exception { + LookupQueue queue = createQueue(8, 2, "10s"); + TableBucket tableBucket = new TableBucket(1, 0); + + // All queries target the same physical bucket, but they form four batches: normal lookup, + // historical lookup, lookup-with-insert, and prefix lookup. + queue.appendLookup(newLookup(tableBucket)); + queue.appendLookup(newLookup(tableBucket)); + queue.appendLookup( + new LookupQuery( + DATA1_TABLE_PATH_PK, tableBucket, new byte[] {1}, false, "dt=20200101")); + queue.appendLookup( + new LookupQuery( + DATA1_TABLE_PATH_PK, tableBucket, new byte[] {2}, false, "dt=20200102")); + queue.appendLookup( + new LookupQuery(DATA1_TABLE_PATH_PK, tableBucket, new byte[] {3}, true, null)); + queue.appendLookup( + new LookupQuery(DATA1_TABLE_PATH_PK, tableBucket, new byte[] {4}, true, null)); + queue.appendLookup(new PrefixLookupQuery(DATA1_TABLE_PATH_PK, tableBucket, new byte[] {5})); + queue.appendLookup(new PrefixLookupQuery(DATA1_TABLE_PATH_PK, tableBucket, new byte[] {6})); + + // Lookup type and historical mode are part of the batching key, so incompatible operations + // must never be merged even when their table bucket is identical. + List batches = queue.drain(); + assertThat(batches).hasSize(4); + assertThat(batches) + .extracting(LookupBatch::lookupType) + .containsExactlyInAnyOrder( + LookupType.LOOKUP, + LookupType.LOOKUP, + LookupType.LOOKUP_WITH_INSERT_IF_NOT_EXISTS, + LookupType.PREFIX_LOOKUP); + + LookupBatch historicalBatch = null; + for (LookupBatch batch : batches) { + if (batch.historical()) { + historicalBatch = batch; + break; + } + } + assertThat(historicalBatch).isNotNull(); + assertThat(historicalBatch.lookups()) + .extracting(AbstractLookupQuery::originalPartitionName) + .containsExactly("dt=20200101", "dt=20200102"); + + for (LookupBatch batch : batches) { + queue.releaseBatchCapacity(batch); + } + } + + @Test + void testQueueSizeMustFitOneBatch() { + Configuration conf = new Configuration(); + conf.set(CLIENT_LOOKUP_QUEUE_SIZE, 4); + conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 5); + + assertThatThrownBy(() -> new LookupQueue(conf)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("must not be smaller"); + } + + private static LookupQueue createQueue(int queueSize, int batchSize, String timeout) { + Configuration conf = new Configuration(); + conf.set(CLIENT_LOOKUP_QUEUE_SIZE, queueSize); + conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, batchSize); + conf.setString(CLIENT_LOOKUP_BATCH_TIMEOUT.key(), timeout); + return new LookupQueue(conf); } - private static void appendLookups(LookupQueue queue, int count) { + private static void appendLookups(LookupQueue queue, TableBucket tableBucket, int count) { for (int i = 0; i < count; i++) { - queue.appendLookup( - new LookupQuery(DATA1_TABLE_PATH_PK, new TableBucket(1, 1), new byte[] {0})); + queue.appendLookup(newLookup(tableBucket)); } } + + private static LookupQuery newLookup(TableBucket tableBucket) { + return new LookupQuery(DATA1_TABLE_PATH_PK, tableBucket, new byte[] {0}); + } } diff --git a/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupSenderTest.java b/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupSenderTest.java index 03b62d52e65..b8788dee059 100644 --- a/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupSenderTest.java +++ b/fluss-client/src/test/java/org/apache/fluss/client/lookup/LookupSenderTest.java @@ -58,10 +58,14 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import static org.apache.fluss.record.TestData.DATA1_TABLE_ID_PK; import static org.apache.fluss.record.TestData.DATA1_TABLE_INFO_PK; import static org.apache.fluss.record.TestData.DATA1_TABLE_PATH_PK; +import static org.apache.fluss.record.TestData.DATA2_TABLE_ID; +import static org.apache.fluss.record.TestData.DATA2_TABLE_INFO; +import static org.apache.fluss.record.TestData.DATA2_TABLE_PATH; import static org.apache.fluss.testutils.common.CommonTestUtils.waitUntil; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -96,7 +100,7 @@ void setup() { .build(); Configuration conf = new Configuration(); - conf.set(ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE, 5); + conf.set(ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE, 20); conf.set(ConfigOptions.CLIENT_LOOKUP_MAX_BATCH_SIZE, 10); lookupQueue = new LookupQueue(conf); @@ -236,6 +240,119 @@ void testNormalLookupsKeepExistingBatching() throws Exception { assertThat(request.getBucketsReqAt(0).hasOriginalPartitionName()).isFalse(); } + @Test + void testReadyBucketBatchesArePackedUpToMaxBatchSize() { + // The sender may combine ready batches from different buckets into one RPC, but the total + // number of lookups in an RPC must not exceed maxBatchSize (10 in setup()). + LookupBatch batchWithFourLookups = createLookupBatch(new TableBucket(1, 0), 4); + LookupBatch batchWithSixLookups = createLookupBatch(new TableBucket(1, 1), 6); + LookupBatch batchWithOneLookup = createLookupBatch(new TableBucket(1, 2), 1); + + List> requestGroups = + lookupSender.packLookupBatches( + Arrays.asList( + batchWithFourLookups, batchWithSixLookups, batchWithOneLookup)); + + // The first two batches exactly fill one RPC. The third batch must be sent in another RPC. + assertThat(requestGroups).hasSize(2); + assertThat(requestGroups.get(0)).containsExactly(batchWithFourLookups, batchWithSixLookups); + assertThat(requestGroups.get(1)).containsExactly(batchWithOneLookup); + assertThat(totalLookupCount(requestGroups.get(0))).isEqualTo(10); + assertThat(totalLookupCount(requestGroups.get(1))).isEqualTo(1); + } + + @Test + void testPerBucketInFlightLimitDoesNotBlockOtherBucketOnSameServer() throws Exception { + Configuration conf = new Configuration(); + conf.set(ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE, 10); + conf.set(ConfigOptions.CLIENT_LOOKUP_MAX_BATCH_SIZE, 1); + conf.set(ConfigOptions.CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET, 5); + conf.setString(ConfigOptions.CLIENT_LOOKUP_BATCH_TIMEOUT.key(), "1ms"); + LookupQueue limitedQueue = new LookupQueue(conf); + + ConfigurableTestTabletServerGateway node1Gateway = + new ConfigurableTestTabletServerGateway(); + List node1Requests = Collections.synchronizedList(new ArrayList<>()); + List> node1Responses = + Collections.synchronizedList(new ArrayList<>()); + AtomicInteger otherBucketRequestCount = new AtomicInteger(); + + node1Gateway.setLookupHandler( + request -> { + if (request.getTableId() == DATA2_TABLE_ID) { + otherBucketRequestCount.incrementAndGet(); + return createSuccessResponse(request, bytes("other-bucket")); + } + node1Requests.add(request); + CompletableFuture response = new CompletableFuture<>(); + node1Responses.add(response); + return response; + }); + + Map tableInfos = new HashMap<>(); + tableInfos.put(DATA1_TABLE_PATH_PK, DATA1_TABLE_INFO_PK); + tableInfos.put(DATA2_TABLE_PATH, DATA2_TABLE_INFO); + TestingMetadataUpdater limitedMetadataUpdater = + TestingMetadataUpdater.builder(tableInfos) + .withTabletServerGateway(1, node1Gateway) + .build(); + LookupSender limitedSender = + new LookupSender( + limitedMetadataUpdater, limitedQueue, 128, MAX_RETRIES, (short) -1, 1000); + Thread limitedSenderThread = new Thread(limitedSender); + limitedSenderThread.start(); + + try { + List node1Queries = new ArrayList<>(); + for (int i = 0; i < 6; i++) { + LookupQuery query = + new LookupQuery(DATA1_TABLE_PATH_PK, TABLE_BUCKET, bytes("node-1-" + i)); + node1Queries.add(query); + limitedQueue.appendLookup(query); + } + + waitUntil( + () -> node1Requests.size() == 5, + Duration.ofSeconds(5), + "five requests for the capped bucket"); + assertThat(limitedQueue.hasUnDrained()).isTrue(); + + LookupQuery otherBucketQuery = + new LookupQuery( + DATA2_TABLE_PATH, + new TableBucket(DATA2_TABLE_ID, 0), + bytes("other-bucket")); + limitedQueue.appendLookup(otherBucketQuery); + + assertThat(otherBucketQuery.future().get(1, TimeUnit.SECONDS)) + .isEqualTo(bytes("other-bucket")); + assertThat(otherBucketRequestCount).hasValue(1); + assertThat(node1Requests).hasSize(5); + + node1Responses + .get(0) + .complete(createSuccessResponse(node1Requests.get(0), bytes("node-1")).join()); + waitUntil( + () -> node1Requests.size() == 6, + Duration.ofSeconds(5), + "sixth request after one bucket permit is released"); + + for (int i = 1; i < node1Responses.size(); i++) { + node1Responses + .get(i) + .complete( + createSuccessResponse(node1Requests.get(i), bytes("node-1")) + .join()); + } + for (LookupQuery query : node1Queries) { + assertThat(query.future().get(1, TimeUnit.SECONDS)).isEqualTo(bytes("node-1")); + } + } finally { + limitedSender.forceClose(); + limitedSenderThread.join(5000); + } + } + @Test void testSendLookupRequestWithNotLeaderOrFollowerException() { assertThat(metadataUpdater.getBucketLocation(tb1)) @@ -599,6 +716,98 @@ void testMultipleConcurrentLookupsWithRetries() throws Exception { .isGreaterThanOrEqualTo(2); // at least 1 failure + 1 success for the batch } + @Test + void testRunFutureCompletionsAfterReleasingOriginalBatchCapacity() throws Exception { + Configuration conf = new Configuration(); + // Keep exactly one batch slot so the first in-flight lookup keeps the queue full until its + // original batch capacity is released. + conf.set(ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE, 1); + conf.set(ConfigOptions.CLIENT_LOOKUP_MAX_BATCH_SIZE, 1); + conf.setString(ConfigOptions.CLIENT_LOOKUP_BATCH_TIMEOUT.key(), "1ms"); + LookupQueue singleSlotQueue = new LookupQueue(conf); + + ConfigurableTestTabletServerGateway singleSlotGateway = + new ConfigurableTestTabletServerGateway(); + CompletableFuture firstResponse = new CompletableFuture<>(); + AtomicReference firstRequest = new AtomicReference<>(); + AtomicReference continuationQuery = new AtomicReference<>(); + CountDownLatch continuationAppended = new CountDownLatch(1); + AtomicInteger requestCount = new AtomicInteger(); + singleSlotGateway.setLookupHandler( + request -> { + if (requestCount.incrementAndGet() == 1) { + firstRequest.set(request); + return firstResponse; + } + return createSuccessResponse(request, bytes("second-value")); + }); + + TestingMetadataUpdater singleSlotMetadataUpdater = + TestingMetadataUpdater.builder( + Collections.singletonMap(DATA1_TABLE_PATH_PK, DATA1_TABLE_INFO_PK)) + .withTabletServerGateway(1, singleSlotGateway) + .build(); + LookupSender singleSlotSender = + new LookupSender( + singleSlotMetadataUpdater, + singleSlotQueue, + MAX_INFLIGHT_REQUESTS, + MAX_RETRIES, + (short) -1, + 1000); + Thread singleSlotSenderThread = new Thread(singleSlotSender); + singleSlotSenderThread.start(); + + try { + LookupQuery query = + new LookupQuery(DATA1_TABLE_PATH_PK, TABLE_BUCKET, bytes("first-key")); + // thenRun is a synchronous continuation. It runs inside future.complete(...) + // and tries to append another lookup while the first lookup completion is still + // on the stack. + query.future() + .thenRun( + () -> { + LookupQuery nextQuery = + new LookupQuery( + DATA1_TABLE_PATH_PK, + TABLE_BUCKET, + bytes("second-key")); + continuationQuery.set(nextQuery); + singleSlotQueue.appendLookup(nextQuery); + continuationAppended.countDown(); + }); + singleSlotQueue.appendLookup(query); + + waitUntil( + () -> firstRequest.get() != null, + Duration.ofSeconds(5), + "first lookup request"); + // Complete the first response from a separate thread. If user future completion happens + // before the original batch capacity is released, the synchronous continuation above + // blocks on appendLookup and this thread never exits. + Thread completer = + new Thread( + () -> + firstResponse.complete( + createSuccessResponse( + firstRequest.get(), + bytes("first-value")) + .join())); + completer.start(); + + assertThat(query.future().get(5, TimeUnit.SECONDS)).isEqualTo(bytes("first-value")); + assertThat(continuationAppended.await(5, TimeUnit.SECONDS)).isTrue(); + completer.join(5000); + assertThat(completer.isAlive()).isFalse(); + assertThat(continuationQuery.get().future().get(5, TimeUnit.SECONDS)) + .isEqualTo(bytes("second-value")); + assertThat(requestCount.get()).isEqualTo(2); + } finally { + singleSlotSender.forceClose(); + singleSlotSenderThread.join(5000); + } + } + // Helper methods private CompletableFuture createPartitionNameEchoResponse( @@ -631,6 +840,26 @@ private CompletableFuture createPartitionNameEchoResponse( return CompletableFuture.completedFuture(response); } + private static LookupBatch createLookupBatch(TableBucket tableBucket, int lookupCount) { + LookupBatch lookupBatch = + new LookupBatch( + new LookupQuery(DATA1_TABLE_PATH_PK, tableBucket, new byte[] {0}), + System.nanoTime()); + for (int i = 1; i < lookupCount; i++) { + lookupBatch.addLookup( + new LookupQuery(DATA1_TABLE_PATH_PK, tableBucket, new byte[] {(byte) i})); + } + return lookupBatch; + } + + private static int totalLookupCount(List lookupBatches) { + int count = 0; + for (LookupBatch lookupBatch : lookupBatches) { + count += lookupBatch.size(); + } + return count; + } + private static byte[] bytes(String value) { return value.getBytes(StandardCharsets.UTF_8); } @@ -682,19 +911,18 @@ private void testException(Exception exception, boolean shouldRetry, int expecte private CompletableFuture createSuccessResponse( LookupRequest request, byte[] value) { LookupResponse response = new LookupResponse(); - PbLookupRespForBucket bucketResp = response.addBucketsResp(); - bucketResp.setBucketId(TABLE_BUCKET.getBucket()); - if (TABLE_BUCKET.getPartitionId() != null) { - bucketResp.setPartitionId(TABLE_BUCKET.getPartitionId()); - } - if (request.getBucketsReqAt(0).hasOriginalPartitionName()) { - bucketResp.setOriginalPartitionName( - request.getBucketsReqAt(0).getOriginalPartitionName()); - } - // Add value for each key in the request - int keyCount = request.getBucketsReqAt(0).getKeysCount(); - for (int i = 0; i < keyCount; i++) { - bucketResp.addValue().setValues(value); + for (PbLookupReqForBucket bucketRequest : request.getBucketsReqsList()) { + PbLookupRespForBucket bucketResp = response.addBucketsResp(); + bucketResp.setBucketId(bucketRequest.getBucketId()); + if (bucketRequest.hasPartitionId()) { + bucketResp.setPartitionId(bucketRequest.getPartitionId()); + } + if (bucketRequest.hasOriginalPartitionName()) { + bucketResp.setOriginalPartitionName(bucketRequest.getOriginalPartitionName()); + } + for (int i = 0; i < bucketRequest.getKeysCount(); i++) { + bucketResp.addValue().setValues(value); + } } return CompletableFuture.completedFuture(response); } diff --git a/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java b/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java index 455096713a7..421763c2b35 100644 --- a/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java +++ b/fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java @@ -1547,14 +1547,19 @@ public class ConfigOptions { key("client.lookup.queue-size") .intType() .defaultValue(25600) - .withDescription("The maximum number of pending lookup operations."); + .withDescription( + "The capacity of pending and in-flight lookup operations. Each active lookup batch " + + "reserves 'client.lookup.max-batch-size' slots until that request attempt " + + "succeeds or fails. This value must not be smaller than 'client.lookup.max-batch-size'."); public static final ConfigOption CLIENT_LOOKUP_MAX_BATCH_SIZE = key("client.lookup.max-batch-size") .intType() .defaultValue(128) .withDescription( - "The maximum batch size of merging lookup operations to one lookup request."); + "The maximum number of lookup operations merged into one lookup request. Lookup " + + "operations are accumulated independently per table bucket, and ready bucket " + + "batches may share one request without exceeding this limit."); public static final ConfigOption CLIENT_LOOKUP_MAX_INFLIGHT_SIZE = key("client.lookup.max-inflight-requests") @@ -1563,6 +1568,14 @@ public class ConfigOptions { .withDescription( "The maximum number of unacknowledged lookup requests for lookup operations."); + public static final ConfigOption CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET = + key("client.lookup.max-inflight-requests-per-bucket") + .intType() + .defaultValue(5) + .withDescription( + "The maximum number of unacknowledged lookup requests per bucket. " + + "When a bucket reaches this limit, lookup requests for other buckets can still be sent."); + public static final ConfigOption CLIENT_LOOKUP_BATCH_TIMEOUT = key("client.lookup.batch-timeout") .durationType()