This is an automated email from the ASF dual-hosted git repository.
capistrant 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 640609453e5 feat: For partial loads, hook into existing strategic
assigner code to better distribute partial load requests (#20155)
640609453e5 is described below
commit 640609453e58ee64ae2af133584d14e8374c43f5
Author: Lucas Capistrant <[email protected]>
AuthorDate: Wed Aug 26 16:51:04 2026 -0500
feat: For partial loads, hook into existing strategic assigner code to
better distribute partial load requests (#20155)
* Use existing strategic assignment code for partial loads to improve
segment distribution
* clarify doc on why using the supplied iterator works for fresh loads
---
.../loading/StrategicSegmentAssigner.java | 71 +++++--
.../duty/RunRulesPartialLoadPlacementTest.java | 207 ++++++++++++++++++++-
.../StrategicSegmentAssignerPartialTest.java | 45 +++++
3 files changed, 296 insertions(+), 27 deletions(-)
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 8982dbb400d..ab95ed09cae 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
@@ -19,6 +19,7 @@
package org.apache.druid.server.coordinator.loading;
+import com.google.common.collect.Iterators;
import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
import org.apache.druid.client.DruidServer;
import org.apache.druid.server.coordinator.DruidCluster;
@@ -488,10 +489,16 @@ public class StrategicSegmentAssigner implements
SegmentActionHandler
}
/**
- * Queues fresh partial-load requests on up to {@code numToLoad} eligible
servers. Preference order: empty
- * (fresh-load) servers first; then servers whose stale-fingerprint
in-flight loads were just canceled (their slot
- * is now free); then stale-loaded servers (additive reload; the historical
fills missing parts in place). The last
- * fallback is what mitigates the "tier saturated with stale" stuck state.
+ * Queues fresh partial-load requests on up to {@code numToLoad} eligible
servers. Preference order: servers that can
+ * take a fresh load, then stale-loaded servers (additive reload; the
historical fills missing parts in place).
+ * <p>
+ * The fresh-load candidates are the classifier's empty servers together
with {@code canceledStaleServers}.
+ * {@link ServerHolder#cancelOperation} clears the queued action and
restores the projected size, so a server whose
+ * stale in-flight load {@link #cancelLoadsOnServers} canceled can take a
fresh load. {@link #serversToLoadSegment}
+ * returns an iterator over fresh load candidates.
+ * <p>
+ * An iterator over candidates that can additive reload the segment is there
for backup in case we can't fully
+ * replicate on our priority one fresh load path.
*/
private int loadPartialReplicas(
int numToLoad,
@@ -509,25 +516,28 @@ public class StrategicSegmentAssigner implements
SegmentActionHandler
return 0;
}
- final List<ServerHolder> destinations = new ArrayList<>(
- status.getEligibleForFreshLoad().size()
- + canceledStaleServers.size()
- + status.getEligibleForAdditiveReload().size()
+ // The classifier's list is already the complete candidate set when
nothing was canceled.
+ final List<ServerHolder> freshCandidates;
+ if (canceledStaleServers.isEmpty()) {
+ freshCandidates = status.getEligibleForFreshLoad();
+ } else {
+ freshCandidates = new ArrayList<>(status.getEligibleForFreshLoad());
+ freshCandidates.addAll(canceledStaleServers);
+ }
+
+ final Iterator<ServerHolder> destinations = Iterators.concat(
+ serversToLoadSegment(segment, tier, freshCandidates),
+ status.getEligibleForAdditiveReload().iterator()
);
- destinations.addAll(status.getEligibleForFreshLoad());
- destinations.addAll(canceledStaleServers);
- destinations.addAll(status.getEligibleForAdditiveReload());
- if (destinations.isEmpty()) {
+ if (!destinations.hasNext()) {
incrementSkipStat(Stats.Segments.ASSIGN_SKIPPED, "No eligible server",
segment, tier);
return 0;
}
int numLoadsQueued = 0;
- for (ServerHolder server : destinations) {
- if (numLoadsQueued >= numToLoad) {
- break;
- }
+ while (numLoadsQueued < numToLoad && destinations.hasNext()) {
+ final ServerHolder server = destinations.next();
final boolean queuedSuccessfully = isAlreadyLoadedOnTier
? replicateSegment(segment, server,
profile)
: loadSegment(segment, server,
profile);
@@ -538,6 +548,30 @@ public class StrategicSegmentAssigner implements
SegmentActionHandler
return numLoadsQueued;
}
+ /**
+ * Orders {@code eligibleServers} into the sequence a load should try them
in: round robin across the tier when
+ * round-robin assignment is enabled, else by the balancer strategy.
+ * <p>
+ * The round-robin branch ignores {@code eligibleServers} and derives its
candidates from the tier, keeping those
+ * that pass {@link ServerHolder#canLoadSegment} at the moment each one is
taken. Callers pass their complete
+ * eligible set: the two branches otherwise disagree on which servers a load
may target, and an empty
+ * {@code eligibleServers} does not imply an empty iterator.
+ * <p>
+ * Consume the result lazily. {@link RoundRobinServerSelector} advances a
per-tier cursor on every element taken, so
+ * draining the iterator for a segment that needs one replica advances the
cursor a full lap and hands the next
+ * segment the same starting server.
+ */
+ private Iterator<ServerHolder> serversToLoadSegment(
+ DataSegment segment,
+ String tier,
+ List<ServerHolder> eligibleServers
+ )
+ {
+ return useRoundRobinAssignment
+ ? serverSelector.getServersInTierToLoadSegment(tier, segment)
+ : strategy.findServersToLoadSegment(segment, eligibleServers);
+ }
+
/**
* Cancels up to {@code numToCancel} in-flight load operations across the
given list of servers. Successfully
* canceled servers are appended to {@code canceledOut} so the caller can
re-target them as fresh-load
@@ -951,10 +985,7 @@ public class StrategicSegmentAssigner implements
SegmentActionHandler
return 0;
}
- final Iterator<ServerHolder> serverIterator =
- useRoundRobinAssignment
- ? serverSelector.getServersInTierToLoadSegment(tier, segment)
- : strategy.findServersToLoadSegment(segment, eligibleServers);
+ final Iterator<ServerHolder> serverIterator =
serversToLoadSegment(segment, tier, eligibleServers);
if (!serverIterator.hasNext()) {
incrementSkipStat(Stats.Segments.ASSIGN_SKIPPED, "No strategic server",
segment, tier);
return 0;
diff --git
a/server/src/test/java/org/apache/druid/server/coordinator/duty/RunRulesPartialLoadPlacementTest.java
b/server/src/test/java/org/apache/druid/server/coordinator/duty/RunRulesPartialLoadPlacementTest.java
index 46e2284503e..bc2de75a1a6 100644
---
a/server/src/test/java/org/apache/druid/server/coordinator/duty/RunRulesPartialLoadPlacementTest.java
+++
b/server/src/test/java/org/apache/druid/server/coordinator/duty/RunRulesPartialLoadPlacementTest.java
@@ -30,12 +30,14 @@ import org.apache.druid.metadata.MetadataRuleManagerConfig;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.column.RowSignature;
import org.apache.druid.server.coordination.ServerType;
-import org.apache.druid.server.coordinator.CoordinatorDynamicConfig;
import org.apache.druid.server.coordinator.DruidCluster;
import org.apache.druid.server.coordinator.DruidCoordinatorRuntimeParams;
import org.apache.druid.server.coordinator.ServerHolder;
import org.apache.druid.server.coordinator.balancer.BalancerStrategy;
import org.apache.druid.server.coordinator.balancer.CostBalancerStrategy;
+import org.apache.druid.server.coordinator.loading.PartialLoadProfile;
+import org.apache.druid.server.coordinator.loading.SegmentAction;
+import org.apache.druid.server.coordinator.loading.SegmentHolder;
import org.apache.druid.server.coordinator.loading.SegmentLoadQueueManager;
import org.apache.druid.server.coordinator.loading.TestLoadQueuePeon;
import org.apache.druid.server.coordinator.rules.CannotMatchBehavior;
@@ -51,6 +53,7 @@ import org.apache.druid.timeline.DataSegment;
import org.apache.druid.timeline.SegmentId;
import org.apache.druid.timeline.partition.NumberedShardSpec;
import org.joda.time.DateTime;
+import org.joda.time.Duration;
import org.joda.time.Interval;
import org.joda.time.Period;
import org.junit.jupiter.api.AfterEach;
@@ -58,7 +61,9 @@ import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import java.util.ArrayList;
import java.util.Collections;
+import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
@@ -73,6 +78,13 @@ public class RunRulesPartialLoadPlacementTest
private static final String TIER = "tier1";
private static final Interval CHUNK = Intervals.of("2026-01-01/2026-01-02");
+ /**
+ * Announced profile for a replica loaded under some earlier rule. Its
fingerprint matches no rule used here, so
+ * the reconciler classifies any replica carrying it as stale.
+ */
+ private static final PartialLoadProfile STALE_PROFILE =
+ PartialLoadProfile.forLoaded(Map.of("type", "partialClusterGroup"),
"v1:previous-rule", 1L);
+
private ListeningExecutorService exec;
private BalancerStrategy balancerStrategy;
private SegmentLoadQueueManager loadQueueManager;
@@ -188,6 +200,121 @@ public class RunRulesPartialLoadPlacementTest
Assertions.assertFalse(stats.hasStat(Stats.Segments.ASSIGNED), "no
partition is fully downloaded");
}
+ /**
+ * Fresh partial loads must be distributed across the historicals of a tier.
With three empty historicals and one
+ * required replica, each historical should end up holding roughly a third
of the assignments.
+ */
+ @Test
+ public void test_freshPartialLoads_areSpreadAcrossHistoricals()
+ {
+ final List<DataSegment> segments = dailySegments(12, "acme");
+
+ final Map<String, TestLoadQueuePeon> peons = new LinkedHashMap<>();
+ final DruidCluster cluster = tierOf(peons, "hist1", "hist2", "hist3");
+
+ runRules(cluster, matchTenantForeverRule("acme"), segments.toArray(new
DataSegment[0]));
+
+ assertAssignmentsSpreadEvenly(peons, segments.size());
+ }
+
+ /**
+ * A rule change that invalidates every loaded replica must spread the
replacement loads across the tier. Each
+ * historical starts out holding a third of the replicas under a fingerprint
the new rule does not match; the
+ * reloads that reconcile them should again land roughly a third on each
historical.
+ */
+ @Test
+ public void
test_ruleChangeInvalidatingEveryReplica_spreadsReloadsAcrossHistoricals()
+ {
+ final List<DataSegment> segments = dailySegments(12, "acme");
+
+ // Deal the segments evenly across the three historicals under a
fingerprint the new rule will not match.
+ final Map<String, List<DataSegment>> preloaded = new LinkedHashMap<>();
+ final List<String> names = List.of("hist1", "hist2", "hist3");
+ names.forEach(name -> preloaded.put(name, new ArrayList<>()));
+ for (int i = 0; i < segments.size(); i++) {
+ preloaded.get(names.get(i % names.size())).add(segments.get(i));
+ }
+
+ final Map<String, TestLoadQueuePeon> peons = new LinkedHashMap<>();
+ final DruidCluster cluster = tierOfPreloadedHistoricals(peons, preloaded);
+
+ runRules(cluster, matchTenantForeverRule("acme"), segments.toArray(new
DataSegment[0]));
+
+ assertAssignmentsSpreadEvenly(peons, segments.size());
+ }
+
+ /**
+ * A deficit larger than the number of historicals that can take a fresh
load must still queue each historical at
+ * most once. The historical whose stale in-flight load is cancelled becomes
a fresh-load candidate, and the
+ * stale-loaded historical takes an additive reload; a skip stat here means
one of them was offered twice.
+ */
+ @Test
+ public void
test_ruleChangeWithStaleInFlightAndStaleLoaded_queuesEachHistoricalOnce()
+ {
+ final DataSegment segment = dailySegments(1, "acme").getFirst();
+
+ // hist1 is mid-load under a fingerprint the new rule does not match, so
cancelling frees it for a fresh load.
+ final TestLoadQueuePeon inFlightPeon = new TestLoadQueuePeon();
+ inFlightPeon.addInFlightHolder(
+ new SegmentHolder(segment, SegmentAction.LOAD, STALE_PROFILE,
Duration.standardSeconds(10), null)
+ );
+ final ServerHolder hist1 = new
ServerHolder(server("hist1").toImmutableDruidServer(), inFlightPeon);
+
+ // hist2 already serves the segment under that same stale fingerprint, so
it can only reload additively.
+ final DruidServer staleServer = server("hist2");
+ staleServer.addDataSegment(segment, STALE_PROFILE);
+ final TestLoadQueuePeon stalePeon = new TestLoadQueuePeon();
+ final ServerHolder hist2 = new
ServerHolder(staleServer.toImmutableDruidServer(), stalePeon);
+
+ final DruidCluster cluster = DruidCluster.builder().addTier(TIER, hist1,
hist2).build();
+
+ final CoordinatorRunStats stats = runRules(cluster,
matchTenantForeverRule("acme", 2), segment);
+
+ Assertions.assertEquals(
+ 1L,
+ stats.getSegmentStat(Stats.Segments.PARTIAL_STALE_CANCELLED, TIER,
DATASOURCE),
+ "the stale in-flight load is cancelled"
+ );
+ Assertions.assertEquals(
+ 2L,
+ stats.getSegmentStat(Stats.Segments.PARTIAL_ASSIGNED, TIER,
DATASOURCE),
+ "both required replicas are queued"
+ );
+ Assertions.assertFalse(
+ stats.hasStat(Stats.Segments.ASSIGN_SKIPPED),
+ "each historical is offered once, so no assignment is skipped"
+ );
+ Assertions.assertNotEquals(
+ STALE_PROFILE,
+ inFlightPeon.getProfileFor(segment),
+ "the cancelled historical is reloaded under the rule's fingerprint"
+ );
+ Assertions.assertNotNull(
+ stalePeon.getProfileFor(segment),
+ "the stale-loaded historical is queued an additive reload"
+ );
+ }
+
+ /**
+ * Asserts that every segment was assigned somewhere and that no historical
carries more than one assignment more
+ * than any other. The failure message reports the per-historical counts so
an uneven split is readable directly.
+ */
+ private static void assertAssignmentsSpreadEvenly(Map<String,
TestLoadQueuePeon> peons, int expectedTotal)
+ {
+ final Map<String, Integer> counts = new LinkedHashMap<>();
+ peons.forEach((name, peon) -> counts.put(name,
peon.getSegmentsToLoad().size()));
+
+ final int total =
counts.values().stream().mapToInt(Integer::intValue).sum();
+ Assertions.assertEquals(expectedTotal, total, "every segment is assigned
somewhere, but counts were " + counts);
+
+ final int max =
counts.values().stream().mapToInt(Integer::intValue).max().orElse(0);
+ final int min =
counts.values().stream().mapToInt(Integer::intValue).min().orElse(0);
+ Assertions.assertTrue(
+ max - min <= 1,
+ "assignments are spread across the tier's historicals, but counts were
" + counts
+ );
+ }
+
private ForeverPartialLoadRule matchNobodyForeverRule()
{
// Include pattern resolves against the "tenant" clustering column
(compatible) but matches none of the segments'
@@ -221,12 +348,6 @@ public class RunRulesPartialLoadPlacementTest
)
.withUsedSegments(segments)
.withBalancerStrategy(balancerStrategy)
- .withDynamicConfigs(
- CoordinatorDynamicConfig.builder()
- .withSmartSegmentLoading(false)
- .withUseRoundRobinSegmentAssignment(false)
- .build()
- )
.withSegmentAssignerUsing(loadQueueManager)
.build();
@@ -234,11 +355,62 @@ public class RunRulesPartialLoadPlacementTest
return params.getCoordinatorStats();
}
+ private static ForeverPartialLoadRule matchTenantForeverRule(String tenant)
+ {
+ return matchTenantForeverRule(tenant, 1);
+ }
+
+ private static ForeverPartialLoadRule matchTenantForeverRule(String tenant,
int replicas)
+ {
+ return new ForeverPartialLoadRule(
+ ImmutableMap.of(TIER, replicas),
+ null,
+ new WildcardClusterGroupPartialLoadMatcher(List.of(Map.of("tenant",
tenant)), null),
+ CannotMatchBehavior.FULL_LOAD
+ );
+ }
+
private static DruidCluster singleTierCluster()
{
return DruidCluster.builder().addTier(TIER, historical("hist1",
TIER)).build();
}
+ /**
+ * Builds a tier of empty historicals, exposing each one's peon by name so a
test can read back what was assigned
+ * to it.
+ */
+ private static DruidCluster tierOf(Map<String, TestLoadQueuePeon> peonsOut,
String... names)
+ {
+ final DruidCluster.Builder cluster = DruidCluster.builder();
+ for (String name : names) {
+ final TestLoadQueuePeon peon = new TestLoadQueuePeon();
+ peonsOut.put(name, peon);
+ cluster.add(new ServerHolder(server(name).toImmutableDruidServer(),
peon));
+ }
+ return cluster.build();
+ }
+
+ /**
+ * Builds a tier whose historicals already serve the given segments under
{@link #STALE_PROFILE}, so a rule
+ * resolving to any other fingerprint sees every replica as stale.
+ */
+ private static DruidCluster tierOfPreloadedHistoricals(
+ Map<String, TestLoadQueuePeon> peonsOut,
+ Map<String, List<DataSegment>> nameToLoadedSegments
+ )
+ {
+ final DruidCluster.Builder cluster = DruidCluster.builder();
+ nameToLoadedSegments.forEach((name, loaded) -> {
+ final DruidServer server = server(name);
+ loaded.forEach(segment -> server.addDataSegment(segment, STALE_PROFILE));
+
+ final TestLoadQueuePeon peon = new TestLoadQueuePeon();
+ peonsOut.put(name, peon);
+ cluster.add(new ServerHolder(server.toImmutableDruidServer(), peon));
+ });
+ return cluster.build();
+ }
+
private static ServerHolder historical(String name, String tier)
{
final DruidServer server =
@@ -246,6 +418,27 @@ public class RunRulesPartialLoadPlacementTest
return new ServerHolder(server.toImmutableDruidServer(), new
TestLoadQueuePeon());
}
+ private static DruidServer server(String name)
+ {
+ return new DruidServer(name, name, null, 10L << 30, null,
ServerType.HISTORICAL, TIER, 0);
+ }
+
+ /** One single-partition segment per day, all carrying the same tenant
tuple. */
+ private static List<DataSegment> dailySegments(int count, String tenant)
+ {
+ final ClusterGroupTuples groups = new ClusterGroupTuples(
+ RowSignature.builder().add("tenant", ColumnType.STRING).build(),
+ List.of(Collections.singletonList(tenant))
+ );
+
+ final List<DataSegment> segments = new ArrayList<>(count);
+ for (int i = 0; i < count; i++) {
+ final Interval day = new Interval(CHUNK.getStart().plusDays(i),
CHUNK.getStart().plusDays(i + 1));
+ segments.add(segment(day, new NumberedShardSpec(0, 1), groups));
+ }
+ return segments;
+ }
+
private static DataSegment clusteredSegment(NumberedShardSpec shardSpec,
String tenant)
{
final ClusterGroupTuples groups = new ClusterGroupTuples(
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 a9cfb41bbc5..f6bd341fd86 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
@@ -615,6 +615,51 @@ public class StrategicSegmentAssignerPartialTest
Assertions.assertNull(((TestLoadQueuePeon)
decommServer.getPeon()).getProfileFor(segment));
}
+ @Test
+ public void
testCancelledStaleInFlightIsNotReloadedOntoDecommissioningServer()
+ {
+ // The tier's only server is decommissioning and carries a
stale-fingerprint in-flight load. Cancelling that load
+ // frees its slot, but canLoadSegment still rejects the server, so nothing
is queued back onto it.
+ final DataSegment segment = createSegment();
+ final PartialLoadProfile staleInFlightProfile =
PartialLoadProfile.forRequest(
+ Map.of(
+ "type", "partialProjection",
+ "projections", List.of("users"),
+ "fingerprint", FP_USERS
+ ),
+ FP_USERS
+ );
+ final TestLoadQueuePeon peon = new TestLoadQueuePeon();
+ peon.addInFlightHolder(new SegmentHolder(
+ segment,
+ SegmentAction.LOAD,
+ staleInFlightProfile,
+ Duration.standardSeconds(10),
+ null
+ ));
+ final ServerHolder decommServer =
+ new ServerHolder(createDruidServer(TIER1).toImmutableDruidServer(),
peon, true);
+ final DruidCluster cluster = DruidCluster.builder().addTier(TIER1,
decommServer).build();
+
+ final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster,
segment);
+ params.getSegmentAssigner()
+ .replicateSegmentPartially(segment, profileForRevenue(),
ImmutableMap.of(TIER1, 1));
+
+ final CoordinatorRunStats stats = params.getCoordinatorStats();
+ Assertions.assertEquals(
+ 1L,
+ stats.getSegmentStat(Stats.Segments.PARTIAL_STALE_CANCELLED, TIER1,
segment.getDataSource()),
+ "the stale in-flight load is still cancelled"
+ );
+ Assertions.assertEquals(
+ 0L,
+ stats.getSegmentStat(Stats.Segments.PARTIAL_ASSIGNED, TIER1,
segment.getDataSource()),
+ "no partial load is queued on a decommissioning server"
+ );
+ Assertions.assertTrue(decommServer.getLoadingSegments().isEmpty());
+ Assertions.assertNull(peon.getProfileFor(segment));
+ }
+
@Test
public void testStaleInFlightCancelledAndReplaced()
{
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
