Merge remote-tracking branch 'upstream/master' into cold_ds_schema

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

@@ -60,12 +60,18 @@ public interface CoordinatorClient
   ListenableFuture<List<DataSourceInformation>> fetchDataSourceInformation(Set<String> datasources);
 
   /**
-   * Retrieves list of used datasources.
+   * Fetch bootstrap segments from the coordinator. The results must be streamed back to the caller as the
+   * result set can be large.
    */
-  ListenableFuture<List<String>> fetchDataSources();
+  ListenableFuture<BootstrapSegmentsResponse> fetchBootstrapSegments();
 
   /**
    * Returns a new instance backed by a ServiceClient which follows the provided retryPolicy
    */
   CoordinatorClient withRetryPolicy(ServiceRetryPolicy retryPolicy);
+
+  /**
+   * Retrieves list of used datasources.
+   */
+  ListenableFuture<List<String>> fetchDataSources();
 }

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

@@ -162,15 +162,24 @@ public ListenableFuture<List<DataSourceInformation>> fetchDataSourceInformation(
   }
 
   @Override
-  public ListenableFuture<List<String>> fetchDataSources()
+  public ListenableFuture<BootstrapSegmentsResponse> fetchBootstrapSegments()
   {
-    final String path = "/druid/coordinator/v1/metadata/datasources";
+    final String path = "/druid/coordinator/v1/metadata/bootstrapSegments";
     return FutureUtils.transform(
         client.asyncRequest(
-            new RequestBuilder(HttpMethod.GET, path),
-            new BytesFullResponseHandler()
+            new RequestBuilder(HttpMethod.POST, path),
+            new InputStreamResponseHandler()
         ),
-        holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), new TypeReference<List<String>>() {})
+        in -> new BootstrapSegmentsResponse(
+            new JsonParserIterator<>(
+                // Some servers, like the Broker, may have PruneLoadSpec set to true for optimization reasons.
+                // We specifically use LoadableDataSegment here instead of DataSegment so the callers can still correctly
+                // load the bootstrap segments, as the load specs are guaranteed not to be pruned.
+                jsonMapper.getTypeFactory().constructType(LoadableDataSegment.class),
+                Futures.immediateFuture(in),
+                jsonMapper
+            )
+        )
     );
   }
 
@@ -179,4 +188,17 @@ public CoordinatorClientImpl withRetryPolicy(ServiceRetryPolicy retryPolicy)
   {
     return new CoordinatorClientImpl(client.withRetryPolicy(retryPolicy), jsonMapper);
   }
+
+  @Override
+  public ListenableFuture<List<String>> fetchDataSources()
+  {
+    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>>() {})
+    );
+  }
 }

server/src/main/java/org/apache/druid/server/coordinator/DruidCoordinator.java

@@ -242,16 +242,6 @@ public Map<String, Double> getDatasourceToLoadStatus()
     return loadStatus;
   }
 
-  /**
-   * @return Set of broadcast segments determined by the latest run of the {@link RunRules} duty.
-   * If the coordinator runs haven't triggered or are delayed, this information may be stale.
-   */
-  @Nullable
-  public Set<DataSegment> getBroadcastSegments()
-  {
-    return broadcastSegments;
-  }
-
   @Nullable
   public Long getTotalSizeOfSegmentsAwaitingCompaction(String dataSource)
   {
@@ -701,6 +691,4 @@ public String toString()
       return "DutiesRunnable{group='" + dutyGroupName + '\'' + '}';
     }
   }
-
-
 }

server/src/main/java/org/apache/druid/server/coordinator/SegmentReplicationStatusManager.java

@@ -40,6 +40,7 @@
 import javax.inject.Inject;
 import java.util.Collections;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * Manages information about replication status of segments in a cluster.
@@ -49,6 +50,12 @@ public class SegmentReplicationStatusManager
 {
   private final MetadataManager metadataManager;
 
+  /**
+   * Set of broadcast segments determined in the latest coordinator run of the {@link RunRules} duty.
+   * This might contain stale information if the Coordinator duties haven't run or are delayed.
+   */
+  private volatile Set<DataSegment> broadcastSegments = null;
+
   /**
    * Used to determine count of under-replicated or unavailable segments.
    * Updated in each coordinator run in the {@link SegmentReplicationStatusManager.UpdateReplicationStatus} duty.
@@ -68,6 +75,16 @@ public SegmentReplicationStatusManager(MetadataManager metadataManager)
     this.metadataManager = metadataManager;
   }
 
+  /**
+   * @return Set of broadcast segments determined by the latest run of the {@link RunRules} duty.
+   * If the coordinator runs haven't triggered or are delayed, this information may be stale.
+   */
+  @Nullable
+  public Set<DataSegment> getBroadcastSegments()
+  {
+    return broadcastSegments;
+  }
+
   public Object2IntMap<String> getDatasourceToUnavailableSegmentCount()
   {
     if (segmentReplicationStatus == null) {
@@ -155,6 +172,7 @@ class UpdateReplicationStatus implements CoordinatorDuty
     @Override
     public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params)
     {
+      broadcastSegments = params.getBroadcastSegments();
       segmentReplicationStatus = params.getSegmentReplicationStatus();
 
       // Collect stats for unavailable and under-replicated segments

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

@@ -481,7 +481,7 @@ public Response getDataSourceInformation(
   @ResourceFilters(DatasourceResourceFilter.class)
   public Response getBootstrapSegments()
   {
-    final Set<DataSegment> broadcastSegments = coordinator.getBroadcastSegments();
+    final Set<DataSegment> broadcastSegments = segmentReplicationStatusManager.getBroadcastSegments();
     if (broadcastSegments == null) {
       return Response.status(Response.Status.SERVICE_UNAVAILABLE)
                      .entity("Bootstrap segments are not initialized yet."

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

@@ -64,7 +64,7 @@ public ListenableFuture<List<DataSourceInformation>> fetchDataSourceInformation(
   }
 
   @Override
-  public ListenableFuture<List<String>> fetchDataSources()
+  public ListenableFuture<BootstrapSegmentsResponse> fetchBootstrapSegments()
   {
     throw new UnsupportedOperationException();
   }
@@ -75,4 +75,10 @@ public CoordinatorClient withRetryPolicy(ServiceRetryPolicy retryPolicy)
     // Ignore retryPolicy for the test client.
     return this;
   }
+
+  @Override
+  public ListenableFuture<List<String>> fetchDataSources()
+  {
+    throw new UnsupportedOperationException();
+  }
 }

server/src/test/java/org/apache/druid/server/coordinator/DruidCoordinatorTest.java

@@ -271,6 +271,7 @@ public void testCoordinatorRun() throws Exception
     coordinator.start();
 
     Assert.assertNull(segmentReplicationStatusManager.getReplicationFactor(dataSegment.getId()));
+    Assert.assertNull(segmentReplicationStatusManager.getBroadcastSegments());
 
     // Wait for this coordinator to become leader
     leaderAnnouncerLatch.await();
@@ -298,7 +299,7 @@ public void testCoordinatorRun() throws Exception
         segmentReplicationStatusManager.getDatasourceToUnavailableSegmentCount();
     Assert.assertEquals(1, numsUnavailableUsedSegmentsPerDataSource.size());
     Assert.assertEquals(0, numsUnavailableUsedSegmentsPerDataSource.getInt(dataSource));
-    Assert.assertEquals(0, coordinator.getBroadcastSegments().size());
+    Assert.assertEquals(0, segmentReplicationStatusManager.getBroadcastSegments().size());
 
     Map<String, Object2LongMap<String>> underReplicationCountsPerDataSourcePerTier =
         segmentReplicationStatusManager.getTierToDatasourceToUnderReplicatedCount(false);
@@ -581,7 +582,7 @@ public void testComputeUnderReplicationCountsPerDataSourcePerTierForSegmentsWith
     coordinatorRunLatch.await();
 
     Assert.assertEquals(ImmutableMap.of(dataSource, 100.0), coordinator.getDatasourceToLoadStatus());
-    Assert.assertEquals(new HashSet<>(dataSegments.values()), coordinator.getBroadcastSegments());
+    Assert.assertEquals(new HashSet<>(dataSegments.values()), segmentReplicationStatusManager.getBroadcastSegments());
 
     // Under-replicated counts are updated only after the next coordinator run
     Map<String, Object2LongMap<String>> underReplicationCountsPerDataSourcePerTier =