diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/PartialLoadHistoricalCloningTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/PartialLoadHistoricalCloningTest.java
new file mode 100644
index 000000000000..4ab23c6814e2
--- /dev/null
+++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/PartialLoadHistoricalCloningTest.java
@@ -0,0 +1,329 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.testing.embedded.server;
+
+import org.apache.druid.common.utils.IdUtils;
+import org.apache.druid.data.input.impl.AggregateProjectionSpec;
+import org.apache.druid.data.input.impl.ClusteredValueGroupsBaseTableProjectionSpec;
+import org.apache.druid.data.input.impl.LongDimensionSchema;
+import org.apache.druid.data.input.impl.StringDimensionSchema;
+import org.apache.druid.data.input.impl.TimestampSpec;
+import org.apache.druid.indexer.granularity.SegmentGranularitySpec;
+import org.apache.druid.indexing.common.task.TaskBuilder;
+import org.apache.druid.indexing.common.task.batch.parallel.ParallelIndexSupervisorTask;
+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.aggregation.LongMinAggregatorFactory;
+import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
+import org.apache.druid.server.coordinator.CoordinatorDynamicConfig;
+import org.apache.druid.server.coordinator.rules.CannotMatchBehavior;
+import org.apache.druid.server.coordinator.rules.ForeverPartialLoadRule;
+import org.apache.druid.server.coordinator.rules.WildcardProjectionPartialLoadMatcher;
+import org.apache.druid.testing.embedded.EmbeddedBroker;
+import org.apache.druid.testing.embedded.EmbeddedCoordinator;
+import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
+import org.apache.druid.testing.embedded.EmbeddedHistorical;
+import org.apache.druid.testing.embedded.EmbeddedIndexer;
+import org.apache.druid.testing.embedded.EmbeddedOverlord;
+import org.apache.druid.testing.embedded.EmbeddedRouter;
+import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase;
+import org.apache.druid.testing.embedded.utils.ITRetryUtil;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * End-to-end coverage for cloning a historical that loads segments partially. The source historical loads only the
+ * bundles a {@link ForeverPartialLoadRule} selects, so it reports a footprint smaller than the full segment. Its
+ * clone is expected to hold and report the same footprint: cloning copies the source's load state, and under a
+ * partial-load rule that state includes which parts of the segment were loaded.
+ *
+ * Both historicals are configured identically for partial downloads, so any difference in reported {@code curr_size}
+ * comes from the load request itself rather than from node configuration.
+ */
+public class PartialLoadHistoricalCloningTest extends EmbeddedClusterTestBase
+{
+ private static final String PROJECTION_NAME = "country_delta";
+ // Ingested alongside country_delta but not selected by the rule, so its container bytes stay off the historical's
+ // disk. That is what makes the rule-loaded footprint measurably smaller than the full segment size.
+ private static final String UNMATCHED_PROJECTION_NAME = "country_min_delta";
+
+ private static final long CACHE_SIZE = HumanReadableBytes.parse("1MiB");
+ private static final long MAX_SIZE = HumanReadableBytes.parse("100MiB");
+ private static final long ESTIMATE_SIZE = HumanReadableBytes.parse("2KiB");
+
+ private static final String CLONE_PORT = "7083";
+
+ private final EmbeddedBroker broker = new EmbeddedBroker();
+ private final EmbeddedIndexer indexer = new EmbeddedIndexer();
+ private final EmbeddedOverlord overlord = new EmbeddedOverlord();
+ private final EmbeddedCoordinator coordinator = new EmbeddedCoordinator();
+ private final EmbeddedRouter router = new EmbeddedRouter();
+
+ private final EmbeddedHistorical sourceHistorical = new EmbeddedHistorical();
+ private final EmbeddedHistorical cloneHistorical =
+ new EmbeddedHistorical().addProperty("druid.plaintextPort", CLONE_PORT);
+
+ @Override
+ public EmbeddedDruidCluster createCluster()
+ {
+ configureForPartialDownloads(sourceHistorical);
+ configureForPartialDownloads(cloneHistorical);
+
+ broker.setServerMemory(200_000_000)
+ .addProperty("druid.sql.planner.enableSysQueriesTable", "true");
+
+ coordinator.addProperty("druid.manager.segments.useIncrementalCache", "always");
+
+ overlord.addProperty("druid.manager.segments.useIncrementalCache", "always")
+ .addProperty("druid.manager.segments.pollDuration", "PT0.1s");
+
+ indexer.setServerMemory(300_000_000)
+ .addProperty("druid.worker.capacity", "2")
+ .addProperty("druid.processing.numThreads", "2")
+ .addProperty("druid.segment.handoff.pollDuration", "PT0.1s");
+
+ return EmbeddedDruidCluster
+ .withEmbeddedDerbyAndZookeeper()
+ .useLatchableEmitter()
+ .useDefaultTimeoutForLatchableEmitter(60)
+ .addCommonProperty("druid.indexer.task.buildV10", "true")
+ .addCommonProperty("druid.storage.type", "local")
+ .addCommonProperty("druid.storage.zip", "false")
+ .addServer(coordinator)
+ .addServer(overlord)
+ .addServer(indexer)
+ .addServer(sourceHistorical)
+ .addServer(cloneHistorical)
+ .addServer(broker)
+ .addServer(router);
+ }
+
+ private void configureForPartialDownloads(EmbeddedHistorical historical)
+ {
+ historical.setServerMemory(500_000_000)
+ .addProperty("druid.segmentCache.virtualStorage", "true")
+ .addProperty("druid.segmentCache.virtualStoragePartialDownloadsEnabled", "true")
+ .addProperty(
+ "druid.segmentCache.virtualStorageMetadataReservationEstimate",
+ String.valueOf(ESTIMATE_SIZE)
+ )
+ .addProperty(
+ "druid.segmentCache.virtualStorageLoadThreads",
+ String.valueOf(Runtime.getRuntime().availableProcessors())
+ )
+ .addBeforeStartHook(
+ (cluster, self) -> self.addProperty(
+ "druid.segmentCache.locations",
+ StringUtils.format(
+ "[{\"path\":\"%s\",\"maxSize\":\"%s\"}]",
+ cluster.getTestFolder().newFolder().getAbsolutePath(),
+ CACHE_SIZE
+ )
+ )
+ )
+ .addProperty("druid.server.maxSize", String.valueOf(MAX_SIZE));
+ }
+
+ @BeforeAll
+ void loadDataAndConfigureCloning() throws IOException
+ {
+ dataSource = "partial-clone-" + IdUtils.getRandomId();
+
+ // The rule and the clone mapping are both configured before ingestion so the first coordinator run already sees
+ // the clone target as unmanaged: rule-driven assignment can only pick the source, and everything the clone gets
+ // comes from the cloning duty.
+ 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
+ )
+ )
+ )
+ );
+ cluster.callApi().onLeaderCoordinator(
+ c -> c.updateCoordinatorDynamicConfig(
+ CoordinatorDynamicConfig
+ .builder()
+ .withCloneServers(Map.of(cloneHost(), sourceHost()))
+ .build()
+ )
+ );
+
+ ingestClusteredSegmentWithProjection();
+ }
+
+ @Override
+ protected void refreshDatasourceName()
+ {
+ // Fixed datasource across tests — rule, clone mapping and ingest are one-time setup.
+ }
+
+ @Test
+ void testCloneReportsTheSamePartialFootprintAsItsSource()
+ {
+ coordinator.latchableEmitter().waitForEventAggregate(
+ event -> event.hasMetricName("segment/clone/assigned/count")
+ .hasDimension("server", cloneHost()),
+ agg -> agg.hasSumAtLeast(1)
+ );
+ coordinator.latchableEmitter().waitForEventAggregate(
+ event -> event.hasMetricName("segment/loadQueue/success")
+ .hasDimension("server", cloneHost())
+ .hasDimension(DruidMetrics.DATASOURCE, dataSource),
+ agg -> agg.hasSumAtLeast(1)
+ );
+
+ // The load announcement reaches the broker's inventory asynchronously; wait until both historicals have reported
+ // a footprint before comparing them.
+ ITRetryUtil.retryUntilTrue(
+ () -> currSizeOf(sourceHost()) > 0 && currSizeOf(cloneHost()) > 0,
+ "both historicals to report a non-zero curr_size"
+ );
+
+ final long fullSize = Long.parseLong(
+ cluster.callApi().runSql(
+ "SELECT \"size\" FROM sys.segments WHERE datasource = '" + dataSource + "'"
+ ).trim()
+ );
+ final long sourceSize = currSizeOf(sourceHost());
+ final long cloneSize = currSizeOf(cloneHost());
+
+ Assertions.assertTrue(
+ sourceSize < fullSize,
+ StringUtils.format(
+ "source should hold only the rule-selected parts; got curr_size=%d, full segment size=%d",
+ sourceSize,
+ fullSize
+ )
+ );
+ Assertions.assertEquals(
+ sourceSize,
+ cloneSize,
+ StringUtils.format(
+ "clone should hold the same parts as its source; source curr_size=%d, clone curr_size=%d, "
+ + "full segment size=%d (a clone loaded without the source's partial-load profile downloads the whole "
+ + "segment and reports its full size)",
+ sourceSize,
+ cloneSize,
+ fullSize
+ )
+ );
+ }
+
+ private long currSizeOf(String host)
+ {
+ final String result = cluster.callApi().runSql(
+ "SELECT curr_size FROM sys.servers WHERE server_type = 'historical' AND server = '" + host + "'"
+ ).trim();
+ return result.isEmpty() ? 0L : Long.parseLong(result);
+ }
+
+ private String sourceHost()
+ {
+ return sourceHistorical.bindings().selfNode().getHostAndPort();
+ }
+
+ private String cloneHost()
+ {
+ return cloneHistorical.bindings().selfNode().getHostAndPort();
+ }
+
+ /**
+ * Ingests a single clustered base-table segment (clustered by {@code channel}) with a {@code country_delta}
+ * aggregate projection (group by {@code countryName}, sum {@code delta}) plus a second projection the rule does
+ * not select.
+ */
+ private void ingestClusteredSegmentWithProjection() throws IOException
+ {
+ final File tmpDir = cluster.getTestFolder().newFolder();
+ final File inputFile = new File(tmpDir, "clustered-input.json");
+ final String inputData =
+ "{\"time\":\"2024-01-01T00:10:00Z\",\"channel\":\"#en\",\"countryName\":\"US\",\"delta\":10}\n"
+ + "{\"time\":\"2024-01-01T00:20:00Z\",\"channel\":\"#en\",\"countryName\":\"US\",\"delta\":5}\n"
+ + "{\"time\":\"2024-01-01T00:30:00Z\",\"channel\":\"#en\",\"countryName\":\"CA\",\"delta\":3}\n"
+ + "{\"time\":\"2024-01-01T00:40:00Z\",\"channel\":\"#fr\",\"countryName\":\"FR\",\"delta\":7}\n"
+ + "{\"time\":\"2024-01-01T00:50:00Z\",\"channel\":\"#fr\",\"countryName\":\"US\",\"delta\":2}\n";
+ Files.write(inputFile.toPath(), inputData.getBytes(StandardCharsets.UTF_8));
+
+ final ClusteredValueGroupsBaseTableProjectionSpec clusterSpec =
+ ClusteredValueGroupsBaseTableProjectionSpec.builder()
+ .columns(
+ new StringDimensionSchema("channel"),
+ new StringDimensionSchema("countryName"),
+ new LongDimensionSchema("delta"),
+ new LongDimensionSchema("__time")
+ )
+ .clusteringColumns("channel")
+ .build();
+
+ final AggregateProjectionSpec projection =
+ AggregateProjectionSpec.builder(PROJECTION_NAME)
+ .groupingColumns(new StringDimensionSchema("countryName"))
+ .aggregators(new LongSumAggregatorFactory("sumDelta", "delta"))
+ .build();
+
+ final AggregateProjectionSpec unmatchedProjection =
+ AggregateProjectionSpec.builder(UNMATCHED_PROJECTION_NAME)
+ .groupingColumns(new StringDimensionSchema("countryName"))
+ .aggregators(new LongMinAggregatorFactory("minDelta", "delta"))
+ .build();
+
+ final SegmentGranularitySpec segmentGranularitySpec = new SegmentGranularitySpec(
+ Granularities.HOUR,
+ List.of(Intervals.of("2024-01-01/2024-01-02"))
+ );
+
+ final String taskId = IdUtils.getRandomId();
+ final ParallelIndexSupervisorTask task = TaskBuilder
+ .ofTypeIndexParallel()
+ .jsonInputFormat()
+ .localInputSourceWithFiles(inputFile)
+ .dataSchema(
+ builder -> builder
+ .withDataSource(dataSource)
+ .withTimestamp(new TimestampSpec("time", "iso", null))
+ .withSegmentGranularity(segmentGranularitySpec)
+ .withBaseTable(clusterSpec)
+ .withProjections(List.of(projection, unmatchedProjection))
+ )
+ .tuningConfig(t -> t.withMaxNumConcurrentSubTasks(1))
+ .withId(taskId);
+
+ cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task));
+ cluster.callApi().waitForTaskToSucceed(taskId, overlord);
+ cluster.callApi().waitForAllSegmentsToBeAvailable(dataSource, coordinator, broker);
+ }
+}
diff --git a/server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java b/server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java
index 8ee13001f237..4b442bfd2f00 100644
--- a/server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java
+++ b/server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java
@@ -465,13 +465,13 @@ private void deleteSegmentInfoFile(DataSegment segment)
}
/**
- * Write the info file for a partial-load segment, overwriting any existing content atomically. Distinct from
- * {@link #storeInfoFile} which skips the write when the file already exists, for partial segments we must
- * unconditionally rewrite so an incoming rule swap (new {@code fingerprint}/{@code delegate} inside the
- * wrapped load spec) reaches disk. Otherwise bootstrap after a restart would restore the segment using the
- * prior wrapper and re-announce the old rule until the coordinator resyncs.
+ * Write the info file for a segment, overwriting any existing content atomically. Distinct from
+ * {@link #storeInfoFile}, which skips the write when the file already exists: a partial-load transition must reach
+ * disk unconditionally, whether it is a rule swap (new {@code fingerprint}/{@code delegate} inside the wrapped load
+ * spec) or a return to a regular full load (no wrapper at all). Otherwise bootstrap after a restart would restore
+ * the segment using the prior wrapper and re-announce the old rule until the coordinator resyncs.
*/
- private void writePartialInfoFile(DataSegment segment) throws IOException
+ private void rewriteInfoFile(DataSegment segment) throws IOException
{
final File segmentInfoCacheFile = new File(getEffectiveInfoDir(), segment.getId().toString());
FileUtils.mkdirp(getEffectiveInfoDir());
@@ -903,7 +903,7 @@ private ReservedPartial reservePartial(DataSegment dataSegment, SegmentRangeRead
location.getPath()
);
}
- writePartialInfoFile(dataSegment);
+ rewriteInfoFile(dataSegment);
partial.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment));
return new ReservedPartial(partial, location, hold);
}
@@ -1041,7 +1041,7 @@ private DataSegment loadPartial(DataSegment dataSegment) throws SegmentLoadingEx
// branch. On the find-existing branch the info file on disk still carries the PRIOR rule's wrapped
// load spec, so a rule swap here would apply in memory only. Rewrite unconditionally before mount.
try {
- writePartialInfoFile(dataSegment);
+ rewriteInfoFile(dataSegment);
}
catch (IOException e) {
throw new SegmentLoadingException(
@@ -1338,15 +1338,21 @@ public DataSegment load(final DataSegment dataSegment) throws SegmentLoadingExce
return loadPartial(dataSegment);
}
// virtual storage doesn't do anything with loading immediately, but check to see if the segment is already cached
- // and if so, clear out the onUnmount action
+ // and if so, clear out the onUnmount action. Reaching here with a rule applied means the coordinator asked for
+ // the whole segment again: a rule is only ever applied on the loadPartial path above, so the request that got
+ // here carries no partial-load wrapper for this segment. Release the rule.
final ReferenceCountingLock lock = lock(dataSegment);
synchronized (lock) {
try {
final SegmentCacheEntryIdentifier cacheEntryIdentifier = new SegmentCacheEntryIdentifier(dataSegment.getId());
for (StorageLocation location : locations) {
final SegmentCacheEntry cacheEntry = location.getCacheEntry(cacheEntryIdentifier);
- if (cacheEntry != null) {
- cacheEntry.setOnUnmount(null);
+ if (cacheEntry == null) {
+ continue;
+ }
+ cacheEntry.setOnUnmount(null);
+ if (cacheEntry instanceof PartialSegmentMetadataCacheEntry partial && partial.isRuleHeld()) {
+ releaseRuleForFullLoad(dataSegment, partial);
}
}
}
@@ -1426,7 +1432,7 @@ public DataSegment bootstrap(
reapplyRuleFromInfoFile(dataSegment, partial);
loadedProfile = PartialLoadProfile.forLoaded(
dataSegment.getLoadSpec(),
- (String) dataSegment.getLoadSpec().get("fingerprint"),
+ (String) dataSegment.getLoadSpec().get(PartialLoadSpec.FINGERPRINT_FIELD),
partial.getRealizedBytes()
);
}
@@ -1526,6 +1532,48 @@ public void drop(final DataSegment segment)
}
}
+ /**
+ * Releases the partial-load rule applied to {@code dataSegment} in response to an unwrapped load request: the
+ * coordinator has stopped asking for parts of the segment, so the metadata entry and the rule's bundles are unpinned.
+ * That is what a full load means under virtual storage — nothing is pinned, each part is fetched on demand — and
+ * reclaim of the partial state on disk is left to eviction, as it is for {@link #drop}.
+ *
+ * The info file is rewritten before the rule is cleared, and a failed rewrite fails the load. Nothing is left half
+ * converted: releasing the holds cannot fail, and a load failure sends the historical down its drop path, which
+ * clears the rule and removes the info file, so there is no stale rule for a restart to reinstate. Leaving the rule
+ * applied and carrying on is not an option, because an unwrapped request announces as a full load either way, so the
+ * coordinator would record a replica with no profile and never ask again.
+ *
+ * Callers must hold this segment's {@link #lock(DataSegment)}, which is the external lock that
+ * {@link PartialSegmentMetadataCacheEntry#clearRule} requires to be serialized against
+ * {@link PartialSegmentMetadataCacheEntry#applyRule}.
+ */
+ private void releaseRuleForFullLoad(DataSegment dataSegment, PartialSegmentMetadataCacheEntry partial)
+ throws SegmentLoadingException
+ {
+ // Snapshot both before clearRule zeroes out the rule state so the log can describe what was released.
+ final String priorFingerprint = partial.getRuleFingerprint();
+ final long priorRealizedBytes = partial.getRealizedBytes();
+ try {
+ rewriteInfoFile(dataSegment);
+ }
+ catch (IOException e) {
+ throw new SegmentLoadingException(
+ e,
+ "Failed to rewrite info file for segment[%s] while releasing partial-load rule[fingerprint=%s]",
+ dataSegment.getId(),
+ priorFingerprint
+ );
+ }
+ partial.clearRule();
+ log.info(
+ "Released partial-load rule[fingerprint=%s, realizedBytes=%d] for segment[%s]; it is a regular full load now.",
+ priorFingerprint,
+ priorRealizedBytes,
+ dataSegment.getId()
+ );
+ }
+
/**
* Reapply the persisted partial-load rule to a bootstrap-restored metadata entry. Reads the wrapper from the
* segment's info-file {@code loadSpec}, resolves the selected bundle names against the just-parsed on-disk
diff --git a/server/src/main/java/org/apache/druid/server/coordination/SegmentChangeRequestLoad.java b/server/src/main/java/org/apache/druid/server/coordination/SegmentChangeRequestLoad.java
index c6b5f4ac2904..d6f4a04289e5 100644
--- a/server/src/main/java/org/apache/druid/server/coordination/SegmentChangeRequestLoad.java
+++ b/server/src/main/java/org/apache/druid/server/coordination/SegmentChangeRequestLoad.java
@@ -74,7 +74,11 @@ public static SegmentChangeRequestLoad forAnnouncement(DataSegment segment)
final Map loadSpec = segment.getLoadSpec();
if (PartialLoadSpec.detectPartialLoadSpec(loadSpec)) {
// Historical didn't wrap, treat as full-fallback: fingerprint from the loadSpec, loadedBytes = full size.
- return new SegmentChangeRequestLoad(segment, (String) loadSpec.get("fingerprint"), segment.getSize());
+ return new SegmentChangeRequestLoad(
+ segment,
+ (String) loadSpec.get(PartialLoadSpec.FINGERPRINT_FIELD),
+ segment.getSize()
+ );
}
if (PartialLoadSpec.hasPartialTypePrefix(loadSpec)) {
// Type name claims partial-load but the wire form is malformed, the PartialLoadSpec subtype's @JsonProperty
@@ -85,8 +89,8 @@ public static SegmentChangeRequestLoad forAnnouncement(DataSegment segment)
+ "announcing as a regular load.",
segment.getId(),
loadSpec.get("type"),
- loadSpec.get("fingerprint"),
- loadSpec.get("delegate")
+ loadSpec.get(PartialLoadSpec.FINGERPRINT_FIELD),
+ loadSpec.get(PartialLoadSpec.DELEGATE_FIELD)
);
}
return new SegmentChangeRequestLoad(segment);
diff --git a/server/src/main/java/org/apache/druid/server/coordinator/ServerHolder.java b/server/src/main/java/org/apache/druid/server/coordinator/ServerHolder.java
index 2b9362e82741..0e137caed149 100644
--- a/server/src/main/java/org/apache/druid/server/coordinator/ServerHolder.java
+++ b/server/src/main/java/org/apache/druid/server/coordinator/ServerHolder.java
@@ -468,6 +468,24 @@ public PartialLoadProfile getInFlightProfile(DataSegment segment)
return inFlightProfiles.get(segment);
}
+ /**
+ * The {@link PartialLoadProfile} this server is expected to hold {@code segment} under once its queued operations
+ * finish: the profile of an in-flight load if one is queued, else the profile announced for the loaded replica.
+ * Returns null when the replica is (or is becoming) a regular full load, including an in-flight load that carries
+ * no profile.
+ *
+ * Read this before {@link #cancelOperation}, which clears the in-flight profile.
+ */
+ @Nullable
+ public PartialLoadProfile getProjectedProfile(DataSegment segment)
+ {
+ final SegmentAction action = getActionOnSegment(segment);
+ if (action != null && action.isLoad()) {
+ return getInFlightProfile(segment);
+ }
+ return server.getPartialLoadProfile(segment.getId());
+ }
+
private boolean hasSegmentLoaded(SegmentId segmentId)
{
return server.getSegment(segmentId) != null;
diff --git a/server/src/main/java/org/apache/druid/server/coordinator/duty/CloneHistoricals.java b/server/src/main/java/org/apache/druid/server/coordinator/duty/CloneHistoricals.java
index c12193c9b0e5..902ed31a5733 100644
--- a/server/src/main/java/org/apache/druid/server/coordinator/duty/CloneHistoricals.java
+++ b/server/src/main/java/org/apache/druid/server/coordinator/duty/CloneHistoricals.java
@@ -27,6 +27,7 @@
import org.apache.druid.server.coordinator.DruidCoordinatorRuntimeParams;
import org.apache.druid.server.coordinator.ServerCloneStatus;
import org.apache.druid.server.coordinator.ServerHolder;
+import org.apache.druid.server.coordinator.loading.PartialLoadProfile;
import org.apache.druid.server.coordinator.loading.SegmentAction;
import org.apache.druid.server.coordinator.loading.SegmentLoadQueueManager;
import org.apache.druid.server.coordinator.stats.Dimension;
@@ -38,6 +39,7 @@
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
+import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
@@ -45,6 +47,13 @@
* Handles cloning of historicals. Given the historical to historical clone mappings, based on
* {@link CoordinatorDynamicConfig#getCloneServers()}, copies any segments load or unload requests from the source
* historical to the target historical.
+ *
+ * Under a partial-load rule the source holds only part of a segment, so copying its load state means copying the
+ * {@link PartialLoadProfile} it holds the segment under, not just the segment id. Replicas are therefore compared by
+ * profile fingerprint: a clone whose replica was loaded under a different profile than the source's is re-loaded with
+ * the source's profile, including a source that has stopped loading partially, for which the clone is re-loaded
+ * without one. Clone targets are excluded from rule-driven assignment
+ * ({@link DruidCluster#getManagedHistoricals()}), so this duty is the only thing that can correct them.
*/
public class CloneHistoricals implements CoordinatorDuty
{
@@ -103,10 +112,13 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params)
final Set sourceProjectedSegments = sourceServer.getProjectedSegments();
final Set targetProjectedSegments = targetServer.getProjectedSegments();
- // Load any segments missing in the clone target.
+ // Load any segment that the clone target is missing, or that it holds under a different partial-load profile
+ // than the source. Segment identity alone can't tell those apart: two replicas of the same segment id may hold
+ // different parts of it.
for (DataSegment segment : sourceProjectedSegments) {
- if (!targetProjectedSegments.contains(segment)) {
- loadSegmentOnTargetServer(segment, targetServer, params);
+ final PartialLoadProfile sourceProfile = sourceServer.getProjectedProfile(segment);
+ if (shouldLoadSegmentOnTargetServer(segment, sourceProfile, targetServer, targetProjectedSegments)) {
+ loadSegmentOnTargetServer(segment, sourceProfile, targetServer, params);
}
}
@@ -124,8 +136,13 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params)
return params;
}
+ /**
+ * Queues a load of {@code segment} on the clone target, asking for the same parts of the segment that the source
+ * holds. A null {@code sourceProfile} means the source holds the whole segment, which is the regular full load.
+ */
private void loadSegmentOnTargetServer(
DataSegment segment,
+ @Nullable PartialLoadProfile sourceProfile,
ServerHolder targetServer,
DruidCoordinatorRuntimeParams params
)
@@ -141,7 +158,12 @@ private void loadSegmentOnTargetServer(
rowKey.and(Dimension.DESCRIPTION, "Segment not found in metadata cache"),
1L
);
- } else if (loadQueueManager.loadSegment(loadableSegment, targetServer, SegmentAction.LOAD)) {
+ } else if (loadQueueManager.loadSegment(
+ loadableSegment,
+ targetServer,
+ SegmentAction.LOAD,
+ sourceProfile == null ? null : sourceProfile.asCloneRequest()
+ )) {
params.getCoordinatorStats().add(
Stats.Segments.ASSIGNED_TO_CLONE,
rowKey.build(),
@@ -150,6 +172,12 @@ private void loadSegmentOnTargetServer(
}
}
+ @Nullable
+ private static String fingerprintOf(@Nullable PartialLoadProfile profile)
+ {
+ return profile == null ? null : profile.fingerprint();
+ }
+
private void dropSegmentFromTargetServer(
DataSegment segment,
ServerHolder targetServer,
@@ -234,4 +262,27 @@ private Map createCurrentStatusMap(
return newStatusMap;
}
+
+ /**
+ * Determine whether a segment should be loaded on the target server.
+ *
+ * If the target server does not have the segment, it should be loaded. If the target server has the segment but with
+ * a different partial load profile than the source server, it should also be loaded.
+ *
+ * The two conditions are separate because a null profile does not identify a missing replica: a target that does not
+ * have the segment and a target holding it as a regular full load both project no profile at all.
+ */
+ private boolean shouldLoadSegmentOnTargetServer(
+ DataSegment segment,
+ @Nullable PartialLoadProfile sourceProfile,
+ ServerHolder targetServer,
+ Set targetProjectedSegments
+ )
+ {
+ if (!targetProjectedSegments.contains(segment)) {
+ return true;
+ }
+ final PartialLoadProfile targetProfile = targetServer.getProjectedProfile(segment);
+ return !Objects.equals(fingerprintOf(sourceProfile), fingerprintOf(targetProfile));
+ }
}
diff --git a/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialLoadProfile.java b/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialLoadProfile.java
index 04e1f86f658e..27defc7801c6 100644
--- a/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialLoadProfile.java
+++ b/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialLoadProfile.java
@@ -93,6 +93,18 @@ public static PartialLoadProfile forLoaded(Map wrappedLoadSpec,
return intern(new PartialLoadProfile(wrappedLoadSpec, fingerprint, loadedBytes));
}
+ /**
+ * This profile in request form, for asking another server for the same partial load that produced it: a clone
+ * catching up with its source, or a move handing a replica to its destination. The wrapped load spec and the
+ * fingerprint identify the request and carry over as they are; {@code loadedBytes} is dropped, because a profile read
+ * back off a server carries the footprint that server realized, which belongs to that server's announcement rather
+ * than to a request. Returns {@code this} when the profile is already a request.
+ */
+ public PartialLoadProfile asCloneRequest()
+ {
+ return loadedBytes == null ? this : forRequest(wrappedLoadSpec, fingerprint);
+ }
+
private static PartialLoadProfile intern(PartialLoadProfile profile)
{
return INTERNER.intern(profile);
diff --git a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentLoadQueueManager.java b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentLoadQueueManager.java
index 2db28f2b98a2..346a0a6f8f62 100644
--- a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentLoadQueueManager.java
+++ b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentLoadQueueManager.java
@@ -105,10 +105,15 @@ public boolean dropSegment(DataSegment segment, ServerHolder server)
}
}
+ /**
+ * Moves the segment from serverA to serverB, optionally carrying the partial-load profile that serverA holds it
+ * under so that serverB is asked for the same parts of the segment rather than the whole of it.
+ */
public boolean moveSegment(
DataSegment segment,
ServerHolder serverA,
- ServerHolder serverB
+ ServerHolder serverB,
+ @Nullable PartialLoadProfile profile
)
{
final LoadQueuePeon peonA = serverA.getPeon();
@@ -117,7 +122,7 @@ public boolean moveSegment(
if (!serverA.startOperation(SegmentAction.MOVE_FROM, segment)) {
return false;
}
- if (!serverB.startOperation(SegmentAction.MOVE_TO, segment)) {
+ if (!serverB.startOperation(SegmentAction.MOVE_TO, segment, profile)) {
serverA.cancelOperation(SegmentAction.MOVE_FROM, segment);
return false;
}
@@ -132,6 +137,7 @@ public boolean moveSegment(
peonB.loadSegment(
segment,
SegmentAction.MOVE_TO,
+ profile,
success -> {
// Drop segment only if:
// (1) segment load was successful on serverB
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 f65ed4c46f47..17c4ac500706 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
@@ -180,26 +180,33 @@ public boolean moveSegment(
/**
* Moves the given segment from serverA to serverB.
+ *
+ * If the segment on serverA is a partial load, the partial loadprofile is used to create the request to load
+ * the segment on serverB, ensuring that an equivalent partial load is loaded on serverB.
*/
private boolean moveSegment(DataSegment segment, ServerHolder serverA, ServerHolder serverB)
{
final String tier = serverA.getServer().getTier();
+
+ final PartialLoadProfile profile = serverA.getProjectedProfile(segment);
+ final PartialLoadProfile request = profile == null ? null : profile.asCloneRequest();
+
if (serverA.isLoadingSegment(segment)) {
// Cancel the load on serverA and load on serverB instead
if (serverA.cancelOperation(SegmentAction.LOAD, segment)) {
int loadedCountOnTier = replicaCountMap.get(segment.getId(), tier)
.loadedNotDropping();
if (loadedCountOnTier >= 1) {
- return replicateSegment(segment, serverB, null);
+ return replicateSegment(segment, serverB, request);
} else {
- return loadSegment(segment, serverB, null);
+ return loadSegment(segment, serverB, request);
}
}
// Could not cancel load, let the segment load on serverA and count it as unmoved
return false;
} else if (serverA.isServingSegment(segment)) {
- return loadQueueManager.moveSegment(segment, serverA, serverB);
+ return loadQueueManager.moveSegment(segment, serverA, serverB, request);
} else {
return false;
}
diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
index e6c6c5854516..87be5378926d 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
+++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
@@ -553,6 +553,93 @@ void testRuleSwapRewritesInfoFileOnDisk() throws Exception
);
}
+ @Test
+ void testFullLoadRequestReleasesRuleAndRewritesInfoFile() throws Exception
+ {
+ // An unwrapped load request for a segment held under a rule is the coordinator asking for the whole segment again.
+ // The rule's holds have to come off, so the pinned parts become evictable like any other virtual-storage full load,
+ // and the info file has to stop describing the rule so a restart doesn't reinstate it.
+ manager = makeManager(true, true);
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+ Assertions.assertEquals(FINGERPRINT, manager.getRuleFingerprintForSegment(SEGMENT_ID));
+
+ manager.load(plainSegment());
+
+ Assertions.assertNull(
+ manager.getRuleFingerprintForSegment(SEGMENT_ID),
+ "a full load request must release the applied rule"
+ );
+ final File infoFile = new File(new File(cacheRoot, "info_dir"), SEGMENT_ID.toString());
+ final DataSegment onDisk = jsonMapper.readValue(infoFile, DataSegment.class);
+ Assertions.assertFalse(
+ PartialLoadSpec.detectPartialLoadSpec(onDisk.getLoadSpec()),
+ "info file on disk must no longer carry a partial-load wrapper"
+ );
+ }
+
+ @Test
+ void testFullLoadRequestFailsWhenTheInfoFileCannotBeRewritten() throws Exception
+ {
+ // The release has to reach disk or not happen at all: an unwrapped request announces as a full load either way, so
+ // a rule released only in memory would leave the coordinator recording a replica with no profile while a restart
+ // reinstates the rule. Failing the load instead sends the historical down its drop path, which cleans both up.
+ manager = makeManager(true, true);
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+
+ final File infoDir = new File(cacheRoot, "info_dir");
+ Assertions.assertTrue(infoDir.setReadOnly(), "test setup must be able to make the info dir read-only");
+ try {
+ Assertions.assertThrows(
+ SegmentLoadingException.class,
+ () -> manager.load(plainSegment())
+ );
+ Assertions.assertEquals(
+ FINGERPRINT,
+ manager.getRuleFingerprintForSegment(SEGMENT_ID),
+ "the rule must still be applied when the release could not be persisted"
+ );
+ }
+ finally {
+ Assertions.assertTrue(infoDir.setWritable(true), "test teardown must restore write permission");
+ }
+ }
+
+ @Test
+ void testRestartAfterFullLoadRequestDoesNotReinstateRule() throws Exception
+ {
+ // The released rule has to stay released across a restart: bootstrap reads the rewritten info file, so it restores
+ // the partial layout that is still on disk without reapplying the rule, and announces no profile for it.
+ manager = makeManager(true, true);
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+ manager.load(plainSegment());
+ manager.shutdown();
+ manager = null;
+
+ final SegmentLocalCacheManager restarted = makeManager(true, true);
+ try {
+ final DataSegment cached = restarted.getCachedSegments()
+ .stream()
+ .filter(s -> s.getId().equals(SEGMENT_ID))
+ .findFirst()
+ .orElse(null);
+ Assertions.assertNotNull(cached, "restarted historical must rediscover the segment via its info file");
+
+ final DataSegment bootstrapped = restarted.bootstrap(cached, SegmentLazyLoadFailCallback.NOOP);
+
+ Assertions.assertNull(
+ restarted.getRuleFingerprintForSegment(SEGMENT_ID),
+ "bootstrap must not reapply a rule that a full load request released"
+ );
+ Assertions.assertFalse(
+ bootstrapped instanceof DataSegmentAndLoadProfile,
+ "bootstrap must not announce a partial-load profile for a released rule"
+ );
+ }
+ finally {
+ restarted.shutdown();
+ }
+ }
+
@Test
void testDropClearsRule() throws Exception
{
@@ -840,6 +927,19 @@ private DataSegment compositeWrapperSegment(List projectionPerMember, St
.build();
}
+ /**
+ * The same segment as {@link #partialWrapperSegment}, but with the plain deep-storage load spec the coordinator sends
+ * for a regular full load.
+ */
+ private DataSegment plainSegment()
+ {
+ return DataSegment.builder(SEGMENT_ID)
+ .shardSpec(NoneShardSpec.instance())
+ .loadSpec(Map.of("type", "local", "path", DEEP_STORAGE_DIR.getAbsolutePath()))
+ .size(0)
+ .build();
+ }
+
/**
* A wrapper whose inner LoadSpec resolves via {@code LocalLoadSpec} against a directory that holds no V10 file, so
* {@code openRangeReader()} returns {@code null}. Simulates the "backend doesn't support range reads" case.
diff --git a/server/src/test/java/org/apache/druid/server/coordinator/duty/CloneHistoricalsTest.java b/server/src/test/java/org/apache/druid/server/coordinator/duty/CloneHistoricalsTest.java
new file mode 100644
index 000000000000..9904a319eb57
--- /dev/null
+++ b/server/src/test/java/org/apache/druid/server/coordinator/duty/CloneHistoricalsTest.java
@@ -0,0 +1,332 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.server.coordinator.duty;
+
+import org.apache.druid.client.DruidServer;
+import org.apache.druid.java.util.common.DateTimes;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.segment.TestDataSource;
+import org.apache.druid.server.coordination.ServerType;
+import org.apache.druid.server.coordinator.CloneStatusManager;
+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.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.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.partition.NumberedShardSpec;
+import org.joda.time.Duration;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import javax.annotation.Nullable;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Verifies that {@link CloneHistoricals} reproduces the source historical's load state on its clone target, including
+ * the {@link PartialLoadProfile} a partial-load rule resolved to. A clone that receives the segment without the
+ * profile loads the whole segment and announces the full segment size, so the two servers diverge in both on-disk
+ * footprint and reported {@code curr_size}. The duty therefore compares replicas by partial-load fingerprint rather
+ * than by segment id, and threads the source's wrapped load spec into the clone's load request.
+ */
+public class CloneHistoricalsTest
+{
+ private static final String TIER = "tier1";
+ private static final String SOURCE_HOST = "source_host:8083";
+ private static final String TARGET_HOST = "target_host:8083";
+
+ private static final String FP_REVENUE = "v1:deadbeefcafebabe";
+ private static final String FP_USERS = "v1:0123456789abcdef";
+
+ private static final long SEGMENT_SIZE = 1000L;
+ private static final long REALIZED_BYTES = 250L;
+
+ private SegmentLoadQueueManager loadQueueManager;
+ private CloneHistoricals duty;
+
+ @BeforeEach
+ public void setUp()
+ {
+ loadQueueManager = new SegmentLoadQueueManager(null, null);
+ duty = new CloneHistoricals(loadQueueManager, new CloneStatusManager());
+ }
+
+ @Test
+ public void testCloneLoadsSegmentWithTheSourcePartialLoadProfile()
+ {
+ final DataSegment segment = createSegment();
+ final ServerHolder source = createServer(SOURCE_HOST, segment, loadedProfile(FP_REVENUE, "revenue"));
+ final ServerHolder target = createServer(TARGET_HOST);
+
+ runDuty(source, target, segment);
+
+ final PartialLoadProfile queued = peonOf(target).getProfileFor(segment);
+ Assertions.assertNotNull(queued, "Clone must be asked to load the same parts as its source");
+ Assertions.assertEquals(FP_REVENUE, queued.fingerprint());
+ Assertions.assertEquals(loadedProfile(FP_REVENUE, "revenue").wrappedLoadSpec(), queued.wrappedLoadSpec());
+ Assertions.assertNull(queued.loadedBytes(), "Outbound request profile must not carry loadedBytes");
+ }
+
+ @Test
+ public void testCloneLoadsSegmentWithTheProfileOfAnInFlightSourceLoad()
+ {
+ // The source's own load is still queued, so the profile lives only on the peon's in-flight holder. The clone
+ // must follow the state the source is heading towards, not the state it is in.
+ final DataSegment segment = createSegment();
+ final PartialLoadProfile inFlight = requestProfile(FP_REVENUE, "revenue");
+
+ final TestLoadQueuePeon sourcePeon = new TestLoadQueuePeon();
+ sourcePeon.addInFlightHolder(
+ new SegmentHolder(segment, SegmentAction.LOAD, inFlight, Duration.standardSeconds(10), null)
+ );
+ final ServerHolder source = new ServerHolder(createDruidServer(SOURCE_HOST).toImmutableDruidServer(), sourcePeon);
+ final ServerHolder target = createServer(TARGET_HOST);
+
+ runDuty(source, target, segment);
+
+ final PartialLoadProfile queued = peonOf(target).getProfileFor(segment);
+ Assertions.assertNotNull(queued, "Clone must follow an in-flight partial load on the source");
+ Assertions.assertEquals(FP_REVENUE, queued.fingerprint());
+ }
+
+ @Test
+ public void testCloneReloadsWhenItsFingerprintDiffersFromTheSource()
+ {
+ // Both servers hold the segment, but under different rules. The segment ids are identical, so only the
+ // fingerprint distinguishes them.
+ final DataSegment segment = createSegment();
+ final ServerHolder source = createServer(SOURCE_HOST, segment, loadedProfile(FP_USERS, "users"));
+ final ServerHolder target = createServer(TARGET_HOST, segment, loadedProfile(FP_REVENUE, "revenue"));
+
+ runDuty(source, target, segment);
+
+ final PartialLoadProfile queued = peonOf(target).getProfileFor(segment);
+ Assertions.assertNotNull(queued, "Clone holding a different set of parts must be re-loaded");
+ Assertions.assertEquals(FP_USERS, queued.fingerprint());
+ }
+
+ @Test
+ public void testCloneReloadsAsFullLoadWhenSourceNoLongerLoadsPartially()
+ {
+ // Source moved off the partial-load rule and now holds the whole segment; the clone must follow it back. The
+ // request goes out with no profile even though the clone is already serving the segment: the historical releases
+ // the partial-load rule it holds the replica under when it receives an unwrapped load request.
+ final DataSegment segment = createSegment();
+ final ServerHolder source = createServer(SOURCE_HOST, segment, null);
+ final ServerHolder target = createServer(TARGET_HOST, segment, loadedProfile(FP_REVENUE, "revenue"));
+
+ runDuty(source, target, segment);
+
+ Assertions.assertTrue(
+ peonOf(target).getSegmentsToLoad().contains(segment),
+ "Clone must be re-loaded when the source stops loading partially"
+ );
+ Assertions.assertNull(
+ peonOf(target).getProfileFor(segment),
+ "A full-load source must not thread a profile to the clone"
+ );
+ Assertions.assertTrue(peonOf(target).getSegmentsToDrop().isEmpty());
+ }
+
+ @Test
+ public void testCloneWithAPartialLoadStillQueuedIsConvertedOnALaterRun()
+ {
+ // A partial load can be queued on top of a replica the clone already serves under a different profile, which is how
+ // the historical is asked to fill in the missing parts in place. A segment with an operation already queued cannot
+ // take another one, so the queued load is left to complete and the next run converts the replica it produces.
+ final DataSegment segment = createSegment();
+ final ServerHolder source = createServer(SOURCE_HOST, segment, null);
+
+ final TestLoadQueuePeon targetPeon = new TestLoadQueuePeon();
+ targetPeon.addInFlightHolder(new SegmentHolder(
+ segment,
+ SegmentAction.LOAD,
+ requestProfile(FP_USERS, "users"),
+ Duration.standardSeconds(10),
+ null
+ ));
+ final DruidServer targetDruidServer = createDruidServer(TARGET_HOST);
+ targetDruidServer.addDataSegment(segment, loadedProfile(FP_REVENUE, "revenue"));
+ final ServerHolder target = new ServerHolder(targetDruidServer.toImmutableDruidServer(), targetPeon);
+
+ runDuty(source, target, segment);
+
+ Assertions.assertEquals(
+ requestProfile(FP_USERS, "users"),
+ targetPeon.getProfileFor(segment),
+ "The queued partial load must be left alone"
+ );
+ Assertions.assertTrue(targetPeon.getSegmentsToDrop().isEmpty());
+ }
+
+ @Test
+ public void testFullLoadSourceQueuesPlainLoadOnClone()
+ {
+ final DataSegment segment = createSegment();
+ final ServerHolder source = createServer(SOURCE_HOST, segment, null);
+ final ServerHolder target = createServer(TARGET_HOST);
+
+ runDuty(source, target, segment);
+
+ Assertions.assertTrue(peonOf(target).getSegmentsToLoad().contains(segment));
+ Assertions.assertNull(peonOf(target).getProfileFor(segment));
+ }
+
+ @Test
+ public void testNothingIsQueuedWhenCloneFingerprintMatchesTheSource()
+ {
+ final DataSegment segment = createSegment();
+ final ServerHolder source = createServer(SOURCE_HOST, segment, loadedProfile(FP_REVENUE, "revenue"));
+ final ServerHolder target = createServer(TARGET_HOST, segment, loadedProfile(FP_REVENUE, "revenue"));
+
+ runDuty(source, target, segment);
+
+ Assertions.assertTrue(peonOf(target).getSegmentsToLoad().isEmpty());
+ Assertions.assertTrue(peonOf(target).getSegmentsToDrop().isEmpty());
+ }
+
+ @Test
+ public void testNothingIsQueuedWhenCloneFellBackToAFullDownloadOfTheSameRequest()
+ {
+ // A clone whose historical cannot honour partial downloads announces the requested fingerprint with the full
+ // segment size as its footprint. The request was satisfied, so the duty must leave it alone rather than
+ // re-queueing the load on every run.
+ final DataSegment segment = createSegment();
+ final ServerHolder source = createServer(SOURCE_HOST, segment, loadedProfile(FP_REVENUE, "revenue"));
+ final ServerHolder target = createServer(
+ TARGET_HOST,
+ segment,
+ PartialLoadProfile.forLoaded(wrappedLoadSpec(FP_REVENUE, "revenue"), FP_REVENUE, SEGMENT_SIZE)
+ );
+
+ runDuty(source, target, segment);
+
+ Assertions.assertTrue(peonOf(target).getSegmentsToLoad().isEmpty());
+ Assertions.assertTrue(peonOf(target).getSegmentsToDrop().isEmpty());
+ }
+
+ @Test
+ public void testSegmentMissingFromSourceIsDroppedFromClone()
+ {
+ final DataSegment segment = createSegment();
+ final ServerHolder source = createServer(SOURCE_HOST);
+ final ServerHolder target = createServer(TARGET_HOST, segment, loadedProfile(FP_REVENUE, "revenue"));
+
+ runDuty(source, target, segment);
+
+ Assertions.assertTrue(peonOf(target).getSegmentsToDrop().contains(segment));
+ Assertions.assertTrue(peonOf(target).getSegmentsToLoad().isEmpty());
+ }
+
+ private void runDuty(ServerHolder source, ServerHolder target, DataSegment... usedSegments)
+ {
+ final DruidCluster cluster = DruidCluster.builder().addTier(TIER, source, target).build();
+ final DruidCoordinatorRuntimeParams params =
+ DruidCoordinatorRuntimeParams
+ .builder()
+ .withDruidCluster(cluster)
+ .withUsedSegments(usedSegments)
+ .withDynamicConfigs(
+ CoordinatorDynamicConfig.builder()
+ .withCloneServers(Map.of(TARGET_HOST, SOURCE_HOST))
+ .build()
+ )
+ .build();
+
+ duty.run(params);
+ }
+
+ private static TestLoadQueuePeon peonOf(ServerHolder server)
+ {
+ return (TestLoadQueuePeon) server.getPeon();
+ }
+
+ private static DruidServer createDruidServer(String host)
+ {
+ return new DruidServer(host, host, null, 10L << 30, null, ServerType.HISTORICAL, TIER, 0);
+ }
+
+ /**
+ * Creates a server holder that serves each of the given segments. A non-null profile announces the segment as a
+ * partial load with that profile; a null profile announces it as a regular full load.
+ */
+ private static ServerHolder createServer(String host, DataSegment segment, @Nullable PartialLoadProfile profile)
+ {
+ final DruidServer server = createDruidServer(host);
+ server.addDataSegment(segment, profile);
+ return new ServerHolder(server.toImmutableDruidServer(), new TestLoadQueuePeon());
+ }
+
+ private static ServerHolder createServer(String host)
+ {
+ return new ServerHolder(createDruidServer(host).toImmutableDruidServer(), new TestLoadQueuePeon());
+ }
+
+ private static DataSegment createSegment()
+ {
+ return DataSegment
+ .builder(
+ SegmentId.of(
+ TestDataSource.WIKI,
+ Intervals.of("2024/2025"),
+ DateTimes.nowUtc().toString(),
+ new NumberedShardSpec(0, 0)
+ )
+ )
+ .loadSpec(Map.of("type", "local", "path", "/var/druid/segments/foo"))
+ .projections(List.of("revenue", "users"))
+ .size(SEGMENT_SIZE)
+ .build();
+ }
+
+ private static Map wrappedLoadSpec(String fingerprint, String projection)
+ {
+ return Map.of(
+ "type", "partialProjection",
+ "delegate", Map.of("type", "local", "path", "/var/druid/segments/foo"),
+ "projections", List.of(projection),
+ "fingerprint", fingerprint
+ );
+ }
+
+ /**
+ * The profile shape a historical announces after completing a partial load: the request it was given, plus the
+ * footprint it actually materialized.
+ */
+ private static PartialLoadProfile loadedProfile(String fingerprint, String projection)
+ {
+ return PartialLoadProfile.forLoaded(wrappedLoadSpec(fingerprint, projection), fingerprint, REALIZED_BYTES);
+ }
+
+ /**
+ * The profile shape the coordinator sends out with a load request: no footprint is known yet.
+ */
+ private static PartialLoadProfile requestProfile(String fingerprint, String projection)
+ {
+ return PartialLoadProfile.forRequest(wrappedLoadSpec(fingerprint, projection), fingerprint);
+ }
+}
diff --git a/server/src/test/java/org/apache/druid/server/coordinator/loading/PartialLoadProfileTest.java b/server/src/test/java/org/apache/druid/server/coordinator/loading/PartialLoadProfileTest.java
index ce0f42b7c078..85ea1ae284be 100644
--- a/server/src/test/java/org/apache/druid/server/coordinator/loading/PartialLoadProfileTest.java
+++ b/server/src/test/java/org/apache/druid/server/coordinator/loading/PartialLoadProfileTest.java
@@ -115,6 +115,27 @@ public void testDefensiveCopyOfWrappedLoadSpec()
Assertions.assertFalse(profile.wrappedLoadSpec().containsKey("extra"));
}
+ @Test
+ public void testAsCloneRequestDropsTheAnnouncedFootprint()
+ {
+ // The realized footprint belongs to the announcement of the server that loaded the segment, not to the request the
+ // clone target or move destination is about to get. Everything that identifies the request carries over as-is.
+ final PartialLoadProfile loaded = PartialLoadProfile.forLoaded(WRAPPED, FINGERPRINT, 12345L);
+
+ final PartialLoadProfile request = loaded.asCloneRequest();
+
+ Assertions.assertNull(request.loadedBytes(), "a request carries no realized footprint");
+ Assertions.assertEquals(WRAPPED, request.wrappedLoadSpec());
+ Assertions.assertEquals(FINGERPRINT, request.fingerprint());
+ }
+
+ @Test
+ public void testAsCloneRequestOfARequestIsItself()
+ {
+ final PartialLoadProfile request = PartialLoadProfile.forRequest(WRAPPED, FINGERPRINT);
+ Assertions.assertSame(request, request.asCloneRequest());
+ }
+
@Test
public void testEquals()
{
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 e1cc1fb0d39c..c5ef0280cf3c 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
@@ -404,6 +404,73 @@ public void testForeverPartialLoadRuleEndToEndFullLoadFallback()
);
}
+ @Test
+ public void testMoveOfPartialReplicaCarriesProfileToDestination()
+ {
+ // Balancing a partial replica must move the same parts of the segment. Without the profile the destination
+ // downloads the whole segment, and the reconciler then has to replace it with a partial replica on a later run.
+ final DataSegment segment = createSegment();
+ final ServerHolder source = createDecommissioningServerWithLoaded(TIER1, segment, profileForRevenue());
+ final ServerHolder destination = createServer(TIER1);
+ final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, source, destination).build();
+
+ final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, segment);
+ final boolean moved = params.getSegmentAssigner().moveSegment(segment, source, List.of(destination));
+
+ Assert.assertTrue(moved);
+ Assert.assertEquals(SegmentAction.MOVE_TO, destination.getActionOnSegment(segment));
+ final PartialLoadProfile queued = ((TestLoadQueuePeon) destination.getPeon()).getProfileFor(segment);
+ Assert.assertNotNull("Move destination should be asked for the same parts the source holds", queued);
+ Assert.assertEquals(FP_REVENUE, queued.fingerprint());
+ }
+
+ @Test
+ public void testMoveOfInFlightPartialLoadCarriesProfileToDestination()
+ {
+ // The source's load has not completed yet, so the move cancels it and loads on the destination instead. The
+ // profile lives on the peon's in-flight holder, and cancelling the operation clears it from the source, so it
+ // has to be read before the cancellation.
+ final DataSegment segment = createSegment();
+ final TestLoadQueuePeon sourcePeon = new TestLoadQueuePeon();
+ sourcePeon.addInFlightHolder(new SegmentHolder(
+ segment,
+ SegmentAction.LOAD,
+ profileForRevenue(),
+ org.joda.time.Duration.standardSeconds(10),
+ null
+ ));
+ final ServerHolder source = new ServerHolder(createDruidServer(TIER1).toImmutableDruidServer(), sourcePeon, true);
+ final ServerHolder destination = createServer(TIER1);
+ final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, source, destination).build();
+
+ final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, segment);
+ final boolean moved = params.getSegmentAssigner().moveSegment(segment, source, List.of(destination));
+
+ Assert.assertTrue(moved);
+ final PartialLoadProfile queued = ((TestLoadQueuePeon) destination.getPeon()).getProfileFor(segment);
+ Assert.assertNotNull("Cancelled in-flight partial load should be reissued to the destination", queued);
+ Assert.assertEquals(FP_REVENUE, queued.fingerprint());
+ }
+
+ @Test
+ public void testMoveOfFullLoadReplicaCarriesNoProfile()
+ {
+ final DataSegment segment = createSegment();
+ final ServerHolder source = createDecommissioningServerWithLoaded(TIER1, segment, null);
+ final ServerHolder destination = createServer(TIER1);
+ final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, source, destination).build();
+
+ final DruidCoordinatorRuntimeParams params = makeRuntimeParams(cluster, segment);
+ final boolean moved = params.getSegmentAssigner().moveSegment(segment, source, List.of(destination));
+
+ Assert.assertTrue(moved);
+ Assert.assertEquals(SegmentAction.MOVE_TO, destination.getActionOnSegment(segment));
+ Assert.assertNull(
+ "Moving a regular full-load replica must not thread a profile",
+ ((TestLoadQueuePeon) destination.getPeon()).getProfileFor(segment)
+ );
+ }
+
private DruidCoordinatorRuntimeParams makeRuntimeParams(DruidCluster cluster, DataSegment... segments)
{
return DruidCoordinatorRuntimeParams
@@ -453,6 +520,22 @@ private ServerHolder createDecommissioningServer(String tier)
return new ServerHolder(createDruidServer(tier).toImmutableDruidServer(), new TestLoadQueuePeon(), true);
}
+ /**
+ * Creates a decommissioning server that already serves the given segment, announced with {@code profile} when it is
+ * non-null. Decommissioning keeps the server out of its own move-destination candidates, so the move has exactly
+ * one place to go.
+ */
+ private ServerHolder createDecommissioningServerWithLoaded(
+ String tier,
+ DataSegment segment,
+ @Nullable PartialLoadProfile profile
+ )
+ {
+ final DruidServer server = createDruidServer(tier);
+ server.addDataSegment(segment, profile);
+ return new ServerHolder(server.toImmutableDruidServer(), new TestLoadQueuePeon(), true);
+ }
+
private static DataSegment createSegment()
{
return DataSegment