From 536d3eab1ba9fe65cdf1ccabca2f2a6a2b945c75 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Mon, 6 Jul 2026 11:23:26 -0700 Subject: [PATCH 1/3] feat: wire up partial load rules to segment cache --- .../file/PartialSegmentFileMapperV10.java | 4 +- .../loading/PartialClusterGroupLoadSpec.java | 65 ++ .../segment/loading/PartialLoadSpec.java | 14 + .../loading/PartialProjectionLoadSpec.java | 38 + .../PartialClusterGroupLoadSpecTest.java | 206 +++++ .../PartialProjectionLoadSpecTest.java | 98 +++ .../PartialSegmentBundleCacheEntry.java | 10 + .../loading/PartialSegmentCacheBootstrap.java | 118 ++- .../PartialSegmentMetadataCacheEntry.java | 179 +++- .../loading/SegmentLocalCacheManager.java | 578 +++++++++++- .../PartialSegmentBundleCacheEntryTest.java | 10 + .../PartialSegmentCacheBootstrapTest.java | 228 +++++ .../PartialSegmentMetadataCacheEntryTest.java | 62 +- ...ntLocalCacheManagerPartialAcquireTest.java | 6 +- ...gmentLocalCacheManagerPartialDropTest.java | 3 + ...tLocalCacheManagerPartialRuleLoadTest.java | 829 ++++++++++++++++++ 16 files changed, 2373 insertions(+), 75 deletions(-) create mode 100644 server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java diff --git a/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java b/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java index ea4895a7f2e4..62d1b964b147 100644 --- a/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java +++ b/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java @@ -919,8 +919,10 @@ private static void fetchAndPersistHeader( /** * Parse metadata from the header file. Throws on any parse failure (corrupt JSON, truncated header, etc.). + * Public so external callers (e.g. the partial-load-rule bootstrap path in {@code SegmentLocalCacheManager}) can + * share the header-parse convention with this class rather than re-implementing FileInputStream + reader plumbing. */ - private static SegmentFileMetadataReader.Result parseHeaderFile( + public static SegmentFileMetadataReader.Result parseHeaderFile( File headerFile, ObjectMapper jsonMapper ) throws IOException 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 636de7bc776b..fad2b36d81e4 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 @@ -25,7 +25,16 @@ import com.fasterxml.jackson.annotation.JsonTypeName; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Preconditions; +import org.apache.druid.error.DruidException; +import org.apache.druid.segment.file.SegmentFileMetadata; +import org.apache.druid.segment.projections.ClusteredValueGroupsBaseTableSchema; +import org.apache.druid.segment.projections.ProjectionMetadata; +import org.apache.druid.segment.projections.Projections; +import org.apache.druid.segment.projections.TableClusterGroupSpec; +import org.apache.druid.timeline.ClusterGroupTuples; +import org.apache.druid.timeline.DataSegment; +import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Objects; @@ -83,6 +92,62 @@ public List getClusterGroupIndices() return clusterGroupIndices; } + /** + * Translates each selected {@code clusterGroupIndex} to the bundle name of the corresponding on-disk cluster + * group, by finding the {@link TableClusterGroupSpec} at that position in the parsed base-table schema and + * formatting its {@link TableClusterGroupSpec#getClusteringValueIds()} via + * {@link Projections#getClusterGroupBundleName(List)}. + *

+ * Positional alignment between {@link DataSegment#getClusterGroups()} tuples and the metadata's + * {@link ClusteredValueGroupsBaseTableSchema#getClusterGroups()} list is invariant of the writer: both are populated + * from the same source array at segment-build time. A size-mismatch indicates a writer/reader contract violation + * and trips a defensive throw. + */ + @Override + public List getSelectedBundleNames(DataSegment segment, SegmentFileMetadata metadata) + { + if (clusterGroupIndices.isEmpty()) { + return List.of(); + } + final List projections = metadata.getProjections(); + if (projections == null || projections.isEmpty()) { + throw DruidException.defensive( + "Cannot resolve cluster-group bundle names for segment[%s]: metadata has no projections", + segment.getId() + ); + } + if (!(projections.getFirst().getSchema() instanceof ClusteredValueGroupsBaseTableSchema clusteredSummary)) { + throw DruidException.defensive( + "Cannot resolve cluster-group bundle names for segment[%s]: base projection is not clustered", + segment.getId() + ); + } + final List metadataGroups = clusteredSummary.getClusterGroups(); + final ClusterGroupTuples segmentTuples = segment.getClusterGroups(); + final int segmentTupleCount = segmentTuples == null ? 0 : segmentTuples.tuples().size(); + if (segmentTupleCount != metadataGroups.size()) { + throw DruidException.defensive( + "Cluster-group count mismatch for segment[%s]: DataSegment has [%s] tuples, metadata has [%s] groups", + segment.getId(), + segmentTupleCount, + metadataGroups.size() + ); + } + final List bundleNames = new ArrayList<>(clusterGroupIndices.size()); + for (int idx : clusterGroupIndices) { + if (idx < 0 || idx >= metadataGroups.size()) { + throw DruidException.defensive( + "Cluster-group index [%s] is out of range [0, %s) for segment[%s]", + idx, + metadataGroups.size(), + segment.getId() + ); + } + bundleNames.add(Projections.getClusterGroupBundleName(metadataGroups.get(idx).getClusteringValueIds())); + } + return bundleNames; + } + @Override public boolean equals(Object o) { 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 f1924dd6c4db..96900e36d91d 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 @@ -24,10 +24,13 @@ import com.google.common.base.Preconditions; import com.google.common.base.Supplier; import com.google.common.base.Suppliers; +import org.apache.druid.segment.file.SegmentFileMetadata; +import org.apache.druid.timeline.DataSegment; import javax.annotation.Nullable; import java.io.File; import java.io.IOException; +import java.util.List; import java.util.Map; /** @@ -136,4 +139,15 @@ public SegmentRangeReader openRangeReader() throws IOException { return materializedDelegate().openRangeReader(); } + + /** + * Returns the cache-layer bundle names that this partial-load spec selects for {@code segment}, resolving the + * scheme-specific identifiers carried in the load spec against the segment {@code metadata}. The historical's cache + * layer uses the returned names to eagerly mount and pin the matching bundles as static entries; other bundles + * remain weak and on-demand. + *

+ * Returns an empty list when the load spec carries an empty selection, where only the metadata header is sticky and + * no bundle is pinned. + */ + public abstract List getSelectedBundleNames(DataSegment segment, SegmentFileMetadata metadata); } 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 737a10655898..ebc37ca90c85 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 @@ -25,11 +25,17 @@ import com.fasterxml.jackson.annotation.JsonTypeName; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Preconditions; +import org.apache.druid.error.DruidException; +import org.apache.druid.segment.file.SegmentFileMetadata; +import org.apache.druid.segment.projections.ProjectionMetadata; +import org.apache.druid.timeline.DataSegment; import org.apache.druid.utils.CollectionUtils; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; /** * A {@link PartialLoadSpec} that requests partial loading of a segment's projections. The base class carries the @@ -84,6 +90,38 @@ public List getProjections() return projections; } + /** + * Projection names double as bundle names in the V10 partial-segment layout (each projection's containers are + * tagged with its name as the bundle prefix), so the load spec selection maps to bundle names verbatim after + * validating that each requested name refers to a projection actually present on the segment. + */ + @Override + public List getSelectedBundleNames(DataSegment segment, SegmentFileMetadata metadata) + { + final List segmentProjections = metadata.getProjections(); + if (segmentProjections == null || segmentProjections.isEmpty()) { + throw DruidException.defensive( + "Cannot resolve projection bundles for segment[%s]: metadata has no projections", + segment.getId() + ); + } + final Set known = segmentProjections.stream() + .map(pm -> pm.getSchema().getName()) + .collect(Collectors.toSet()); + for (String projection : projections) { + if (!known.contains(projection)) { + throw DruidException.defensive( + "Segment[%s] does not contain projection[%s]; matcher/reader drift or writer/reader contract violation." + + " Known projections: %s", + segment.getId(), + projection, + known + ); + } + } + return projections; + } + @Override public boolean equals(Object o) { diff --git a/processing/src/test/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpecTest.java b/processing/src/test/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpecTest.java index 28906eba94d9..47528c77b3d3 100644 --- a/processing/src/test/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpecTest.java +++ b/processing/src/test/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpecTest.java @@ -27,7 +27,26 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.module.SimpleModule; import com.google.common.collect.ImmutableMap; +import org.apache.druid.error.DruidException; import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.query.OrderBy; +import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.column.ColumnHolder; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.file.SegmentFileMetadata; +import org.apache.druid.segment.projections.ClusteredValueGroupsBaseTableSchema; +import org.apache.druid.segment.projections.ClusteringDictionaries; +import org.apache.druid.segment.projections.ProjectionMetadata; +import org.apache.druid.segment.projections.ProjectionSchema; +import org.apache.druid.segment.projections.Projections; +import org.apache.druid.segment.projections.TableClusterGroupSpec; +import org.apache.druid.segment.projections.TableProjectionSchema; +import org.apache.druid.timeline.ClusterGroupTuples; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.timeline.partition.NumberedShardSpec; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -163,6 +182,193 @@ void testRejectsNullFingerprint() ); } + @Test + void testGetSelectedBundleNamesResolvesIndicesToBundleNames() + { + // Three cluster groups in the metadata; pick groups 0 and 2. + final SegmentFileMetadata metadata = clusteredMetadata( + List.of( + new TableClusterGroupSpec(List.of(0), 10), + new TableClusterGroupSpec(List.of(1), 20), + new TableClusterGroupSpec(List.of(2), 30) + ) + ); + final DataSegment segment = clusteredSegment( + List.of(List.of("acme"), List.of("globex"), List.of("initech")) + ); + final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec( + DELEGATE, + List.of(0, 2), + FINGERPRINT, + jsonMapper + ); + Assertions.assertEquals( + List.of( + Projections.getClusterGroupBundleName(List.of(0)), + Projections.getClusterGroupBundleName(List.of(2)) + ), + spec.getSelectedBundleNames(segment, metadata) + ); + } + + @Test + void testGetSelectedBundleNamesEmptyIndicesReturnsEmpty() + { + // Sibling-empty case: matcher applied but no positive content. Caller pins only the metadata. + final SegmentFileMetadata metadata = clusteredMetadata( + List.of(new TableClusterGroupSpec(List.of(0), 1)) + ); + final DataSegment segment = clusteredSegment(List.of(List.of("acme"))); + final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec( + DELEGATE, + List.of(), + FINGERPRINT, + jsonMapper + ); + Assertions.assertEquals(List.of(), spec.getSelectedBundleNames(segment, metadata)); + } + + @Test + void testGetSelectedBundleNamesOutOfRangeIndexThrows() + { + final SegmentFileMetadata metadata = clusteredMetadata( + List.of(new TableClusterGroupSpec(List.of(0), 1)) + ); + final DataSegment segment = clusteredSegment(List.of(List.of("acme"))); + final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec( + DELEGATE, + List.of(5), + FINGERPRINT, + jsonMapper + ); + final DruidException thrown = Assertions.assertThrows( + DruidException.class, + () -> spec.getSelectedBundleNames(segment, metadata) + ); + Assertions.assertTrue( + thrown.getMessage().contains("Cluster-group index [5] is out of range"), + "unexpected message: " + thrown.getMessage() + ); + } + + @Test + void testGetSelectedBundleNamesSizeMismatchThrows() + { + // Metadata has 2 groups, segment has 1 tuple — writer/reader contract violation; tripwire fires. + final SegmentFileMetadata metadata = clusteredMetadata( + List.of( + new TableClusterGroupSpec(List.of(0), 1), + new TableClusterGroupSpec(List.of(1), 1) + ) + ); + final DataSegment segment = clusteredSegment(List.of(List.of("acme"))); + final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec( + DELEGATE, + List.of(0), + FINGERPRINT, + jsonMapper + ); + final DruidException thrown = Assertions.assertThrows( + DruidException.class, + () -> spec.getSelectedBundleNames(segment, metadata) + ); + Assertions.assertTrue( + thrown.getMessage().contains("Cluster-group count mismatch"), + "unexpected message: " + thrown.getMessage() + ); + } + + @Test + void testGetSelectedBundleNamesNonClusteredBaseThrows() + { + // Base projection is a regular (non-clustered) table — partial cluster-group load is nonsensical. + final ProjectionSchema baseSchema = new TableProjectionSchema( + VirtualColumns.EMPTY, + List.of(ColumnHolder.TIME_COLUMN_NAME, "tenant"), + null, + List.of(OrderBy.ascending(ColumnHolder.TIME_COLUMN_NAME)) + ); + final SegmentFileMetadata metadata = new SegmentFileMetadata( + List.of(), + Map.of(), + null, + null, + List.of(new ProjectionMetadata(1, baseSchema)), + null + ); + final DataSegment segment = clusteredSegment(List.of(List.of("acme"))); + final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec( + DELEGATE, + List.of(0), + FINGERPRINT, + jsonMapper + ); + final DruidException thrown = Assertions.assertThrows( + DruidException.class, + () -> spec.getSelectedBundleNames(segment, metadata) + ); + Assertions.assertTrue( + thrown.getMessage().contains("base projection is not clustered"), + "unexpected message: " + thrown.getMessage() + ); + } + + @Test + void testGetSelectedBundleNamesNoProjectionsThrows() + { + final SegmentFileMetadata metadata = new SegmentFileMetadata(List.of(), Map.of(), null, null, null, null); + final DataSegment segment = clusteredSegment(List.of(List.of("acme"))); + final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec( + DELEGATE, + List.of(0), + FINGERPRINT, + jsonMapper + ); + final DruidException thrown = Assertions.assertThrows( + DruidException.class, + () -> spec.getSelectedBundleNames(segment, metadata) + ); + Assertions.assertTrue( + thrown.getMessage().contains("metadata has no projections"), + "unexpected message: " + thrown.getMessage() + ); + } + + private static final RowSignature CLUSTERING_TENANT = RowSignature.builder() + .add("tenant", ColumnType.STRING) + .build(); + + private static SegmentFileMetadata clusteredMetadata(List groups) + { + final ClusteredValueGroupsBaseTableSchema baseSchema = new ClusteredValueGroupsBaseTableSchema( + VirtualColumns.EMPTY, + List.of(ColumnHolder.TIME_COLUMN_NAME, "tenant", "metric"), + List.of(OrderBy.ascending("tenant"), OrderBy.ascending(ColumnHolder.TIME_COLUMN_NAME)), + CLUSTERING_TENANT, + null, + new ClusteringDictionaries(List.of("acme", "globex", "initech"), null, null, null), + groups + ); + return new SegmentFileMetadata( + List.of(), + Map.of(), + null, + null, + List.of(new ProjectionMetadata(groups.stream().mapToInt(TableClusterGroupSpec::getNumRows).sum(), baseSchema)), + null + ); + } + + private static DataSegment clusteredSegment(List> tuples) + { + return DataSegment.builder( + SegmentId.of("ds", Intervals.ETERNITY, "v1", new NumberedShardSpec(0, 1)) + ) + .size(0) + .clusterGroups(new ClusterGroupTuples(CLUSTERING_TENANT, tuples)) + .build(); + } + /** * Stub LoadSpec used to verify delegation. Uses the same JSON "type"=="stub" key as the test {@link #DELEGATE}. */ diff --git a/processing/src/test/java/org/apache/druid/segment/loading/PartialProjectionLoadSpecTest.java b/processing/src/test/java/org/apache/druid/segment/loading/PartialProjectionLoadSpecTest.java index b64a0d0b7c1a..8c3eaeee1579 100644 --- a/processing/src/test/java/org/apache/druid/segment/loading/PartialProjectionLoadSpecTest.java +++ b/processing/src/test/java/org/apache/druid/segment/loading/PartialProjectionLoadSpecTest.java @@ -27,7 +27,19 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.module.SimpleModule; import com.google.common.collect.ImmutableMap; +import org.apache.druid.error.DruidException; import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.query.OrderBy; +import org.apache.druid.query.aggregation.AggregatorFactory; +import org.apache.druid.query.aggregation.CountAggregatorFactory; +import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.file.SegmentFileMetadata; +import org.apache.druid.segment.projections.AggregateProjectionSchema; +import org.apache.druid.segment.projections.ProjectionMetadata; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.timeline.partition.NumberedShardSpec; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -167,6 +179,92 @@ void testRejectsNullFingerprint() ); } + @Test + void testGetSelectedBundleNamesReturnsProjections() + { + // Projection names double as bundle names; each requested name must appear in the segment's projections. + PartialProjectionLoadSpec spec = new PartialProjectionLoadSpec( + DELEGATE, + List.of("user_daily", "user_hourly"), + FINGERPRINT, + jsonMapper + ); + final SegmentFileMetadata metadata = metadataWithProjections("user_daily", "user_hourly", "some_other"); + Assertions.assertEquals(List.of("user_daily", "user_hourly"), spec.getSelectedBundleNames(anySegment(), metadata)); + } + + @Test + void testGetSelectedBundleNamesThrowsWhenProjectionMissing() + { + // Defensive: a wire form referring to a projection this segment doesn't have. Should only happen on + // matcher/reader drift or a coding bug + PartialProjectionLoadSpec spec = new PartialProjectionLoadSpec( + DELEGATE, + List.of("user_daily", "vanished"), + FINGERPRINT, + jsonMapper + ); + final SegmentFileMetadata metadata = metadataWithProjections("user_daily"); + final DruidException thrown = Assertions.assertThrows( + DruidException.class, + () -> spec.getSelectedBundleNames(anySegment(), metadata) + ); + Assertions.assertTrue( + thrown.getMessage().contains("does not contain projection[vanished]"), + "unexpected message: " + thrown.getMessage() + ); + } + + @Test + void testGetSelectedBundleNamesThrowsWhenMetadataHasNoProjections() + { + // Defensive: a segment with no projections in metadata can't satisfy any projection rule. + PartialProjectionLoadSpec spec = new PartialProjectionLoadSpec( + DELEGATE, + List.of("user_daily"), + FINGERPRINT, + jsonMapper + ); + final SegmentFileMetadata metadata = new SegmentFileMetadata(List.of(), Map.of(), null, null, null, null); + final DruidException thrown = Assertions.assertThrows( + DruidException.class, + () -> spec.getSelectedBundleNames(anySegment(), metadata) + ); + Assertions.assertTrue( + thrown.getMessage().contains("metadata has no projections"), + "unexpected message: " + thrown.getMessage() + ); + } + + private static DataSegment anySegment() + { + return DataSegment.builder(SegmentId.of("ds", Intervals.ETERNITY, "v1", new NumberedShardSpec(0, 1))) + .size(0) + .build(); + } + + private static SegmentFileMetadata metadataWithProjections(String... names) + { + final List projections = new java.util.ArrayList<>(names.length); + for (String name : names) { + projections.add(new ProjectionMetadata(1, projectionSchemaNamed(name))); + } + return new SegmentFileMetadata(List.of(), Map.of(), null, null, projections, null); + } + + private static AggregateProjectionSchema projectionSchemaNamed(String name) + { + return new AggregateProjectionSchema( + name, + null, + null, + VirtualColumns.EMPTY, + List.of("dim1"), + new AggregatorFactory[]{new CountAggregatorFactory("cnt")}, + List.of(OrderBy.ascending("dim1")) + ); + } + /** * Stub LoadSpec used to verify delegation. Uses the same JSON "type"=="stub" key as the test {@link #DELEGATE}. */ diff --git a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java index c2dc29c75164..d9751bcc68be 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java +++ b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java @@ -513,6 +513,16 @@ public Closeable acquireReference() )); } + /** + * Current outstanding reference count (excluding the internal party released by {@link #unmount}). Returns 0 when + * the bundle has never been mounted or has already been fully cleaned up. + */ + public int getNumReferences() + { + final ReferenceCountingCloseableObject current = references.get(); + return (current == null || current.isClosed()) ? 0 : current.getNumReferences(); + } + /** * The actual unmount work, invoked by the reference-counted gate's {@code onAdvance} once every outstanding * reference (plus the wrapper's own initial party) has been released. Evicts owned containers, releases parent diff --git a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java index 8a53c1b5f139..e0255256ecbd 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java +++ b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java @@ -56,11 +56,12 @@ * {@link #restoreBundlesFromDisk} to discover, reserve, and mount any bundles whose container files survived. The * same call from the fresh acquire path is a no-op (no on-disk containers to restore). * - * Parent-set inference is delegated to {@link PartialSegmentMetadataCacheEntry#inferParentBundles}. A bundle whose - * inferred parent isn't itself present on disk is treated as orphaned: its on-disk container files are deleted - * (via {@link PartialSegmentFileMapperV10#evictContainer}, which also clears the relevant bitmap bits) and the bundle - * is not restored. The next access through the cache manager acquire path then triggers a clean cold re-fetch, the - * same fall-back as when the cache manager finds a segment listed in the info directory but missing on disk. + * Dependency inference is delegated to {@link PartialSegmentMetadataCacheEntry#inferBundleDependencies}. A bundle + * whose inferred dependency isn't itself present on disk is treated as orphaned: its on-disk container files + * are deleted (via {@link PartialSegmentFileMapperV10#evictContainer}, which also clears the relevant bitmap bits) and + * the bundle is not restored. The next access through the cache manager acquire path then triggers a clean cold + * re-fetch, the same fall-back as when the cache manager finds a segment listed in the info directory but missing on + * disk. */ public final class PartialSegmentCacheBootstrap { @@ -68,16 +69,23 @@ public final class PartialSegmentCacheBootstrap /** * Reserve a partial segment's metadata cache entry on the supplied location from the on-disk header. Light-weight: - * no range read, no file mapper, no bundle work. The entry is registered as a weak cache entry (consistent with - * the runtime acquire path), so it is evictable once mounted; bootstrap-restored data is treated as a cache - * optimization, not a permanent fixture. The caller is expected to drive the actual mount through + * no range read, no file mapper, no bundle work. The caller is expected to drive the actual mount through * {@code metadata.mount(location)} later (typically via {@code SegmentCacheManager#bootstrap}), which builds the * file mapper (parsing the header from local disk, no fetch) and cascades into {@link #restoreBundlesFromDisk}. + *

+ * A non-null {@code fingerprint} means the segment was persisted via a partial-load rule and the metadata entry is + * installed in {@link StorageLocation}'s static map (kept sticky across the segment's lifetime on the historical); + * null means an ordinary on-demand partial load, which uses the weak path. {@code staticBundleNames} is the bundle- + * name subset that should likewise be reserved as static by {@link #restoreBundlesFromDisk} (remaining on-disk + * bundles restore as weak); empty for non-partial-rule loads. * * @param segmentId the segment whose entries are being restored * @param localCacheDir the per-segment directory containing the header + container files * @param targetFilename the V10 entry-point filename * @param externalFilenames any external segment file names that were registered as children of the entry-point file + * @param staticBundleNames bundle names the rule has pinned as static; empty for ordinary partial loads + * @param fingerprint the partial-load-rule fingerprint from the persisted load spec; null for non-rule loads. + * Non-null also selects the static reservation path for the metadata entry * @param rangeReader the segment's deep-storage range reader, retained for later on-demand fetches * @param jsonMapper used by the metadata entry's mount path to parse the header * @param storagePool thread pool the async cursor path submits on-demand column downloads to (which bounds @@ -91,6 +99,8 @@ public static PartialSegmentMetadataCacheEntry reserveFromDisk( File localCacheDir, String targetFilename, List externalFilenames, + Set staticBundleNames, + @Nullable String fingerprint, SegmentRangeReader rangeReader, ObjectMapper jsonMapper, @Nullable StorageLoadingThreadPool storagePool, @@ -113,13 +123,16 @@ public static PartialSegmentMetadataCacheEntry reserveFromDisk( localCacheDir, targetFilename, externalFilenames, + staticBundleNames, + fingerprint, rangeReader, jsonMapper, storagePool, actualMetadataSize ); - if (!location.reserveWeak(metadata)) { + final boolean reserved = fingerprint != null ? location.reserve(metadata) : location.reserveWeak(metadata); + if (!reserved) { throw DruidException.defensive( "Failed to reserve metadata entry for partial segment[%s] at location[%s]", segmentId, @@ -165,16 +178,16 @@ static void restoreBundlesFromDisk(PartialSegmentMetadataCacheEntry metadata, St return; } - // Classify each present bundle as either mountable or orphaned. A bundle is orphaned when its inferred parent - // set includes a bundle that isn't itself present on disk; restoring it would only produce a degenerate state - // where column reads that resolve into the missing parent would fail at query time. Instead, delete the - // orphan's on-disk containers so the next access triggers a clean cold re-fetch from deep storage. + // Classify each present bundle as either mountable or orphaned. A bundle is orphaned when its inferred + // dependency set includes a bundle that isn't itself present on disk; restoring it would only produce a + // degenerate state where column reads that resolve into the missing dependency would fail at query time. + // Instead, delete the orphan's on-disk containers so the next access triggers a clean cold re-fetch. final List mountableBundleNames = new ArrayList<>(); final Set orphanedBundleNames = new HashSet<>(); for (String name : presentBundleNames) { boolean orphaned = false; - for (PartialSegmentBundleCacheEntryIdentifier parent : metadata.inferParentBundles(name)) { - if (!presentBundleNames.contains(parent.bundleName())) { + for (PartialSegmentBundleCacheEntryIdentifier dep : metadata.inferBundleDependencies(name)) { + if (!presentBundleNames.contains(dep.bundleName())) { orphaned = true; break; } @@ -192,41 +205,79 @@ static void restoreBundlesFromDisk(PartialSegmentMetadataCacheEntry metadata, St fileMapper.mapperForContainer(ref.externalFilename()).evictContainer(ref.containerIndex()); } LOG.debug( - "Deleted on-disk state of orphaned bundle[%s] for segment[%s] (parent unrestorable); next access " + "Deleted on-disk state of orphaned bundle[%s] for segment[%s] (dependency unrestorable); next access " + "will trigger cold re-fetch", orphanName, segmentId ); } - // mount base bundle before any dependent bundle so its hold is available when dependents acquire parent holds + // Mount the base bundle before any dependent bundle so its hold is available when dependents acquire deps. mountableBundleNames.sort(Comparator.comparing(name -> !Projections.BASE_TABLE_PROJECTION_NAME.equals(name))); + // Expand the rule's direct selection to include transitive dependencies so a pinned child bundle's mount can + // take (no-op) holds against its static dependencies. + final Set staticBundleNames = new HashSet<>( + metadata.bundlesInMountOrder(metadata.getStaticBundleNames()) + ); final List mountedBundles = new ArrayList<>(); boolean success = false; try { for (String bundleName : mountableBundleNames) { - // Mountable bundles have all parents present by construction (orphans were filtered out above), so the - // inferred parent set is exactly what we want, no further filtering needed. - final List parentIds = metadata.inferParentBundles(bundleName); + // Mountable bundles have all dependencies present by construction (orphans were filtered out above), so the + // inferred dependency set is exactly what we want, no further filtering needed. + final List parentIds = metadata.inferBundleDependencies(bundleName); final PartialSegmentBundleCacheEntry bundle = PartialSegmentBundleCacheEntry.forBundle( metadata, bundleName, parentIds ); - // weak-reserve with a temporary hold so the mount call's own parent-hold acquisition can succeed; release the - // bootstrap hold immediately after, if the entry should remain alive for query-side access, the runtime - // hold chain (transitive parents from aggregates, segment-level holds from acquire APIs) keeps it pinned. - try (StorageLocation.ReservationHold bootstrapHold = - location.addWeakReservationHold(bundle.getId(), () -> bundle)) { - if (bootstrapHold == null) { + if (staticBundleNames.contains(bundleName)) { + // Rule-pinned bundle: reserve static. mount() acquires parent holds via addWeakReservationHoldIfExists, + // which returns no-op holds for static parents, so static + weak parents compose correctly. + if (!location.reserve(bundle)) { throw DruidException.defensive( - "Failed to reserve bundle entry[%s] in location[%s] during bootstrap", + "Failed to reserve static bundle entry[%s] in location[%s] during bootstrap", bundle.getId(), location.getPath() ); } - bundle.mount(location); + try { + bundle.mount(location); + } + catch (Throwable t) { + // Release the reservation ourselves before rethrowing: the outer rollback only walks bundles that made it + // into mountedBundles, so a mount-failed static bundle would otherwise leak the reservation in + // staticCacheEntries. (The weak branch below uses try-with-resources on the hold, which closes the hold + // on failure and lets StorageLocation's release runnable evict the unmounted entry.) + try { + location.release(bundle); + } + catch (Throwable releaseFailure) { + LOG.warn( + releaseFailure, + "Failed to release static bundle[%s] during bootstrap mount rollback for [%s]", + bundle.getId(), + segmentId + ); + } + throw t; + } + } else { + // weak-reserve with a temporary hold so the mount call's own parent-hold acquisition can succeed; release the + // bootstrap hold immediately after, if the entry should remain alive for query-side access, the runtime + // hold chain (transitive parents from aggregates, segment-level holds from acquire APIs) keeps it pinned. + try (StorageLocation.ReservationHold bootstrapHold = + location.addWeakReservationHold(bundle.getId(), () -> bundle)) { + if (bootstrapHold == null) { + throw DruidException.defensive( + "Failed to reserve bundle entry[%s] in location[%s] during bootstrap", + bundle.getId(), + location.getPath() + ); + } + bundle.mount(location); + } } mountedBundles.add(bundle); } @@ -243,9 +294,18 @@ static void restoreBundlesFromDisk(PartialSegmentMetadataCacheEntry metadata, St if (!success) { // Reverse-dependency rollback for bundles only; the metadata entry's own rollback is the caller's // responsibility (it will fire when the propagated throw escapes doMount's try/catch). + // + // Cleanup covers both entry types: + // - Static bundles: location.release() removes from staticCacheEntries AND unmounts. Without the release, + // a previously-successful static bundle would leak its reservation, since the caller's metadata-side + // rollback (location.release(metadata) → unmount → deleteHeaderFiles) never cascades to linked bundles. + // - Weak bundles: their bootstrap hold was closed inside the loop after successful mount. release() is a + // no-op for weak entries (only touches staticCacheEntries), so we also call removeUnheldWeakEntry to + // evict the now-unheld weak entry from weakCacheEntries synchronously. for (PartialSegmentBundleCacheEntry bundle : mountedBundles) { try { - bundle.unmount(); + location.release(bundle); + location.removeUnheldWeakEntry(bundle.getId()); } catch (Throwable t) { LOG.warn(t, "Failed to roll back bundle[%s] during bootstrap failure for [%s]", bundle.getId(), segmentId); diff --git a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntry.java b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntry.java index d597b656da57..f5bd1aaa40b3 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntry.java +++ b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntry.java @@ -48,7 +48,10 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Collection; +import java.util.HashMap; +import java.util.LinkedHashSet; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.concurrent.Callable; @@ -94,6 +97,15 @@ public class PartialSegmentMetadataCacheEntry implements SegmentCacheEntry, Resi private final File localCacheDir; private final String targetFilename; private final List externalFilenames; + private final Set staticBundleNames; + /** + * The partial-load-rule fingerprint carried by the load spec wrapper that produced this entry, or {@code null} for + * ordinary on-demand partial loads (no rule). {@link SegmentLocalCacheManager#loadPartial} compares this against the + * incoming wrapper's fingerprint to distinguish an idempotent re-load from a rule change; a mismatch triggers + * drop-and-reload so the new rule's bundle selection takes effect. + */ + @Nullable + private final String fingerprint; private final SegmentRangeReader rangeReader; private final ObjectMapper jsonMapper; private final long reservationEstimate; @@ -152,6 +164,8 @@ public PartialSegmentMetadataCacheEntry( File localCacheDir, String targetFilename, List externalFilenames, + Set staticBundleNames, + @Nullable String fingerprint, SegmentRangeReader rangeReader, ObjectMapper jsonMapper, @Nullable StorageLoadingThreadPool storagePool, @@ -169,7 +183,9 @@ public PartialSegmentMetadataCacheEntry( this.id = new SegmentCacheEntryIdentifier(segmentId); this.localCacheDir = localCacheDir; this.targetFilename = targetFilename; - this.externalFilenames = List.copyOf(externalFilenames); + this.externalFilenames = externalFilenames == null ? List.of() : List.copyOf(externalFilenames); + this.staticBundleNames = staticBundleNames == null ? Set.of() : Set.copyOf(staticBundleNames); + this.fingerprint = fingerprint; this.rangeReader = rangeReader; this.jsonMapper = jsonMapper; this.storagePool = storagePool; @@ -199,6 +215,25 @@ File getLocalCacheDir() return localCacheDir; } + /** + * Bundle names pinned as static by the partial-load rule that produced this entry. Empty for ordinary on-demand + * partial loads (every bundle restores as weak). + */ + public Set getStaticBundleNames() + { + return staticBundleNames; + } + + /** + * The partial-load-rule fingerprint carried by the load spec wrapper that produced this entry, or {@code null} for + * ordinary on-demand partial loads. Used by {@link SegmentLocalCacheManager#loadPartial} to detect rule changes. + */ + @Nullable + public String getFingerprint() + { + return fingerprint; + } + @Override public long getSize() { @@ -463,6 +498,11 @@ private static void awaitMount(SettableFuture future) throws IOException * outstanding, the actual unmap-and-delete work is deferred until the last reference releases; in that case this * method returns immediately and {@link #doActualUnmount} will fire later on the thread that closes the last * reference. With no outstanding references, cleanup runs synchronously on the caller's thread. + *

+ * If this entry was reserved but never mounted, there is no reference-counted gate to close; instead run the + * {@link #setOnUnmount onUnmount} hook directly so external cleanup (e.g. info-file deletion) still fires. Without + * this fall-through a caller that sets the hook before {@code mount()} and releases the reservation on a mount + * failure would leak whatever resource the hook was intended to clean up. */ @Override public void unmount() @@ -470,6 +510,28 @@ public void unmount() final ReferenceCountingCloseableObject current = references.get(); if (current != null && !current.isClosed()) { current.close(); + return; + } + // Never-mounted (or already-completed-cleanup) release: doActualUnmount won't run, so fire the hook here. + runOnUnmountHookOnce(); + } + + /** + * Atomically extract and run the {@link #setOnUnmount onUnmount} hook if one is set. Idempotent across concurrent + * callers; the {@code getAndSet(null)} ensures exactly one caller ever observes a non-null hook. Called from both + * {@link #doActualUnmount} (post-cleanup on the mounted path) and {@link #unmount} (never-mounted fallback), and + * safe to call multiple times. + */ + private void runOnUnmountHookOnce() + { + final Runnable hook = onUnmount.getAndSet(null); + if (hook != null) { + try { + hook.run(); + } + catch (Throwable t) { + LOG.warn(t, "onUnmount hook failed for partial segment metadata entry[%s]", segmentId); + } } } @@ -502,6 +564,56 @@ public Closeable acquireMetadataReference() )); } + /** + * Whether releasing this entry + all its linked bundles via a drop cascade would leave any cleanup deferred by an + * external reference (a running query holding a {@link Segment} or a raw {@link #acquireMetadataReference} / + * {@link PartialSegmentBundleCacheEntry#acquireReference}). Returns false when the cascade would run to completion + * synchronously; safe to reconcile. + *

+ * Distinguishes external refs from internal cascade-tracked refs by subtracting the known cross-entry structure: + *

+ * Any {@link ReferenceCountingCloseableObject#getNumReferences} count in excess of that structure means an outside + * consumer holds the entry and drop's cascade would defer to it. + *

+ * Used by reconciliation paths to decide BEFORE drop whether it's safe to proceed. Checking post-drop is unsafe: + * drop already removed entries from {@code staticCacheEntries} synchronously, so a "yes it's deferred, throw" + * response would leave a window where the static entry is gone while a query still holds it — any subsequent + * loadPartial or on-demand acquire would install a fresh entry alongside the deferred one, racing on the shared + * per-segment on-disk directory. + */ + public boolean cascadeReleaseWouldDeferCleanup() + { + final ReferenceCountingCloseableObject metaRefs = references.get(); + if (metaRefs == null || metaRefs.isClosed()) { + return false; + } + // Snapshot linked bundles before doing any arithmetic — the ConcurrentHashMap-backed set can churn, but + // reconciliation paths hold the segment lock so no bundle mount/unmount runs concurrently with this call. + final Collection bundles = snapshotLinkedBundles(); + // Metadata refs above what's accounted for by linked-bundle mount-time refs → external consumer. + if (metaRefs.getNumReferences() > bundles.size()) { + return true; + } + final Map childCountByParentId = new HashMap<>(); + for (PartialSegmentBundleCacheEntry child : bundles) { + for (PartialSegmentBundleCacheEntryIdentifier parentId : child.getParentEntryIds()) { + childCountByParentId.merge(parentId, 1, Integer::sum); + } + } + for (PartialSegmentBundleCacheEntry bundle : bundles) { + final int expectedChildHolds = childCountByParentId.getOrDefault(bundle.getId(), 0); + if (bundle.getNumReferences() > expectedChildHolds) { + return true; + } + } + return false; + } + /** * Build a {@link Segment} backed by this entry's mounted file mapper. Returns {@link Optional#empty()} when the * entry isn't mounted. The segment internally acquires one {@link #acquireMetadataReference} so that closing the @@ -634,7 +746,6 @@ public Optional acquireFullReference(@Nullable Closeable extraOnClose) */ private void doActualUnmount() { - final Runnable hook; entryLock.lock(); try { if (fileMapper == null) { @@ -654,21 +765,13 @@ private void doActualUnmount() // Clear the mount-dedup gate so a subsequent mount() on this same instance starts a fresh attempt. mountFuture.set(null); deleteHeaderFiles(); - hook = onUnmount.getAndSet(null); } finally { entryLock.unlock(); } // Run the hook outside entryLock so it can touch the file system / cache manager without contending with // concurrent status reads, and so a slow or buggy hook can't deadlock against acquireReference paths. - if (hook != null) { - try { - hook.run(); - } - catch (Throwable t) { - LOG.warn(t, "onUnmount hook failed for partial segment metadata entry[%s]", segmentId); - } - } + runOnUnmountHookOnce(); } /** @@ -804,7 +907,7 @@ public Closeable acquire(String requestedBundleName) () -> PartialSegmentBundleCacheEntry.forBundle( PartialSegmentMetadataCacheEntry.this, bundleName, - inferParentBundles(bundleName) + inferBundleDependencies(bundleName) ) ); if (hold == null) { @@ -861,13 +964,13 @@ private void mountBundleOrClose( // holds + references on each parent and fails with a defensive error if a parent isn't registered+mounted at // the location. The bootstrap restore path orders base-before-dependents; this is the equivalent ordering for // the runtime acquire path, which would otherwise reach a projection bundle directly with no __base mounted. - // The bundle keeps its own parent holds/refs for its lifetime, so we hold these transient ones only across the - // mount and release them immediately after a successful mount (acquire() is acyclic: base/root have no - // parents, so the recursion terminates). + // The bundle keeps its own dependency holds/refs for its lifetime, so we hold these transient ones only across + // the mount and release them immediately after a successful mount (acquire() is acyclic: base/root have no + // dependencies, so the recursion terminates). final Closer parentHolds = Closer.create(); try { - for (PartialSegmentBundleCacheEntryIdentifier parentId : inferParentBundles(bundleName)) { - parentHolds.register(acquire(parentId.bundleName())); + for (PartialSegmentBundleCacheEntryIdentifier depId : inferBundleDependencies(bundleName)) { + parentHolds.register(acquire(depId.bundleName())); } bundle.mount(loc); } @@ -885,17 +988,19 @@ private void mountBundleOrClose( } /** - * Inference of the parent bundles that the given {@code bundleName} depends on within this segment. + * The bundles that {@code bundleName} depends on within this segment. Depending here means: at mount time, this + * bundle's file mapper must be able to acquire a hold on each dependency, and the dependency's containers must + * remain resident for as long as this bundle is mounted. *

* The rule is uniform: the base bundle and the {@link SegmentFileBuilder#ROOT_BUNDLE_NAME root bundle} have no - * parents (the root bundle owns everything written without an explicit {@code startFileBundle} call, for older - * fileGroup-less segments, or any future shared internal metadata and is structurally a peer of the base); - * every other bundle depends on the base bundle, but only if this segment actually carries one. + * dependencies (the root bundle owns everything written without an explicit {@code startFileBundle} call, for older + * fileGroup-less segments, or any future shared internal metadata and is structurally a peer of the base); every + * other bundle depends on the base bundle, but only if this segment actually carries one. *

* If future writers introduce richer dependency graphs, the rule will need to grow, likely by reading dependency * metadata the writer records explicitly. */ - public List inferParentBundles(String bundleName) + public List inferBundleDependencies(String bundleName) { if (Projections.BASE_TABLE_PROJECTION_NAME.equals(bundleName) || SegmentFileBuilder.ROOT_BUNDLE_NAME.equals(bundleName) @@ -910,6 +1015,36 @@ public List inferParentBundles(String ); } + /** + * Return {@code roots} plus every transitive dependency (as inferred by {@link #inferBundleDependencies}) in + * dependency-first order — i.e. the order to reserve+mount them in. Used both by the fresh-load path (to reserve + * a rule-pinned bundle's dependencies before the bundle itself) and by the bootstrap-restore path (to expand the + * rule's direct selection into the full static set the restore should install). + *

+ * No cycle guard: {@link #inferBundleDependencies} only returns {@code []} or {@code [__base]}, and {@code __base} + * itself has no dependencies, so the graph is at most one level deep by construction. A richer future dependency + * graph should revisit this. + */ + public List bundlesInMountOrder(Iterable roots) + { + final LinkedHashSet ordered = new LinkedHashSet<>(); + for (String root : roots) { + visitDependenciesFirst(root, ordered); + } + return List.copyOf(ordered); + } + + private void visitDependenciesFirst(String bundleName, LinkedHashSet ordered) + { + if (ordered.contains(bundleName)) { + return; + } + for (PartialSegmentBundleCacheEntryIdentifier dep : inferBundleDependencies(bundleName)) { + visitDependenciesFirst(dep.bundleName(), ordered); + } + ordered.add(bundleName); + } + /** * Whether this segment carries a {@code __base} bundle (shared base-table column data). Probed from the mounted file * mapper's actual bundle set; returns false when the entry is not mounted. 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 f5d2452d9c38..eac1a777c177 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 @@ -42,6 +42,7 @@ import org.apache.druid.segment.Segment; import org.apache.druid.segment.SegmentLazyLoadFailCallback; import org.apache.druid.segment.file.PartialSegmentFileMapperV10; +import org.apache.druid.segment.file.SegmentFileMetadata; import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.SegmentId; import org.apache.druid.utils.CloseableUtils; @@ -54,15 +55,20 @@ import java.io.InputStream; import java.nio.file.Files; import java.nio.file.StandardCopyOption; +import java.util.ArrayList; +import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Objects; import java.util.Optional; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.Lock; @@ -350,12 +356,57 @@ private void addFilesToCachedSegments( continue; } removeInfo = false; + // When the segment's loadSpec is a partial-load-rule wrapper, parse the on-disk header and resolve the + // rule-selected bundle names so the bootstrap restore can reserve metadata + selected bundles as static. + // The fingerprint from the persisted wrapper is threaded through so a later load() call can detect a rule + // change and reconcile via drop-and-reload. + Set staticBundleNames = Set.of(); + String fingerprint = null; + if (PartialLoadSpec.detectPartialLoadSpec(segment.getLoadSpec())) { + try { + final LoadSpec materializedLoadSpec = jsonMapper.convertValue(segment.getLoadSpec(), LoadSpec.class); + if (materializedLoadSpec instanceof PartialLoadSpec wrapper) { + final File headerFile = new File( + partialDir, + IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX + ); + final SegmentFileMetadata parsed = PartialSegmentFileMapperV10.parseHeaderFile( + headerFile, + jsonMapper + ).getMetadata(); + staticBundleNames = Set.copyOf(wrapper.getSelectedBundleNames(segment, parsed)); + fingerprint = wrapper.getFingerprint(); + } + } + catch (Throwable t) { + // ANY failure here means we can't honor bootstrap's "restore from disk" contract for this segment, so nuke + // everything and the coordinator will re-issue load and handle everything. + if (t instanceof DruidException) { + log.makeAlert( + t, + "Partial-load rule contract violation for segment[%s] at bootstrap; deleting on-disk cache state", + segment.getId() + ).emit(); + } else { + log.warn( + t, + "Failed to resolve partial-load wrapper for segment[%s] at bootstrap; deleting on-disk cache state", + segment.getId() + ); + } + atomicMoveAndDeleteCacheEntryDirectory(partialDir); + removeInfo = true; + continue; + } + } try { PartialSegmentCacheBootstrap.reserveFromDisk( segment.getId(), partialDir, IndexIO.V10_FILE_NAME, List.of(), + staticBundleNames, + fingerprint, rangeReader, jsonMapper, virtualStorageLoadingThreadPool, @@ -799,6 +850,8 @@ private ReservedPartial reservePartial(DataSegment dataSegment, SegmentRangeRead partialDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, rangeReader, jsonMapper, virtualStorageLoadingThreadPool, @@ -872,6 +925,454 @@ private record ReservedPartial( { } + /** + * Handle the partial-load-rule case at {@link #load} time: reserve a {@link PartialSegmentMetadataCacheEntry} as + * static, persist the info file, mount the metadata header, and reserve+mount each rule-selected bundle as static + * with its data eagerly downloaded. Static reservations are not subject to SIEVE eviction, so the rule-pinned + * subset stays resident until the segment is explicitly dropped; non-selected bundles are downloaded on-demand at + * query time via the existing weak path. + *

+ * The entire flow runs synchronously on {@link #virtualStorageLoadingThreadPool} (via submit + get) so it respects + * the pool's concurrency cap and virtual-thread runtime while giving + * {@link org.apache.druid.server.coordination.SegmentLoadDropHandler} a truthful success/failure signal. On any + * failure (mount, bundle reservation, download) the metadata + any partially-reserved bundles are cascade-released + * before the {@link SegmentLoadingException} propagates, so no orphan static reservation survives the load call. + */ + private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException + { + final LoadSpec materializedLoadSpec; + try { + materializedLoadSpec = jsonMapper.convertValue(dataSegment.getLoadSpec(), LoadSpec.class); + } + catch (Exception e) { + throw new SegmentLoadingException( + e, + "Failed to materialize partial load spec for segment[%s]", + dataSegment.getId() + ); + } + if (!(materializedLoadSpec instanceof PartialLoadSpec wrapper)) { + throw DruidException.defensive( + "Segment[%s] load spec was detected as partial but materialized to non-partial type[%s]", + dataSegment.getId(), + materializedLoadSpec.getClass().getSimpleName() + ); + } + + final SegmentRangeReader rangeReader; + try { + rangeReader = wrapper.openRangeReader(); + } + catch (IOException e) { + throw new SegmentLoadingException( + e, + "Failed to open range reader for segment[%s]", + dataSegment.getId() + ); + } + if (rangeReader == null) { + // Backend doesn't support range reads (e.g. zipped deep storage). The rule can't be honored as a partial load; + // fall through to the on-demand weak-full-load path at query time. But if a prior loadPartial installed a + // static entry with a different rule fingerprint, we cannot silently keep serving under the old rule — + // reconcile the static side so the coordinator's rule change takes effect, then return so the next query + // installs a weak entry via the on-demand path. + log.warn( + "Backend for segment[%s] does not support range reads; partial-load rule[fingerprint=%s] cannot be honored, " + + "falling back to weak full-load at query time", + dataSegment.getId(), + wrapper.getFingerprint() + ); + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(dataSegment.getId()); + final ReferenceCountingLock lock = lock(dataSegment); + synchronized (lock) { + try { + // Static-side reconciliation only. Weak entries (from bootstrap without a rule, or a prior on-demand + // acquire) already represent the fallback behavior we're falling back to — leave them alone. + boolean dropStaleStatic = false; + for (StorageLocation location : locations) { + if (!location.isReserved(id)) { + continue; + } + final SegmentCacheEntry existing = location.getCacheEntry(id); + if (existing instanceof PartialSegmentMetadataCacheEntry existingPartial + && Objects.equals(existingPartial.getFingerprint(), wrapper.getFingerprint())) { + // Same-rule re-issue: preserve the existing entry's info file across its eventual drop by clearing + // the onUnmount hook, matching the fingerprint-match short-circuit in the main reconciliation branch. + existing.setOnUnmount(null); + return; + } + dropStaleStatic = true; + } + if (dropStaleStatic) { + log.info( + "Dropping stale static partial entry for segment[%s]: current rule cannot be applied without range reads", + dataSegment.getId() + ); + drop(dataSegment); + // drop() only deleted tracked files; no reserve follows on this path, so nuke leftover per-segment dirs + // across every location or an empty dir would linger indefinitely. + nukeStaleSegmentDirsOnAllLocations(dataSegment); + } + } + finally { + unlock(dataSegment, lock); + } + } + return; + } + + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(dataSegment.getId()); + final ReferenceCountingLock lock = lock(dataSegment); + synchronized (lock) { + try { + // Reconcile against any existing entry. The end-state invariant this block establishes is: on any return + // path below, there is exactly one static partial entry for this segment (matching the wrapper fingerprint) + // and no weak entry alongside it. + // + // - Same-fingerprint static partial with no weak entry alongside: idempotent re-load. Clear onUnmount so the + // info file isn't re-deleted, return. + // - Any other state (different-fingerprint static partial, non-partial static entry, weak partial from + // bootstrap without a persisted fingerprint, weak full entry from a prior on-demand acquireSegment): drop + // the static (cascade-releases linked bundles + fires the info-file cleanup hook) AND evict the weak + // entry, then fall through to a fresh reserve. + // + // The weak eviction step is critical: {@link #drop} only touches {@link StorageLocation#staticCacheEntries}, + // so without it a lingering weak entry from bootstrap or on-demand load would coexist with the new static + // entry — same on-disk directory, both counted against capacity, and the eventual bundle downloads would + // race the weak entry's own reads. When the weak entry has active holds we cannot evict it synchronously, + // so throw and let the coordinator retry once the running query completes. drop() takes no per-segment lock + // so calling it from inside our synchronized(lock) block is safe. + PartialSegmentMetadataCacheEntry sameFingerprintStatic = null; + final List oldStaticEntries = new ArrayList<>(); + for (StorageLocation location : locations) { + if (!location.isReserved(id)) { + continue; + } + final SegmentCacheEntry existing = location.getCacheEntry(id); + if (existing == null) { + continue; + } + oldStaticEntries.add(existing); + if (existing instanceof PartialSegmentMetadataCacheEntry existingPartial + && Objects.equals(existingPartial.getFingerprint(), wrapper.getFingerprint())) { + sameFingerprintStatic = existingPartial; + } + } + final boolean weakExists = locations.stream().anyMatch(loc -> loc.isWeakReserved(id)); + if (sameFingerprintStatic != null && !weakExists && oldStaticEntries.size() == 1) { + sameFingerprintStatic.setOnUnmount(null); + return; + } + if (!oldStaticEntries.isEmpty()) { + // Precheck BEFORE drop: refuse the reconciliation if any old partial metadata entry (or any of its linked + // bundles) has an external reference. drop() would remove entries from staticCacheEntries synchronously + // but leave their doActualUnmount (onUnmount hook + deleteHeaderFiles + evictContainer) deferred until + // the pinning query releases. Between the throw and the next coordinator retry, other threads (queries, + // on-demand acquires) would observe the segment as absent, install a fresh partial entry via the acquire + // path, and race the deferred cleanup on the shared per-segment on-disk directory. Checking BEFORE drop + // keeps the old entry fully installed on the failure path so a subsequent retry finds it via + // reconciliation logic as if this load never happened. + // + // cascadeReleaseWouldDeferCleanup subtracts the known cross-entry structure (bundle→metadata refs, + // parent-bundle→child holds) from raw reference counts to isolate external consumers; those are the only + // refs drop's cascade cannot release. + for (SegmentCacheEntry old : oldStaticEntries) { + if (old instanceof PartialSegmentMetadataCacheEntry oldPartial + && oldPartial.cascadeReleaseWouldDeferCleanup()) { + throw new SegmentLoadingException( + "Cannot reconcile partial-load entry for segment[%s]: outstanding references on the old entry " + + "would defer drop's cleanup and race the new entry's on-disk state; coordinator will retry", + dataSegment.getId() + ); + } + } + log.info( + "Reconciling partial-load entry for segment[%s]: fingerprint or type changed, dropping and reloading", + dataSegment.getId() + ); + drop(dataSegment); + } + for (StorageLocation location : locations) { + if (!location.isWeakReserved(id)) { + continue; + } + location.removeUnheldWeakEntry(id); + if (location.isWeakReserved(id)) { + throw new SegmentLoadingException( + "Cannot install partial-load rule for segment[%s]: an existing weak entry has active holds; " + + "coordinator will retry", + dataSegment.getId() + ); + } + } + + // After any reconciliation drop/eviction above, nuke leftover per-segment directories on every location. + // The drop cascade deleted TRACKED files (header via metadata's doActualUnmount, containers via each bundle's + // evictContainer), but not the enclosing per-segment directory nor any orphan files. If the new reserve below + // lands on a DIFFERENT location than the dropped entry, the old location's now-empty dir would accumulate on + // disk indefinitely; if the new reserve lands on the SAME location, the mkdirp inside the reserve loop + // recreates the dir fresh. + nukeStaleSegmentDirsOnAllLocations(dataSegment); + + // Reserve the metadata entry as static on the first location that accepts BOTH the reservation and a + // per-segment mkdirp. mkdirp failures are location-local (inode exhaustion, permission mismatch, transient + // EIO on one disk) — a different location may accept, so release-and-continue rather than abort. Collect + // per-location failures to attach as suppressed exceptions on the final alert. + StorageLocation reservedLocation = null; + PartialSegmentMetadataCacheEntry metadata = null; + final List perLocationFailures = new ArrayList<>(); + final Iterator iterator = strategy.getLocations(); + while (iterator.hasNext()) { + final StorageLocation candidate = iterator.next(); + final File partialDir = new File(candidate.getPath(), dataSegment.getId().toString()); + final PartialSegmentMetadataCacheEntry candidateEntry = new PartialSegmentMetadataCacheEntry( + dataSegment.getId(), + partialDir, + IndexIO.V10_FILE_NAME, + List.of(), + // staticBundleNames left empty: the fresh-load path below resolves selected bundles AFTER metadata mount + // (since the bundle names depend on the parsed header) and reserves them static directly via + // reserveAndMountStaticBundle. The field is only consulted by bootstrap restore, which has a parsed + // on-disk header available before constructing the entry. + Set.of(), + wrapper.getFingerprint(), + rangeReader, + jsonMapper, + virtualStorageLoadingThreadPool, + config.getVirtualStorageMetadataReservationEstimate() + ); + if (!candidate.reserve(candidateEntry)) { + continue; + } + try { + FileUtils.mkdirp(partialDir); + } + catch (IOException mkdirpFailure) { + candidate.release(candidateEntry); + log.warn( + mkdirpFailure, + "Failed to create partial cache dir on location[%s] for segment[%s]; trying next location", + candidate.getPath(), + dataSegment.getId() + ); + perLocationFailures.add(mkdirpFailure); + continue; + } + reservedLocation = candidate; + metadata = candidateEntry; + break; + } + + if (reservedLocation == null) { + final SegmentLoadingException noCapacity = new SegmentLoadingException( + "Unable to reserve static partial metadata entry for segment[%s] on any location; ensure enough disk " + + "space and that per-segment cache directories can be created", + dataSegment.getId() + ); + perLocationFailures.forEach(noCapacity::addSuppressed); + log.makeAlert( + noCapacity, + "Failed to reserve partial-load metadata on any location for segment[%s]", + dataSegment.getId() + ).emit(); + throw noCapacity; + } + + // Persist the info file so bootstrap can rediscover this segment on restart. The info dir is shared across + // locations (single configured path), so failure here is not a per-disk problem — no fall-through, alert + // and abort. + try { + storeInfoFile(dataSegment); + } + catch (IOException e) { + reservedLocation.release(metadata); + final SegmentLoadingException storeFailure = new SegmentLoadingException( + e, + "Failed to write info file for segment[%s]", + dataSegment.getId() + ); + log.makeAlert( + storeFailure, + "Failed to persist partial-load info file for segment[%s]; check druid.segmentCache.infoDir", + dataSegment.getId() + ).emit(); + throw storeFailure; + } + metadata.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); + + // Mount metadata + reserve+mount+download each rule-selected bundle as static. Runs on the loading thread pool + // (submit + get) so the flow respects the pool's concurrency cap while giving the calling thread a truthful + // success/failure signal. Any exception cascade-releases the metadata + every bundle reserved so far, so no + // orphan static reservation survives the load call, the reservation's onUnmount hook then deletes the info + // file, mirroring drop() semantics. + final StorageLocation theLocation = reservedLocation; + final PartialSegmentMetadataCacheEntry theMetadata = metadata; + // All rollback responsibility (bundles AND metadata) lives INSIDE the callable so no state mutation of + // theMetadata crosses threads. This avoids the FutureTask cancel-vs-cleanup race: cancel(true) transitions + // the future to CANCELLED immediately, so a subsequent get() returns CancellationException without waiting + // for the callable's cleanup to unwind — if the caller then released theMetadata itself, that release would + // race the still-running pool-thread rollback. With cleanup owned by the pool thread, the caller only ever + // observes the outcome; there is no shared state to race. + final Future future = virtualStorageLoadingThreadPool.getExecutorService().submit(() -> { + final List mounted = new ArrayList<>(); + try { + theMetadata.mount(theLocation); + final PartialSegmentFileMapperV10 mapper = theMetadata.getFileMapper(); + if (mapper == null) { + throw DruidException.defensive( + "Partial metadata for segment[%s] mounted without a file mapper", + dataSegment.getId() + ); + } + final List selectedBundleNames = wrapper.getSelectedBundleNames( + dataSegment, + mapper.getSegmentFileMetadata() + ); + // A rule-pinned bundle's dependencies must also be pinned: the bundle's own mount call acquires holds on + // each dependency via addWeakReservationHoldIfExists, which requires the dependency to be registered. + // Iterate in dependency-first order so children see their dependencies resident. + for (String bundleName : theMetadata.bundlesInMountOrder(selectedBundleNames)) { + mounted.add(reserveAndMountStaticBundle(theLocation, theMetadata, mapper, bundleName)); + } + return null; + } + catch (Throwable t) { + // Cascade rollback on the SAME thread: bundles first in reverse mount order, then the metadata (whose + // release-triggered doActualUnmount hook deletes the info file). + for (int i = mounted.size() - 1; i >= 0; i--) { + try { + theLocation.release(mounted.get(i)); + } + catch (Throwable releaseFailure) { + log.warn( + releaseFailure, + "Failed to release bundle during rollback for segment[%s]", + dataSegment.getId() + ); + } + } + releaseMetadataQuietly(theLocation, theMetadata, dataSegment); + throw t; + } + }); + + try { + future.get(); + } + catch (InterruptedException e) { + // Caller interrupted while waiting. Restore the interrupt flag and propagate, but do NOT cancel the pool + // task: FutureTask.cancel(true) transitions to CANCELLED synchronously, but the callable's run() may still + // be executing its catch/finally on the pool thread — so a subsequent metadata release from the caller + // would race that cleanup. Leaving the task alone lets its own cleanup complete on the pool thread + // naturally; on eventual shutdown the pool's shutdownNow will interrupt the worker and drive the same + // rollback via the callable's catch. The coordinator's retry will find the segment in one of three + // consistent end-states: fully mounted (task raced to success — reconciliation short-circuits), fully + // rolled-back (task failed after interrupt — fresh reserve), or still in flight (retry throws until the + // outstanding-references guard sees the task's own metadata release fire). + Thread.currentThread().interrupt(); + throw new SegmentLoadingException(e, "Interrupted while loading partial segment[%s]", dataSegment.getId()); + } + catch (Throwable t) { + // Callable ran to a terminal state. On failure it already released metadata + bundles on the pool thread + // via its catch block. We only propagate the diagnostic here; no cleanup remains for the caller. Unwrap + // ExecutionException for a clean cause; handle a null cause defensively (rare, but permitted by + // ExecutionException's contract). Restore the interrupt flag if the pool worker was interrupted mid-run + // so downstream shutdown handling still observes it. + final Throwable executionCause = (t instanceof ExecutionException) ? t.getCause() : t; + final Throwable cause = (executionCause != null) ? executionCause : t; + if (cause instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } + throw new SegmentLoadingException( + cause, + "Failed to load partial segment[%s] with rule-pinned bundles", + dataSegment.getId() + ); + } + } + finally { + unlock(dataSegment, lock); + } + } + } + + /** + * Reserve a {@link PartialSegmentBundleCacheEntry} as static on {@code location}, mount it (which initializes the + * sparse containers and links it to the metadata entry), and eagerly download all of the bundle's containers via + * the supplied {@link PartialSegmentFileMapperV10}. Returns the mounted bundle on success. + *

+ * Any failure ({@code forBundle} throwing, capacity-exceeded on reserve, or a mid-mount / mid-download exception) + * releases the bundle (if it was reserved) and rethrows so {@link #loadPartial}'s outer cascade can roll back the + * metadata reservation. + */ + private PartialSegmentBundleCacheEntry reserveAndMountStaticBundle( + StorageLocation location, + PartialSegmentMetadataCacheEntry metadata, + PartialSegmentFileMapperV10 mapper, + String bundleName + ) throws IOException + { + final PartialSegmentBundleCacheEntry bundle = PartialSegmentBundleCacheEntry.forBundle( + metadata, + bundleName, + metadata.inferBundleDependencies(bundleName) + ); + if (!location.reserve(bundle)) { + throw DruidException.forPersona(DruidException.Persona.OPERATOR) + .ofCategory(DruidException.Category.CAPACITY_EXCEEDED) + .build( + "Unable to reserve static bundle[%s] for segment[%s]; ensure enough disk space", + bundleName, + metadata.getSegmentId() + ); + } + try { + bundle.mount(location); + mapper.ensureBundleDownloaded(bundleName); + return bundle; + } + catch (Throwable t) { + try { + location.release(bundle); + } + catch (Throwable releaseFailure) { + log.warn(releaseFailure, "Failed to release bundle[%s] during mount rollback", bundleName); + } + throw t; + } + } + + /** + * Release a metadata entry, swallowing any failure with a warn log. Used from {@link #loadPartial}'s rollback path + * where a release exception should not mask the original failure being propagated. + */ + private void releaseMetadataQuietly( + StorageLocation location, + PartialSegmentMetadataCacheEntry metadata, + DataSegment segment + ) + { + try { + location.release(metadata); + } + catch (Throwable releaseFailure) { + log.warn(releaseFailure, "Failed to release metadata during rollback for segment[%s]", segment.getId()); + } + } + + /** + * Move-and-delete any leftover per-segment cache directory for {@code dataSegment} on every configured location. + * Called from the partial-load reconciliation paths after drop() so an empty partial dir on the old location doesn't + * linger when the new reserve happens to land elsewhere. {@link #atomicMoveAndDeleteCacheEntryDirectory} is a no-op + * on missing paths, so this walks all locations unconditionally. + */ + private void nukeStaleSegmentDirsOnAllLocations(DataSegment dataSegment) + { + for (StorageLocation loc : locations) { + atomicMoveAndDeleteCacheEntryDirectory(new File(loc.getPath(), dataSegment.getId().toString())); + } + } + @Nullable private AcquireSegmentAction acquireExistingSegment(SegmentCacheEntryIdentifier identifier) { @@ -927,6 +1428,15 @@ public void load(final DataSegment dataSegment) throws SegmentLoadingException "load() should not be called when virtualStorageIsEphemeral is true" ); } + // Partial-load rule routing: a wrapped load spec carrying a fingerprint + per-spec selection (cluster groups, + // projections, ...) means the coordinator wants this segment loaded as a sticky partial. Reserve the metadata + // entry as STATIC (never SIEVE-evicted) and eagerly download the rule-selected bundles as STATIC. Other bundles + // are left for the existing weak/on-demand path at query time. + if (config.isVirtualStoragePartialDownloadsEnabled() + && PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec())) { + loadPartial(dataSegment); + return; + } // 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 final ReferenceCountingLock lock = lock(dataSegment); @@ -993,12 +1503,60 @@ public void bootstrap( partial.mount(location); } catch (IOException e) { + // A mount failure on a bootstrap-restored partial entry must release the reservation so it doesn't + // leak. doMount's own catch already calls removeUnheldWeakEntry for the weak variant; here we + // additionally call location.release which handles the static variant introduced by the partial-load- + // rule path (release is a no-op for weak entries but removes static ones from staticCacheEntries). + // Without this, a static rule-pinned segment whose header parse fails at mount time would occupy the + // metadata reservation forever. + try { + location.release(partial); + } + catch (Throwable releaseFailure) { + log.warn( + releaseFailure, + "Failed to release partial metadata for segment[%s] after bootstrap mount failure", + dataSegment.getId() + ); + } throw new SegmentLoadingException( e, "Failed to mount partial metadata for segment[%s]", dataSegment.getId() ); } + // Post-mount rule-integrity check: if any direct rule-pinned bundle didn't make it into the linked set + // (JVM crash mid-loadPartial, or orphan cleanup deleted a bundle whose parent was missing), the segment + // isn't fully restored to the rule's intent. Drop it entirely so the coordinator's reconciliation re- + // issues load with a clean cold-fetch — operators don't drive drop/load directly, so a self-healing + // pathway through the coordinator is the right recovery. + final Set directRulePins = partial.getStaticBundleNames(); + if (!directRulePins.isEmpty()) { + final Set mounted = new HashSet<>(); + for (PartialSegmentBundleCacheEntry b : partial.snapshotLinkedBundles()) { + mounted.add(b.getBundleName()); + } + final List missing = directRulePins.stream() + .filter(name -> !mounted.contains(name)) + .toList(); + if (!missing.isEmpty()) { + log.makeAlert( + "Partial segment[%s] restored missing rule-pinned bundles [%s]; dropping so coordinator can" + + " re-issue load with a fresh cold-fetch", + dataSegment.getId(), + missing + ).emit(); + // Cascade-release everything (linked bundles + metadata → fires info-file cleanup hook), then nuke + // the on-disk partial layout to catch any orphan container files that no bundle entry owned. + drop(dataSegment); + final File partialDir = new File(location.getPath(), dataSegment.getId().toString()); + atomicMoveAndDeleteCacheEntryDirectory(partialDir); + throw new SegmentLoadingException( + "Partial segment[%s] rule-integrity check failed at bootstrap; dropped for re-load", + dataSegment.getId() + ); + } + } } else { throw DruidException.defensive( "Unexpected cache entry type[%s] for segment[%s] during bootstrap", @@ -1069,9 +1627,19 @@ public void drop(final DataSegment segment) final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segment.getId()); for (StorageLocation location : locations) { final CacheEntry entry = location.getCacheEntry(id); - if (entry != null) { - location.release(entry); + if (entry == null) { + continue; + } + // Cascade-release linked bundle entries before the metadata. {@link StorageLocation#release} only operates on + // entries in {@code staticCacheEntries}, so this is a no-op for weak/on-demand bundles (they continue to ride + // SIEVE eviction + hold-release cleanup); for static rule-pinned bundles installed by {@link #loadPartial}, + // this is the only path that removes them from the static map. + if (entry instanceof PartialSegmentMetadataCacheEntry partial) { + for (PartialSegmentBundleCacheEntry bundle : partial.snapshotLinkedBundles()) { + location.release(bundle); + } } + location.release(entry); } } @@ -1310,10 +1878,14 @@ private static CompleteSegmentCacheEntry checkComplete(@Nullable CacheEntry entr /** * Performs an atomic move to a sibling {@link #DROP_PATH} directory, and then deletes the directory and logs about - * it. This method should only be called under the lock of a {@link #segmentLocks}. + * it. No-op if {@code path} does not exist — callers can invoke unconditionally without a preceding {@code exists()} + * guard. Should only be called under the lock of a {@link #segmentLocks}. */ private static void atomicMoveAndDeleteCacheEntryDirectory(final File path) { + if (!path.exists()) { + return; + } final File parent = path.getParentFile(); final File tempLocation = new File(parent, DROP_PATH); try { diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java index 1d94cab8d2f0..5873bfb12fde 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java @@ -519,6 +519,8 @@ void testForBundleAcceptsBundleNameContainingSlash() throws IOException cache, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(deepDir), JSON_MAPPER, null, @@ -567,6 +569,8 @@ void testForBundleSpansMainAndExternalContainers() throws IOException cache, IndexIO.V10_FILE_NAME, List.of(externalName), + Set.of(), + null, new DirectoryBackedRangeReader(deepDir), JSON_MAPPER, null, @@ -610,6 +614,8 @@ void testForBundleRootOwnsAllUngroupedContainers() throws IOException cache, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(deepDir), JSON_MAPPER, null, @@ -641,6 +647,8 @@ void testResolveBundleNameFallsBackToRootWhenRootIsSoleBundle() throws IOExcepti cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(deepDir), JSON_MAPPER, null, @@ -887,6 +895,8 @@ private PartialSegmentMetadataCacheEntry newMetadataEntry() cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(deepStorageDir), JSON_MAPPER, null, diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java index c71e209d004e..6d2b725ad4aa 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java @@ -334,6 +334,8 @@ void testMountFailureRemovesLingeringWeakEntry() throws IOException cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, failingReader, JSON_MAPPER, null, @@ -402,6 +404,8 @@ void testReserveFailsWhenHeaderMissing() cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(deepStorageDir), JSON_MAPPER, null, @@ -485,6 +489,8 @@ private void primeOnDiskState() throws IOException cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(deepStorageDir), JSON_MAPPER, null, @@ -533,6 +539,8 @@ private PartialSegmentMetadataCacheEntry restoreFromDisk(StorageLocation locatio cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(deepStorageDir), JSON_MAPPER, null, @@ -542,6 +550,226 @@ private PartialSegmentMetadataCacheEntry restoreFromDisk(StorageLocation locatio return metadata; } + /** + * Static-restore variant: reserve metadata + selected bundles as static entries. + */ + private PartialSegmentMetadataCacheEntry restoreFromDiskAsStatic(StorageLocation location, Set staticBundles) + throws IOException + { + final PartialSegmentMetadataCacheEntry metadata = PartialSegmentCacheBootstrap.reserveFromDisk( + SEGMENT_ID, + cacheDir, + IndexIO.V10_FILE_NAME, + List.of(), + staticBundles, + "v1:test-fingerprint", + new DirectoryBackedRangeReader(deepStorageDir), + JSON_MAPPER, + null, + location + ); + metadata.mount(location); + return metadata; + } + + @Test + void testStaticReserveFromDiskRegistersMetadataInStaticMap() throws IOException + { + primeOnDiskState(); + final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); + restoreFromDiskAsStatic(location, Set.of(AGG_BUNDLE)); + + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + Assertions.assertTrue(location.isReserved(metaId), "metadata entry should be static when reserveAsStatic=true"); + Assertions.assertFalse(location.isWeakReserved(metaId), "metadata entry should NOT be weak"); + } + + @Test + void testStaticReserveFromDiskMakesSelectedBundleStatic() throws IOException + { + primeOnDiskState(); + final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); + restoreFromDiskAsStatic(location, Set.of(AGG_BUNDLE)); + + final PartialSegmentBundleCacheEntryIdentifier aggId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); + Assertions.assertTrue(location.isReserved(aggId), "rule-selected bundle should be in staticCacheEntries"); + Assertions.assertFalse(location.isWeakReserved(aggId)); + } + + @Test + void testStaticReserveFromDiskExpandsToParentBundles() throws IOException + { + // Selecting the aggregate projection must also pin its __base parent as static: otherwise a (no-op) parent hold + // taken by the aggregate's mount call would resolve to a weak entry the cache could later SIEVE-evict, breaking + // the rule's stickiness guarantee. + primeOnDiskState(); + final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); + restoreFromDiskAsStatic(location, Set.of(AGG_BUNDLE)); + + final PartialSegmentBundleCacheEntryIdentifier baseId = new PartialSegmentBundleCacheEntryIdentifier( + SEGMENT_ID, + Projections.BASE_TABLE_PROJECTION_NAME + ); + Assertions.assertTrue(location.isReserved(baseId), "transitive parent bundle should be static"); + Assertions.assertFalse(location.isWeakReserved(baseId)); + } + + @Test + void testNonStaticReserveFromDiskKeepsBundlesWeak() throws IOException + { + // Sanity: when no rule selection is passed and reserveAsStatic=false, both metadata and bundles are weak — the + // pre-existing behavior for ordinary on-demand partials. + primeOnDiskState(); + final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); + restoreFromDisk(location); + + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + final PartialSegmentBundleCacheEntryIdentifier aggId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); + final PartialSegmentBundleCacheEntryIdentifier baseId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME); + Assertions.assertTrue(location.isWeakReserved(metaId)); + Assertions.assertTrue(location.isWeakReserved(aggId)); + Assertions.assertTrue(location.isWeakReserved(baseId)); + Assertions.assertFalse(location.isReserved(metaId)); + Assertions.assertFalse(location.isReserved(aggId)); + Assertions.assertFalse(location.isReserved(baseId)); + } + + @Test + void testStaticRestoreSkipsRulePinnedBundleAbsentOnDisk() throws IOException + { + // JVM-crash-mid-loadPartial scenario at the reserveFromDisk + mount layer: metadata + info file are on disk, and + // __base is on disk, but the rule-pinned aggregate bundle's containers never finished downloading. This helper + // level restores metadata + __base as static and leaves the missing aggregate absent — no phantom static entry. + // The higher-level SegmentLocalCacheManager.bootstrap() layer then detects the missing rule-pinned bundle and + // drops the segment entirely so the coordinator can re-issue a clean cold-fetch. + primeOnDiskState(); + + // Delete the aggregate bundle's container file(s) but keep base's containers intact. + final PartialSegmentFileMapperV10 introspect = createMapper(deepStorageDir, cacheDir); + final Set aggContainers = new HashSet<>(); + final Set baseContainers = new HashSet<>(); + for (var entry : introspect.getSegmentFileMetadata().getFiles().entrySet()) { + final String name = entry.getKey(); + final int slash = name.indexOf('/'); + if (slash < 0) { + continue; + } + final String group = name.substring(0, slash); + if (AGG_BUNDLE.equals(group)) { + aggContainers.add(entry.getValue().getContainer()); + } else if (Projections.BASE_TABLE_PROJECTION_NAME.equals(group)) { + baseContainers.add(entry.getValue().getContainer()); + } + } + introspect.close(); + final Set aggOnly = new HashSet<>(aggContainers); + aggOnly.removeAll(baseContainers); + if (aggOnly.isEmpty()) { + // Tiny test segment happens to share all containers between base and agg — the scenario isn't reachable. + return; + } + for (Integer ci : aggOnly) { + final File cf = new File(cacheDir, StringUtils.format("%s.container.%05d", IndexIO.V10_FILE_NAME, ci)); + Assertions.assertTrue(cf.exists()); + Assertions.assertTrue(cf.delete()); + } + + final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); + restoreFromDiskAsStatic(location, Set.of(AGG_BUNDLE)); + + // Metadata + __base are restored as STATIC (rule intent preserved for what's available on disk). + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + final PartialSegmentBundleCacheEntryIdentifier baseId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME); + Assertions.assertTrue(location.isReserved(metaId), "metadata should be static"); + Assertions.assertTrue(location.isReserved(baseId), "__base (present on disk) should be static"); + + // The missing rule-pinned bundle must NOT be pinned as static (no phantom entry) and must NOT be silently + // materialized as weak either — bootstrap doesn't do deep-storage fetches. The next runtime acquire will pick it + // up on the weak/on-demand path. + final PartialSegmentBundleCacheEntryIdentifier aggId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); + Assertions.assertNull( + location.getCacheEntry(aggId), + "missing rule-pinned bundle must not be in the cache after bootstrap; runtime weak-acquire handles it" + ); + } + + @Test + void testStaticRestoreReleasesReservationsWhenBundleMountFails() throws IOException + { + // Regression guard: when a static bundle's mount() throws after its reservation was taken via location.reserve(), + // the reservation must be released. Also — any static bundle mounted successfully earlier in the loop must be + // released too; the outer rollback previously only called bundle.unmount(), which doesn't remove from + // staticCacheEntries. Without both, mount-time failures leave orphan static reservations after bootstrap. + primeOnDiskState(); + + // Identify the aggregate bundle's exclusive containers (those NOT shared with __base), and replace one with a + // directory of the same name. filterByContainerPresence uses File#exists which returns true for directories, so the + // agg bundle passes classification, but PartialSegmentFileMapperV10#ensureContainerInitialized opens the file with + // new RandomAccessFile(..., "rw") which throws when the path is a directory — forcing bundle.mount to throw AFTER + // location.reserve(agg) succeeded. __base's mount succeeds and lands in mountedBundles first, exercising both the + // inline reserve-then-mount catch (agg) and the outer mountedBundles rollback (base). + final PartialSegmentFileMapperV10 introspect = createMapper(deepStorageDir, cacheDir); + final Set aggContainers = new HashSet<>(); + final Set baseContainers = new HashSet<>(); + for (var entry : introspect.getSegmentFileMetadata().getFiles().entrySet()) { + final String name = entry.getKey(); + final int slash = name.indexOf('/'); + if (slash < 0) { + continue; + } + final String group = name.substring(0, slash); + if (AGG_BUNDLE.equals(group)) { + aggContainers.add(entry.getValue().getContainer()); + } else if (Projections.BASE_TABLE_PROJECTION_NAME.equals(group)) { + baseContainers.add(entry.getValue().getContainer()); + } + } + introspect.close(); + final Set aggOnly = new HashSet<>(aggContainers); + aggOnly.removeAll(baseContainers); + if (aggOnly.isEmpty()) { + // Tiny test segment happens to share all containers between base and agg — the scenario isn't reachable. + return; + } + final Integer aggOnlyIdx = aggOnly.iterator().next(); + final File aggContainer = new File( + cacheDir, + StringUtils.format("%s.container.%05d", IndexIO.V10_FILE_NAME, aggOnlyIdx) + ); + Assertions.assertTrue(aggContainer.delete()); + Assertions.assertTrue(aggContainer.mkdir(), "must replace container file with a directory to force mount failure"); + + final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); + Assertions.assertThrows(Throwable.class, () -> restoreFromDiskAsStatic(location, Set.of(AGG_BUNDLE))); + + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + final PartialSegmentBundleCacheEntryIdentifier baseId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME); + final PartialSegmentBundleCacheEntryIdentifier aggId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); + + // The specific finding: agg's reservation was released by the inline catch after mount threw. + Assertions.assertFalse(location.isReserved(aggId), "failed static bundle mount must release its own reservation"); + // Broader cleanup: the previously-succeeded base bundle in mountedBundles is also released by the outer rollback, + // not left as an orphan (bundle.unmount() alone would only clear the mount state, not staticCacheEntries). + Assertions.assertFalse( + location.isReserved(baseId), + "previously-succeeded static bundle must be released on outer rollback" + ); + // Metadata cleanup is the caller's responsibility, but metadata.mount()'s catch calls removeUnheldWeakEntry, which + // is a no-op for static entries. This test drives mount directly (no wrapping SegmentLocalCacheManager) so verify + // only bundle-level cleanup here; the manager-level bootstrap catch is exercised by other tests. + Assertions.assertTrue( + location.isReserved(metaId), + "static metadata reservation persists at this layer; the caller (manager bootstrap) is responsible for release" + ); + } + private static PartialSegmentFileMapperV10 createMapper(File deepStorageDir, File cacheDir) throws IOException { return PartialSegmentFileMapperV10.create( diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java index f09ca7ba1be1..3dd033734237 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java @@ -46,6 +46,7 @@ import java.io.IOException; import java.util.List; import java.util.Optional; +import java.util.Set; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; @@ -219,6 +220,25 @@ void testOnUnmountHookRunsAfterStorageLocationCleanup() throws IOException ); } + @Test + void testOnUnmountHookRunsOnReleaseBeforeMount() + { + // Regression guard: the SegmentLocalCacheManager.loadPartial flow persists the segment info file BEFORE mounting + // the entry, and registers a hook that deletes the info file on unmount. If mount then fails (or is never called) + // and the reservation is released, unmount() must still fire the hook — otherwise the info file leaks on disk. + final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE); + final AtomicReference hookFired = new AtomicReference<>(false); + entry.setOnUnmount(() -> hookFired.set(true)); + + entry.unmount(); + Assertions.assertTrue(hookFired.get(), "onUnmount hook must run even when the entry was never mounted"); + + // Second unmount is a no-op — the hook must not double-fire. + hookFired.set(false); + entry.unmount(); + Assertions.assertFalse(hookFired.get(), "onUnmount hook must fire exactly once across repeated unmount() calls"); + } + @Test void testConstructorRejectsNonPositiveEstimate() { @@ -229,6 +249,8 @@ void testConstructorRejectsNonPositiveEstimate() cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(segmentFile.getParentFile()), JSON_MAPPER, null, @@ -284,6 +306,8 @@ void testConcurrentMountIsDeduplicated() throws Exception cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, rangeReader, JSON_MAPPER, null, @@ -385,27 +409,27 @@ void testAcquireReferenceBeforeMountReturnsEmpty() } @Test - void testInferParentBundlesForBaseReturnsEmpty() + void testInferBundleDependenciesForBaseReturnsEmpty() { final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE); Assertions.assertEquals( List.of(), - entry.inferParentBundles(Projections.BASE_TABLE_PROJECTION_NAME) + entry.inferBundleDependencies(Projections.BASE_TABLE_PROJECTION_NAME) ); } @Test - void testInferParentBundlesForRootReturnsEmpty() + void testInferBundleDependenciesForRootReturnsEmpty() { final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE); Assertions.assertEquals( List.of(), - entry.inferParentBundles(SegmentFileBuilder.ROOT_BUNDLE_NAME) + entry.inferBundleDependencies(SegmentFileBuilder.ROOT_BUNDLE_NAME) ); } @Test - void testInferParentBundlesDependsOnBaseWhenBaseBundlePresent() throws IOException + void testInferBundleDependenciesIncludesBaseWhenBaseBundlePresent() throws IOException { // A segment that carries a __base bundle (the non-clustered base+projection shape): every non-base/root bundle // depends on it. Asserted uniformly for an aggregate-projection bundle and for a cluster-group bundle name (the @@ -414,20 +438,20 @@ void testInferParentBundlesDependsOnBaseWhenBaseBundlePresent() throws IOExcepti buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "some_projection") ); for (String dependent : List.of("some_projection", Projections.getClusterGroupBundleName(List.of(0, 1)))) { - final List parents = entry.inferParentBundles(dependent); - Assertions.assertEquals(1, parents.size(), "expected a __base parent for bundle[" + dependent + "]"); - Assertions.assertEquals(SEGMENT_ID, parents.getFirst().segmentId()); - Assertions.assertEquals(Projections.BASE_TABLE_PROJECTION_NAME, parents.getFirst().bundleName()); + final List deps = entry.inferBundleDependencies(dependent); + Assertions.assertEquals(1, deps.size(), "expected a __base dependency for bundle[" + dependent + "]"); + Assertions.assertEquals(SEGMENT_ID, deps.getFirst().segmentId()); + Assertions.assertEquals(Projections.BASE_TABLE_PROJECTION_NAME, deps.getFirst().bundleName()); } } @Test - void testInferParentBundlesEmptyWhenSegmentHasNoBaseBundle() throws IOException + void testInferBundleDependenciesEmptyWhenSegmentHasNoBaseBundle() throws IOException { // A clustered + aggregate-projection segment with no shared columns has no __base bundle: the base data lives in // per-group __base$ bundles and the aggregate projection is self-contained. So neither a cluster group nor - // the aggregate projection has a parent to depend on. (Pre-shared-columns; the old "aggregate always depends on - // __base" rule would have wrongly tried to mount a nonexistent __base for the projection bundle.) + // the aggregate projection has a dependency. (Pre-shared-columns; the old "aggregate always depends on __base" + // rule would have wrongly tried to mount a nonexistent __base for the projection bundle.) final PartialSegmentMetadataCacheEntry entry = mountedEntryOver( buildSegmentWithBundles( Projections.getClusterGroupBundleName(List.of(0)), @@ -437,9 +461,9 @@ void testInferParentBundlesEmptyWhenSegmentHasNoBaseBundle() throws IOException ); Assertions.assertEquals( List.of(), - entry.inferParentBundles(Projections.getClusterGroupBundleName(List.of(0))) + entry.inferBundleDependencies(Projections.getClusterGroupBundleName(List.of(0))) ); - Assertions.assertEquals(List.of(), entry.inferParentBundles("some_projection")); + Assertions.assertEquals(List.of(), entry.inferBundleDependencies("some_projection")); } private PartialSegmentMetadataCacheEntry newEntry(long estimate) @@ -449,6 +473,8 @@ private PartialSegmentMetadataCacheEntry newEntry(long estimate) cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(segmentFile.getParentFile()), JSON_MAPPER, null, @@ -472,8 +498,8 @@ private File buildTestSegment(int numFiles) throws IOException /** * Build a V10 segment whose containers are tagged with exactly the given bundle names (one column file per bundle), - * so {@link PartialSegmentMetadataCacheEntry#inferParentBundles} can be exercised against a known bundle set without - * a full ingestion. Returns the deep-storage directory containing the V10 file. + * so {@link PartialSegmentMetadataCacheEntry#inferBundleDependencies} can be exercised against a known bundle set + * without a full ingestion. Returns the deep-storage directory containing the V10 file. */ private File buildSegmentWithBundles(String... bundleNames) throws IOException { @@ -493,7 +519,7 @@ private File buildSegmentWithBundles(String... bundleNames) throws IOException /** * Reserve and mount a fresh metadata entry over the segment in {@code deepStorageDir}, into a per-call cache - * directory. The mounted entry's file mapper is what {@code inferParentBundles} probes for the base-bundle existence. + * directory. The mounted entry's file mapper is what {@code inferBundleDependencies} probes for base-bundle presence. */ private PartialSegmentMetadataCacheEntry mountedEntryOver(File deepStorageDir) throws IOException { @@ -505,6 +531,8 @@ private PartialSegmentMetadataCacheEntry mountedEntryOver(File deepStorageDir) t cache, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(deepStorageDir), JSON_MAPPER, null, diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java index 01862eb2b9a0..7ca9d71cc78b 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java @@ -492,9 +492,9 @@ void testPartialAcquireClusteredWithProjectionMountsProjectionBundleWithoutBase( Assertions.assertEquals(CLUSTERED_SEGMENT_ID, segment.getId()); // group-by tenant + sum(x) matches the aggregate projection. Building this cursor drives the 'proj' bundle to - // mount through the real acquire path. inferParentBundles must return no parent for it (the clustered segment - // has no __base bundle); the old "aggregate always depends on __base" rule would have tried to mount a - // nonexistent __base here and failed. + // mount through the real acquire path. inferBundleDependencies must return no dep for it (the clustered + // segment has no __base bundle); the old "aggregate always depends on __base" rule would have tried to mount + // a nonexistent __base here and failed. final CursorBuildSpec aggSpec = CursorBuildSpec.builder() .setGroupingColumns(List.of("tenant")) .setAggregators(List.of(new LongSumAggregatorFactory("sum_x", "x"))) diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialDropTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialDropTest.java index 6613b49dc1b4..9718294afb80 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialDropTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialDropTest.java @@ -59,6 +59,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Set; import java.util.concurrent.ThreadLocalRandom; /** @@ -197,6 +198,8 @@ private HeldMetadata registerMountedMetadata() cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(deepStorageDir), JSON_MAPPER, null, 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 new file mode 100644 index 000000000000..238a9f4e3ad6 --- /dev/null +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java @@ -0,0 +1,829 @@ +/* + * 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.segment.loading; + +import com.fasterxml.jackson.databind.InjectableValues; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.jsontype.NamedType; +import org.apache.druid.data.input.InputRow; +import org.apache.druid.data.input.ListBasedInputRow; +import org.apache.druid.data.input.impl.AggregateProjectionSpec; +import org.apache.druid.data.input.impl.DimensionsSpec; +import org.apache.druid.data.input.impl.LongDimensionSchema; +import org.apache.druid.data.input.impl.StringDimensionSchema; +import org.apache.druid.guice.LocalDataStorageDruidModule; +import org.apache.druid.jackson.SegmentizerModule; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.FileUtils; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.math.expr.ExprMacroTable; +import org.apache.druid.query.aggregation.CountAggregatorFactory; +import org.apache.druid.query.aggregation.LongSumAggregatorFactory; +import org.apache.druid.query.expression.TestExprMacroTable; +import org.apache.druid.segment.IndexBuilder; +import org.apache.druid.segment.IndexIO; +import org.apache.druid.segment.IndexSpec; +import org.apache.druid.segment.TestHelper; +import org.apache.druid.segment.column.ColumnConfig; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.data.CompressionStrategy; +import org.apache.druid.segment.file.PartialSegmentFileMapperV10; +import org.apache.druid.segment.incremental.IncrementalIndexSchema; +import org.apache.druid.segment.projections.Projections; +import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory; +import org.apache.druid.server.metrics.NoopServiceEmitter; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.timeline.partition.NoneShardSpec; +import org.joda.time.DateTime; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.File; +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ThreadLocalRandom; + +/** + * Manager-level tests for the partial-load-rule {@link SegmentLocalCacheManager#load} path: a segment whose + * {@code loadSpec} is a {@link PartialLoadSpec} wrapper should mount the metadata as a STATIC cache entry, eagerly + * download the rule-selected bundles as STATIC, and leave non-selected bundles unmounted (subject to weak/on-demand + * acquisition at query time). Drop on a static partial cleanly removes both metadata and bundles. + */ +class SegmentLocalCacheManagerPartialRuleLoadTest +{ + private static final SegmentId SEGMENT_ID = SegmentId.of("test", Intervals.of("2025/2026"), "v1", 0); + private static final DateTime TIME = DateTimes.of("2025-01-01"); + private static final String AGG_BUNDLE = "dim1_metric1_sum"; + private static final String OTHER_AGG_BUNDLE = "dim1_count"; + private static final String FINGERPRINT = "v1:rule-bundle-test"; + + private static final RowSignature ROW_SIGNATURE = RowSignature.builder() + .add("dim1", ColumnType.STRING) + .add("metric1", ColumnType.LONG) + .build(); + + private static final List PROJECTIONS = Arrays.asList( + AggregateProjectionSpec.builder(AGG_BUNDLE) + .groupingColumns(new StringDimensionSchema("dim1")) + .aggregators( + new LongSumAggregatorFactory("_metric1_sum", "metric1"), + new CountAggregatorFactory("_count") + ) + .build(), + // A second aggregate projection so tests can verify that non-selected bundles are NOT pre-mounted while + // selected + parent bundles ARE. + AggregateProjectionSpec.builder(OTHER_AGG_BUNDLE) + .groupingColumns(new StringDimensionSchema("dim1")) + .aggregators(new CountAggregatorFactory("_count")) + .build() + ); + + private static final List ROWS = Arrays.asList( + new ListBasedInputRow(ROW_SIGNATURE, TIME, ROW_SIGNATURE.getColumnNames(), Arrays.asList("a", 1L)), + new ListBasedInputRow(ROW_SIGNATURE, TIME.plusMinutes(1), ROW_SIGNATURE.getColumnNames(), Arrays.asList("a", 2L)), + new ListBasedInputRow(ROW_SIGNATURE, TIME.plusMinutes(2), ROW_SIGNATURE.getColumnNames(), Arrays.asList("b", 3L)), + new ListBasedInputRow(ROW_SIGNATURE, TIME.plusMinutes(3), ROW_SIGNATURE.getColumnNames(), Arrays.asList("b", 4L)) + ); + + @TempDir + static File SHARED_TEMP_DIR; + + private static File DEEP_STORAGE_DIR; + + @TempDir + File perTestTempDir; + + private ObjectMapper jsonMapper; + private File cacheRoot; + private SegmentLocalCacheManager manager; + + @BeforeAll + static void buildSegment() + { + final File tmp = new File(SHARED_TEMP_DIR, "build_" + ThreadLocalRandom.current().nextInt()); + DEEP_STORAGE_DIR = IndexBuilder.create() + .useV10() + .tmpDir(tmp) + .segmentWriteOutMediumFactory(OffHeapMemorySegmentWriteOutMediumFactory.instance()) + .schema( + IncrementalIndexSchema.builder() + .withDimensionsSpec( + DimensionsSpec.builder() + .setDimensions( + List.of( + new StringDimensionSchema("dim1"), + new LongDimensionSchema("metric1") + ) + ) + .build() + ) + .withRollup(false) + .withMinTimestamp(TIME.getMillis()) + .withProjections(PROJECTIONS) + .build() + ) + .indexSpec(IndexSpec.builder() + .withMetadataCompression(CompressionStrategy.NONE) + .build()) + .rows(ROWS) + .buildMMappedIndexFile(); + EmittingLogger.registerEmitter(new NoopServiceEmitter()); + } + + @BeforeEach + void setup() throws IOException + { + jsonMapper = TestHelper.makeJsonMapper(); + jsonMapper.registerSubtypes(new NamedType(LocalLoadSpec.class, "local")); + jsonMapper.registerSubtypes(new NamedType(PartialProjectionLoadSpec.class, PartialProjectionLoadSpec.TYPE)); + jsonMapper.registerModule(new SegmentizerModule()); + jsonMapper.registerModules(new LocalDataStorageDruidModule().getJacksonModules()); + jsonMapper.setInjectableValues( + new InjectableValues.Std() + .addValue(LocalDataSegmentPuller.class, new LocalDataSegmentPuller()) + .addValue(IndexIO.class, TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT)) + .addValue(ObjectMapper.class, jsonMapper) + .addValue(DataSegment.PruneSpecsHolder.class, DataSegment.PruneSpecsHolder.DEFAULT) + .addValue(ExprMacroTable.class, TestExprMacroTable.INSTANCE) + ); + + cacheRoot = new File(perTestTempDir, "cache_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)); + FileUtils.mkdirp(cacheRoot); + } + + @AfterEach + void tearDown() + { + if (manager != null) { + manager.shutdown(); + } + } + + private SegmentLocalCacheManager makeManager(boolean virtualStorage, boolean partialDownloadsEnabled) + { + return makeManagerAtLocations(virtualStorage, partialDownloadsEnabled, List.of(cacheRoot)); + } + + private SegmentLocalCacheManager makeManagerAtLocations( + boolean virtualStorage, + boolean partialDownloadsEnabled, + List locationRoots + ) + { + final List locConfigs = locationRoots.stream() + .map(root -> new StorageLocationConfig(root, 1024L * 1024L * 1024L, null)) + .toList(); + final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig() + .setLocations(locConfigs) + .setVirtualStorage(virtualStorage) + .setVirtualStoragePartialDownloadsEnabled(partialDownloadsEnabled); + final List storageLocations = loaderConfig.toStorageLocations(); + return new SegmentLocalCacheManager( + storageLocations, + loaderConfig, + StorageLoadingThreadPool.createFromConfig(loaderConfig), + new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations), + TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT), + jsonMapper + ); + } + + private DataSegment partialWrapperSegment(List selectedProjections) + { + return partialWrapperSegment(selectedProjections, FINGERPRINT); + } + + private DataSegment partialWrapperSegment(List selectedProjections, String fingerprint) + { + final Map delegate = Map.of( + "type", "local", + "path", DEEP_STORAGE_DIR.getAbsolutePath() + ); + final Map wrapperWire = + PartialProjectionLoadSpec.wireForm(delegate, selectedProjections, fingerprint); + return DataSegment.builder(SEGMENT_ID) + .shardSpec(NoneShardSpec.instance()) + .loadSpec(wrapperWire) + .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 (e.g. a + * historical that received a partial-load rule pointing at zipped deep storage). + */ + private DataSegment partialWrapperSegmentWithNullRangeReader(List selectedProjections, String fingerprint) + throws IOException + { + final File noRangeReaderDir = new File( + perTestTempDir, + "no_range_reader_" + fingerprint.replace(':', '_').replace('.', '_') + ); + FileUtils.mkdirp(noRangeReaderDir); + final Map delegate = Map.of( + "type", "local", + "path", noRangeReaderDir.getAbsolutePath() + ); + final Map wrapperWire = + PartialProjectionLoadSpec.wireForm(delegate, selectedProjections, fingerprint); + return DataSegment.builder(SEGMENT_ID) + .shardSpec(NoneShardSpec.instance()) + .loadSpec(wrapperWire) + .size(0) + .build(); + } + + /** + * Post-load lookup of the mounted partial metadata entry. loadPartial is synchronous — by the time load() returns, + * the entry is either fully mounted or the call threw; tests can inspect state immediately. + */ + private static PartialSegmentMetadataCacheEntry mountedMetadata(StorageLocation location, SegmentId segmentId) + { + final CacheEntry entry = location.getStaticCacheEntry(new SegmentCacheEntryIdentifier(segmentId)); + Assertions.assertInstanceOf(PartialSegmentMetadataCacheEntry.class, entry); + final PartialSegmentMetadataCacheEntry partial = (PartialSegmentMetadataCacheEntry) entry; + Assertions.assertTrue(partial.isMounted(), "metadata entry for " + segmentId + " should be mounted after load()"); + return partial; + } + + private static PartialSegmentBundleCacheEntry mountedBundle( + StorageLocation location, + SegmentId segmentId, + String bundleName + ) + { + final CacheEntry entry = location.getStaticCacheEntry( + new PartialSegmentBundleCacheEntryIdentifier(segmentId, bundleName) + ); + Assertions.assertInstanceOf(PartialSegmentBundleCacheEntry.class, entry); + final PartialSegmentBundleCacheEntry bundle = (PartialSegmentBundleCacheEntry) entry; + Assertions.assertTrue(bundle.isMounted(), "bundle " + bundleName + " should be mounted after load()"); + return bundle; + } + + @Test + void testLoadPinsMetadataAndSelectedBundleAsStatic() throws Exception + { + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE)); + + manager.load(segment); + + final PartialSegmentMetadataCacheEntry metadata = mountedMetadata(location, SEGMENT_ID); + final PartialSegmentBundleCacheEntry agg = mountedBundle(location, SEGMENT_ID, AGG_BUNDLE); + + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + final PartialSegmentBundleCacheEntryIdentifier aggId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); + Assertions.assertTrue(location.isReserved(metaId), "metadata entry should be in staticCacheEntries"); + Assertions.assertTrue(location.isReserved(aggId), "selected bundle should be in staticCacheEntries"); + Assertions.assertFalse(location.isWeakReserved(metaId), "metadata entry should NOT be in weakCacheEntries"); + Assertions.assertFalse(location.isWeakReserved(aggId), "selected bundle should NOT be in weakCacheEntries"); + + Assertions.assertFalse(agg.getContainerRefs().isEmpty(), "selected bundle must own at least one container"); + // staticBundleNames on the metadata entry is intentionally empty on the fresh-load path (loadPartial reserves + // statics directly after mount); the field communicates rule intent only across the bootstrap-restore boundary. + Assertions.assertTrue(metadata.getStaticBundleNames().isEmpty()); + } + + @Test + void testLoadDoesNotPreMountNonSelectedBundles() throws Exception + { + // Select one of the two aggregate projections. The selected one + its parent (__base) must be mounted; the other + // aggregate is neither selected nor a parent so it must stay unmounted. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE)); + + manager.load(segment); + mountedMetadata(location, SEGMENT_ID); + mountedBundle(location, SEGMENT_ID, AGG_BUNDLE); + mountedBundle(location, SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME); + + final PartialSegmentBundleCacheEntryIdentifier otherAggId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, OTHER_AGG_BUNDLE); + Assertions.assertNull( + location.getCacheEntry(otherAggId), + "non-selected, non-parent bundle must not be mounted at load time" + ); + } + + @Test + void testLoadEagerlyDownloadsSelectedBundleContainers() throws Exception + { + // With sync loadPartial, ensureBundleDownloaded has already run on the calling thread by the time load() returns. + // A second ensureBundleDownloaded call must be a pure no-op — no new bytes fetched from deep storage. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE)); + + manager.load(segment); + final PartialSegmentMetadataCacheEntry metadata = mountedMetadata(location, SEGMENT_ID); + mountedBundle(location, SEGMENT_ID, AGG_BUNDLE); + + final PartialSegmentFileMapperV10 fileMapper = metadata.getFileMapper(); + Assertions.assertNotNull(fileMapper, "metadata mount should produce a file mapper"); + final long downloadedBefore = fileMapper.getDownloadedBytes(); + fileMapper.ensureBundleDownloaded(AGG_BUNDLE); + Assertions.assertEquals( + downloadedBefore, + fileMapper.getDownloadedBytes(), + "ensureBundleDownloaded should be a no-op since loadPartial already downloaded the bundle eagerly" + ); + } + + @Test + void testLoadWithoutVirtualStorageFallsThroughToEagerFullLoad() throws Exception + { + // virtualStorage=false: the partial wrapper unwraps to its inner delegate (full download via CompleteSegmentCacheEntry, + // static). No partial machinery involved. + manager = makeManager(false, false); + final StorageLocation location = manager.getLocations().get(0); + final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE)); + + manager.load(segment); + + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + final CacheEntry entry = location.getCacheEntry(metaId); + Assertions.assertNotNull(entry, "complete-load static entry must exist"); + Assertions.assertFalse( + entry instanceof PartialSegmentMetadataCacheEntry, + "virtualStorage=false must NOT take the partial machinery path" + ); + Assertions.assertTrue(location.isReserved(metaId), "entry should be in staticCacheEntries"); + } + + @Test + void testLoadWithPartialsDisabledNoOps() throws Exception + { + // virtualStorage=true but partialDownloads=false: today's current behavior, load() is a noop and the segment will + // be weakly-loaded at query time via the existing PartialLoadSpec.loadSegment delegate-to-full fallback. + manager = makeManager(true, false); + final StorageLocation location = manager.getLocations().get(0); + final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE)); + + manager.load(segment); + + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + Assertions.assertNull( + location.getCacheEntry(metaId), + "with partial downloads disabled, load() should not eagerly reserve anything" + ); + } + + @Test + void testDropReleasesStaticMetadataAndBundles() throws Exception + { + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE)); + + manager.load(segment); + final PartialSegmentMetadataCacheEntry metadata = mountedMetadata(location, SEGMENT_ID); + mountedBundle(location, SEGMENT_ID, AGG_BUNDLE); + + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + final PartialSegmentBundleCacheEntryIdentifier aggId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); + Assertions.assertTrue(location.isReserved(metaId)); + Assertions.assertTrue(location.isReserved(aggId)); + + manager.drop(segment); + + Assertions.assertFalse(location.isReserved(metaId), "metadata entry should be released from staticCacheEntries"); + Assertions.assertFalse(location.isReserved(aggId), "selected bundle should be released from staticCacheEntries"); + Assertions.assertFalse(metadata.isMounted(), "metadata entry should be unmounted after drop cascade"); + + final File headerFile = new File( + cacheRoot, + SEGMENT_ID.toString() + "/" + IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX + ); + Assertions.assertFalse(headerFile.exists(), "header file should be deleted after drop"); + } + + @Test + void testSecondLoadWithSameFingerprintIsIdempotent() throws Exception + { + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE)); + + manager.load(segment); + final PartialSegmentMetadataCacheEntry firstEntry = mountedMetadata(location, SEGMENT_ID); + + // Second load() with the same fingerprint must not release + re-reserve the metadata; the operation is a no-op. + manager.load(segment); + final PartialSegmentMetadataCacheEntry secondEntry = mountedMetadata(location, SEGMENT_ID); + Assertions.assertSame(firstEntry, secondEntry, "same-fingerprint re-load should not create a new metadata entry"); + } + + @Test + void testSecondLoadWithChangedFingerprintReconcilesViaDropAndReload() throws Exception + { + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + + // First load: rule pins only the aggregate projection. + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original")); + final PartialSegmentMetadataCacheEntry firstEntry = mountedMetadata(location, SEGMENT_ID); + Assertions.assertEquals("v1:rule-original", firstEntry.getFingerprint()); + + // Second load: rule now pins the OTHER aggregate. Different fingerprint → reconciliation drops the first entry + // and reserves a fresh one with the new selection. + manager.load(partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), "v1:rule-updated")); + final PartialSegmentMetadataCacheEntry secondEntry = mountedMetadata(location, SEGMENT_ID); + Assertions.assertNotSame(firstEntry, secondEntry, "changed fingerprint should trigger a fresh metadata entry"); + Assertions.assertEquals("v1:rule-updated", secondEntry.getFingerprint()); + Assertions.assertFalse(firstEntry.isMounted(), "old entry should be unmounted after reconciliation"); + + // The old rule's bundle is no longer pinned; the new rule's bundle is. + final PartialSegmentBundleCacheEntryIdentifier oldBundleId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); + final PartialSegmentBundleCacheEntryIdentifier newBundleId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, OTHER_AGG_BUNDLE); + Assertions.assertNull(location.getCacheEntry(oldBundleId), "old rule's bundle should be released"); + Assertions.assertTrue(location.isReserved(newBundleId), "new rule's bundle should be statically pinned"); + } + + @Test + void testRangeReaderNullDropsStaleStaticEntryOnFingerprintChange() throws Exception + { + // Regression guard: without range reads a partial load cannot be honored, but a stale static entry from a prior + // rule must still be reconciled — otherwise the historical silently keeps serving the OLD rule after the + // coordinator changed the rule fingerprint. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + // First load: normal wrapper, real range reader, static entry installed under fingerprint v1. + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original")); + mountedMetadata(location, SEGMENT_ID); + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + Assertions.assertTrue(location.isReserved(metaId)); + + // Second load: same segment, different rule fingerprint, but the delegate now points at a directory without a V10 + // file so LocalLoadSpec.openRangeReader() returns null. The rule can't be applied — but the stale static entry + // for the OLD rule must be dropped so the coordinator's rule change isn't silently ignored. + manager.load(partialWrapperSegmentWithNullRangeReader(List.of(AGG_BUNDLE), "v2:rule-updated")); + + Assertions.assertFalse( + location.isReserved(metaId), + "stale static partial entry for the old rule should be dropped when the new rule cannot be applied" + ); + } + + @Test + void testRangeReaderNullWithMatchingFingerprintPreservesEntry() throws Exception + { + // The rangeReader-null path must still short-circuit on a same-fingerprint re-issue — otherwise a transient + // range-reader failure on a coordinator retry would drop a still-valid static entry. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original")); + final PartialSegmentMetadataCacheEntry first = mountedMetadata(location, SEGMENT_ID); + + manager.load(partialWrapperSegmentWithNullRangeReader(List.of(AGG_BUNDLE), "v1:rule-original")); + + final PartialSegmentMetadataCacheEntry second = mountedMetadata(location, SEGMENT_ID); + Assertions.assertSame( + first, + second, + "same-fingerprint re-issue on the rangeReader-null path should preserve the existing entry" + ); + } + + @Test + void testReconciliationNukesLeftoverPartialDirectoryAcrossLocations() throws Exception + { + // Regression guard: reconciliation must move-and-delete stale per-segment directories at every location before + // the new reserve picks a location. Without this, an initial load on location A followed by a reconciliation + // whose new reserve lands on location B leaves A's now-empty partial dir on disk indefinitely — drop()'s cascade + // only deletes tracked files (header via metadata unmount, containers via bundle eviction), not the enclosing dir. + final File locationA = new File(perTestTempDir, "loc_a_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)); + final File locationB = new File(perTestTempDir, "loc_b_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)); + FileUtils.mkdirp(locationA); + FileUtils.mkdirp(locationB); + manager = makeManagerAtLocations(true, true, List.of(locationA, locationB)); + + // First load: LeastBytesUsedStorageLocationSelectorStrategy picks locationA (both empty, first wins). Verify A + // received the segment dir, and neither the __drop dir nor B's segment dir exist yet. + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original")); + final File aSegDir = new File(locationA, SEGMENT_ID.toString()); + final File bSegDir = new File(locationB, SEGMENT_ID.toString()); + Assertions.assertTrue(aSegDir.isDirectory(), "first load must create the segment dir on locationA"); + Assertions.assertFalse(bSegDir.exists(), "locationB should be untouched by the first load"); + Assertions.assertFalse(new File(locationA, "__drop").exists()); + Assertions.assertFalse(new File(locationB, "__drop").exists()); + + // Second load: same segment, different fingerprint. Reconciliation drops the entry on locationA, then the nuke + // step atomicMoveAndDeletes A's leftover dir. The new reserve picks the least-used location — which is A again + // in this configuration, so mkdirp recreates A/segId fresh. The observable signal that the nuke ran is that A's + // __drop directory now exists (atomicMoveAndDeleteCacheEntryDirectory creates it lazily on first use). + manager.load(partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), "v2:rule-updated")); + Assertions.assertTrue( + new File(locationA, "__drop").exists(), + "reconciliation must have moved the stale dir into locationA's __drop dir, proving the nuke was invoked" + ); + // New entry present at whichever location the strategy picked — verify the segment is loaded somewhere. + Assertions.assertTrue( + aSegDir.exists() || bSegDir.exists(), + "reconciliation must have installed the new entry on one of the locations" + ); + } + + @Test + void testReserveMkdirpFailureFallsThroughToNextLocation() throws Exception + { + // Regression guard: a per-location mkdirp failure (inode exhaustion, permission drop, transient EIO on one disk) + // must not abort the whole load when another location would accept the segment. Setup: + // - writable first in list so info_dir defaults to writable/info_dir (write-tolerant path). + // - Pre-fill writable with a dummy reservation so LeastBytesUsed picks the read-only location FIRST. + // - Read-only location: mkdirp of the per-segment subdir throws IOException → release the reservation and + // fall through to the writable location → mount lands there. + final File writable = new File(perTestTempDir, "loc_rw_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)); + final File readOnly = new File(perTestTempDir, "loc_ro_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)); + FileUtils.mkdirp(writable); + FileUtils.mkdirp(readOnly); + manager = makeManagerAtLocations(true, true, List.of(writable, readOnly)); + + // Bump writable's currentSize above 0 so LeastBytesUsed iterates readOnly (usage 0) first. + final SegmentId dummyId = SegmentId.of("dummy", Intervals.of("2020/2021"), "v", 0); + Assertions.assertTrue( + manager.getLocations().get(0).reserveWeak(stubCacheEntry(new SegmentCacheEntryIdentifier(dummyId), 4096L)), + "dummy pre-fill on writable must succeed" + ); + + Assertions.assertTrue(readOnly.setReadOnly(), "test setup must be able to make readOnly location read-only"); + try { + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + + final File roSegDir = new File(readOnly, SEGMENT_ID.toString()); + final File rwSegDir = new File(writable, SEGMENT_ID.toString()); + Assertions.assertFalse(roSegDir.exists(), "read-only location must NOT receive the segment"); + Assertions.assertTrue(rwSegDir.isDirectory(), "load must fall through to the writable location"); + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + Assertions.assertTrue( + manager.getLocations().get(0).isReserved(metaId), + "static metadata reservation must land on the writable location after readOnly's mkdirp fails" + ); + Assertions.assertFalse( + manager.getLocations().get(1).isReserved(metaId), + "the reservation attempted on the readOnly location must have been released on mkdirp failure" + ); + } + finally { + Assertions.assertTrue(readOnly.setWritable(true), "test teardown must restore write permission"); + } + } + + @Test + void testLoadThrowsWhenOldStaticEntryHasOutstandingReferences() throws Exception + { + // Regression guard for the deferred-cleanup race: a running query holding a metadata reference on the OLD static + // partial would defer its doActualUnmount past drop(). Its onUnmount hook (deleteSegmentInfoFile) plus + // deleteHeaderFiles would fire later — after the new entry's storeInfoFile + mount have written fresh on-disk + // state — and corrupt the new entry (segment ID + partial dir are shared). Fail the load explicitly BEFORE drop + // so the old entry stays fully installed for a subsequent coordinator retry. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original")); + final PartialSegmentMetadataCacheEntry oldEntry = mountedMetadata(location, SEGMENT_ID); + + // Pin the old entry with an outstanding metadata reference (a running query would hold one of these via a Segment). + final java.io.Closeable heldRef = oldEntry.acquireMetadataReference(); + try { + final SegmentLoadingException thrown = Assertions.assertThrows( + SegmentLoadingException.class, + () -> manager.load(partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), "v1:rule-updated")) + ); + Assertions.assertTrue( + thrown.getMessage().contains("outstanding references"), + "unexpected message: " + thrown.getMessage() + ); + // Old entry survives the failed reconciliation WITHOUT dropping — the precheck refused to touch static state, + // so the old entry is fully installed and mounted, not just deferred-mounted. + Assertions.assertTrue( + oldEntry.isMounted(), + "held reference should keep the old entry mounted" + ); + Assertions.assertTrue( + location.isReserved(metaId), + "reconciliation must not remove the old static entry when it refuses to proceed" + ); + } + finally { + heldRef.close(); + } + // The old entry is still fully installed (no drop happened). Once the external reference releases, retry + // succeeds — the precheck sees no external references and proceeds with drop+reserve as usual. + Assertions.assertTrue(oldEntry.isMounted(), "old entry stays mounted between failed reconciliation and retry"); + manager.load(partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), "v1:rule-updated")); + final PartialSegmentMetadataCacheEntry newEntry = mountedMetadata(location, SEGMENT_ID); + Assertions.assertNotSame(oldEntry, newEntry, "retry after reference release should install a fresh entry"); + Assertions.assertEquals("v1:rule-updated", newEntry.getFingerprint()); + Assertions.assertFalse(oldEntry.isMounted(), "old entry finally unmounts as part of the successful retry drop"); + } + + @Test + void testLoadThrowsWhenOldStaticBundleHasOutstandingReferences() throws Exception + { + // Regression guard: outstanding references on a LINKED BUNDLE (not just the metadata directly) also defer + // cleanup past drop(). This is caught transitively because PartialSegmentBundleCacheEntry.doMount holds a + // metadataEntry.acquireMetadataReference for the bundle's lifetime — so while any bundle has outstanding + // references, metadata.isMounted() still returns true and the reconciliation guard fires. This test proves + // the transitive relationship survives a code change that decouples bundle-and-metadata refs. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original")); + final PartialSegmentMetadataCacheEntry oldMeta = mountedMetadata(location, SEGMENT_ID); + final PartialSegmentBundleCacheEntry oldAgg = mountedBundle(location, SEGMENT_ID, AGG_BUNDLE); + + // Pin the BUNDLE only (not the metadata directly). A running query walking the aggregate projection would hold + // this reference via its Segment; the mount-time metadata ref carried by the bundle keeps the metadata alive + // transitively. + final java.io.Closeable heldBundleRef = oldAgg.acquireReference(); + try { + final SegmentLoadingException thrown = Assertions.assertThrows( + SegmentLoadingException.class, + () -> manager.load(partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), "v2:rule-updated")) + ); + Assertions.assertTrue( + thrown.getMessage().contains("outstanding references"), + "unexpected message: " + thrown.getMessage() + ); + // Both the pinned bundle and its metadata still report mounted (both cleanup paths are deferred). + Assertions.assertTrue(oldAgg.isMounted(), "held bundle reference must keep bundle mounted"); + Assertions.assertTrue( + oldMeta.isMounted(), + "bundle's mount-time metadata reference must keep metadata mounted transitively" + ); + } + finally { + heldBundleRef.close(); + } + // Once the bundle reference releases, its dependency-ref release chains into the metadata's cleanup. Retry. + manager.load(partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), "v2:rule-updated")); + final PartialSegmentMetadataCacheEntry newEntry = mountedMetadata(location, SEGMENT_ID); + Assertions.assertNotSame(oldMeta, newEntry, "retry after bundle reference release should install a fresh entry"); + Assertions.assertEquals("v2:rule-updated", newEntry.getFingerprint()); + } + + @Test + void testLoadEvictsPreExistingUnheldWeakEntry() throws Exception + { + // Regression guard for the "stale weak entry" leak: a weak reservation for this segment ID (e.g., installed by + // bootstrap when no rule fingerprint was persisted, or by a prior on-demand acquireSegment) must be evicted + // before loadPartial reserves the new static entry. drop() only clears staticCacheEntries, so without explicit + // weak-eviction the two would coexist on the same on-disk directory. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(SEGMENT_ID); + Assertions.assertTrue(location.reserveWeak(stubCacheEntry(id, 1024L))); + Assertions.assertTrue(location.isWeakReserved(id), "precondition: weak entry should be installed"); + + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + + Assertions.assertFalse( + location.isWeakReserved(id), + "reconciliation must evict the pre-existing weak entry before installing the static partial" + ); + Assertions.assertTrue(location.isReserved(id), "static partial entry should be installed after reconciliation"); + mountedMetadata(location, SEGMENT_ID); + } + + @Test + void testLoadThrowsWhenPreExistingWeakEntryHasActiveHold() throws Exception + { + // A held weak entry (a running query) cannot be evicted synchronously without breaking that query's on-disk + // reads. Fail the load explicitly so the coordinator retries once the query completes, rather than stacking a + // static reservation on top of a live weak entry. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(SEGMENT_ID); + final CacheEntry stub = stubCacheEntry(id, 1024L); + final StorageLocation.ReservationHold hold = location.addWeakReservationHold(id, () -> stub); + Assertions.assertNotNull(hold, "precondition: weak reservation with hold should succeed"); + try { + final SegmentLoadingException thrown = Assertions.assertThrows( + SegmentLoadingException.class, + () -> manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))) + ); + Assertions.assertTrue( + thrown.getMessage().contains("active holds"), + "unexpected message: " + thrown.getMessage() + ); + Assertions.assertTrue( + location.isWeakReserved(id), + "held weak entry must survive the failed reconciliation attempt" + ); + Assertions.assertFalse( + location.isReserved(id), + "no static partial entry should be installed alongside a held weak entry" + ); + } + finally { + hold.close(); + } + } + + private static CacheEntry stubCacheEntry(CacheEntryIdentifier id, long size) + { + return new CacheEntry() + { + @Override + public CacheEntryIdentifier getId() + { + return id; + } + + @Override + public long getSize() + { + return size; + } + + @Override + public boolean isMounted() + { + return false; + } + + @Override + public void mount(StorageLocation location) + { + } + + @Override + public void unmount() + { + } + }; + } + + @Test + void testLoadFailureCascadeReleasesEverything() throws Exception + { + // Defensive check exercise: a wire form referring to a projection this segment doesn't have would only happen + // on matcher/reader drift or a coding bug, but the sync loadPartial flow must still cascade-release everything + // and let load() throw so the failure is visible instead of leaving a silent zombie entry. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + final DataSegment segment = partialWrapperSegment(List.of("does_not_exist")); + + final SegmentLoadingException thrown = Assertions.assertThrows( + SegmentLoadingException.class, + () -> manager.load(segment) + ); + Assertions.assertTrue( + thrown.getMessage().contains("Failed to load partial segment"), + "unexpected message: " + thrown.getMessage() + ); + // The underlying cause carries the specific validation failure — a rule projection name not in this segment. + Assertions.assertNotNull(thrown.getCause(), "expected an underlying cause"); + Assertions.assertTrue( + thrown.getCause().getMessage().contains("does not contain projection[does_not_exist]"), + "unexpected cause message: " + thrown.getCause().getMessage() + ); + + final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); + Assertions.assertFalse( + location.isReserved(metaId), + "metadata reservation should be rolled back after a rule failure" + ); + Assertions.assertFalse( + location.isWeakReserved(metaId), + "no weak fallback should be left behind either" + ); + // The info file was written before the async work started; the metadata's onUnmount hook should have deleted it + // as part of the cascade release. + final File infoFile = new File(new File(cacheRoot, "info_dir"), SEGMENT_ID.toString()); + Assertions.assertFalse(infoFile.exists(), "info file should be deleted by the metadata's onUnmount hook"); + } +} From 2b877376e05674dadaba87d998b306e69956c3e9 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Mon, 6 Jul 2026 11:56:20 -0700 Subject: [PATCH 2/3] tidy up --- .../file/PartialSegmentFileMapperV10.java | 2 - .../loading/SegmentLocalCacheManager.java | 99 +++++++------------ .../PartialSegmentCacheBootstrapTest.java | 6 +- .../PartialSegmentMetadataCacheEntryTest.java | 6 +- ...tLocalCacheManagerPartialRuleLoadTest.java | 49 ++++----- 5 files changed, 63 insertions(+), 99 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java b/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java index 62d1b964b147..79820bf30e9d 100644 --- a/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java +++ b/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java @@ -919,8 +919,6 @@ private static void fetchAndPersistHeader( /** * Parse metadata from the header file. Throws on any parse failure (corrupt JSON, truncated header, etc.). - * Public so external callers (e.g. the partial-load-rule bootstrap path in {@code SegmentLocalCacheManager}) can - * share the header-parse convention with this class rather than re-implementing FileInputStream + reader plumbing. */ public static SegmentFileMetadataReader.Result parseHeaderFile( File headerFile, 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 eac1a777c177..a3d501811749 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 @@ -928,9 +928,8 @@ private record ReservedPartial( /** * Handle the partial-load-rule case at {@link #load} time: reserve a {@link PartialSegmentMetadataCacheEntry} as * static, persist the info file, mount the metadata header, and reserve+mount each rule-selected bundle as static - * with its data eagerly downloaded. Static reservations are not subject to SIEVE eviction, so the rule-pinned - * subset stays resident until the segment is explicitly dropped; non-selected bundles are downloaded on-demand at - * query time via the existing weak path. + * with its data eagerly downloaded (non-selected bundles are downloaded on-demand at query time via the existing + * weak path). *

* The entire flow runs synchronously on {@link #virtualStorageLoadingThreadPool} (via submit + get) so it respects * the pool's concurrency cap and virtual-thread runtime while giving @@ -973,7 +972,7 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException if (rangeReader == null) { // Backend doesn't support range reads (e.g. zipped deep storage). The rule can't be honored as a partial load; // fall through to the on-demand weak-full-load path at query time. But if a prior loadPartial installed a - // static entry with a different rule fingerprint, we cannot silently keep serving under the old rule — + // static entry with a different rule fingerprint, we cannot silently keep serving under the old rule so // reconcile the static side so the coordinator's rule change takes effect, then return so the next query // installs a weak entry via the on-demand path. log.warn( @@ -987,7 +986,7 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException synchronized (lock) { try { // Static-side reconciliation only. Weak entries (from bootstrap without a rule, or a prior on-demand - // acquire) already represent the fallback behavior we're falling back to — leave them alone. + // acquire) already represent the fallback behavior we're falling back to, leave them alone. boolean dropStaleStatic = false; for (StorageLocation location : locations) { if (!location.isReserved(id)) { @@ -1011,7 +1010,7 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException drop(dataSegment); // drop() only deleted tracked files; no reserve follows on this path, so nuke leftover per-segment dirs // across every location or an empty dir would linger indefinitely. - nukeStaleSegmentDirsOnAllLocations(dataSegment); + removeStaleSegmentDirsOnAllLocations(dataSegment); } } finally { @@ -1038,7 +1037,7 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException // // The weak eviction step is critical: {@link #drop} only touches {@link StorageLocation#staticCacheEntries}, // so without it a lingering weak entry from bootstrap or on-demand load would coexist with the new static - // entry — same on-disk directory, both counted against capacity, and the eventual bundle downloads would + // entry, same on-disk directory, both counted against capacity, and the eventual bundle downloads would // race the weak entry's own reads. When the weak entry has active holds we cannot evict it synchronously, // so throw and let the coordinator retry once the running query completes. drop() takes no per-segment lock // so calling it from inside our synchronized(lock) block is safe. @@ -1072,10 +1071,6 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException // path, and race the deferred cleanup on the shared per-segment on-disk directory. Checking BEFORE drop // keeps the old entry fully installed on the failure path so a subsequent retry finds it via // reconciliation logic as if this load never happened. - // - // cascadeReleaseWouldDeferCleanup subtracts the known cross-entry structure (bundle→metadata refs, - // parent-bundle→child holds) from raw reference counts to isolate external consumers; those are the only - // refs drop's cascade cannot release. for (SegmentCacheEntry old : oldStaticEntries) { if (old instanceof PartialSegmentMetadataCacheEntry oldPartial && oldPartial.cascadeReleaseWouldDeferCleanup()) { @@ -1107,16 +1102,11 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException } // After any reconciliation drop/eviction above, nuke leftover per-segment directories on every location. - // The drop cascade deleted TRACKED files (header via metadata's doActualUnmount, containers via each bundle's - // evictContainer), but not the enclosing per-segment directory nor any orphan files. If the new reserve below - // lands on a DIFFERENT location than the dropped entry, the old location's now-empty dir would accumulate on - // disk indefinitely; if the new reserve lands on the SAME location, the mkdirp inside the reserve loop - // recreates the dir fresh. - nukeStaleSegmentDirsOnAllLocations(dataSegment); + removeStaleSegmentDirsOnAllLocations(dataSegment); // Reserve the metadata entry as static on the first location that accepts BOTH the reservation and a // per-segment mkdirp. mkdirp failures are location-local (inode exhaustion, permission mismatch, transient - // EIO on one disk) — a different location may accept, so release-and-continue rather than abort. Collect + // EIO on one disk) and a different location may accept, so release-and-continue rather than abort. Collect // per-location failures to attach as suppressed exceptions on the final alert. StorageLocation reservedLocation = null; PartialSegmentMetadataCacheEntry metadata = null; @@ -1178,9 +1168,7 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException throw noCapacity; } - // Persist the info file so bootstrap can rediscover this segment on restart. The info dir is shared across - // locations (single configured path), so failure here is not a per-disk problem — no fall-through, alert - // and abort. + // Persist the info file so bootstrap can rediscover this segment on restart try { storeInfoFile(dataSegment); } @@ -1200,19 +1188,12 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException } metadata.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); - // Mount metadata + reserve+mount+download each rule-selected bundle as static. Runs on the loading thread pool - // (submit + get) so the flow respects the pool's concurrency cap while giving the calling thread a truthful - // success/failure signal. Any exception cascade-releases the metadata + every bundle reserved so far, so no - // orphan static reservation survives the load call, the reservation's onUnmount hook then deletes the info - // file, mirroring drop() semantics. + // Mount metadata + reserve+mount+download each rule-selected bundle as static final StorageLocation theLocation = reservedLocation; final PartialSegmentMetadataCacheEntry theMetadata = metadata; // All rollback responsibility (bundles AND metadata) lives INSIDE the callable so no state mutation of - // theMetadata crosses threads. This avoids the FutureTask cancel-vs-cleanup race: cancel(true) transitions - // the future to CANCELLED immediately, so a subsequent get() returns CancellationException without waiting - // for the callable's cleanup to unwind — if the caller then released theMetadata itself, that release would - // race the still-running pool-thread rollback. With cleanup owned by the pool thread, the caller only ever - // observes the outcome; there is no shared state to race. + // theMetadata crosses threads. With cleanup owned by the pool thread, the caller only ever observes the + // outcome; there is no shared state to race. final Future future = virtualStorageLoadingThreadPool.getExecutorService().submit(() -> { final List mounted = new ArrayList<>(); try { @@ -1228,9 +1209,8 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException dataSegment, mapper.getSegmentFileMetadata() ); - // A rule-pinned bundle's dependencies must also be pinned: the bundle's own mount call acquires holds on - // each dependency via addWeakReservationHoldIfExists, which requires the dependency to be registered. - // Iterate in dependency-first order so children see their dependencies resident. + // A rule-pinned bundle's dependencies must also be pinned. Iterate in dependency-first order so children + // see their dependencies resident. for (String bundleName : theMetadata.bundlesInMountOrder(selectedBundleNames)) { mounted.add(reserveAndMountStaticBundle(theLocation, theMetadata, mapper, bundleName)); } @@ -1262,22 +1242,17 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException catch (InterruptedException e) { // Caller interrupted while waiting. Restore the interrupt flag and propagate, but do NOT cancel the pool // task: FutureTask.cancel(true) transitions to CANCELLED synchronously, but the callable's run() may still - // be executing its catch/finally on the pool thread — so a subsequent metadata release from the caller - // would race that cleanup. Leaving the task alone lets its own cleanup complete on the pool thread - // naturally; on eventual shutdown the pool's shutdownNow will interrupt the worker and drive the same - // rollback via the callable's catch. The coordinator's retry will find the segment in one of three - // consistent end-states: fully mounted (task raced to success — reconciliation short-circuits), fully - // rolled-back (task failed after interrupt — fresh reserve), or still in flight (retry throws until the - // outstanding-references guard sees the task's own metadata release fire). + // be executing its catch/finally on the pool thread, so a subsequent metadata release from the caller would + // race that cleanup. Leaving the task alone lets its own cleanup complete on the pool thread naturally; on + // eventual shutdown the pool's shutdownNow will interrupt the worker and drive the same + // rollback via the callable's catch. Thread.currentThread().interrupt(); throw new SegmentLoadingException(e, "Interrupted while loading partial segment[%s]", dataSegment.getId()); } catch (Throwable t) { // Callable ran to a terminal state. On failure it already released metadata + bundles on the pool thread - // via its catch block. We only propagate the diagnostic here; no cleanup remains for the caller. Unwrap - // ExecutionException for a clean cause; handle a null cause defensively (rare, but permitted by - // ExecutionException's contract). Restore the interrupt flag if the pool worker was interrupted mid-run - // so downstream shutdown handling still observes it. + // via its catch block. We only propagate the diagnostic here; no cleanup remains for the caller. Restore the + // interrupt flag if set so downstream shutdown handling still observes it. final Throwable executionCause = (t instanceof ExecutionException) ? t.getCause() : t; final Throwable cause = (executionCause != null) ? executionCause : t; if (cause instanceof InterruptedException) { @@ -1366,7 +1341,7 @@ private void releaseMetadataQuietly( * linger when the new reserve happens to land elsewhere. {@link #atomicMoveAndDeleteCacheEntryDirectory} is a no-op * on missing paths, so this walks all locations unconditionally. */ - private void nukeStaleSegmentDirsOnAllLocations(DataSegment dataSegment) + private void removeStaleSegmentDirsOnAllLocations(DataSegment dataSegment) { for (StorageLocation loc : locations) { atomicMoveAndDeleteCacheEntryDirectory(new File(loc.getPath(), dataSegment.getId().toString())); @@ -1428,17 +1403,16 @@ public void load(final DataSegment dataSegment) throws SegmentLoadingException "load() should not be called when virtualStorageIsEphemeral is true" ); } - // Partial-load rule routing: a wrapped load spec carrying a fingerprint + per-spec selection (cluster groups, - // projections, ...) means the coordinator wants this segment loaded as a sticky partial. Reserve the metadata - // entry as STATIC (never SIEVE-evicted) and eagerly download the rule-selected bundles as STATIC. Other bundles - // are left for the existing weak/on-demand path at query time. + // a wrapped load spec carrying a fingerprint + per-spec selection means the coordinator wants this segment + // loaded as a sticky partial. Reserve the metadata entry as static and eagerly download the rule-selected + // bundles as static. Other bundles are left for the existing weak/on-demand path at query time. if (config.isVirtualStoragePartialDownloadsEnabled() && PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec())) { loadPartial(dataSegment); return; } - // 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 + // .. otherwise, 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 final ReferenceCountingLock lock = lock(dataSegment); synchronized (lock) { try { @@ -1503,12 +1477,7 @@ public void bootstrap( partial.mount(location); } catch (IOException e) { - // A mount failure on a bootstrap-restored partial entry must release the reservation so it doesn't - // leak. doMount's own catch already calls removeUnheldWeakEntry for the weak variant; here we - // additionally call location.release which handles the static variant introduced by the partial-load- - // rule path (release is a no-op for weak entries but removes static ones from staticCacheEntries). - // Without this, a static rule-pinned segment whose header parse fails at mount time would occupy the - // metadata reservation forever. + // release the reservation so it doesn't leak try { location.release(partial); } @@ -1525,11 +1494,10 @@ public void bootstrap( dataSegment.getId() ); } - // Post-mount rule-integrity check: if any direct rule-pinned bundle didn't make it into the linked set - // (JVM crash mid-loadPartial, or orphan cleanup deleted a bundle whose parent was missing), the segment - // isn't fully restored to the rule's intent. Drop it entirely so the coordinator's reconciliation re- - // issues load with a clean cold-fetch — operators don't drive drop/load directly, so a self-healing - // pathway through the coordinator is the right recovery. + // if any direct rule-pinned bundle didn't make it into the linked set (JVM crash mid-loadPartial, + // or orphan cleanup deleted a bundle whose parent was missing), the segment isn't fully restored to the + // rule's intent. Drop it entirely so the coordinator's reconciliation re-issues load with a clean + // cold-fetch. final Set directRulePins = partial.getStaticBundleNames(); if (!directRulePins.isEmpty()) { final Set mounted = new HashSet<>(); @@ -1631,9 +1599,8 @@ public void drop(final DataSegment segment) continue; } // Cascade-release linked bundle entries before the metadata. {@link StorageLocation#release} only operates on - // entries in {@code staticCacheEntries}, so this is a no-op for weak/on-demand bundles (they continue to ride - // SIEVE eviction + hold-release cleanup); for static rule-pinned bundles installed by {@link #loadPartial}, - // this is the only path that removes them from the static map. + // entries in {@code staticCacheEntries}, so this is a no-op for weak/on-demand bundles; for static rule-pinned + // bundles installed by {@link #loadPartial}, this is the only path that removes them from the static map. if (entry instanceof PartialSegmentMetadataCacheEntry partial) { for (PartialSegmentBundleCacheEntry bundle : partial.snapshotLinkedBundles()) { location.release(bundle); diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java index 6d2b725ad4aa..328b68841eb7 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java @@ -701,10 +701,9 @@ void testStaticRestoreSkipsRulePinnedBundleAbsentOnDisk() throws IOException @Test void testStaticRestoreReleasesReservationsWhenBundleMountFails() throws IOException { - // Regression guard: when a static bundle's mount() throws after its reservation was taken via location.reserve(), + // when a static bundle's mount() throws after its reservation was taken via location.reserve(), // the reservation must be released. Also — any static bundle mounted successfully earlier in the loop must be - // released too; the outer rollback previously only called bundle.unmount(), which doesn't remove from - // staticCacheEntries. Without both, mount-time failures leave orphan static reservations after bootstrap. + // released too; primeOnDiskState(); // Identify the aggregate bundle's exclusive containers (those NOT shared with __base), and replace one with a @@ -753,7 +752,6 @@ void testStaticRestoreReleasesReservationsWhenBundleMountFails() throws IOExcept final PartialSegmentBundleCacheEntryIdentifier aggId = new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); - // The specific finding: agg's reservation was released by the inline catch after mount threw. Assertions.assertFalse(location.isReserved(aggId), "failed static bundle mount must release its own reservation"); // Broader cleanup: the previously-succeeded base bundle in mountedBundles is also released by the outer rollback, // not left as an orphan (bundle.unmount() alone would only clear the mount state, not staticCacheEntries). diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java index 3dd033734237..585fbd5e02a5 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java @@ -223,9 +223,9 @@ void testOnUnmountHookRunsAfterStorageLocationCleanup() throws IOException @Test void testOnUnmountHookRunsOnReleaseBeforeMount() { - // Regression guard: the SegmentLocalCacheManager.loadPartial flow persists the segment info file BEFORE mounting - // the entry, and registers a hook that deletes the info file on unmount. If mount then fails (or is never called) - // and the reservation is released, unmount() must still fire the hook — otherwise the info file leaks on disk. + // the SegmentLocalCacheManager.loadPartial flow persists the segment info file BEFORE mounting the entry, and + // registers a hook that deletes the info file on unmount. If mount then fails (or is never called) and the + // reservation is released, unmount() must still fire the hook — otherwise the info file leaks on disk. final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE); final AtomicReference hookFired = new AtomicReference<>(false); entry.setOnUnmount(() -> hookFired.set(true)); 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 238a9f4e3ad6..16b3d599fcd2 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 @@ -62,6 +62,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import java.io.Closeable; import java.io.File; import java.io.IOException; import java.util.Arrays; @@ -476,9 +477,9 @@ void testSecondLoadWithChangedFingerprintReconcilesViaDropAndReload() throws Exc @Test void testRangeReaderNullDropsStaleStaticEntryOnFingerprintChange() throws Exception { - // Regression guard: without range reads a partial load cannot be honored, but a stale static entry from a prior - // rule must still be reconciled — otherwise the historical silently keeps serving the OLD rule after the - // coordinator changed the rule fingerprint. + // without range reads a partial load cannot be honored, but a stale static entry from a prior rule must still + // be reconciled — otherwise the historical silently keeps serving the OLD rule after the coordinator changed the + // rule fingerprint. manager = makeManager(true, true); final StorageLocation location = manager.getLocations().get(0); // First load: normal wrapper, real range reader, static entry installed under fingerprint v1. @@ -521,10 +522,10 @@ void testRangeReaderNullWithMatchingFingerprintPreservesEntry() throws Exception @Test void testReconciliationNukesLeftoverPartialDirectoryAcrossLocations() throws Exception { - // Regression guard: reconciliation must move-and-delete stale per-segment directories at every location before - // the new reserve picks a location. Without this, an initial load on location A followed by a reconciliation - // whose new reserve lands on location B leaves A's now-empty partial dir on disk indefinitely — drop()'s cascade - // only deletes tracked files (header via metadata unmount, containers via bundle eviction), not the enclosing dir. + // reconciliation must move-and-delete stale per-segment directories at every location before the new reserve picks + // a location. Without this, an initial load on location A followed by a reconciliation whose new reserve lands on + // location B leaves A's now-empty partial dir on disk indefinitely — drop()'s cascade only deletes tracked files + // (header via metadata unmount, containers via bundle eviction), not the enclosing dir. final File locationA = new File(perTestTempDir, "loc_a_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)); final File locationB = new File(perTestTempDir, "loc_b_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)); FileUtils.mkdirp(locationA); @@ -560,8 +561,8 @@ void testReconciliationNukesLeftoverPartialDirectoryAcrossLocations() throws Exc @Test void testReserveMkdirpFailureFallsThroughToNextLocation() throws Exception { - // Regression guard: a per-location mkdirp failure (inode exhaustion, permission drop, transient EIO on one disk) - // must not abort the whole load when another location would accept the segment. Setup: + // a per-location mkdirp failure (inode exhaustion, permission drop, transient EIO on one disk) must not abort the + // whole load when another location would accept the segment. Setup: // - writable first in list so info_dir defaults to writable/info_dir (write-tolerant path). // - Pre-fill writable with a dummy reservation so LeastBytesUsed picks the read-only location FIRST. // - Read-only location: mkdirp of the per-segment subdir throws IOException → release the reservation and @@ -605,11 +606,11 @@ void testReserveMkdirpFailureFallsThroughToNextLocation() throws Exception @Test void testLoadThrowsWhenOldStaticEntryHasOutstandingReferences() throws Exception { - // Regression guard for the deferred-cleanup race: a running query holding a metadata reference on the OLD static - // partial would defer its doActualUnmount past drop(). Its onUnmount hook (deleteSegmentInfoFile) plus - // deleteHeaderFiles would fire later — after the new entry's storeInfoFile + mount have written fresh on-disk - // state — and corrupt the new entry (segment ID + partial dir are shared). Fail the load explicitly BEFORE drop - // so the old entry stays fully installed for a subsequent coordinator retry. + // a running query holding a metadata reference on the OLD static partial would defer its doActualUnmount past + // drop(). Its onUnmount hook (deleteSegmentInfoFile) plus deleteHeaderFiles would fire later — after the new + // entry's storeInfoFile + mount have written fresh on-disk state — and corrupt the new entry + // (segment ID + partial dir are shared). Fail the load explicitly BEFORE drop so the old entry stays fully + // installed for a subsequent coordinator retry. manager = makeManager(true, true); final StorageLocation location = manager.getLocations().get(0); final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID); @@ -617,7 +618,7 @@ void testLoadThrowsWhenOldStaticEntryHasOutstandingReferences() throws Exception final PartialSegmentMetadataCacheEntry oldEntry = mountedMetadata(location, SEGMENT_ID); // Pin the old entry with an outstanding metadata reference (a running query would hold one of these via a Segment). - final java.io.Closeable heldRef = oldEntry.acquireMetadataReference(); + final Closeable heldRef = oldEntry.acquireMetadataReference(); try { final SegmentLoadingException thrown = Assertions.assertThrows( SegmentLoadingException.class, @@ -654,11 +655,11 @@ void testLoadThrowsWhenOldStaticEntryHasOutstandingReferences() throws Exception @Test void testLoadThrowsWhenOldStaticBundleHasOutstandingReferences() throws Exception { - // Regression guard: outstanding references on a LINKED BUNDLE (not just the metadata directly) also defer - // cleanup past drop(). This is caught transitively because PartialSegmentBundleCacheEntry.doMount holds a + // outstanding references on a LINKED BUNDLE (not just the metadata directly) also defer cleanup past drop(). This + // is caught transitively because PartialSegmentBundleCacheEntry.doMount holds a // metadataEntry.acquireMetadataReference for the bundle's lifetime — so while any bundle has outstanding - // references, metadata.isMounted() still returns true and the reconciliation guard fires. This test proves - // the transitive relationship survives a code change that decouples bundle-and-metadata refs. + // references, metadata.isMounted() still returns true and the reconciliation guard fires. This test proves the + // transitive relationship survives a code change that decouples bundle-and-metadata refs. manager = makeManager(true, true); final StorageLocation location = manager.getLocations().get(0); manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original")); @@ -668,7 +669,7 @@ void testLoadThrowsWhenOldStaticBundleHasOutstandingReferences() throws Exceptio // Pin the BUNDLE only (not the metadata directly). A running query walking the aggregate projection would hold // this reference via its Segment; the mount-time metadata ref carried by the bundle keeps the metadata alive // transitively. - final java.io.Closeable heldBundleRef = oldAgg.acquireReference(); + final Closeable heldBundleRef = oldAgg.acquireReference(); try { final SegmentLoadingException thrown = Assertions.assertThrows( SegmentLoadingException.class, @@ -698,10 +699,10 @@ void testLoadThrowsWhenOldStaticBundleHasOutstandingReferences() throws Exceptio @Test void testLoadEvictsPreExistingUnheldWeakEntry() throws Exception { - // Regression guard for the "stale weak entry" leak: a weak reservation for this segment ID (e.g., installed by - // bootstrap when no rule fingerprint was persisted, or by a prior on-demand acquireSegment) must be evicted - // before loadPartial reserves the new static entry. drop() only clears staticCacheEntries, so without explicit - // weak-eviction the two would coexist on the same on-disk directory. + // a weak reservation for this segment ID (e.g., installed by bootstrap when no rule fingerprint was persisted, + // or by a prior on-demand acquireSegment) must be evicted before loadPartial reserves the new static entry. + // drop() only clears staticCacheEntries, so without explicit weak-eviction the two would coexist on the same + // on-disk directory. manager = makeManager(true, true); final StorageLocation location = manager.getLocations().get(0); final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(SEGMENT_ID); From 85654f8a42455c573cc5b38e06998a925d6545fc Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Tue, 7 Jul 2026 03:48:14 -0700 Subject: [PATCH 3/3] better --- .../loading/SegmentLocalCacheManager.java | 610 ++++++++++-------- .../PartialSegmentCacheBootstrapTest.java | 190 +++--- ...tLocalCacheManagerPartialRuleLoadTest.java | 282 ++++---- 3 files changed, 577 insertions(+), 505 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 a3d501811749..4dd948730138 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 @@ -938,6 +938,37 @@ private record ReservedPartial( * before the {@link SegmentLoadingException} propagates, so no orphan static reservation survives the load call. */ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException + { + final PartialLoadSpec wrapper = materializePartialLoadSpec(dataSegment); + final SegmentRangeReader rangeReader = openPartialRangeReader(dataSegment, wrapper); + if (rangeReader == null) { + handleUnsupportedRangeReaderFallback(dataSegment, wrapper); + return; + } + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(dataSegment.getId()); + final ReferenceCountingLock lock = lock(dataSegment); + synchronized (lock) { + try { + if (reconcileExistingEntries(dataSegment, wrapper, id)) { + return; + } + final PartialReservation reservation = + reserveMetadataOnFirstAvailableLocation(dataSegment, wrapper, rangeReader); + persistInfoFileAndAttachCleanupHook(dataSegment, reservation); + mountAndAwaitPartial(dataSegment, wrapper, reservation); + } + finally { + unlock(dataSegment, lock); + } + } + } + + /** + * Materialize the segment's wrapped load spec to a {@link PartialLoadSpec}. Fails with a diagnostic when Jackson + * can't convert (missing subtype registration, malformed wire form) or when the materialized type isn't a partial + * wrapper despite the dispatch check upstream. + */ + private PartialLoadSpec materializePartialLoadSpec(DataSegment dataSegment) throws SegmentLoadingException { final LoadSpec materializedLoadSpec; try { @@ -957,320 +988,353 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException materializedLoadSpec.getClass().getSimpleName() ); } + return wrapper; + } - final SegmentRangeReader rangeReader; + /** + * Open a range reader for the segment's deep storage, wrapping any {@link IOException} as a + * {@link SegmentLoadingException}. Returns {@code null} if the backend doesn't support range reads (e.g. zipped + * deep storage) so caller falls back to weak-full-load semantics. + */ + @Nullable + private SegmentRangeReader openPartialRangeReader(DataSegment dataSegment, PartialLoadSpec wrapper) + throws SegmentLoadingException + { try { - rangeReader = wrapper.openRangeReader(); + return wrapper.openRangeReader(); } catch (IOException e) { - throw new SegmentLoadingException( - e, - "Failed to open range reader for segment[%s]", - dataSegment.getId() - ); - } - if (rangeReader == null) { - // Backend doesn't support range reads (e.g. zipped deep storage). The rule can't be honored as a partial load; - // fall through to the on-demand weak-full-load path at query time. But if a prior loadPartial installed a - // static entry with a different rule fingerprint, we cannot silently keep serving under the old rule so - // reconcile the static side so the coordinator's rule change takes effect, then return so the next query - // installs a weak entry via the on-demand path. - log.warn( - "Backend for segment[%s] does not support range reads; partial-load rule[fingerprint=%s] cannot be honored, " - + "falling back to weak full-load at query time", - dataSegment.getId(), - wrapper.getFingerprint() - ); - final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(dataSegment.getId()); - final ReferenceCountingLock lock = lock(dataSegment); - synchronized (lock) { - try { - // Static-side reconciliation only. Weak entries (from bootstrap without a rule, or a prior on-demand - // acquire) already represent the fallback behavior we're falling back to, leave them alone. - boolean dropStaleStatic = false; - for (StorageLocation location : locations) { - if (!location.isReserved(id)) { - continue; - } - final SegmentCacheEntry existing = location.getCacheEntry(id); - if (existing instanceof PartialSegmentMetadataCacheEntry existingPartial - && Objects.equals(existingPartial.getFingerprint(), wrapper.getFingerprint())) { - // Same-rule re-issue: preserve the existing entry's info file across its eventual drop by clearing - // the onUnmount hook, matching the fingerprint-match short-circuit in the main reconciliation branch. - existing.setOnUnmount(null); - return; - } - dropStaleStatic = true; - } - if (dropStaleStatic) { - log.info( - "Dropping stale static partial entry for segment[%s]: current rule cannot be applied without range reads", - dataSegment.getId() - ); - drop(dataSegment); - // drop() only deleted tracked files; no reserve follows on this path, so nuke leftover per-segment dirs - // across every location or an empty dir would linger indefinitely. - removeStaleSegmentDirsOnAllLocations(dataSegment); - } - } - finally { - unlock(dataSegment, lock); - } - } - return; + throw new SegmentLoadingException(e, "Failed to open range reader for segment[%s]", dataSegment.getId()); } + } + /** + * Handle the rangeReader-null path: the backend can't support partial loads, so the rule can't be honored. Do + * static-side reconciliation only; a stale static entry from a prior rule with a different fingerprint would keep + * serving the old selection unless dropped. + */ + private void handleUnsupportedRangeReaderFallback(DataSegment dataSegment, PartialLoadSpec wrapper) + { + log.warn( + "Backend for segment[%s] does not support range reads; partial-load rule[fingerprint=%s] cannot be honored, " + + "falling back to weak full-load at query time", + dataSegment.getId(), + wrapper.getFingerprint() + ); final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(dataSegment.getId()); final ReferenceCountingLock lock = lock(dataSegment); synchronized (lock) { try { - // Reconcile against any existing entry. The end-state invariant this block establishes is: on any return - // path below, there is exactly one static partial entry for this segment (matching the wrapper fingerprint) - // and no weak entry alongside it. - // - // - Same-fingerprint static partial with no weak entry alongside: idempotent re-load. Clear onUnmount so the - // info file isn't re-deleted, return. - // - Any other state (different-fingerprint static partial, non-partial static entry, weak partial from - // bootstrap without a persisted fingerprint, weak full entry from a prior on-demand acquireSegment): drop - // the static (cascade-releases linked bundles + fires the info-file cleanup hook) AND evict the weak - // entry, then fall through to a fresh reserve. - // - // The weak eviction step is critical: {@link #drop} only touches {@link StorageLocation#staticCacheEntries}, - // so without it a lingering weak entry from bootstrap or on-demand load would coexist with the new static - // entry, same on-disk directory, both counted against capacity, and the eventual bundle downloads would - // race the weak entry's own reads. When the weak entry has active holds we cannot evict it synchronously, - // so throw and let the coordinator retry once the running query completes. drop() takes no per-segment lock - // so calling it from inside our synchronized(lock) block is safe. - PartialSegmentMetadataCacheEntry sameFingerprintStatic = null; - final List oldStaticEntries = new ArrayList<>(); + boolean dropStaleStatic = false; for (StorageLocation location : locations) { if (!location.isReserved(id)) { continue; } final SegmentCacheEntry existing = location.getCacheEntry(id); - if (existing == null) { - continue; - } - oldStaticEntries.add(existing); if (existing instanceof PartialSegmentMetadataCacheEntry existingPartial && Objects.equals(existingPartial.getFingerprint(), wrapper.getFingerprint())) { - sameFingerprintStatic = existingPartial; + // Same-rule re-issue: preserve the existing entry's info file across its eventual drop by clearing + // the onUnmount hook, matching the fingerprint-match short-circuit in the main reconciliation branch. + existing.setOnUnmount(null); + return; } + dropStaleStatic = true; } - final boolean weakExists = locations.stream().anyMatch(loc -> loc.isWeakReserved(id)); - if (sameFingerprintStatic != null && !weakExists && oldStaticEntries.size() == 1) { - sameFingerprintStatic.setOnUnmount(null); - return; - } - if (!oldStaticEntries.isEmpty()) { - // Precheck BEFORE drop: refuse the reconciliation if any old partial metadata entry (or any of its linked - // bundles) has an external reference. drop() would remove entries from staticCacheEntries synchronously - // but leave their doActualUnmount (onUnmount hook + deleteHeaderFiles + evictContainer) deferred until - // the pinning query releases. Between the throw and the next coordinator retry, other threads (queries, - // on-demand acquires) would observe the segment as absent, install a fresh partial entry via the acquire - // path, and race the deferred cleanup on the shared per-segment on-disk directory. Checking BEFORE drop - // keeps the old entry fully installed on the failure path so a subsequent retry finds it via - // reconciliation logic as if this load never happened. - for (SegmentCacheEntry old : oldStaticEntries) { - if (old instanceof PartialSegmentMetadataCacheEntry oldPartial - && oldPartial.cascadeReleaseWouldDeferCleanup()) { - throw new SegmentLoadingException( - "Cannot reconcile partial-load entry for segment[%s]: outstanding references on the old entry " - + "would defer drop's cleanup and race the new entry's on-disk state; coordinator will retry", - dataSegment.getId() - ); - } - } + if (dropStaleStatic) { log.info( - "Reconciling partial-load entry for segment[%s]: fingerprint or type changed, dropping and reloading", + "Dropping stale static partial entry for segment[%s]: current rule cannot be applied without range reads", dataSegment.getId() ); drop(dataSegment); + // drop() only deleted tracked files; no reserve follows on this path, so nuke leftover per-segment dirs + // across every location or an empty dir would linger indefinitely. + removeStaleSegmentDirsOnAllLocations(dataSegment); } - for (StorageLocation location : locations) { - if (!location.isWeakReserved(id)) { - continue; - } - location.removeUnheldWeakEntry(id); - if (location.isWeakReserved(id)) { - throw new SegmentLoadingException( - "Cannot install partial-load rule for segment[%s]: an existing weak entry has active holds; " - + "coordinator will retry", - dataSegment.getId() - ); - } - } - - // After any reconciliation drop/eviction above, nuke leftover per-segment directories on every location. - removeStaleSegmentDirsOnAllLocations(dataSegment); - - // Reserve the metadata entry as static on the first location that accepts BOTH the reservation and a - // per-segment mkdirp. mkdirp failures are location-local (inode exhaustion, permission mismatch, transient - // EIO on one disk) and a different location may accept, so release-and-continue rather than abort. Collect - // per-location failures to attach as suppressed exceptions on the final alert. - StorageLocation reservedLocation = null; - PartialSegmentMetadataCacheEntry metadata = null; - final List perLocationFailures = new ArrayList<>(); - final Iterator iterator = strategy.getLocations(); - while (iterator.hasNext()) { - final StorageLocation candidate = iterator.next(); - final File partialDir = new File(candidate.getPath(), dataSegment.getId().toString()); - final PartialSegmentMetadataCacheEntry candidateEntry = new PartialSegmentMetadataCacheEntry( - dataSegment.getId(), - partialDir, - IndexIO.V10_FILE_NAME, - List.of(), - // staticBundleNames left empty: the fresh-load path below resolves selected bundles AFTER metadata mount - // (since the bundle names depend on the parsed header) and reserves them static directly via - // reserveAndMountStaticBundle. The field is only consulted by bootstrap restore, which has a parsed - // on-disk header available before constructing the entry. - Set.of(), - wrapper.getFingerprint(), - rangeReader, - jsonMapper, - virtualStorageLoadingThreadPool, - config.getVirtualStorageMetadataReservationEstimate() - ); - if (!candidate.reserve(candidateEntry)) { - continue; - } - try { - FileUtils.mkdirp(partialDir); - } - catch (IOException mkdirpFailure) { - candidate.release(candidateEntry); - log.warn( - mkdirpFailure, - "Failed to create partial cache dir on location[%s] for segment[%s]; trying next location", - candidate.getPath(), - dataSegment.getId() - ); - perLocationFailures.add(mkdirpFailure); - continue; - } - reservedLocation = candidate; - metadata = candidateEntry; - break; - } + } + finally { + unlock(dataSegment, lock); + } + } + } - if (reservedLocation == null) { - final SegmentLoadingException noCapacity = new SegmentLoadingException( - "Unable to reserve static partial metadata entry for segment[%s] on any location; ensure enough disk " - + "space and that per-segment cache directories can be created", + /** + * Reconcile any existing cache entry for this segment against the incoming wrapper. + *

+ * End-state invariant on any non-throw return: there is at most one static partial entry for this segment + * (matching the wrapper fingerprint) and no weak entry alongside it. + *

+ * The weak eviction step is critical: {@link #drop} only touches {@link StorageLocation#staticCacheEntries}, so + * without it a lingering weak entry would coexist with the new static entry on the same on-disk directory. + *

+ * Deferred-cleanup precheck. If any old partial has external references (queries holding + * {@link PartialSegmentMetadataCacheEntry#acquireMetadataReference} or a linked bundle's + * {@link PartialSegmentBundleCacheEntry#acquireReference}), {@link #drop} would leave doActualUnmount deferred + * until they release. That deferred cleanup (onUnmount hook + deleteHeaderFiles + evictContainer) would race the + * new entry's on-disk state, so throw pre-drop and let the coordinator retry. Same for a held weak entry: no safe + * way to evict it synchronously without breaking the running query's reads. + */ + private boolean reconcileExistingEntries( + DataSegment dataSegment, + PartialLoadSpec wrapper, + SegmentCacheEntryIdentifier id + ) throws SegmentLoadingException + { + PartialSegmentMetadataCacheEntry sameFingerprintStatic = null; + final List oldStaticEntries = new ArrayList<>(); + for (StorageLocation location : locations) { + if (!location.isReserved(id)) { + continue; + } + final SegmentCacheEntry existing = location.getCacheEntry(id); + if (existing == null) { + continue; + } + oldStaticEntries.add(existing); + if (existing instanceof PartialSegmentMetadataCacheEntry existingPartial + && Objects.equals(existingPartial.getFingerprint(), wrapper.getFingerprint())) { + sameFingerprintStatic = existingPartial; + } + } + final boolean weakExists = locations.stream().anyMatch(loc -> loc.isWeakReserved(id)); + if (sameFingerprintStatic != null && !weakExists && oldStaticEntries.size() == 1) { + sameFingerprintStatic.setOnUnmount(null); + return true; + } + if (!oldStaticEntries.isEmpty()) { + for (SegmentCacheEntry old : oldStaticEntries) { + if (old instanceof PartialSegmentMetadataCacheEntry oldPartial + && oldPartial.cascadeReleaseWouldDeferCleanup()) { + throw new SegmentLoadingException( + "Cannot reconcile partial-load entry for segment[%s]: outstanding references on the old entry would " + + "defer drop's cleanup and race the new entry's on-disk state; coordinator will retry", dataSegment.getId() ); - perLocationFailures.forEach(noCapacity::addSuppressed); - log.makeAlert( - noCapacity, - "Failed to reserve partial-load metadata on any location for segment[%s]", - dataSegment.getId() - ).emit(); - throw noCapacity; } + } + log.info( + "Reconciling partial-load entry for segment[%s]: fingerprint or type changed, dropping and reloading", + dataSegment.getId() + ); + drop(dataSegment); + } + for (StorageLocation location : locations) { + if (!location.isWeakReserved(id)) { + continue; + } + location.removeUnheldWeakEntry(id); + if (location.isWeakReserved(id)) { + throw new SegmentLoadingException( + "Cannot install partial-load rule for segment[%s]: an existing weak entry has active holds; coordinator " + + "will retry", + dataSegment.getId() + ); + } + } + removeStaleSegmentDirsOnAllLocations(dataSegment); + return false; + } - // Persist the info file so bootstrap can rediscover this segment on restart - try { - storeInfoFile(dataSegment); - } - catch (IOException e) { - reservedLocation.release(metadata); - final SegmentLoadingException storeFailure = new SegmentLoadingException( - e, - "Failed to write info file for segment[%s]", + /** + * Reserve the static partial metadata on the first location that accepts BOTH the reservation and a per-segment + * mkdirp. mkdirp failures are location-local (inode exhaustion, permission mismatch, transient EIO on one disk) so + * release-and-continue rather than abort. Per-location failures are gathered as suppressed exceptions on the alert + * emitted when no location succeeds. + */ + private PartialReservation reserveMetadataOnFirstAvailableLocation( + DataSegment dataSegment, + PartialLoadSpec wrapper, + SegmentRangeReader rangeReader + ) throws SegmentLoadingException + { + final List perLocationFailures = new ArrayList<>(); + final Iterator iterator = strategy.getLocations(); + while (iterator.hasNext()) { + final StorageLocation candidate = iterator.next(); + final File partialDir = new File(candidate.getPath(), dataSegment.getId().toString()); + final PartialSegmentMetadataCacheEntry candidateEntry = new PartialSegmentMetadataCacheEntry( + dataSegment.getId(), + partialDir, + IndexIO.V10_FILE_NAME, + List.of(), + // staticBundleNames left empty: the fresh-load path below resolves selected bundles AFTER metadata mount + // (since the bundle names depend on the parsed header) and reserves them static directly via + // reserveAndMountStaticBundle. The field is only consulted by bootstrap restore, which has a parsed + // on-disk header available before constructing the entry. + Set.of(), + wrapper.getFingerprint(), + rangeReader, + jsonMapper, + virtualStorageLoadingThreadPool, + config.getVirtualStorageMetadataReservationEstimate() + ); + if (!candidate.reserve(candidateEntry)) { + continue; + } + try { + FileUtils.mkdirp(partialDir); + } + catch (IOException mkdirpFailure) { + candidate.release(candidateEntry); + // Partial mkdirp may have left directory fragments behind (parent created, sub-dir failed). No-op if fully + // absent; otherwise moves to __drop and deletes. + atomicMoveAndDeleteCacheEntryDirectory(partialDir); + log.warn( + mkdirpFailure, + "Failed to create partial cache dir on location[%s] for segment[%s]; trying next location", + candidate.getPath(), + dataSegment.getId() + ); + perLocationFailures.add(mkdirpFailure); + continue; + } + return new PartialReservation(candidate, candidateEntry); + } + final SegmentLoadingException noCapacity = new SegmentLoadingException( + "Unable to reserve static partial metadata entry for segment[%s] on any location; ensure enough disk space " + + "and that per-segment cache directories can be created", + dataSegment.getId() + ); + perLocationFailures.forEach(noCapacity::addSuppressed); + log.makeAlert( + noCapacity, + "Failed to reserve partial-load metadata on any location for segment[%s]", + dataSegment.getId() + ).emit(); + throw noCapacity; + } + + /** + * Persist the segment's info file so bootstrap can rediscover it on restart, then attach the info-file-delete hook + * to the metadata entry so it fires on drop. The info dir is shared across locations (single configured path), so + * failure here is not a per-disk problem — no fall-through: release the reservation, alert, abort. + */ + private void persistInfoFileAndAttachCleanupHook(DataSegment dataSegment, PartialReservation reservation) + throws SegmentLoadingException + { + try { + storeInfoFile(dataSegment); + } + catch (IOException e) { + reservation.location().release(reservation.metadata()); + // release() went through the never-mounted branch and ran no hook (none set yet), so the mkdirp'd partial dir + // is still on disk with no entry backing it. Nuke it so bootstrap doesn't see a stale directory later. + atomicMoveAndDeleteCacheEntryDirectory(reservation.metadata().getLocalCacheDir()); + final SegmentLoadingException storeFailure = new SegmentLoadingException( + e, + "Failed to write info file for segment[%s]", + dataSegment.getId() + ); + log.makeAlert( + storeFailure, + "Failed to persist partial-load info file for segment[%s]; check druid.segmentCache.infoDir", + dataSegment.getId() + ).emit(); + throw storeFailure; + } + reservation.metadata().setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); + } + + /** + * Submit the metadata + rule-selected-bundle mount to the loading thread pool and await its outcome. All rollback + * (bundles AND metadata) lives INSIDE the callable so no state mutation of {@code reservation.metadata()} crosses + * threads: the caller only ever observes the outcome and never touches shared state on the failure path. This + * avoids the FutureTask cancel-vs-cleanup race — {@code cancel(true)} transitions the future to CANCELLED + * synchronously while the callable's cleanup may still be executing on the pool thread. + */ + private void mountAndAwaitPartial( + DataSegment dataSegment, + PartialLoadSpec wrapper, + PartialReservation reservation + ) throws SegmentLoadingException + { + final StorageLocation theLocation = reservation.location(); + final PartialSegmentMetadataCacheEntry theMetadata = reservation.metadata(); + final Future future = virtualStorageLoadingThreadPool.getExecutorService().submit(() -> { + final List mounted = new ArrayList<>(); + try { + theMetadata.mount(theLocation); + final PartialSegmentFileMapperV10 mapper = theMetadata.getFileMapper(); + if (mapper == null) { + throw DruidException.defensive( + "Partial metadata for segment[%s] mounted without a file mapper", dataSegment.getId() ); - log.makeAlert( - storeFailure, - "Failed to persist partial-load info file for segment[%s]; check druid.segmentCache.infoDir", - dataSegment.getId() - ).emit(); - throw storeFailure; } - metadata.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); - - // Mount metadata + reserve+mount+download each rule-selected bundle as static - final StorageLocation theLocation = reservedLocation; - final PartialSegmentMetadataCacheEntry theMetadata = metadata; - // All rollback responsibility (bundles AND metadata) lives INSIDE the callable so no state mutation of - // theMetadata crosses threads. With cleanup owned by the pool thread, the caller only ever observes the - // outcome; there is no shared state to race. - final Future future = virtualStorageLoadingThreadPool.getExecutorService().submit(() -> { - final List mounted = new ArrayList<>(); + final List selectedBundleNames = + wrapper.getSelectedBundleNames(dataSegment, mapper.getSegmentFileMetadata()); + // A rule-pinned bundle's dependencies must also be pinned. Iterate in dependency-first order so children + // see their dependencies resident. + for (String bundleName : theMetadata.bundlesInMountOrder(selectedBundleNames)) { + mounted.add(reserveAndMountStaticBundle(theLocation, theMetadata, mapper, bundleName)); + } + return null; + } + catch (Throwable t) { + // Cascade rollback on the SAME thread: bundles first in reverse mount order, then the metadata (whose + // release-triggered doActualUnmount hook deletes the info file + header files). Finally nuke the per-segment + // dir itself: doActualUnmount only deletes tracked files (headers, containers via bundle eviction), not the + // enclosing directory that mkdirp created up in reserveMetadataOnFirstAvailableLocation. Without this the + // empty dir would linger until the next reconciliation for this segment ID. + for (int i = mounted.size() - 1; i >= 0; i--) { try { - theMetadata.mount(theLocation); - final PartialSegmentFileMapperV10 mapper = theMetadata.getFileMapper(); - if (mapper == null) { - throw DruidException.defensive( - "Partial metadata for segment[%s] mounted without a file mapper", - dataSegment.getId() - ); - } - final List selectedBundleNames = wrapper.getSelectedBundleNames( - dataSegment, - mapper.getSegmentFileMetadata() - ); - // A rule-pinned bundle's dependencies must also be pinned. Iterate in dependency-first order so children - // see their dependencies resident. - for (String bundleName : theMetadata.bundlesInMountOrder(selectedBundleNames)) { - mounted.add(reserveAndMountStaticBundle(theLocation, theMetadata, mapper, bundleName)); - } - return null; - } - catch (Throwable t) { - // Cascade rollback on the SAME thread: bundles first in reverse mount order, then the metadata (whose - // release-triggered doActualUnmount hook deletes the info file). - for (int i = mounted.size() - 1; i >= 0; i--) { - try { - theLocation.release(mounted.get(i)); - } - catch (Throwable releaseFailure) { - log.warn( - releaseFailure, - "Failed to release bundle during rollback for segment[%s]", - dataSegment.getId() - ); - } - } - releaseMetadataQuietly(theLocation, theMetadata, dataSegment); - throw t; + theLocation.release(mounted.get(i)); } - }); - - try { - future.get(); - } - catch (InterruptedException e) { - // Caller interrupted while waiting. Restore the interrupt flag and propagate, but do NOT cancel the pool - // task: FutureTask.cancel(true) transitions to CANCELLED synchronously, but the callable's run() may still - // be executing its catch/finally on the pool thread, so a subsequent metadata release from the caller would - // race that cleanup. Leaving the task alone lets its own cleanup complete on the pool thread naturally; on - // eventual shutdown the pool's shutdownNow will interrupt the worker and drive the same - // rollback via the callable's catch. - Thread.currentThread().interrupt(); - throw new SegmentLoadingException(e, "Interrupted while loading partial segment[%s]", dataSegment.getId()); - } - catch (Throwable t) { - // Callable ran to a terminal state. On failure it already released metadata + bundles on the pool thread - // via its catch block. We only propagate the diagnostic here; no cleanup remains for the caller. Restore the - // interrupt flag if set so downstream shutdown handling still observes it. - final Throwable executionCause = (t instanceof ExecutionException) ? t.getCause() : t; - final Throwable cause = (executionCause != null) ? executionCause : t; - if (cause instanceof InterruptedException) { - Thread.currentThread().interrupt(); + catch (Throwable releaseFailure) { + log.warn(releaseFailure, "Failed to release bundle during rollback for segment[%s]", dataSegment.getId()); } - throw new SegmentLoadingException( - cause, - "Failed to load partial segment[%s] with rule-pinned bundles", - dataSegment.getId() - ); } + releaseMetadataQuietly(theLocation, theMetadata, dataSegment); + atomicMoveAndDeleteCacheEntryDirectory(theMetadata.getLocalCacheDir()); + throw t; } - finally { - unlock(dataSegment, lock); + }); + + try { + future.get(); + } + catch (InterruptedException e) { + // Caller interrupted while waiting. Restore the interrupt flag and propagate, but do NOT cancel the pool task: + // FutureTask.cancel(true) transitions to CANCELLED synchronously, but the callable's run() may still be + // executing its catch/finally on the pool thread, so a subsequent metadata release from the caller would race + // that cleanup. Leaving the task alone lets its own cleanup complete on the pool thread naturally; on eventual + // shutdown the pool's shutdownNow will interrupt the worker and drive the same rollback via the callable's + // catch. + Thread.currentThread().interrupt(); + throw new SegmentLoadingException(e, "Interrupted while loading partial segment[%s]", dataSegment.getId()); + } + catch (Throwable t) { + // Callable ran to a terminal state. On failure it already released metadata + bundles on the pool thread via + // its catch block. We only propagate the diagnostic here; no cleanup remains for the caller. Restore the + // interrupt flag if set so downstream shutdown handling still observes it. + final Throwable executionCause = (t instanceof ExecutionException) ? t.getCause() : t; + final Throwable cause = (executionCause != null) ? executionCause : t; + if (cause instanceof InterruptedException) { + Thread.currentThread().interrupt(); } + throw new SegmentLoadingException( + cause, + "Failed to load partial segment[%s] with rule-pinned bundles", + dataSegment.getId() + ); } } + /** + * Pair of storage location + reserved metadata entry, returned by + * {@link #reserveMetadataOnFirstAvailableLocation} and threaded through the info-file persist and async mount + * steps. + */ + private record PartialReservation(StorageLocation location, PartialSegmentMetadataCacheEntry metadata) + { + } + /** * Reserve a {@link PartialSegmentBundleCacheEntry} as static on {@code location}, mount it (which initializes the * sparse containers and links it to the metadata entry), and eagerly download all of the bundle's containers via diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java index 328b68841eb7..a020b7846886 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java @@ -477,101 +477,6 @@ void testBitmapRepairClearsBitsForMissingContainers() throws IOException } } - /** - * Populate the per-segment cache dir with the on-disk artifacts a previous historical run would have left behind: - * the V10 header file plus sparse-allocated container files for the base and aggregate bundles. - */ - private void primeOnDiskState() throws IOException - { - final StorageLocation seedLocation = new StorageLocation(cacheDir, ESTIMATE * 8, null); - final PartialSegmentMetadataCacheEntry seedMeta = new PartialSegmentMetadataCacheEntry( - SEGMENT_ID, - cacheDir, - IndexIO.V10_FILE_NAME, - List.of(), - Set.of(), - null, - new DirectoryBackedRangeReader(deepStorageDir), - JSON_MAPPER, - null, - ESTIMATE - ); - Assertions.assertTrue(seedLocation.reserve(seedMeta)); - seedMeta.mount(seedLocation); - - final PartialSegmentBundleCacheEntry base = PartialSegmentBundleCacheEntry.forBundle( - seedMeta, - Projections.BASE_TABLE_PROJECTION_NAME, - List.of() - ); - final var baseHold = seedLocation.addWeakReservationHold(base.getId(), () -> base); - Assertions.assertNotNull(baseHold); - base.mount(seedLocation); - - final PartialSegmentBundleCacheEntry agg = PartialSegmentBundleCacheEntry.forBundle( - seedMeta, - AGG_BUNDLE, - List.of(base.getId()) - ); - final var aggHold = seedLocation.addWeakReservationHold(agg.getId(), () -> agg); - Assertions.assertNotNull(aggHold); - agg.mount(seedLocation); - - // Leave on-disk state behind: unmount the bundles (which deletes container files!), that's the wrong final - // state. Instead, we want containers ON disk, so leave bundles mounted but close the file mapper. Since the - // restore path re-opens via PartialSegmentFileMapperV10.create which is idempotent w.r.t. on-disk files, - // un-mount on the SEED side AFTER files are sparse-allocated would also delete them. So we just leave the - // seed mounted: at test end @TempDir cleans up. - aggHold.close(); - baseHold.close(); - } - - /** - * Two-step restore helper that mirrors what {@code SegmentLocalCacheManager} does in production: reserve the metadata - * entry via {@link PartialSegmentCacheBootstrap#reserveFromDisk}, then drive {@link PartialSegmentMetadataCacheEntry#mount} - * to trigger the file-mapper build + bundle restore. On a mount failure, {@code mount}'s own rollback removes the - * (unheld) weak entry from the location, so tests can assert on a clean location without any extra cleanup here. - */ - private PartialSegmentMetadataCacheEntry restoreFromDisk(StorageLocation location) throws IOException - { - final PartialSegmentMetadataCacheEntry metadata = PartialSegmentCacheBootstrap.reserveFromDisk( - SEGMENT_ID, - cacheDir, - IndexIO.V10_FILE_NAME, - List.of(), - Set.of(), - null, - new DirectoryBackedRangeReader(deepStorageDir), - JSON_MAPPER, - null, - location - ); - metadata.mount(location); - return metadata; - } - - /** - * Static-restore variant: reserve metadata + selected bundles as static entries. - */ - private PartialSegmentMetadataCacheEntry restoreFromDiskAsStatic(StorageLocation location, Set staticBundles) - throws IOException - { - final PartialSegmentMetadataCacheEntry metadata = PartialSegmentCacheBootstrap.reserveFromDisk( - SEGMENT_ID, - cacheDir, - IndexIO.V10_FILE_NAME, - List.of(), - staticBundles, - "v1:test-fingerprint", - new DirectoryBackedRangeReader(deepStorageDir), - JSON_MAPPER, - null, - location - ); - metadata.mount(location); - return metadata; - } - @Test void testStaticReserveFromDiskRegistersMetadataInStaticMap() throws IOException { @@ -768,6 +673,101 @@ void testStaticRestoreReleasesReservationsWhenBundleMountFails() throws IOExcept ); } + /** + * Populate the per-segment cache dir with the on-disk artifacts a previous historical run would have left behind: + * the V10 header file plus sparse-allocated container files for the base and aggregate bundles. + */ + private void primeOnDiskState() throws IOException + { + final StorageLocation seedLocation = new StorageLocation(cacheDir, ESTIMATE * 8, null); + final PartialSegmentMetadataCacheEntry seedMeta = new PartialSegmentMetadataCacheEntry( + SEGMENT_ID, + cacheDir, + IndexIO.V10_FILE_NAME, + List.of(), + Set.of(), + null, + new DirectoryBackedRangeReader(deepStorageDir), + JSON_MAPPER, + null, + ESTIMATE + ); + Assertions.assertTrue(seedLocation.reserve(seedMeta)); + seedMeta.mount(seedLocation); + + final PartialSegmentBundleCacheEntry base = PartialSegmentBundleCacheEntry.forBundle( + seedMeta, + Projections.BASE_TABLE_PROJECTION_NAME, + List.of() + ); + final var baseHold = seedLocation.addWeakReservationHold(base.getId(), () -> base); + Assertions.assertNotNull(baseHold); + base.mount(seedLocation); + + final PartialSegmentBundleCacheEntry agg = PartialSegmentBundleCacheEntry.forBundle( + seedMeta, + AGG_BUNDLE, + List.of(base.getId()) + ); + final var aggHold = seedLocation.addWeakReservationHold(agg.getId(), () -> agg); + Assertions.assertNotNull(aggHold); + agg.mount(seedLocation); + + // Leave on-disk state behind: unmount the bundles (which deletes container files!), that's the wrong final + // state. Instead, we want containers ON disk, so leave bundles mounted but close the file mapper. Since the + // restore path re-opens via PartialSegmentFileMapperV10.create which is idempotent w.r.t. on-disk files, + // un-mount on the SEED side AFTER files are sparse-allocated would also delete them. So we just leave the + // seed mounted: at test end @TempDir cleans up. + aggHold.close(); + baseHold.close(); + } + + /** + * Two-step restore helper that mirrors what {@code SegmentLocalCacheManager} does in production: reserve the metadata + * entry via {@link PartialSegmentCacheBootstrap#reserveFromDisk}, then drive {@link PartialSegmentMetadataCacheEntry#mount} + * to trigger the file-mapper build + bundle restore. On a mount failure, {@code mount}'s own rollback removes the + * (unheld) weak entry from the location, so tests can assert on a clean location without any extra cleanup here. + */ + private PartialSegmentMetadataCacheEntry restoreFromDisk(StorageLocation location) throws IOException + { + final PartialSegmentMetadataCacheEntry metadata = PartialSegmentCacheBootstrap.reserveFromDisk( + SEGMENT_ID, + cacheDir, + IndexIO.V10_FILE_NAME, + List.of(), + Set.of(), + null, + new DirectoryBackedRangeReader(deepStorageDir), + JSON_MAPPER, + null, + location + ); + metadata.mount(location); + return metadata; + } + + /** + * Static-restore variant: reserve metadata + selected bundles as static entries. + */ + private PartialSegmentMetadataCacheEntry restoreFromDiskAsStatic(StorageLocation location, Set staticBundles) + throws IOException + { + final PartialSegmentMetadataCacheEntry metadata = PartialSegmentCacheBootstrap.reserveFromDisk( + SEGMENT_ID, + cacheDir, + IndexIO.V10_FILE_NAME, + List.of(), + staticBundles, + "v1:test-fingerprint", + new DirectoryBackedRangeReader(deepStorageDir), + JSON_MAPPER, + null, + location + ); + metadata.mount(location); + return metadata; + } + private static PartialSegmentFileMapperV10 createMapper(File deepStorageDir, File cacheDir) throws IOException { return PartialSegmentFileMapperV10.create( 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 16b3d599fcd2..7b082005c24d 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 @@ -186,109 +186,6 @@ void tearDown() } } - private SegmentLocalCacheManager makeManager(boolean virtualStorage, boolean partialDownloadsEnabled) - { - return makeManagerAtLocations(virtualStorage, partialDownloadsEnabled, List.of(cacheRoot)); - } - - private SegmentLocalCacheManager makeManagerAtLocations( - boolean virtualStorage, - boolean partialDownloadsEnabled, - List locationRoots - ) - { - final List locConfigs = locationRoots.stream() - .map(root -> new StorageLocationConfig(root, 1024L * 1024L * 1024L, null)) - .toList(); - final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig() - .setLocations(locConfigs) - .setVirtualStorage(virtualStorage) - .setVirtualStoragePartialDownloadsEnabled(partialDownloadsEnabled); - final List storageLocations = loaderConfig.toStorageLocations(); - return new SegmentLocalCacheManager( - storageLocations, - loaderConfig, - StorageLoadingThreadPool.createFromConfig(loaderConfig), - new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations), - TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT), - jsonMapper - ); - } - - private DataSegment partialWrapperSegment(List selectedProjections) - { - return partialWrapperSegment(selectedProjections, FINGERPRINT); - } - - private DataSegment partialWrapperSegment(List selectedProjections, String fingerprint) - { - final Map delegate = Map.of( - "type", "local", - "path", DEEP_STORAGE_DIR.getAbsolutePath() - ); - final Map wrapperWire = - PartialProjectionLoadSpec.wireForm(delegate, selectedProjections, fingerprint); - return DataSegment.builder(SEGMENT_ID) - .shardSpec(NoneShardSpec.instance()) - .loadSpec(wrapperWire) - .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 (e.g. a - * historical that received a partial-load rule pointing at zipped deep storage). - */ - private DataSegment partialWrapperSegmentWithNullRangeReader(List selectedProjections, String fingerprint) - throws IOException - { - final File noRangeReaderDir = new File( - perTestTempDir, - "no_range_reader_" + fingerprint.replace(':', '_').replace('.', '_') - ); - FileUtils.mkdirp(noRangeReaderDir); - final Map delegate = Map.of( - "type", "local", - "path", noRangeReaderDir.getAbsolutePath() - ); - final Map wrapperWire = - PartialProjectionLoadSpec.wireForm(delegate, selectedProjections, fingerprint); - return DataSegment.builder(SEGMENT_ID) - .shardSpec(NoneShardSpec.instance()) - .loadSpec(wrapperWire) - .size(0) - .build(); - } - - /** - * Post-load lookup of the mounted partial metadata entry. loadPartial is synchronous — by the time load() returns, - * the entry is either fully mounted or the call threw; tests can inspect state immediately. - */ - private static PartialSegmentMetadataCacheEntry mountedMetadata(StorageLocation location, SegmentId segmentId) - { - final CacheEntry entry = location.getStaticCacheEntry(new SegmentCacheEntryIdentifier(segmentId)); - Assertions.assertInstanceOf(PartialSegmentMetadataCacheEntry.class, entry); - final PartialSegmentMetadataCacheEntry partial = (PartialSegmentMetadataCacheEntry) entry; - Assertions.assertTrue(partial.isMounted(), "metadata entry for " + segmentId + " should be mounted after load()"); - return partial; - } - - private static PartialSegmentBundleCacheEntry mountedBundle( - StorageLocation location, - SegmentId segmentId, - String bundleName - ) - { - final CacheEntry entry = location.getStaticCacheEntry( - new PartialSegmentBundleCacheEntryIdentifier(segmentId, bundleName) - ); - Assertions.assertInstanceOf(PartialSegmentBundleCacheEntry.class, entry); - final PartialSegmentBundleCacheEntry bundle = (PartialSegmentBundleCacheEntry) entry; - Assertions.assertTrue(bundle.isMounted(), "bundle " + bundleName + " should be mounted after load()"); - return bundle; - } - @Test void testLoadPinsMetadataAndSelectedBundleAsStatic() throws Exception { @@ -754,40 +651,6 @@ void testLoadThrowsWhenPreExistingWeakEntryHasActiveHold() throws Exception } } - private static CacheEntry stubCacheEntry(CacheEntryIdentifier id, long size) - { - return new CacheEntry() - { - @Override - public CacheEntryIdentifier getId() - { - return id; - } - - @Override - public long getSize() - { - return size; - } - - @Override - public boolean isMounted() - { - return false; - } - - @Override - public void mount(StorageLocation location) - { - } - - @Override - public void unmount() - { - } - }; - } - @Test void testLoadFailureCascadeReleasesEverything() throws Exception { @@ -826,5 +689,150 @@ void testLoadFailureCascadeReleasesEverything() throws Exception // as part of the cascade release. final File infoFile = new File(new File(cacheRoot, "info_dir"), SEGMENT_ID.toString()); Assertions.assertFalse(infoFile.exists(), "info file should be deleted by the metadata's onUnmount hook"); + // The per-segment cache dir was created by mkdirp in reserveMetadataOnFirstAvailableLocation; the mount-failure + // rollback in mountAndAwaitPartial must nuke it (doActualUnmount deletes only header files, not the enclosing + // directory). Otherwise an empty dir would linger on disk until the next reconciliation for this segment. + final File partialDir = new File(cacheRoot, SEGMENT_ID.toString()); + Assertions.assertFalse( + partialDir.exists(), + "per-segment cache dir must be nuked after mount-failure rollback" + ); + } + + private SegmentLocalCacheManager makeManager(boolean virtualStorage, boolean partialDownloadsEnabled) + { + return makeManagerAtLocations(virtualStorage, partialDownloadsEnabled, List.of(cacheRoot)); + } + + private SegmentLocalCacheManager makeManagerAtLocations( + boolean virtualStorage, + boolean partialDownloadsEnabled, + List locationRoots + ) + { + final List locConfigs = locationRoots.stream() + .map(root -> new StorageLocationConfig(root, 1024L * 1024L * 1024L, null)) + .toList(); + final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig() + .setLocations(locConfigs) + .setVirtualStorage(virtualStorage) + .setVirtualStoragePartialDownloadsEnabled(partialDownloadsEnabled); + final List storageLocations = loaderConfig.toStorageLocations(); + return new SegmentLocalCacheManager( + storageLocations, + loaderConfig, + StorageLoadingThreadPool.createFromConfig(loaderConfig), + new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations), + TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT), + jsonMapper + ); + } + + private DataSegment partialWrapperSegment(List selectedProjections) + { + return partialWrapperSegment(selectedProjections, FINGERPRINT); + } + + private DataSegment partialWrapperSegment(List selectedProjections, String fingerprint) + { + final Map delegate = Map.of( + "type", "local", + "path", DEEP_STORAGE_DIR.getAbsolutePath() + ); + final Map wrapperWire = + PartialProjectionLoadSpec.wireForm(delegate, selectedProjections, fingerprint); + return DataSegment.builder(SEGMENT_ID) + .shardSpec(NoneShardSpec.instance()) + .loadSpec(wrapperWire) + .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 (e.g. a + * historical that received a partial-load rule pointing at zipped deep storage). + */ + private DataSegment partialWrapperSegmentWithNullRangeReader(List selectedProjections, String fingerprint) + throws IOException + { + final File noRangeReaderDir = new File( + perTestTempDir, + "no_range_reader_" + fingerprint.replace(':', '_').replace('.', '_') + ); + FileUtils.mkdirp(noRangeReaderDir); + final Map delegate = Map.of( + "type", "local", + "path", noRangeReaderDir.getAbsolutePath() + ); + final Map wrapperWire = + PartialProjectionLoadSpec.wireForm(delegate, selectedProjections, fingerprint); + return DataSegment.builder(SEGMENT_ID) + .shardSpec(NoneShardSpec.instance()) + .loadSpec(wrapperWire) + .size(0) + .build(); + } + + /** + * Post-load lookup of the mounted partial metadata entry. loadPartial is synchronous — by the time load() returns, + * the entry is either fully mounted or the call threw; tests can inspect state immediately. + */ + private static PartialSegmentMetadataCacheEntry mountedMetadata(StorageLocation location, SegmentId segmentId) + { + final CacheEntry entry = location.getStaticCacheEntry(new SegmentCacheEntryIdentifier(segmentId)); + Assertions.assertInstanceOf(PartialSegmentMetadataCacheEntry.class, entry); + final PartialSegmentMetadataCacheEntry partial = (PartialSegmentMetadataCacheEntry) entry; + Assertions.assertTrue(partial.isMounted(), "metadata entry for " + segmentId + " should be mounted after load()"); + return partial; + } + + private static PartialSegmentBundleCacheEntry mountedBundle( + StorageLocation location, + SegmentId segmentId, + String bundleName + ) + { + final CacheEntry entry = location.getStaticCacheEntry( + new PartialSegmentBundleCacheEntryIdentifier(segmentId, bundleName) + ); + Assertions.assertInstanceOf(PartialSegmentBundleCacheEntry.class, entry); + final PartialSegmentBundleCacheEntry bundle = (PartialSegmentBundleCacheEntry) entry; + Assertions.assertTrue(bundle.isMounted(), "bundle " + bundleName + " should be mounted after load()"); + return bundle; + } + + private static CacheEntry stubCacheEntry(CacheEntryIdentifier id, long size) + { + return new CacheEntry() + { + @Override + public CacheEntryIdentifier getId() + { + return id; + } + + @Override + public long getSize() + { + return size; + } + + @Override + public boolean isMounted() + { + return false; + } + + @Override + public void mount(StorageLocation location) + { + } + + @Override + public void unmount() + { + } + }; } }