Introduce published segment cache in broker

benchmarks/src/main/java/org/apache/druid/benchmark/query/SqlBenchmark.java

@@ -114,7 +114,7 @@ public void setup()
     final QueryRunnerFactoryConglomerate conglomerate = conglomerateCloserPair.lhs;
     final PlannerConfig plannerConfig = new PlannerConfig();
     final DruidSchema druidSchema = CalciteTests.createMockSchema(conglomerate, walker, plannerConfig);
-    final SystemSchema systemSchema = CalciteTests.createMockSystemSchema(druidSchema, walker);
+    final SystemSchema systemSchema = CalciteTests.createMockSystemSchema(druidSchema, walker, plannerConfig);
     this.walker = new SpecificSegmentsQuerySegmentWalker(conglomerate).add(dataSegment, index);
     final PlannerFactory plannerFactory = new PlannerFactory(
         druidSchema,

docs/content/querying/sql.md

@@ -698,6 +698,8 @@ The Druid SQL server is configured through the following properties on the Broke
 |`druid.sql.planner.useFallback`|Whether to evaluate operations on the Broker when they cannot be expressed as Druid queries. This option is not recommended for production since it can generate unscalable query plans. If false, SQL queries that cannot be translated to Druid queries will fail.|false|
 |`druid.sql.planner.requireTimeCondition`|Whether to require SQL to have filter conditions on __time column so that all generated native queries will have user specified intervals. If true, all queries wihout filter condition on __time column will fail|false|
 |`druid.sql.planner.sqlTimeZone`|Sets the default time zone for the server, which will affect how time functions and timestamp literals behave. Should be a time zone name like "America/Los_Angeles" or offset like "-08:00".|UTC|
+|`druid.sql.planner.metadataSegmentCacheEnable`|Whether to keep a cache of published segments in broker. If true, broker polls coordinator in background to get segments from metadata store and maintains a local cache. If false, coordinator's REST api will be invoked when broker needs published segments info.|false|
+|`druid.sql.planner.metadataSegmentPollPeriod`|How often to poll coordinator for published segments list if `druid.sql.planner.metadataSegmentCacheEnable` is set to true. Poll period is in milliseconds. |60000|
 
 ## SQL Metrics
 

extensions-core/druid-bloom-filter/src/test/java/org/apache/druid/query/aggregation/bloom/sql/BloomFilterSqlAggregatorTest.java

@@ -192,7 +192,7 @@ public void setUp() throws Exception
 
     final PlannerConfig plannerConfig = new PlannerConfig();
     final DruidSchema druidSchema = CalciteTests.createMockSchema(conglomerate, walker, plannerConfig);
-    final SystemSchema systemSchema = CalciteTests.createMockSystemSchema(druidSchema, walker);
+    final SystemSchema systemSchema = CalciteTests.createMockSystemSchema(druidSchema, walker, plannerConfig);
     final DruidOperatorTable operatorTable = new DruidOperatorTable(
         ImmutableSet.of(new BloomFilterSqlAggregator()),
         ImmutableSet.of()

extensions-core/histogram/src/test/java/org/apache/druid/query/aggregation/histogram/sql/QuantileSqlAggregatorTest.java

@@ -157,7 +157,7 @@ public void setUp() throws Exception
 
     final PlannerConfig plannerConfig = new PlannerConfig();
     final DruidSchema druidSchema = CalciteTests.createMockSchema(conglomerate, walker, plannerConfig);
-    final SystemSchema systemSchema = CalciteTests.createMockSystemSchema(druidSchema, walker);
+    final SystemSchema systemSchema = CalciteTests.createMockSystemSchema(druidSchema, walker, plannerConfig);
     final DruidOperatorTable operatorTable = new DruidOperatorTable(
         ImmutableSet.of(new QuantileSqlAggregator()),
         ImmutableSet.of()

server/src/main/java/org/apache/druid/client/DataSegmentInterner.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.client;
+
+import com.google.common.collect.Interner;
+import com.google.common.collect.Interners;
+import org.apache.druid.timeline.DataSegment;
+
+/**
+ * Interns the DataSegment object in order to share the reference for same DataSegment.
+ * It uses two separate interners for realtime and historical segments to prevent
+ * overwriting the size of a segment which was served by a historical and later served
+ * by another realtime server, since realtime server always publishes with size 0.
+ */
+public class DataSegmentInterner
+{
+  private static final Interner<DataSegment> REALTIME_INTERNER = Interners.newWeakInterner();
+  private static final Interner<DataSegment> HISTORICAL_INTERNER = Interners.newWeakInterner();
+
+  private DataSegmentInterner()
+  {
+    //No instantiation
+  }
+
+  public static DataSegment intern(DataSegment segment)
+  {
+    // A segment learns it's size and dimensions when it moves from a relatime to historical server
+    // for that reason, we are using it's size as the indicator to decide whether to use REALTIME or
+    // HISTORICAL interner.
+    return segment.getSize() > 0 ? HISTORICAL_INTERNER.intern(segment) : REALTIME_INTERNER.intern(segment);
+  }
+}

server/src/main/java/org/apache/druid/client/selector/ServerSelector.java

@@ -20,6 +20,7 @@
 package org.apache.druid.client.selector;
 
 import it.unimi.dsi.fastutil.ints.Int2ObjectRBTreeMap;
+import org.apache.druid.client.DataSegmentInterner;
 import org.apache.druid.server.coordination.DruidServerMetadata;
 import org.apache.druid.server.coordination.ServerType;
 import org.apache.druid.timeline.DataSegment;
@@ -50,7 +51,7 @@ public ServerSelector(
       TierSelectorStrategy strategy
   )
   {
-    this.segment = new AtomicReference<>(segment);
+    this.segment = new AtomicReference<>(DataSegmentInterner.intern(segment));
     this.strategy = strategy;
     this.historicalServers = new Int2ObjectRBTreeMap<>(strategy.getComparator());
     this.realtimeServers = new Int2ObjectRBTreeMap<>(strategy.getComparator());

server/src/main/java/org/apache/druid/server/http/MetadataResource.java

@@ -57,6 +57,7 @@
 import java.util.List;
 import java.util.Set;
 import java.util.TreeSet;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 /**
@@ -148,14 +149,22 @@ public Response getDatabaseSegmentDataSource(@PathParam("dataSourceName") final
   @GET
   @Path("/segments")
   @Produces(MediaType.APPLICATION_JSON)
-  public Response getDatabaseSegments(@Context final HttpServletRequest req)
+  public Response getDatabaseSegments(
+      @Context final HttpServletRequest req,
+      @QueryParam("datasources") final Set<String> datasources
+  )
   {
-    final Collection<ImmutableDruidDataSource> druidDataSources = metadataSegmentManager.getDataSources();
+    Collection<ImmutableDruidDataSource> druidDataSources = metadataSegmentManager.getDataSources();
+    if (datasources != null && !datasources.isEmpty()) {
+      druidDataSources = druidDataSources.stream()
+                                         .filter(src -> datasources.contains(src.getName()))
+                                         .collect(Collectors.toSet());
+    }
     final Stream<DataSegment> metadataSegments = druidDataSources
         .stream()
         .flatMap(t -> t.getSegments().stream());
 
-    Function<DataSegment, Iterable<ResourceAction>> raGenerator = segment -> Collections.singletonList(
+    final Function<DataSegment, Iterable<ResourceAction>> raGenerator = segment -> Collections.singletonList(
         AuthorizationUtils.DATASOURCE_READ_RA_GENERATOR.apply(segment.getDataSource()));
 
     final Iterable<DataSegment> authorizedSegments =

server/src/test/java/org/apache/druid/client/CachingClusteredClientTest.java

@@ -2213,13 +2213,13 @@ private List<Map<DruidServer, ServerExpectations>> populateTimeline(
             expectedResults.get(k).get(j)
         );
         serverExpectations.get(lastServer).addExpectation(expectation);
-
+        EasyMock.expect(mockSegment.getSize()).andReturn(0L).anyTimes();
+        EasyMock.replay(mockSegment);
         ServerSelector selector = new ServerSelector(
             expectation.getSegment(),
             new HighestPriorityTierSelectorStrategy(new RandomServerSelectorStrategy())
         );
         selector.addServerAndUpdateSegment(new QueryableDruidServer(lastServer, null), selector.getSegment());
-
         final ShardSpec shardSpec;
         if (numChunks == 1) {
           shardSpec = new SingleDimensionShardSpec("dimAll", null, null, 0);
@@ -2234,6 +2234,7 @@ private List<Map<DruidServer, ServerExpectations>> populateTimeline(
           }
           shardSpec = new SingleDimensionShardSpec("dim" + k, start, end, j);
         }
+        EasyMock.reset(mockSegment);
         EasyMock.expect(mockSegment.getShardSpec())
                 .andReturn(shardSpec)
                 .anyTimes();

sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerConfig.java

@@ -67,6 +67,21 @@ public class PlannerConfig
   private DateTimeZone sqlTimeZone = DateTimeZone.UTC;
 
   @JsonProperty
+  private boolean metadataSegmentCacheEnable = false;
+
+  @JsonProperty
+  private long metadataSegmentPollPeriod = 60000;
+
+  public long getMetadataSegmentPollPeriod()
+  {
+    return metadataSegmentPollPeriod;
+  }
+
+  public boolean isMetadataSegmentCacheEnable()
+  {
+    return metadataSegmentCacheEnable;
+  }
+
   private boolean serializeComplexValues = true;
 
   public Period getMetadataRefreshPeriod()
@@ -159,6 +174,8 @@ public PlannerConfig withOverrides(final Map<String, Object> context)
     newConfig.requireTimeCondition = isRequireTimeCondition();
     newConfig.sqlTimeZone = getSqlTimeZone();
     newConfig.awaitInitializationOnStart = isAwaitInitializationOnStart();
+    newConfig.metadataSegmentCacheEnable = isMetadataSegmentCacheEnable();
+    newConfig.metadataSegmentPollPeriod = getMetadataSegmentPollPeriod();
     newConfig.serializeComplexValues = shouldSerializeComplexValues();
     return newConfig;
   }
@@ -200,6 +217,8 @@ public boolean equals(final Object o)
            useFallback == that.useFallback &&
            requireTimeCondition == that.requireTimeCondition &&
            awaitInitializationOnStart == that.awaitInitializationOnStart &&
+           metadataSegmentCacheEnable == that.metadataSegmentCacheEnable &&
+           metadataSegmentPollPeriod == that.metadataSegmentPollPeriod &&
            serializeComplexValues == that.serializeComplexValues &&
            Objects.equals(metadataRefreshPeriod, that.metadataRefreshPeriod) &&
            Objects.equals(sqlTimeZone, that.sqlTimeZone);
@@ -221,6 +240,8 @@ public int hashCode()
         requireTimeCondition,
         awaitInitializationOnStart,
         sqlTimeZone,
+        metadataSegmentCacheEnable,
+        metadataSegmentPollPeriod,
         serializeComplexValues
     );
   }
@@ -239,6 +260,8 @@ public String toString()
            ", useFallback=" + useFallback +
            ", requireTimeCondition=" + requireTimeCondition +
            ", awaitInitializationOnStart=" + awaitInitializationOnStart +
+           ", metadataSegmentCacheEnable=" + metadataSegmentCacheEnable +
+           ", metadataSegmentPollPeriod=" + metadataSegmentPollPeriod +
            ", sqlTimeZone=" + sqlTimeZone +
            ", serializeComplexValues=" + serializeComplexValues +
            '}';