From adf1332d7a4879a8fd627c065d2e67c95cfabd84 Mon Sep 17 00:00:00 2001 From: Lucas Capistrant Date: Thu, 30 Jul 2026 23:38:07 -0500 Subject: [PATCH 1/7] fix bug in historical clones when source and/or target enable partial loads --- .../PartialLoadHistoricalCloningTest.java | 329 ++++++++++++++++++ .../coordinator/duty/CloneHistoricals.java | 69 +++- .../duty/CloneHistoricalsTest.java | 298 ++++++++++++++++ 3 files changed, 692 insertions(+), 4 deletions(-) create mode 100644 embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/PartialLoadHistoricalCloningTest.java create mode 100644 server/src/test/java/org/apache/druid/server/coordinator/duty/CloneHistoricalsTest.java 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/server/coordinator/duty/CloneHistoricals.java b/server/src/main/java/org/apache/druid/server/coordinator/duty/CloneHistoricals.java index c12193c9b0e5..e7f809c5a5be 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,12 @@ * 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. 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 +111,17 @@ 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 = projectedProfile(sourceServer, segment); + if (!targetProjectedSegments.contains(segment) + || !Objects.equals( + fingerprintOf(sourceProfile), + fingerprintOf(projectedProfile(targetServer, segment)) + )) { + loadSegmentOnTargetServer(segment, sourceProfile, targetServer, params); } } @@ -124,8 +139,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 +161,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, + toRequestProfile(sourceProfile) + )) { params.getCoordinatorStats().add( Stats.Segments.ASSIGNED_TO_CLONE, rowKey.build(), @@ -150,6 +175,42 @@ private void loadSegmentOnTargetServer( } } + /** + * The {@link PartialLoadProfile} that {@code server} is expected to hold {@code segment} under once its queued + * operations finish: the profile of an in-flight load if one is queued (a null profile there means a regular full + * load is on its way), else the profile the server announced for the loaded replica. Returns null when the replica + * is a regular full load. Mirrors the branch order in + * {@link org.apache.druid.server.coordinator.loading.PartialSegmentStatusInTier}, which classifies rule-managed + * replicas the same way. + */ + @Nullable + private static PartialLoadProfile projectedProfile(ServerHolder server, DataSegment segment) + { + final SegmentAction action = server.getActionOnSegment(segment); + if (action != null && action.isLoad()) { + return server.getInFlightProfile(segment); + } + return server.getServer().getPartialLoadProfile(segment.getId()); + } + + @Nullable + private static String fingerprintOf(@Nullable PartialLoadProfile profile) + { + return profile == null ? null : profile.fingerprint(); + } + + /** + * Rebuilds a profile read off a server as an outbound load request. A profile announced by a historical carries + * the footprint that historical realized, which is its own to report and not part of the request. + */ + @Nullable + private static PartialLoadProfile toRequestProfile(@Nullable PartialLoadProfile profile) + { + return profile == null + ? null + : PartialLoadProfile.forRequest(profile.wrappedLoadSpec(), profile.fingerprint()); + } + private void dropSegmentFromTargetServer( DataSegment segment, ServerHolder targetServer, 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..1b6d53c94e44 --- /dev/null +++ b/server/src/test/java/org/apache/druid/server/coordinator/duty/CloneHistoricalsTest.java @@ -0,0 +1,298 @@ +/* + * 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. + 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" + ); + } + + @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); + } +} From 9fcf1d04b3569836912c0bcf4e59f41e7aaf7952 Mon Sep 17 00:00:00 2001 From: Lucas Capistrant Date: Fri, 31 Jul 2026 00:03:17 -0500 Subject: [PATCH 2/7] fix MOVE_TO as well --- .../server/coordinator/ServerHolder.java | 18 ++++ .../coordinator/duty/CloneHistoricals.java | 36 +------- .../loading/PartialLoadProfile.java | 11 +++ .../loading/SegmentLoadQueueManager.java | 10 ++- .../loading/StrategicSegmentAssigner.java | 12 ++- .../StrategicSegmentAssignerPartialTest.java | 83 +++++++++++++++++++ 6 files changed, 132 insertions(+), 38 deletions(-) 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 e7f809c5a5be..2d71affddad6 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 @@ -115,11 +115,11 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params) // 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) { - final PartialLoadProfile sourceProfile = projectedProfile(sourceServer, segment); + final PartialLoadProfile sourceProfile = sourceServer.getProjectedProfile(segment); if (!targetProjectedSegments.contains(segment) || !Objects.equals( fingerprintOf(sourceProfile), - fingerprintOf(projectedProfile(targetServer, segment)) + fingerprintOf(targetServer.getProjectedProfile(segment)) )) { loadSegmentOnTargetServer(segment, sourceProfile, targetServer, params); } @@ -165,7 +165,7 @@ private void loadSegmentOnTargetServer( loadableSegment, targetServer, SegmentAction.LOAD, - toRequestProfile(sourceProfile) + sourceProfile == null ? null : sourceProfile.asRequest() )) { params.getCoordinatorStats().add( Stats.Segments.ASSIGNED_TO_CLONE, @@ -175,42 +175,12 @@ private void loadSegmentOnTargetServer( } } - /** - * The {@link PartialLoadProfile} that {@code server} is expected to hold {@code segment} under once its queued - * operations finish: the profile of an in-flight load if one is queued (a null profile there means a regular full - * load is on its way), else the profile the server announced for the loaded replica. Returns null when the replica - * is a regular full load. Mirrors the branch order in - * {@link org.apache.druid.server.coordinator.loading.PartialSegmentStatusInTier}, which classifies rule-managed - * replicas the same way. - */ - @Nullable - private static PartialLoadProfile projectedProfile(ServerHolder server, DataSegment segment) - { - final SegmentAction action = server.getActionOnSegment(segment); - if (action != null && action.isLoad()) { - return server.getInFlightProfile(segment); - } - return server.getServer().getPartialLoadProfile(segment.getId()); - } - @Nullable private static String fingerprintOf(@Nullable PartialLoadProfile profile) { return profile == null ? null : profile.fingerprint(); } - /** - * Rebuilds a profile read off a server as an outbound load request. A profile announced by a historical carries - * the footprint that historical realized, which is its own to report and not part of the request. - */ - @Nullable - private static PartialLoadProfile toRequestProfile(@Nullable PartialLoadProfile profile) - { - return profile == null - ? null - : PartialLoadProfile.forRequest(profile.wrappedLoadSpec(), profile.fingerprint()); - } - private void dropSegmentFromTargetServer( DataSegment segment, ServerHolder targetServer, 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..117374252f83 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,17 @@ public static PartialLoadProfile forLoaded(Map wrappedLoadSpec, return intern(new PartialLoadProfile(wrappedLoadSpec, fingerprint, loadedBytes)); } + /** + * This profile in request form, for reissuing to another server the same partial load that produced it. A profile + * read back off a server carries the footprint that server realized, which belongs to that server's announcement + * and not to a request; the wrapped load spec and fingerprint are what identify the request. Returns {@code this} + * when the profile is already a request. + */ + public PartialLoadProfile asRequest() + { + 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..2120cb71b2ed 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 @@ -184,22 +184,28 @@ public boolean moveSegment( private boolean moveSegment(DataSegment segment, ServerHolder serverA, ServerHolder serverB) { final String tier = serverA.getServer().getTier(); + + // A replica loaded under a partial-load rule holds only part of the segment, so the destination has to be asked + // for the same parts. Read the profile up front: cancelling the load below clears serverA's in-flight profile. + final PartialLoadProfile profile = serverA.getProjectedProfile(segment); + final PartialLoadProfile request = profile == null ? null : profile.asRequest(); + 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/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 From 5925fc99e410c3985d45e62daee1d3c8a5ad76c4 Mon Sep 17 00:00:00 2001 From: Lucas Capistrant Date: Fri, 31 Jul 2026 09:45:52 -0500 Subject: [PATCH 3/7] Address p2 from review - avoid stale load profile application --- .../loading/PartialClusterGroupLoadSpec.java | 4 +- .../segment/loading/PartialLoadSpec.java | 16 ++++- .../loading/PartialProjectionLoadSpec.java | 4 +- .../loading/SegmentLocalCacheManager.java | 2 +- .../SegmentChangeRequestLoad.java | 10 +++- .../coordinator/duty/CloneHistoricals.java | 2 +- .../loading/PartialLoadProfile.java | 35 +++++++++-- .../loading/StrategicSegmentAssigner.java | 4 +- .../duty/CloneHistoricalsTest.java | 20 +++++++ .../loading/PartialLoadProfileTest.java | 58 ++++++++++++++++++- .../StrategicSegmentAssignerPartialTest.java | 22 +++++++ 11 files changed, 158 insertions(+), 19 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpec.java b/processing/src/main/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpec.java index cbc27ad3cf99..46178db5c053 100644 --- a/processing/src/main/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpec.java +++ b/processing/src/main/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpec.java @@ -63,9 +63,9 @@ public static Map wireForm( { return Map.of( "type", TYPE, - "delegate", delegate, + DELEGATE_FIELD, delegate, "clusterGroupIndices", clusterGroupIndices, - "fingerprint", fingerprint + FINGERPRINT_FIELD, fingerprint ); } diff --git a/processing/src/main/java/org/apache/druid/segment/loading/PartialLoadSpec.java b/processing/src/main/java/org/apache/druid/segment/loading/PartialLoadSpec.java index c9b0e82acdd2..c6adabf4bd83 100644 --- a/processing/src/main/java/org/apache/druid/segment/loading/PartialLoadSpec.java +++ b/processing/src/main/java/org/apache/druid/segment/loading/PartialLoadSpec.java @@ -64,6 +64,18 @@ public abstract class PartialLoadSpec implements LoadSpec */ public static final String TYPE_PREFIX = "partial"; + /** + * Wire-form field holding the inner load spec that says where the segment data lives. Named here so code that + * builds or inspects raw {@link Map}-form wrappers agrees with {@link #getDelegate()}. + */ + public static final String DELEGATE_FIELD = "delegate"; + + /** + * Wire-form field holding the fingerprint of the request that produced the wrapper. Named here so code that builds + * or inspects raw {@link Map}-form wrappers agrees with {@link #getFingerprint()}. + */ + public static final String FINGERPRINT_FIELD = "fingerprint"; + /** * Returns {@code true} if {@code loadSpec} matches the shape of the {@link PartialLoadSpec} subtype. * Convention-based detection (no subtype allowlist): the {@code type} field must be a {@link String} starting with @@ -76,8 +88,8 @@ public static boolean detectPartialLoadSpec(@Nullable Map loadSp return loadSpec != null && loadSpec.get("type") instanceof String typeString && typeString.startsWith(TYPE_PREFIX) - && loadSpec.get("fingerprint") instanceof String - && loadSpec.get("delegate") instanceof Map; + && loadSpec.get(FINGERPRINT_FIELD) instanceof String + && loadSpec.get(DELEGATE_FIELD) instanceof Map; } /** diff --git a/processing/src/main/java/org/apache/druid/segment/loading/PartialProjectionLoadSpec.java b/processing/src/main/java/org/apache/druid/segment/loading/PartialProjectionLoadSpec.java index c063f33db8e1..4a88cbd42784 100644 --- a/processing/src/main/java/org/apache/druid/segment/loading/PartialProjectionLoadSpec.java +++ b/processing/src/main/java/org/apache/druid/segment/loading/PartialProjectionLoadSpec.java @@ -60,9 +60,9 @@ public static Map wireForm( { return Map.of( "type", TYPE, - "delegate", delegate, + DELEGATE_FIELD, delegate, "projections", projections, - "fingerprint", fingerprint + FINGERPRINT_FIELD, fingerprint ); } 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..3ec62cd31ccf 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 @@ -1426,7 +1426,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() ); } 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/duty/CloneHistoricals.java b/server/src/main/java/org/apache/druid/server/coordinator/duty/CloneHistoricals.java index 2d71affddad6..c2e15209e8b0 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 @@ -165,7 +165,7 @@ private void loadSegmentOnTargetServer( loadableSegment, targetServer, SegmentAction.LOAD, - sourceProfile == null ? null : sourceProfile.asRequest() + sourceProfile == null ? null : sourceProfile.asRequestFor(loadableSegment) )) { params.getCoordinatorStats().add( Stats.Segments.ASSIGNED_TO_CLONE, 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 117374252f83..3c4bcd16a33b 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 @@ -22,8 +22,11 @@ import com.google.common.collect.Interner; import com.google.common.collect.Interners; import org.apache.druid.error.InvalidInput; +import org.apache.druid.segment.loading.PartialLoadSpec; +import org.apache.druid.timeline.DataSegment; import javax.annotation.Nullable; +import java.util.HashMap; import java.util.Map; import java.util.Objects; @@ -94,14 +97,34 @@ public static PartialLoadProfile forLoaded(Map wrappedLoadSpec, } /** - * This profile in request form, for reissuing to another server the same partial load that produced it. A profile - * read back off a server carries the footprint that server realized, which belongs to that server's announcement - * and not to a request; the wrapped load spec and fingerprint are what identify the request. Returns {@code this} - * when the profile is already a request. + * This profile in request form for {@code segment}, for reissuing to another server the same partial load that + * produced it (clone catch-up, balancer move). Two things are normalized: + *

+ * {@code segment} must be the current metadata view of the segment, which is what the coordinator's data-sources + * snapshot hands back. The wrapper's delegate is left alone when there is nothing better to point it at: a segment + * carrying no load spec, or one whose load spec is already a partial-load wrapper (an outbound request segment + * rather than the metadata view, which would otherwise nest one wrapper inside another). */ - public PartialLoadProfile asRequest() + public PartialLoadProfile asRequestFor(DataSegment segment) { - return loadedBytes == null ? this : forRequest(wrappedLoadSpec, fingerprint); + final Map currentLoadSpec = segment.getLoadSpec(); + if (currentLoadSpec == null + || currentLoadSpec.isEmpty() + || PartialLoadSpec.hasPartialTypePrefix(currentLoadSpec) + || currentLoadSpec.equals(wrappedLoadSpec.get(PartialLoadSpec.DELEGATE_FIELD))) { + return loadedBytes == null ? this : forRequest(wrappedLoadSpec, fingerprint); + } + final Map rebased = new HashMap<>(wrappedLoadSpec); + rebased.put(PartialLoadSpec.DELEGATE_FIELD, currentLoadSpec); + return forRequest(rebased, fingerprint); } private static PartialLoadProfile intern(PartialLoadProfile profile) 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 2120cb71b2ed..a7628f5b574c 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 @@ -187,8 +187,10 @@ private boolean moveSegment(DataSegment segment, ServerHolder serverA, ServerHol // A replica loaded under a partial-load rule holds only part of the segment, so the destination has to be asked // for the same parts. Read the profile up front: cancelling the load below clears serverA's in-flight profile. + // `segment` is the metadata-resolved segment (see TierSegmentBalancer.getLoadableSegment), which is what + // asRequestFor needs to rebase the request onto the segment's current location. final PartialLoadProfile profile = serverA.getProjectedProfile(segment); - final PartialLoadProfile request = profile == null ? null : profile.asRequest(); + final PartialLoadProfile request = profile == null ? null : profile.asRequestFor(segment); if (serverA.isLoadingSegment(segment)) { // Cancel the load on serverA and load on serverB instead 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 index 1b6d53c94e44..8d52a9965d1c 100644 --- 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 @@ -195,6 +195,26 @@ public void testNothingIsQueuedWhenCloneFellBackToAFullDownloadOfTheSameRequest( Assertions.assertTrue(peonOf(target).getSegmentsToDrop().isEmpty()); } + @Test + public void testCloneLoadIsRebasedOntoTheSegmentsCurrentLoadSpec() + { + // The source announced its profile when the segment lived somewhere else, and the wrapped load spec replaces the + // outbound segment's load spec wholesale. The clone must be pointed at where the segment lives now, otherwise it + // can never catch up once the old object is gone. + final DataSegment segment = createSegment(); + final DataSegment relocated = segment.withLoadSpec(Map.of("type", "local", "path", "/mnt/relocated/foo")); + final ServerHolder source = createServer(SOURCE_HOST, segment, loadedProfile(FP_REVENUE, "revenue")); + final ServerHolder target = createServer(TARGET_HOST); + + runDuty(source, target, relocated); + + final PartialLoadProfile queued = peonOf(target).getProfileFor(relocated); + Assertions.assertNotNull(queued); + Assertions.assertEquals(relocated.getLoadSpec(), queued.wrappedLoadSpec().get("delegate")); + Assertions.assertEquals(FP_REVENUE, queued.fingerprint()); + Assertions.assertEquals(List.of("revenue"), queued.wrappedLoadSpec().get("projections")); + } + @Test public void testSegmentMissingFromSourceIsDroppedFromClone() { 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..560a6957ccbe 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 @@ -23,10 +23,15 @@ import nl.jqno.equalsverifier.EqualsVerifier; import org.apache.druid.error.DruidException; import org.apache.druid.error.DruidExceptionMatcher; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.timeline.partition.NumberedShardSpec; import org.hamcrest.MatcherAssert; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import javax.annotation.Nullable; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -35,9 +40,15 @@ public class PartialLoadProfileTest { private static final String FINGERPRINT = "v1:0123456789abcdef"; + private static final Map ORIGINAL_DELEGATE = + ImmutableMap.of("type", "local", "path", "/var/druid/segments/foo"); + + private static final Map MIGRATED_DELEGATE = + ImmutableMap.of("type", "local", "path", "/mnt/relocated/segments/foo"); + private static final Map WRAPPED = ImmutableMap.of( "type", "partialProjection", - "delegate", ImmutableMap.of("type", "local", "path", "/var/druid/segments/foo"), + "delegate", ORIGINAL_DELEGATE, "projections", List.of("user_daily", "user_hourly"), "fingerprint", FINGERPRINT ); @@ -115,6 +126,42 @@ public void testDefensiveCopyOfWrappedLoadSpec() Assertions.assertFalse(profile.wrappedLoadSpec().containsKey("extra")); } + @Test + public void testAsRequestForRebasesDelegateOntoTheSegmentsCurrentLoadSpec() + { + // A profile read off a server carries the location the segment had when that server was asked to load. Reissuing + // it must point at where the segment lives now, keeping the selection and fingerprint that identify the request. + final PartialLoadProfile loaded = PartialLoadProfile.forLoaded(WRAPPED, FINGERPRINT, 12345L); + + final PartialLoadProfile request = loaded.asRequestFor(segmentWithLoadSpec(MIGRATED_DELEGATE)); + + Assertions.assertEquals(MIGRATED_DELEGATE, request.wrappedLoadSpec().get("delegate")); + Assertions.assertEquals(FINGERPRINT, request.fingerprint()); + Assertions.assertEquals(List.of("user_daily", "user_hourly"), request.wrappedLoadSpec().get("projections")); + Assertions.assertEquals("partialProjection", request.wrappedLoadSpec().get("type")); + Assertions.assertNull(request.loadedBytes(), "a request carries no realized footprint"); + } + + @Test + public void testAsRequestForKeepsDelegateWhenSegmentIsUnmoved() + { + final PartialLoadProfile request = PartialLoadProfile.forRequest(WRAPPED, FINGERPRINT); + Assertions.assertSame(request, request.asRequestFor(segmentWithLoadSpec(ORIGINAL_DELEGATE))); + } + + @Test + public void testAsRequestForKeepsDelegateWhenSegmentHasNoLoadSpec() + { + // Nothing better to point the wrapper at, so the existing delegate rides through rather than being replaced by an + // empty one the historical would reject. + final PartialLoadProfile loaded = PartialLoadProfile.forLoaded(WRAPPED, FINGERPRINT, 12345L); + + final PartialLoadProfile request = loaded.asRequestFor(segmentWithLoadSpec(null)); + + Assertions.assertEquals(ORIGINAL_DELEGATE, request.wrappedLoadSpec().get("delegate")); + Assertions.assertNull(request.loadedBytes()); + } + @Test public void testEquals() { @@ -144,4 +191,13 @@ public void testInterningSharesReferenceForEquivalentProfiles() PartialLoadProfile pd = PartialLoadProfile.forLoaded(WRAPPED, "v1:differentfingerprint", 12345L); Assertions.assertNotSame(pa, pd); } + + private static DataSegment segmentWithLoadSpec(@Nullable Map loadSpec) + { + return DataSegment.builder(SegmentId.of("wiki", Intervals.of("2025/2026"), "v1", 0)) + .shardSpec(new NumberedShardSpec(0, 1)) + .loadSpec(loadSpec) + .size(100L) + .build(); + } } 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..98c47e6bcbb3 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 @@ -452,6 +452,28 @@ public void testMoveOfInFlightPartialLoadCarriesProfileToDestination() Assert.assertEquals(FP_REVENUE, queued.fingerprint()); } + @Test + public void testMoveOfPartialReplicaIsRebasedOntoTheSegmentsCurrentLoadSpec() + { + // The source announced its profile when the segment lived somewhere else, and the wrapped load spec replaces the + // outbound segment's load spec wholesale. The destination must be pointed at where the segment lives now, + // otherwise the move can never complete once the old object is gone. + final DataSegment segment = createSegment(); + final DataSegment relocated = segment.withLoadSpec(Map.of("type", "local", "path", "/mnt/relocated/foo")); + 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, relocated); + Assert.assertTrue(params.getSegmentAssigner().moveSegment(relocated, source, List.of(destination))); + + final PartialLoadProfile queued = ((TestLoadQueuePeon) destination.getPeon()).getProfileFor(relocated); + Assert.assertNotNull(queued); + Assert.assertEquals(relocated.getLoadSpec(), queued.wrappedLoadSpec().get("delegate")); + Assert.assertEquals(FP_REVENUE, queued.fingerprint()); + Assert.assertEquals(List.of("revenue"), queued.wrappedLoadSpec().get("projections")); + } + @Test public void testMoveOfFullLoadReplicaCarriesNoProfile() { From 52260093177469dfc50fd4c056280c53ac687cf1 Mon Sep 17 00:00:00 2001 From: Lucas Capistrant Date: Fri, 31 Jul 2026 09:46:47 -0500 Subject: [PATCH 4/7] Address P1 from review. if a clone has to full load something it has partial loaded. drop and let next cycle load it --- .../coordinator/duty/CloneHistoricals.java | 50 ++++++++++++++++--- .../duty/CloneHistoricalsTest.java | 43 +++++++++++++--- 2 files changed, 80 insertions(+), 13 deletions(-) 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 c2e15209e8b0..834ea42ac241 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 @@ -51,8 +51,9 @@ * 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. Clone targets are excluded from rule-driven assignment - * ({@link DruidCluster#getManagedHistoricals()}), so this duty is the only thing that can correct them. + * the source's profile. The one transition that is not a re-load is a clone going from partial back to full, which + * takes a drop first (see {@link #convertCloneReplicaToFullLoad}). 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 { @@ -116,11 +117,17 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params) // different parts of it. for (DataSegment segment : sourceProjectedSegments) { final PartialLoadProfile sourceProfile = sourceServer.getProjectedProfile(segment); - if (!targetProjectedSegments.contains(segment) - || !Objects.equals( - fingerprintOf(sourceProfile), - fingerprintOf(targetServer.getProjectedProfile(segment)) - )) { + if (!targetProjectedSegments.contains(segment)) { + loadSegmentOnTargetServer(segment, sourceProfile, targetServer, params); + continue; + } + final PartialLoadProfile targetProfile = targetServer.getProjectedProfile(segment); + if (Objects.equals(fingerprintOf(sourceProfile), fingerprintOf(targetProfile))) { + continue; + } + if (sourceProfile == null) { + convertCloneReplicaToFullLoad(segment, targetServer, params); + } else { loadSegmentOnTargetServer(segment, sourceProfile, targetServer, params); } } @@ -175,6 +182,35 @@ private void loadSegmentOnTargetServer( } } + /** + * Returns the clone target to a full load of {@code segment}, for when the source has stopped holding it partially. + *

+ * A plain load request on top of the existing replica does not achieve this. A historical that receives an unwrapped + * load request for a segment it already holds under a partial-load rule keeps that rule applied: its holds go on + * pinning the parts the rule selected, and the segment's info file goes on describing a partial load, which the + * historical reapplies and re-announces on its next restart. Dropping the replica does release the rule and retire + * the info file, so the next coordinator run sees a clone that is missing the segment and queues the ordinary full + * load. + *

+ * A partial load that is still queued is cancelled and replaced by the full load within this run, since nothing has + * been applied on the historical yet. If the request has already gone out, the cancel fails and that load runs to + * completion; the drop path then converts the replica on a later run. + */ + private void convertCloneReplicaToFullLoad( + DataSegment segment, + ServerHolder targetServer, + DruidCoordinatorRuntimeParams params + ) + { + if (targetServer.isLoadingSegment(segment)) { + if (targetServer.cancelOperation(SegmentAction.LOAD, segment)) { + loadSegmentOnTargetServer(segment, null, targetServer, params); + } + } else { + dropSegmentFromTargetServer(segment, targetServer, params); + } + } + @Nullable private static String fingerprintOf(@Nullable PartialLoadProfile profile) { 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 index 8d52a9965d1c..8bb2b5538dfa 100644 --- 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 @@ -130,9 +130,12 @@ public void testCloneReloadsWhenItsFingerprintDiffersFromTheSource() } @Test - public void testCloneReloadsAsFullLoadWhenSourceNoLongerLoadsPartially() + public void testClonePartialReplicaIsDroppedWhenSourceNoLongerLoadsPartially() { - // Source moved off the partial-load rule and now holds the whole segment; the clone must follow it back. + // Source moved off the partial-load rule and now holds the whole segment, so the clone has to follow it back to a + // full load. That takes a drop rather than a load on top: a historical asked to load a segment it already holds + // under a partial-load rule keeps the rule, its holds and its info file, so the clone would report a full replica + // while still holding only the rule's parts. The full load follows on the next run, once the replica is gone. final DataSegment segment = createSegment(); final ServerHolder source = createServer(SOURCE_HOST, segment, null); final ServerHolder target = createServer(TARGET_HOST, segment, loadedProfile(FP_REVENUE, "revenue")); @@ -140,13 +143,41 @@ public void testCloneReloadsAsFullLoadWhenSourceNoLongerLoadsPartially() runDuty(source, target, segment); Assertions.assertTrue( - peonOf(target).getSegmentsToLoad().contains(segment), - "Clone must be re-loaded when the source stops loading partially" + peonOf(target).getSegmentsToDrop().contains(segment), + "Clone holding a partial replica must be dropped when the source stops loading partially" ); + Assertions.assertTrue( + peonOf(target).getSegmentsToLoad().isEmpty(), + "The full load has to wait for the drop to release the historical's partial-load rule" + ); + } + + @Test + public void testCloneWithInFlightPartialLoadIsSwitchedToAFullLoadInTheSameRun() + { + // Nothing has been applied on the historical while the partial load is still queued, so there is no rule to + // release: cancel that load and queue the full one right away instead of waiting for a drop. + 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_REVENUE, "revenue"), + Duration.standardSeconds(10), + null + )); + final ServerHolder target = new ServerHolder(createDruidServer(TARGET_HOST).toImmutableDruidServer(), targetPeon); + + runDuty(source, target, segment); + + Assertions.assertTrue(targetPeon.getSegmentsToLoad().contains(segment)); Assertions.assertNull( - peonOf(target).getProfileFor(segment), - "A full-load source must not thread a profile to the clone" + targetPeon.getProfileFor(segment), + "The cancelled partial load must be replaced by a plain full load" ); + Assertions.assertTrue(targetPeon.getSegmentsToDrop().isEmpty()); } @Test From a5fafe4a1da61b42e567a8093db80234be8cff66 Mon Sep 17 00:00:00 2001 From: Lucas Capistrant Date: Mon, 3 Aug 2026 10:15:00 -0500 Subject: [PATCH 5/7] Attempt to fixup the workaround for partial to full transition --- .../coordinator/duty/CloneHistoricals.java | 18 ++++++---- .../duty/CloneHistoricalsTest.java | 33 +++++++++++++++++++ 2 files changed, 44 insertions(+), 7 deletions(-) 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 834ea42ac241..985f76efe0ee 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 @@ -192,9 +192,12 @@ private void loadSegmentOnTargetServer( * the info file, so the next coordinator run sees a clone that is missing the segment and queues the ordinary full * load. *

- * A partial load that is still queued is cancelled and replaced by the full load within this run, since nothing has - * been applied on the historical yet. If the request has already gone out, the cancel fails and that load runs to - * completion; the drop path then converts the replica on a later run. + * A queued partial load is cancelled first, and then what the target is actually serving decides the rest. + * Cancelling a load says nothing about that: a partial load can be queued on top of a replica the target already + * serves under a different profile, which is how the historical is asked to fill in missing parts in place. So a + * served replica still has to be dropped, and only a target that serves nothing can take the full load right away. + * When the cancel fails because the request has already gone to the historical, that load runs to completion and a + * later run converts the replica it produces. */ private void convertCloneReplicaToFullLoad( DataSegment segment, @@ -203,11 +206,12 @@ private void convertCloneReplicaToFullLoad( ) { if (targetServer.isLoadingSegment(segment)) { - if (targetServer.cancelOperation(SegmentAction.LOAD, segment)) { - loadSegmentOnTargetServer(segment, null, targetServer, params); - } - } else { + targetServer.cancelOperation(SegmentAction.LOAD, segment); + } + if (targetServer.isServingSegment(segment)) { dropSegmentFromTargetServer(segment, targetServer, params); + } else { + loadSegmentOnTargetServer(segment, null, targetServer, params); } } 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 index 8bb2b5538dfa..ec194d0e0df8 100644 --- 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 @@ -180,6 +180,39 @@ public void testCloneWithInFlightPartialLoadIsSwitchedToAFullLoadInTheSameRun() Assertions.assertTrue(targetPeon.getSegmentsToDrop().isEmpty()); } + @Test + public void testCloneServingAPartialReplicaIsDroppedEvenWithAReloadQueued() + { + // 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. Cancelling that queued load leaves the served + // replica and its rule behind, so the conversion still has to go through a drop. + 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.assertTrue( + targetPeon.getSegmentsToDrop().contains(segment), + "The served partial replica must be dropped, not loaded over" + ); + Assertions.assertTrue( + targetPeon.getSegmentsToLoad().isEmpty(), + "The queued reload must be cancelled and no full load queued while the replica is still served" + ); + } + @Test public void testFullLoadSourceQueuesPlainLoadOnClone() { From 624cdaa6c5fdc1eea92980071d532f6c485cd269 Mon Sep 17 00:00:00 2001 From: Lucas Capistrant Date: Mon, 3 Aug 2026 10:16:03 -0500 Subject: [PATCH 6/7] Fix the underlying blocker for clones to be able to cleanly go from partial to full load of a segment --- .../loading/SegmentLocalCacheManager.java | 71 +++++++++++++++--- .../coordinator/duty/CloneHistoricals.java | 55 ++------------ ...tLocalCacheManagerPartialRuleLoadTest.java | 73 +++++++++++++++++++ .../duty/CloneHistoricalsTest.java | 66 +++++------------ 4 files changed, 159 insertions(+), 106 deletions(-) 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 3ec62cd31ccf..aa69013f1dcb 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,23 @@ 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. An unwrapped request for a segment currently held under a + // partial-load rule is the coordinator asking for the whole segment again, so release the rule as well. + final boolean isFullLoadRequest = !PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec()); 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 (isFullLoadRequest + && cacheEntry instanceof PartialSegmentMetadataCacheEntry partial + && partial.isRuleHeld()) { + releaseRuleForFullLoad(dataSegment, partial); } } } @@ -1526,6 +1534,47 @@ 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 rule is cleared before the info file is rewritten because clearing cannot fail, so the in-memory state and the + * load announcement come out right either way. A failed rewrite leaves the info file describing the released rule, + * which a restart reapplies and re-announces until the coordinator's next load request converts the segment 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) + { + // 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(); + 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() + ); + try { + rewriteInfoFile(dataSegment); + } + catch (IOException e) { + log.warn( + e, + "Failed to rewrite info file for segment[%s] after releasing partial-load rule[fingerprint=%s]. The rule is " + + "released on this historical, but the info file still describes it, so a restart will reapply and " + + "re-announce it until the coordinator's next load request releases it again.", + dataSegment.getId(), + priorFingerprint + ); + } + } + /** * 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/coordinator/duty/CloneHistoricals.java b/server/src/main/java/org/apache/druid/server/coordinator/duty/CloneHistoricals.java index 985f76efe0ee..5b200ca11cbd 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 @@ -51,9 +51,9 @@ * 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. The one transition that is not a re-load is a clone going from partial back to full, which - * takes a drop first (see {@link #convertCloneReplicaToFullLoad}). Clone targets are excluded from rule-driven - * assignment ({@link DruidCluster#getManagedHistoricals()}), so this duty is the only thing that can correct them. + * 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 { @@ -117,17 +117,11 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params) // different parts of it. for (DataSegment segment : sourceProjectedSegments) { final PartialLoadProfile sourceProfile = sourceServer.getProjectedProfile(segment); - if (!targetProjectedSegments.contains(segment)) { - loadSegmentOnTargetServer(segment, sourceProfile, targetServer, params); - continue; - } - final PartialLoadProfile targetProfile = targetServer.getProjectedProfile(segment); - if (Objects.equals(fingerprintOf(sourceProfile), fingerprintOf(targetProfile))) { - continue; - } - if (sourceProfile == null) { - convertCloneReplicaToFullLoad(segment, targetServer, params); - } else { + if (!targetProjectedSegments.contains(segment) + || !Objects.equals( + fingerprintOf(sourceProfile), + fingerprintOf(targetServer.getProjectedProfile(segment)) + )) { loadSegmentOnTargetServer(segment, sourceProfile, targetServer, params); } } @@ -182,39 +176,6 @@ private void loadSegmentOnTargetServer( } } - /** - * Returns the clone target to a full load of {@code segment}, for when the source has stopped holding it partially. - *

- * A plain load request on top of the existing replica does not achieve this. A historical that receives an unwrapped - * load request for a segment it already holds under a partial-load rule keeps that rule applied: its holds go on - * pinning the parts the rule selected, and the segment's info file goes on describing a partial load, which the - * historical reapplies and re-announces on its next restart. Dropping the replica does release the rule and retire - * the info file, so the next coordinator run sees a clone that is missing the segment and queues the ordinary full - * load. - *

- * A queued partial load is cancelled first, and then what the target is actually serving decides the rest. - * Cancelling a load says nothing about that: a partial load can be queued on top of a replica the target already - * serves under a different profile, which is how the historical is asked to fill in missing parts in place. So a - * served replica still has to be dropped, and only a target that serves nothing can take the full load right away. - * When the cancel fails because the request has already gone to the historical, that load runs to completion and a - * later run converts the replica it produces. - */ - private void convertCloneReplicaToFullLoad( - DataSegment segment, - ServerHolder targetServer, - DruidCoordinatorRuntimeParams params - ) - { - if (targetServer.isLoadingSegment(segment)) { - targetServer.cancelOperation(SegmentAction.LOAD, segment); - } - if (targetServer.isServingSegment(segment)) { - dropSegmentFromTargetServer(segment, targetServer, params); - } else { - loadSegmentOnTargetServer(segment, null, targetServer, params); - } - } - @Nullable private static String fingerprintOf(@Nullable PartialLoadProfile profile) { 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..6fa1fdd76075 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,66 @@ 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 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 +900,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 index ec194d0e0df8..d2d4a30f00c3 100644 --- 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 @@ -130,12 +130,11 @@ public void testCloneReloadsWhenItsFingerprintDiffersFromTheSource() } @Test - public void testClonePartialReplicaIsDroppedWhenSourceNoLongerLoadsPartially() + public void testCloneReloadsAsFullLoadWhenSourceNoLongerLoadsPartially() { - // Source moved off the partial-load rule and now holds the whole segment, so the clone has to follow it back to a - // full load. That takes a drop rather than a load on top: a historical asked to load a segment it already holds - // under a partial-load rule keeps the rule, its holds and its info file, so the clone would report a full replica - // while still holding only the rule's parts. The full load follows on the next run, once the replica is gone. + // 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")); @@ -143,49 +142,22 @@ public void testClonePartialReplicaIsDroppedWhenSourceNoLongerLoadsPartially() runDuty(source, target, segment); Assertions.assertTrue( - peonOf(target).getSegmentsToDrop().contains(segment), - "Clone holding a partial replica must be dropped when the source stops loading partially" + peonOf(target).getSegmentsToLoad().contains(segment), + "Clone must be re-loaded when the source stops loading partially" ); - Assertions.assertTrue( - peonOf(target).getSegmentsToLoad().isEmpty(), - "The full load has to wait for the drop to release the historical's partial-load rule" - ); - } - - @Test - public void testCloneWithInFlightPartialLoadIsSwitchedToAFullLoadInTheSameRun() - { - // Nothing has been applied on the historical while the partial load is still queued, so there is no rule to - // release: cancel that load and queue the full one right away instead of waiting for a drop. - 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_REVENUE, "revenue"), - Duration.standardSeconds(10), - null - )); - final ServerHolder target = new ServerHolder(createDruidServer(TARGET_HOST).toImmutableDruidServer(), targetPeon); - - runDuty(source, target, segment); - - Assertions.assertTrue(targetPeon.getSegmentsToLoad().contains(segment)); Assertions.assertNull( - targetPeon.getProfileFor(segment), - "The cancelled partial load must be replaced by a plain full load" + peonOf(target).getProfileFor(segment), + "A full-load source must not thread a profile to the clone" ); - Assertions.assertTrue(targetPeon.getSegmentsToDrop().isEmpty()); + Assertions.assertTrue(peonOf(target).getSegmentsToDrop().isEmpty()); } @Test - public void testCloneServingAPartialReplicaIsDroppedEvenWithAReloadQueued() + 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. Cancelling that queued load leaves the served - // replica and its rule behind, so the conversion still has to go through a drop. + // 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); @@ -203,14 +175,12 @@ public void testCloneServingAPartialReplicaIsDroppedEvenWithAReloadQueued() runDuty(source, target, segment); - Assertions.assertTrue( - targetPeon.getSegmentsToDrop().contains(segment), - "The served partial replica must be dropped, not loaded over" - ); - Assertions.assertTrue( - targetPeon.getSegmentsToLoad().isEmpty(), - "The queued reload must be cancelled and no full load queued while the replica is still served" + Assertions.assertEquals( + requestProfile(FP_USERS, "users"), + targetPeon.getProfileFor(segment), + "The queued partial load must be left alone" ); + Assertions.assertTrue(targetPeon.getSegmentsToDrop().isEmpty()); } @Test From a7a7e58344d4345cfe158be839191ed550a4f32f Mon Sep 17 00:00:00 2001 From: Lucas Capistrant Date: Mon, 3 Aug 2026 14:28:01 -0500 Subject: [PATCH 7/7] Improvements, cleanups, and fixes based on review comments --- .../loading/SegmentLocalCacheManager.java | 39 +++++++------- .../coordinator/duty/CloneHistoricals.java | 31 ++++++++--- .../loading/PartialLoadProfile.java | 36 +++---------- .../loading/StrategicSegmentAssigner.java | 9 ++-- ...tLocalCacheManagerPartialRuleLoadTest.java | 27 ++++++++++ .../duty/CloneHistoricalsTest.java | 20 ------- .../loading/PartialLoadProfileTest.java | 53 ++++--------------- .../StrategicSegmentAssignerPartialTest.java | 22 -------- 8 files changed, 91 insertions(+), 146 deletions(-) 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 aa69013f1dcb..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 @@ -1338,9 +1338,9 @@ 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. An unwrapped request for a segment currently held under a - // partial-load rule is the coordinator asking for the whole segment again, so release the rule as well. - final boolean isFullLoadRequest = !PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec()); + // 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 { @@ -1351,9 +1351,7 @@ public DataSegment load(final DataSegment dataSegment) throws SegmentLoadingExce continue; } cacheEntry.setOnUnmount(null); - if (isFullLoadRequest - && cacheEntry instanceof PartialSegmentMetadataCacheEntry partial - && partial.isRuleHeld()) { + if (cacheEntry instanceof PartialSegmentMetadataCacheEntry partial && partial.isRuleHeld()) { releaseRuleForFullLoad(dataSegment, partial); } } @@ -1540,39 +1538,40 @@ public void drop(final DataSegment segment) * 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 rule is cleared before the info file is rewritten because clearing cannot fail, so the in-memory state and the - * load announcement come out right either way. A failed rewrite leaves the info file describing the released rule, - * which a restart reapplies and re-announces until the coordinator's next load request converts the segment again. + * 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(); - 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() - ); try { rewriteInfoFile(dataSegment); } catch (IOException e) { - log.warn( + throw new SegmentLoadingException( e, - "Failed to rewrite info file for segment[%s] after releasing partial-load rule[fingerprint=%s]. The rule is " - + "released on this historical, but the info file still describes it, so a restart will reapply and " - + "re-announce it until the coordinator's next load request releases it again.", + "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() + ); } /** 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 5b200ca11cbd..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 @@ -117,11 +117,7 @@ public DruidCoordinatorRuntimeParams run(DruidCoordinatorRuntimeParams params) // different parts of it. for (DataSegment segment : sourceProjectedSegments) { final PartialLoadProfile sourceProfile = sourceServer.getProjectedProfile(segment); - if (!targetProjectedSegments.contains(segment) - || !Objects.equals( - fingerprintOf(sourceProfile), - fingerprintOf(targetServer.getProjectedProfile(segment)) - )) { + if (shouldLoadSegmentOnTargetServer(segment, sourceProfile, targetServer, targetProjectedSegments)) { loadSegmentOnTargetServer(segment, sourceProfile, targetServer, params); } } @@ -166,7 +162,7 @@ private void loadSegmentOnTargetServer( loadableSegment, targetServer, SegmentAction.LOAD, - sourceProfile == null ? null : sourceProfile.asRequestFor(loadableSegment) + sourceProfile == null ? null : sourceProfile.asCloneRequest() )) { params.getCoordinatorStats().add( Stats.Segments.ASSIGNED_TO_CLONE, @@ -266,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 3c4bcd16a33b..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 @@ -22,11 +22,8 @@ import com.google.common.collect.Interner; import com.google.common.collect.Interners; import org.apache.druid.error.InvalidInput; -import org.apache.druid.segment.loading.PartialLoadSpec; -import org.apache.druid.timeline.DataSegment; import javax.annotation.Nullable; -import java.util.HashMap; import java.util.Map; import java.util.Objects; @@ -97,34 +94,15 @@ public static PartialLoadProfile forLoaded(Map wrappedLoadSpec, } /** - * This profile in request form for {@code segment}, for reissuing to another server the same partial load that - * produced it (clone catch-up, balancer move). Two things are normalized: - *

    - *
  • {@code loadedBytes} is dropped. A profile read back off a server carries the footprint that server - * realized, which belongs to that server's announcement and not to a request.
    • - *
  • The wrapper's {@link PartialLoadSpec#DELEGATE_FIELD} is replaced with {@code segment}'s load spec. The - * wrapper was built when the source server was asked to load, so it carries whatever deep-storage location - * the segment had then; if the payload has since been corrected or migrated, that location may no longer - * exist. The scheme-specific selection and the fingerprint are what identify the request and are preserved, - * so the reissued load still reconciles against the same rule.
    • - *
- * {@code segment} must be the current metadata view of the segment, which is what the coordinator's data-sources - * snapshot hands back. The wrapper's delegate is left alone when there is nothing better to point it at: a segment - * carrying no load spec, or one whose load spec is already a partial-load wrapper (an outbound request segment - * rather than the metadata view, which would otherwise nest one wrapper inside another). + * 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 asRequestFor(DataSegment segment) + public PartialLoadProfile asCloneRequest() { - final Map currentLoadSpec = segment.getLoadSpec(); - if (currentLoadSpec == null - || currentLoadSpec.isEmpty() - || PartialLoadSpec.hasPartialTypePrefix(currentLoadSpec) - || currentLoadSpec.equals(wrappedLoadSpec.get(PartialLoadSpec.DELEGATE_FIELD))) { - return loadedBytes == null ? this : forRequest(wrappedLoadSpec, fingerprint); - } - final Map rebased = new HashMap<>(wrappedLoadSpec); - rebased.put(PartialLoadSpec.DELEGATE_FIELD, currentLoadSpec); - return forRequest(rebased, fingerprint); + return loadedBytes == null ? this : forRequest(wrappedLoadSpec, fingerprint); } private static PartialLoadProfile intern(PartialLoadProfile profile) 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 a7628f5b574c..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,17 +180,16 @@ 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(); - // A replica loaded under a partial-load rule holds only part of the segment, so the destination has to be asked - // for the same parts. Read the profile up front: cancelling the load below clears serverA's in-flight profile. - // `segment` is the metadata-resolved segment (see TierSegmentBalancer.getLoadableSegment), which is what - // asRequestFor needs to rebase the request onto the segment's current location. final PartialLoadProfile profile = serverA.getProjectedProfile(segment); - final PartialLoadProfile request = profile == null ? null : profile.asRequestFor(segment); + final PartialLoadProfile request = profile == null ? null : profile.asCloneRequest(); if (serverA.isLoadingSegment(segment)) { // Cancel the load on serverA and load on serverB instead 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 6fa1fdd76075..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 @@ -577,6 +577,33 @@ void testFullLoadRequestReleasesRuleAndRewritesInfoFile() throws Exception ); } + @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 { 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 index d2d4a30f00c3..9904a319eb57 100644 --- 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 @@ -229,26 +229,6 @@ public void testNothingIsQueuedWhenCloneFellBackToAFullDownloadOfTheSameRequest( Assertions.assertTrue(peonOf(target).getSegmentsToDrop().isEmpty()); } - @Test - public void testCloneLoadIsRebasedOntoTheSegmentsCurrentLoadSpec() - { - // The source announced its profile when the segment lived somewhere else, and the wrapped load spec replaces the - // outbound segment's load spec wholesale. The clone must be pointed at where the segment lives now, otherwise it - // can never catch up once the old object is gone. - final DataSegment segment = createSegment(); - final DataSegment relocated = segment.withLoadSpec(Map.of("type", "local", "path", "/mnt/relocated/foo")); - final ServerHolder source = createServer(SOURCE_HOST, segment, loadedProfile(FP_REVENUE, "revenue")); - final ServerHolder target = createServer(TARGET_HOST); - - runDuty(source, target, relocated); - - final PartialLoadProfile queued = peonOf(target).getProfileFor(relocated); - Assertions.assertNotNull(queued); - Assertions.assertEquals(relocated.getLoadSpec(), queued.wrappedLoadSpec().get("delegate")); - Assertions.assertEquals(FP_REVENUE, queued.fingerprint()); - Assertions.assertEquals(List.of("revenue"), queued.wrappedLoadSpec().get("projections")); - } - @Test public void testSegmentMissingFromSourceIsDroppedFromClone() { 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 560a6957ccbe..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 @@ -23,15 +23,10 @@ import nl.jqno.equalsverifier.EqualsVerifier; import org.apache.druid.error.DruidException; import org.apache.druid.error.DruidExceptionMatcher; -import org.apache.druid.java.util.common.Intervals; -import org.apache.druid.timeline.DataSegment; -import org.apache.druid.timeline.SegmentId; -import org.apache.druid.timeline.partition.NumberedShardSpec; import org.hamcrest.MatcherAssert; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; -import javax.annotation.Nullable; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -40,15 +35,9 @@ public class PartialLoadProfileTest { private static final String FINGERPRINT = "v1:0123456789abcdef"; - private static final Map ORIGINAL_DELEGATE = - ImmutableMap.of("type", "local", "path", "/var/druid/segments/foo"); - - private static final Map MIGRATED_DELEGATE = - ImmutableMap.of("type", "local", "path", "/mnt/relocated/segments/foo"); - private static final Map WRAPPED = ImmutableMap.of( "type", "partialProjection", - "delegate", ORIGINAL_DELEGATE, + "delegate", ImmutableMap.of("type", "local", "path", "/var/druid/segments/foo"), "projections", List.of("user_daily", "user_hourly"), "fingerprint", FINGERPRINT ); @@ -127,39 +116,24 @@ public void testDefensiveCopyOfWrappedLoadSpec() } @Test - public void testAsRequestForRebasesDelegateOntoTheSegmentsCurrentLoadSpec() + public void testAsCloneRequestDropsTheAnnouncedFootprint() { - // A profile read off a server carries the location the segment had when that server was asked to load. Reissuing - // it must point at where the segment lives now, keeping the selection and fingerprint that identify the request. + // 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.asRequestFor(segmentWithLoadSpec(MIGRATED_DELEGATE)); + final PartialLoadProfile request = loaded.asCloneRequest(); - Assertions.assertEquals(MIGRATED_DELEGATE, request.wrappedLoadSpec().get("delegate")); - Assertions.assertEquals(FINGERPRINT, request.fingerprint()); - Assertions.assertEquals(List.of("user_daily", "user_hourly"), request.wrappedLoadSpec().get("projections")); - Assertions.assertEquals("partialProjection", request.wrappedLoadSpec().get("type")); Assertions.assertNull(request.loadedBytes(), "a request carries no realized footprint"); + Assertions.assertEquals(WRAPPED, request.wrappedLoadSpec()); + Assertions.assertEquals(FINGERPRINT, request.fingerprint()); } @Test - public void testAsRequestForKeepsDelegateWhenSegmentIsUnmoved() + public void testAsCloneRequestOfARequestIsItself() { final PartialLoadProfile request = PartialLoadProfile.forRequest(WRAPPED, FINGERPRINT); - Assertions.assertSame(request, request.asRequestFor(segmentWithLoadSpec(ORIGINAL_DELEGATE))); - } - - @Test - public void testAsRequestForKeepsDelegateWhenSegmentHasNoLoadSpec() - { - // Nothing better to point the wrapper at, so the existing delegate rides through rather than being replaced by an - // empty one the historical would reject. - final PartialLoadProfile loaded = PartialLoadProfile.forLoaded(WRAPPED, FINGERPRINT, 12345L); - - final PartialLoadProfile request = loaded.asRequestFor(segmentWithLoadSpec(null)); - - Assertions.assertEquals(ORIGINAL_DELEGATE, request.wrappedLoadSpec().get("delegate")); - Assertions.assertNull(request.loadedBytes()); + Assertions.assertSame(request, request.asCloneRequest()); } @Test @@ -191,13 +165,4 @@ public void testInterningSharesReferenceForEquivalentProfiles() PartialLoadProfile pd = PartialLoadProfile.forLoaded(WRAPPED, "v1:differentfingerprint", 12345L); Assertions.assertNotSame(pa, pd); } - - private static DataSegment segmentWithLoadSpec(@Nullable Map loadSpec) - { - return DataSegment.builder(SegmentId.of("wiki", Intervals.of("2025/2026"), "v1", 0)) - .shardSpec(new NumberedShardSpec(0, 1)) - .loadSpec(loadSpec) - .size(100L) - .build(); - } } 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 98c47e6bcbb3..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 @@ -452,28 +452,6 @@ public void testMoveOfInFlightPartialLoadCarriesProfileToDestination() Assert.assertEquals(FP_REVENUE, queued.fingerprint()); } - @Test - public void testMoveOfPartialReplicaIsRebasedOntoTheSegmentsCurrentLoadSpec() - { - // The source announced its profile when the segment lived somewhere else, and the wrapped load spec replaces the - // outbound segment's load spec wholesale. The destination must be pointed at where the segment lives now, - // otherwise the move can never complete once the old object is gone. - final DataSegment segment = createSegment(); - final DataSegment relocated = segment.withLoadSpec(Map.of("type", "local", "path", "/mnt/relocated/foo")); - 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, relocated); - Assert.assertTrue(params.getSegmentAssigner().moveSegment(relocated, source, List.of(destination))); - - final PartialLoadProfile queued = ((TestLoadQueuePeon) destination.getPeon()).getProfileFor(relocated); - Assert.assertNotNull(queued); - Assert.assertEquals(relocated.getLoadSpec(), queued.wrappedLoadSpec().get("delegate")); - Assert.assertEquals(FP_REVENUE, queued.fingerprint()); - Assert.assertEquals(List.of("revenue"), queued.wrappedLoadSpec().get("projections")); - } - @Test public void testMoveOfFullLoadReplicaCarriesNoProfile() {