Use MetadataStoreCacheLoader replace MetadataStoreCache in KopBrokerLookupManager

docs/configuration.md

@@ -80,6 +80,7 @@ This section lists configurations that may affect the performance.
 | requestTimeoutMs  | Limit the timeout in milliseconds for request, like `request.timeout.ms` in Kafka client.<br>If a request was not processed in the timeout, KoP would return an error response to client. | 30000   |
 | connectionMaxIdleMs | The idle connection timeout in milliseconds. If the idle connection timeout (such as `connections.max.idle.ms` used in the Kafka server) is reached, the server handler will close this idle connection.<br>**Note**: If it is set to `-1`, it indicates that the idle connection timeout is disabled. | 600000 |
 | failedAuthenticationDelayMs | Connection close delay on failed authentication: this is the time (in milliseconds) by which connection close will be delayed on authentication failure, like `connection.failed.authentication.delay.ms` in Kafka server. | 300 |
+| brokerLookupTimeoutMs | The timeout for broker lookups (in milliseconds). | 30000 |
 
 > **NOTE**
 > 

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaProtocolHandler.java

@@ -448,8 +448,14 @@ public void start(BrokerService service) {
         offsetTopicClient = new SystemTopicClient(brokerService.pulsar(), kafkaConfig);
         txnTopicClient = new SystemTopicClient(brokerService.pulsar(), kafkaConfig);
 
-        kopBrokerLookupManager = new KopBrokerLookupManager(
-                brokerService.getPulsar(), kafkaConfig.getKafkaAdvertisedListeners());
+        try {
+            kopBrokerLookupManager = new KopBrokerLookupManager(
+                    brokerService.getPulsar(), kafkaConfig.getKafkaAdvertisedListeners(),
+                    kafkaConfig.getBrokerLookupTimeoutMs());
+        } catch (Exception ex) {
+            log.error("Failed to get kopBrokerLookupManager", ex);
+            throw new IllegalStateException(ex);
+        }
 
         brokerService.pulsar()
                 .getNamespaceService()

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaRequestHandler.java

@@ -186,7 +186,6 @@
 import org.apache.pulsar.common.util.FutureUtil;
 import org.apache.pulsar.common.util.Murmur3_32Hash;
 import org.apache.pulsar.metadata.api.extended.MetadataStoreExtended;
-import org.apache.pulsar.policies.data.loadbalancer.ServiceLookupData;
 
 /**
  * This class contains all the request handling methods.
@@ -2658,14 +2657,6 @@ static AbstractResponse failedResponse(KafkaHeaderAndRequest requestHar, Throwab
         return requestHar.getRequest().getErrorResponse(((Integer) THROTTLE_TIME_MS.defaultValue), e);
     }
 
-    // whether a ServiceLookupData contains wanted address.
-    static boolean lookupDataContainsAddress(ServiceLookupData data, String hostAndPort) {
-        return (data.getPulsarServiceUrl() != null && data.getPulsarServiceUrl().contains(hostAndPort))
-            || (data.getPulsarServiceUrlTls() != null && data.getPulsarServiceUrlTls().contains(hostAndPort))
-            || (data.getWebServiceUrl() != null && data.getWebServiceUrl().contains(hostAndPort))
-            || (data.getWebServiceUrlTls() != null && data.getWebServiceUrlTls().contains(hostAndPort));
-    }
-
     private static MemoryRecords validateRecords(short version, TopicPartition topicPartition, MemoryRecords records) {
         if (version >= 3) {
             Iterator<MutableRecordBatch> iterator = records.batches().iterator();

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaServiceConfiguration.java

@@ -231,6 +231,12 @@ public class KafkaServiceConfiguration extends ServiceConfiguration {
     )
     private int failedAuthenticationDelayMs = 300;
 
+    @FieldContext(
+            category = CATEGORY_KOP,
+            doc = "The timeout for broker lookups (in milliseconds)"
+    )
+    private int brokerLookupTimeoutMs = 30_000;
+
     // Kafka SSL configs
     @FieldContext(
         category = CATEGORY_KOP_SSL,

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KopBrokerLookupManager.java

@@ -13,25 +13,22 @@
  */
 package io.streamnative.pulsar.handlers.kop;
 
-import static io.streamnative.pulsar.handlers.kop.KafkaRequestHandler.lookupDataContainsAddress;
 
-import com.google.common.collect.Lists;
+import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.pulsar.broker.PulsarService;
-import org.apache.pulsar.broker.loadbalance.LoadManager;
+import org.apache.pulsar.broker.resources.MetadataStoreCacheLoader;
 import org.apache.pulsar.common.naming.TopicName;
-import org.apache.pulsar.common.util.FutureUtil;
-import org.apache.pulsar.metadata.api.MetadataCache;
-import org.apache.pulsar.policies.data.loadbalancer.LocalBrokerData;
+import org.apache.pulsar.policies.data.loadbalancer.LoadManagerReport;
 import org.apache.pulsar.policies.data.loadbalancer.ServiceLookupData;
 
 
@@ -41,10 +38,9 @@
 @Slf4j
 public class KopBrokerLookupManager {
 
-    private final PulsarService pulsarService;
     private final String advertisedListeners;
     private final LookupClient lookupClient;
-    private final MetadataCache<LocalBrokerData> localBrokerDataCache;
+    private final MetadataStoreCacheLoader metadataStoreCacheLoader;
 
     private final AtomicBoolean closed = new AtomicBoolean(false);
 
@@ -54,11 +50,12 @@ public class KopBrokerLookupManager {
     public static final ConcurrentHashMap<String, CompletableFuture<Optional<String>>>
             KOP_ADDRESS_CACHE = new ConcurrentHashMap<>();
 
-    public KopBrokerLookupManager(PulsarService pulsarService, String advertisedListeners) {
-        this.pulsarService = pulsarService;
+    public KopBrokerLookupManager(
+            PulsarService pulsarService, String advertisedListeners, int brokerLookupTimeoutMs) throws Exception {
         this.advertisedListeners = advertisedListeners;
         this.lookupClient = KafkaProtocolHandler.getLookupClient(pulsarService);
-        this.localBrokerDataCache = pulsarService.getLocalMetadataStore().getMetadataCache(LocalBrokerData.class);
+        this.metadataStoreCacheLoader = new MetadataStoreCacheLoader(pulsarService.getPulsarResources(),
+                brokerLookupTimeoutMs);
     }
 
     public CompletableFuture<Optional<InetSocketAddress>> findBroker(@NonNull TopicName topic,
@@ -184,83 +181,39 @@ private CompletableFuture<Optional<String>> getProtocolDataToAdvertise(
             return returnFuture;
         }
 
-        // advertised data is write in  /loadbalance/brokers/advertisedAddress:webServicePort
-        // here we get the broker url, need to find related webServiceUrl.
-        pulsarService.getPulsarResources()
-                .getDynamicConfigResources()
-                .getChildrenAsync(LoadManager.LOADBALANCE_BROKERS_ROOT)
-                .whenComplete((set, throwable) -> {
-                    if (throwable != null) {
-                        log.error("Error in getChildrenAsync(zk://loadbalance) for {}", pulsarAddress, throwable);
-                        returnFuture.complete(Optional.empty());
-                        return;
-                    }
-
-                    String hostAndPort = pulsarAddress.getHostName() + ":" + pulsarAddress.getPort();
-                    List<String> matchBrokers = Lists.newArrayList();
-                    // match host part of url
-                    for (String activeBroker : set) {
-                        if (activeBroker.startsWith(pulsarAddress.getHostName() + ":")) {
-                            matchBrokers.add(activeBroker);
-                        }
-                    }
-
-                    if (matchBrokers.isEmpty()) {
-                        log.error("No node for broker {} under zk://loadbalance", pulsarAddress);
-                        returnFuture.complete(Optional.empty());
-                        removeTopicManagerCache(topic.toString());
-                        return;
-                    }
+        List<LoadManagerReport> availableBrokers = metadataStoreCacheLoader.getAvailableBrokers();
+        if (log.isDebugEnabled()) {
+            availableBrokers.forEach(loadManagerReport ->
+                    log.debug("Handle getProtocolDataToAdvertise for {}, pulsarUrl: {}, "
+                                    + "pulsarUrlTls: {}, webUrl: {}, webUrlTls: {} kafka: {}",
+                            topic,
+                            loadManagerReport.getPulsarServiceUrl(),
+                            loadManagerReport.getPulsarServiceUrlTls(),
+                            loadManagerReport.getWebServiceUrl(),
+                            loadManagerReport.getWebServiceUrlTls(),
+                            loadManagerReport.getProtocol(KafkaProtocolHandler.PROTOCOL_NAME)));
+        }
 
-                    // Get a list of ServiceLookupData for each matchBroker.
-                    List<CompletableFuture<Optional<LocalBrokerData>>> list = matchBrokers.stream()
-                            .map(matchBroker -> localBrokerDataCache.get(
-                                    String.format("%s/%s", LoadManager.LOADBALANCE_BROKERS_ROOT, matchBroker)))
-                            .collect(Collectors.toList());
-
-                    FutureUtil.waitForAll(list).whenComplete((ignore, th) -> {
-                        if (th != null) {
-                            log.error("Error in getDataAsync() for {}", pulsarAddress, th);
-                            returnFuture.complete(Optional.empty());
-                            removeTopicManagerCache(topic.toString());
-                            return;
-                        }
-
-                        try {
-                            for (CompletableFuture<Optional<LocalBrokerData>> lookupData : list) {
-                                ServiceLookupData data = lookupData.get().get();
-                                if (log.isDebugEnabled()) {
-                                    log.debug("Handle getProtocolDataToAdvertise for {}, pulsarUrl: {}, "
-                                                    + "pulsarUrlTls: {}, webUrl: {}, webUrlTls: {} kafka: {}",
-                                            topic,
-                                            data.getPulsarServiceUrl(),
-                                            data.getPulsarServiceUrlTls(),
-                                            data.getWebServiceUrl(),
-                                            data.getWebServiceUrlTls(),
-                                            data.getProtocol(KafkaProtocolHandler.PROTOCOL_NAME));
-                                }
-
-                                if (lookupDataContainsAddress(data, hostAndPort)) {
-                                    KOP_ADDRESS_CACHE.put(topic.toString(), returnFuture);
-                                    returnFuture.complete(data.getProtocol(KafkaProtocolHandler.PROTOCOL_NAME));
-                                    return;
-                                }
-                            }
-                        } catch (Exception e) {
-                            log.error("Error in {} lookupFuture get: ", pulsarAddress, e);
-                            returnFuture.complete(Optional.empty());
-                            removeTopicManagerCache(topic.toString());
-                            return;
-                        }
-
-                        // no matching lookup data in all matchBrokers.
-                        log.error("Not able to search {} in all child of zk://loadbalance", pulsarAddress);
-                        returnFuture.complete(Optional.empty());
-                    });
-                });
+        String hostAndPort = pulsarAddress.getHostName() + ":" + pulsarAddress.getPort();
+        Optional<LoadManagerReport> serviceLookupData = availableBrokers.stream()
+                .filter(loadManagerReport -> lookupDataContainsAddress(loadManagerReport, hostAndPort)).findAny();
+        if (serviceLookupData.isPresent()) {
+            KOP_ADDRESS_CACHE.put(topic.toString(), returnFuture);
+            returnFuture.complete(serviceLookupData.get().getProtocol(KafkaProtocolHandler.PROTOCOL_NAME));
+        } else {
+            log.error("No node for broker {} under loadBalance", pulsarAddress);
+            removeTopicManagerCache(topic.toString());
+            returnFuture.complete(Optional.empty());
+        }
         return returnFuture;
     }
 
+    // whether a ServiceLookupData contains wanted address.
+    private static boolean lookupDataContainsAddress(ServiceLookupData data, String hostAndPort) {
+        return StringUtils.endsWith(data.getPulsarServiceUrl(), hostAndPort)
+                || StringUtils.endsWith(data.getPulsarServiceUrlTls(), hostAndPort);
+    }
+
     public static void removeTopicManagerCache(String topicName) {
         LOOKUP_CACHE.remove(topicName);
         KOP_ADDRESS_CACHE.remove(topicName);
@@ -279,6 +232,11 @@ public void close() {
             return;
         }
         clear();
+        try {
+            metadataStoreCacheLoader.close();
+        } catch (IOException e) {
+            log.error("Close metadataStoreCacheLoader failed.", e);
+        }
     }
 
 }