Enable smartSegmentLoading on the Coordinator

benchmarks/src/test/java/org/apache/druid/server/coordinator/BalancerStrategyBenchmark.java

@@ -20,17 +20,14 @@
 package org.apache.druid.server.coordinator;
 
 import com.google.common.collect.ImmutableMap;
-import com.google.common.util.concurrent.MoreExecutors;
 import org.apache.druid.client.ImmutableDruidDataSource;
 import org.apache.druid.client.ImmutableDruidServer;
 import org.apache.druid.java.util.common.Intervals;
-import org.apache.druid.java.util.common.concurrent.Execs;
 import org.apache.druid.server.coordination.DruidServerMetadata;
 import org.apache.druid.server.coordination.ServerType;
 import org.apache.druid.server.coordinator.balancer.BalancerSegmentHolder;
-import org.apache.druid.server.coordinator.balancer.BalancerStrategy;
-import org.apache.druid.server.coordinator.balancer.CostBalancerStrategy;
-import org.apache.druid.server.coordinator.loadqueue.LoadQueuePeonTester;
+import org.apache.druid.server.coordinator.balancer.ReservoirSegmentSampler;
+import org.apache.druid.server.coordinator.loading.LoadQueuePeonTester;
 import org.apache.druid.timeline.DataSegment;
 import org.joda.time.Interval;
 import org.openjdk.jmh.annotations.Benchmark;
@@ -67,7 +64,7 @@ public class BalancerStrategyBenchmark
   private static final Interval TEST_SEGMENT_INTERVAL = Intervals.of("2012-03-15T00:00:00.000/2012-03-16T00:00:00.000");
   private static final int NUMBER_OF_SERVERS = 20;
 
-  @Param({"default", "50percentOfSegmentsToConsiderPerMove", "useBatchedSegmentSampler"})
+  @Param({"default", "useBatchedSegmentSampler"})
   private String mode;
 
   @Param({"10000", "100000", "1000000"})
@@ -77,26 +74,13 @@ public class BalancerStrategyBenchmark
   private int maxSegmentsToMove;
 
   private final List<ServerHolder> serverHolders = new ArrayList<>();
-  private boolean useBatchedSegmentSampler;
   private int reservoirSize = 1;
-  private double percentOfSegmentsToConsider = 100;
-  private final BalancerStrategy balancerStrategy = new CostBalancerStrategy(
-      MoreExecutors.listeningDecorator(Execs.multiThreaded(1, "BalancerStrategyBenchmark-%d"))
-  );
 
   @Setup(Level.Trial)
   public void setup()
   {
-    switch (mode) {
-      case "50percentOfSegmentsToConsiderPerMove":
-        percentOfSegmentsToConsider = 50;
-        useBatchedSegmentSampler = false;
-        break;
-      case "useBatchedSegmentSampler":
-        reservoirSize = maxSegmentsToMove;
-        useBatchedSegmentSampler = true;
-        break;
-      default:
+    if ("useBatchedSegmentSampler".equals(mode)) {
+      reservoirSize = maxSegmentsToMove;
     }
 
     List<List<DataSegment>> segmentList = new ArrayList<>(NUMBER_OF_SERVERS);
@@ -135,21 +119,9 @@ public void setup()
   @Benchmark
   public void pickSegmentsToMove(Blackhole blackhole)
   {
-    Iterator<BalancerSegmentHolder> iterator;
-    if (useBatchedSegmentSampler) {
-      iterator = balancerStrategy.pickSegmentsToMove(
-          serverHolders,
-          Collections.emptySet(),
-          reservoirSize,
-          false
-      );
-    } else {
-      iterator = balancerStrategy.pickSegmentsToMove(
-          serverHolders,
-          Collections.emptySet(),
-          percentOfSegmentsToConsider
-      );
-    }
+    Iterator<BalancerSegmentHolder> iterator = ReservoirSegmentSampler
+        .pickMovableLoadedSegmentsFrom(serverHolders, reservoirSize, Collections.emptySet())
+        .iterator();
 
     for (int i = 0; i < maxSegmentsToMove && iterator.hasNext(); i++) {
       blackhole.consume(iterator.next());

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

@@ -27,7 +27,7 @@
 import org.apache.druid.common.config.JacksonConfigManager;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.server.coordinator.duty.KillUnusedSegments;
-import org.apache.druid.server.coordinator.loadqueue.LoadQueuePeon;
+import org.apache.druid.server.coordinator.loading.LoadQueuePeon;
 import org.apache.druid.server.coordinator.stats.Dimension;
 
 import javax.annotation.Nonnull;
@@ -37,7 +37,6 @@
 import javax.validation.constraints.NotNull;
 import java.util.Collection;
 import java.util.EnumMap;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
@@ -236,20 +235,13 @@ private Map<Dimension, String> validateDebugDimensions(Map<String, String> debug
       return validDebugDimensions;
     }
 
-    final Map<String, Dimension> nameToDimension = new HashMap<>();
     for (Dimension dimension : Dimension.values()) {
-      nameToDimension.put(dimension.reportedName(), dimension);
+      final String dimensionValue = debugDimensions.get(dimension.reportedName());
+      if (dimensionValue != null) {
+        validDebugDimensions.put(dimension, dimensionValue);
+      }
     }
 
-    debugDimensions.forEach(
-        (dimensionName, value) -> {
-          Dimension dimension = nameToDimension.get(dimensionName);
-          if (dimension != null && value != null) {
-            validDebugDimensions.put(dimension, value);
-          }
-        }
-    );
-
     return validDebugDimensions;
   }
 
@@ -511,7 +503,8 @@ public boolean equals(Object o)
            && Objects.equals(
                dataSourcesToNotKillStalePendingSegmentsIn,
                that.dataSourcesToNotKillStalePendingSegmentsIn)
-           && Objects.equals(decommissioningNodes, that.decommissioningNodes);
+           && Objects.equals(decommissioningNodes, that.decommissioningNodes)
+           && Objects.equals(debugDimensions, that.debugDimensions);
   }
 
   @Override
@@ -534,7 +527,8 @@ public int hashCode()
         decommissioningNodes,
         decommissioningMaxPercentOfMaxSegmentsToMove,
         pauseCoordination,
-        maxNonPrimaryReplicantsToLoad
+        maxNonPrimaryReplicantsToLoad,
+        debugDimensions
     );
   }
 

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

@@ -70,9 +70,13 @@
 import org.apache.druid.server.coordinator.duty.MarkOvershadowedSegmentsAsUnused;
 import org.apache.druid.server.coordinator.duty.RunRules;
 import org.apache.druid.server.coordinator.duty.UnloadUnusedSegments;
-import org.apache.druid.server.coordinator.loadqueue.LoadQueuePeon;
-import org.apache.druid.server.coordinator.loadqueue.LoadQueueTaskMaster;
-import org.apache.druid.server.coordinator.loadqueue.SegmentLoadQueueManager;
+import org.apache.druid.server.coordinator.loading.LoadQueuePeon;
+import org.apache.druid.server.coordinator.loading.LoadQueueTaskMaster;
+import org.apache.druid.server.coordinator.loading.ReplicationThrottler;
+import org.apache.druid.server.coordinator.loading.SegmentLoadQueueManager;
+import org.apache.druid.server.coordinator.loading.SegmentLoadingConfig;
+import org.apache.druid.server.coordinator.loading.SegmentReplicaCount;
+import org.apache.druid.server.coordinator.loading.SegmentReplicantLookup;
 import org.apache.druid.server.coordinator.rules.LoadRule;
 import org.apache.druid.server.coordinator.stats.CoordinatorRunStats;
 import org.apache.druid.server.coordinator.stats.CoordinatorStat;
@@ -258,11 +262,11 @@ public Object2IntMap<String> getDatasourceToUnavailableSegmentCount()
 
     final Iterable<DataSegment> dataSegments = segmentsMetadataManager.iterateAllUsedSegments();
     for (DataSegment segment : dataSegments) {
-      if (segmentReplicantLookup.getLoadedReplicas(segment.getId(), true) == 0) {
-        datasourceToUnavailableSegments.addTo(segment.getDataSource(), 1);
-      } else {
-        datasourceToUnavailableSegments.addTo(segment.getDataSource(), 0);
-      }
+      SegmentReplicaCount replicaCount = segmentReplicantLookup.getReplicaCountsInCluster(segment.getId());
+      datasourceToUnavailableSegments.addTo(
+          segment.getDataSource(),
+          replicaCount.loaded() == 0 ? 1 : 0
+      );
     }
 
     return datasourceToUnavailableSegments;
@@ -801,22 +805,23 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params)
       startPeonsForNewServers(currentServers);
       stopPeonsForDisappearedServers(currentServers);
 
-      final DruidCluster cluster = prepareCluster(params.getCoordinatorDynamicConfig(), currentServers);
-      cancelLoadsOnDecommissioningServers(cluster);
-
       final CoordinatorDynamicConfig dynamicConfig = params.getCoordinatorDynamicConfig();
+      final SegmentLoadingConfig segmentLoadingConfig = params.getSegmentLoadingConfig();
+
+      final DruidCluster cluster = prepareCluster(dynamicConfig, segmentLoadingConfig, currentServers);
+      cancelLoadsOnDecommissioningServers(cluster);
 
       initBalancerExecutor();
       final BalancerStrategy balancerStrategy = balancerStrategyFactory.createBalancerStrategy(balancerExec);
       log.info(
           "Using balancer strategy [%s] with round-robin assignment [%s] and debug dimensions [%s].",
           balancerStrategy.getClass().getSimpleName(),
-          dynamicConfig.isUseRoundRobinSegmentAssignment(), dynamicConfig.getDebugDimensions()
+          segmentLoadingConfig.isUseRoundRobinSegmentAssignment(),
+          dynamicConfig.getDebugDimensions()
       );
 
       params = params.buildFromExisting()
                      .withDruidCluster(cluster)
-                     .withDynamicConfigs(recomputeDynamicConfig(params))
                      .withBalancerStrategy(balancerStrategy)
                      .withSegmentAssignerUsing(loadQueueManager)
                      .build();
@@ -826,44 +831,6 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params)
       return params;
     }
 
-    /**
-     * Recomputes dynamic config values if {@code smartLoadQueue} is enabled.
-     */
-    private CoordinatorDynamicConfig recomputeDynamicConfig(DruidCoordinatorRuntimeParams params)
-    {
-      final CoordinatorDynamicConfig dynamicConfig = params.getCoordinatorDynamicConfig();
-      if (!dynamicConfig.isSmartSegmentLoading()) {
-        return dynamicConfig;
-      }
-
-      // Impose a lower bound on both replicationThrottleLimit and maxSegmentsToMove
-      final int throttlePercentage = 2;
-      final int replicationThrottleLimit = Math.max(
-          100,
-          params.getUsedSegments().size() * throttlePercentage / 100
-      );
-
-      // Impose an upper bound on maxSegmentsToMove to ensure that coordinator
-      // run times are bounded. This limit can be relaxed as performance of
-      // the CostBalancerStrategy.computeCost() is improved.
-      final int maxSegmentsToMove = Math.min(1000, replicationThrottleLimit);
-
-      log.info(
-          "Smart segment loading is enabled. Recomputed replicationThrottleLimit"
-          + " [%d] (%d%% of used segments) and maxSegmentsToMove [%d].",
-          replicationThrottleLimit, throttlePercentage, maxSegmentsToMove
-      );
-
-      return CoordinatorDynamicConfig.builder()
-                                     .withMaxSegmentsInNodeLoadingQueue(0)
-                                     .withReplicationThrottleLimit(replicationThrottleLimit)
-                                     .withMaxSegmentsToMove(maxSegmentsToMove)
-                                     .withUseRoundRobinSegmentAssignment(true)
-                                     .withUseBatchedSegmentSampler(true)
-                                     .withEmitBalancingStats(false)
-                                     .build(dynamicConfig);
-    }
-
     /**
      * Cancels all load/move operations on decommissioning servers. This should
      * be done before initializing the SegmentReplicantLookup so that
@@ -931,7 +898,11 @@ void startPeonsForNewServers(List<ImmutableDruidServer> currentServers)
       }
     }
 
-    DruidCluster prepareCluster(CoordinatorDynamicConfig dynamicConfig, List<ImmutableDruidServer> currentServers)
+    DruidCluster prepareCluster(
+        CoordinatorDynamicConfig dynamicConfig,
+        SegmentLoadingConfig segmentLoadingConfig,
+        List<ImmutableDruidServer> currentServers
+    )
     {
       final Set<String> decommissioningServers = dynamicConfig.getDecommissioningNodes();
       final DruidCluster.Builder cluster = DruidCluster.builder();
@@ -941,8 +912,8 @@ DruidCluster prepareCluster(CoordinatorDynamicConfig dynamicConfig, List<Immutab
                 server,
                 loadManagementPeons.get(server.getName()),
                 decommissioningServers.contains(server.getHost()),
-                dynamicConfig.getMaxSegmentsInNodeLoadingQueue(),
-                dynamicConfig.getReplicantLifetime()
+                segmentLoadingConfig.getMaxSegmentsInLoadQueue(),
+                segmentLoadingConfig.getMaxLifetimeInLoadQueue()
             )
         );
       }