From 9e1c412af3931114e731ee18bf6c24f2bb1237ed Mon Sep 17 00:00:00 2001 From: TessaIO Date: Tue, 1 Oct 2024 16:46:54 +0200 Subject: [PATCH 1/2] fix loading lookup extension Signed-off-by: TessaIO --- .../druid/server/lookup/LoadingLookup.java | 44 ++++++------- .../server/lookup/LoadingLookupTest.java | 62 +++++++++++++------ 2 files changed, 62 insertions(+), 44 deletions(-) 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 508b07b9330e..9712c628e504 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,8 +28,10 @@ import javax.annotation.Nullable; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicBoolean; @@ -73,15 +75,22 @@ public String apply(@Nullable final String key) return null; } - final String presentVal; - try { - presentVal = loadingCache.get(keyEquivalent, new ApplyCallable(keyEquivalent)); + final String 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 map = new HashMap<>(); + map.put(keyEquivalent, val); + + this.loadingCache.putAll(map); + + return NullHandling.emptyToNullIfNeeded(val); } @Override @@ -108,13 +117,16 @@ public List unapply(@Nullable final String value) @Override public boolean supportsAsMap() { - return false; + return true; } @Override public Map asMap() { - throw new UnsupportedOperationException("Cannot get map view"); + final Map map = new HashMap<>(); + Optional.ofNullable(this.dataFetcher.fetchAll()) + .ifPresent(data -> data.forEach(entry -> map.put(entry.getKey(), entry.getValue()))); + return map; } @Override @@ -123,24 +135,6 @@ public byte[] getCacheKey() return LookupExtractionModule.getRandomCacheKey(); } - private class ApplyCallable implements Callable - { - private final String key; - - public ApplyCallable(String key) - { - this.key = key; - } - - @Override - public String call() - { - // When SQL compatible null handling is disabled, - // avoid returning null and return an empty string to cache it. - return NullHandling.nullToEmptyIfNeeded(dataFetcher.fetch(key)); - } - } - public synchronized void close() { if (isOpen.getAndSet(false)) { 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 43588bf8c476..0d2f217be7dd 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 @@ -31,15 +31,18 @@ import org.junit.Test; import org.junit.rules.ExpectedException; +import java.util.AbstractMap; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; +import java.util.Map; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; public class LoadingLookupTest extends InitializedNullHandlingTest { DataFetcher dataFetcher = EasyMock.createMock(DataFetcher.class); - LoadingCache lookupCache = EasyMock.createStrictMock(LoadingCache.class); + LoadingCache lookupCache = EasyMock.createMock(LoadingCache.class); LoadingCache reverseLookupCache = EasyMock.createStrictMock(LoadingCache.class); LoadingLookup loadingLookup = new LoadingLookup(dataFetcher, lookupCache, reverseLookupCache); @@ -47,9 +50,9 @@ public class LoadingLookupTest extends InitializedNullHandlingTest public ExpectedException expectedException = ExpectedException.none(); @Test - public void testApplyEmptyOrNull() throws ExecutionException + public void testApplyEmptyOrNull() { - 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("")); @@ -73,14 +76,43 @@ public void testUnapplyNull() } @Test - public void testApply() throws ExecutionException + public void testApply() { - 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); } + @Test + public void testApplyWithNullValue() + { + EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("key"))).andReturn(null).once(); + EasyMock.expect(dataFetcher.fetch("key")).andReturn(null).once(); + EasyMock.replay(lookupCache, dataFetcher); + Assert.assertNull(loadingLookup.apply("key")); + EasyMock.verify(lookupCache, dataFetcher); + } + + @Test + public void testApplyWithCacheCheck() + { + // This test that we have a cache miss that will trigger a data fetch call + // and then the cache gets populated, so in the second apply call we get the + // value directly from the cache + Map map = new HashMap<>(); + map.put("key", "value"); + EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("key"))).andReturn(null).once(); + EasyMock.expect(dataFetcher.fetch("key")).andReturn("value").once(); + lookupCache.putAll(map); + EasyMock.expectLastCall().andVoid(); + EasyMock.expect(lookupCache.getIfPresent("key")).andReturn("value").once(); + EasyMock.replay(lookupCache, dataFetcher); + Assert.assertEquals(loadingLookup.apply("key"), "value"); + Assert.assertEquals(loadingLookup.apply("key"), "value"); + EasyMock.verify(lookupCache, dataFetcher); + } + @Test public void testUnapplyAll() throws ExecutionException { @@ -105,17 +137,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 +157,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 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); } } From 875e66c99a1731c3f8bd20236616c5f76e99c61d Mon Sep 17 00:00:00 2001 From: TessaIO Date: Tue, 15 Oct 2024 17:42:03 +0200 Subject: [PATCH 2/2] fix code nits and cover more test cases Signed-off-by: TessaIO --- .../apache/druid/server/lookup/LoadingLookup.java | 7 ++----- .../druid/server/lookup/LoadingLookupTest.java | 15 +++++++-------- 2 files changed, 9 insertions(+), 13 deletions(-) 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 9712c628e504..edd19ca415ad 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 @@ -80,15 +80,12 @@ public String apply(@Nullable final String key) return NullHandling.emptyToNullIfNeeded(presentVal); } - String val = this.dataFetcher.fetch(keyEquivalent); + final String val = this.dataFetcher.fetch(keyEquivalent); if (val == null) { return null; } - Map map = new HashMap<>(); - map.put(keyEquivalent, val); - - this.loadingCache.putAll(map); + this.loadingCache.putAll(Collections.singletonMap(keyEquivalent, val)); return NullHandling.emptyToNullIfNeeded(val); } 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 0d2f217be7dd..180a9338bf26 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 @@ -31,7 +31,6 @@ import org.junit.Test; import org.junit.rules.ExpectedException; -import java.util.AbstractMap; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; @@ -95,11 +94,8 @@ public void testApplyWithNullValue() } @Test - public void testApplyWithCacheCheck() + public void testApplyTriggersCacheMissAndSubsequentCacheHit() { - // This test that we have a cache miss that will trigger a data fetch call - // and then the cache gets populated, so in the second apply call we get the - // value directly from the cache Map map = new HashMap<>(); map.put("key", "value"); EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("key"))).andReturn(null).once(); @@ -163,10 +159,13 @@ public void testSupportsAsMap() @Test public void testAsMap() { - Map.Entry fetchedData = new AbstractMap.SimpleEntry<>("dummy", "test"); - EasyMock.expect(dataFetcher.fetchAll()).andReturn(Collections.singletonList(fetchedData)); + final Map fetchedData = new HashMap<>(); + fetchedData.put("dummy", "test"); + fetchedData.put("key", null); + fetchedData.put(null, "value"); + EasyMock.expect(dataFetcher.fetchAll()).andReturn(fetchedData.entrySet()); EasyMock.replay(dataFetcher); - Assert.assertEquals(loadingLookup.asMap().size(), 1); + Assert.assertEquals(loadingLookup.asMap(), fetchedData); EasyMock.verify(dataFetcher); } }