fix: Prevent excessive segment assignment churn when using partial load rules

server/src/main/java/org/apache/druid/server/coordinator/loading/PartialSegmentStatusInTier.java

@@ -89,18 +89,19 @@ public List<ServerHolder> getStaleLoaded()
   }
 
   /**
-   * Servers with an in-flight load whose profile fingerprint matches the request. Counts toward projected matching
-   * replicas, the load is on its way to satisfying the rule.
+   * Servers with an in-flight load, or an incoming balancer move, whose profile fingerprint matches the request.
+   * Counts toward projected matching replicas, the load is on its way to satisfying the rule.
    */
   public List<ServerHolder> getMatchingInFlight()
   {
     return matchingInFlight;
   }
 
   /**
-   * Servers with an in-flight load whose profile fingerprint differs from the request (e.g., the rule changed
-   * mid-flight, or an in-flight regular full-load against a partial rule). Cancel-and-replace targets when there is a
-   * matching deficit.
+   * Servers with an in-flight load, or an incoming balancer move, whose profile fingerprint differs from the request
+   * (e.g., the rule changed mid-flight, or an in-flight regular full-load against a partial rule). Cancel-and-replace
+   * targets when there is a matching deficit; an incoming move refuses the cancellation and is instead reconciled
+   * once it lands and reclassifies as stale-loaded.
    */
   public List<ServerHolder> getStaleInFlight()
   {
@@ -130,15 +131,23 @@ public List<ServerHolder> getEligibleForAdditiveReload()
   /**
    * Mechanical classification of one server against the request fingerprint. Branches are mutually exclusive in
    * order: <b>loaded</b> ({@link ServerHolder#isServingSegment}: matching / stale, with stale optionally also added
-   * to {@link #eligibleForAdditiveReload}), <b>in-flight LOAD/REPLICATE</b> (matching / stale based on the peon's
-   * queued profile), <b>empty-and-loadable</b> ({@link #eligibleForFreshLoad}).
+   * to {@link #eligibleForAdditiveReload}), <b>in-flight LOAD/REPLICATE/MOVE_TO</b> (matching / stale based on the
+   * peon's queued profile), <b>empty-and-loadable</b> ({@link #eligibleForFreshLoad}).
    * <p>
-   * Servers with a queued {@link SegmentAction#DROP}, {@link SegmentAction#MOVE_TO}, or
-   * {@link SegmentAction#MOVE_FROM} fall through all branches by design, they're already accounted for in
+   * A balancer move is counted at its destination: the {@link SegmentAction#MOVE_TO} carries the profile cloned from
+   * the source, so it classifies by fingerprint like any other in-flight load, and once it lands the destination
+   * classifies as loaded. The source ({@link SegmentAction#MOVE_FROM}) is deliberately left unclassified, so that the
+   * two endpoints of a move count as the single replica they will settle into, in both phases of the move and
+   * whether the moving replica is matching or stale. Counting the move at the source instead would require the
+   * {@link SegmentReplicaCount#moveCompletedPendingDrop()} correction the full-load path uses, which cannot be
+   * applied to a fingerprint-partitioned count: it does not know whether the move it is netting out was of a
+   * matching or a stale replica.
+   * <p>
+   * Servers with a queued {@link SegmentAction#DROP} fall through all branches as well, they're accounted for in
    * {@link SegmentReplicaCount} totals and {@link StrategicSegmentAssigner}'s cross-tier drop budget.
-   * The {@code isLoaded} branch is in particular gated by {@link ServerHolder#isServingSegment}, which requires
-   * <em>no</em> action queued, so stale-loaded servers added to {@link #eligibleForAdditiveReload} are guaranteed
-   * to be action-free at snapshot time.
+   * The {@code isLoaded} branch is gated by {@link ServerHolder#isServingSegment}, which requires <em>no</em> action
+   * queued, so stale-loaded servers added to {@link #eligibleForAdditiveReload} are guaranteed to be action-free at
+   * snapshot time.
    */
   private void classify(ServerHolder server, DataSegment segment, String requestedFingerprint)
   {
@@ -155,7 +164,9 @@ private void classify(ServerHolder server, DataSegment segment, String requested
           eligibleForAdditiveReload.add(server);
         }
       }
-    } else if (action == SegmentAction.LOAD || action == SegmentAction.REPLICATE) {
+    } else if (action == SegmentAction.LOAD
+               || action == SegmentAction.REPLICATE
+               || action == SegmentAction.MOVE_TO) {
       final PartialLoadProfile inFlight = server.getInFlightProfile(segment);
       if (inFlight != null && Objects.equals(inFlight.fingerprint(), requestedFingerprint)) {
         matchingInFlight.add(server);

server/src/main/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssigner.java

@@ -345,10 +345,13 @@ public void replicateSegmentPartially(
    * <ol>
    *   <li><b>Classify.</b> Build {@link PartialSegmentStatusInTier} for this tier; every server falls into at most
    *       one of: matching-loaded, stale-loaded (optionally also eligible-for-additive-reload),
-   *       matching-in-flight, stale-in-flight, eligible-for-fresh-load, or unclassified (drop/move pending; see
-   *       {@link PartialSegmentStatusInTier#classify} for why). Matching means the announced fingerprint equals
-   *       this request's fingerprint; stale is anything else, including a non-profile regular full-load replica.</li>
-   *   <li><b>Compute matching count:</b> matching-loaded + matching-in-flight − pending-move-drop.</li>
+   *       matching-in-flight, stale-in-flight, eligible-for-fresh-load, or unclassified (drop or move source
+   *       pending; see {@link PartialSegmentStatusInTier#classify} for why). Matching means the announced
+   *       fingerprint equals this request's fingerprint; stale is anything else, including a non-profile regular
+   *       full-load replica. A balancer move counts once, at its destination.</li>
+   *   <li><b>Compute matching count:</b> matching-loaded + matching-in-flight. Unlike
+   *       {@link #updateReplicasInTier}, no {@link SegmentReplicaCount#moveCompletedPendingDrop()} correction is
+   *       needed, because the classification never counts both endpoints of a move.</li>
    *   <li><b>If matching count is short of {@code requiredReplicas}</b> (deficit):
    *     <ol type="a">
    *       <li>Cancel stale-in-flight loads to free their slots. Canceled servers become same-run fresh-load
@@ -389,7 +392,6 @@ private int updateReplicasInTierPartial(
   {
     final SegmentReplicaCount replicaCountOnTier = replicaCountMap.get(segment.getId(), tier);
     final int movingReplicas = replicaCountOnTier.moving();
-    final int moveCompletedPendingDrop = Math.max(0, replicaCountOnTier.moveCompletedPendingDrop());
 
     final PartialSegmentStatusInTier status = new PartialSegmentStatusInTier(
         segment,
@@ -398,8 +400,7 @@ private int updateReplicasInTierPartial(
     );
 
     final int matchingProjected = status.getMatchingLoaded().size()
-                                  + status.getMatchingInFlight().size()
-                                  - moveCompletedPendingDrop;
+                                  + status.getMatchingInFlight().size();
     final boolean shouldCancelMoves = requiredReplicas == 0 && movingReplicas > 0;
 
     // If everything's already in shape and no stale work, fast-exit.

server/src/test/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssignerPartialTest.java

@@ -50,8 +50,10 @@
 
 import javax.annotation.Nullable;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -471,6 +473,83 @@ public void testMoveOfFullLoadReplicaCarriesNoProfile()
     );
   }
 
+  @Test
+  public void testNoExtraReplicaWhileMoveIsInFlight()
+  {
+    // A move queued by a previous run is still in flight: the source is marked to drop (MOVE_FROM) and the
+    // destination has a MOVE_TO carrying the profile cloned from the source. The move is the tier's one replica, so
+    // the reconciler must leave the spare server alone; assigning to it here would make every balancer move churn
+    // an extra load and then a drop.
+    final DataSegment segment = createSegment();
+    final ServerHolder source = moveSourceOf(segment, profileForRevenue());
+    final ServerHolder destination = moveDestinationOf(segment, profileForRevenue().asCloneRequest());
+    final ServerHolder spare = createServer(TIER1);
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, source, destination, spare).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, segment);
+    params.getSegmentAssigner()
+          .replicateSegmentPartially(segment, profileForRevenue(), ImmutableMap.of(TIER1, 1));
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertFalse(stats.hasStat(Stats.Segments.PARTIAL_ASSIGNED));
+    Assert.assertFalse(stats.hasStat(Stats.Segments.DROPPED));
+    Assert.assertFalse(stats.hasStat(Stats.Segments.PARTIAL_STALE_DROPPED));
+    Assert.assertTrue(spare.getLoadingSegments().isEmpty());
+  }
+
+  @Test
+  public void testNoExtraReplicaWhenMoveIsPendingItsDrop()
+  {
+    // During a segment move, after the destination has finished loading and announced the matching fingerprint,
+    // the source is still marked to drop. The destination alone satisfies the replica requirement and no further action
+    // should be taken.
+    final DataSegment segment = createSegment();
+    final ServerHolder source = moveSourceOf(segment, profileForRevenue());
+    final ServerHolder destination = createServerWithLoaded(TIER1, segment, profileForRevenue());
+    final ServerHolder spare = createServer(TIER1);
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, source, destination, spare).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, segment);
+    params.getSegmentAssigner()
+          .replicateSegmentPartially(segment, profileForRevenue(), ImmutableMap.of(TIER1, 1));
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertFalse(stats.hasStat(Stats.Segments.PARTIAL_ASSIGNED));
+    Assert.assertFalse(stats.hasStat(Stats.Segments.DROPPED));
+    Assert.assertFalse(stats.hasStat(Stats.Segments.PARTIAL_STALE_DROPPED));
+    Assert.assertTrue(spare.getLoadingSegments().isEmpty());
+  }
+
+  @Test
+  public void testMoveOfStaleReplicaDoesNotDisturbMatchingReplica()
+  {
+    // A stale replica is being moved (the rule changed after the move was queued) while a matching replica serves
+    // the requirement elsewhere. The move is counted at its destination and by its own fingerprint, so it neither
+    // satisfies nor subtracts from the matching count: nothing to load, and nothing to drop until the move lands.
+    final DataSegment segment = createSegment();
+    final PartialLoadProfile usersProfile = PartialLoadProfile.forRequest(
+        Map.of("type", "partialProjection", "projections", List.of("users"), "fingerprint", FP_USERS),
+        FP_USERS
+    );
+    final ServerHolder source = moveSourceOf(segment, usersProfile);
+    final ServerHolder destination = moveDestinationOf(segment, usersProfile);
+    final ServerHolder matching = createServerWithLoaded(TIER1, segment, profileForRevenue());
+    final ServerHolder spare = createServer(TIER1);
+    final DruidCluster cluster =
+        DruidCluster.builder().addTier(TIER1, source, destination, matching, spare).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, segment);
+    params.getSegmentAssigner()
+          .replicateSegmentPartially(segment, profileForRevenue(), ImmutableMap.of(TIER1, 1));
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertFalse(stats.hasStat(Stats.Segments.PARTIAL_ASSIGNED));
+    Assert.assertFalse(stats.hasStat(Stats.Segments.DROPPED));
+    Assert.assertFalse(stats.hasStat(Stats.Segments.PARTIAL_STALE_DROPPED));
+    Assert.assertTrue(spare.getLoadingSegments().isEmpty());
+    Assert.assertTrue(matching.getPeon().getSegmentsToDrop().isEmpty());
+  }
+
   private DruidCoordinatorRuntimeParams makeRuntimeParams(DruidCluster cluster, DataSegment... segments)
   {
     return DruidCoordinatorRuntimeParams
@@ -515,6 +594,32 @@ private ServerHolder createServerWithLoaded(String tier, DataSegment segment, @N
     return new ServerHolder(server.toImmutableDruidServer(), new TestLoadQueuePeon());
   }
 
+  /**
+   * Creates the source server of a move queued by a previous coordinator run: it still serves the segment (under
+   * {@code profile}) and is marked to drop, which the ServerHolder reads back as a MOVE_FROM.
+   */
+  private ServerHolder moveSourceOf(DataSegment segment, @Nullable PartialLoadProfile profile)
+  {
+    final DruidServer server = createDruidServer(TIER1);
+    server.addDataSegment(segment, profile);
+    final MarkToDropPeon peon = new MarkToDropPeon();
+    peon.markSegmentToDrop(segment);
+    return new ServerHolder(server.toImmutableDruidServer(), peon);
+  }
+
+  /**
+   * Creates the destination server of a move queued by a previous coordinator run: a MOVE_TO sits in its queue,
+   * carrying the profile cloned from the move's source.
+   */
+  private ServerHolder moveDestinationOf(DataSegment segment, @Nullable PartialLoadProfile profile)
+  {
+    final TestLoadQueuePeon peon = new TestLoadQueuePeon();
+    peon.addInFlightHolder(
+        new SegmentHolder(segment, SegmentAction.MOVE_TO, profile, org.joda.time.Duration.standardSeconds(10), null)
+    );
+    return new ServerHolder(createDruidServer(TIER1).toImmutableDruidServer(), peon);
+  }
+
   private ServerHolder createDecommissioningServer(String tier)
   {
     return new ServerHolder(createDruidServer(tier).toImmutableDruidServer(), new TestLoadQueuePeon(), true);
@@ -558,6 +663,33 @@ private static DataSegment segmentWithProjections(List<String> projections)
         .build();
   }
 
+  /**
+   * Peon that records the segments marked to drop, so that a move source shows up as MOVE_FROM in the
+   * {@link ServerHolder}'s queue the way the real peon reports it.
+   */
+  private static class MarkToDropPeon extends TestLoadQueuePeon
+  {
+    private final Set<DataSegment> markedToDrop = new HashSet<>();
+
+    @Override
+    public void markSegmentToDrop(DataSegment segment)
+    {
+      markedToDrop.add(segment);
+    }
+
+    @Override
+    public void unmarkSegmentToDrop(DataSegment segment)
+    {
+      markedToDrop.remove(segment);
+    }
+
+    @Override
+    public Set<DataSegment> getSegmentsMarkedToDrop()
+    {
+      return markedToDrop;
+    }
+  }
+
   private static PartialLoadProfile profileForRevenue()
   {
     return PartialLoadProfile.forRequest(