From 63e8679a5401e6455386dbe4f07dec6ca0181004 Mon Sep 17 00:00:00 2001
From: Clint Wylie <cwylie@apache.org>
Date: Thu, 6 Aug 2026 16:50:36 -0700
Subject: [PATCH] fix: replica counting fix for partial load reverts

---
 .../loading/SegmentReplicaCountMap.java       | 30 +++++++++++++++----
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCountMap.java b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCountMap.java
index df694086885c..7058c261d88e 100644
--- a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCountMap.java
+++ b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCountMap.java
@@ -23,6 +23,7 @@
 import org.apache.druid.client.DataSegmentAndLoadProfile;
 import org.apache.druid.client.ImmutableDruidServer;
 import org.apache.druid.server.coordinator.DruidCluster;
+import org.apache.druid.server.coordinator.ServerHolder;
 import org.apache.druid.timeline.DataSegment;
 import org.apache.druid.timeline.SegmentId;
 
@@ -55,13 +56,13 @@ private void initReplicaCounts(DruidCluster cluster)
             serverHolder -> {
               // Add segments already loaded on this server.
               final Collection<DataSegment> servedSegments = serverHolder.getServedSegments();
-              final Set<SegmentId> servedSegmentIds = Sets.newHashSetWithExpectedSize(servedSegments.size());
+              final Set<SegmentId> partiallyLoadedIds = Sets.newHashSetWithExpectedSize(servedSegments.size());
               for (DataSegment segment : servedSegments) {
-                servedSegmentIds.add(segment.getId());
                 final SegmentReplicaCount replicaCount = computeIfAbsent(segment.getId(), tier);
                 if (DataSegmentAndLoadProfile.profileOf(segment) == null) {
                   replicaCount.incrementLoaded();
                 } else {
+                  partiallyLoadedIds.add(segment.getId());
                   replicaCount.incrementLoadedWithPartialProfile();
                 }
               }
@@ -69,10 +70,7 @@ private void initReplicaCounts(DruidCluster cluster)
               // Add segments queued for load, drop or move on this server
               serverHolder.getQueuedSegments().forEach(
                   (segment, state) -> {
-                    // A load queued on a server that is already serving the segment is an in-place reload, not an
-                    // additional replica: the replica exists and was counted above, and the reload only changes which
-                    // parts of it the server holds.
-                    if (state == SegmentAction.LOAD && servedSegmentIds.contains(segment.getId())) {
+                    if (isPartialLoadRevert(serverHolder, segment, state, partiallyLoadedIds)) {
                       return;
                     }
                     computeIfAbsent(segment.getId(), tier).incrementQueued(state);
@@ -99,6 +97,26 @@ private void initReplicaCounts(DruidCluster cluster)
     });
   }
 
+  /**
+   * Whether a queued operation is the in-place reload that
+   * {@link StrategicSegmentAssigner#revertPartialProfileReplica} queues to release a partial-load rule that no longer
+   * applies. Such a reload refreshes an existing replica rather than adding one, so counting it as {@code loading}
+   * would push {@code projectedReplicas} past the requirement; {@code updateReplicasInTier} would then "correct" the
+   * phantom surplus by canceling the very reload it queued last run, and requeue it in the same run, churning
+   * forever against a backlogged peon.
+   */
+  private static boolean isPartialLoadRevert(
+      ServerHolder serverHolder,
+      DataSegment segment,
+      SegmentAction state,
+      Set<SegmentId> partiallyLoadedIds
+  )
+  {
+    return state == SegmentAction.LOAD
+           && partiallyLoadedIds.contains(segment.getId())
+           && serverHolder.getInFlightProfile(segment) == null;
+  }
+
   SegmentReplicaCount get(SegmentId segmentId, String tier)
   {
     SegmentReplicaCount count = replicaCounts.getOrDefault(segmentId, Collections.emptyMap())