You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by ch...@apache.org on 2024/11/18 06:45:12 UTC
(kafka) branch trunk updated: KAFKA-17663 Add metadata caching in PartitionLeaderStrategy (#17367)
This is an automated email from the ASF dual-hosted git repository.
chia7712 pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new a592912ec94 KAFKA-17663 Add metadata caching in PartitionLeaderStrategy (#17367)
a592912ec94 is described below
commit a592912ec94b021a2bd7485abb9ebda02dc01766
Author: Andrew Schofield <as...@confluent.io>
AuthorDate: Mon Nov 18 06:45:06 2024 +0000
KAFKA-17663 Add metadata caching in PartitionLeaderStrategy (#17367)
Admin API operations have two phases: lookup and fulfilment. The lookup phase involves a METADATA request whose details depend upon the operation being performed.
For some operations, the METADATA request can be quite expensive to serve. For example, if the user calls Admin.listOffsets for 1000 topics, the METADATA request will include all 1000 topics and the response will contain the leader information for all of these topics. And then the actual fulfilment phase does the real work of the operation.
In cases where a long-running application is performing repeated admin operations which need the same metadata information about partition leadership, it is not necessary to send the METADATA request for every single admin operation.
This PR adds a cache of the mapping from topic-partition to leader id to the admin client. The cache doesn't need to be very sophisticated because the admin client will retry if the information becomes stale, and the cache can be updated as a result of the retry.
Reviewers: Chia-Ping Tsai <ch...@gmail.com>
---
.../kafka/clients/admin/KafkaAdminClient.java | 18 +-
.../admin/internals/AbortTransactionHandler.java | 8 +-
.../clients/admin/internals/AdminApiDriver.java | 8 +-
.../clients/admin/internals/AdminApiFuture.java | 20 +
.../admin/internals/DeleteRecordsHandler.java | 8 +-
.../admin/internals/DescribeProducersHandler.java | 7 +-
.../admin/internals/ListOffsetsHandler.java | 8 +-
.../admin/internals/PartitionLeaderStrategy.java | 92 ++++
.../PartitionLeaderStrategyIntegrationTest.java | 491 +++++++++++++++++++++
9 files changed, 638 insertions(+), 22 deletions(-)
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java b/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java
index 7e8f42091fc..1b6d3efc8e1 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java
@@ -58,6 +58,7 @@ import org.apache.kafka.clients.admin.internals.FenceProducersHandler;
import org.apache.kafka.clients.admin.internals.ListConsumerGroupOffsetsHandler;
import org.apache.kafka.clients.admin.internals.ListOffsetsHandler;
import org.apache.kafka.clients.admin.internals.ListTransactionsHandler;
+import org.apache.kafka.clients.admin.internals.PartitionLeaderStrategy;
import org.apache.kafka.clients.admin.internals.RemoveMembersFromConsumerGroupHandler;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.clients.consumer.internals.ConsumerProtocol;
@@ -408,6 +409,7 @@ public class KafkaAdminClient extends AdminClient {
private final ExponentialBackoff retryBackoff;
private final long rebootstrapTriggerMs;
private final MetadataRecoveryStrategy metadataRecoveryStrategy;
+ private final Map<TopicPartition, Integer> partitionLeaderCache;
private final AdminFetchMetricsManager adminFetchMetricsManager;
private final Optional<ClientTelemetryReporter> clientTelemetryReporter;
@@ -635,6 +637,7 @@ public class KafkaAdminClient extends AdminClient {
this.clientTelemetryReporter.ifPresent(reporters::add);
this.rebootstrapTriggerMs = config.getLong(AdminClientConfig.METADATA_RECOVERY_REBOOTSTRAP_TRIGGER_MS_CONFIG);
this.metadataRecoveryStrategy = MetadataRecoveryStrategy.forName(config.getString(AdminClientConfig.METADATA_RECOVERY_STRATEGY_CONFIG));
+ this.partitionLeaderCache = new HashMap<>();
this.adminFetchMetricsManager = new AdminFetchMetricsManager(metrics);
config.logUnused();
AppInfoParser.registerAppInfo(JMX_PREFIX, clientId, metrics, time.milliseconds());
@@ -3359,7 +3362,8 @@ public class KafkaAdminClient extends AdminClient {
@Override
public DeleteRecordsResult deleteRecords(final Map<TopicPartition, RecordsToDelete> recordsToDelete,
final DeleteRecordsOptions options) {
- SimpleAdminApiFuture<TopicPartition, DeletedRecords> future = DeleteRecordsHandler.newFuture(recordsToDelete.keySet());
+ PartitionLeaderStrategy.PartitionLeaderFuture<DeletedRecords> future =
+ DeleteRecordsHandler.newFuture(recordsToDelete.keySet(), partitionLeaderCache);
int timeoutMs = defaultApiTimeoutMs;
if (options.timeoutMs() != null) {
timeoutMs = options.timeoutMs();
@@ -4372,8 +4376,8 @@ public class KafkaAdminClient extends AdminClient {
@Override
public ListOffsetsResult listOffsets(Map<TopicPartition, OffsetSpec> topicPartitionOffsets,
ListOffsetsOptions options) {
- AdminApiFuture.SimpleAdminApiFuture<TopicPartition, ListOffsetsResultInfo> future =
- ListOffsetsHandler.newFuture(topicPartitionOffsets.keySet());
+ PartitionLeaderStrategy.PartitionLeaderFuture<ListOffsetsResultInfo> future =
+ ListOffsetsHandler.newFuture(topicPartitionOffsets.keySet(), partitionLeaderCache);
Map<TopicPartition, Long> offsetQueriesByPartition = topicPartitionOffsets.entrySet().stream()
.collect(Collectors.toMap(Map.Entry::getKey, e -> getOffsetFromSpec(e.getValue())));
ListOffsetsHandler handler = new ListOffsetsHandler(offsetQueriesByPartition, options, logContext, defaultApiTimeoutMs);
@@ -4918,8 +4922,8 @@ public class KafkaAdminClient extends AdminClient {
@Override
public DescribeProducersResult describeProducers(Collection<TopicPartition> topicPartitions, DescribeProducersOptions options) {
- AdminApiFuture.SimpleAdminApiFuture<TopicPartition, DescribeProducersResult.PartitionProducerState> future =
- DescribeProducersHandler.newFuture(topicPartitions);
+ PartitionLeaderStrategy.PartitionLeaderFuture<DescribeProducersResult.PartitionProducerState> future =
+ DescribeProducersHandler.newFuture(topicPartitions, partitionLeaderCache);
DescribeProducersHandler handler = new DescribeProducersHandler(options, logContext);
invokeDriver(handler, future, options.timeoutMs);
return new DescribeProducersResult(future.all());
@@ -4936,8 +4940,8 @@ public class KafkaAdminClient extends AdminClient {
@Override
public AbortTransactionResult abortTransaction(AbortTransactionSpec spec, AbortTransactionOptions options) {
- AdminApiFuture.SimpleAdminApiFuture<TopicPartition, Void> future =
- AbortTransactionHandler.newFuture(Collections.singleton(spec.topicPartition()));
+ PartitionLeaderStrategy.PartitionLeaderFuture<Void> future =
+ AbortTransactionHandler.newFuture(Collections.singleton(spec.topicPartition()), partitionLeaderCache);
AbortTransactionHandler handler = new AbortTransactionHandler(spec, logContext);
invokeDriver(handler, future, options.timeoutMs);
return new AbortTransactionResult(future.all());
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/AbortTransactionHandler.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/AbortTransactionHandler.java
index 0f5f4781080..f0b6d28be6b 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/internals/AbortTransactionHandler.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/AbortTransactionHandler.java
@@ -34,6 +34,7 @@ import org.apache.kafka.common.utils.LogContext;
import org.slf4j.Logger;
import java.util.List;
+import java.util.Map;
import java.util.Set;
import static java.util.Collections.singleton;
@@ -53,10 +54,11 @@ public class AbortTransactionHandler extends AdminApiHandler.Batched<TopicPartit
this.lookupStrategy = new PartitionLeaderStrategy(logContext);
}
- public static AdminApiFuture.SimpleAdminApiFuture<TopicPartition, Void> newFuture(
- Set<TopicPartition> topicPartitions
+ public static PartitionLeaderStrategy.PartitionLeaderFuture<Void> newFuture(
+ Set<TopicPartition> topicPartitions,
+ Map<TopicPartition, Integer> partitionLeaderCache
) {
- return AdminApiFuture.forKeys(topicPartitions);
+ return new PartitionLeaderStrategy.PartitionLeaderFuture<>(topicPartitions, partitionLeaderCache);
}
@Override
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminApiDriver.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminApiDriver.java
index 92e724e74f5..19e3f13ebef 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminApiDriver.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminApiDriver.java
@@ -110,7 +110,13 @@ public class AdminApiDriver<K, V> {
retryBackoffMaxMs,
CommonClientConfigs.RETRY_BACKOFF_JITTER);
this.log = logContext.logger(AdminApiDriver.class);
- retryLookup(future.lookupKeys());
+
+ // For any lookup keys for which we do not have cached information, we will need to look up
+ // metadata. For all cached keys, they can proceed straight to the fulfillment map.
+ // Note that the cache is only used on the initial calls, and any errors that result
+ // in additional lookups use the full set of lookup keys.
+ retryLookup(future.uncachedLookupKeys());
+ future.cachedKeyBrokerIdMapping().forEach((key, brokerId) -> fulfillmentMap.put(new FulfillmentScope(brokerId), key));
}
/**
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminApiFuture.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminApiFuture.java
index b0294d86166..322d116a3df 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminApiFuture.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/AdminApiFuture.java
@@ -19,6 +19,7 @@ package org.apache.kafka.clients.admin.internals;
import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.internals.KafkaFutureImpl;
+import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
@@ -37,6 +38,25 @@ public interface AdminApiFuture<K, V> {
*/
Set<K> lookupKeys();
+ /**
+ * The set of request keys that do not have cached key-broker id mappings. If there
+ * is no cached key mapping, this will be the same as the lookup keys.
+ * Can be empty, but only if the cached key mapping is not empty.
+ */
+ default Set<K> uncachedLookupKeys() {
+ return lookupKeys();
+ }
+
+ /**
+ * The cached key-broker id mapping. For lookup strategies that do not make use of a
+ * cache of metadata, this will be empty.
+ *
+ * @return mapping of keys to broker ids
+ */
+ default Map<K, Integer> cachedKeyBrokerIdMapping() {
+ return Collections.emptyMap();
+ }
+
/**
* Complete the futures associated with the given keys.
*
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandler.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandler.java
index 836f2bc2a59..4afef617cb2 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandler.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/DeleteRecordsHandler.java
@@ -18,7 +18,6 @@ package org.apache.kafka.clients.admin.internals;
import org.apache.kafka.clients.admin.DeletedRecords;
import org.apache.kafka.clients.admin.RecordsToDelete;
-import org.apache.kafka.clients.admin.internals.AdminApiFuture.SimpleAdminApiFuture;
import org.apache.kafka.clients.admin.internals.AdminApiHandler.Batched;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicPartition;
@@ -72,10 +71,11 @@ public final class DeleteRecordsHandler extends Batched<TopicPartition, DeletedR
return this.lookupStrategy;
}
- public static SimpleAdminApiFuture<TopicPartition, DeletedRecords> newFuture(
- Collection<TopicPartition> topicPartitions
+ public static PartitionLeaderStrategy.PartitionLeaderFuture<DeletedRecords> newFuture(
+ Collection<TopicPartition> topicPartitions,
+ Map<TopicPartition, Integer> partitionLeaderCache
) {
- return AdminApiFuture.forKeys(new HashSet<>(topicPartitions));
+ return new PartitionLeaderStrategy.PartitionLeaderFuture<>(new HashSet<>(topicPartitions), partitionLeaderCache);
}
@Override
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/DescribeProducersHandler.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/DescribeProducersHandler.java
index e4b203545bd..84338feb9e4 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/internals/DescribeProducersHandler.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/DescribeProducersHandler.java
@@ -66,10 +66,11 @@ public class DescribeProducersHandler extends AdminApiHandler.Batched<TopicParti
}
}
- public static AdminApiFuture.SimpleAdminApiFuture<TopicPartition, PartitionProducerState> newFuture(
- Collection<TopicPartition> topicPartitions
+ public static PartitionLeaderStrategy.PartitionLeaderFuture<PartitionProducerState> newFuture(
+ Collection<TopicPartition> topicPartitions,
+ Map<TopicPartition, Integer> partitionLeaderCache
) {
- return AdminApiFuture.forKeys(new HashSet<>(topicPartitions));
+ return new PartitionLeaderStrategy.PartitionLeaderFuture<>(new HashSet<>(topicPartitions), partitionLeaderCache);
}
@Override
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/ListOffsetsHandler.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/ListOffsetsHandler.java
index e36330a9924..f7c495d7fd8 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/internals/ListOffsetsHandler.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/ListOffsetsHandler.java
@@ -18,7 +18,6 @@ package org.apache.kafka.clients.admin.internals;
import org.apache.kafka.clients.admin.ListOffsetsOptions;
import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
-import org.apache.kafka.clients.admin.internals.AdminApiFuture.SimpleAdminApiFuture;
import org.apache.kafka.clients.admin.internals.AdminApiHandler.Batched;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicPartition;
@@ -217,9 +216,10 @@ public final class ListOffsetsHandler extends Batched<TopicPartition, ListOffset
}
}
- public static SimpleAdminApiFuture<TopicPartition, ListOffsetsResultInfo> newFuture(
- Collection<TopicPartition> topicPartitions
+ public static PartitionLeaderStrategy.PartitionLeaderFuture<ListOffsetsResultInfo> newFuture(
+ Collection<TopicPartition> topicPartitions,
+ Map<TopicPartition, Integer> partitionLeaderCache
) {
- return AdminApiFuture.forKeys(new HashSet<>(topicPartitions));
+ return new PartitionLeaderStrategy.PartitionLeaderFuture<>(new HashSet<>(topicPartitions), partitionLeaderCache);
}
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategy.java b/clients/src/main/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategy.java
index 9d52327b3c4..ff7dff2db8e 100644
--- a/clients/src/main/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategy.java
+++ b/clients/src/main/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategy.java
@@ -16,9 +16,11 @@
*/
package org.apache.kafka.clients.admin.internals;
+import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.InvalidTopicException;
import org.apache.kafka.common.errors.TopicAuthorizationException;
+import org.apache.kafka.common.internals.KafkaFutureImpl;
import org.apache.kafka.common.message.MetadataRequestData;
import org.apache.kafka.common.message.MetadataResponseData;
import org.apache.kafka.common.protocol.Errors;
@@ -31,9 +33,11 @@ import org.slf4j.Logger;
import java.util.Collections;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
+import java.util.stream.Collectors;
/**
* Base driver implementation for APIs which target partition leaders.
@@ -195,4 +199,92 @@ public class PartitionLeaderStrategy implements AdminApiLookupStrategy<TopicPart
return new LookupResult<>(failed, mapped);
}
+ /**
+ * This subclass of {@link AdminApiFuture} starts with a pre-fetched map for keys to broker ids which can be
+ * used to optimise the request. The map is kept up to date as metadata is fetching as this request is processed.
+ * This is useful for situations in which {@link PartitionLeaderStrategy} is used
+ * repeatedly, such as a sequence of identical calls to
+ * {@link org.apache.kafka.clients.admin.Admin#listOffsets(Map, org.apache.kafka.clients.admin.ListOffsetsOptions)}.
+ */
+ public static class PartitionLeaderFuture<V> implements AdminApiFuture<TopicPartition, V> {
+ private final Set<TopicPartition> requestKeys;
+ private final Map<TopicPartition, Integer> partitionLeaderCache;
+ private final Map<TopicPartition, KafkaFuture<V>> futures;
+
+ public PartitionLeaderFuture(Set<TopicPartition> requestKeys, Map<TopicPartition, Integer> partitionLeaderCache) {
+ this.requestKeys = requestKeys;
+ this.partitionLeaderCache = partitionLeaderCache;
+ this.futures = requestKeys.stream().collect(Collectors.toUnmodifiableMap(
+ Function.identity(),
+ k -> new KafkaFutureImpl<>()
+ ));
+ }
+
+ @Override
+ public Set<TopicPartition> lookupKeys() {
+ return futures.keySet();
+ }
+
+ @Override
+ public Set<TopicPartition> uncachedLookupKeys() {
+ Set<TopicPartition> keys = new HashSet<>();
+ requestKeys.forEach(tp -> {
+ if (!partitionLeaderCache.containsKey(tp)) {
+ keys.add(tp);
+ }
+ });
+ return keys;
+ }
+
+ @Override
+ public Map<TopicPartition, Integer> cachedKeyBrokerIdMapping() {
+ Map<TopicPartition, Integer> mapping = new HashMap<>();
+ requestKeys.forEach(tp -> {
+ Integer brokerId = partitionLeaderCache.get(tp);
+ if (brokerId != null) {
+ mapping.put(tp, brokerId);
+ }
+ });
+ return mapping;
+ }
+
+ public Map<TopicPartition, KafkaFuture<V>> all() {
+ return futures;
+ }
+
+ @Override
+ public void complete(Map<TopicPartition, V> values) {
+ values.forEach(this::complete);
+ }
+
+ private void complete(TopicPartition key, V value) {
+ futureOrThrow(key).complete(value);
+ }
+
+ @Override
+ public void completeLookup(Map<TopicPartition, Integer> brokerIdMapping) {
+ partitionLeaderCache.putAll(brokerIdMapping);
+ }
+
+ @Override
+ public void completeExceptionally(Map<TopicPartition, Throwable> errors) {
+ errors.forEach(this::completeExceptionally);
+ }
+
+ private void completeExceptionally(TopicPartition key, Throwable t) {
+ partitionLeaderCache.remove(key);
+ futureOrThrow(key).completeExceptionally(t);
+ }
+
+ private KafkaFutureImpl<V> futureOrThrow(TopicPartition key) {
+ // The below typecast is safe because we initialise futures using only KafkaFutureImpl.
+ KafkaFutureImpl<V> future = (KafkaFutureImpl<V>) futures.get(key);
+ if (future == null) {
+ throw new IllegalArgumentException("Attempt to complete future for " + key +
+ ", which was not requested");
+ } else {
+ return future;
+ }
+ }
+ }
}
diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategyIntegrationTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategyIntegrationTest.java
new file mode 100644
index 00000000000..4e03ae7d952
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/admin/internals/PartitionLeaderStrategyIntegrationTest.java
@@ -0,0 +1,491 @@
+/*
+ * 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.kafka.clients.admin.internals;
+
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.DisconnectException;
+import org.apache.kafka.common.errors.LeaderNotAvailableException;
+import org.apache.kafka.common.errors.NotLeaderOrFollowerException;
+import org.apache.kafka.common.errors.RetriableException;
+import org.apache.kafka.common.errors.UnknownServerException;
+import org.apache.kafka.common.message.ListOffsetsResponseData;
+import org.apache.kafka.common.message.MetadataRequestData;
+import org.apache.kafka.common.message.MetadataResponseData;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.AbstractRequest;
+import org.apache.kafka.common.requests.AbstractResponse;
+import org.apache.kafka.common.requests.ListOffsetsResponse;
+import org.apache.kafka.common.requests.MetadataRequest;
+import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.test.TestUtils;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.OptionalInt;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class PartitionLeaderStrategyIntegrationTest {
+ private static final long TIMEOUT_MS = 5000;
+ private static final long RETRY_BACKOFF_MS = 100;
+
+ private static final Node NODE_1 = new Node(1, "host1", 9092);
+ private static final Node NODE_2 = new Node(2, "host2", 9092);
+
+ private final LogContext logContext = new LogContext();
+ private final MockTime time = new MockTime();
+
+ private AdminApiDriver<TopicPartition, Void> buildDriver(
+ PartitionLeaderStrategy.PartitionLeaderFuture<Void> result
+ ) {
+ return new AdminApiDriver<>(
+ new MockApiHandler(),
+ result,
+ time.milliseconds() + TIMEOUT_MS,
+ RETRY_BACKOFF_MS,
+ RETRY_BACKOFF_MS,
+ logContext
+ );
+ }
+
+ @Test
+ public void testCachingRepeatedRequest() {
+ Map<TopicPartition, Integer> partitionLeaderCache = new HashMap<>();
+
+ TopicPartition tp0 = new TopicPartition("T", 0);
+ TopicPartition tp1 = new TopicPartition("T", 1);
+ Set<TopicPartition> requestKeys = Set.of(tp0, tp1);
+
+ // First, the lookup stage needs to obtain leadership data because the cache is empty
+ PartitionLeaderStrategy.PartitionLeaderFuture<Void> result =
+ new PartitionLeaderStrategy.PartitionLeaderFuture<>(requestKeys, partitionLeaderCache);
+ AdminApiDriver<TopicPartition, Void> driver = buildDriver(result);
+
+ List<AdminApiDriver.RequestSpec<TopicPartition>> requestSpecs = driver.poll();
+ assertEquals(1, requestSpecs.size());
+
+ assertEquals(OptionalInt.empty(), requestSpecs.get(0).scope.destinationBrokerId());
+ assertEquals(requestKeys, requestSpecs.get(0).keys);
+
+ // The cache will be populated using the leader information from this metadata response
+ Map<TopicPartition, Integer> leaders = Map.of(tp0, 1, tp1, 2);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), metadataResponseWithPartitionLeaders(leaders), Node.noNode());
+ assertFalse(result.all().get(tp0).isDone());
+ assertFalse(result.all().get(tp1).isDone());
+
+ assertEquals(1, partitionLeaderCache.get(tp0));
+ assertEquals(2, partitionLeaderCache.get(tp1));
+
+ // Second, the fulfillment stage makes the actual requests
+ requestSpecs = driver.poll();
+ assertEquals(2, requestSpecs.size());
+
+ assertEquals(OptionalInt.of(1), requestSpecs.get(0).scope.destinationBrokerId());
+ assertEquals(OptionalInt.of(2), requestSpecs.get(1).scope.destinationBrokerId());
+
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), listOffsetsResponseSuccess(requestSpecs.get(0).keys), NODE_1);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(1), listOffsetsResponseSuccess(requestSpecs.get(1).keys), NODE_2);
+ assertTrue(result.all().get(tp0).isDone());
+ assertTrue(result.all().get(tp1).isDone());
+
+ // On the second request, the partition leader cache already contains all the leadership
+ // data so the request goes straight to the fulfillment stage
+ result = new PartitionLeaderStrategy.PartitionLeaderFuture<>(requestKeys, partitionLeaderCache);
+ driver = buildDriver(result);
+
+ requestSpecs = driver.poll();
+ assertEquals(2, requestSpecs.size());
+
+ // We can tell this is the fulfillment stage by the destination broker id being set
+ assertEquals(OptionalInt.of(1), requestSpecs.get(0).scope.destinationBrokerId());
+ assertEquals(OptionalInt.of(2), requestSpecs.get(1).scope.destinationBrokerId());
+
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), listOffsetsResponseSuccess(requestSpecs.get(0).keys), NODE_1);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(1), listOffsetsResponseSuccess(requestSpecs.get(1).keys), NODE_2);
+ assertTrue(result.all().get(tp0).isDone());
+ assertTrue(result.all().get(tp1).isDone());
+ }
+
+ @Test
+ public void testCachingOverlappingRequests() {
+ // This test uses several requests to exercise the caching in various ways:
+ // 1) for T-0 and T-1 (initially the cache is empty)
+ // 2) for T-1 and T-2 (leadership data for T-1 should be cached from previous request)
+ // 3) for T-0, T-1 and T-2 (all leadership data should be cached already)
+ // 4) for T-0, T-1, T-2 and T-3 (just T-3 needs to be looked up)
+ Map<TopicPartition, Integer> partitionLeaderCache = new HashMap<>();
+
+ TopicPartition tp0 = new TopicPartition("T", 0);
+ TopicPartition tp1 = new TopicPartition("T", 1);
+ TopicPartition tp2 = new TopicPartition("T", 2);
+ TopicPartition tp3 = new TopicPartition("T", 3);
+
+ //
+ // Request 1 - T-0 and T-1
+ //
+ Set<TopicPartition> requestKeys = Set.of(tp0, tp1);
+
+ // First, the lookup stage needs to obtain leadership data because the cache is empty
+ PartitionLeaderStrategy.PartitionLeaderFuture<Void> result =
+ new PartitionLeaderStrategy.PartitionLeaderFuture<>(requestKeys, partitionLeaderCache);
+ AdminApiDriver<TopicPartition, Void> driver = buildDriver(result);
+
+ List<AdminApiDriver.RequestSpec<TopicPartition>> requestSpecs = driver.poll();
+ assertEquals(1, requestSpecs.size());
+
+ assertEquals(OptionalInt.empty(), requestSpecs.get(0).scope.destinationBrokerId());
+ assertEquals(requestKeys, requestSpecs.get(0).keys);
+
+ // The cache will be populated using the leader information from this metadata response
+ Map<TopicPartition, Integer> leaders = Map.of(tp0, 1, tp1, 2);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), metadataResponseWithPartitionLeaders(leaders), Node.noNode());
+ assertFalse(result.all().get(tp0).isDone());
+ assertFalse(result.all().get(tp1).isDone());
+
+ assertEquals(1, partitionLeaderCache.get(tp0));
+ assertEquals(2, partitionLeaderCache.get(tp1));
+
+ // Second, the fulfillment stage makes the actual requests
+ requestSpecs = driver.poll();
+ assertEquals(2, requestSpecs.size());
+
+ assertEquals(OptionalInt.of(1), requestSpecs.get(0).scope.destinationBrokerId());
+ assertEquals(OptionalInt.of(2), requestSpecs.get(1).scope.destinationBrokerId());
+
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), listOffsetsResponseSuccess(requestSpecs.get(0).keys), NODE_1);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(1), listOffsetsResponseSuccess(requestSpecs.get(1).keys), NODE_2);
+ assertTrue(result.all().get(tp0).isDone());
+ assertTrue(result.all().get(tp1).isDone());
+
+ //
+ // Request 2 - T-1 and T-2
+ //
+ // On the second request, the partition leader cache already contains some of the leadership data.
+ // Now the lookup and fulfillment stages overlap.
+ requestKeys = Set.of(tp1, tp2);
+ result = new PartitionLeaderStrategy.PartitionLeaderFuture<>(requestKeys, partitionLeaderCache);
+ driver = buildDriver(result);
+
+ requestSpecs = driver.poll();
+ assertEquals(2, requestSpecs.size());
+
+ assertEquals(OptionalInt.empty(), requestSpecs.get(0).scope.destinationBrokerId());
+ assertEquals(Collections.singleton(tp2), requestSpecs.get(0).keys);
+ assertEquals(OptionalInt.of(2), requestSpecs.get(1).scope.destinationBrokerId());
+
+ // The cache will be populated using the leader information from this metadata response
+ leaders = Map.of(tp2, 1);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), metadataResponseWithPartitionLeaders(leaders), Node.noNode());
+ driver.onResponse(time.milliseconds(), requestSpecs.get(1), listOffsetsResponseSuccess(requestSpecs.get(1).keys), NODE_2);
+ assertTrue(result.all().get(tp1).isDone()); // Already fulfilled
+ assertFalse(result.all().get(tp2).isDone());
+
+ assertEquals(1, partitionLeaderCache.get(tp0));
+ assertEquals(2, partitionLeaderCache.get(tp1));
+ assertEquals(1, partitionLeaderCache.get(tp2));
+
+ // Finally, the fulfillment stage makes the actual request for the uncached topic-partition
+ requestSpecs = driver.poll();
+ assertEquals(1, requestSpecs.size());
+
+ assertEquals(OptionalInt.of(1), requestSpecs.get(0).scope.destinationBrokerId());
+
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), listOffsetsResponseSuccess(requestSpecs.get(0).keys), NODE_1);
+ assertTrue(result.all().get(tp1).isDone());
+ assertTrue(result.all().get(tp2).isDone());
+
+ //
+ // Request 3 - T-0, T-1 and T-2
+ //
+ // On the third request, the partition leader cache contains all the leadership data
+ requestKeys = Set.of(tp0, tp1, tp2);
+ result = new PartitionLeaderStrategy.PartitionLeaderFuture<>(requestKeys, partitionLeaderCache);
+ driver = buildDriver(result);
+
+ requestSpecs = driver.poll();
+ assertEquals(2, requestSpecs.size());
+
+ // We can tell this is the fulfillment stage by the destination broker id being set
+ assertEquals(OptionalInt.of(1), requestSpecs.get(0).scope.destinationBrokerId());
+ assertEquals(OptionalInt.of(2), requestSpecs.get(1).scope.destinationBrokerId());
+
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), listOffsetsResponseSuccess(requestSpecs.get(0).keys), NODE_1);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(1), listOffsetsResponseSuccess(requestSpecs.get(1).keys), NODE_2);
+ assertTrue(result.all().get(tp0).isDone());
+ assertTrue(result.all().get(tp1).isDone());
+ assertTrue(result.all().get(tp2).isDone());
+
+ //
+ // Request 4 - T-0, T-1, T-2 and T-3
+ //
+ // On the fourth request, the partition leader cache already contains some of the leadership data.
+ // Now the lookup and fulfillment stages overlap.
+ requestKeys = Set.of(tp0, tp1, tp2, tp3);
+ result = new PartitionLeaderStrategy.PartitionLeaderFuture<>(requestKeys, partitionLeaderCache);
+ driver = buildDriver(result);
+
+ requestSpecs = driver.poll();
+ assertEquals(3, requestSpecs.size());
+
+ assertEquals(OptionalInt.empty(), requestSpecs.get(0).scope.destinationBrokerId());
+ assertEquals(Collections.singleton(tp3), requestSpecs.get(0).keys);
+ assertEquals(OptionalInt.of(1), requestSpecs.get(1).scope.destinationBrokerId());
+ assertEquals(OptionalInt.of(2), requestSpecs.get(2).scope.destinationBrokerId());
+
+ // The cache will be populated using the leader information from this metadata response
+ leaders = Map.of(tp3, 2);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), metadataResponseWithPartitionLeaders(leaders), Node.noNode());
+ driver.onResponse(time.milliseconds(), requestSpecs.get(1), listOffsetsResponseSuccess(requestSpecs.get(1).keys), NODE_1);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(2), listOffsetsResponseSuccess(requestSpecs.get(2).keys), NODE_2);
+ assertTrue(result.all().get(tp0).isDone()); // Already fulfilled
+ assertTrue(result.all().get(tp1).isDone()); // Already fulfilled
+ assertTrue(result.all().get(tp2).isDone()); // Already fulfilled
+ assertFalse(result.all().get(tp3).isDone());
+
+ assertEquals(1, partitionLeaderCache.get(tp0));
+ assertEquals(2, partitionLeaderCache.get(tp1));
+ assertEquals(1, partitionLeaderCache.get(tp2));
+ assertEquals(2, partitionLeaderCache.get(tp3));
+
+ // Finally, the fulfillment stage makes the actual request for the uncached topic-partition
+ requestSpecs = driver.poll();
+ assertEquals(1, requestSpecs.size());
+
+ assertEquals(OptionalInt.of(2), requestSpecs.get(0).scope.destinationBrokerId());
+
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), listOffsetsResponseSuccess(requestSpecs.get(0).keys), NODE_2);
+ assertTrue(result.all().get(tp0).isDone());
+ assertTrue(result.all().get(tp1).isDone());
+ assertTrue(result.all().get(tp2).isDone());
+ assertTrue(result.all().get(tp3).isDone());
+ }
+
+ @Test
+ public void testNotLeaderFulfillmentError() {
+ Map<TopicPartition, Integer> partitionLeaderCache = new HashMap<>();
+
+ TopicPartition tp0 = new TopicPartition("T", 0);
+ TopicPartition tp1 = new TopicPartition("T", 1);
+ Set<TopicPartition> requestKeys = Set.of(tp0, tp1);
+
+ // First, the lookup stage needs to obtain leadership data because the cache is empty
+ PartitionLeaderStrategy.PartitionLeaderFuture<Void> result =
+ new PartitionLeaderStrategy.PartitionLeaderFuture<>(requestKeys, partitionLeaderCache);
+ AdminApiDriver<TopicPartition, Void> driver = buildDriver(result);
+
+ List<AdminApiDriver.RequestSpec<TopicPartition>> requestSpecs = driver.poll();
+ assertEquals(1, requestSpecs.size());
+
+ assertEquals(OptionalInt.empty(), requestSpecs.get(0).scope.destinationBrokerId());
+ assertEquals(requestKeys, requestSpecs.get(0).keys);
+
+ // The cache will be populated using the leader information from this metadata response
+ Map<TopicPartition, Integer> leaders = Map.of(tp0, 1, tp1, 2);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), metadataResponseWithPartitionLeaders(leaders), Node.noNode());
+ assertFalse(result.all().get(tp0).isDone());
+ assertFalse(result.all().get(tp1).isDone());
+
+ assertEquals(1, partitionLeaderCache.get(tp0));
+ assertEquals(2, partitionLeaderCache.get(tp1));
+
+ // Second, the fulfillment stage makes the actual requests
+ requestSpecs = driver.poll();
+ assertEquals(2, requestSpecs.size());
+
+ assertEquals(OptionalInt.of(1), requestSpecs.get(0).scope.destinationBrokerId());
+ assertEquals(OptionalInt.of(2), requestSpecs.get(1).scope.destinationBrokerId());
+
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), listOffsetsResponseSuccess(requestSpecs.get(0).keys), NODE_1);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(1), listOffsetsResponseFailure(requestSpecs.get(1).keys, Errors.NOT_LEADER_OR_FOLLOWER), NODE_2);
+ assertTrue(result.all().get(tp0).isDone());
+ assertFalse(result.all().get(tp1).isDone());
+
+ // Now the lookup occurs again - change leadership to node 1
+ requestSpecs = driver.poll();
+ assertEquals(1, requestSpecs.size());
+
+ assertEquals(OptionalInt.empty(), requestSpecs.get(0).scope.destinationBrokerId());
+
+ leaders = Map.of(tp1, 1);
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), metadataResponseWithPartitionLeaders(leaders), Node.noNode());
+ assertTrue(result.all().get(tp0).isDone());
+ assertFalse(result.all().get(tp1).isDone());
+
+ assertEquals(1, partitionLeaderCache.get(tp0));
+ assertEquals(1, partitionLeaderCache.get(tp1));
+
+ // And the fulfillment stage makes the actual request
+ requestSpecs = driver.poll();
+ assertEquals(1, requestSpecs.size());
+
+ assertEquals(OptionalInt.of(1), requestSpecs.get(0).scope.destinationBrokerId());
+
+ driver.onResponse(time.milliseconds(), requestSpecs.get(0), listOffsetsResponseSuccess(requestSpecs.get(0).keys), NODE_1);
+ assertTrue(result.all().get(tp0).isDone());
+ assertTrue(result.all().get(tp1).isDone());
+ }
+
+ @Test
+ public void testFatalLookupError() {
+ TopicPartition tp0 = new TopicPartition("T", 0);
+ Map<TopicPartition, Integer> partitionLeaderCache = new HashMap<>();
+ PartitionLeaderStrategy.PartitionLeaderFuture<Void> result =
+ new PartitionLeaderStrategy.PartitionLeaderFuture<>(Collections.singleton(tp0), partitionLeaderCache);
+ AdminApiDriver<TopicPartition, Void> driver = buildDriver(result);
+
+ List<AdminApiDriver.RequestSpec<TopicPartition>> requestSpecs = driver.poll();
+ assertEquals(1, requestSpecs.size());
+
+ AdminApiDriver.RequestSpec<TopicPartition> spec = requestSpecs.get(0);
+ assertEquals(Collections.singleton(tp0), spec.keys);
+
+ driver.onFailure(time.milliseconds(), spec, new UnknownServerException());
+ assertTrue(result.all().get(tp0).isDone());
+ TestUtils.assertFutureThrows(result.all().get(tp0), UnknownServerException.class);
+ assertEquals(Collections.emptyList(), driver.poll());
+ }
+
+ @Test
+ public void testRetryLookupAfterDisconnect() {
+ TopicPartition tp0 = new TopicPartition("T", 0);
+ Map<TopicPartition, Integer> partitionLeaderCache = new HashMap<>();
+ PartitionLeaderStrategy.PartitionLeaderFuture<Void> result =
+ new PartitionLeaderStrategy.PartitionLeaderFuture<>(Collections.singleton(tp0), partitionLeaderCache);
+ AdminApiDriver<TopicPartition, Void> driver = buildDriver(result);
+
+ List<AdminApiDriver.RequestSpec<TopicPartition>> requestSpecs = driver.poll();
+ assertEquals(1, requestSpecs.size());
+
+ AdminApiDriver.RequestSpec<TopicPartition> spec = requestSpecs.get(0);
+ assertEquals(Collections.singleton(tp0), spec.keys);
+
+ driver.onFailure(time.milliseconds(), spec, new DisconnectException());
+ List<AdminApiDriver.RequestSpec<TopicPartition>> retrySpecs = driver.poll();
+ assertEquals(1, retrySpecs.size());
+
+ AdminApiDriver.RequestSpec<TopicPartition> retrySpec = retrySpecs.get(0);
+ assertEquals(Collections.singleton(tp0), retrySpec.keys);
+ assertEquals(time.milliseconds(), retrySpec.nextAllowedTryMs);
+ assertEquals(Collections.emptyList(), driver.poll());
+ }
+
+ private MetadataResponse metadataResponseWithPartitionLeaders(Map<TopicPartition, Integer> mapping) {
+ MetadataResponseData response = new MetadataResponseData();
+ mapping.forEach((tp, brokerId) -> response.topics().add(new MetadataResponseData.MetadataResponseTopic()
+ .setName(tp.topic())
+ .setPartitions(Collections.singletonList(new MetadataResponseData.MetadataResponsePartition()
+ .setPartitionIndex(tp.partition())
+ .setLeaderId(brokerId)))));
+ return new MetadataResponse(response, ApiKeys.METADATA.latestVersion());
+ }
+
+ private ListOffsetsResponse listOffsetsResponseSuccess(Set<TopicPartition> keys) {
+ // This structure is not quite how Kafka does it, but it works for the MockApiHandler
+ ListOffsetsResponseData response = new ListOffsetsResponseData();
+ keys.forEach(tp -> {
+ ListOffsetsResponseData.ListOffsetsPartitionResponse partResponse =
+ new ListOffsetsResponseData.ListOffsetsPartitionResponse()
+ .setPartitionIndex(tp.partition());
+ ListOffsetsResponseData.ListOffsetsTopicResponse topicResponse =
+ new ListOffsetsResponseData.ListOffsetsTopicResponse()
+ .setName(tp.topic())
+ .setPartitions(Collections.singletonList(partResponse));
+ response.topics().add(topicResponse);
+ });
+ return new ListOffsetsResponse(response);
+ }
+
+ private ListOffsetsResponse listOffsetsResponseFailure(Set<TopicPartition> keys, Errors error) {
+ // This structure is not quite how Kafka does it, but it works for the MockApiHandler
+ ListOffsetsResponseData response = new ListOffsetsResponseData();
+ keys.forEach(tp -> {
+ ListOffsetsResponseData.ListOffsetsPartitionResponse partResponse =
+ new ListOffsetsResponseData.ListOffsetsPartitionResponse()
+ .setPartitionIndex(tp.partition())
+ .setErrorCode(error.code());
+ ListOffsetsResponseData.ListOffsetsTopicResponse topicResponse =
+ new ListOffsetsResponseData.ListOffsetsTopicResponse()
+ .setName(tp.topic())
+ .setPartitions(Collections.singletonList(partResponse));
+ response.topics().add(topicResponse);
+ });
+ return new ListOffsetsResponse(response);
+ }
+
+ private class MockApiHandler extends AdminApiHandler.Batched<TopicPartition, Void> {
+ private final PartitionLeaderStrategy partitionLeaderStrategy = new PartitionLeaderStrategy(logContext);
+
+ @Override
+ public String apiName() {
+ return "mock-api";
+ }
+
+ @Override
+ public AbstractRequest.Builder<?> buildBatchedRequest(
+ int brokerId,
+ Set<TopicPartition> keys
+ ) {
+ return new MetadataRequest.Builder(new MetadataRequestData());
+ }
+
+ @Override
+ public ApiResult<TopicPartition, Void> handleResponse(
+ Node broker,
+ Set<TopicPartition> keys,
+ AbstractResponse abstractResponse
+ ) {
+ ListOffsetsResponse response = (ListOffsetsResponse) abstractResponse;
+
+ Map<TopicPartition, Void> completed = new HashMap<>();
+ Map<TopicPartition, Throwable> failed = new HashMap<>();
+ List<TopicPartition> unmapped = new ArrayList<>();
+
+ response.topics().forEach(topic -> topic.partitions().forEach(partition -> {
+ TopicPartition tp = new TopicPartition(topic.name(), partition.partitionIndex());
+ if (partition.errorCode() != Errors.NONE.code()) {
+ Exception exception = Errors.forCode(partition.errorCode()).exception();
+ if (exception instanceof NotLeaderOrFollowerException || exception instanceof LeaderNotAvailableException) {
+ unmapped.add(tp);
+ } else if (!(exception instanceof RetriableException)) {
+ failed.put(tp, Errors.forCode(partition.errorCode()).exception());
+ }
+ } else {
+ completed.put(tp, null);
+ }
+ }));
+
+ return new ApiResult<>(completed, failed, unmapped);
+ }
+
+ @Override
+ public PartitionLeaderStrategy lookupStrategy() {
+ return partitionLeaderStrategy;
+ }
+ }
+}