fix loading lookup extension

Signed-off-by: TessaIO <[email protected]>

extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java

@@ -28,6 +28,7 @@
 
 import javax.annotation.Nullable;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Callable;
@@ -74,14 +75,22 @@ public String apply(@Nullable final String key)
     }
 
     final String presentVal;
-    try {
-      presentVal = loadingCache.get(keyEquivalent, new ApplyCallable(keyEquivalent));
+    presentVal = this.loadingCache.getIfPresent(keyEquivalent);
+    if (presentVal != null) {
       return NullHandling.emptyToNullIfNeeded(presentVal);
     }
-    catch (ExecutionException e) {
-      LOGGER.debug("value not found for key [%s]", key);
+
+    String val = this.dataFetcher.fetch(keyEquivalent);
+    if (val == null) {
       return null;
     }
+
+    Map<String, String> map = new HashMap<>();
+    map.put(keyEquivalent, val);
+
+    this.loadingCache.putAll(map);
+
+    return NullHandling.emptyToNullIfNeeded(val);
   }
 
   @Override
@@ -108,13 +117,18 @@ public List<String> unapply(@Nullable final String value)
   @Override
   public boolean supportsAsMap()
   {
-    return false;
+    return true;
   }
 
   @Override
   public Map<String, String> asMap()
   {
-    throw new UnsupportedOperationException("Cannot get map view");
+    Iterable<Map.Entry<String, String>> data = this.dataFetcher.fetchAll();
+    Map<String, String> map = new HashMap<>();
+    if (data != null) {
+      data.forEach(each -> map.put(each.getKey(), each.getValue()));
+    }
+    return map;
   }
 
   @Override

extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java

@@ -31,8 +31,10 @@
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
+import java.util.AbstractMap;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 
@@ -49,7 +51,7 @@ public class LoadingLookupTest extends InitializedNullHandlingTest
   @Test
   public void testApplyEmptyOrNull() throws ExecutionException
   {
-    EasyMock.expect(lookupCache.get(EasyMock.eq(""), EasyMock.anyObject(Callable.class)))
+    EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("")))
             .andReturn("empty").atLeastOnce();
     EasyMock.replay(lookupCache);
     Assert.assertEquals("empty", loadingLookup.apply(""));
@@ -75,7 +77,7 @@ public void testUnapplyNull()
   @Test
   public void testApply() throws ExecutionException
   {
-    EasyMock.expect(lookupCache.get(EasyMock.eq("key"), EasyMock.anyObject(Callable.class))).andReturn("value").once();
+    EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("key"))).andReturn("value").once();
     EasyMock.replay(lookupCache);
     Assert.assertEquals(ImmutableMap.of("key", "value"), loadingLookup.applyAll(ImmutableSet.of("key")));
     EasyMock.verify(lookupCache);
@@ -105,17 +107,6 @@ public void testClose()
     EasyMock.verify(lookupCache, reverseLookupCache);
   }
 
-  @Test
-  public void testApplyWithExecutionError() throws ExecutionException
-  {
-    EasyMock.expect(lookupCache.get(EasyMock.eq("key"), EasyMock.anyObject(Callable.class)))
-            .andThrow(new ExecutionException(null))
-            .once();
-    EasyMock.replay(lookupCache);
-    Assert.assertNull(loadingLookup.apply("key"));
-    EasyMock.verify(lookupCache);
-  }
-
   @Test
   public void testUnApplyWithExecutionError() throws ExecutionException
   {
@@ -136,13 +127,16 @@ public void testGetCacheKey()
   @Test
   public void testSupportsAsMap()
   {
-    Assert.assertFalse(loadingLookup.supportsAsMap());
+    Assert.assertTrue(loadingLookup.supportsAsMap());
   }
 
   @Test
   public void testAsMap()
   {
-    expectedException.expect(UnsupportedOperationException.class);
-    loadingLookup.asMap();
+    Map.Entry<String, String> fetchedData = new AbstractMap.SimpleEntry<>("dummy", "test");
+    EasyMock.expect(dataFetcher.fetchAll()).andReturn(Collections.singletonList(fetchedData));
+    EasyMock.replay(dataFetcher);
+    Assert.assertEquals(loadingLookup.asMap().size(), 1);
+    EasyMock.verify(dataFetcher);
   }
 }

processing/src/test/java/org/apache/druid/query/MetricsEmittingMergingBlockingPoolTest.java

@@ -0,0 +1,49 @@
+
+/*
+ * 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.druid.query;
+
+import org.apache.druid.java.util.metrics.StubServiceEmitter;
+import org.apache.druid.query.DruidProcessingConfig;
+import com.google.common.base.Supplier;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.nio.ByteBuffer;
+
+@SuppressWarnings("DoNotMock")
+public class MetricsEmittingMergingBlockingPoolTest
+{
+  @Test
+  public void testPrioritizedExecutorDelegate()
+  {
+    PrioritizedExecutorService service = Mockito.mock(PrioritizedExecutorService.class);
+    Mockito.when(service.getQueueSize()).thenReturn(10);
+    Mockito.when(service.getActiveTasks()).thenReturn(2);
+    ExecutorServiceMonitor monitor = new ExecutorServiceMonitor();
+
+    final StubServiceEmitter serviceEmitter = new StubServiceEmitter("service", "host");
+    monitor.doMonitor(serviceEmitter);
+
+    serviceEmitter.verifyValue("segment/scan/pending", 10);
+    serviceEmitter.verifyValue("segment/scan/active", 2);
+  }
+}