Improve CostBalancerStrategy, deprecate cachingCost

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

@@ -20,6 +20,7 @@
 package org.apache.druid.server.coordinator;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Stopwatch;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Ordering;
 import com.google.common.collect.Sets;
@@ -98,6 +99,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
@@ -672,7 +674,7 @@ public void run()
     {
       try {
         log.info("Starting coordinator run for group [%s]", dutyGroupName);
-        final long globalStart = System.currentTimeMillis();
+        final Stopwatch groupRunTime = Stopwatch.createStarted();
 
         synchronized (lock) {
           if (!coordLeaderSelector.isLeader()) {
@@ -719,23 +721,25 @@ public void run()
           log.info("Coordination has been paused. Duties will not run until coordination is resumed.");
         }
 
+        final Stopwatch dutyRunTime = Stopwatch.createUnstarted();
         for (CoordinatorDuty duty : duties) {
           // Don't read state and run state in the same duty otherwise racy conditions may exist
           if (!coordinationPaused
               && coordLeaderSelector.isLeader()
               && startingLeaderCounter == coordLeaderSelector.localTerm()) {
 
-            final long start = System.currentTimeMillis();
+            dutyRunTime.reset().start();
             params = duty.run(params);
-            final long end = System.currentTimeMillis();
+            dutyRunTime.stop();
 
             final String dutyName = duty.getClass().getName();
             if (params == null) {
               log.info("Stopping run for group [%s] on request of duty [%s].", dutyGroupName, dutyName);
               return;
             } else {
-              final RowKey rowKey = RowKey.builder().add(Dimension.DUTY, dutyName).build();
-              params.getCoordinatorStats().add(Stats.CoordinatorRun.DUTY_RUN_TIME, rowKey, end - start);
+              final RowKey rowKey = RowKey.of(Dimension.DUTY, dutyName);
+              final long dutyRunMillis = dutyRunTime.elapsed(TimeUnit.MILLISECONDS);
+              params.getCoordinatorStats().add(Stats.CoordinatorRun.DUTY_RUN_TIME, rowKey, dutyRunMillis);
             }
           }
         }
@@ -745,9 +749,9 @@ public void run()
         if (allStats.rowCount() > 0) {
           final AtomicInteger emittedCount = new AtomicInteger();
           allStats.forEachStat(
-              (dimensionValues, stat, value) -> {
+              (stat, dimensions, value) -> {
                 if (stat.shouldEmit()) {
-                  emitStat(stat, dimensionValues, value);
+                  emitStat(stat, dimensions.getValues(), value);
                   emittedCount.incrementAndGet();
                 }
               }
@@ -760,7 +764,7 @@ public void run()
         }
 
         // Emit the runtime of the full DutiesRunnable
-        final long runMillis = System.currentTimeMillis() - globalStart;
+        final long runMillis = groupRunTime.stop().elapsed(TimeUnit.MILLISECONDS);
         emitStat(Stats.CoordinatorRun.GROUP_RUN_TIME, Collections.emptyMap(), runMillis);
         log.info("Finished coordinator run for group [%s] in [%d] ms", dutyGroupName, runMillis);
       }
@@ -771,10 +775,6 @@ public void run()
 
     private void emitStat(CoordinatorStat stat, Map<Dimension, String> dimensionValues, long value)
     {
-      if (stat.equals(Stats.Balancer.NORMALIZED_COST_X_1000)) {
-        value = value / 1000;
-      }
-
       ServiceMetricEvent.Builder eventBuilder = new ServiceMetricEvent.Builder()
           .setDimension(Dimension.DUTY_GROUP.reportedName(), dutyGroupName);
       dimensionValues.forEach(

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

@@ -0,0 +1,67 @@
+/*
+ * 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.server.coordinator;
+
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntMaps;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+import org.apache.druid.timeline.DataSegment;
+import org.joda.time.Interval;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Maintains a count of segments for each datasource and interval.
+ */
+public class SegmentCountsPerInterval
+{
+  private final Map<String, Object2IntMap<Interval>>
+      datasourceIntervalToSegmentCount = new HashMap<>();
+  private final Object2IntMap<Interval> intervalToTotalSegmentCount = new Object2IntOpenHashMap<>();
+
+  public void addSegment(DataSegment segment)
+  {
+    updateCountInInterval(segment, 1);
+  }
+
+  public void removeSegment(DataSegment segment)
+  {
+    updateCountInInterval(segment, -1);
+  }
+
+  public Object2IntMap<Interval> getIntervalToSegmentCount(String datasource)
+  {
+    return datasourceIntervalToSegmentCount.getOrDefault(datasource, Object2IntMaps.emptyMap());
+  }
+
+  public Object2IntMap<Interval> getIntervalToTotalSegmentCount()
+  {
+    return intervalToTotalSegmentCount;
+  }
+
+  private void updateCountInInterval(DataSegment segment, int delta)
+  {
+    intervalToTotalSegmentCount.merge(segment.getInterval(), delta, Integer::sum);
+    datasourceIntervalToSegmentCount
+        .computeIfAbsent(segment.getDataSource(), ds -> new Object2IntOpenHashMap<>())
+        .merge(segment.getInterval(), delta, Integer::sum);
+  }
+}

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

@@ -32,11 +32,9 @@
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -75,11 +73,7 @@ public class ServerHolder implements Comparable<ServerHolder>
    */
   private final Map<DataSegment, SegmentAction> queuedSegments = new HashMap<>();
 
-  /**
-   * Segments that are expected to be loaded on this server once all the
-   * operations in progress have completed.
-   */
-  private final Set<DataSegment> projectedSegments = new HashSet<>();
+  private final SegmentCountsPerInterval projectedSegments = new SegmentCountsPerInterval();
 
   public ServerHolder(ImmutableDruidServer server, LoadQueuePeon peon)
   {
@@ -133,7 +127,7 @@ private void initializeQueuedSegments(
       AtomicInteger loadingReplicaCount
   )
   {
-    projectedSegments.addAll(server.iterateAllSegments());
+    server.iterateAllSegments().forEach(projectedSegments::addSegment);
 
     final List<SegmentHolder> expiredSegments = new ArrayList<>();
     peon.getSegmentsInQueue().forEach(
@@ -251,11 +245,21 @@ public Map<DataSegment, SegmentAction> getQueuedSegments()
    * Segments that are expected to be loaded on this server once all the
    * operations in progress have completed.
    */
-  public Set<DataSegment> getProjectedSegments()
+  public SegmentCountsPerInterval getProjectedSegments()
   {
     return projectedSegments;
   }
 
+  public boolean isProjectedSegment(DataSegment segment)
+  {
+    SegmentAction action = getActionOnSegment(segment);
+    if (action == null) {
+      return hasSegmentLoaded(segment.getId());
+    } else {
+      return action.isLoad();
+    }
+  }
+
   /**
    * Segments that are currently in the queue for being loaded on this server.
    * This does not include segments that are being moved to this server.
@@ -362,10 +366,10 @@ private void addToQueuedSegments(DataSegment segment, SegmentAction action)
 
     // Add to projected if load is started, remove from projected if drop has started
     if (action.isLoad()) {
-      projectedSegments.add(segment);
+      projectedSegments.addSegment(segment);
       sizeOfLoadingSegments += segment.getSize();
     } else {
-      projectedSegments.remove(segment);
+      projectedSegments.removeSegment(segment);
       if (action == SegmentAction.DROP) {
         sizeOfDroppingSegments += segment.getSize();
       }
@@ -379,10 +383,10 @@ private void removeFromQueuedSegments(DataSegment segment, SegmentAction action)
     queuedSegments.remove(segment);
 
     if (action.isLoad()) {
-      projectedSegments.remove(segment);
+      projectedSegments.removeSegment(segment);
       sizeOfLoadingSegments -= segment.getSize();
     } else {
-      projectedSegments.add(segment);
+      projectedSegments.addSegment(segment);
       if (action == SegmentAction.DROP) {
         sizeOfDroppingSegments -= segment.getSize();
       }

server/src/main/java/org/apache/druid/server/coordinator/balancer/BalancerStrategy.java

@@ -20,70 +20,60 @@
 package org.apache.druid.server.coordinator.balancer;
 
 import org.apache.druid.server.coordinator.ServerHolder;
-import org.apache.druid.server.coordinator.duty.BalanceSegments;
 import org.apache.druid.server.coordinator.stats.CoordinatorRunStats;
 import org.apache.druid.timeline.DataSegment;
 
 import javax.annotation.Nullable;
 import java.util.Iterator;
 import java.util.List;
-import java.util.NavigableSet;
 
 /**
- * This interface describes the coordinator balancing strategy, which is responsible for making decisions on where
- * to place {@link DataSegment}s on historical servers (described by {@link ServerHolder}). The balancing strategy
- * is used by {@link org.apache.druid.server.coordinator.rules.LoadRule} to assign and drop segments, and by
- * {@link BalanceSegments} to migrate segments between historicals.
+ * Segment balancing strategy, used in every coordinator run by
+ * {@link org.apache.druid.server.coordinator.loading.StrategicSegmentAssigner}
+ * to choose optimal servers to load, move or drop a segment.
  */
 public interface BalancerStrategy
 {
 
   /**
-   * Finds the best server to move a segment to according to the balancing strategy.
+   * Finds the best server from the list of {@code destinationServers} to load
+   * the {@code segmentToMove}, if it is moved from the {@code sourceServer}.
+   * <p>
+   * In order to avoid unnecessary moves when the segment is already optimally placed,
+   * include the {@code sourceServer} in the list of {@code destinationServers}.
    *
-   * @param proposalSegment    segment to move
-   * @param sourceServer       Server the segment is currently placed on.
-   * @param destinationServers servers to consider as move destinations
-   * @return The server to move to, or null if no move should be made or no server is suitable
+   * @return The server to move to, or null if the segment is already optimally placed.
    */
   @Nullable
   ServerHolder findDestinationServerToMoveSegment(
-      DataSegment proposalSegment,
+      DataSegment segmentToMove,
       ServerHolder sourceServer,
       List<ServerHolder> destinationServers
   );
 
   /**
-   * Finds the best servers on which to place the {@code proposalSegment}.
-   * This method can be used both for placing the first copy of a segment
-   * in the tier or a replica of the segment.
+   * Finds the best servers to load the given segment. This method can be used
+   * both for placing the first copy of a segment in a tier or a replica of an
+   * already available segment.
    *
-   * @param proposalSegment segment to place on servers
-   * @param serverHolders   servers to consider as segment homes
-   * @return Iterator over the best servers (in order) on which the segment
-   * can be placed.
+   * @return Iterator over the best servers (in order of preference) to load
+   * the segment.
    */
   Iterator<ServerHolder> findServersToLoadSegment(
-      DataSegment proposalSegment,
+      DataSegment segmentToLoad,
       List<ServerHolder> serverHolders
   );
 
   /**
-   * Returns an iterator for a set of servers to drop from, ordered by preference of which server to drop from first
-   * for a given drop strategy. One or more segments may be dropped, depending on how much the segment is
-   * over-replicated.
-   * @param toDropSegment segment to drop from one or more servers
-   * @param serverHolders set of historicals to consider dropping from
-   * @return Iterator for set of historicals, ordered by drop preference
+   * Finds the best servers to drop the given segment.
+   *
+   * @return Iterator over the servers (in order of preference) to drop the segment
    */
-  Iterator<ServerHolder> pickServersToDropSegment(DataSegment toDropSegment, NavigableSet<ServerHolder> serverHolders);
+  Iterator<ServerHolder> findServersToDropSegment(DataSegment segmentToDrop, List<ServerHolder> serverHolders);
 
   /**
-   * Add balancing strategy stats during the 'balanceTier' operation of
-   * {@link BalanceSegments} to be included
-   * @param tier historical tier being balanced
-   * @param stats stats object to add balancing strategy stats to
-   * @param serverHolderList servers in tier being balanced
+   * Returns the stats collected by the strategy in the current run and resets
+   * the stats collector for the next run.
    */
-  void emitStats(String tier, CoordinatorRunStats stats, List<ServerHolder> serverHolderList);
+  CoordinatorRunStats getAndResetStats();
 }

server/src/main/java/org/apache/druid/server/coordinator/balancer/CachingCostBalancerStrategy.java

@@ -28,10 +28,8 @@
 import java.util.Collections;
 import java.util.Set;
 
-
 public class CachingCostBalancerStrategy extends CostBalancerStrategy
 {
-
   private final ClusterCostCache clusterCostCache;
 
   public CachingCostBalancerStrategy(ClusterCostCache clusterCostCache, ListeningExecutorService exec)
@@ -41,13 +39,8 @@ public CachingCostBalancerStrategy(ClusterCostCache clusterCostCache, ListeningE
   }
 
   @Override
-  protected double computeCost(DataSegment proposalSegment, ServerHolder server, boolean includeCurrentServer)
+  protected double computePlacementCost(DataSegment proposalSegment, ServerHolder server)
   {
-    // (optional) Don't include server if it cannot load the segment
-    if (!includeCurrentServer && !server.canLoadSegment(proposalSegment)) {
-      return Double.POSITIVE_INFINITY;
-    }
-
     final String serverName = server.getServer().getName();
 
     double cost = clusterCostCache.computeCost(serverName, proposalSegment);

server/src/main/java/org/apache/druid/server/coordinator/balancer/CachingCostBalancerStrategyFactory.java

@@ -125,6 +125,10 @@ private boolean isInitialized()
   @Override
   public BalancerStrategy createBalancerStrategy(final ListeningExecutorService exec)
   {
+    LOG.warn(
+        "'cachingCost' balancer strategy has been deprecated as it can lead to"
+        + " unbalanced clusters. Use 'cost' strategy instead."
+    );
     if (!isInitialized() && config.isAwaitInitialization()) {
       try {
         final long startMillis = System.currentTimeMillis();