This is an automated email from the ASF dual-hosted git repository.

clintropolis pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git


The following commit(s) were added to refs/heads/master by this push:
     new ccb09096887 fix: partial rule to regular rule transitions (#19884)
ccb09096887 is described below

commit ccb09096887ab7dbdc02d1db796eab9305570e37
Author: Clint Wylie <[email protected]>
AuthorDate: Thu Aug 6 14:02:54 2026 -0700

    fix: partial rule to regular rule transitions (#19884)
---
 .../query/PartialProjectionLoadRuleQueryTest.java  | 162 ++++++++--
 .../coordinator/loading/SegmentReplicaCount.java   |  23 ++
 .../loading/SegmentReplicaCountMap.java            |  29 +-
 .../loading/StrategicSegmentAssigner.java          |  79 ++++-
 .../druid/server/coordinator/stats/Stats.java      |   2 +
 .../StrategicSegmentAssignerPartialTest.java       | 338 ++++++++++++++++++++-
 6 files changed, 599 insertions(+), 34 deletions(-)

diff --git 
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
 
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
index 0215e63fa63..925fb78722a 100644
--- 
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
+++ 
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
@@ -37,11 +37,14 @@ import org.apache.druid.java.util.common.HumanReadableBytes;
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.granularity.Granularities;
+import org.apache.druid.query.DruidMetrics;
 import org.apache.druid.query.QueryContexts;
 import org.apache.druid.query.aggregation.LongMinAggregatorFactory;
 import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
 import org.apache.druid.server.coordinator.rules.CannotMatchBehavior;
+import org.apache.druid.server.coordinator.rules.ForeverLoadRule;
 import org.apache.druid.server.coordinator.rules.ForeverPartialLoadRule;
+import org.apache.druid.server.coordinator.rules.Rule;
 import 
org.apache.druid.server.coordinator.rules.WildcardProjectionPartialLoadMatcher;
 import org.apache.druid.server.metrics.LatchableEmitter;
 import org.apache.druid.server.metrics.StorageMonitor;
@@ -69,6 +72,7 @@ import java.nio.file.Files;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.function.LongPredicate;
 
 /**
  * End-to-end coverage for the partial-load-rule wiring: a segment with a 
clustered base + aggregate projection
@@ -97,6 +101,10 @@ class PartialProjectionLoadRuleQueryTest extends 
EmbeddedClusterTestBase
   // Quiescence wait for the storage monitor: a few times its PT1s emission 
period so a single missed tick
   // doesn't falsely read as "idle" while still bounding how long we wait once 
activity has actually stopped.
   private static final long MONITOR_QUIESCE_TIMEOUT_MILLIS = 3_000L;
+  // A rule change has to travel through a coordinator run, the load queue, 
the historical's reload and its
+  // re-announcement before the new footprint is visible in sys.servers, so 
this is generous relative to the
+  // coordinator's run cadence.
+  private static final long TRANSITION_TIMEOUT_MILLIS = 60_000L;
 
   private final EmbeddedBroker broker = new EmbeddedBroker();
   private final EmbeddedIndexer indexer = new EmbeddedIndexer();
@@ -202,19 +210,7 @@ class PartialProjectionLoadRuleQueryTest extends 
EmbeddedClusterTestBase
     // Configure the partial-load rule BEFORE ingestion so the initial load 
applies the rule directly, rather than
     // needing a rule change afterward. The matcher selects only the 
country_delta projection; base-table bundles
     // (cluster groups) are NOT rule-loaded.
-    cluster.callApi().onLeaderCoordinator(
-        c -> c.updateRulesForDatasource(
-            dataSource,
-            List.of(
-                new ForeverPartialLoadRule(
-                    Map.of("_default_tier", 1),
-                    null,
-                    new 
WildcardProjectionPartialLoadMatcher(List.of(PROJECTION_NAME), null),
-                    CannotMatchBehavior.FALL_THROUGH
-                )
-            )
-        )
-    );
+    applyRule(partialProjectionRule());
 
     ingestClusteredSegmentWithProjection();
   }
@@ -263,18 +259,10 @@ class PartialProjectionLoadRuleQueryTest extends 
EmbeddedClusterTestBase
     // sys.servers.curr_size for the historical reflects the partial footprint 
— strictly less than the segment's full
     // size (which includes the unmatched projection's bytes). Without the 
fix, forAnnouncement would stamp
     // segment.getSize() and curr_size would equal the full size.
-    final long fullSize = Long.parseLong(
-        cluster.callApi().runSql(
-            "SELECT \"size\" FROM sys.segments WHERE datasource = '" + 
dataSource + "'"
-        ).trim()
-    );
+    final long fullSize = queryFullSegmentSize();
     Assertions.assertTrue(fullSize > 0, "sys.segments.size must be populated 
for the ingested segment");
 
-    final long currSize = Long.parseLong(
-        cluster.callApi().runSql(
-            "SELECT curr_size FROM sys.servers WHERE server_type = 
'historical'"
-        ).trim()
-    );
+    final long currSize = queryHistoricalCurrSize();
 
     Assertions.assertTrue(
         currSize > 0,
@@ -319,6 +307,134 @@ class PartialProjectionLoadRuleQueryTest extends 
EmbeddedClusterTestBase
     );
   }
 
+  @Test
+  void testRuleTransitionBetweenPartialAndFullLoad()
+  {
+    // Rule churn in both directions against an already-loaded segment.
+    //
+    // Observed through sys.servers.curr_size, which is driven by the 
loadedBytes the historical announces: the
+    // partial footprint excludes the unmatched projection's bytes, while a 
full-load announcement carries no profile
+    // at all and falls back to segment.getSize().
+    final long fullSize = queryFullSegmentSize();
+    Assertions.assertTrue(fullSize > 0, "sys.segments.size must be populated 
for the ingested segment");
+
+    // Precondition from loadDataAndConfigureRule(): the segment is already 
loaded under the partial rule.
+    MatcherAssert.assertThat(
+        "partial-load rule should already be in effect at the start of this 
test",
+        queryHistoricalCurrSize(),
+        Matchers.lessThan(fullSize)
+    );
+
+    boolean restored = false;
+    try {
+      final LatchableEmitter coordinatorEmitter = 
coordinator.latchableEmitter();
+      coordinatorEmitter.flush();
+
+      // Leg 1: partial -> full.
+      applyRule(new ForeverLoadRule(Map.of("_default_tier", 1), null));
+      coordinatorEmitter.waitForEvent(
+          event -> event.hasMetricName("segment/partial/ruleReverted/count")
+                        .hasDimension(DruidMetrics.DATASOURCE, dataSource)
+      );
+      Assertions.assertEquals(
+          fullSize,
+          awaitHistoricalCurrSize(
+              size -> size == fullSize,
+              "full segment size after reverting to a regular load rule"
+          ),
+          "a reverted replica must announce with no partial-load profile, so 
loadedBytes falls back to segment size"
+      );
+
+      // Leg 2: full -> partial. Already covered by the initial load, but 
asserting it here proves the transition is
+      // symmetric on an already-loaded segment and restores the state the 
other tests in this class depend on.
+      applyRule(partialProjectionRule());
+      restored = true;
+      MatcherAssert.assertThat(
+          "re-applying the partial rule must shrink the announced footprint 
again",
+          awaitHistoricalCurrSize(size -> size < fullSize, "partial footprint 
after re-applying the partial rule"),
+          Matchers.lessThan(fullSize)
+      );
+    }
+    finally {
+      // Best effort only, and deliberately silent: on the success path leg 2 
already restored and verified the rule.
+      // This exists so a failure above can't leave the wrong rule behind for 
the rest of the class, and it must never
+      // mask the original failure.
+      if (!restored) {
+        try {
+          applyRule(partialProjectionRule());
+        }
+        catch (Exception ignored) {
+          // fall through; the failure that got us here is the one worth 
reporting
+        }
+      }
+    }
+  }
+
+  /**
+   * The partial-load rule this class runs under: selects only {@link 
#PROJECTION_NAME}, leaving the base-table
+   * cluster-group bundles and {@link #UNMATCHED_PROJECTION_NAME} off the 
historical's disk.
+   */
+  private static ForeverPartialLoadRule partialProjectionRule()
+  {
+    return new ForeverPartialLoadRule(
+        Map.of("_default_tier", 1),
+        null,
+        new WildcardProjectionPartialLoadMatcher(List.of(PROJECTION_NAME), 
null),
+        CannotMatchBehavior.FALL_THROUGH
+    );
+  }
+
+  private void applyRule(Rule rule)
+  {
+    cluster.callApi().onLeaderCoordinator(c -> 
c.updateRulesForDatasource(dataSource, List.of(rule)));
+  }
+
+  /**
+   * The ingested segment's full on-disk size, i.e. what the historical would 
report if it held everything.
+   */
+  private long queryFullSegmentSize()
+  {
+    return Long.parseLong(
+        cluster.callApi().runSql(
+            "SELECT \"size\" FROM sys.segments WHERE datasource = '" + 
dataSource + "'"
+        ).trim()
+    );
+  }
+
+  /**
+   * The footprint the coordinator currently attributes to the historical, 
i.e. the {@code loadedBytes} carried by the
+   * segment's most recent announcement.
+   */
+  private long queryHistoricalCurrSize()
+  {
+    return Long.parseLong(
+        cluster.callApi().runSql(
+            "SELECT curr_size FROM sys.servers WHERE server_type = 
'historical'"
+        ).trim()
+    );
+  }
+
+  /**
+   * Polls {@link #queryHistoricalCurrSize()} until it satisfies {@code 
matcher}. A rule change has to travel through a
+   * coordinator run, the load queue, the historical's reload and its 
re-announcement before the inventory catches up,
+   * and none of those emit a single metric that means "the new footprint is 
now visible in sys.servers".
+   */
+  private long awaitHistoricalCurrSize(LongPredicate matcher, String 
expectation)
+  {
+    try {
+      return cluster.callApi()
+                    .waitForResult(this::queryHistoricalCurrSize, size -> 
matcher.test(size))
+                    .withTimeoutMillis(TRANSITION_TIMEOUT_MILLIS)
+                    .go();
+    }
+    catch (Exception e) {
+      throw new AssertionError(
+          StringUtils.format("Timed out waiting for the historical's curr_size 
to reach [%s]", expectation),
+          e
+      );
+    }
+  }
+
   /**
    * Waits for any lingering storage-monitor activity from ingest / 
rule-application to settle, then flushes the
    * emitter's event queue so subsequent {@code getMetricEventLongSum(...)} 
calls only reflect activity that
diff --git 
a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCount.java
 
b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCount.java
index 8f0bf11b7d4..d2261080e5c 100644
--- 
a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCount.java
+++ 
b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentReplicaCount.java
@@ -30,6 +30,7 @@ public class SegmentReplicaCount
 
   private int loaded;
   private int loadedNonHistorical;
+  private int loadedWithPartialProfile;
 
   private int loading;
   private int dropping;
@@ -44,6 +45,16 @@ public class SegmentReplicaCount
     ++loaded;
   }
 
+  /**
+   * Increments number of replicas loaded on historical servers, additionally 
counting this replica as one that
+   * announced a {@link PartialLoadProfile}, i.e. it is pinned on its 
historical by a partial-load rule.
+   */
+  void incrementLoadedWithPartialProfile()
+  {
+    ++loaded;
+    ++loadedWithPartialProfile;
+  }
+
   /**
    * Increments number of replicas loaded on non-historical servers. This value
    * is used only for computing level of under-replication of broadcast 
segments.
@@ -145,6 +156,17 @@ public class SegmentReplicaCount
     return loaded - dropping;
   }
 
+  /**
+   * Number of loaded replicas that announced a {@link PartialLoadProfile}. 
Under a partial-load rule this is only
+   * bookkeeping, that rule's own reconciler classifies replicas by 
fingerprint. Under a regular load rule a non-zero
+   * value means those historicals are still pinned by a partial-load rule 
that no longer applies; see
+   * {@link StrategicSegmentAssigner#updateReplicasInTier}.
+   */
+  int loadedWithPartialProfile()
+  {
+    return loadedWithPartialProfile;
+  }
+
   /**
    * Number of replicas that are required to be loaded but are missing.
    * This includes replicas that may be in excess of the cluster capacity.
@@ -183,6 +205,7 @@ public class SegmentReplicaCount
 
     this.loaded += other.loaded;
     this.loadedNonHistorical += other.loadedNonHistorical;
+    this.loadedWithPartialProfile += other.loadedWithPartialProfile;
 
     this.loading += other.loading;
     this.dropping += other.dropping;
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 241759f8b95..df694086885 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
@@ -19,14 +19,18 @@
 
 package org.apache.druid.server.coordinator.loading;
 
+import com.google.common.collect.Sets;
+import org.apache.druid.client.DataSegmentAndLoadProfile;
 import org.apache.druid.client.ImmutableDruidServer;
 import org.apache.druid.server.coordinator.DruidCluster;
 import org.apache.druid.timeline.DataSegment;
 import org.apache.druid.timeline.SegmentId;
 
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * Contains a mapping from tier to {@link SegmentReplicaCount}s.
@@ -49,15 +53,30 @@ public class SegmentReplicaCountMap
     cluster.getManagedHistoricals().forEach(
         (tier, historicals) -> historicals.forEach(
             serverHolder -> {
-              // Add segments already loaded on this server
-              for (DataSegment segment : serverHolder.getServedSegments()) {
-                computeIfAbsent(segment.getId(), tier).incrementLoaded();
+              // Add segments already loaded on this server.
+              final Collection<DataSegment> servedSegments = 
serverHolder.getServedSegments();
+              final Set<SegmentId> servedSegmentIds = 
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 {
+                  replicaCount.incrementLoadedWithPartialProfile();
+                }
               }
 
               // Add segments queued for load, drop or move on this server
               serverHolder.getQueuedSegments().forEach(
-                  (segment, state) -> computeIfAbsent(segment.getId(), tier)
-                      .incrementQueued(state)
+                  (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())) {
+                      return;
+                    }
+                    computeIfAbsent(segment.getId(), 
tier).incrementQueued(state);
+                  }
               );
             }
         )
diff --git 
a/server/src/main/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssigner.java
 
b/server/src/main/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssigner.java
index 72920c78d9d..f8bae67cbe4 100644
--- 
a/server/src/main/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssigner.java
+++ 
b/server/src/main/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssigner.java
@@ -615,8 +615,16 @@ public class StrategicSegmentAssigner implements 
SegmentActionHandler
     final int movingReplicas = replicaCountOnTier.moving();
     final boolean shouldCancelMoves = requiredReplicas == 0 && movingReplicas 
> 0;
 
+    // A replica serving under a partial-load profile is pinned by a 
partial-load rule, but we got here through the
+    // regular full-load path, so that rule no longer applies: the 
datasource's partial-load rule was replaced or
+    // shadowed by a higher-priority load rule, or its matcher stopped 
resolving and fell through to FULL_LOAD. That
+    // replica needs an in-place reload carrying the plain unwrapped load spec 
so the historical releases its rule
+    // holds. When the tier wants no replicas at all we skip it, the drops 
below are already on their way and dropping
+    // clears the rule on the historical.
+    final int replicasToRevert = requiredReplicas > 0 ? 
replicaCountOnTier.loadedWithPartialProfile() : 0;
+
     // Check if there is any action required on this tier
-    if (projectedReplicas == requiredReplicas && !shouldCancelMoves) {
+    if (projectedReplicas == requiredReplicas && !shouldCancelMoves && 
replicasToRevert <= 0) {
       return 0;
     }
 
@@ -645,6 +653,7 @@ public class StrategicSegmentAssigner implements 
SegmentActionHandler
     }
 
     // Cancel loads and queue drops if the projected count exceeds the 
requirement
+    int dropsQueuedOnTier = 0;
     if (projectedReplicas > requiredReplicas) {
       int replicaSurplus = projectedReplicas - requiredReplicas;
       int canceledLoads =
@@ -652,13 +661,66 @@ public class StrategicSegmentAssigner implements 
SegmentActionHandler
 
       int numReplicasToDrop = Math.min(replicaSurplus - canceledLoads, 
maxReplicasToDrop);
       if (numReplicasToDrop > 0) {
-        int dropsQueuedOnTier = dropReplicas(numReplicasToDrop, segment, tier, 
segmentStatus);
+        dropsQueuedOnTier = dropReplicas(numReplicasToDrop, segment, tier, 
segmentStatus);
         incrementStat(Stats.Segments.DROPPED, segment, tier, 
dropsQueuedOnTier);
-        return dropsQueuedOnTier;
       }
     }
 
-    return 0;
+    // Release partial-load rules that no longer apply. Done last so the 
load/drop decisions above claim their
+    // servers first: a replica that just picked up an action is no longer 
`isServingSegment`, so it is skipped here
+    // and reverted on a later run if it is still around.
+    if (replicasToRevert > 0) {
+      final int reverted = revertPartialProfileReplicas(segment, tier);
+      if (reverted > 0) {
+        incrementStat(Stats.Segments.PARTIAL_RULE_REVERTED, segment, tier, 
reverted);
+      }
+    }
+
+    return dropsQueuedOnTier;
+  }
+
+  /**
+   * Queues an in-place reload on every server in {@code tier} that serves 
{@code segment} under a
+   * {@link PartialLoadProfile}. The request carries the plain unwrapped 
{@code segment}, which is what tells the
+   * historical to release its rule holds rather than apply or swap one.
+   * <p>
+   * The replica count is deliberately left alone: these servers <em>are</em> 
serving, so they still satisfy the
+   * rule's replication requirement and must not be double-counted as a 
deficit. This only refreshes what they hold.
+   * <p>
+   * Servers are skipped when:
+   * <ul>
+   *   <li>they have any queued action, via {@link 
ServerHolder#isServingSegment}, which covers both the load/drop
+   *       decisions made earlier in this run and operations left over from a 
previous one;</li>
+   *   <li>their load queue is already at the configured {@code 
maxSegmentsInNodeLoadingQueue} budget for this run.</li>
+   *   <li>they are decommissioning, since their replicas are on the way out 
and reloading them is wasted work.</li>
+   * </ul>
+   * These are the same two eligibility conditions {@code 
PartialSegmentStatusInTier.canReloadAdditively} applies to
+   * the partial-load reconciler's in-place reload. {@link 
ServerHolder#canLoadSegment} is not usable here because it
+   * requires the server to <em>not</em> already have the segment, which is 
precisely the case being handled.
+   */
+  private int revertPartialProfileReplicas(DataSegment segment, String tier)
+  {
+    int numReverted = 0;
+    for (ServerHolder server : cluster.getManagedHistoricalsByTier(tier)) {
+      if (revertPartialProfileReplica(segment, server)) {
+        ++numReverted;
+      }
+    }
+    return numReverted;
+  }
+
+  /**
+   * Queues the in-place reload described by {@link 
#revertPartialProfileReplicas} on a single server, if that server
+   * is holding {@code segment} under a partial-load rule and has room in its 
load queue. Returns whether a reload was
+   * queued.
+   */
+  private boolean revertPartialProfileReplica(DataSegment segment, 
ServerHolder server)
+  {
+    return server.isServingSegment(segment)
+           && !server.isDecommissioning()
+           && !server.isLoadQueueFull()
+           && server.getServer().getPartialLoadProfile(segment.getId()) != null
+           && loadQueueManager.loadSegment(segment, server, 
SegmentAction.LOAD, null);
   }
 
   private void reportTierCapacityStats(DataSegment segment, int 
requiredReplicas, String tier)
@@ -697,11 +759,17 @@ public class StrategicSegmentAssigner implements 
SegmentActionHandler
       // Drop from decommissioning servers and load on active servers
       int numDropsQueued = 0;
       int numLoadsQueued = 0;
+      int numRevertsQueued = 0;
       if (server.isDecommissioning()) {
         numDropsQueued += dropBroadcastSegment(segment, server) ? 1 : 0;
       } else {
         tierToRequiredReplicas.addTo(tier, 1);
         numLoadsQueued += loadBroadcastSegment(segment, server) ? 1 : 0;
+        // A broadcast rule wants the whole segment on every target, so a 
replica still pinned by a partial-load rule
+        // has to be reloaded unwrapped, exactly as on the replication path. 
loadBroadcastSegment cannot do this
+        // itself: it returns early for a server that is already serving, 
which is precisely the case here. The two
+        // are mutually exclusive, since a server that just had a genuine load 
queued is no longer `isServingSegment`.
+        numRevertsQueued += revertPartialProfileReplica(segment, server) ? 1 : 
0;
       }
 
       if (numLoadsQueued > 0) {
@@ -710,6 +778,9 @@ public class StrategicSegmentAssigner implements 
SegmentActionHandler
       if (numDropsQueued > 0) {
         incrementStat(Stats.Segments.DROPPED, segment, tier, numDropsQueued);
       }
+      if (numRevertsQueued > 0) {
+        incrementStat(Stats.Segments.PARTIAL_RULE_REVERTED, segment, tier, 
numRevertsQueued);
+      }
     }
 
     // Update required replica counts
diff --git 
a/server/src/main/java/org/apache/druid/server/coordinator/stats/Stats.java 
b/server/src/main/java/org/apache/druid/server/coordinator/stats/Stats.java
index fe2031899a0..80368ba7558 100644
--- a/server/src/main/java/org/apache/druid/server/coordinator/stats/Stats.java
+++ b/server/src/main/java/org/apache/druid/server/coordinator/stats/Stats.java
@@ -79,6 +79,8 @@ public class Stats
         = CoordinatorStat.toDebugAndEmit("partialStaleDropped", 
"segment/partial/staleDropped/count");
     public static final CoordinatorStat PARTIAL_STALE_CANCELLED
         = CoordinatorStat.toDebugAndEmit("partialStaleCancelled", 
"segment/partial/staleCancelled/count");
+    public static final CoordinatorStat PARTIAL_RULE_REVERTED
+        = CoordinatorStat.toDebugAndEmit("partialRuleReverted", 
"segment/partial/ruleReverted/count");
   }
 
   public static class SegmentQueue
diff --git 
a/server/src/test/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssignerPartialTest.java
 
b/server/src/test/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssignerPartialTest.java
index c5ef0280cf3..c6290b1f776 100644
--- 
a/server/src/test/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssignerPartialTest.java
+++ 
b/server/src/test/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssignerPartialTest.java
@@ -43,6 +43,8 @@ import org.apache.druid.server.coordinator.stats.Stats;
 import org.apache.druid.timeline.DataSegment;
 import org.apache.druid.timeline.SegmentId;
 import org.apache.druid.timeline.partition.NumberedShardSpec;
+import org.joda.time.Duration;
+import org.joda.time.Interval;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -59,6 +61,10 @@ import java.util.concurrent.atomic.AtomicInteger;
  * through the load queue and reconciles fingerprint state correctly: matching 
replicas count toward the requirement,
  * stale-fingerprint replicas (and full-load replicas under a partial rule) 
are treated as "loaded but not satisfying"
  * and follow the load-then-drop swap so the cluster never goes unavailable 
during reconciliation.
+ * <p>
+ * Also covers the opposite direction, where a segment stops being governed by 
a partial-load rule and reconciliation
+ * arrives through {@link StrategicSegmentAssigner#replicateSegment}: replicas 
still announcing a profile are pinned
+ * by a rule that no longer applies and have to be reloaded in place with the 
plain unwrapped load spec.
  */
 public class StrategicSegmentAssignerPartialTest
 {
@@ -151,6 +157,296 @@ public class StrategicSegmentAssignerPartialTest
     Assert.assertTrue(s1.getPeon().getSegmentsToDrop().isEmpty());
   }
 
+  @Test
+  public void testFullLoadRuleRevertsPartialProfileReplicaInPlace()
+  {
+    // The partial-load rule was replaced (or shadowed) by a regular load 
rule, so reconciliation arrives through
+    // replicateSegment instead of replicateSegmentPartially. s1 is still 
pinned by the old rule, and it is the only
+    // replica, so replica counts alone say "satisfied" and nothing would ever 
be queued. The assigner must notice the
+    // announced profile and queue an in-place reload carrying the plain 
unwrapped segment.
+    final DataSegment segment = createSegment();
+    final ServerHolder s1 = createServerWithLoaded(TIER1, segment, 
profileForRevenue());
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, 
s1).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().replicateSegment(segment, 
ImmutableMap.of(TIER1, 1));
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertEquals(
+        1L,
+        stats.getSegmentStat(Stats.Segments.PARTIAL_RULE_REVERTED, TIER1, 
segment.getDataSource())
+    );
+    Assert.assertTrue("in-place reload must be queued on the pinned server", 
s1.getLoadingSegments().contains(segment));
+    Assert.assertNull(
+        "revert must carry no profile so the historical receives the plain 
unwrapped load spec",
+        ((TestLoadQueuePeon) s1.getPeon()).getProfileFor(segment)
+    );
+    Assert.assertTrue("the replica is still serving and must not be dropped", 
s1.getPeon().getSegmentsToDrop().isEmpty());
+  }
+
+  @Test
+  public void testFullLoadRuleLeavesProfilelessReplicaAlone()
+  {
+    // Ordinary full-load replica under a full-load rule: nothing to 
reconcile, and the tier must still fast-exit.
+    final DataSegment segment = createSegment();
+    final ServerHolder s1 = createServerWithLoaded(TIER1, segment, null);
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, 
s1).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().replicateSegment(segment, 
ImmutableMap.of(TIER1, 1));
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertFalse(stats.hasStat(Stats.Segments.PARTIAL_RULE_REVERTED));
+    Assert.assertFalse(stats.hasStat(Stats.Segments.ASSIGNED));
+    Assert.assertTrue(s1.getLoadingSegments().isEmpty());
+    Assert.assertTrue(s1.getPeon().getSegmentsToDrop().isEmpty());
+  }
+
+  @Test
+  public void testFullLoadRuleDoesNotRevertWhenTierWantsNoReplicas()
+  {
+    // The tier is being emptied, so the drop is already on its way and 
dropping clears the rule on the historical.
+    // Queueing a reload here would be wasted work on a replica that is about 
to disappear.
+    final DataSegment segment = createSegment();
+    final ServerHolder s1 = createServerWithLoaded(TIER1, segment, 
profileForRevenue());
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, 
s1).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().replicateSegment(segment, 
ImmutableMap.of(TIER1, 0));
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertFalse(stats.hasStat(Stats.Segments.PARTIAL_RULE_REVERTED));
+    Assert.assertTrue(s1.getLoadingSegments().isEmpty());
+    Assert.assertTrue(s1.getPeon().getSegmentsToDrop().contains(segment));
+  }
+
+  @Test
+  public void testFullLoadRuleRevertsSurvivorWhileDroppingSurplus()
+  {
+    // Two pinned replicas but the new full-load rule only wants one. The 
surplus drop is fingerprint-blind so it may
+    // pick either server; whichever survives must still get reverted, and the 
one being dropped must not also be
+    // asked to reload.
+    final DataSegment segment = createSegment();
+    final ServerHolder s1 = createServerWithLoaded(TIER1, segment, 
profileForRevenue());
+    final ServerHolder s2 = createServerWithLoaded(TIER1, segment, 
profileForRevenue());
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, s1, 
s2).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().replicateSegment(segment, 
ImmutableMap.of(TIER1, 1));
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertEquals(1L, stats.getSegmentStat(Stats.Segments.DROPPED, 
TIER1, segment.getDataSource()));
+    Assert.assertEquals(
+        1L,
+        stats.getSegmentStat(Stats.Segments.PARTIAL_RULE_REVERTED, TIER1, 
segment.getDataSource())
+    );
+
+    final ServerHolder dropped = 
s1.getPeon().getSegmentsToDrop().contains(segment) ? s1 : s2;
+    final ServerHolder survivor = dropped == s1 ? s2 : s1;
+    Assert.assertTrue(dropped.getPeon().getSegmentsToDrop().contains(segment));
+    Assert.assertTrue("a server queued for drop must not also be reloaded", 
dropped.getLoadingSegments().isEmpty());
+    Assert.assertTrue(survivor.getLoadingSegments().contains(segment));
+  }
+
+  @Test
+  public void testFullLoadRuleRevertsPinnedReplicaWhileLoadingMissingOne()
+  {
+    // The new full-load rule wants two replicas and only the pinned one 
exists. Both things have to happen: a genuine
+    // new replica on the empty server, and a revert of the pinned one. The 
revert must not be mistaken for the
+    // deficit load, and it must not consume the deficit.
+    final DataSegment segment = createSegment();
+    final ServerHolder s1 = createServerWithLoaded(TIER1, segment, 
profileForRevenue());
+    final ServerHolder s2 = createServer(TIER1);
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, s1, 
s2).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().replicateSegment(segment, 
ImmutableMap.of(TIER1, 2));
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertEquals(1L, stats.getSegmentStat(Stats.Segments.ASSIGNED, 
TIER1, segment.getDataSource()));
+    Assert.assertEquals(
+        1L,
+        stats.getSegmentStat(Stats.Segments.PARTIAL_RULE_REVERTED, TIER1, 
segment.getDataSource())
+    );
+    Assert.assertTrue("deficit must be filled on the empty server", 
s2.getLoadingSegments().contains(segment));
+    Assert.assertTrue("pinned replica must still be reverted", 
s1.getLoadingSegments().contains(segment));
+    Assert.assertNull(((TestLoadQueuePeon) 
s1.getPeon()).getProfileFor(segment));
+    Assert.assertNull(((TestLoadQueuePeon) 
s2.getPeon()).getProfileFor(segment));
+  }
+
+  @Test
+  public void testRevertsRespectTheLoadQueueBudgetAcrossSegments()
+  {
+    // Switching a whole datasource back to a full-load rule means every one 
of its segments is a revert candidate on
+    // every historical holding it. ServerHolder.startOperation does not 
enforce maxSegmentsInNodeLoadingQueue on its
+    // own, so without an explicit budget check the reverts for an entire 
datasource would all be enqueued in a single
+    // coordinator run and flood the peon.
+    final DataSegment segment1 = 
createSegment(Intervals.of("2024-01-01/2024-01-02"));
+    final DataSegment segment2 = 
createSegment(Intervals.of("2024-01-02/2024-01-03"));
+    final ServerHolder s1 = createServerWithLoadedAndQueueLimit(TIER1, 1, 
profileForRevenue(), segment1, segment2);
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, 
s1).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment1, segment2);
+    params.getSegmentAssigner().replicateSegment(segment1, 
ImmutableMap.of(TIER1, 1));
+    params.getSegmentAssigner().replicateSegment(segment2, 
ImmutableMap.of(TIER1, 1));
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertEquals(
+        "only one revert fits in this run's queue budget; the rest wait for a 
later run",
+        1L,
+        stats.getSegmentStat(Stats.Segments.PARTIAL_RULE_REVERTED, TIER1, 
segment1.getDataSource())
+    );
+    Assert.assertEquals(1, s1.getLoadingSegments().size());
+  }
+
+  @Test
+  public void testPendingRevertSurvivesTheNextCoordinatorRun()
+  {
+    // Second run with the revert from the first still sitting in the queue. 
The server both serves the segment and
+    // has a LOAD queued for it, so counting that LOAD as a replica would put 
projectedReplicas at 2 against a
+    // requirement of 1. updateReplicasInTier would then read the phantom 
surplus and cancel the very reload it queued
+    // last run — and since the revert would be requeued at the end of the 
same run, a backlogged peon could be made
+    // to cancel and requeue forever, never completing the transition.
+    final DataSegment segment = createSegment();
+    final TestLoadQueuePeon peon = new TestLoadQueuePeon();
+    peon.addInFlightHolder(new SegmentHolder(
+        segment,
+        SegmentAction.LOAD,
+        // The revert carries no profile; that is what tells the historical to 
release its rule holds.
+        null,
+        Duration.standardSeconds(10),
+        null
+    ));
+    final DruidServer server = createDruidServer(TIER1);
+    // The historical has not finished the reload yet, so it is still 
announcing the old partial profile.
+    server.addDataSegment(segment, profileForRevenue());
+    final ServerHolder s1 = new ServerHolder(server.toImmutableDruidServer(), 
peon);
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, 
s1).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().replicateSegment(segment, 
ImmutableMap.of(TIER1, 1));
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertTrue(
+        "the in-flight revert must still be queued, not cancelled as surplus",
+        s1.getLoadingSegments().contains(segment)
+    );
+    Assert.assertFalse("a pending revert must not be mistaken for a surplus 
replica", stats.hasStat(Stats.Segments.DROPPED));
+    Assert.assertTrue(s1.getPeon().getSegmentsToDrop().isEmpty());
+    Assert.assertFalse(
+        "the revert is already in flight, so this run has nothing more to 
queue",
+        stats.hasStat(Stats.Segments.PARTIAL_RULE_REVERTED)
+    );
+  }
+
+  @Test
+  public void testRevertSkippedWhenLoadQueueIsAlreadyFull()
+  {
+    // The budget is the whole load queue, not just what this run has added to 
it: a queue already at the limit from a
+    // previous run leaves no room, so nothing is enqueued at all.
+    final DataSegment segment = createSegment();
+    final DataSegment queuedFromPriorRun = 
createSegment(Intervals.of("2023/2024"));
+    final TestLoadQueuePeon peon = new TestLoadQueuePeon();
+    peon.addInFlightHolder(new SegmentHolder(
+        queuedFromPriorRun,
+        SegmentAction.LOAD,
+        null,
+        Duration.standardSeconds(10),
+        null
+    ));
+    final DruidServer server = createDruidServer(TIER1);
+    server.addDataSegment(segment, profileForRevenue());
+    final ServerHolder s1 = new ServerHolder(server.toImmutableDruidServer(), 
peon, false, 1, 1);
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, 
s1).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().replicateSegment(segment, 
ImmutableMap.of(TIER1, 1));
+
+    
Assert.assertFalse(params.getCoordinatorStats().hasStat(Stats.Segments.PARTIAL_RULE_REVERTED));
+    Assert.assertFalse(s1.getLoadingSegments().contains(segment));
+  }
+
+  @Test
+  public void testBroadcastRuleRevertsPartialProfileReplicaInPlace()
+  {
+    // Replacing a partial-load rule with a broadcast rule reconciles through 
broadcastSegment, not replicateSegment.
+    // loadBroadcastSegment returns early for a server that is already 
serving, so without an explicit revert the old
+    // profile and its rule holds would survive indefinitely.
+    final DataSegment segment = createSegment();
+    final ServerHolder s1 = createServerWithLoaded(TIER1, segment, 
profileForRevenue());
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, 
s1).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().broadcastSegment(segment);
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertEquals(
+        1L,
+        stats.getSegmentStat(Stats.Segments.PARTIAL_RULE_REVERTED, TIER1, 
segment.getDataSource())
+    );
+    Assert.assertTrue(s1.getLoadingSegments().contains(segment));
+    Assert.assertNull(
+        "the broadcast revert must carry no profile, same as on the 
replication path",
+        ((TestLoadQueuePeon) s1.getPeon()).getProfileFor(segment)
+    );
+    Assert.assertFalse(
+        "an in-place revert is not a new assignment",
+        stats.hasStat(Stats.Segments.ASSIGNED)
+    );
+  }
+
+  @Test
+  public void testBroadcastRuleLeavesProfilelessReplicaAlone()
+  {
+    // The ordinary broadcast no-op must survive: a replica already serving 
without a profile needs nothing.
+    final DataSegment segment = createSegment();
+    final ServerHolder s1 = createServerWithLoaded(TIER1, segment, null);
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, 
s1).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().broadcastSegment(segment);
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertFalse(stats.hasStat(Stats.Segments.PARTIAL_RULE_REVERTED));
+    Assert.assertFalse(stats.hasStat(Stats.Segments.ASSIGNED));
+    Assert.assertTrue(s1.getLoadingSegments().isEmpty());
+  }
+
+  @Test
+  public void testBroadcastRuleLoadsMissingReplicaWithoutRevert()
+  {
+    // A broadcast target that does not have the segment at all takes an 
ordinary assignment, and must not be counted
+    // as a revert: the two branches are mutually exclusive because a freshly 
queued load stops the server from being
+    // `isServingSegment`.
+    final DataSegment segment = createSegment();
+    final ServerHolder s1 = createServer(TIER1);
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, 
s1).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().broadcastSegment(segment);
+
+    final CoordinatorRunStats stats = params.getCoordinatorStats();
+    Assert.assertEquals(1L, stats.getSegmentStat(Stats.Segments.ASSIGNED, 
TIER1, segment.getDataSource()));
+    Assert.assertFalse(stats.hasStat(Stats.Segments.PARTIAL_RULE_REVERTED));
+    Assert.assertTrue(s1.getLoadingSegments().contains(segment));
+  }
+
+  @Test
+  public void testRevertSkippedOnDecommissioningServer()
+  {
+    // A decommissioning server's replicas are on their way off it, so 
reloading them just to release rule holds is
+    // wasted work on a server that is trying to shed load.
+    final DataSegment segment = createSegment();
+    final ServerHolder s1 = createDecommissioningServerWithLoaded(TIER1, 
segment, profileForRevenue());
+    final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, 
s1).build();
+
+    final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, 
segment);
+    params.getSegmentAssigner().replicateSegment(segment, 
ImmutableMap.of(TIER1, 1));
+
+    
Assert.assertFalse(params.getCoordinatorStats().hasStat(Stats.Segments.PARTIAL_RULE_REVERTED));
+    Assert.assertTrue(s1.getLoadingSegments().isEmpty());
+  }
+
   @Test
   public void testFullFallbackLoadedBytesCountsAsMatching()
   {
@@ -320,7 +616,7 @@ public class StrategicSegmentAssignerPartialTest
         segment,
         SegmentAction.LOAD,
         staleInFlightProfile,
-        org.joda.time.Duration.standardSeconds(10),
+        Duration.standardSeconds(10),
         null
     ));
     final DruidServer druidServer = createDruidServer(TIER1);
@@ -436,7 +732,7 @@ public class StrategicSegmentAssignerPartialTest
         segment,
         SegmentAction.LOAD,
         profileForRevenue(),
-        org.joda.time.Duration.standardSeconds(10),
+        Duration.standardSeconds(10),
         null
     ));
     final ServerHolder source = new 
ServerHolder(createDruidServer(TIER1).toImmutableDruidServer(), sourcePeon, 
true);
@@ -515,6 +811,30 @@ public class StrategicSegmentAssignerPartialTest
     return new ServerHolder(server.toImmutableDruidServer(), new 
TestLoadQueuePeon());
   }
 
+  /**
+   * Creates a server that already serves every given segment under {@code 
profile}, with its load queue capped at
+   * {@code maxSegmentsInLoadQueue} so tests can observe the budget being 
exhausted.
+   */
+  private ServerHolder createServerWithLoadedAndQueueLimit(
+      String tier,
+      int maxSegmentsInLoadQueue,
+      @Nullable PartialLoadProfile profile,
+      DataSegment... segments
+  )
+  {
+    final DruidServer server = createDruidServer(tier);
+    for (DataSegment segment : segments) {
+      server.addDataSegment(segment, profile);
+    }
+    return new ServerHolder(
+        server.toImmutableDruidServer(),
+        new TestLoadQueuePeon(),
+        false,
+        maxSegmentsInLoadQueue,
+        1
+    );
+  }
+
   private ServerHolder createDecommissioningServer(String tier)
   {
     return new ServerHolder(createDruidServer(tier).toImmutableDruidServer(), 
new TestLoadQueuePeon(), true);
@@ -547,6 +867,20 @@ public class StrategicSegmentAssignerPartialTest
         .build();
   }
 
+  /**
+   * A segment over a specific interval, for tests that need several distinct 
segments of the same datasource.
+   */
+  private static DataSegment createSegment(Interval interval)
+  {
+    return DataSegment
+        .builder(SegmentId.of(TestDataSource.WIKI, interval, "v1", null))
+        .loadSpec(Map.of("type", "local", "path", "/var/druid/segments/foo"))
+        .dimensions(Collections.emptyList())
+        .metrics(Collections.emptyList())
+        .size(0)
+        .build();
+  }
+
   private static DataSegment segmentWithProjections(List<String> projections)
   {
     return DataSegment


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to