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..00c047171f1 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 @@ -24,140 +24,364 @@ import javax.annotation.concurrent.ThreadSafe; +import java.util.ArrayDeque; import java.util.ArrayList; +import java.util.Deque; +import java.util.HashMap; +import java.util.HashSet; 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.Set; import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.ReentrantLock; + +import static org.apache.fluss.utils.Preconditions.checkArgument; +import static org.apache.fluss.utils.Preconditions.checkState; /** - * A queue that buffers the pending lookup operations and provides a list of {@link LookupQuery} - * when call method {@link #drain()}. + * A queue that buffers pending lookup operations by lookup queue key and drains globally bounded + * batches. + * + *

Lookups within a queue key preserve FIFO order. A drain consumes one key continuously before + * moving to the next key. If a key still has pending lookups after the global batch is full, it is + * moved to the tail so the next drain starts from another key. Before a drained batch is sent, its + * keys are counted as in-flight until the corresponding requests complete. */ @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(); + private final Condition appendCondition = stateLock.newCondition(); + private final Condition drainCondition = stateLock.newCondition(); + + private final Map>> lookupQueues; + private final Deque lookupOrder; + private final Deque> reEnqueuedLookups; + // Counts started send batches, including batches waiting to be submitted to the network. + private final Map inFlightRequestsByKey; + private final int queueSize; private final int maxBatchSize; + private final int maxInFlightRequestsPerKey; private final long batchTimeoutNanos; + private boolean closed; + private boolean forceClosed; + private int queuedSize; + 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.maxInFlightRequestsPerKey = + conf.get(ConfigOptions.CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET); this.batchTimeoutNanos = conf.get(ConfigOptions.CLIENT_LOOKUP_BATCH_TIMEOUT).toNanos(); - this.closed = false; + checkArgument(queueSize > 0, "Lookup queue size must be greater than 0."); + checkArgument(maxBatchSize > 0, "Lookup batch size must be greater than 0."); + checkArgument( + maxInFlightRequestsPerKey > 0, + "Maximum in-flight lookup requests per lookup queue key must be greater than 0."); + + this.lookupQueues = new HashMap<>(); + this.lookupOrder = new ArrayDeque<>(); + this.reEnqueuedLookups = new ArrayDeque<>(); + this.inFlightRequestsByKey = new HashMap<>(); } void appendLookup(AbstractLookupQuery lookup) { - if (closed) { - throw new IllegalStateException( - "Can not append lookup operation since the LookupQueue is closed."); + InterruptedException interruptedException = null; + stateLock.lock(); + try { + while (queuedSize >= queueSize && !closed) { + try { + appendCondition.await(); + } catch (InterruptedException e) { + interruptedException = e; + break; + } + } + + if (interruptedException == null) { + if (closed) { + throw new IllegalStateException( + "Can not append lookup operation since the LookupQueue is closed."); + } + + LookupQueueKey lookupQueueKey = LookupQueueKey.fromLookup(lookup); + Deque> lookupQueue = lookupQueues.get(lookupQueueKey); + if (lookupQueue == null) { + lookupQueue = new ArrayDeque<>(); + lookupQueues.put(lookupQueueKey, lookupQueue); + lookupOrder.addLast(lookupQueueKey); + } + lookupQueue.addLast(lookup); + queuedSize++; + drainCondition.signal(); + } + } finally { + stateLock.unlock(); } - try { - lookupQueue.put(lookup); - } catch (InterruptedException e) { - lookup.future().completeExceptionally(e); + if (interruptedException != null) { + Thread.currentThread().interrupt(); + lookup.future().completeExceptionally(interruptedException); } } + /** Re-enqueues a retry without blocking an RPC callback thread on regular queue capacity. */ void reEnqueue(AbstractLookupQuery lookup) { - if (closed) { - throw new IllegalStateException( - "Can not re-enqueue lookup operation since the LookupQueue is closed."); + stateLock.lock(); + try { + if (closed) { + throw new IllegalStateException( + "Can not re-enqueue lookup operation since the LookupQueue is closed."); + } + reEnqueuedLookups.addLast(lookup); + drainCondition.signal(); + } finally { + stateLock.unlock(); } + } + boolean hasUnDrained() { + stateLock.lock(); try { - reEnqueuedLookupQueue.put(lookup); - } catch (InterruptedException e) { - lookup.future().completeExceptionally(e); + return hasUnDrainedUnsafe(); + } finally { + stateLock.unlock(); } } - boolean hasUnDrained() { - return !lookupQueue.isEmpty() || !reEnqueuedLookupQueue.isEmpty(); + /** Drain a globally bounded batch of lookup operations. */ + List> drain() throws InterruptedException { + return drain(false); } - /** 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; + /** Drain all lookup operations without waiting for the batch timeout. */ + List> drainAll() throws InterruptedException { + return drain(true); + } + + void startInFlightRequests(Set lookupQueueKeys) { + stateLock.lock(); + try { + for (LookupQueueKey lookupQueueKey : lookupQueueKeys) { + inFlightRequestsByKey.merge(lookupQueueKey, 1, Integer::sum); } + } finally { + stateLock.unlock(); + } + } - 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 lookupQueueKeys) { + stateLock.lock(); + try { + boolean keyBecameSendable = false; + for (LookupQueueKey lookupQueueKey : lookupQueueKeys) { + int inFlightRequests = inFlightRequestsByKey.getOrDefault(lookupQueueKey, 0); + checkState( + inFlightRequests > 0, + "No in-flight lookup request exists for lookup queue key %s.", + lookupQueueKey); + if (inFlightRequests == maxInFlightRequestsPerKey) { + keyBecameSendable = true; } - long retryDelayMs = lookup.nextRetryTimeMs() - System.currentTimeMillis(); - if (retryDelayMs <= 0) { - lookupOperations.add(lookup); - count++; + if (inFlightRequests == 1) { + inFlightRequestsByKey.remove(lookupQueueKey); } else { - nextRetryDelayNanos = - Math.min( - nextRetryDelayNanos, - TimeUnit.MILLISECONDS.toNanos(retryDelayMs)); - reEnqueuedLookupQueue.add(lookup); + inFlightRequestsByKey.put(lookupQueueKey, inFlightRequests - 1); } - reEnqueuedToCheck--; } + if (keyBecameSendable) { + drainCondition.signal(); + } + } finally { + stateLock.unlock(); + } + } + + public void close() { + close(false); + } - long lookupWaitNanos = waitNanos; - if (count == 0 && nextRetryDelayNanos != Long.MAX_VALUE) { - lookupWaitNanos = Math.min(waitNanos, Math.max(1L, nextRetryDelayNanos)); + void forceClose() { + close(true); + } + + private void close(boolean forceClose) { + stateLock.lock(); + try { + closed = true; + forceClosed |= forceClose; + appendCondition.signalAll(); + drainCondition.signalAll(); + } finally { + stateLock.unlock(); + } + } + + private List> drain(boolean drainAll) throws InterruptedException { + final long startNanos = System.nanoTime(); + final int drainLimit = drainAll ? Integer.MAX_VALUE : maxBatchSize; + List> lookupOperations = new ArrayList<>(maxBatchSize); + Set drainedKeys = new HashSet<>(); + stateLock.lock(); + try { + while (lookupOperations.size() < drainLimit) { + if (forceClosed) { + lookupOperations.clear(); + return lookupOperations; + } + + long nowNanos = System.nanoTime(); + long nextRetryDelayNanos = + drainReEnqueuedLookups( + lookupOperations, + drainedKeys, + drainLimit, + drainAll, + System.currentTimeMillis()); + drainLookups(lookupOperations, drainedKeys, drainLimit); + + if (lookupOperations.size() >= drainLimit) { + return lookupOperations; + } + if (drainAll) { + if (!hasUnDrainedUnsafe()) { + return lookupOperations; + } + drainCondition.await(); + continue; + } + if (closed) { + return lookupOperations; + } + + long waitNanos = batchTimeoutNanos - (nowNanos - startNanos); + if (waitNanos <= 0) { + return lookupOperations; + } + if (nextRetryDelayNanos != Long.MAX_VALUE) { + waitNanos = Math.min(waitNanos, Math.max(1L, nextRetryDelayNanos)); + } + drainCondition.awaitNanos(waitNanos); } - AbstractLookupQuery lookup = lookupQueue.poll(lookupWaitNanos, TimeUnit.NANOSECONDS); - if (lookup == null) { - break; + return lookupOperations; + } finally { + stateLock.unlock(); + } + } + + private long drainReEnqueuedLookups( + List> lookupOperations, + Set drainedKeys, + int drainLimit, + boolean drainAll, + long nowMs) { + long nextRetryDelayNanos = Long.MAX_VALUE; + int retriesToCheck = reEnqueuedLookups.size(); + while (retriesToCheck > 0 && lookupOperations.size() < drainLimit) { + AbstractLookupQuery lookup = reEnqueuedLookups.removeFirst(); + long retryDelayMs = lookup.nextRetryTimeMs() - nowMs; + if (!drainAll && retryDelayMs > 0) { + nextRetryDelayNanos = + Math.min(nextRetryDelayNanos, TimeUnit.MILLISECONDS.toNanos(retryDelayMs)); + reEnqueuedLookups.addLast(lookup); + } else if (!tryDrainKeyUnsafe(LookupQueueKey.fromLookup(lookup), drainedKeys)) { + reEnqueuedLookups.addLast(lookup); + } else { + lookupOperations.add(lookup); } - lookupOperations.add(lookup); - count++; - int transferred = lookupQueue.drainTo(lookupOperations, maxBatchSize - count); - count += transferred; - if (count >= maxBatchSize) { - break; + retriesToCheck--; + } + return nextRetryDelayNanos; + } + + private void drainLookups( + List> lookupOperations, + Set drainedKeys, + int drainLimit) { + int keysToCheck = lookupOrder.size(); + int drainedLookups = 0; + while (keysToCheck > 0 && lookupOperations.size() < drainLimit) { + LookupQueueKey lookupQueueKey = lookupOrder.removeFirst(); + Deque> lookupQueue = lookupQueues.get(lookupQueueKey); + checkState( + lookupQueue != null && !lookupQueue.isEmpty(), + "Lookup queue key %s is active without pending lookups.", + lookupQueueKey); + + if (!tryDrainKeyUnsafe(lookupQueueKey, drainedKeys)) { + lookupOrder.addLast(lookupQueueKey); + keysToCheck--; + continue; } + + while (!lookupQueue.isEmpty() && lookupOperations.size() < drainLimit) { + lookupOperations.add(lookupQueue.removeFirst()); + queuedSize--; + drainedLookups++; + } + if (lookupQueue.isEmpty()) { + lookupQueues.remove(lookupQueueKey); + } else { + lookupOrder.addLast(lookupQueueKey); + } + keysToCheck--; + } + + if (drainedLookups > 0) { + appendCondition.signalAll(); } - return lookupOperations; } - /** 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 tryDrainKeyUnsafe( + LookupQueueKey lookupQueueKey, Set drainedKeys) { + if (drainedKeys.contains(lookupQueueKey)) { + return true; + } + if (!canSendMoreRequestsUnsafe(lookupQueueKey)) { + return false; + } + drainedKeys.add(lookupQueueKey); + return true; } - public void close() { - closed = true; + private boolean canSendMoreRequestsUnsafe(LookupQueueKey lookupQueueKey) { + return inFlightRequestsByKey.getOrDefault(lookupQueueKey, 0) < maxInFlightRequestsPerKey; + } + + private boolean hasUnDrainedUnsafe() { + return queuedSize > 0 || !reEnqueuedLookups.isEmpty(); } @VisibleForTesting - ArrayBlockingQueue> getLookupQueue() { - return lookupQueue; + int queuedSize() { + stateLock.lock(); + try { + return queuedSize; + } finally { + stateLock.unlock(); + } } @VisibleForTesting - BlockingQueue> getReEnqueuedLookupQueue() { - return reEnqueuedLookupQueue; + int reEnqueuedLookupCount() { + stateLock.lock(); + try { + return reEnqueuedLookups.size(); + } finally { + stateLock.unlock(); + } + } + + @VisibleForTesting + int inFlightRequestCount(LookupQueueKey lookupQueueKey) { + stateLock.lock(); + try { + return inFlightRequestsByKey.getOrDefault(lookupQueueKey, 0); + } finally { + stateLock.unlock(); + } } } diff --git a/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQueueKey.java b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQueueKey.java new file mode 100644 index 00000000000..16e642a3071 --- /dev/null +++ b/fluss-client/src/main/java/org/apache/fluss/client/lookup/LookupQueueKey.java @@ -0,0 +1,83 @@ +/* + * 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.metadata.TableBucket; + +import java.util.Objects; + +/** Identifies lookup operations that can share one queue and one RPC request. */ +final class LookupQueueKey { + private final TableBucket tableBucket; + private final LookupType lookupType; + private final boolean historical; + + private LookupQueueKey(TableBucket tableBucket, LookupType lookupType, boolean historical) { + this.tableBucket = tableBucket; + this.lookupType = lookupType; + this.historical = historical; + } + + static LookupQueueKey of(TableBucket tableBucket, LookupType lookupType, boolean historical) { + return new LookupQueueKey(tableBucket, lookupType, historical); + } + + static LookupQueueKey fromLookup(AbstractLookupQuery lookup) { + return of( + lookup.tableBucket(), lookup.lookupType(), lookup.originalPartitionName() != null); + } + + TableBucket tableBucket() { + return tableBucket; + } + + LookupType lookupType() { + return lookupType; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof LookupQueueKey)) { + return false; + } + LookupQueueKey that = (LookupQueueKey) o; + return historical == that.historical + && tableBucket.equals(that.tableBucket) + && lookupType == that.lookupType; + } + + @Override + public int hashCode() { + return Objects.hash(tableBucket, lookupType, historical); + } + + @Override + public String toString() { + return "LookupQueueKey{" + + "tableBucket=" + + tableBucket + + ", lookupType=" + + lookupType + + ", historical=" + + historical + + '}'; + } +} 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..a5840ee4987 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 @@ -51,11 +51,13 @@ import java.util.Collection; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Collectors; import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeLookupRequest; @@ -78,7 +80,7 @@ class LookupSender implements Runnable { private final LookupQueue lookupQueue; - private final Semaphore maxInFlightReuqestsSemaphore; + private final Semaphore maxInFlightRequestsSemaphore; private final int maxRetries; @@ -91,13 +93,13 @@ class LookupSender implements Runnable { LookupSender( MetadataUpdater metadataUpdater, LookupQueue lookupQueue, - int maxFlightRequests, + int maxInFlightRequests, int maxRetries, short acks, int maxRequestTimeoutMs) { this.metadataUpdater = metadataUpdater; this.lookupQueue = lookupQueue; - this.maxInFlightReuqestsSemaphore = new Semaphore(maxFlightRequests); + this.maxInFlightRequestsSemaphore = new Semaphore(maxInFlightRequests); this.maxRetries = maxRetries; this.running = true; this.acks = acks; @@ -123,7 +125,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) { @@ -139,13 +141,14 @@ public void run() { private void runOnce(boolean drainAll) throws Exception { List> lookups = drainAll ? lookupQueue.drainAll() : lookupQueue.drain(); + if (lookups.isEmpty() || forceClose) { + return; + } + sendLookups(lookups); } private void sendLookups(List> lookups) throws Exception { - if (lookups.isEmpty()) { - return; - } // group by to lookup batches Map, List>> lookupBatches = groupByLeaderAndType(lookups); @@ -165,33 +168,50 @@ private void sendLookups(List> lookups) throws Exception private Map, List>> groupByLeaderAndType( List> lookups) { + Map>> lookupsByQueueKey = new LinkedHashMap<>(); + for (AbstractLookupQuery lookup : lookups) { + lookupsByQueueKey + .computeIfAbsent(LookupQueueKey.fromLookup(lookup), key -> new ArrayList<>()) + .add(lookup); + } + // -> lookup batches Map, List>> lookupBatchesByLeader = new HashMap<>(); - for (AbstractLookupQuery lookup : lookups) { + for (Map.Entry>> entry : + lookupsByQueueKey.entrySet()) { + LookupQueueKey lookupQueueKey = entry.getKey(); + List> lookupsForKey = entry.getValue(); + AbstractLookupQuery representativeLookup = lookupsForKey.get(0); int leader; // 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); + leader = + metadataUpdater.leaderFor( + representativeLookup.tablePath(), lookupQueueKey.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); + // Metadata refresh confirmed that the queued lookups carry a deleted partition id. + // Complete them instead of repeatedly enqueueing the stale TableBucket; a primary + // key lookuper can then reroute by partition name. + lookupsForKey.forEach(lookup -> 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); + // if leader is not found, re-enqueue the lookups to send again. + LOG.warn( + "Failed to lookup the leader for {} when lookup", + lookupQueueKey.tableBucket(), + e); + lookupsForKey.forEach(this::reEnqueueLookup); continue; } lookupBatchesByLeader - .computeIfAbsent(Tuple2.of(leader, lookup.lookupType()), k -> new ArrayList<>()) - .add(lookup); + .computeIfAbsent( + Tuple2.of(leader, lookupQueueKey.lookupType()), + key -> new ArrayList<>()) + .addAll(lookupsForKey); } return lookupBatchesByLeader; } @@ -225,37 +245,48 @@ private void sendLookupRequest( .addLookup(lookup); } - TabletServerGateway gateway = metadataUpdater.newTabletServerClientForNode(destination); + TabletServerGateway gateway; + Throwable gatewayFailure; + try { + gateway = metadataUpdater.newTabletServerClientForNode(destination); + gatewayFailure = null; + } catch (Throwable t) { + gateway = null; + gatewayFailure = t; + } if (gateway == null) { + if (gatewayFailure == null) { + gatewayFailure = + new LeaderNotAvailableException( + "Server " + destination + " is not found in metadata cache."); + } + final Throwable requestFailure = gatewayFailure; lookupByTableId.forEach( (tableId, lookupsByBatchKey) -> handleLookupRequestException( - new LeaderNotAvailableException( - "Server " - + destination - + " is not found in metadata cache."), - destination, - lookupsByBatchKey.values())); + requestFailure, destination, lookupsByBatchKey.values())); return; } + final TabletServerGateway requestGateway = gateway; lookupByTableId.forEach( (tableId, lookupsByBatchKey) -> { List> lookupRequestGroups = packLookupRequestGroups(lookupsByBatchKey.values()); for (Map lookupsByBatchKeyInRequest : lookupRequestGroups) { + InFlightBatch inFlightBatch = + new InFlightBatch( + lookupQueue, + lookupQueueKeysFromLookupBatches( + lookupsByBatchKeyInRequest.values())); sendLookupRequestAndHandleResponse( destination, - gateway, - makeLookupRequest( - tableId, - lookupsByBatchKeyInRequest.values(), - insertIfNotExists, - acks, - maxRequestTimeoutMs), + requestGateway, tableId, - lookupsByBatchKeyInRequest); + lookupsByBatchKeyInRequest, + insertIfNotExists, + inFlightBatch); } }); } @@ -303,100 +334,130 @@ private void sendPrefixLookupRequest( .addLookup(prefixLookup); } - TabletServerGateway gateway = metadataUpdater.newTabletServerClientForNode(destination); + TabletServerGateway gateway; + Throwable gatewayFailure; + try { + gateway = metadataUpdater.newTabletServerClientForNode(destination); + gatewayFailure = null; + } catch (Throwable t) { + gateway = null; + gatewayFailure = t; + } if (gateway == null) { + if (gatewayFailure == null) { + gatewayFailure = + new LeaderNotAvailableException( + "Server " + destination + " is not found in metadata cache."); + } + final Throwable requestFailure = gatewayFailure; lookupByTableId.forEach( (tableId, lookupsByBucket) -> handlePrefixLookupException( - new LeaderNotAvailableException( - "Server " - + destination - + " is not found in metadata cache."), - destination, - lookupsByBucket)); + requestFailure, destination, lookupsByBucket)); return; } + final TabletServerGateway requestGateway = gateway; lookupByTableId.forEach( - (tableId, prefixLookupBatch) -> - sendPrefixLookupRequestAndHandleResponse( - destination, - gateway, - makePrefixLookupRequest(tableId, prefixLookupBatch.values()), - tableId, - prefixLookupBatch)); + (tableId, prefixLookupBatch) -> { + InFlightBatch inFlightBatch = + new InFlightBatch( + lookupQueue, + lookupQueueKeysFromPrefixLookupBatches( + prefixLookupBatch.values())); + sendPrefixLookupRequestAndHandleResponse( + destination, requestGateway, tableId, prefixLookupBatch, inFlightBatch); + }); } private void sendLookupRequestAndHandleResponse( int destination, 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, + boolean insertIfNotExists, + InFlightBatch inFlightBatch) { try { - maxInFlightReuqestsSemaphore.acquire(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new FlussRuntimeException("interrupted:", e); - } - 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(); - } - }); + acquireInFlightRequest(inFlightBatch); + LookupRequest lookupRequest = + makeLookupRequest( + tableId, + lookupsByBatchKey.values(), + insertIfNotExists, + acks, + maxRequestTimeoutMs); + gateway.lookup(lookupRequest) + .whenComplete( + (lookupResponse, e) -> { + try { + if (e != null) { + handleLookupRequestException( + e, destination, lookupsByBatchKey.values()); + } else { + try { + handleLookupResponse( + tableId, + destination, + lookupResponse, + lookupsByBatchKey); + } catch (Throwable t) { + handleLookupRequestException( + t, destination, lookupsByBatchKey.values()); + } + } + } finally { + releaseInFlightRequest(inFlightBatch); + } + }); + } catch (Throwable t) { + try { + handleLookupRequestException(t, destination, lookupsByBatchKey.values()); + } finally { + releaseInFlightRequest(inFlightBatch); + } + } } private void sendPrefixLookupRequestAndHandleResponse( int destination, TabletServerGateway gateway, - PrefixLookupRequest prefixLookupRequest, long tableId, - Map lookupsByBucket) { + Map lookupsByBucket, + InFlightBatch inFlightBatch) { try { - maxInFlightReuqestsSemaphore.acquire(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new FlussRuntimeException("interrupted:", 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(); - } - }); + acquireInFlightRequest(inFlightBatch); + PrefixLookupRequest prefixLookupRequest = + makePrefixLookupRequest(tableId, lookupsByBucket.values()); + gateway.prefixLookup(prefixLookupRequest) + .whenComplete( + (prefixLookupResponse, e) -> { + try { + if (e != null) { + handlePrefixLookupException( + e, destination, lookupsByBucket); + } else { + try { + handlePrefixLookupResponse( + tableId, + destination, + prefixLookupResponse, + lookupsByBucket); + } catch (Throwable t) { + handlePrefixLookupException( + t, destination, lookupsByBucket); + } + } + } finally { + releaseInFlightRequest(inFlightBatch); + } + }); + } catch (Throwable t) { + try { + handlePrefixLookupException(t, destination, lookupsByBucket); + } finally { + releaseInFlightRequest(inFlightBatch); + } + } } private void handleLookupResponse( @@ -500,6 +561,49 @@ private void handlePrefixLookupException( } } + private static Set lookupQueueKeysFromLookupBatches( + Collection lookupBatches) { + Set lookupQueueKeys = new HashSet<>(); + for (LookupBatch lookupBatch : lookupBatches) { + for (LookupQuery lookup : lookupBatch.lookups()) { + lookupQueueKeys.add(LookupQueueKey.fromLookup(lookup)); + } + } + return lookupQueueKeys; + } + + private static Set lookupQueueKeysFromPrefixLookupBatches( + Collection prefixLookupBatches) { + Set lookupQueueKeys = new HashSet<>(); + for (PrefixLookupBatch prefixLookupBatch : prefixLookupBatches) { + for (PrefixLookupQuery lookup : prefixLookupBatch.lookups()) { + lookupQueueKeys.add(LookupQueueKey.fromLookup(lookup)); + } + } + return lookupQueueKeys; + } + + private void acquireInFlightRequest(InFlightBatch inFlightBatch) { + try { + maxInFlightRequestsSemaphore.acquire(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new FlussRuntimeException("Interrupted while sending lookup request.", e); + } + try { + inFlightBatch.startInFlightRequests(); + } catch (Throwable t) { + maxInFlightRequestsSemaphore.release(); + throw t; + } + } + + private void releaseInFlightRequest(InFlightBatch inFlightBatch) { + if (inFlightBatch.requestCompleted()) { + maxInFlightRequestsSemaphore.release(); + } + } + private void reEnqueueLookup(AbstractLookupQuery lookup) { lookupQueue.reEnqueue(lookup); } @@ -589,7 +693,8 @@ private void handleLookupError( void forceClose() { forceClose = true; - initiateClose(); + lookupQueue.forceClose(); + running = false; } void initiateClose() { @@ -599,6 +704,36 @@ void initiateClose() { running = false; } + static class InFlightBatch { + private final LookupQueue lookupQueue; + private final Set lookupQueueKeys; + private final AtomicBoolean started = new AtomicBoolean(); + + InFlightBatch(LookupQueue lookupQueue, Set lookupQueueKeys) { + this.lookupQueue = lookupQueue; + this.lookupQueueKeys = lookupQueueKeys; + } + + void startInFlightRequests() { + if (started.compareAndSet(false, true)) { + try { + lookupQueue.startInFlightRequests(lookupQueueKeys); + } catch (Throwable t) { + started.set(false); + throw t; + } + } + } + + boolean requestCompleted() { + if (started.compareAndSet(true, false)) { + lookupQueue.completeInFlightRequests(lookupQueueKeys); + return true; + } + return false; + } + } + /** A helper class to hold table ids or table partitions. */ private static class TableOrPartitions { private final @Nullable Set tableIds; 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..addf3d42f94 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,13 +22,22 @@ import org.junit.jupiter.api.Test; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; 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; /** Tests for {@link LookupQueue}. */ @@ -41,8 +50,7 @@ void testDrainMaxBatchSize() throws Exception { conf.setString(CLIENT_LOOKUP_BATCH_TIMEOUT.key(), "1ms"); LookupQueue queue = new LookupQueue(conf); - // drain empty - assertThat(queue.drain()).hasSize(0); + assertThat(queue.drain()).isEmpty(); appendLookups(queue, 1); assertThat(queue.drain()).hasSize(1); @@ -56,72 +64,368 @@ void testDrainMaxBatchSize() throws Exception { assertThat(queue.drain()).hasSize(10); assertThat(queue.hasUnDrained()).isFalse(); - appendLookups(queue, 20); + appendLookups(queue, 30); assertThat(queue.drain()).hasSize(10); assertThat(queue.hasUnDrained()).isTrue(); - assertThat(queue.drainAll()).hasSize(10); + assertThat(queue.drainAll()).hasSize(20); assertThat(queue.hasUnDrained()).isFalse(); } @Test - void testAppendLookupBlocksWhenQueueIsFull() throws Exception { + void testDrainByLookupQueueKeyAndRotateStartKey() throws Exception { Configuration conf = new Configuration(); - conf.set(CLIENT_LOOKUP_QUEUE_SIZE, 5); + conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 4); + conf.setString(CLIENT_LOOKUP_BATCH_TIMEOUT.key(), "1ms"); LookupQueue queue = new LookupQueue(conf); + TableBucket bucketA = new TableBucket(1, 0); + TableBucket bucketB = new TableBucket(1, 1); + TableBucket bucketC = new TableBucket(1, 2); - appendLookups(queue, 5); - assertThat(queue.getLookupQueue()).hasSize(5); + appendLookup(queue, bucketA, "A0"); + appendLookup(queue, bucketB, "B0"); + appendLookup(queue, bucketA, "A1"); + appendLookup(queue, bucketC, "C0"); + appendLookup(queue, bucketA, "A2"); + appendLookup(queue, bucketB, "B1"); + appendLookup(queue, bucketA, "A3"); + appendLookup(queue, bucketC, "C1"); + appendLookup(queue, bucketA, "A4"); + appendLookup(queue, bucketB, "B2"); + appendLookup(queue, bucketA, "A5"); + appendLookup(queue, bucketC, "C2"); + + assertThat(keys(queue.drain())).containsExactly("A0", "A1", "A2", "A3"); + assertThat(keys(queue.drain())).containsExactly("B0", "B1", "B2", "C0"); + assertThat(keys(queue.drain())).containsExactly("A4", "A5", "C1", "C2"); + assertThat(queue.hasUnDrained()).isFalse(); + } + + @Test + void testRotateByLookupQueueKey() throws Exception { + Configuration conf = new Configuration(); + conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 1); + conf.setString(CLIENT_LOOKUP_BATCH_TIMEOUT.key(), "1ms"); + LookupQueue queue = new LookupQueue(conf); + TableBucket bucketA = new TableBucket(1, 0); + TableBucket bucketB = new TableBucket(1, 1); + + appendLookup(queue, bucketA, "A-N0"); + appendPrefixLookup(queue, bucketA, "A-P0"); + appendLookup(queue, bucketB, "B-N0"); + appendLookup(queue, bucketA, "A-N1"); + appendPrefixLookup(queue, bucketA, "A-P1"); + appendLookup(queue, bucketB, "B-N1"); - CompletableFuture future = + assertThat(keys(queue.drain())).containsExactly("A-N0"); + assertThat(keys(queue.drain())).containsExactly("A-P0"); + assertThat(keys(queue.drain())).containsExactly("B-N0"); + assertThat(keys(queue.drain())).containsExactly("A-N1"); + assertThat(keys(queue.drain())).containsExactly("A-P1"); + assertThat(keys(queue.drain())).containsExactly("B-N1"); + } + + @Test + void testAppendSameBucketWhileDrainWaitsForBatchTimeout() throws Exception { + Configuration conf = new Configuration(); + conf.set(CLIENT_LOOKUP_QUEUE_SIZE, 2); + conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 4); + conf.set(CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET, 1); + conf.set(CLIENT_LOOKUP_BATCH_TIMEOUT, Duration.ofSeconds(10)); + LookupQueue queue = new LookupQueue(conf); + TableBucket bucketA = new TableBucket(1, 0); + appendLookup(queue, bucketA, "A0"); + appendLookup(queue, bucketA, "A1"); + + CompletableFuture>> drainFuture = + CompletableFuture.supplyAsync( + () -> { + try { + return queue.drain(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new CompletionException(e); + } + }); + waitUntil( + () -> queue.queuedSize() == 0, + Duration.ofSeconds(1), + "the first two lookups to be moved into the sender batch"); + + CompletableFuture appendFuture = CompletableFuture.runAsync( () -> { - appendLookups(queue, 1); // will be blocked. + appendLookup(queue, bucketA, "A2"); + appendLookup(queue, bucketA, "A3"); }); - // appendLookup should block and not complete immediately. - assertThat(future.isDone()).isFalse(); + appendFuture.get(1, TimeUnit.SECONDS); + assertThat(keys(drainFuture.get(1, TimeUnit.SECONDS))) + .containsExactly("A0", "A1", "A2", "A3"); + } + + @Test + void testAppendLookupBlocksWhenQueueIsFull() 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.queuedSize()).isEqualTo(5); + CompletableFuture future = CompletableFuture.runAsync(() -> appendLookups(queue, 1)); + + assertThat(future.isDone()).isFalse(); Thread.sleep(100); - // Still blocked after 100ms. assertThat(future.isDone()).isFalse(); - queue.drain(); + assertThat(queue.drain()).hasSize(5); future.get(1, TimeUnit.SECONDS); - assertThat(future.isDone()).isTrue(); + assertThat(queue.queuedSize()).isEqualTo(1); } @Test - void testReEnqueueNotBlock() throws Exception { + void testReEnqueueDoesNotBlock() 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.queuedSize()).isEqualTo(5); + assertThat(queue.reEnqueuedLookupCount()).isZero(); - assertThat(queue.drain()).hasSize(5); - // drain re-enqueued lookup first. - assertThat(queue.getReEnqueuedLookupQueue().isEmpty()).isTrue(); - assertThat(queue.getLookupQueue()).hasSize(1); - assertThat(queue.hasUnDrained()).isTrue(); + queue.reEnqueue(lookup(new TableBucket(1, 1), "retry")); + assertThat(queue.queuedSize()).isEqualTo(5); + assertThat(queue.reEnqueuedLookupCount()).isEqualTo(1); - assertThat(queue.drain()).hasSize(1); + assertThat(keys(queue.drain())) + .containsExactly("retry", "lookup-0", "lookup-1", "lookup-2", "lookup-3"); + assertThat(queue.reEnqueuedLookupCount()).isZero(); + assertThat(queue.queuedSize()).isEqualTo(1); + + assertThat(keys(queue.drain())).containsExactly("lookup-4"); assertThat(queue.hasUnDrained()).isFalse(); } + @Test + void testDrainHonorsMaxInFlightBatchesPerLookupQueueKey() throws Exception { + Configuration conf = new Configuration(); + conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 1); + conf.set(CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET, 2); + conf.setString(CLIENT_LOOKUP_BATCH_TIMEOUT.key(), "1ms"); + LookupQueue queue = new LookupQueue(conf); + TableBucket bucketA = new TableBucket(1, 0); + LookupQueueKey normalKey = LookupQueueKey.of(bucketA, LookupType.LOOKUP, false); + + appendLookup(queue, bucketA, "A0"); + appendLookup(queue, bucketA, "A1"); + appendLookup(queue, bucketA, "A2"); + + assertThat(keys(queue.drain())).containsExactly("A0"); + LookupSender.InFlightBatch batch0 = startInFlightBatch(queue, normalKey); + assertThat(queue.inFlightRequestCount(normalKey)).isEqualTo(1); + assertThat(keys(queue.drain())).containsExactly("A1"); + startInFlightBatch(queue, normalKey); + assertThat(queue.inFlightRequestCount(normalKey)).isEqualTo(2); + assertThat(queue.drain()).isEmpty(); + assertThat(queue.queuedSize()).isEqualTo(1); + + batch0.requestCompleted(); + assertThat(batch0.requestCompleted()).isFalse(); + assertThat(queue.inFlightRequestCount(normalKey)).isEqualTo(1); + assertThat(keys(queue.drain())).containsExactly("A2"); + startInFlightBatch(queue, normalKey); + assertThat(queue.inFlightRequestCount(normalKey)).isEqualTo(2); + } + + @Test + void testDrainSkipsReservedKeyAndContinuesNextKey() throws Exception { + Configuration conf = new Configuration(); + conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 2); + conf.set(CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET, 1); + conf.setString(CLIENT_LOOKUP_BATCH_TIMEOUT.key(), "1ms"); + LookupQueue queue = new LookupQueue(conf); + TableBucket bucketA = new TableBucket(1, 0); + TableBucket bucketB = new TableBucket(1, 1); + LookupQueueKey bucketANormalKey = LookupQueueKey.of(bucketA, LookupType.LOOKUP, false); + + appendLookup(queue, bucketA, "A0"); + appendLookup(queue, bucketA, "A1"); + appendLookup(queue, bucketA, "A2"); + appendLookup(queue, bucketB, "B0"); + + assertThat(keys(queue.drain())).containsExactly("A0", "A1"); + LookupSender.InFlightBatch bucketABatch = startInFlightBatch(queue, bucketANormalKey); + assertThat(queue.inFlightRequestCount(bucketANormalKey)).isEqualTo(1); + assertThat(keys(queue.drain())).containsExactly("B0"); + assertThat(queue.queuedSize()).isEqualTo(1); + + bucketABatch.requestCompleted(); + assertThat(keys(queue.drain())).containsExactly("A2"); + } + + @Test + void testInFlightLimitIsIndependentForLookupQueueKeys() throws Exception { + Configuration conf = new Configuration(); + 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(), "1ms"); + LookupQueue queue = new LookupQueue(conf); + TableBucket bucket = new TableBucket(1, 0); + LookupQueueKey normalKey = LookupQueueKey.of(bucket, LookupType.LOOKUP, false); + LookupQueueKey historicalKey = LookupQueueKey.of(bucket, LookupType.LOOKUP, true); + LookupQueueKey prefixKey = LookupQueueKey.of(bucket, LookupType.PREFIX_LOOKUP, false); + LookupQueueKey insertKey = + LookupQueueKey.of(bucket, LookupType.LOOKUP_WITH_INSERT_IF_NOT_EXISTS, false); + + appendLookup(queue, bucket, "N0"); + appendLookup(queue, bucket, "N1"); + appendHistoricalLookup(queue, bucket, "H0", "dt=20200101"); + appendHistoricalLookup(queue, bucket, "H1", "dt=20200102"); + appendPrefixLookup(queue, bucket, "P0"); + appendInsertLookup(queue, bucket, "I0"); + + assertThat(keys(queue.drain())).containsExactly("N0"); + LookupSender.InFlightBatch normalBatch = startInFlightBatch(queue, normalKey); + assertThat(keys(queue.drain())).containsExactly("H0"); + LookupSender.InFlightBatch historicalBatch = startInFlightBatch(queue, historicalKey); + assertThat(keys(queue.drain())).containsExactly("P0"); + startInFlightBatch(queue, prefixKey); + assertThat(keys(queue.drain())).containsExactly("I0"); + startInFlightBatch(queue, insertKey); + assertThat(queue.drain()).isEmpty(); + assertThat(queue.queuedSize()).isEqualTo(2); + assertThat(queue.inFlightRequestCount(normalKey)).isEqualTo(1); + assertThat(queue.inFlightRequestCount(historicalKey)).isEqualTo(1); + assertThat(queue.inFlightRequestCount(prefixKey)).isEqualTo(1); + assertThat(queue.inFlightRequestCount(insertKey)).isEqualTo(1); + + normalBatch.requestCompleted(); + assertThat(keys(queue.drain())).containsExactly("N1"); + startInFlightBatch(queue, normalKey); + assertThat(queue.drain()).isEmpty(); + + historicalBatch.requestCompleted(); + assertThat(keys(queue.drain())).containsExactly("H1"); + } + + @Test + void testLookupQueueKeyClassification() { + TableBucket bucket = new TableBucket(1, 0); + LookupQueueKey normalKey = LookupQueueKey.fromLookup(lookup(bucket, "normal")); + LookupQueueKey historicalKey1 = + LookupQueueKey.fromLookup(historicalLookup(bucket, "historical-1", "dt=20200101")); + LookupQueueKey historicalKey2 = + LookupQueueKey.fromLookup(historicalLookup(bucket, "historical-2", "dt=20200102")); + LookupQueueKey prefixKey = LookupQueueKey.fromLookup(prefixLookup(bucket, "prefix")); + LookupQueueKey insertKey = LookupQueueKey.fromLookup(insertLookup(bucket, "insert")); + + assertThat(historicalKey1).isEqualTo(historicalKey2); + assertThat(Arrays.asList(normalKey, historicalKey1, prefixKey, insertKey)) + .doesNotHaveDuplicates(); + } + + @Test + void testForceCloseUnblocksDrainAllAtInFlightLimit() throws Exception { + Configuration conf = new Configuration(); + conf.set(CLIENT_LOOKUP_MAX_BATCH_SIZE, 1); + conf.set(CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET, 1); + LookupQueue queue = new LookupQueue(conf); + TableBucket bucketA = new TableBucket(1, 0); + + appendLookup(queue, bucketA, "A0"); + appendLookup(queue, bucketA, "A1"); + assertThat(keys(queue.drain())).containsExactly("A0"); + startInFlightBatch(queue, LookupQueueKey.of(bucketA, LookupType.LOOKUP, false)); + + CompletableFuture>> drainFuture = + CompletableFuture.supplyAsync( + () -> { + try { + return queue.drainAll(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new CompletionException(e); + } + }); + Thread.sleep(100); + assertThat(drainFuture).isNotDone(); + + queue.forceClose(); + assertThat(drainFuture.get(1, TimeUnit.SECONDS)).isEmpty(); + assertThat(queue.queuedSize()).isEqualTo(1); + } + + private static LookupSender.InFlightBatch startInFlightBatch( + LookupQueue queue, LookupQueueKey lookupQueueKey) { + LookupSender.InFlightBatch inFlightBatch = + new LookupSender.InFlightBatch(queue, Collections.singleton(lookupQueueKey)); + inFlightBatch.startInFlightRequests(); + return inFlightBatch; + } + private static void appendLookups(LookupQueue queue, int count) { + TableBucket tableBucket = new TableBucket(1, 0); for (int i = 0; i < count; i++) { - queue.appendLookup( - new LookupQuery(DATA1_TABLE_PATH_PK, new TableBucket(1, 1), new byte[] {0})); + appendLookup(queue, tableBucket, "lookup-" + i); } } + + private static void appendLookup(LookupQueue queue, TableBucket tableBucket, String lookupKey) { + queue.appendLookup(lookup(tableBucket, lookupKey)); + } + + private static void appendHistoricalLookup( + LookupQueue queue, + TableBucket tableBucket, + String lookupKey, + String originalPartitionName) { + queue.appendLookup(historicalLookup(tableBucket, lookupKey, originalPartitionName)); + } + + private static void appendPrefixLookup( + LookupQueue queue, TableBucket tableBucket, String lookupKey) { + queue.appendLookup(prefixLookup(tableBucket, lookupKey)); + } + + private static void appendInsertLookup( + LookupQueue queue, TableBucket tableBucket, String lookupKey) { + queue.appendLookup(insertLookup(tableBucket, lookupKey)); + } + + private static LookupQuery lookup(TableBucket tableBucket, String lookupKey) { + return new LookupQuery( + DATA1_TABLE_PATH_PK, tableBucket, lookupKey.getBytes(StandardCharsets.UTF_8)); + } + + private static LookupQuery historicalLookup( + TableBucket tableBucket, String lookupKey, String originalPartitionName) { + return new LookupQuery( + DATA1_TABLE_PATH_PK, + tableBucket, + lookupKey.getBytes(StandardCharsets.UTF_8), + false, + originalPartitionName); + } + + private static PrefixLookupQuery prefixLookup(TableBucket tableBucket, String lookupKey) { + return new PrefixLookupQuery( + DATA1_TABLE_PATH_PK, tableBucket, lookupKey.getBytes(StandardCharsets.UTF_8)); + } + + private static LookupQuery insertLookup(TableBucket tableBucket, String lookupKey) { + return new LookupQuery( + DATA1_TABLE_PATH_PK, + tableBucket, + lookupKey.getBytes(StandardCharsets.UTF_8), + true, + null); + } + + private static List keys(List> lookups) { + return lookups.stream() + .map(lookup -> new String(lookup.key(), StandardCharsets.UTF_8)) + .collect(Collectors.toList()); + } } 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..c7835e8cd4b 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 @@ -19,6 +19,7 @@ import org.apache.fluss.client.metadata.TestingMetadataUpdater; import org.apache.fluss.cluster.BucketLocation; +import org.apache.fluss.cluster.Cluster; import org.apache.fluss.config.ConfigOptions; import org.apache.fluss.config.Configuration; import org.apache.fluss.exception.HistoricalPartitionThrottledException; @@ -58,6 +59,7 @@ 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; @@ -93,11 +95,14 @@ void setup() { metadataUpdater = TestingMetadataUpdater.builder(tableInfos) .withTabletServerGateway(1, gateway) + .withTabletServerGateway(2, gateway) + .withTabletServerGateway(3, gateway) .build(); Configuration conf = new Configuration(); conf.set(ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE, 5); conf.set(ConfigOptions.CLIENT_LOOKUP_MAX_BATCH_SIZE, 10); + conf.set(ConfigOptions.CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET, 1); lookupQueue = new LookupQueue(conf); lookupSender = @@ -149,6 +154,9 @@ void testHistoricalLookupsBatchDifferentPartitionsForSameBucket() throws Excepti false, "dt=20200103"); + assertThat(LookupQueueKey.fromLookup(sameBucketQuery1)) + .isEqualTo(LookupQueueKey.fromLookup(sameBucketQuery2)); + lookupSender.sendLookups( 1, LookupType.LOOKUP, @@ -188,11 +196,14 @@ void testNormalAndHistoricalLookupsSplitRequests() throws Exception { LookupQuery historicalQuery = new LookupQuery( DATA1_TABLE_PATH_PK, - new TableBucket(DATA1_TABLE_ID_PK, 1), + TABLE_BUCKET, bytes("historical-key"), false, "dt=20200101"); + assertThat(LookupQueueKey.fromLookup(normalQuery)) + .isNotEqualTo(LookupQueueKey.fromLookup(historicalQuery)); + lookupSender.sendLookups(1, LookupType.LOOKUP, Arrays.asList(normalQuery, historicalQuery)); assertThat(normalQuery.future().get(5, TimeUnit.SECONDS)) @@ -208,7 +219,7 @@ void testNormalAndHistoricalLookupsSplitRequests() throws Exception { LookupRequest historicalRequest = receivedRequests.get(1); assertThat(historicalRequest.getBucketsReqsCount()).isEqualTo(1); - assertThat(historicalRequest.getBucketsReqAt(0).getBucketId()).isEqualTo(1); + assertThat(historicalRequest.getBucketsReqAt(0).getBucketId()).isEqualTo(0); assertThat(historicalRequest.getBucketsReqAt(0).getOriginalPartitionName()) .isEqualTo("dt=20200101"); } @@ -236,6 +247,252 @@ void testNormalLookupsKeepExistingBatching() throws Exception { assertThat(request.getBucketsReqAt(0).hasOriginalPartitionName()).isFalse(); } + @Test + void testInFlightLimitDoesNotBlockOtherLookupQueueKeys() throws Exception { + TableBucket bucketA = TABLE_BUCKET; + TableBucket bucketB = new TableBucket(DATA1_TABLE_ID_PK, 1); + AtomicInteger bucketARequestCount = new AtomicInteger(); + AtomicInteger bucketBRequestCount = new AtomicInteger(); + AtomicReference firstBucketARequest = new AtomicReference<>(); + CompletableFuture firstBucketAResponse = new CompletableFuture<>(); + gateway.setLookupHandler( + request -> { + int bucketId = request.getBucketsReqAt(0).getBucketId(); + if (bucketId == bucketA.getBucket()) { + if (bucketARequestCount.incrementAndGet() == 1) { + firstBucketARequest.set(request); + return firstBucketAResponse; + } + } else if (bucketId == bucketB.getBucket()) { + bucketBRequestCount.incrementAndGet(); + } + return createPartitionNameEchoResponse(request); + }); + + LookupQuery bucketAQuery1 = + new LookupQuery(DATA1_TABLE_PATH_PK, bucketA, bytes("bucket-a-1")); + lookupQueue.appendLookup(bucketAQuery1); + waitUntil( + () -> bucketARequestCount.get() == 1, + Duration.ofSeconds(5), + "the first bucket A request to be sent"); + + LookupQuery bucketAQuery2 = + new LookupQuery(DATA1_TABLE_PATH_PK, bucketA, bytes("bucket-a-2")); + LookupQuery bucketBQuery = new LookupQuery(DATA1_TABLE_PATH_PK, bucketB, bytes("bucket-b")); + lookupQueue.appendLookup(bucketAQuery2); + lookupQueue.appendLookup(bucketBQuery); + + assertThat(bucketBQuery.future().get(1, TimeUnit.SECONDS)) + .isEqualTo(responseValue("", "bucket-b")); + assertThat(bucketARequestCount.get()).isEqualTo(1); + assertThat(bucketBRequestCount.get()).isEqualTo(1); + assertThat(bucketAQuery2.future()).isNotDone(); + + firstBucketAResponse.complete( + createPartitionNameEchoResponse(firstBucketARequest.get()).join()); + + assertThat(bucketAQuery1.future().get(5, TimeUnit.SECONDS)) + .isEqualTo(responseValue("", "bucket-a-1")); + assertThat(bucketAQuery2.future().get(5, TimeUnit.SECONDS)) + .isEqualTo(responseValue("", "bucket-a-2")); + assertThat(bucketARequestCount.get()).isEqualTo(2); + } + + @Test + void testRpcCompletionReleasesAllLookupQueueKeys() throws Exception { + TableBucket bucketA = TABLE_BUCKET; + TableBucket bucketB = new TableBucket(DATA1_TABLE_ID_PK, 1); + Cluster cluster = metadataUpdater.getCluster(); + PhysicalTablePath physicalTablePath = PhysicalTablePath.of(DATA1_TABLE_PATH_PK); + Map> bucketLocationsByPath = + new HashMap<>(cluster.getBucketLocationsByPath()); + List bucketLocations = + new ArrayList<>(bucketLocationsByPath.get(physicalTablePath)); + bucketLocations.set( + 1, new BucketLocation(physicalTablePath, bucketB, 1, new int[] {1, 2, 3})); + bucketLocationsByPath.put(physicalTablePath, bucketLocations); + metadataUpdater.updateCluster( + new Cluster( + new HashMap<>(cluster.getAliveTabletServers()), + cluster.getCoordinatorServer(), + bucketLocationsByPath, + new HashMap<>(cluster.getTableIdByPath()), + new HashMap<>(cluster.getPartitionIdByPath()))); + + AtomicReference sentRequest = new AtomicReference<>(); + CompletableFuture response = new CompletableFuture<>(); + gateway.setLookupHandler( + request -> { + sentRequest.set(request); + return response; + }); + + Configuration conf = new Configuration(); + conf.set(ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE, 10); + conf.set(ConfigOptions.CLIENT_LOOKUP_MAX_BATCH_SIZE, 2); + conf.set(ConfigOptions.CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET, 1); + conf.set(ConfigOptions.CLIENT_LOOKUP_BATCH_TIMEOUT, Duration.ofSeconds(1)); + LookupQueue localQueue = new LookupQueue(conf); + LookupSender localSender = + new LookupSender( + metadataUpdater, + localQueue, + MAX_INFLIGHT_REQUESTS, + MAX_RETRIES, + (short) -1, + 1000); + + LookupQuery queryA = new LookupQuery(DATA1_TABLE_PATH_PK, bucketA, bytes("key-a")); + LookupQuery queryB = new LookupQuery(DATA1_TABLE_PATH_PK, bucketB, bytes("key-b")); + LookupQueueKey keyA = LookupQueueKey.fromLookup(queryA); + LookupQueueKey keyB = LookupQueueKey.fromLookup(queryB); + localQueue.appendLookup(queryA); + localQueue.appendLookup(queryB); + + Thread localSenderThread = new Thread(localSender); + localSenderThread.start(); + try { + waitUntil( + () -> sentRequest.get() != null, + Duration.ofSeconds(5), + "one RPC containing both lookup queue keys to be sent"); + assertThat(sentRequest.get().getBucketsReqsCount()).isEqualTo(2); + assertThat(localQueue.inFlightRequestCount(keyA)).isEqualTo(1); + assertThat(localQueue.inFlightRequestCount(keyB)).isEqualTo(1); + + response.complete(createPartitionNameEchoResponse(sentRequest.get()).join()); + assertThat(queryA.future().get(5, TimeUnit.SECONDS)) + .isEqualTo(responseValue("", "key-a")); + assertThat(queryB.future().get(5, TimeUnit.SECONDS)) + .isEqualTo(responseValue("", "key-b")); + waitUntil( + () -> + localQueue.inFlightRequestCount(keyA) == 0 + && localQueue.inFlightRequestCount(keyB) == 0, + Duration.ofSeconds(5), + "all lookup queue keys in the RPC to be released"); + } finally { + localSender.forceClose(); + localSenderThread.join(5000); + } + } + + @Test + void testSplitRequestsReleaseLookupQueueKeysIndependently() throws Exception { + AtomicReference normalRequest = new AtomicReference<>(); + AtomicReference historicalRequest = new AtomicReference<>(); + CompletableFuture normalResponse = new CompletableFuture<>(); + CompletableFuture historicalResponse = new CompletableFuture<>(); + AtomicInteger normalRequestCount = new AtomicInteger(); + AtomicInteger historicalRequestCount = new AtomicInteger(); + gateway.setLookupHandler( + request -> { + PbLookupReqForBucket bucketRequest = request.getBucketsReqAt(0); + String key = new String(bucketRequest.getKeyAt(0), StandardCharsets.UTF_8); + if (bucketRequest.hasOriginalPartitionName()) { + historicalRequestCount.incrementAndGet(); + if (key.equals("historical-key")) { + historicalRequest.set(request); + return historicalResponse; + } + } else { + normalRequestCount.incrementAndGet(); + if (key.equals("normal-key")) { + normalRequest.set(request); + return normalResponse; + } + } + return createPartitionNameEchoResponse(request); + }); + + Configuration conf = new Configuration(); + conf.set(ConfigOptions.CLIENT_LOOKUP_QUEUE_SIZE, 10); + conf.set(ConfigOptions.CLIENT_LOOKUP_MAX_BATCH_SIZE, 2); + conf.set(ConfigOptions.CLIENT_LOOKUP_MAX_INFLIGHT_REQUESTS_PER_BUCKET, 1); + conf.set(ConfigOptions.CLIENT_LOOKUP_BATCH_TIMEOUT, Duration.ofMillis(10)); + LookupQueue localQueue = new LookupQueue(conf); + LookupSender localSender = + new LookupSender( + metadataUpdater, + localQueue, + MAX_INFLIGHT_REQUESTS, + MAX_RETRIES, + (short) -1, + 1000); + + LookupQuery normalQuery = + new LookupQuery(DATA1_TABLE_PATH_PK, TABLE_BUCKET, bytes("normal-key")); + LookupQuery historicalQuery = + new LookupQuery( + DATA1_TABLE_PATH_PK, + TABLE_BUCKET, + bytes("historical-key"), + false, + "dt=20200101"); + LookupQueueKey normalKey = LookupQueueKey.fromLookup(normalQuery); + LookupQueueKey historicalKey = LookupQueueKey.fromLookup(historicalQuery); + localQueue.appendLookup(normalQuery); + localQueue.appendLookup(historicalQuery); + + Thread localSenderThread = new Thread(localSender); + localSenderThread.start(); + try { + waitUntil( + () -> normalRequest.get() != null && historicalRequest.get() != null, + Duration.ofSeconds(5), + "normal and historical requests from the first drain to be sent"); + assertThat(normalRequestCount.get()).isEqualTo(1); + assertThat(historicalRequestCount.get()).isEqualTo(1); + assertThat(localQueue.inFlightRequestCount(normalKey)).isEqualTo(1); + assertThat(localQueue.inFlightRequestCount(historicalKey)).isEqualTo(1); + + LookupQuery nextNormalQuery = + new LookupQuery(DATA1_TABLE_PATH_PK, TABLE_BUCKET, bytes("next-normal-key")); + LookupQuery nextHistoricalQuery = + new LookupQuery( + DATA1_TABLE_PATH_PK, + TABLE_BUCKET, + bytes("next-historical-key"), + false, + "dt=20200102"); + localQueue.appendLookup(nextNormalQuery); + localQueue.appendLookup(nextHistoricalQuery); + + normalResponse.complete(createPartitionNameEchoResponse(normalRequest.get()).join()); + assertThat(normalQuery.future().get(5, TimeUnit.SECONDS)) + .isEqualTo(responseValue("", "normal-key")); + assertThat(nextNormalQuery.future().get(5, TimeUnit.SECONDS)) + .isEqualTo(responseValue("", "next-normal-key")); + + waitUntil( + () -> localQueue.inFlightRequestCount(normalKey) == 0, + Duration.ofSeconds(5), + "the normal lookup queue key to be released"); + assertThat(historicalQuery.future()).isNotDone(); + assertThat(nextHistoricalQuery.future()).isNotDone(); + assertThat(normalRequestCount.get()).isEqualTo(2); + assertThat(historicalRequestCount.get()).isEqualTo(1); + assertThat(localQueue.queuedSize()).isEqualTo(1); + assertThat(localQueue.inFlightRequestCount(historicalKey)).isEqualTo(1); + + historicalResponse.complete( + createPartitionNameEchoResponse(historicalRequest.get()).join()); + assertThat(historicalQuery.future().get(5, TimeUnit.SECONDS)) + .isEqualTo(responseValue("dt=20200101", "historical-key")); + assertThat(nextHistoricalQuery.future().get(5, TimeUnit.SECONDS)) + .isEqualTo(responseValue("dt=20200102", "next-historical-key")); + waitUntil( + () -> localQueue.inFlightRequestCount(historicalKey) == 0, + Duration.ofSeconds(5), + "the historical lookup queue key to be released"); + assertThat(historicalRequestCount.get()).isEqualTo(2); + } finally { + localSender.forceClose(); + localSenderThread.join(5000); + } + } + @Test void testSendLookupRequestWithNotLeaderOrFollowerException() { assertThat(metadataUpdater.getBucketLocation(tb1)) 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..0aa8d976bc7 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 @@ -1554,7 +1554,8 @@ public class ConfigOptions { .intType() .defaultValue(128) .withDescription( - "The maximum batch size of merging lookup operations to one lookup request."); + "The maximum number of lookup operations drained into one send batch. " + + "The send batch may be split into multiple RPC requests by destination and lookup type."); public static final ConfigOption CLIENT_LOOKUP_MAX_INFLIGHT_SIZE = key("client.lookup.max-inflight-requests") @@ -1563,6 +1564,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 send batches per bucket that have been drained but have not completed, including batches waiting to be sent. " + + "A bucket at this limit is skipped while other buckets can still be drained."); + public static final ConfigOption CLIENT_LOOKUP_BATCH_TIMEOUT = key("client.lookup.batch-timeout") .durationType()