diff --git a/extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java b/extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java index 58cbfb4e32e2..6deb0d74fcf1 100644 --- a/extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java +++ b/extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java @@ -28,10 +28,11 @@ import javax.annotation.Nullable; import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; -import java.util.HashSet; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicBoolean; @@ -74,16 +75,21 @@ public String apply(@Nullable final String key) // otherwise null will be replaced with empty string in nullToEmptyIfNeeded above. return null; } - - final String presentVal; - try { - presentVal = loadingCache.get(keyEquivalent, new ApplyCallable(keyEquivalent)); + // the usage of getIfPresent instead of get is to mitigate the issue that can arise + // when the data exists in druid but not in the data fetcher. + String presentVal = 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 map = new HashMap<>(); + map.put(keyEquivalent, val); + loadingCache.putAll(map); + return NullHandling.emptyToNullIfNeeded(val); } @Override @@ -131,9 +137,7 @@ public Set keySet() Iterable> data = this.dataFetcher.fetchAll(); Set set = new HashSet<>(); if (data != null) { - data.forEach(each -> { - set.add(each.getKey()); - }); + data.forEach(each -> set.add(each.getKey())); } return set; } diff --git a/extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java b/extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java index 4360b9178691..5cfc5cc44cfd 100644 --- a/extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java +++ b/extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java @@ -19,7 +19,6 @@ package org.apache.druid.server.lookup; -import com.amazonaws.transform.MapEntry; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import org.apache.druid.common.config.NullHandling; @@ -31,7 +30,11 @@ import org.junit.Test; import org.junit.rules.ExpectedException; -import java.util.*; +import java.util.AbstractMap; +import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; +import java.util.Map; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; @@ -48,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("")); @@ -74,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); @@ -101,17 +104,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 { @@ -161,9 +153,13 @@ public void testFetchAll() EasyMock.verify(dataFetcher); } - public int getIteratorSize(Iterator> it) { - int i = 0; - for ( ; it.hasNext() ; ++i ) it.next(); - return i; + public int getIteratorSize(Iterator> it) + { + int sum = 0; + while (it.hasNext()) { + sum++; + it.next(); + } + return sum; } }