Fix refresh condition on Broker, add UTs

server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClient.java

@@ -73,5 +73,5 @@ public interface CoordinatorClient
   /**
    * Retrieves list of used datasources.
    */
-  ListenableFuture<List<String>> fetchDataSources();
+  ListenableFuture<Set<String>> fetchUsedDataSources();
 }

server/src/main/java/org/apache/druid/client/coordinator/CoordinatorClientImpl.java

@@ -190,15 +190,15 @@ public CoordinatorClientImpl withRetryPolicy(ServiceRetryPolicy retryPolicy)
   }
 
   @Override
-  public ListenableFuture<List<String>> fetchDataSources()
+  public ListenableFuture<Set<String>> fetchUsedDataSources()
   {
     final String path = "/druid/coordinator/v1/metadata/datasources";
     return FutureUtils.transform(
         client.asyncRequest(
             new RequestBuilder(HttpMethod.GET, path),
             new BytesFullResponseHandler()
         ),
-        holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), new TypeReference<List<String>>() {})
+        holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), new TypeReference<Set<String>>() {})
     );
   }
 }

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

@@ -273,7 +273,7 @@ protected void cacheExecLoop()
 
               if (isServerViewInitialized &&
                   !wasRecentFailure &&
-                  refreshCondition() &&
+                  shouldRefresh() &&
                   (refreshImmediately || nextRefresh < System.currentTimeMillis())) {
                 // We need to do a refresh. Break out of the waiting loop.
                 break;
@@ -364,11 +364,10 @@ public void refreshWaitCondition() throws InterruptedException
     // noop
   }
 
-  public boolean refreshCondition()
+  @SuppressWarnings("GuardedBy")
+  public boolean shouldRefresh()
   {
-    synchronized (lock) {
-      return (!segmentsNeedingRefresh.isEmpty() || !dataSourcesNeedingRebuild.isEmpty());
-    }
+    return (!segmentsNeedingRefresh.isEmpty() || !dataSourcesNeedingRebuild.isEmpty());
   }
 
   public void awaitInitialization() throws InterruptedException

server/src/test/java/org/apache/druid/client/coordinator/NoopCoordinatorClient.java

@@ -77,7 +77,7 @@ public CoordinatorClient withRetryPolicy(ServiceRetryPolicy retryPolicy)
   }
 
   @Override
-  public ListenableFuture<List<String>> fetchDataSources()
+  public ListenableFuture<Set<String>> fetchUsedDataSources()
   {
     throw new UnsupportedOperationException();
   }

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

@@ -174,12 +174,13 @@ public void stop()
   }
 
   /**
-   * Execute refresh on the broker in each cycle if CentralizedDatasourceSchema is enabled.
+   * Execute refresh on the broker in each cycle if CentralizedDatasourceSchema is enabled
+   * else if there are segments or datasources to be refreshed.
    */
   @Override
-  public boolean refreshCondition()
+  public boolean shouldRefresh()
   {
-    return centralizedDatasourceSchemaConfig.isEnabled();
+    return centralizedDatasourceSchemaConfig.isEnabled() || super.shouldRefresh();
   }
 
   /**
@@ -206,7 +207,7 @@ public void refresh(final Set<SegmentId> segmentsToRefresh, final Set<String> da
     final Set<String> dataSourcesToQuery = new HashSet<>(segmentMetadataInfo.keySet());
 
     // this is the complete set of datasources polled from the Coordinator
-    final List<String> polledDatasources = queryDataSources();
+    final Set<String> polledDatasources = queryDataSources();
 
     dataSourcesToQuery.addAll(polledDatasources);
 
@@ -274,12 +275,15 @@ protected void removeSegmentAction(SegmentId segmentId)
     // noop, no additional action needed when segment is removed.
   }
 
-  private List<String> queryDataSources()
+  private Set<String> queryDataSources()
   {
-    List<String> dataSources = null;
+    Set<String> dataSources = new HashSet<>();
 
     try {
-      dataSources = FutureUtils.getUnchecked(coordinatorClient.fetchDataSources(), true);
+      Set<String> polled = FutureUtils.getUnchecked(coordinatorClient.fetchUsedDataSources(), true);
+      if (polled != null) {
+        dataSources.addAll(polled);
+      }
     }
     catch (Exception e) {
       log.debug(e, "Failed to query datasources from the Coordinator.");

sql/src/test/java/org/apache/druid/sql/calcite/schema/BrokerSegmentMetadataCacheTest.java

@@ -74,11 +74,13 @@
 import org.apache.druid.sql.calcite.table.DatasourceTable;
 import org.apache.druid.sql.calcite.table.DruidTable;
 import org.apache.druid.sql.calcite.util.CalciteTests;
+import org.apache.druid.sql.calcite.util.TestTimelineServerView;
 import org.apache.druid.timeline.DataSegment;
 import org.apache.druid.timeline.SegmentId;
 import org.apache.druid.timeline.partition.LinearShardSpec;
 import org.apache.druid.timeline.partition.NumberedShardSpec;
 import org.easymock.EasyMock;
+import org.joda.time.Period;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -332,6 +334,9 @@ public void testBrokerPollsAllDSSchema() throws InterruptedException
     ArgumentCaptor<Set<String>> argumentCaptor = ArgumentCaptor.forClass(Set.class);
     CoordinatorClient coordinatorClient = Mockito.mock(CoordinatorClient.class);
     Mockito.when(coordinatorClient.fetchDataSourceInformation(argumentCaptor.capture())).thenReturn(Futures.immediateFuture(null));
+
+    Set<String> datsources = Sets.newHashSet(DATASOURCE1, DATASOURCE2, DATASOURCE3, SOME_DATASOURCE, "xyz", "coldDS");
+    Mockito.when(coordinatorClient.fetchUsedDataSources()).thenReturn(Futures.immediateFuture(datsources));
     BrokerSegmentMetadataCache schema = new BrokerSegmentMetadataCache(
         CalciteTests.createMockQueryLifecycleFactory(walker, conglomerate),
         serverView,
@@ -347,15 +352,95 @@ public void testBrokerPollsAllDSSchema() throws InterruptedException
     schema.start();
     schema.awaitInitialization();
 
-    Assert.assertEquals(Sets.newHashSet(DATASOURCE1, DATASOURCE2, DATASOURCE3, SOME_DATASOURCE), argumentCaptor.getValue());
+    Assert.assertEquals(datsources, argumentCaptor.getValue());
 
     refreshLatch = new CountDownLatch(1);
     serverView.addSegment(newSegment("xyz", 0), ServerType.HISTORICAL);
 
     refreshLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS);
 
     // verify that previously refreshed are included in the last coordinator poll
-    Assert.assertEquals(Sets.newHashSet(DATASOURCE1, DATASOURCE2, DATASOURCE3, SOME_DATASOURCE, "xyz"), argumentCaptor.getValue());
+    Assert.assertEquals(datsources, argumentCaptor.getValue());
+  }
+
+  @Test
+  public void testRefreshOnEachCycleCentralizedDatasourceSchemaEnabled() throws InterruptedException
+  {
+    CentralizedDatasourceSchemaConfig config = CentralizedDatasourceSchemaConfig.create();
+    config.setEnabled(true);
+
+    serverView = new TestTimelineServerView(walker.getSegments(), Collections.emptyList());
+    druidServers = serverView.getDruidServers();
+
+    BrokerSegmentMetadataCacheConfig metadataCacheConfig = BrokerSegmentMetadataCacheConfig.create("PT1S");
+    metadataCacheConfig.setMetadataRefreshPeriod(Period.parse("PT0.001S"));
+    BrokerSegmentMetadataCache schema = new BrokerSegmentMetadataCache(
+        CalciteTests.createMockQueryLifecycleFactory(walker, conglomerate),
+        serverView,
+        SEGMENT_CACHE_CONFIG_DEFAULT,
+        new NoopEscalator(),
+        new InternalQueryConfig(),
+        new NoopServiceEmitter(),
+        new PhysicalDatasourceMetadataFactory(globalTableJoinable, segmentManager),
+        new NoopCoordinatorClient(),
+        config
+    ) {
+      @Override
+      public void refresh(Set<SegmentId> segmentsToRefresh, Set<String> dataSourcesToRebuild)
+          throws IOException
+      {
+        super.refresh(segmentsToRefresh, dataSourcesToRebuild);
+        refreshLatch.countDown();
+      }
+    };
+
+    // refresh should be executed more than once, with the feature disabled refresh should be executed only once
+    refreshLatch = new CountDownLatch(3);
+    schema.start();
+    schema.awaitInitialization();
+
+    refreshLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS);
+
+    Assert.assertEquals(0, refreshLatch.getCount());
+  }
+
+  @Test
+  public void testRefreshOnEachCycleCentralizedDatasourceSchemaDisabled() throws InterruptedException
+  {
+    BrokerSegmentMetadataCacheConfig metadataCacheConfig = BrokerSegmentMetadataCacheConfig.create("PT1S");
+    metadataCacheConfig.setMetadataRefreshPeriod(Period.parse("PT0.001S"));
+
+    serverView = new TestTimelineServerView(walker.getSegments(), Collections.emptyList());
+    druidServers = serverView.getDruidServers();
+
+    BrokerSegmentMetadataCache schema = new BrokerSegmentMetadataCache(
+        CalciteTests.createMockQueryLifecycleFactory(walker, conglomerate),
+        serverView,
+        SEGMENT_CACHE_CONFIG_DEFAULT,
+        new NoopEscalator(),
+        new InternalQueryConfig(),
+        new NoopServiceEmitter(),
+        new PhysicalDatasourceMetadataFactory(globalTableJoinable, segmentManager),
+        new NoopCoordinatorClient(),
+        CentralizedDatasourceSchemaConfig.create()
+    ) {
+      @Override
+      public void refresh(Set<SegmentId> segmentsToRefresh, Set<String> dataSourcesToRebuild)
+          throws IOException
+      {
+        super.refresh(segmentsToRefresh, dataSourcesToRebuild);
+        refreshLatch.countDown();
+      }
+    };
+
+    // refresh should be executed only once
+    refreshLatch = new CountDownLatch(3);
+    schema.start();
+    schema.awaitInitialization();
+
+    refreshLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS);
+
+    Assert.assertEquals(2, refreshLatch.getCount());
   }
 
   @Test