feat: add flag to include datasource and intervals in the result-level cache key

docs/configuration/index.md

@@ -1968,6 +1968,7 @@ You can optionally only configure caching to be enabled on the Broker by setting
 |`druid.broker.cache.populateCache`|true, false|Populate the cache on the Broker.|false|
 |`druid.broker.cache.useResultLevelCache`|true, false|Enable result level caching on the Broker.|false|
 |`druid.broker.cache.populateResultLevelCache`|true, false|Populate the result level cache on the Broker.|false|
+|`druid.broker.cache.includeDatasourceAndIntervalsInResultLevelCacheKey`|true, false|Include the query's datasource and intervals in the result level cache key. By default the key has neither, so queries that differ only by interval or datasource share one cache entry and overwrite each other. Enabling this gives each of them its own entry, which helps when the same queries repeat over a few different intervals. Trade-offs: the cache holds more entries, and a query whose interval changes on every run, such as a rolling time window, adds a new entry each run that is never read again. Queries whose intervals cover the same data but differ elsewhere, such as an end time past the latest data, no longer reuse each other's entry. Changing this setting means existing result level cache entries are no longer read.|false|
 |`druid.broker.cache.resultLevelCacheLimit`|positive integer|Maximum size of query response that can be cached.|`Integer.MAX_VALUE`|
 |`druid.broker.cache.unCacheable`|All druid query types|All query types to not cache.|`[scan]`|
 |`druid.broker.cache.cacheBulkMergeLimit`|positive integer or 0|Queries with more segments than this number will not attempt to fetch from cache at the broker level, leaving potential caching fetches (and cache result merging) to the Historicals|`Integer.MAX_VALUE`|

server/src/main/java/org/apache/druid/client/cache/CacheConfig.java

@@ -46,6 +46,14 @@ public class CacheConfig
   @JsonProperty
   private boolean populateResultLevelCache = false;
 
+  /**
+   * When true, the Broker's result-level cache key includes the query's datasource and intervals, so that queries which
+   * differ only by interval or datasource no longer share (and overwrite) one cache entry. Off by default, since
+   * changing the key invalidates existing result-level cache entries.
+   */
+  @JsonProperty
+  private boolean includeDatasourceAndIntervalsInResultLevelCacheKey = false;
+
   @JsonProperty
   @Min(0)
   private int numBackgroundThreads = 0;
@@ -87,6 +95,11 @@ public boolean isUseResultLevelCache()
     return useResultLevelCache;
   }
 
+  public boolean isIncludeDatasourceAndIntervalsInResultLevelCacheKey()
+  {
+    return includeDatasourceAndIntervalsInResultLevelCacheKey;
+  }
+
   public int getNumBackgroundThreads()
   {
     return numBackgroundThreads;

server/src/main/java/org/apache/druid/query/ResultLevelCachingQueryRunner.java

@@ -39,8 +39,10 @@
 import org.apache.druid.java.util.common.jackson.JacksonUtils;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
+import org.apache.druid.query.cache.CacheKeyBuilder;
 import org.apache.druid.query.context.ResponseContext;
 import org.apache.druid.server.QueryResource;
+import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
 import java.io.ByteArrayOutputStream;
@@ -51,6 +53,7 @@
 public class ResultLevelCachingQueryRunner<T> implements QueryRunner<T>
 {
   private static final Logger log = new Logger(ResultLevelCachingQueryRunner.class);
+  private static final byte RESULT_LEVEL_CACHE_KEY_WITH_DATASOURCE_AND_INTERVALS = 0x1;
   private final QueryRunner baseRunner;
   private ObjectMapper objectMapper;
   private final Cache cache;
@@ -95,7 +98,7 @@ public Sequence<T> run(QueryPlus queryPlus, ResponseContext responseContext)
   {
     if (useResultCache || populateResultCache) {
 
-      final byte[] queryCacheKey = strategy.computeResultLevelCacheKey(query);
+      final byte[] queryCacheKey = computeResultLevelCacheKey();
       final Cache.NamedKey cacheKey = CacheUtil.computeResultLevelCacheKey(queryCacheKey);
       final byte[] cachedResultSet = fetchResultsFromResultLevelCache(cacheKey);
       String existingResultSetId = extractEtagFromResults(cachedResultSet);
@@ -206,6 +209,29 @@ private byte[] fetchResultsFromResultLevelCache(
     return null;
   }
 
+  private byte[] computeResultLevelCacheKey()
+  {
+    final byte[] queryCacheKey = strategy.computeResultLevelCacheKey(query);
+    if (!cacheConfig.isIncludeDatasourceAndIntervalsInResultLevelCacheKey()) {
+      return queryCacheKey;
+    }
+    // The cached value still starts with the ETag, which is checked before the entry is used. Adding the datasource
+    // and intervals to the key gives queries that differ only by interval or datasource their own entries, instead of
+    // one shared entry that each of them overwrites. The byte arrays are length-prefixed so that where one ends and the
+    // next begins is unambiguous.
+    final CacheKeyBuilder builder = new CacheKeyBuilder(RESULT_LEVEL_CACHE_KEY_WITH_DATASOURCE_AND_INTERVALS)
+        .appendInt(queryCacheKey.length)
+        .appendByteArray(queryCacheKey);
+    final byte[] dataSourceCacheKey = query.getDataSource().getCacheKey();
+    if (dataSourceCacheKey != null) {
+      builder.appendInt(dataSourceCacheKey.length).appendByteArray(dataSourceCacheKey);
+    }
+    for (Interval interval : query.getIntervals()) {
+      builder.appendLong(interval.getStartMillis()).appendLong(interval.getEndMillis());
+    }
+    return builder.build();
+  }
+
   private String extractEtagFromResults(
       final byte[] cachedResult
   )

server/src/test/java/org/apache/druid/query/ResultLevelCachingQueryRunnerTest.java

@@ -37,6 +37,7 @@
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatchers;
 import org.mockito.Mockito;
 import org.mockito.stubbing.Answer;
@@ -404,6 +405,97 @@ public void testPopulateCacheThrowsException()
     emitter.verifyValue(DefaultQueryMetrics.QUERY_RESULT_CACHE_HIT, 0);
   }
 
+  @Test
+  public void testDifferentIntervalsShareOneEntryByDefault()
+  {
+    prepareCluster(10);
+    final Query<Result<TimeseriesResultValue>> query1 = timeseriesQuery(Intervals.of("2000-01-01T00/PT2H"));
+    final Query<Result<TimeseriesResultValue>> query2 = timeseriesQuery(Intervals.of("2000-01-01T01/PT2H"));
+    final CacheConfig config = newCacheConfig(true, true, DEFAULT_CACHE_ENTRY_MAX_SIZE);
+
+    createQueryRunner(config, query1).run(QueryPlus.wrap(query1), responseContext()).toList();
+    createQueryRunner(config, query2).run(QueryPlus.wrap(query2), responseContext()).toList();
+    Assertions.assertEquals(1, cache.getStats().getNumEntries());
+    emitter.flush();
+
+    createQueryRunner(config, query1).run(QueryPlus.wrap(query1), responseContext()).toList();
+
+    // query2 overwrote query1's entry, so the last run finds an entry (a hit for the cache) but can't use it.
+    Assertions.assertEquals(1, cache.getStats().getNumEntries());
+    Assertions.assertEquals(2, cache.getStats().getNumHits());
+    emitter.verifyEmitted(DefaultQueryMetrics.QUERY_RESULT_CACHE_HIT, 1);
+    emitter.verifyValue(DefaultQueryMetrics.QUERY_RESULT_CACHE_HIT, 0);
+  }
+
+  @Test
+  public void testDifferentIntervalsUseSeparateEntriesWhenFlagEnabled()
+  {
+    prepareCluster(10);
+    final Query<Result<TimeseriesResultValue>> query1 = timeseriesQuery(Intervals.of("2000-01-01T00/PT2H"));
+    final Query<Result<TimeseriesResultValue>> query2 = timeseriesQuery(Intervals.of("2000-01-01T01/PT2H"));
+    final CacheConfig config = newCacheConfig(true, true, DEFAULT_CACHE_ENTRY_MAX_SIZE, true);
+
+    final List<Result<TimeseriesResultValue>> results1 =
+        createQueryRunner(config, query1).run(QueryPlus.wrap(query1), responseContext()).toList();
+    final List<Result<TimeseriesResultValue>> results2 =
+        createQueryRunner(config, query2).run(QueryPlus.wrap(query2), responseContext()).toList();
+    Assertions.assertEquals(2, cache.getStats().getNumEntries());
+    emitter.flush();
+
+    Assertions.assertEquals(
+        results1,
+        createQueryRunner(config, query1).run(QueryPlus.wrap(query1), responseContext()).toList()
+    );
+    emitter.verifyEmitted(DefaultQueryMetrics.QUERY_RESULT_CACHE_HIT, 1);
+    emitter.verifyValue(DefaultQueryMetrics.QUERY_RESULT_CACHE_HIT, 1);
+    emitter.flush();
+
+    Assertions.assertEquals(
+        results2,
+        createQueryRunner(config, query2).run(QueryPlus.wrap(query2), responseContext()).toList()
+    );
+    emitter.verifyEmitted(DefaultQueryMetrics.QUERY_RESULT_CACHE_HIT, 1);
+    emitter.verifyValue(DefaultQueryMetrics.QUERY_RESULT_CACHE_HIT, 1);
+    Assertions.assertEquals(2, cache.getStats().getNumEntries());
+  }
+
+  @Test
+  public void testDatasourceAndIntervalsInCacheKeyOnlyWhenFlagEnabled()
+  {
+    cache = Mockito.spy(cache);
+    prepareCluster(10);
+    final Query<Result<TimeseriesResultValue>> query = timeseriesQuery(Intervals.of("2000-01-01T00/PT2H"));
+    final Query<Result<TimeseriesResultValue>> sameQuery = timeseriesQuery(Intervals.of("2000-01-01T00/PT2H"));
+    final Query<Result<TimeseriesResultValue>> otherInterval = timeseriesQuery(Intervals.of("2000-01-01T01/PT2H"));
+    final Query<Result<TimeseriesResultValue>> otherDatasource = query.withDataSource(TableDataSource.create("other"));
+
+    final Cache.NamedKey key = resultLevelCacheKey(query, false);
+    Assertions.assertEquals(key, resultLevelCacheKey(otherInterval, false));
+    Assertions.assertEquals(key, resultLevelCacheKey(otherDatasource, false));
+
+    final Cache.NamedKey keyWithFlag = resultLevelCacheKey(query, true);
+    Assertions.assertNotEquals(key, keyWithFlag);
+    Assertions.assertEquals(keyWithFlag, resultLevelCacheKey(sameQuery, true));
+    Assertions.assertNotEquals(keyWithFlag, resultLevelCacheKey(otherInterval, true));
+    Assertions.assertNotEquals(keyWithFlag, resultLevelCacheKey(otherDatasource, true));
+  }
+
+  /**
+   * Runs the query with the result-level cache in use-only mode and returns the key it looked up. Requires
+   * {@link #cache} to be a Mockito spy.
+   */
+  private <T> Cache.NamedKey resultLevelCacheKey(Query<T> query, boolean includeDatasourceAndIntervals)
+  {
+    Mockito.clearInvocations(cache);
+    createQueryRunner(
+        newCacheConfig(false, true, DEFAULT_CACHE_ENTRY_MAX_SIZE, includeDatasourceAndIntervals),
+        query
+    ).run(QueryPlus.wrap(query), responseContext()).toList();
+    final ArgumentCaptor<Cache.NamedKey> captor = ArgumentCaptor.forClass(Cache.NamedKey.class);
+    Mockito.verify(cache).get(captor.capture());
+    return captor.getValue();
+  }
+
   private <T> ResultLevelCachingQueryRunner<T> createQueryRunner(
       CacheConfig cacheConfig,
       Query<T> query
@@ -431,9 +523,25 @@ private CacheConfig newCacheConfig(
       boolean useResultLevelCache,
       int resultLevelCacheLimit
   )
+  {
+    return newCacheConfig(populateResultLevelCache, useResultLevelCache, resultLevelCacheLimit, false);
+  }
+
+  private CacheConfig newCacheConfig(
+      boolean populateResultLevelCache,
+      boolean useResultLevelCache,
+      int resultLevelCacheLimit,
+      boolean includeDatasourceAndIntervals
+  )
   {
     return new CacheConfig()
     {
+      @Override
+      public boolean isIncludeDatasourceAndIntervalsInResultLevelCacheKey()
+      {
+        return includeDatasourceAndIntervals;
+      }
+
       @Override
       public boolean isPopulateResultLevelCache()
       {