[GitHub] [druid] imply-cheddar commented on a diff in pull request #13197: Segment loading: Allow cancellation and prioritization of load queue items

["imply-cheddar (via GitHub)" <gi...@apache.org>]

imply-cheddar commented on code in PR #13197:
URL: https://github.com/apache/druid/pull/13197#discussion_r1231715441


##########
server/src/main/java/org/apache/druid/server/coordinator/balancer/BalancerStrategy.java:
##########
@@ -36,77 +39,65 @@
  */
 public interface BalancerStrategy
 {
+
   /**
-   * Find the best server to move a {@link DataSegment} to according the balancing strategy.
-   * @param proposalSegment segment to move
-   * @param serverHolders servers to consider as move destinations
+   * Finds the best server to move a segment to according to the balancing strategy.
+   *
+   * @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
    */
   @Nullable
-  ServerHolder findNewSegmentHomeBalancer(DataSegment proposalSegment, List<ServerHolder> serverHolders);
+  ServerHolder findDestinationServerToMoveSegment(
+      DataSegment proposalSegment,
+      ServerHolder sourceServer,
+      List<ServerHolder> destinationServers
+  );
 
   /**
-   * Finds the best servers on which to place a replica of the {@code proposalSegment}
-   * according to the balancing strategy.
+   * 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.
    *
-   * @param proposalSegment segment to replicate
-   * @param serverHolders   servers to consider as replica holders
-   * @return Iterator over the best servers (in order) on which the replica(s)
+   * @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.
    */
-  Iterator<ServerHolder> findNewSegmentHomeReplicator(
+  Iterator<ServerHolder> findServersToLoadSegment(
       DataSegment proposalSegment,
       List<ServerHolder> serverHolders
   );
 
   /**
-   * Pick the best segments to move from one of the supplied set of servers according to the balancing strategy.
+   * Picks segments from the given set of servers based on the balancing strategy.
+   * Default behaviour is to pick segments using reservoir sampling.
    *
-   * @param serverHolders set of historicals to consider for moving segments
-   * @param broadcastDatasources Datasources that contain segments which were loaded via broadcast rules.
-   *                             Balancing strategies should avoid rebalancing segments for such datasources, since
-   *                             they should be loaded on all servers anyway.
-   *                             NOTE: this should really be handled on a per-segment basis, to properly support
-   *                                   the interval or period-based broadcast rules. For simplicity of the initial
-   *                                   implementation, only forever broadcast rules are supported.
-   * @param reservoirSize the reservoir size maintained by the Reservoir Sampling algorithm.
-   * @return Iterator for set of {@link BalancerSegmentHolder} containing segment to move and server they currently
-   * reside on, or empty if there are no segments to pick from (i. e. all provided serverHolders are empty).
+   * @param serverHolders        Set of historicals to consider for picking segments
+   * @param broadcastDatasources Segments belonging to these datasources will not
+   *                             be picked for balancing, since they should be
+   *                             loaded on all servers anyway.
+   * @param maxSegmentsToPick    Maximum number of segments to pick
+   * @param pickLoadingSegments  If true, picks only segments currently being
+   *                             loaded on a server. If false, picks segments
+   *                             already loaded on a server.
+   * @return Iterator over {@link BalancerSegmentHolder}s, each of which contains
+   * a segment picked for moving and the server currently serving/loading it.
    */
   default Iterator<BalancerSegmentHolder> pickSegmentsToMove(
       List<ServerHolder> serverHolders,
       Set<String> broadcastDatasources,
-      int reservoirSize
+      int maxSegmentsToPick,
+      boolean pickLoadingSegments

Review Comment:
   This API definition with a boolean that says to pick loading segments or not is not really a good structure.  I initially thought it meant "also include loading segments", but it turns out that it means "either pick only loading or pick only non-loading".  If we want those semantics, we should have 2 different methods, one for picking loading segments and one for picking non loading.
   
   From what I can tell, at least with what is committed to the Druid code base, both of the "pickSegmentsToMove" methods are never overridden.  I don't know that this method actually makes sense as an extension point (it would likely be better to just replace the CoordinatorDuty entirely), so I think that the best thing to do in this PR is to actually completely eliminate this method *and* the other `pickSegmentsToMove` method and just inline the code.
   
   That will then make it the whole decision between loading segments versus not much easier to ascertain (you can basically grab the list that you care about and only use that, no need to push down the boolean).



##########
server/src/main/java/org/apache/druid/server/coordinator/CoordinatorDynamicConfig.java:
##########
@@ -194,21 +211,46 @@ public CoordinatorDynamicConfig(
           + "reflect this configuration being added to Druid in a recent release. Druid is defaulting the value "
           + "to the Druid default of %d. It is recommended that you re-submit your dynamic config with your "
           + "desired value for maxNonPrimaryReplicantsToLoad",
-          Builder.DEFAULT_MAX_NON_PRIMARY_REPLICANTS_TO_LOAD
+          Defaults.MAX_NON_PRIMARY_REPLICANTS_TO_LOAD
       );
-      maxNonPrimaryReplicantsToLoad = Builder.DEFAULT_MAX_NON_PRIMARY_REPLICANTS_TO_LOAD;
+      maxNonPrimaryReplicantsToLoad = Defaults.MAX_NON_PRIMARY_REPLICANTS_TO_LOAD;
     }
     Preconditions.checkArgument(
         maxNonPrimaryReplicantsToLoad >= 0,
         "maxNonPrimaryReplicantsToLoad must be greater than or equal to 0."
     );
     this.maxNonPrimaryReplicantsToLoad = maxNonPrimaryReplicantsToLoad;
 
-    if (useRoundRobinSegmentAssignment == null) {
-      this.useRoundRobinSegmentAssignment = Builder.DEFAULT_USE_ROUND_ROBIN_ASSIGNMENT;
-    } else {
-      this.useRoundRobinSegmentAssignment = useRoundRobinSegmentAssignment;
+    this.useRoundRobinSegmentAssignment = Builder.valueOrDefault(
+        useRoundRobinSegmentAssignment,
+        Defaults.USE_ROUND_ROBIN_ASSIGNMENT
+    );
+    this.debugDimensions = debugDimensions;
+    this.validDebugDimensions = validateDebugDimensions(debugDimensions);
+  }
+
+  private Map<Dimension, String> validateDebugDimensions(Map<String, String> debugDimensions)
+  {
+    final Map<Dimension, String> validDebugDimensions = new EnumMap<>(Dimension.class);
+    if (debugDimensions == null || debugDimensions.isEmpty()) {
+      return validDebugDimensions;
+    }
+
+    final Map<String, Dimension> nameToDimension = new HashMap<>();
+    for (Dimension dimension : Dimension.values()) {
+      nameToDimension.put(dimension.reportedName(), dimension);
     }
+
+    debugDimensions.forEach(
+        (dimensionName, value) -> {
+          Dimension dimension = nameToDimension.get(dimensionName);
+          if (dimension != null && value != null) {
+            validDebugDimensions.put(dimension, value);
+          }
+        }
+    );

Review Comment:
   This `forEach` seems egregious to me.  Can't you just use the name from the for loop above to lookup whether there's an entry in `debugDimensions` and then add it if it exists?



##########
server/src/main/java/org/apache/druid/server/coordinator/ServerHolder.java:
##########
@@ -155,128 +211,161 @@ public long getAvailableSize()
    * <p>
    * A load is possible only if the server meets all of the following criteria:
    * <ul>
-   *   <li>is not already serving or loading the segment</li>
    *   <li>is not being decommissioned</li>
+   *   <li>is not already serving the segment</li>
+   *   <li>is not performing any other action on the segment</li>
    *   <li>has not already exceeded the load queue limit in this run</li>
    *   <li>has available disk space</li>
    * </ul>
    */
   public boolean canLoadSegment(DataSegment segment)
   {
-    final SegmentState state = getSegmentState(segment);
     return !isDecommissioning
-           && (maxSegmentsInLoadQueue == 0 || maxSegmentsInLoadQueue > segmentsQueuedForLoad)
-           && getAvailableSize() >= segment.getSize()
-           && state == SegmentState.NONE;
+           && !hasSegmentLoaded(segment.getId())
+           && getActionOnSegment(segment) == null
+           && totalAssignmentsInRun < maxAssignmentsInRun
+           && getAvailableSize() >= segment.getSize();
   }
 
-  public SegmentState getSegmentState(DataSegment segment)
+  public SegmentAction getActionOnSegment(DataSegment segment)
   {
-    SegmentState state = queuedSegments.get(segment.getId());
-    if (state != null) {
-      return state;
-    }
-
-    return isServingSegment(segment) ? SegmentState.LOADED : SegmentState.NONE;
+    return queuedSegments.get(segment);
   }
 
   /**
    * Segments queued for load, drop or move on this server.
+   * <ul>
+   * <li>Contains segments present in the queue when the current coordinator run started.</li>
+   * <li>Contains segments added to the queue during the current run.</li>
+   * <li>Maps replicating segments to LOAD rather than REPLICATE for simplicity.</li>
+   * <li>Does not contain segments whose actions were cancelled.</li>
+   * </ul>
    */
-  public Map<SegmentId, SegmentState> getQueuedSegments()
+  public Map<DataSegment, SegmentAction> getQueuedSegments()
   {
-    return Collections.unmodifiableMap(queuedSegments);
+    return new HashMap<>(queuedSegments);
   }
 
+  /**
+   * Segments that are expected to be loaded on this server once all the
+   * operations in progress have completed.
+   */
+  public Set<DataSegment> getProjectedSegments()
+  {
+    return projectedSegments;
+  }
+
+  /**
+   * 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.
+   */
+  public List<DataSegment> getLoadingSegments()
+  {
+    final List<DataSegment> loadingSegments = new ArrayList<>();
+    queuedSegments.forEach((segment, action) -> {
+      if (action == SegmentAction.LOAD) {
+        loadingSegments.add(segment);
+      }
+    });
+
+    return loadingSegments;
+  }
+
+  /**
+   * Returns true if this server has the segment loaded and is not dropping it.
+   */
   public boolean isServingSegment(DataSegment segment)
   {
-    return isServingSegment(segment.getId());
+    return hasSegmentLoaded(segment.getId()) && getActionOnSegment(segment) == null;
   }
 
   public boolean isLoadingSegment(DataSegment segment)
   {
-    return getSegmentState(segment) == SegmentState.LOADING;
+    return getActionOnSegment(segment) == SegmentAction.LOAD;
   }
 
   public boolean isDroppingSegment(DataSegment segment)
   {
-    return getSegmentState(segment) == SegmentState.DROPPING;
+    return getActionOnSegment(segment) == SegmentAction.DROP;
   }
 
-  public boolean startOperation(DataSegment segment, SegmentState newState)
+  public int getNumMovingSegments()
   {
-    if (queuedSegments.containsKey(segment.getId())) {
-      return false;
-    }
+    return movingSegmentCount;
+  }
 
-    if (newState == SegmentState.LOADING || newState == SegmentState.MOVING_TO) {
-      ++segmentsQueuedForLoad;
-      sizeOfLoadingSegments += segment.getSize();
-    }
-    queuedSegments.put(segment.getId(), newState);
-    return true;
+  public int getNumLoadingReplicas()
+  {
+    return loadingReplicaCount;
   }
 
-  public boolean cancelOperation(DataSegment segment, SegmentState currentState)
+  public boolean startOperation(SegmentAction action, DataSegment segment)
   {
-    SegmentState observedState = queuedSegments.get(segment.getId());
-    if (observedState != currentState) {
+    if (queuedSegments.containsKey(segment)) {
       return false;
     }
 
-    if (currentState == SegmentState.LOADING || currentState == SegmentState.MOVING_TO) {
-      --segmentsQueuedForLoad;
-      sizeOfLoadingSegments -= segment.getSize();
+    if (action.isLoad()) {
+      ++totalAssignmentsInRun;
     }
-    queuedSegments.remove(segment.getId());
+
+    updateQueuedSegments(segment, simplify(action), true);
     return true;
   }
 
-  public boolean isServingSegment(SegmentId segmentId)
+  public boolean cancelOperation(SegmentAction action, DataSegment segment)
+  {
+    final SegmentAction queuedAction = queuedSegments.get(segment);
+    return queuedAction == simplify(action)
+           && (queuedAction == SegmentAction.MOVE_FROM || peon.cancelOperation(segment))
+           && updateQueuedSegments(segment, queuedAction, false);
+  }
+
+  public boolean hasSegmentLoaded(SegmentId segmentId)
   {
     return server.getSegment(segmentId) != null;
   }
 
-  /**
-   * Checks if the server can load the given segment.
-   * <p>
-   * A load is possible only if the server meets all of the following criteria:
-   * <ul>
-   *   <li>is not being decommissioned</li>
-   *   <li>is not already serving the segment</li>
-   *   <li>is not performing any other action on the segment</li>
-   *   <li>has not already exceeded the load queue limit in this run</li>
-   *   <li>has available disk space</li>
-   * </ul>
-   */
-  public boolean canLoadSegment(DataSegment segment)
+  public boolean isRealtimeServer()
   {
-    return !isDecommissioning
-           && !isServingSegment(segment.getId())
-           && !isLoadingSegment(segment)
-           && (maxSegmentsInLoadQueue == 0 || maxSegmentsInLoadQueue > peon.getNumberOfSegmentsInQueue())
-           && getAvailableSize() >= segment.getSize();
+    return server.getType() == ServerType.REALTIME
+           || server.getType() == ServerType.INDEXER_EXECUTOR;
   }
 
-  @Override
-  public int compareTo(ServerHolder serverHolder)
+  private SegmentAction simplify(SegmentAction action)
   {
-    int result = Long.compare(getAvailableSize(), serverHolder.getAvailableSize());
-    if (result != 0) {
-      return result;
+    return action == SegmentAction.REPLICATE ? SegmentAction.LOAD : action;
+  }
+
+  private boolean updateQueuedSegments(DataSegment segment, SegmentAction action, boolean addToQueue)
+  {
+    if (addToQueue) {
+      queuedSegments.put(segment, action);
+    } else {
+      queuedSegments.remove(segment);
     }
 
-    result = server.getHost().compareTo(serverHolder.server.getHost());
-    if (result != 0) {
-      return result;
+    final long sizeDelta = addToQueue ? segment.getSize() : -segment.getSize();
+    if (action.isLoad()) {
+      sizeOfLoadingSegments += sizeDelta;
+    } else if (action == SegmentAction.DROP) {
+      sizeOfDroppingSegments += sizeDelta;
     }
 
-    result = server.getTier().compareTo(serverHolder.server.getTier());
-    if (result != 0) {
-      return result;
+    // Remove from projected if load is cancelled or drop is started, add otherwise
+    if (addToQueue ^ action.isLoad()) {
+      projectedSegments.remove(segment);
+    } else {
+      projectedSegments.add(segment);
     }
 
-    return server.getType().compareTo(serverHolder.server.getType());
+    return true;
+  }

Review Comment:
   nit: the boolean parameter here feels really weird.  Why not just have 2 methods:
   
   ```
   addQueuedSegment()
   removeQueuedSegment()
   ```



##########
server/src/main/java/org/apache/druid/server/coordinator/loadqueue/HttpLoadQueuePeon.java:
##########
@@ -470,134 +517,127 @@
     return Collections.unmodifiableSet(segmentsMarkedToDrop);
   }
 
-  private abstract class SegmentHolder
+  /**
+   * A request is considered to have timed out if the time elapsed since it was
+   * first sent to the server is greater than the configured load timeout.
+   *
+   * @see DruidCoordinatorConfig#getLoadTimeoutDelay()
+   */
+  private boolean hasRequestTimedOut(SegmentHolder holder, long currentTimeMillis)
   {
-    private final DataSegment segment;
-    private final DataSegmentChangeRequest changeRequest;
-    private final List<LoadPeonCallback> callbacks = new ArrayList<>();
-
-    // Time when this request was sent to target server the first time.
-    private volatile long scheduleTime = -1;
-
-    private SegmentHolder(
-        DataSegment segment,
-        DataSegmentChangeRequest changeRequest,
-        LoadPeonCallback callback
-    )
-    {
-      this.segment = segment;
-      this.changeRequest = changeRequest;
+    return holder.isRequestSentToServer()
+           && currentTimeMillis - holder.getFirstRequestMillis()
+              > config.getLoadTimeoutDelay().getMillis();
+  }
 
-      if (callback != null) {
-        this.callbacks.add(callback);
-      }
-    }
+  private void onRequestFailed(SegmentHolder holder, String failureCause)
+  {
+    log.error(
+        "Server[%s] failed segment[%s] request[%s] with cause [%s].",
+        serverId, holder.getSegment().getId(), holder.getAction(), failureCause
+    );
+    onRequestCompleted(holder, QueueStatus.FAILED);
+  }
 
-    public void addCallback(LoadPeonCallback newCallback)
-    {
-      synchronized (callbacks) {
-        if (newCallback != null) {
-          callbacks.add(newCallback);
-        }
-      }
-    }
+  private void onRequestCompleted(SegmentHolder holder, QueueStatus status)
+  {
+    final SegmentAction action = holder.getAction();
+    log.trace(
+        "Server[%s] completed request[%s] on segment[%s] with status[%s].",
+        serverId, action, holder.getSegment().getId(), status
+    );
 
-    public DataSegment getSegment()
-    {
-      return segment;
+    if (holder.isLoad()) {
+      queuedSize.addAndGet(-holder.getSegment().getSize());
     }
+    incrementStat(holder, status);
+    executeCallbacks(holder, status == QueueStatus.SUCCESS);
+  }
 
-    public DataSegmentChangeRequest getChangeRequest()
-    {
-      return changeRequest;
+  private void incrementStat(SegmentHolder holder, QueueStatus status)
+  {
+    stats.add(status.getStatForAction(holder.getAction()), 1);
+    if (status.datasourceStat != null) {
+      stats.addToDatasourceStat(status.datasourceStat, holder.getSegment().getDataSource(), 1);
     }
+  }
 
-    public boolean hasTimedOut()
-    {
-      if (scheduleTime < 0) {
-        scheduleTime = System.currentTimeMillis();
-        return false;
-      } else if (System.currentTimeMillis() - scheduleTime > config.getLoadTimeoutDelay().getMillis()) {
-        return true;
-      } else {
+  private void executeCallbacks(SegmentHolder holder, boolean success)
+  {
+    callBackExecutor.execute(() -> {
+      for (LoadPeonCallback callback : holder.getCallbacks()) {
+        callback.execute(success);
+      }
+    });
+  }
+
+  /**
+   * Tries to cancel a load/drop operation. An load/drop request can be cancelled
+   * only if it has not already been sent to the corresponding server.
+   */
+  @Override
+  public boolean cancelOperation(DataSegment segment)
+  {
+    synchronized (lock) {
+      if (activeRequestSegments.contains(segment)) {
         return false;
       }
-    }
 
-    public void requestSucceeded()
-    {
-      log.trace(
-          "Server[%s] Successfully processed segment[%s] request[%s].",
-          serverId,
-          segment.getId(),
-          changeRequest.getClass().getSimpleName()
-      );
+      // Find the action on this segment, if any
+      final SegmentHolder holder = segmentsToLoad.containsKey(segment)
+                                   ? segmentsToLoad.remove(segment)
+                                   : segmentsToDrop.remove(segment);
+      if (holder == null) {
+        return false;
+      }
 
-      callBackExecutor.execute(() -> {
-        for (LoadPeonCallback callback : callbacks) {
-          if (callback != null) {
-            callback.execute(true);
-          }
-        }
-      });
+      queuedSegments.remove(holder);
+      onRequestCompleted(holder, QueueStatus.CANCELLED);
+      return true;
     }
+  }
 
-    public void requestFailed(String failureCause)
-    {
-      log.error(
-          "Server[%s] Failed segment[%s] request[%s] with cause [%s].",
-          serverId,
-          segment.getId(),
-          changeRequest.getClass().getSimpleName(),
-          failureCause
-      );
-
-      failedAssignCount.getAndIncrement();
+  private enum QueueStatus
+  {
+    ASSIGNED(Stats.SegmentQueue.ASSIGNED_ACTIONS),
+    SUCCESS(Stats.SegmentQueue.COMPLETED_ACTIONS),
+    FAILED(Stats.SegmentQueue.FAILED_ACTIONS),
+    CANCELLED(Stats.SegmentQueue.CANCELLED_ACTIONS);
 
-      callBackExecutor.execute(() -> {
-        for (LoadPeonCallback callback : callbacks) {
-          if (callback != null) {
-            callback.execute(false);
-          }
-        }
-      });
-    }
+    final CoordinatorStat loadStat;
+    final CoordinatorStat moveStat;
+    final CoordinatorStat dropStat;
+    final CoordinatorStat datasourceStat;
 
-    @Override
-    public String toString()
+    QueueStatus()
     {
-      return changeRequest.toString();
+      this(null);
     }
-  }
 
-  private class LoadSegmentHolder extends SegmentHolder
-  {
-    public LoadSegmentHolder(DataSegment segment, LoadPeonCallback callback)
+    QueueStatus(CoordinatorStat datasourceStat)
     {
-      super(segment, new SegmentChangeRequestLoad(segment), callback);
-      queuedSize.addAndGet(segment.getSize());
-    }
+      // These stats are not emitted and are tracked for debugging purposes only
+      final String prefix = StringUtils.toLowerCase(name());
+      this.loadStat = new CoordinatorStat(prefix + "Load");
+      this.moveStat = new CoordinatorStat(prefix + "Move");
+      this.dropStat = new CoordinatorStat(prefix + "Drop");
 
-    @Override
-    public void requestSucceeded()
-    {
-      queuedSize.addAndGet(-getSegment().getSize());
-      super.requestSucceeded();
+      this.datasourceStat = datasourceStat;
     }
 
-    @Override
-    public void requestFailed(String failureCause)
+    CoordinatorStat getStatForAction(SegmentAction action)
     {
-      queuedSize.addAndGet(-getSegment().getSize());
-      super.requestFailed(failureCause);
-    }
-  }
+      switch (action) {

Review Comment:
   This codeQL seems legitimate?



##########
server/src/main/java/org/apache/druid/server/coordinator/SegmentReplicantLookup.java:
##########
@@ -125,49 +117,85 @@ public int getMovingReplicas(SegmentId segmentId, String tier)
    * Number of replicas of the segment which are safely loaded on the given tier
    * and are not being dropped.
    */
-  public int getLoadedReplicas(SegmentId segmentId, String tier)
+  public int getLoadedNotDroppingReplicas(SegmentId segmentId, String tier)
+  {
+    ReplicaCount count = replicaCounts.get(segmentId, tier);
+    return (count == null) ? 0 : count.loadedNotDropping();
+  }
+
+  public int getLoadingReplicas(SegmentId segmentId, String tier)
   {
     ReplicaCount count = replicaCounts.get(segmentId, tier);
-    return (count == null) ? 0 : count.safelyLoaded();
+    return count == null ? 0 : count.loading;
   }
 
   /**
-   * Number of replicas of the segment which are safely loaded on the cluster
-   * and are not being dropped.
+   * Number of replicas of the segment which are loaded on the cluster.
+   *
+   * @param includeDropping Whether segments which are being dropped should be
+   *                        included in the total count.
    */
-  public int getTotalLoadedReplicas(SegmentId segmentId)
+  public int getLoadedReplicas(SegmentId segmentId, boolean includeDropping)
   {
     final Map<String, ReplicaCount> allTiers = replicaCounts.row(segmentId);
     int totalLoaded = 0;
     for (ReplicaCount count : allTiers.values()) {
-      totalLoaded += count.safelyLoaded();
+      totalLoaded += includeDropping ? count.loaded : count.loadedNotDropping();
     }
     return totalLoaded;
   }

Review Comment:
   I wonder if it doesn't make sense to switch the signature of this to
   
   ```
   public ReplicaCount getReplicationStatus(SegmentId segmentId) {}
   ```
   
   And then in the method, create a `ReplicateCount` object that you aggregate all of the counts into.  This allows the caller to determine what they want to do with that information and the added overhead of adding the extra fields that aren't currently added right now is relatively minimal, I think.



##########
server/src/main/java/org/apache/druid/server/coordinator/SegmentReplicantLookup.java:
##########
@@ -125,49 +117,85 @@ public int getMovingReplicas(SegmentId segmentId, String tier)
    * Number of replicas of the segment which are safely loaded on the given tier
    * and are not being dropped.
    */
-  public int getLoadedReplicas(SegmentId segmentId, String tier)
+  public int getLoadedNotDroppingReplicas(SegmentId segmentId, String tier)
+  {
+    ReplicaCount count = replicaCounts.get(segmentId, tier);
+    return (count == null) ? 0 : count.loadedNotDropping();
+  }

Review Comment:
   There's a number of methods that are basically just getting the `ReplicaCount` and then reading a single field from it.  We can collapse them together and just return the `ReplicaCount` object.  Let the caller decide which thing they want to read from it.



##########
server/src/main/java/org/apache/druid/server/coordinator/DruidCoordinator.java:
##########
@@ -888,26 +799,101 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params)
       List<ImmutableDruidServer> currentServers = prepareCurrentServers();
 
       startPeonsForNewServers(currentServers);
+      stopPeonsForDisappearedServers(currentServers);
 
-      cluster = prepareCluster(params, currentServers);
-      segmentReplicantLookup = SegmentReplicantLookup.make(cluster, getDynamicConfigs().getReplicateAfterLoadTimeout());
+      final DruidCluster cluster = prepareCluster(params.getCoordinatorDynamicConfig(), currentServers);
+      cancelLoadsOnDecommissioningServers(cluster);
 
-      stopPeonsForDisappearedServers(currentServers);
+      final CoordinatorDynamicConfig dynamicConfig = params.getCoordinatorDynamicConfig();
 
-      final RoundRobinServerSelector roundRobinServerSelector;
-      if (params.getCoordinatorDynamicConfig().isUseRoundRobinSegmentAssignment()) {
-        roundRobinServerSelector = new RoundRobinServerSelector(cluster);
-        log.info("Using round-robin segment assignment.");
-      } else {
-        roundRobinServerSelector = null;
+      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()
+      );
+
+      params = params.buildFromExisting()
+                     .withDruidCluster(cluster)
+                     .withDynamicConfigs(recomputeDynamicConfig(params))
+                     .withBalancerStrategy(balancerStrategy)
+                     .withSegmentAssignerUsing(loadQueueManager)
+                     .build();
+
+      segmentReplicantLookup = params.getSegmentReplicantLookup();
+
+      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;
       }
 
-      return params.buildFromExisting()
-                   .withDruidCluster(cluster)
-                   .withLoadManagementPeons(loadManagementPeons)
-                   .withSegmentReplicantLookup(segmentReplicantLookup)
-                   .withRoundRobinServerSelector(roundRobinServerSelector)
-                   .build();
+      // 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)

Review Comment:
   Instead of recomputing the dynamic config, I would've preferred that there be a separate config which houses all of the things that are recomputed and that reference gets updated/changed.  The `DynamicConfig` object itself should really only be mutated by users changing config values.



##########
server/src/main/java/org/apache/druid/server/coordinator/CoordinatorDynamicConfig.java:
##########
@@ -428,66 +489,36 @@ public boolean equals(Object o)
 
     CoordinatorDynamicConfig that = (CoordinatorDynamicConfig) o;
 
-    if (leadingTimeMillisBeforeCanMarkAsUnusedOvershadowedSegments !=
-        that.leadingTimeMillisBeforeCanMarkAsUnusedOvershadowedSegments) {
-      return false;
-    }
-    if (mergeBytesLimit != that.mergeBytesLimit) {
-      return false;
-    }
-    if (mergeSegmentsLimit != that.mergeSegmentsLimit) {
-      return false;
-    }
-    if (maxSegmentsToMove != that.maxSegmentsToMove) {
-      return false;
-    }
-    if (percentOfSegmentsToConsiderPerMove != that.percentOfSegmentsToConsiderPerMove) {
-      return false;
-    }
-    if (useBatchedSegmentSampler != that.useBatchedSegmentSampler) {
-      return false;
-    }
-    if (replicantLifetime != that.replicantLifetime) {
-      return false;
-    }
-    if (replicationThrottleLimit != that.replicationThrottleLimit) {
-      return false;
-    }
-    if (balancerComputeThreads != that.balancerComputeThreads) {
-      return false;
-    }
-    if (emitBalancingStats != that.emitBalancingStats) {
-      return false;
-    }
-    if (maxSegmentsInNodeLoadingQueue != that.maxSegmentsInNodeLoadingQueue) {
-      return false;
-    }
-    if (!Objects.equals(specificDataSourcesToKillUnusedSegmentsIn, that.specificDataSourcesToKillUnusedSegmentsIn)) {
-      return false;
-    }
-    if (!Objects.equals(dataSourcesToNotKillStalePendingSegmentsIn, that.dataSourcesToNotKillStalePendingSegmentsIn)) {
-      return false;
-    }
-    if (!Objects.equals(decommissioningNodes, that.decommissioningNodes)) {
-      return false;
-    }
-    if (pauseCoordination != that.pauseCoordination) {
-      return false;
-    }
-    if (replicateAfterLoadTimeout != that.replicateAfterLoadTimeout) {
-      return false;
-    }
-    if (maxNonPrimaryReplicantsToLoad != that.maxNonPrimaryReplicantsToLoad) {
-      return false;
-    }
-    return decommissioningMaxPercentOfMaxSegmentsToMove == that.decommissioningMaxPercentOfMaxSegmentsToMove;
+    return markSegmentAsUnusedDelayMillis == that.markSegmentAsUnusedDelayMillis
+           && mergeBytesLimit == that.mergeBytesLimit
+           && mergeSegmentsLimit == that.mergeSegmentsLimit
+           && maxSegmentsToMove == that.maxSegmentsToMove
+           && percentOfSegmentsToConsiderPerMove == that.percentOfSegmentsToConsiderPerMove
+           && decommissioningMaxPercentOfMaxSegmentsToMove == that.decommissioningMaxPercentOfMaxSegmentsToMove
+           && useBatchedSegmentSampler == that.useBatchedSegmentSampler
+           && balancerComputeThreads == that.balancerComputeThreads
+           && emitBalancingStats == that.emitBalancingStats
+           && replicantLifetime == that.replicantLifetime
+           && replicationThrottleLimit == that.replicationThrottleLimit
+           && replicateAfterLoadTimeout == that.replicateAfterLoadTimeout
+           && maxSegmentsInNodeLoadingQueue == that.maxSegmentsInNodeLoadingQueue
+           && maxNonPrimaryReplicantsToLoad == that.maxNonPrimaryReplicantsToLoad
+           && useRoundRobinSegmentAssignment == that.useRoundRobinSegmentAssignment
+           && pauseCoordination == that.pauseCoordination
+           && Objects.equals(
+               specificDataSourcesToKillUnusedSegmentsIn,
+               that.specificDataSourcesToKillUnusedSegmentsIn)
+           && Objects.equals(
+               dataSourcesToNotKillStalePendingSegmentsIn,
+               that.dataSourcesToNotKillStalePendingSegmentsIn)
+           && Objects.equals(decommissioningNodes, that.decommissioningNodes);

Review Comment:
   This (and hashcode) appear to be ignoring the various debugDimensions things, is that intentional?



##########
server/src/main/java/org/apache/druid/server/coordinator/DruidCoordinator.java:
##########
@@ -888,26 +799,101 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params)
       List<ImmutableDruidServer> currentServers = prepareCurrentServers();
 
       startPeonsForNewServers(currentServers);
+      stopPeonsForDisappearedServers(currentServers);
 
-      cluster = prepareCluster(params, currentServers);
-      segmentReplicantLookup = SegmentReplicantLookup.make(cluster, getDynamicConfigs().getReplicateAfterLoadTimeout());
+      final DruidCluster cluster = prepareCluster(params.getCoordinatorDynamicConfig(), currentServers);
+      cancelLoadsOnDecommissioningServers(cluster);
 
-      stopPeonsForDisappearedServers(currentServers);
+      final CoordinatorDynamicConfig dynamicConfig = params.getCoordinatorDynamicConfig();
 
-      final RoundRobinServerSelector roundRobinServerSelector;
-      if (params.getCoordinatorDynamicConfig().isUseRoundRobinSegmentAssignment()) {
-        roundRobinServerSelector = new RoundRobinServerSelector(cluster);
-        log.info("Using round-robin segment assignment.");
-      } else {
-        roundRobinServerSelector = null;
+      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()
+      );
+
+      params = params.buildFromExisting()
+                     .withDruidCluster(cluster)
+                     .withDynamicConfigs(recomputeDynamicConfig(params))
+                     .withBalancerStrategy(balancerStrategy)
+                     .withSegmentAssignerUsing(loadQueueManager)
+                     .build();
+
+      segmentReplicantLookup = params.getSegmentReplicantLookup();
+
+      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;
       }
 
-      return params.buildFromExisting()
-                   .withDruidCluster(cluster)
-                   .withLoadManagementPeons(loadManagementPeons)
-                   .withSegmentReplicantLookup(segmentReplicantLookup)
-                   .withRoundRobinServerSelector(roundRobinServerSelector)
-                   .build();
+      // 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
+     * under-replicated segments can be assigned in the current run itself.
+     */
+    private void cancelLoadsOnDecommissioningServers(DruidCluster cluster)
+    {
+      final AtomicInteger cancelledCount = new AtomicInteger(0);
+      final List<ServerHolder> decommissioningServers
+          = cluster.getAllServers().stream()
+                   .filter(ServerHolder::isDecommissioning)
+                   .collect(Collectors.toList());
+
+      for (ServerHolder server : decommissioningServers) {
+        server.getQueuedSegments().forEach(
+            (segment, action) -> {
+              // Cancel the operation if it is a type of load
+              if (action.isLoad() && server.cancelOperation(action, segment)) {
+                cancelledCount.incrementAndGet();
+              }
+            }
+        );
+      }
+
+      if (cancelledCount.get() > 0) {
+        log.info(
+            "Cancelled [%d] load/move operations on [%d] decommissioning servers.",
+            cancelledCount.get(), decommissioningServers.size()
+        );
+      }
     }

Review Comment:
   I would've expected that the act of decommissioning would put the queue into a state where it will only accept DROP requests and then dropped the load queue immediately.  It seems weird to me that this is logic handled in this class instead.  The only thing I'd expect a CoordinatorDuty to do about it is to look for decomissioning servers and move the segments away.



##########
server/src/main/java/org/apache/druid/server/coordinator/SegmentReplicantLookup.java:
##########
@@ -125,49 +117,85 @@ public int getMovingReplicas(SegmentId segmentId, String tier)
    * Number of replicas of the segment which are safely loaded on the given tier
    * and are not being dropped.
    */
-  public int getLoadedReplicas(SegmentId segmentId, String tier)
+  public int getLoadedNotDroppingReplicas(SegmentId segmentId, String tier)
+  {
+    ReplicaCount count = replicaCounts.get(segmentId, tier);
+    return (count == null) ? 0 : count.loadedNotDropping();
+  }
+
+  public int getLoadingReplicas(SegmentId segmentId, String tier)
   {
     ReplicaCount count = replicaCounts.get(segmentId, tier);
-    return (count == null) ? 0 : count.safelyLoaded();
+    return count == null ? 0 : count.loading;
   }
 
   /**
-   * Number of replicas of the segment which are safely loaded on the cluster
-   * and are not being dropped.
+   * Number of replicas of the segment which are loaded on the cluster.
+   *
+   * @param includeDropping Whether segments which are being dropped should be
+   *                        included in the total count.
    */
-  public int getTotalLoadedReplicas(SegmentId segmentId)
+  public int getLoadedReplicas(SegmentId segmentId, boolean includeDropping)
   {
     final Map<String, ReplicaCount> allTiers = replicaCounts.row(segmentId);
     int totalLoaded = 0;
     for (ReplicaCount count : allTiers.values()) {
-      totalLoaded += count.safelyLoaded();
+      totalLoaded += includeDropping ? count.loaded : count.loadedNotDropping();
     }
     return totalLoaded;
   }
 
-  public Object2LongMap<String> getBroadcastUnderReplication(SegmentId segmentId)
+  /**
+   * Sets the number of replicas required for the specified segment in the tier.
+   * In a given coordinator run, this method must be called atleast once for
+   * every segment every tier.
+   */
+  public void setRequiredReplicas(SegmentId segmentId, boolean isBroadcast, String tier, int requiredReplicas)
   {
-    Object2LongOpenHashMap<String> perTier = new Object2LongOpenHashMap<>();
-    for (ServerHolder holder : cluster.getAllServers()) {
-      // Only record tier entry for server that is segment broadcast target
-      if (holder.getServer().getType().isSegmentBroadcastTarget()) {
-        // Every broadcast target server should be serving 1 replica of the segment
-        if (!holder.isServingSegment(segmentId)) {
-          perTier.addTo(holder.getServer().getTier(), 1L);
-        } else {
-          perTier.putIfAbsent(holder.getServer().getTier(), 0);
-        }
+    ReplicaCount counts = computeIfAbsent(replicaCounts, segmentId, tier);
+    counts.required = requiredReplicas;
+    if (isBroadcast) {
+      counts.possible = requiredReplicas;
+    } else {
+      counts.possible = tierToHistoricalCount.getOrDefault(tier, 0);
+    }
+  }
+
+  public Map<String, Object2LongMap<String>> getTierToDatasourceToUnderReplicated(
+      Iterable<DataSegment> usedSegments,
+      boolean ignoreMissingServers
+  )
+  {
+    final Map<String, Object2LongMap<String>> tierToUnderReplicated = new HashMap<>();
+
+    for (DataSegment segment : usedSegments) {
+      final Map<String, ReplicaCount> tierToReplicaCount = replicaCounts.row(segment.getId());
+      if (tierToReplicaCount == null) {
+        continue;
       }
+
+      tierToReplicaCount.forEach((tier, counts) -> {
+        final int underReplicated = counts.underReplicated(ignoreMissingServers);
+        if (underReplicated >= 0) {
+          Object2LongOpenHashMap<String> datasourceToUnderReplicated = (Object2LongOpenHashMap<String>)
+              tierToUnderReplicated.computeIfAbsent(tier, ds -> new Object2LongOpenHashMap<>());
+          datasourceToUnderReplicated.addTo(segment.getDataSource(), underReplicated);
+        }
+      });
     }
-    return perTier;
+
+    return tierToUnderReplicated;
   }

Review Comment:
   You might've noticed by now, but any time I see a `boolean` parameter in a function signature, I go and try to figure out why it exists and what to do to remove it.  Most of the time, it's a signal to some code that needs to be re-designed.
   
   This case is no different.  That said, I'm a little less confident about this suggestion, but I *think* that if this method were to compute and return a `Map<String, Map<String, ReplicaCount>>` that was 
   
   ```
   tier -> datasource -> total replication status of that datasource across that tier
   ```
   
   Then that would be enough to effectively eliminate the boolean property from this method.  I think it might also allow for you to compute the value once and reuse it for other purposes as well.
   
   Given that I'm a bit less confident about this one, treat it as just a suggestion.



##########
server/src/main/java/org/apache/druid/server/coordinator/StrategicSegmentAssigner.java:
##########
@@ -0,0 +1,577 @@
+/*
+ * 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 com.google.common.collect.Sets;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+import org.apache.druid.java.util.emitter.EmittingLogger;
+import org.apache.druid.server.coordinator.balancer.BalancerStrategy;
+import org.apache.druid.server.coordinator.loadqueue.SegmentAction;
+import org.apache.druid.server.coordinator.loadqueue.SegmentLoadQueueManager;
+import org.apache.druid.server.coordinator.rules.SegmentActionHandler;
+import org.apache.druid.server.coordinator.stats.CoordinatorRunStats;
+import org.apache.druid.server.coordinator.stats.CoordinatorStat;
+import org.apache.druid.server.coordinator.stats.Dimension;
+import org.apache.druid.server.coordinator.stats.RowKey;
+import org.apache.druid.server.coordinator.stats.Stats;
+import org.apache.druid.timeline.DataSegment;
+
+import java.util.Comparator;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
+/**
+ * Used by the coordinator in each run for segment loading, dropping, balancing
+ * and broadcasting.
+ * <p>
+ * An instance of this class is freshly created for each coordinator run.
+ */
+public class StrategicSegmentAssigner implements SegmentActionHandler
+{
+  private static final EmittingLogger log = new EmittingLogger(StrategicSegmentAssigner.class);
+
+  private final SegmentLoadQueueManager loadQueueManager;
+  private final DruidCluster cluster;
+  private final CoordinatorRunStats stats;
+  private final SegmentReplicantLookup replicantLookup;
+  private final ReplicationThrottler replicationThrottler;
+  private final RoundRobinServerSelector serverSelector;
+  private final BalancerStrategy strategy;
+
+  private final boolean useRoundRobinAssignment;
+
+  private final Set<String> tiersWithNoServer = new HashSet<>();
+
+  public StrategicSegmentAssigner(
+      SegmentLoadQueueManager loadQueueManager,
+      DruidCluster cluster,
+      BalancerStrategy strategy,
+      CoordinatorDynamicConfig dynamicConfig,
+      CoordinatorRunStats stats
+  )
+  {
+    this.stats = stats;
+    this.cluster = cluster;
+    this.strategy = strategy;
+    this.loadQueueManager = loadQueueManager;
+    this.replicantLookup = SegmentReplicantLookup.make(cluster);
+    this.replicationThrottler = createReplicationThrottler(dynamicConfig);
+    this.useRoundRobinAssignment = dynamicConfig.isUseRoundRobinSegmentAssignment();
+    this.serverSelector = useRoundRobinAssignment ? new RoundRobinServerSelector(cluster) : null;
+  }
+
+  public CoordinatorRunStats getStats()
+  {
+    return stats;
+  }
+
+  public SegmentReplicantLookup getReplicantLookup()
+  {
+    return replicantLookup;
+  }
+
+  public void makeAlerts()
+  {
+    if (!tiersWithNoServer.isEmpty()) {
+      log.makeAlert("Tiers [%s] have no servers! Check your cluster configuration.", tiersWithNoServer).emit();
+    }
+  }
+
+  /**
+   * Moves the given segment from the source server to an eligible destination
+   * server.
+   * <p>
+   * An eligible destination server must:
+   * <ul>
+   *   <li>be present in the given list of destination servers</li>
+   *   <li>belong to the same tier as the source server</li>
+   *   <li>not already be serving or loading a replica of the segment</li>
+   *   <li>have enough space to load the segment</li>
+   * </ul>
+   * <p>
+   * The segment is not moved if:
+   * <ul>
+   *   <li>there is no eligible destination server, or</li>
+   *   <li>skipIfOptimallyPlaced is true and segment is already optimally placed, or</li>
+   *   <li>some other error occurs</li>
+   * </ul>
+   */
+  public boolean moveSegment(
+      DataSegment segment,
+      ServerHolder sourceServer,
+      List<ServerHolder> destinationServers
+  )
+  {
+    final String tier = sourceServer.getServer().getTier();
+    final List<ServerHolder> eligibleDestinationServers =
+        destinationServers.stream()
+                          .filter(s -> s.getServer().getTier().equals(tier))
+                          .filter(s -> s.canLoadSegment(segment))
+                          .collect(Collectors.toList());
+
+    if (eligibleDestinationServers.isEmpty()) {
+      incrementSkipStat(Stats.Segments.MOVE_SKIPPED, "No eligible server", segment, tier);
+      return false;
+    }
+
+    // If the source server is not decommissioning, move can be skipped if the
+    // segment is already optimally placed
+    if (!sourceServer.isDecommissioning()) {
+      eligibleDestinationServers.add(sourceServer);
+    }
+
+    final ServerHolder destination =
+        strategy.findDestinationServerToMoveSegment(segment, sourceServer, eligibleDestinationServers);
+
+    if (destination == null || destination.getServer().equals(sourceServer.getServer())) {
+      incrementSkipStat(Stats.Segments.MOVE_SKIPPED, "Optimally placed", segment, tier);
+      return false;
+    } else if (moveSegment(segment, sourceServer, destination)) {
+      incrementStat(Stats.Segments.MOVED, segment, tier, 1);
+      return true;
+    } else {
+      incrementSkipStat(Stats.Segments.MOVE_SKIPPED, "Encountered error", segment, tier);
+      return false;
+    }
+  }
+
+  /**
+   * Moves the given segment from serverA to serverB.
+   */
+  private boolean moveSegment(DataSegment segment, ServerHolder serverA, ServerHolder serverB)
+  {
+    final String tier = serverA.getServer().getTier();
+    if (serverA.isLoadingSegment(segment)) {
+      // Cancel the load on serverA and load on serverB instead
+      if (serverA.cancelOperation(SegmentAction.LOAD, segment)) {
+        int loadedCountOnTier = replicantLookup.getLoadedNotDroppingReplicas(segment.getId(), tier);
+        return loadSegment(segment, serverB, loadedCountOnTier >= 1);
+      }
+
+      // Could not cancel load, let the segment load on serverA and count it as unmoved
+      return false;
+    } else if (serverA.isServingSegment(segment)) {
+      return loadQueueManager.moveSegment(segment, serverA, serverB);
+    } else {
+      return false;
+    }
+  }
+
+  @Override
+  public void updateSegmentReplicasInTiers(DataSegment segment, Map<String, Integer> tierToReplicaCount)
+  {
+    // Identify empty tiers and determine total required replicas
+    final AtomicInteger requiredTotalReplicas = new AtomicInteger(0);
+    final Set<String> allTiers = Sets.newHashSet(cluster.getTierNames());
+    tierToReplicaCount.forEach((tier, requiredReplicas) -> {
+      reportTierCapacityStats(segment, requiredReplicas, tier);
+      replicantLookup.setRequiredReplicas(segment.getId(), false, tier, requiredReplicas);
+      if (allTiers.contains(tier)) {
+        requiredTotalReplicas.addAndGet(requiredReplicas);
+      } else {
+        tiersWithNoServer.add(tier);
+      }
+    });
+
+    final int loadedNotDroppingReplicas = replicantLookup.getLoadedReplicas(segment.getId(), false);
+    final int replicaSurplus = loadedNotDroppingReplicas - requiredTotalReplicas.get();
+
+    // Update replicas in every tier
+    int dropsQueued = 0;
+    for (String tier : allTiers) {
+      dropsQueued += updateReplicasInTier(
+          segment,
+          tier,
+          tierToReplicaCount.getOrDefault(tier, 0),
+          replicaSurplus - dropsQueued
+      );
+    }
+  }
+
+  /**
+   * Queues load or drop operations on this tier based on the required
+   * number of replicas and the current state.
+   * <p>
+   * The {@code maxReplicasToDrop} helps to maintain the required level of
+   * replication in the cluster. This ensures that segment read concurrency does
+   * not suffer during a tier shift or load rule change.
+   * <p>
+   * Returns the number of new drop operations queued on this tier.
+   */
+  private int updateReplicasInTier(
+      DataSegment segment,
+      String tier,
+      int requiredReplicas,
+      int maxReplicasToDrop
+  )
+  {
+    final int loadedNotDroppingReplicas =
+        replicantLookup.getLoadedNotDroppingReplicas(segment.getId(), tier);
+    final int loadingReplicas = replicantLookup.getLoadingReplicas(segment.getId(), tier);
+    final int projectedReplicas = loadedNotDroppingReplicas + loadingReplicas;
+
+    final int movingReplicas = replicantLookup.getMovingReplicas(segment.getId(), tier);
+    final boolean shouldCancelMoves = requiredReplicas == 0 && movingReplicas > 0;
+
+    // Check if there is any action required on this tier
+    if (projectedReplicas == requiredReplicas && !shouldCancelMoves) {
+      return 0;
+    }
+
+    final SegmentStatusInTier segmentStatus =
+        new SegmentStatusInTier(segment, cluster.getHistoricalsByTier(tier));
+
+    // Cancel all moves in this tier if it does not need to have replicas
+    if (shouldCancelMoves) {
+      cancelOperations(SegmentAction.MOVE_TO, movingReplicas, segment, segmentStatus);
+      cancelOperations(SegmentAction.MOVE_FROM, movingReplicas, segment, segmentStatus);
+    }
+
+    // Cancel drops and queue loads if the projected count is below the requirement
+    if (projectedReplicas < requiredReplicas) {
+      int replicaDeficit = requiredReplicas - projectedReplicas;
+      int cancelledDrops =
+          cancelOperations(SegmentAction.DROP, replicaDeficit, segment, segmentStatus);
+
+      // Cancelled drops can be counted as loaded replicas, thus reducing deficit
+      int numReplicasToLoad = replicaDeficit - cancelledDrops;
+      if (numReplicasToLoad > 0) {
+        boolean isAlreadyLoadedOnTier = loadedNotDroppingReplicas + cancelledDrops >= 1;
+        int numLoadsQueued = loadReplicas(numReplicasToLoad, segment, tier, segmentStatus, isAlreadyLoadedOnTier);
+        incrementStat(Stats.Segments.ASSIGNED, segment, tier, numLoadsQueued);
+      }
+    }
+
+    // Cancel loads and queue drops if the projected count exceeds the requirement
+    if (projectedReplicas > requiredReplicas) {
+      int replicaSurplus = projectedReplicas - requiredReplicas;
+      int cancelledLoads =
+          cancelOperations(SegmentAction.LOAD, replicaSurplus, segment, segmentStatus);
+
+      int numReplicasToDrop = Math.min(replicaSurplus - cancelledLoads, maxReplicasToDrop);
+      if (numReplicasToDrop > 0) {
+        int dropsQueuedOnTier = dropReplicas(numReplicasToDrop, segment, tier, segmentStatus);
+        incrementStat(Stats.Segments.DROPPED, segment, tier, dropsQueuedOnTier);
+        return dropsQueuedOnTier;
+      }
+    }
+
+    return 0;
+  }
+
+  private void reportTierCapacityStats(DataSegment segment, int requiredReplicas, String tier)
+  {
+    final RowKey rowKey = RowKey.forTier(tier);
+    stats.updateMax(Stats.Tier.REPLICATION_FACTOR, rowKey, requiredReplicas);
+    stats.add(Stats.Tier.REQUIRED_CAPACITY, rowKey, segment.getSize() * requiredReplicas);
+  }
+
+  @Override
+  public void broadcastSegment(DataSegment segment)
+  {
+    final Object2IntOpenHashMap<String> tierToRequiredReplicas = new Object2IntOpenHashMap<>();
+    for (ServerHolder server : cluster.getAllServers()) {
+      // Ignore servers which are not broadcast targets
+      if (!server.getServer().getType().isSegmentBroadcastTarget()) {
+        continue;
+      }
+
+      final String tier = server.getServer().getTier();
+
+      // Drop from decommissioning servers and load on active servers
+      int numDropsQueued = 0;
+      int numLoadsQueued = 0;
+      if (server.isDecommissioning()) {
+        numDropsQueued += dropBroadcastSegment(segment, server) ? 1 : 0;
+      } else {
+        tierToRequiredReplicas.addTo(tier, 1);
+        numLoadsQueued += loadBroadcastSegment(segment, server) ? 1 : 0;
+      }
+
+      if (numLoadsQueued > 0) {
+        incrementStat(Stats.Segments.ASSIGNED, segment, tier, numLoadsQueued);
+      }
+      if (numDropsQueued > 0) {
+        incrementStat(Stats.Segments.DROPPED, segment, tier, numDropsQueued);
+      }
+    }
+
+    // Update required replica counts
+    tierToRequiredReplicas.object2IntEntrySet().fastForEach(
+        entry -> replicantLookup
+            .setRequiredReplicas(segment.getId(), true, entry.getKey(), entry.getIntValue())
+    );
+  }
+
+  @Override
+  public void deleteSegment(DataSegment segment)
+  {
+    loadQueueManager.deleteSegment(segment);
+    stats.addToDatasourceStat(Stats.Segments.DELETED, segment.getDataSource(), 1);
+  }
+
+  /**
+   * Loads the broadcast segment if it is not loaded on the given server.
+   * Returns true only if the segment was successfully queued for load on the server.
+   */
+  private boolean loadBroadcastSegment(DataSegment segment, ServerHolder server)
+  {
+    if (server.isServingSegment(segment) || server.isLoadingSegment(segment)) {
+      return false;
+    } else if (server.isDroppingSegment(segment)) {
+      return server.cancelOperation(SegmentAction.DROP, segment);
+    }
+
+    if (server.canLoadSegment(segment) && loadSegment(segment, server, false)) {
+      return true;
+    } else {
+      log.makeAlert("Could not assign broadcast segment for datasource [%s]", segment.getDataSource())
+         .addData("segmentId", segment.getId())
+         .addData("segmentSize", segment.getSize())
+         .addData("hostName", server.getServer().getHost())
+         .addData("availableSize", server.getAvailableSize())
+         .emit();
+      return false;
+    }
+  }
+
+  /**
+   * Drops the broadcast segment if it is loaded on the given server.
+   * Returns true only if the segment was successfully queued for drop on the server.
+   */
+  private boolean dropBroadcastSegment(DataSegment segment, ServerHolder server)
+  {
+    if (server.isLoadingSegment(segment)) {
+      return server.cancelOperation(SegmentAction.LOAD, segment);
+    } else if (server.isServingSegment(segment)) {
+      return loadQueueManager.dropSegment(segment, server);
+    } else {
+      return false;
+    }
+  }
+
+  /**
+   * Queues drop of {@code numToDrop} replicas of the segment from a tier.
+   * Tries to drop replicas first from decommissioning servers and then from
+   * active servers.
+   * <p>
+   * Returns the number of successfully queued drop operations.
+   */
+  private int dropReplicas(
+      final int numToDrop,
+      DataSegment segment,
+      String tier,
+      SegmentStatusInTier segmentStatus
+  )
+  {
+    if (numToDrop <= 0) {
+      return 0;
+    }
+
+    final List<ServerHolder> eligibleServers = segmentStatus.getServersEligibleToDrop();
+    if (eligibleServers.isEmpty()) {
+      incrementSkipStat(Stats.Segments.DROP_SKIPPED, "No eligible server", segment, tier);
+      return 0;
+    }
+
+    // Keep eligible servers sorted by most full first
+    final TreeSet<ServerHolder> eligibleLiveServers = new TreeSet<>(Comparator.reverseOrder());
+    final TreeSet<ServerHolder> eligibleDyingServers = new TreeSet<>(Comparator.reverseOrder());
+    for (ServerHolder server : eligibleServers) {
+      if (server.isDecommissioning()) {
+        eligibleDyingServers.add(server);
+      } else {
+        eligibleLiveServers.add(server);
+      }
+    }
+
+    // Drop as many replicas as possible from decommissioning servers
+    int remainingNumToDrop = numToDrop;
+    int numDropsQueued =
+        dropReplicasFromServers(remainingNumToDrop, segment, eligibleDyingServers.iterator(), tier);
+
+    // Drop replicas from active servers if required
+    if (numToDrop > numDropsQueued) {
+      remainingNumToDrop = numToDrop - numDropsQueued;
+      Iterator<ServerHolder> serverIterator =
+          (useRoundRobinAssignment || eligibleLiveServers.size() >= remainingNumToDrop)
+          ? eligibleLiveServers.iterator()
+          : strategy.pickServersToDropSegment(segment, eligibleLiveServers);
+      numDropsQueued += dropReplicasFromServers(remainingNumToDrop, segment, serverIterator, tier);
+    }
+
+    return numDropsQueued;
+  }
+
+  /**
+   * Queues drop of {@code numToDrop} replicas of the segment from the servers.
+   * Returns the number of successfully queued drop operations.
+   */
+  private int dropReplicasFromServers(
+      int numToDrop,
+      DataSegment segment,
+      Iterator<ServerHolder> serverIterator,
+      String tier
+  )
+  {
+    int numDropsQueued = 0;
+    while (numToDrop > numDropsQueued && serverIterator.hasNext()) {
+      ServerHolder holder = serverIterator.next();
+      boolean dropped = loadQueueManager.dropSegment(segment, holder);
+
+      if (dropped) {
+        ++numDropsQueued;
+      } else {
+        incrementSkipStat(Stats.Segments.DROP_SKIPPED, "Encountered error", segment, tier);
+      }
+    }
+
+    return numDropsQueued;
+  }
+
+  /**
+   * Queues load of {@code numToLoad} replicas of the segment on a tier.
+   */
+  private int loadReplicas(
+      int numToLoad,
+      DataSegment segment,
+      String tier,
+      SegmentStatusInTier segmentStatus,
+      boolean isAlreadyLoadedOnTier
+  )
+  {
+    // Do not assign replicas if tier is already busy loading some
+    if (isAlreadyLoadedOnTier && replicationThrottler.isTierLoadingReplicas(tier)) {
+      return 0;
+    }
+
+    final List<ServerHolder> eligibleServers = segmentStatus.getServersEligibleToLoad();
+    if (eligibleServers.isEmpty()) {
+      incrementSkipStat(Stats.Segments.ASSIGN_SKIPPED, "No eligible server", segment, tier);
+      return 0;
+    }
+
+    final Iterator<ServerHolder> serverIterator =
+        useRoundRobinAssignment
+        ? serverSelector.getServersInTierToLoadSegment(tier, segment)
+        : strategy.findServersToLoadSegment(segment, eligibleServers);
+    if (!serverIterator.hasNext()) {
+      incrementSkipStat(Stats.Segments.ASSIGN_SKIPPED, "No server chosen by strategy", segment, tier);
+      return 0;
+    }
+
+    // Load the replicas on this tier
+    int numLoadsQueued = 0;
+    while (numLoadsQueued < numToLoad && serverIterator.hasNext()) {
+      numLoadsQueued += loadSegment(segment, serverIterator.next(), isAlreadyLoadedOnTier)
+                        ? 1 : 0;
+    }
+
+    return numLoadsQueued;
+  }
+
+  private boolean loadSegment(DataSegment segment, ServerHolder server, boolean isAlreadyLoadedOnTier)
+  {
+    final String tier = server.getServer().getTier();
+    if (isAlreadyLoadedOnTier && !replicationThrottler.canAssignReplica(tier)) {
+      incrementSkipStat(Stats.Segments.ASSIGN_SKIPPED, "Throttled replication", segment, tier);
+      return false;
+    }
+
+    final SegmentAction action = isAlreadyLoadedOnTier ? SegmentAction.REPLICATE : SegmentAction.LOAD;
+    final boolean assigned = loadQueueManager.loadSegment(segment, server, action);
+
+    if (!assigned) {
+      incrementSkipStat(Stats.Segments.ASSIGN_SKIPPED, "Encountered error", segment, tier);
+    } else if (isAlreadyLoadedOnTier) {
+      replicationThrottler.incrementAssignedReplicas(tier);
+    }
+
+    return assigned;
+  }

Review Comment:
   boolean police again: is it really better to have this boolean and a bunch of checks strewn through this code than to just have a `replicateSegment` method and a `loadSegment` method?



##########
server/src/main/java/org/apache/druid/server/coordinator/StrategicSegmentAssigner.java:
##########
@@ -0,0 +1,577 @@
+/*
+ * 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 com.google.common.collect.Sets;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+import org.apache.druid.java.util.emitter.EmittingLogger;
+import org.apache.druid.server.coordinator.balancer.BalancerStrategy;
+import org.apache.druid.server.coordinator.loadqueue.SegmentAction;
+import org.apache.druid.server.coordinator.loadqueue.SegmentLoadQueueManager;
+import org.apache.druid.server.coordinator.rules.SegmentActionHandler;
+import org.apache.druid.server.coordinator.stats.CoordinatorRunStats;
+import org.apache.druid.server.coordinator.stats.CoordinatorStat;
+import org.apache.druid.server.coordinator.stats.Dimension;
+import org.apache.druid.server.coordinator.stats.RowKey;
+import org.apache.druid.server.coordinator.stats.Stats;
+import org.apache.druid.timeline.DataSegment;
+
+import java.util.Comparator;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
+/**
+ * Used by the coordinator in each run for segment loading, dropping, balancing
+ * and broadcasting.
+ * <p>
+ * An instance of this class is freshly created for each coordinator run.
+ */
+public class StrategicSegmentAssigner implements SegmentActionHandler
+{
+  private static final EmittingLogger log = new EmittingLogger(StrategicSegmentAssigner.class);
+
+  private final SegmentLoadQueueManager loadQueueManager;
+  private final DruidCluster cluster;
+  private final CoordinatorRunStats stats;
+  private final SegmentReplicantLookup replicantLookup;
+  private final ReplicationThrottler replicationThrottler;
+  private final RoundRobinServerSelector serverSelector;
+  private final BalancerStrategy strategy;
+
+  private final boolean useRoundRobinAssignment;
+
+  private final Set<String> tiersWithNoServer = new HashSet<>();
+
+  public StrategicSegmentAssigner(
+      SegmentLoadQueueManager loadQueueManager,
+      DruidCluster cluster,
+      BalancerStrategy strategy,
+      CoordinatorDynamicConfig dynamicConfig,
+      CoordinatorRunStats stats
+  )
+  {
+    this.stats = stats;
+    this.cluster = cluster;
+    this.strategy = strategy;
+    this.loadQueueManager = loadQueueManager;
+    this.replicantLookup = SegmentReplicantLookup.make(cluster);
+    this.replicationThrottler = createReplicationThrottler(dynamicConfig);
+    this.useRoundRobinAssignment = dynamicConfig.isUseRoundRobinSegmentAssignment();
+    this.serverSelector = useRoundRobinAssignment ? new RoundRobinServerSelector(cluster) : null;
+  }
+
+  public CoordinatorRunStats getStats()
+  {
+    return stats;
+  }
+
+  public SegmentReplicantLookup getReplicantLookup()
+  {
+    return replicantLookup;
+  }
+
+  public void makeAlerts()
+  {
+    if (!tiersWithNoServer.isEmpty()) {
+      log.makeAlert("Tiers [%s] have no servers! Check your cluster configuration.", tiersWithNoServer).emit();
+    }
+  }
+
+  /**
+   * Moves the given segment from the source server to an eligible destination
+   * server.
+   * <p>
+   * An eligible destination server must:
+   * <ul>
+   *   <li>be present in the given list of destination servers</li>
+   *   <li>belong to the same tier as the source server</li>
+   *   <li>not already be serving or loading a replica of the segment</li>
+   *   <li>have enough space to load the segment</li>
+   * </ul>
+   * <p>
+   * The segment is not moved if:
+   * <ul>
+   *   <li>there is no eligible destination server, or</li>
+   *   <li>skipIfOptimallyPlaced is true and segment is already optimally placed, or</li>
+   *   <li>some other error occurs</li>
+   * </ul>
+   */
+  public boolean moveSegment(
+      DataSegment segment,
+      ServerHolder sourceServer,
+      List<ServerHolder> destinationServers
+  )
+  {
+    final String tier = sourceServer.getServer().getTier();
+    final List<ServerHolder> eligibleDestinationServers =
+        destinationServers.stream()
+                          .filter(s -> s.getServer().getTier().equals(tier))
+                          .filter(s -> s.canLoadSegment(segment))
+                          .collect(Collectors.toList());
+
+    if (eligibleDestinationServers.isEmpty()) {
+      incrementSkipStat(Stats.Segments.MOVE_SKIPPED, "No eligible server", segment, tier);
+      return false;
+    }
+
+    // If the source server is not decommissioning, move can be skipped if the
+    // segment is already optimally placed
+    if (!sourceServer.isDecommissioning()) {
+      eligibleDestinationServers.add(sourceServer);
+    }
+
+    final ServerHolder destination =
+        strategy.findDestinationServerToMoveSegment(segment, sourceServer, eligibleDestinationServers);
+
+    if (destination == null || destination.getServer().equals(sourceServer.getServer())) {
+      incrementSkipStat(Stats.Segments.MOVE_SKIPPED, "Optimally placed", segment, tier);
+      return false;
+    } else if (moveSegment(segment, sourceServer, destination)) {
+      incrementStat(Stats.Segments.MOVED, segment, tier, 1);
+      return true;
+    } else {
+      incrementSkipStat(Stats.Segments.MOVE_SKIPPED, "Encountered error", segment, tier);
+      return false;
+    }
+  }
+
+  /**
+   * Moves the given segment from serverA to serverB.
+   */
+  private boolean moveSegment(DataSegment segment, ServerHolder serverA, ServerHolder serverB)
+  {
+    final String tier = serverA.getServer().getTier();
+    if (serverA.isLoadingSegment(segment)) {
+      // Cancel the load on serverA and load on serverB instead
+      if (serverA.cancelOperation(SegmentAction.LOAD, segment)) {
+        int loadedCountOnTier = replicantLookup.getLoadedNotDroppingReplicas(segment.getId(), tier);
+        return loadSegment(segment, serverB, loadedCountOnTier >= 1);
+      }
+
+      // Could not cancel load, let the segment load on serverA and count it as unmoved
+      return false;
+    } else if (serverA.isServingSegment(segment)) {
+      return loadQueueManager.moveSegment(segment, serverA, serverB);
+    } else {
+      return false;
+    }
+  }
+
+  @Override
+  public void updateSegmentReplicasInTiers(DataSegment segment, Map<String, Integer> tierToReplicaCount)
+  {
+    // Identify empty tiers and determine total required replicas
+    final AtomicInteger requiredTotalReplicas = new AtomicInteger(0);
+    final Set<String> allTiers = Sets.newHashSet(cluster.getTierNames());
+    tierToReplicaCount.forEach((tier, requiredReplicas) -> {
+      reportTierCapacityStats(segment, requiredReplicas, tier);
+      replicantLookup.setRequiredReplicas(segment.getId(), false, tier, requiredReplicas);
+      if (allTiers.contains(tier)) {
+        requiredTotalReplicas.addAndGet(requiredReplicas);
+      } else {
+        tiersWithNoServer.add(tier);
+      }
+    });
+
+    final int loadedNotDroppingReplicas = replicantLookup.getLoadedReplicas(segment.getId(), false);
+    final int replicaSurplus = loadedNotDroppingReplicas - requiredTotalReplicas.get();
+
+    // Update replicas in every tier
+    int dropsQueued = 0;
+    for (String tier : allTiers) {
+      dropsQueued += updateReplicasInTier(
+          segment,
+          tier,
+          tierToReplicaCount.getOrDefault(tier, 0),
+          replicaSurplus - dropsQueued
+      );
+    }
+  }
+
+  /**
+   * Queues load or drop operations on this tier based on the required
+   * number of replicas and the current state.
+   * <p>
+   * The {@code maxReplicasToDrop} helps to maintain the required level of
+   * replication in the cluster. This ensures that segment read concurrency does
+   * not suffer during a tier shift or load rule change.
+   * <p>
+   * Returns the number of new drop operations queued on this tier.
+   */
+  private int updateReplicasInTier(
+      DataSegment segment,
+      String tier,
+      int requiredReplicas,
+      int maxReplicasToDrop
+  )
+  {
+    final int loadedNotDroppingReplicas =
+        replicantLookup.getLoadedNotDroppingReplicas(segment.getId(), tier);
+    final int loadingReplicas = replicantLookup.getLoadingReplicas(segment.getId(), tier);
+    final int projectedReplicas = loadedNotDroppingReplicas + loadingReplicas;
+
+    final int movingReplicas = replicantLookup.getMovingReplicas(segment.getId(), tier);
+    final boolean shouldCancelMoves = requiredReplicas == 0 && movingReplicas > 0;
+
+    // Check if there is any action required on this tier
+    if (projectedReplicas == requiredReplicas && !shouldCancelMoves) {
+      return 0;
+    }
+
+    final SegmentStatusInTier segmentStatus =
+        new SegmentStatusInTier(segment, cluster.getHistoricalsByTier(tier));
+
+    // Cancel all moves in this tier if it does not need to have replicas
+    if (shouldCancelMoves) {
+      cancelOperations(SegmentAction.MOVE_TO, movingReplicas, segment, segmentStatus);
+      cancelOperations(SegmentAction.MOVE_FROM, movingReplicas, segment, segmentStatus);
+    }
+
+    // Cancel drops and queue loads if the projected count is below the requirement
+    if (projectedReplicas < requiredReplicas) {
+      int replicaDeficit = requiredReplicas - projectedReplicas;
+      int cancelledDrops =
+          cancelOperations(SegmentAction.DROP, replicaDeficit, segment, segmentStatus);
+
+      // Cancelled drops can be counted as loaded replicas, thus reducing deficit
+      int numReplicasToLoad = replicaDeficit - cancelledDrops;
+      if (numReplicasToLoad > 0) {
+        boolean isAlreadyLoadedOnTier = loadedNotDroppingReplicas + cancelledDrops >= 1;
+        int numLoadsQueued = loadReplicas(numReplicasToLoad, segment, tier, segmentStatus, isAlreadyLoadedOnTier);
+        incrementStat(Stats.Segments.ASSIGNED, segment, tier, numLoadsQueued);
+      }
+    }
+
+    // Cancel loads and queue drops if the projected count exceeds the requirement
+    if (projectedReplicas > requiredReplicas) {
+      int replicaSurplus = projectedReplicas - requiredReplicas;
+      int cancelledLoads =
+          cancelOperations(SegmentAction.LOAD, replicaSurplus, segment, segmentStatus);
+
+      int numReplicasToDrop = Math.min(replicaSurplus - cancelledLoads, maxReplicasToDrop);
+      if (numReplicasToDrop > 0) {
+        int dropsQueuedOnTier = dropReplicas(numReplicasToDrop, segment, tier, segmentStatus);
+        incrementStat(Stats.Segments.DROPPED, segment, tier, dropsQueuedOnTier);
+        return dropsQueuedOnTier;
+      }
+    }
+
+    return 0;
+  }
+
+  private void reportTierCapacityStats(DataSegment segment, int requiredReplicas, String tier)
+  {
+    final RowKey rowKey = RowKey.forTier(tier);
+    stats.updateMax(Stats.Tier.REPLICATION_FACTOR, rowKey, requiredReplicas);
+    stats.add(Stats.Tier.REQUIRED_CAPACITY, rowKey, segment.getSize() * requiredReplicas);
+  }
+
+  @Override
+  public void broadcastSegment(DataSegment segment)
+  {
+    final Object2IntOpenHashMap<String> tierToRequiredReplicas = new Object2IntOpenHashMap<>();
+    for (ServerHolder server : cluster.getAllServers()) {
+      // Ignore servers which are not broadcast targets
+      if (!server.getServer().getType().isSegmentBroadcastTarget()) {
+        continue;
+      }
+
+      final String tier = server.getServer().getTier();
+
+      // Drop from decommissioning servers and load on active servers
+      int numDropsQueued = 0;
+      int numLoadsQueued = 0;
+      if (server.isDecommissioning()) {
+        numDropsQueued += dropBroadcastSegment(segment, server) ? 1 : 0;
+      } else {
+        tierToRequiredReplicas.addTo(tier, 1);
+        numLoadsQueued += loadBroadcastSegment(segment, server) ? 1 : 0;
+      }
+
+      if (numLoadsQueued > 0) {
+        incrementStat(Stats.Segments.ASSIGNED, segment, tier, numLoadsQueued);
+      }
+      if (numDropsQueued > 0) {
+        incrementStat(Stats.Segments.DROPPED, segment, tier, numDropsQueued);
+      }
+    }
+
+    // Update required replica counts
+    tierToRequiredReplicas.object2IntEntrySet().fastForEach(
+        entry -> replicantLookup
+            .setRequiredReplicas(segment.getId(), true, entry.getKey(), entry.getIntValue())
+    );
+  }
+
+  @Override
+  public void deleteSegment(DataSegment segment)
+  {
+    loadQueueManager.deleteSegment(segment);
+    stats.addToDatasourceStat(Stats.Segments.DELETED, segment.getDataSource(), 1);
+  }
+
+  /**
+   * Loads the broadcast segment if it is not loaded on the given server.
+   * Returns true only if the segment was successfully queued for load on the server.
+   */
+  private boolean loadBroadcastSegment(DataSegment segment, ServerHolder server)
+  {
+    if (server.isServingSegment(segment) || server.isLoadingSegment(segment)) {
+      return false;
+    } else if (server.isDroppingSegment(segment)) {
+      return server.cancelOperation(SegmentAction.DROP, segment);
+    }
+
+    if (server.canLoadSegment(segment) && loadSegment(segment, server, false)) {
+      return true;
+    } else {
+      log.makeAlert("Could not assign broadcast segment for datasource [%s]", segment.getDataSource())
+         .addData("segmentId", segment.getId())
+         .addData("segmentSize", segment.getSize())
+         .addData("hostName", server.getServer().getHost())
+         .addData("availableSize", server.getAvailableSize())
+         .emit();
+      return false;
+    }
+  }
+
+  /**
+   * Drops the broadcast segment if it is loaded on the given server.
+   * Returns true only if the segment was successfully queued for drop on the server.
+   */
+  private boolean dropBroadcastSegment(DataSegment segment, ServerHolder server)
+  {
+    if (server.isLoadingSegment(segment)) {
+      return server.cancelOperation(SegmentAction.LOAD, segment);
+    } else if (server.isServingSegment(segment)) {
+      return loadQueueManager.dropSegment(segment, server);
+    } else {
+      return false;
+    }
+  }
+
+  /**
+   * Queues drop of {@code numToDrop} replicas of the segment from a tier.
+   * Tries to drop replicas first from decommissioning servers and then from
+   * active servers.
+   * <p>
+   * Returns the number of successfully queued drop operations.
+   */
+  private int dropReplicas(
+      final int numToDrop,
+      DataSegment segment,
+      String tier,
+      SegmentStatusInTier segmentStatus
+  )
+  {
+    if (numToDrop <= 0) {
+      return 0;
+    }
+
+    final List<ServerHolder> eligibleServers = segmentStatus.getServersEligibleToDrop();
+    if (eligibleServers.isEmpty()) {
+      incrementSkipStat(Stats.Segments.DROP_SKIPPED, "No eligible server", segment, tier);
+      return 0;
+    }
+
+    // Keep eligible servers sorted by most full first
+    final TreeSet<ServerHolder> eligibleLiveServers = new TreeSet<>(Comparator.reverseOrder());
+    final TreeSet<ServerHolder> eligibleDyingServers = new TreeSet<>(Comparator.reverseOrder());
+    for (ServerHolder server : eligibleServers) {
+      if (server.isDecommissioning()) {
+        eligibleDyingServers.add(server);
+      } else {
+        eligibleLiveServers.add(server);
+      }
+    }
+
+    // Drop as many replicas as possible from decommissioning servers
+    int remainingNumToDrop = numToDrop;
+    int numDropsQueued =
+        dropReplicasFromServers(remainingNumToDrop, segment, eligibleDyingServers.iterator(), tier);
+
+    // Drop replicas from active servers if required
+    if (numToDrop > numDropsQueued) {
+      remainingNumToDrop = numToDrop - numDropsQueued;
+      Iterator<ServerHolder> serverIterator =
+          (useRoundRobinAssignment || eligibleLiveServers.size() >= remainingNumToDrop)
+          ? eligibleLiveServers.iterator()
+          : strategy.pickServersToDropSegment(segment, eligibleLiveServers);
+      numDropsQueued += dropReplicasFromServers(remainingNumToDrop, segment, serverIterator, tier);
+    }
+
+    return numDropsQueued;
+  }
+
+  /**
+   * Queues drop of {@code numToDrop} replicas of the segment from the servers.
+   * Returns the number of successfully queued drop operations.
+   */
+  private int dropReplicasFromServers(
+      int numToDrop,
+      DataSegment segment,
+      Iterator<ServerHolder> serverIterator,
+      String tier
+  )
+  {
+    int numDropsQueued = 0;
+    while (numToDrop > numDropsQueued && serverIterator.hasNext()) {
+      ServerHolder holder = serverIterator.next();
+      boolean dropped = loadQueueManager.dropSegment(segment, holder);
+
+      if (dropped) {
+        ++numDropsQueued;
+      } else {
+        incrementSkipStat(Stats.Segments.DROP_SKIPPED, "Encountered error", segment, tier);
+      }
+    }
+
+    return numDropsQueued;
+  }
+
+  /**
+   * Queues load of {@code numToLoad} replicas of the segment on a tier.
+   */
+  private int loadReplicas(
+      int numToLoad,
+      DataSegment segment,
+      String tier,
+      SegmentStatusInTier segmentStatus,
+      boolean isAlreadyLoadedOnTier
+  )
+  {
+    // Do not assign replicas if tier is already busy loading some
+    if (isAlreadyLoadedOnTier && replicationThrottler.isTierLoadingReplicas(tier)) {
+      return 0;
+    }
+
+    final List<ServerHolder> eligibleServers = segmentStatus.getServersEligibleToLoad();
+    if (eligibleServers.isEmpty()) {
+      incrementSkipStat(Stats.Segments.ASSIGN_SKIPPED, "No eligible server", segment, tier);
+      return 0;
+    }
+
+    final Iterator<ServerHolder> serverIterator =
+        useRoundRobinAssignment
+        ? serverSelector.getServersInTierToLoadSegment(tier, segment)
+        : strategy.findServersToLoadSegment(segment, eligibleServers);
+    if (!serverIterator.hasNext()) {
+      incrementSkipStat(Stats.Segments.ASSIGN_SKIPPED, "No server chosen by strategy", segment, tier);
+      return 0;
+    }
+
+    // Load the replicas on this tier
+    int numLoadsQueued = 0;
+    while (numLoadsQueued < numToLoad && serverIterator.hasNext()) {
+      numLoadsQueued += loadSegment(segment, serverIterator.next(), isAlreadyLoadedOnTier)
+                        ? 1 : 0;
+    }
+
+    return numLoadsQueued;
+  }
+
+  private boolean loadSegment(DataSegment segment, ServerHolder server, boolean isAlreadyLoadedOnTier)
+  {
+    final String tier = server.getServer().getTier();
+    if (isAlreadyLoadedOnTier && !replicationThrottler.canAssignReplica(tier)) {
+      incrementSkipStat(Stats.Segments.ASSIGN_SKIPPED, "Throttled replication", segment, tier);
+      return false;
+    }
+
+    final SegmentAction action = isAlreadyLoadedOnTier ? SegmentAction.REPLICATE : SegmentAction.LOAD;
+    final boolean assigned = loadQueueManager.loadSegment(segment, server, action);
+
+    if (!assigned) {
+      incrementSkipStat(Stats.Segments.ASSIGN_SKIPPED, "Encountered error", segment, tier);
+    } else if (isAlreadyLoadedOnTier) {
+      replicationThrottler.incrementAssignedReplicas(tier);
+    }
+
+    return assigned;
+  }
+
+  private ReplicationThrottler createReplicationThrottler(CoordinatorDynamicConfig dynamicConfig)

Review Comment:
   Can this be static?  I think it can and maybe should be.  If it cannot be static, please rename to `initializeReplicationThrottler()` to help indicate that it is intended for initialization and not as a factory to be called multiple times from the same object.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org