From b7727ff605ab8d559b5b34143407083afba03ce8 Mon Sep 17 00:00:00 2001 From: Lucas Capistrant Date: Thu, 6 Aug 2026 11:55:58 -0500 Subject: [PATCH] fix: Prevent excessive segment assignment churn when using partial segment loading by fixing tiered replicant accounting for moving segments --- .../loading/PartialSegmentStatusInTier.java | 37 +++-- .../loading/StrategicSegmentAssigner.java | 15 +- .../StrategicSegmentAssignerPartialTest.java | 132 ++++++++++++++++++ 3 files changed, 164 insertions(+), 20 deletions(-) diff --git a/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialSegmentStatusInTier.java b/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialSegmentStatusInTier.java index 7cb30ebe9538..84dd10f3dd45 100644 --- a/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialSegmentStatusInTier.java +++ b/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialSegmentStatusInTier.java @@ -89,8 +89,8 @@ public List 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 getMatchingInFlight() { @@ -98,9 +98,10 @@ public List getMatchingInFlight() } /** - * 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 getStaleInFlight() { @@ -130,15 +131,23 @@ public List getEligibleForAdditiveReload() /** * Mechanical classification of one server against the request fingerprint. Branches are mutually exclusive in * order: loaded ({@link ServerHolder#isServingSegment}: matching / stale, with stale optionally also added - * to {@link #eligibleForAdditiveReload}), in-flight LOAD/REPLICATE (matching / stale based on the peon's - * queued profile), empty-and-loadable ({@link #eligibleForFreshLoad}). + * to {@link #eligibleForAdditiveReload}), in-flight LOAD/REPLICATE/MOVE_TO (matching / stale based on the + * peon's queued profile), empty-and-loadable ({@link #eligibleForFreshLoad}). *

- * 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. + *

+ * 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 - * no 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 no 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); 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 72920c78d9de..b899925f446a 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 @@ -345,10 +345,13 @@ public void replicateSegmentPartially( *

    *
  1. Classify. 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.
    2. - *
  2. Compute matching count: matching-loaded + matching-in-flight − pending-move-drop.
    3. + * 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. + *
  3. Compute matching count: 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.
    4. *
  4. If matching count is short of {@code requiredReplicas} (deficit): *
      *
    1. 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. 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 c5ef0280cf3c..5a1a87741e7e 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 @@ -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 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 markedToDrop = new HashSet<>(); + + @Override + public void markSegmentToDrop(DataSegment segment) + { + markedToDrop.add(segment); + } + + @Override + public void unmarkSegmentToDrop(DataSegment segment) + { + markedToDrop.remove(segment); + } + + @Override + public Set getSegmentsMarkedToDrop() + { + return markedToDrop; + } + } + private static PartialLoadProfile profileForRevenue() { return PartialLoadProfile.forRequest(