Use CoordinatorClient for fetching tier lookup (apache#18142)

uds5501 · web-flow · commit 6b3cb50cdd8f · 2025-06-27T18:08:40.000+05:30
This patch is part of the effort to phase out `DruidLeaderClient` and utilize `CoordinatorClient`
to interact with Coordinators.

This patch does the following:
- Add API `CoordinatorClient.fetchLookupForTier()`
- Replace use of `DruidLeaderClient` in `LookupReferencesManager` with `CoordinatorClient`
diff --git a/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClient.java b/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClient.java
@@ -23,6 +23,7 @@
 import org.apache.druid.client.BootstrapSegmentsResponse;
 import org.apache.druid.client.ImmutableSegmentLoadInfo;
 import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.query.lookup.LookupExtractorFactoryContainer;
 import org.apache.druid.rpc.ServiceRetryPolicy;
 import org.apache.druid.segment.metadata.DataSourceInformation;
 import org.apache.druid.server.compaction.CompactionStatusResponse;
@@ -32,6 +33,7 @@
 
 import javax.annotation.Nullable;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 public interface CoordinatorClient
@@ -94,4 +96,13 @@ public interface CoordinatorClient
    * API: {@code GET /druid/coordinator/v1/config}
    */
   ListenableFuture<CoordinatorDynamicConfig> getCoordinatorDynamicConfig();
+
+  /**
+   * Gets the lookup configuration for a tier
+   * <p>
+   * API: {@code GET /druid/coordinator/v1/lookups/config/<tier>}
+   *
+   * @param tier         The name of the tier for which the lookup configuration is to be fetched.
+   */
+  ListenableFuture<Map<String, LookupExtractorFactoryContainer>> fetchLookupsForTier(String tier);
 }
diff --git a/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClientImpl.java b/server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClientImpl.java
@@ -30,8 +30,11 @@
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.jackson.JacksonUtils;
 import org.apache.druid.java.util.http.client.response.BytesFullResponseHandler;
+import org.apache.druid.java.util.http.client.response.BytesFullResponseHolder;
 import org.apache.druid.java.util.http.client.response.InputStreamResponseHandler;
 import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.query.lookup.LookupExtractorFactoryContainer;
+import org.apache.druid.query.lookup.LookupUtils;
 import org.apache.druid.rpc.RequestBuilder;
 import org.apache.druid.rpc.ServiceClient;
 import org.apache.druid.rpc.ServiceRetryPolicy;
@@ -46,6 +49,7 @@
 import javax.annotation.Nullable;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 public class CoordinatorClientImpl implements CoordinatorClient
@@ -239,4 +243,36 @@ public ListenableFuture<CoordinatorDynamicConfig> getCoordinatorDynamicConfig()
         )
     );
   }
+
+  @Override
+  public ListenableFuture<Map<String, LookupExtractorFactoryContainer>> fetchLookupsForTier(
+      String tier
+  )
+  {
+    final String path = StringUtils.format(
+        "/druid/coordinator/v1/lookups/config/%s?detailed=true",
+        StringUtils.urlEncode(tier)
+    );
+
+    return FutureUtils.transform(
+        client.asyncRequest(
+            new RequestBuilder(HttpMethod.GET, path),
+            new BytesFullResponseHandler()
+        ),
+        this::extractLookupFactory
+    );
+  }
+
+  private Map<String, LookupExtractorFactoryContainer> extractLookupFactory(BytesFullResponseHolder holder)
+  {
+    Map<String, Object> lookupNameToGenericConfig = JacksonUtils.readValue(
+        jsonMapper,
+        holder.getContent(),
+        new TypeReference<>() {}
+    );
+    return LookupUtils.tryConvertObjectMapToLookupConfigMap(
+        lookupNameToGenericConfig,
+        jsonMapper
+    );
+  }
 }
diff --git a/server/src/main/java/org/apache/druid/client/coordinator/NoopCoordinatorClient.java b/server/src/main/java/org/apache/druid/client/coordinator/NoopCoordinatorClient.java
@@ -23,6 +23,7 @@
 import org.apache.druid.client.BootstrapSegmentsResponse;
 import org.apache.druid.client.ImmutableSegmentLoadInfo;
 import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.query.lookup.LookupExtractorFactoryContainer;
 import org.apache.druid.rpc.ServiceRetryPolicy;
 import org.apache.druid.segment.metadata.DataSourceInformation;
 import org.apache.druid.server.compaction.CompactionStatusResponse;
@@ -32,6 +33,7 @@
 
 import javax.annotation.Nullable;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 public class NoopCoordinatorClient implements CoordinatorClient
@@ -97,4 +99,12 @@ public ListenableFuture<CoordinatorDynamicConfig> getCoordinatorDynamicConfig()
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public ListenableFuture<Map<String, LookupExtractorFactoryContainer>> fetchLookupsForTier(
+      String tier
+  )
+  {
+    throw new UnsupportedOperationException();
+  }
+
 }
diff --git a/server/src/main/java/org/apache/druid/query/lookup/LookupReferencesManager.java b/server/src/main/java/org/apache/druid/query/lookup/LookupReferencesManager.java
@@ -20,34 +20,31 @@
 package org.apache.druid.query.lookup;
 
 
-import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
+import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.inject.Inject;
 import org.apache.commons.lang3.mutable.MutableBoolean;
-import org.apache.druid.client.coordinator.Coordinator;
+import org.apache.druid.client.coordinator.CoordinatorClient;
+import org.apache.druid.common.guava.FutureUtils;
 import org.apache.druid.concurrent.LifecycleLock;
-import org.apache.druid.discovery.DruidLeaderClient;
 import org.apache.druid.guice.ManageLifecycle;
 import org.apache.druid.guice.annotations.Json;
 import org.apache.druid.java.util.common.FileUtils;
-import org.apache.druid.java.util.common.IOE;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.RE;
 import org.apache.druid.java.util.common.RetryUtils;
-import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.concurrent.Execs;
 import org.apache.druid.java.util.common.lifecycle.LifecycleStart;
 import org.apache.druid.java.util.common.lifecycle.LifecycleStop;
 import org.apache.druid.java.util.emitter.EmittingLogger;
-import org.apache.druid.java.util.http.client.response.StringFullResponseHolder;
+import org.apache.druid.rpc.HttpResponseException;
 import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
 import org.apache.druid.server.metrics.DataSourceTaskIdHolder;
-import org.jboss.netty.handler.codec.http.HttpMethod;
 import org.jboss.netty.handler.codec.http.HttpResponseStatus;
 
 import javax.annotation.Nullable;
@@ -78,7 +75,7 @@
  * This class provide a basic {@link LookupExtractorFactory} references manager. It allows basic operations fetching,
  * listing, adding and deleting of {@link LookupExtractor} objects, and can take periodic snap shot of the loaded lookup
  * extractor specifications in order to bootstrap nodes after restart.
- *
+ * <p>
  * It also implements {@link LookupExtractorFactoryContainerProvider}, to supply queries and indexing transformations
  * with a reference to a {@link LookupExtractorFactoryContainer}. This class is a companion of
  * {@link org.apache.druid.server.lookup.cache.LookupCoordinatorManager}, which communicates with
@@ -89,9 +86,6 @@ public class LookupReferencesManager implements LookupExtractorFactoryContainerP
 {
   private static final EmittingLogger LOG = new EmittingLogger(LookupReferencesManager.class);
 
-  private static final TypeReference<Map<String, Object>> LOOKUPS_ALL_GENERIC_REFERENCE =
-      new TypeReference<>() {};
-
   // Lookups state (loaded/to-be-loaded/to-be-dropped etc) is managed by immutable LookupUpdateState instance.
   // Any update to state is done by creating updated LookupUpdateState instance and atomically setting that
   // into the ref here.
@@ -111,9 +105,7 @@ public class LookupReferencesManager implements LookupExtractorFactoryContainerP
   //for unit testing only
   private final boolean testMode;
 
-  private final DruidLeaderClient druidLeaderClient;
-
-  private final ObjectMapper jsonMapper;
+  private final CoordinatorClient coordinatorClient;
 
   private final LookupListeningAnnouncerConfig lookupListeningAnnouncerConfig;
 
@@ -125,18 +117,18 @@ public class LookupReferencesManager implements LookupExtractorFactoryContainerP
   public LookupReferencesManager(
       LookupConfig lookupConfig,
       @Json ObjectMapper objectMapper,
-      @Coordinator DruidLeaderClient druidLeaderClient,
+      CoordinatorClient coordinatorClient,
       LookupListeningAnnouncerConfig lookupListeningAnnouncerConfig
   )
   {
-    this(lookupConfig, objectMapper, druidLeaderClient, lookupListeningAnnouncerConfig, false);
+    this(lookupConfig, objectMapper, coordinatorClient, lookupListeningAnnouncerConfig, false);
   }
 
   @VisibleForTesting
   LookupReferencesManager(
       LookupConfig lookupConfig,
       ObjectMapper objectMapper,
-      DruidLeaderClient druidLeaderClient,
+      CoordinatorClient coordinatorClient,
       LookupListeningAnnouncerConfig lookupListeningAnnouncerConfig,
       boolean testMode
   )
@@ -146,8 +138,7 @@ public LookupReferencesManager(
     } else {
       this.lookupSnapshotTaker = new LookupSnapshotTaker(objectMapper, lookupConfig.getSnapshotWorkingDir());
     }
-    this.druidLeaderClient = druidLeaderClient;
-    this.jsonMapper = objectMapper;
+    this.coordinatorClient = coordinatorClient;
     this.lookupListeningAnnouncerConfig = lookupListeningAnnouncerConfig;
     this.lookupConfig = lookupConfig;
     this.testMode = testMode;
@@ -286,7 +277,11 @@ public void add(String lookupName, LookupExtractorFactoryContainer lookupExtract
     if (lookupLoadingSpec.getMode() == LookupLoadingSpec.Mode.NONE ||
         (lookupLoadingSpec.getMode() == LookupLoadingSpec.Mode.ONLY_REQUIRED
          && !lookupLoadingSpec.getLookupsToLoad().contains(lookupName))) {
-      LOG.info("Skipping notice to add lookup [%s] since current lookup loading mode [%s] does not allow it.", lookupName, lookupLoadingSpec.getMode());
+      LOG.info(
+          "Skipping notice to add lookup[%s] since current lookup loading mode[%s] does not allow it.",
+          lookupName,
+          lookupLoadingSpec.getMode()
+      );
       return;
     }
     addNotice(new LoadNotice(lookupName, lookupExtractorFactoryContainer, lookupConfig.getLookupStartRetries()));
@@ -401,7 +396,8 @@ private void loadLookupsAndInitStateRef()
       lookupBeanList = getLookupsList();
       if (lookupLoadingSpec.getMode() == LookupLoadingSpec.Mode.ONLY_REQUIRED && lookupBeanList != null) {
         lookupBeanList = lookupBeanList.stream()
-                                       .filter(lookupBean -> lookupLoadingSpec.getLookupsToLoad().contains(lookupBean.getName()))
+                                       .filter(lookupBean -> lookupLoadingSpec.getLookupsToLoad()
+                                                                              .contains(lookupBean.getName()))
                                        .collect(Collectors.toList());
       }
     }
@@ -437,7 +433,6 @@ private List<LookupBean> getLookupsList()
    * Returns a list of lookups from the coordinator if the coordinator is available. If it's not available, returns null.
    *
    * @param tier lookup tier name
-   *
    * @return list of LookupBean objects, or null
    */
   @Nullable
@@ -476,37 +471,24 @@ private List<LookupBean> getLookupListFromCoordinator(String tier)
 
   @Nullable
   private Map<String, LookupExtractorFactoryContainer> tryGetLookupListFromCoordinator(String tier)
-      throws IOException, InterruptedException
   {
-    final StringFullResponseHolder response = fetchLookupsForTier(tier);
-    if (response.getStatus().equals(HttpResponseStatus.NOT_FOUND)) {
-      LOG.warn("No lookups found for tier [%s], response [%s]", tier, response);
-      return null;
-    } else if (!response.getStatus().equals(HttpResponseStatus.OK)) {
-      throw new IOE(
-          "Error while fetching lookup code from Coordinator with status[%s] and content[%s]",
-          response.getStatus(),
-          response.getContent()
-      );
+    try {
+      return FutureUtils.getUnchecked(coordinatorClient.fetchLookupsForTier(tier), true);
     }
-
-    // Older version of getSpecificTier returns a list of lookup names.
-    // Lookup loading is performed via snapshot if older version is present.
-    // This check is only for backward compatibility and should be removed in a future release
-    if (response.getContent().startsWith("[")) {
-      LOG.info(
-          "Failed to retrieve lookup information from coordinator, " +
-          "because coordinator appears to be running on older Druid version. " +
-          "Attempting to load lookups using snapshot instead"
-      );
-      return null;
-    } else {
-      Map<String, Object> lookupNameToGenericConfig =
-          jsonMapper.readValue(response.getContent(), LOOKUPS_ALL_GENERIC_REFERENCE);
-      return LookupUtils.tryConvertObjectMapToLookupConfigMap(
-          lookupNameToGenericConfig,
-          jsonMapper
-      );
+    catch (Exception e) {
+      Throwable rootCause = Throwables.getRootCause(e);
+      if (rootCause instanceof HttpResponseException) {
+        final HttpResponseException httpException = (HttpResponseException) rootCause;
+        if (httpException.getResponse().getStatus().equals(HttpResponseStatus.NOT_FOUND)) {
+          LOG.info(
+              "No lookups found for tier [%s], status [%s]",
+              tier,
+              httpException.getResponse().getStatus()
+          );
+          return null;
+        }
+      }
+      throw e;
     }
   }
 
@@ -628,15 +610,6 @@ private LookupUpdateState atomicallyUpdateStateRef(Function<LookupUpdateState, L
     }
   }
 
-  private StringFullResponseHolder fetchLookupsForTier(String tier) throws InterruptedException, IOException
-  {
-    return druidLeaderClient.go(
-        druidLeaderClient.makeRequest(
-            HttpMethod.GET,
-            StringUtils.format("/druid/coordinator/v1/lookups/config/%s?detailed=true", tier)
-        )
-    );
-  }
   private void dropContainer(LookupExtractorFactoryContainer container, String lookupName)
   {
     if (container != null) {
@@ -651,10 +624,12 @@ private void dropContainer(LookupExtractorFactoryContainer container, String loo
       }
     }
   }
+
   @VisibleForTesting
   interface Notice
   {
-    void handle(Map<String, LookupExtractorFactoryContainer> lookupMap, LookupReferencesManager manager) throws Exception;
+    void handle(Map<String, LookupExtractorFactoryContainer> lookupMap, LookupReferencesManager manager)
+        throws Exception;
   }
 
   private static class LoadNotice implements Notice
@@ -741,6 +716,7 @@ public void handle(Map<String, LookupExtractorFactoryContainer> lookupMap, Looku
         }
       });
     }
+
     @Override
     public String toString()
     {
diff --git a/server/src/test/java/org/apache/druid/client/coordinator/CoordinatorClientImplTest.java b/server/src/test/java/org/apache/druid/client/coordinator/CoordinatorClientImplTest.java
diff --git a/server/src/test/java/org/apache/druid/query/lookup/LookupReferencesManagerTest.java b/server/src/test/java/org/apache/druid/query/lookup/LookupReferencesManagerTest.java
