Skip waiting for first lookup version to get initialized (#15598)

apache · Jan 9, 2024 · 747d973 · 747d973
1 parent ea6ba40
commit 747d973
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 1 deletion.
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
@@ -668,7 +668,11 @@ public void handle(Map<String, LookupExtractorFactoryContainer> lookupMap, Looku
           e -> true,
           startRetries
       );
-      if (lookupExtractorFactoryContainer.getLookupExtractorFactory().isInitialized()) {
+      /*
+       if new container is initailized then add it to manager to start serving immediately.
+       if old container is null then it is fresh load, we can skip waiting for initialization and add the container to registry first. Esp for MSQ workers.
+       */
+      if (old == null || lookupExtractorFactoryContainer.getLookupExtractorFactory().isInitialized()) {
         old = lookupMap.put(lookupName, lookupExtractorFactoryContainer);
         LOG.debug("Loaded lookup [%s] with spec [%s].", lookupName, lookupExtractorFactoryContainer);
         manager.dropContainer(old, lookupName);

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
@@ -416,6 +416,7 @@ public void testUpdateWithHigherVersion() throws Exception
     EasyMock.expect(lookupExtractorFactory1.destroy()).andReturn(true).once();
 
     LookupExtractorFactory lookupExtractorFactory2 = EasyMock.createNiceMock(LookupExtractorFactory.class);
+    EasyMock.expect(lookupExtractorFactory2.isInitialized()).andReturn(true).once();
     EasyMock.expect(lookupExtractorFactory2.start()).andReturn(true).once();
 
     EasyMock.replay(lookupExtractorFactory1, lookupExtractorFactory2);
@@ -482,6 +483,35 @@ public void testUpdateWithLowerVersion() throws Exception
     EasyMock.verify(lookupExtractorFactory1, lookupExtractorFactory2);
   }
 
+  @Test
+  public void testAddingNewContainerImmediatelyWithoutWaiting() throws Exception
+  {
+    LookupExtractorFactory lookupExtractorFactory1 = EasyMock.createNiceMock(LookupExtractorFactory.class);
+    EasyMock.expect(lookupExtractorFactory1.start()).andReturn(true).once();
+    EasyMock.expect(lookupExtractorFactory1.isInitialized()).andReturn(false).anyTimes();
+    EasyMock.replay(lookupExtractorFactory1);
+    Map<String, Object> lookupMap = new HashMap<>();
+    String strResult = mapper.writeValueAsString(lookupMap);
+    Request request = new Request(HttpMethod.GET, new URL("http://localhost:1234/xx"));
+    EasyMock.expect(config.getLookupTier()).andReturn(LOOKUP_TIER).anyTimes();
+    EasyMock.replay(config);
+    EasyMock.expect(druidLeaderClient.makeRequest(
+                HttpMethod.GET,
+                "/druid/coordinator/v1/lookups/config/lookupTier?detailed=true"
+            ))
+            .andReturn(request);
+    StringFullResponseHolder responseHolder = new StringFullResponseHolder(
+        newEmptyResponse(HttpResponseStatus.OK),
+        StandardCharsets.UTF_8
+    ).addChunk(strResult);
+    EasyMock.expect(druidLeaderClient.go(request)).andReturn(responseHolder);
+    EasyMock.replay(druidLeaderClient);
+    lookupReferencesManager.start();
+    lookupReferencesManager.add("testName", new LookupExtractorFactoryContainer("1", lookupExtractorFactory1));
+    lookupReferencesManager.handlePendingNotices();
+    Assert.assertTrue(lookupReferencesManager.get("testName").isPresent());
+    EasyMock.verify(lookupExtractorFactory1);
+  }
   @Test
   public void testRemoveNonExisting() throws Exception
   {