minor code changes, add test

server/src/main/java/org/apache/druid/segment/metadata/AbstractSegmentMetadataCache.java

@@ -364,8 +364,11 @@ public void refreshWaitCondition() throws InterruptedException
     // noop
   }
 
+  /**
+   * Refresh is executed only when there are segments or datasources needing refresh.
+   */
   @SuppressWarnings("GuardedBy")
-  public boolean shouldRefresh()
+  protected boolean shouldRefresh()
   {
     return (!segmentsNeedingRefresh.isEmpty() || !dataSourcesNeedingRebuild.isEmpty());
   }

server/src/main/java/org/apache/druid/segment/metadata/CoordinatorSegmentMetadataCache.java

@@ -65,7 +65,6 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -100,11 +99,11 @@ public class CoordinatorSegmentMetadataCache extends AbstractSegmentMetadataCach
   private final SegmentReplicationStatusManager segmentReplicationStatusManager;
 
   // Schema for datasources from cold segments
-  private final ConcurrentMap<String, DataSourceInformation> coldDatasourceSchema = new ConcurrentHashMap<>();
-  private final long coldDatasourceSchemaExecDurationMillis;
-  private final ScheduledExecutorService coldDSScehmaExec;
+  private final ConcurrentHashMap<String, DataSourceInformation> coldSchemaTable = new ConcurrentHashMap<>();
+  private final long coldSchemaExecPeriodMillis;
+  private final ScheduledExecutorService coldScehmaExec;
   private @Nullable Future<?> cacheExecFuture = null;
-  private @Nullable Future<?> coldDSSchemaExecFuture = null;
+  private @Nullable Future<?> coldSchemaExecFuture = null;
 
   @Inject
   public CoordinatorSegmentMetadataCache(
@@ -128,11 +127,11 @@ public CoordinatorSegmentMetadataCache(
     this.segmentSchemaBackfillQueue = segmentSchemaBackfillQueue;
     this.sqlSegmentsMetadataManager = sqlSegmentsMetadataManager;
     this.segmentReplicationStatusManager = segmentReplicationStatusManager;
-    this.coldDatasourceSchemaExecDurationMillis =
+    this.coldSchemaExecPeriodMillis =
         segmentsMetadataManagerConfigSupplier.get().getPollDuration().getMillis();
-    coldDSScehmaExec = Executors.newSingleThreadScheduledExecutor(
+    coldScehmaExec = Executors.newSingleThreadScheduledExecutor(
         new ThreadFactoryBuilder()
-            .setNameFormat("DruidColdDSSchema-ScheduledExecutor-%d")
+            .setNameFormat("DruidColdSchema-ScheduledExecutor-%d")
             .setDaemon(true)
             .build()
     );
@@ -204,12 +203,15 @@ public void stop()
   {
     callbackExec.shutdownNow();
     cacheExec.shutdownNow();
-    coldDSScehmaExec.shutdownNow();
+    coldScehmaExec.shutdownNow();
     segmentSchemaCache.onLeaderStop();
     segmentSchemaBackfillQueue.onLeaderStop();
     if (cacheExecFuture != null) {
       cacheExecFuture.cancel(true);
     }
+    if (coldSchemaExecFuture != null) {
+      coldSchemaExecFuture.cancel(true);
+    }
   }
 
   public void onLeaderStart()
@@ -218,9 +220,9 @@ public void onLeaderStart()
     try {
       segmentSchemaBackfillQueue.onLeaderStart();
       cacheExecFuture = cacheExec.submit(this::cacheExecLoop);
-      coldDSSchemaExecFuture = coldDSScehmaExec.schedule(
+      coldSchemaExecFuture = coldScehmaExec.schedule(
           this::coldDatasourceSchemaExec,
-          coldDatasourceSchemaExecDurationMillis,
+          coldSchemaExecPeriodMillis,
           TimeUnit.MILLISECONDS
       );
 
@@ -239,8 +241,8 @@ public void onLeaderStop()
     if (cacheExecFuture != null) {
       cacheExecFuture.cancel(true);
     }
-    if (coldDSSchemaExecFuture != null) {
-      coldDSSchemaExecFuture.cancel(true);
+    if (coldSchemaExecFuture != null) {
+      coldSchemaExecFuture.cancel(true);
     }
     segmentSchemaCache.onLeaderStop();
     segmentSchemaBackfillQueue.onLeaderStop();
@@ -387,17 +389,19 @@ public DataSourceInformation getDatasource(String name)
   {
     DataSourceInformation dataSourceInformation = tables.get(name);
     if (dataSourceInformation != null) {
+      // implies that the datasource is entirely cold
       return dataSourceInformation;
     }
-    return coldDatasourceSchema.get(name);
+    return coldSchemaTable.get(name);
   }
 
   @Override
   public Map<String, DataSourceInformation> getDataSourceInformationMap()
   {
     Map<String, DataSourceInformation> copy = new HashMap<>(tables);
 
-    for (Map.Entry<String, DataSourceInformation> entry : coldDatasourceSchema.entrySet()) {
+    for (Map.Entry<String, DataSourceInformation> entry : coldSchemaTable.entrySet()) {
+      // add entirely cold datasource schema
       copy.computeIfAbsent(entry.getKey(), value -> entry.getValue());
     }
     return ImmutableMap.copyOf(copy);
@@ -452,7 +456,7 @@ public void refresh(final Set<SegmentId> segmentsToRefresh, final Set<String> da
 
       RowSignature mergedSignature = rowSignature;
 
-      DataSourceInformation coldDatasourceInformation = coldDatasourceSchema.get(dataSource);
+      DataSourceInformation coldDatasourceInformation = coldSchemaTable.get(dataSource);
       if (coldDatasourceInformation != null) {
         if (rowSignature == null) {
           mergedSignature = coldDatasourceInformation.getRowSignature();
@@ -498,6 +502,7 @@ private Set<SegmentId> filterSegmentWithCachedSchema(Set<SegmentId> segmentIds)
     return cachedSegments;
   }
 
+  @VisibleForTesting
   protected void coldDatasourceSchemaExec()
   {
     Collection<ImmutableDruidDataSource> immutableDataSources =
@@ -513,8 +518,8 @@ protected void coldDatasourceSchemaExec()
       Collection<DataSegment> dataSegments = dataSource.getSegments();
 
       for (DataSegment segment : dataSegments) {
-        if (segmentReplicationStatusManager.getReplicationFactor(segment.getId()) != null
-            && segmentReplicationStatusManager.getReplicationFactor(segment.getId()) != 0) {
+        Integer replicationFactor = segmentReplicationStatusManager.getReplicationFactor(segment.getId());
+        if (replicationFactor != null && replicationFactor != 0) {
           continue;
         }
         Optional<SchemaPayloadPlus> optionalSchema = segmentSchemaCache.getSchemaForSegment(segment.getId());
@@ -537,9 +542,10 @@ protected void coldDatasourceSchemaExec()
 
       log.debug("[%s] signature from cold segments is [%s]", dataSourceName, coldSignature);
 
-      coldDatasourceSchema.put(dataSourceName, new DataSourceInformation(dataSourceName, coldSignature));
+      coldSchemaTable.put(dataSourceName, new DataSourceInformation(dataSourceName, coldSignature));
 
-      // update tables map with merged schema
+      // update tables map with merged schema, if signature doesn't exist we do not add entry in this table
+      // schema for entirely cold datasource is maintained separately
       tables.computeIfPresent(
           dataSourceName,
           (ds, info) -> {
@@ -560,14 +566,15 @@ protected void coldDatasourceSchemaExec()
     }
 
     // remove any stale datasource from the map
-    coldDatasourceSchema.keySet().retainAll(dataSources);
+    coldSchemaTable.keySet().retainAll(dataSources);
   }
 
   private RowSignature mergeHotAndColdSchema(RowSignature hot, RowSignature cold)
   {
     final Map<String, ColumnType> columnTypes = new LinkedHashMap<>();
 
     List<RowSignature> signatures = new ArrayList<>();
+    // hot datasource schema takes precedence
     signatures.add(hot);
     signatures.add(cold);
 

server/src/test/java/org/apache/druid/segment/metadata/CoordinatorSegmentMetadataCacheTest.java

@@ -1820,6 +1820,7 @@ private CoordinatorSegmentMetadataCache setupForColdDatasourceSchemaTest()
     schemaPayloadMap.put(
         "fp",
         new SchemaPayload(RowSignature.builder()
+                                      .add("dim1", ColumnType.STRING)
                                       .add("c1", ColumnType.STRING)
                                       .add("c2", ColumnType.LONG)
                                       .build())
@@ -1901,6 +1902,89 @@ public void testColdDatasourceSchema_mergeSchemaInColdSchemaExec() throws IOExce
     Assert.assertEquals(ColumnType.LONG, rowSignature.getColumnType(columnNames.get(7)).get());
   }
 
+  @Test
+  public void testColdDatasourceSchema_verifyStaleDatasourceRemoved()
+  {
+    DataSegment coldSegmentAlpha =
+        DataSegment.builder()
+                   .dataSource("alpha")
+                   .interval(Intervals.of("2000/P2Y"))
+                   .version("1")
+                   .shardSpec(new LinearShardSpec(0))
+                   .size(0)
+                   .build();
+
+    DataSegment coldSegmentBeta =
+        DataSegment.builder()
+                   .dataSource("beta")
+                   .interval(Intervals.of("2000/P2Y"))
+                   .version("1")
+                   .shardSpec(new LinearShardSpec(0))
+                   .size(0)
+                   .build();
+
+    ImmutableMap.Builder<SegmentId, SegmentMetadata> segmentStatsMap = new ImmutableMap.Builder<>();
+    segmentStatsMap.put(coldSegmentAlpha.getId(), new SegmentMetadata(20L, "fp"));
+    segmentStatsMap.put(coldSegmentBeta.getId(), new SegmentMetadata(20L, "fp"));
+
+    ImmutableMap.Builder<String, SchemaPayload> schemaPayloadMap = new ImmutableMap.Builder<>();
+    schemaPayloadMap.put(
+        "fp",
+        new SchemaPayload(RowSignature.builder()
+                                      .add("dim1", ColumnType.STRING)
+                                      .add("c1", ColumnType.STRING)
+                                      .add("c2", ColumnType.LONG)
+                                      .build())
+    );
+    segmentSchemaCache.updateFinalizedSegmentSchema(
+        new SegmentSchemaCache.FinalizedSegmentSchemaInfo(segmentStatsMap.build(), schemaPayloadMap.build())
+    );
+
+    Mockito.when(segmentReplicationStatusManager.getReplicationFactor(ArgumentMatchers.eq(coldSegmentAlpha.getId()))).thenReturn(0);
+    Mockito.when(segmentReplicationStatusManager.getReplicationFactor(ArgumentMatchers.eq(coldSegmentBeta.getId()))).thenReturn(0);
+
+    ImmutableDruidDataSource druidDataSource =
+        new ImmutableDruidDataSource(
+            coldSegmentAlpha.getDataSource(),
+            Collections.emptyMap(),
+            Collections.singletonMap(coldSegmentAlpha.getId(), coldSegmentAlpha)
+        );
+
+    Mockito.when(sqlSegmentsMetadataManager.getImmutableDataSourcesWithAllUsedSegments())
+           .thenReturn(Collections.singletonList(druidDataSource));
+
+    CoordinatorSegmentMetadataCache schema = new CoordinatorSegmentMetadataCache(
+        getQueryLifecycleFactory(walker),
+        serverView,
+        SEGMENT_CACHE_CONFIG_DEFAULT,
+        new NoopEscalator(),
+        new InternalQueryConfig(),
+        new NoopServiceEmitter(),
+        segmentSchemaCache,
+        backFillQueue,
+        sqlSegmentsMetadataManager,
+        segmentReplicationStatusManager,
+        segmentsMetadataManagerConfigSupplier
+    );
+
+    schema.coldDatasourceSchemaExec();
+    Assert.assertNotNull(schema.getDatasource("alpha"));
+
+    druidDataSource =
+        new ImmutableDruidDataSource(
+            coldSegmentBeta.getDataSource(),
+            Collections.emptyMap(),
+            Collections.singletonMap(coldSegmentBeta.getId(), coldSegmentBeta)
+        );
+
+    Mockito.when(sqlSegmentsMetadataManager.getImmutableDataSourcesWithAllUsedSegments())
+           .thenReturn(Collections.singletonList(druidDataSource));
+
+    schema.coldDatasourceSchemaExec();
+    Assert.assertNotNull(schema.getDatasource("beta"));
+    Assert.assertNull(schema.getDatasource("alpha"));
+  }
+
   private void verifyFooDSSchema(CoordinatorSegmentMetadataCache schema, int columns)
   {
     final DataSourceInformation fooDs = schema.getDatasource("foo");

sql/src/main/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCache.java

@@ -178,7 +178,7 @@ public void stop()
    * else if there are segments or datasources to be refreshed.
    */
   @Override
-  public boolean shouldRefresh()
+  protected boolean shouldRefresh()
   {
     return centralizedDatasourceSchemaConfig.isEnabled() || super.shouldRefresh();
   }