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..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 @@ -920,7 +920,7 @@ private static void fetchAndPersistHeader( /** * Parse metadata from the header file. Throws on any parse failure (corrupt JSON, truncated header, etc.). */ - 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..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 @@ -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,493 @@ 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 (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 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 { + 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() + ); + } + return wrapper; + } + + /** + * 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 { + return wrapper.openRangeReader(); + } + catch (IOException e) { + 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 { + 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); + } + } + } + + /** + * 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() + ); + } + } + 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; + } + + /** + * 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() + ); + } + 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 { + 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); + atomicMoveAndDeleteCacheEntryDirectory(theMetadata.getLocalCacheDir()); + 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. + 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 + * 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 removeStaleSegmentDirsOnAllLocations(DataSegment dataSegment) + { + for (StorageLocation loc : locations) { + atomicMoveAndDeleteCacheEntryDirectory(new File(loc.getPath(), dataSegment.getId().toString())); + } + } + @Nullable private AcquireSegmentAction acquireExistingSegment(SegmentCacheEntryIdentifier identifier) { @@ -927,8 +1467,16 @@ public void load(final DataSegment dataSegment) throws SegmentLoadingException "load() should not be called when virtualStorageIsEphemeral is true" ); } - // 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 + // 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; + } + // .. 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 { @@ -993,12 +1541,54 @@ public void bootstrap( partial.mount(location); } catch (IOException e) { + // release the reservation so it doesn't leak + 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() ); } + // 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<>(); + 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 +1659,18 @@ 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; 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 +1909,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..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 @@ -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, @@ -473,6 +477,202 @@ void testBitmapRepairClearsBitsForMissingContainers() throws IOException } } + @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 + { + // 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; + 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); + + 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" + ); + } + /** * 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. @@ -485,6 +685,8 @@ private void primeOnDiskState() throws IOException cacheDir, IndexIO.V10_FILE_NAME, List.of(), + Set.of(), + null, new DirectoryBackedRangeReader(deepStorageDir), JSON_MAPPER, null, @@ -533,6 +735,30 @@ private PartialSegmentMetadataCacheEntry restoreFromDisk(StorageLocation locatio 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, 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..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 @@ -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() + { + // 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..7b082005c24d --- /dev/null +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java @@ -0,0 +1,838 @@ +/* + * 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.Closeable; +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(); + } + } + + @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 + { + // 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 + { + // 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 + { + // 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 + { + // 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 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 + { + // 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 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 + { + // 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(); + } + } + + @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"); + // 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() + { + } + }; + } +}