From 83b1ad0b61e7845a66bf0ced06058814c2ae15a4 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Thu, 9 Jul 2026 11:59:37 -0700 Subject: [PATCH 1/6] feat: wire up partial load rules to segment cache changes: * partial load rules/specs are now wired up to the segment cache when virtual storage mode and partial loads are enabled, allowing eager loading of parts of segments which are pinned in the cache until dropped * when under a partial load spec, partial segment cache entries are still added to the cache as weak entries, but with an additional set of 'rule' holds on all required bundles * partial load spec changes operate on a delta, dropping the rule holds for bundles no longer pinned, and loading and adding new rule holds on things which should be pinned --- .../file/PartialSegmentFileMapperV10.java | 29 +- .../loading/PartialClusterGroupLoadSpec.java | 67 ++ .../segment/loading/PartialLoadSpec.java | 16 + .../loading/PartialProjectionLoadSpec.java | 42 ++ .../PartialClusterGroupLoadSpecTest.java | 207 ++++++ .../PartialProjectionLoadSpecTest.java | 99 +++ .../loading/PartialSegmentCacheBootstrap.java | 25 +- .../PartialSegmentMetadataCacheEntry.java | 526 +++++++++++++- .../loading/SegmentLocalCacheManager.java | 453 +++++++++++- .../PartialSegmentMetadataCacheEntryTest.java | 195 ++++- ...ntLocalCacheManagerPartialAcquireTest.java | 6 +- ...tLocalCacheManagerPartialRuleLoadTest.java | 666 ++++++++++++++++++ 12 files changed, 2257 insertions(+), 74 deletions(-) create mode 100644 server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java diff --git a/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java b/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java index ea4895a7f2e4..78582bd7b6b9 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 @@ -448,7 +448,8 @@ public void ensureAllDownloaded() throws IOException } /** - * Download every container belonging to {@code bundleName} in this mapper, each in a single range read (see + * Download every container belonging to {@code bundleName}, recursing into external mappers to cover bundles that + * span the main file plus one or more externals. Each container is downloaded in a single range read (see * {@link #downloadContainer}). No-op for an unknown bundle. */ public void ensureBundleDownloaded(String bundleName) throws IOException @@ -457,6 +458,32 @@ public void ensureBundleDownloaded(String bundleName) throws IOException for (int containerIndex : getContainerIndicesForBundle(bundleName)) { downloadContainer(containerIndex); } + for (PartialSegmentFileMapperV10 external : externalMappers.values()) { + external.ensureBundleDownloaded(bundleName); + } + } + + /** + * Whether every container belonging to {@code bundleName} has all its files present in {@link #downloadedFiles}, + * across the main mapper AND every attached external mapper. + *

+ * Returns {@code true} for a bundle name unknown to any mapper (no containers to check). + */ + public boolean isBundleFullyDownloaded(String bundleName) + { + checkClosed(); + for (int containerIndex : getContainerIndicesForBundle(bundleName)) { + final List fileNames = containerFileNames.get(containerIndex); + if (!downloadedFiles.containsAll(fileNames)) { + return false; + } + } + for (PartialSegmentFileMapperV10 external : externalMappers.values()) { + if (!external.isBundleFullyDownloaded(bundleName)) { + return false; + } + } + return true; } /** 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..cbc27ad3cf99 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,64 @@ public List getClusterGroupIndices() return clusterGroupIndices; } + /** + * Resolve each {@code clusterGroupIndex} against the segment's {@link ClusteredValueGroupsBaseTableSchema} to a + * {@code __base$} bundle name via {@link Projections#getClusterGroupBundleName}. The base + * projection carries the authoritative group list; each index picks the group at that position and its + * {@code clusteringValueIds} disambiguate the bundle in the V10 layout. + *

+ * Defensive tripwires fire on any structural inconsistency: empty projections list, non-clustered base projection, + * segment's cluster-group tuple count vs. metadata's group count mismatch, or index out of range. These conditions + * are unreachable under a healthy writer/reader contract; a throw here indicates a coding bug. + *

+ * Returns an empty list when {@code clusterGroupIndices} is empty (the "sibling-empty" case where a matcher + * applied to a clustered segment but no configured pattern matched any group). + */ + @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..c9b0e82acdd2 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,17 @@ 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's parsed {@code metadata}. The + * historical's cache layer uses the returned names to acquire rule holds on the matching bundles as part of + * {@link SegmentCacheManager#load} for a partial-load rule; bundles not in the returned list are left to the + * ordinary on-demand acquire path. + *

+ * Returns an empty list when the load spec's selection is empty (the "sibling-empty" case — coordinator matcher + * applied but produced no positive selection; the historical still installs a rule hold on the metadata so + * bootstrap re-discovers the segment). + */ + 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..c063f33db8e1 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,42 @@ 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. + *

+ * These are pure defensive tripwires: the coordinator-side matcher derives the wire form from the same segment + * metadata this method reads, so a mismatch here indicates a coding bug (matcher/reader drift, writer/reader + * contract violation, or serialization corruption). + */ + @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..0c75eb31aed8 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,194 @@ void testRejectsNullFingerprint() ); } + @Test + void testGetSelectedBundleNamesResolvesIndicesToBundleNames() + { + // Three cluster groups in the metadata; pick groups 0 and 2. Each index maps to the group's + // clusteringValueIds → bundle name. + 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: matcher applied but no positive selection. 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..178f767d70be 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,13 +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.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; import javax.annotation.Nullable; import java.io.ByteArrayInputStream; import java.io.File; +import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; @@ -167,6 +180,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 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/PartialSegmentCacheBootstrap.java b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java index 8a53c1b5f139..b4d78ae3cfc3 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 { @@ -173,8 +174,8 @@ static void restoreBundlesFromDisk(PartialSegmentMetadataCacheEntry metadata, St 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,23 +193,23 @@ 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))); 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, 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..9c6856fd8f10 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,11 @@ 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.Objects; import java.util.Optional; import java.util.Set; import java.util.concurrent.Callable; @@ -132,8 +136,8 @@ public class PartialSegmentMetadataCacheEntry implements SegmentCacheEntry, Resi // bundle entries that are currently mounted against this segment, registered by PartialSegmentBundleCacheEntry on // successful mount and removed on unmount. Lets the drop path enumerate bundles for cascade-close without scanning - // the StorageLocation's entry maps. - private final Set linkedBundles = ConcurrentHashMap.newKeySet(); + // the StorageLocation's entry maps. Keyed by bundle name (one bundle per name per segment). + private final ConcurrentHashMap linkedBundles = new ConcurrentHashMap<>(); // Reference-counted gate over the actual cleanup work (close file mapper, delete header files). Set on // successful mount; unmount() closes the wrapper which defers running cleanup until all outstanding references @@ -147,6 +151,33 @@ public class PartialSegmentMetadataCacheEntry implements SegmentCacheEntry, Resi // cleared so retries get a fresh attempt; on success the gate stays set until doActualUnmount clears it. private final AtomicReference> mountFuture = new AtomicReference<>(); + /** + * Rule-holds state machine (all guarded by {@link #entryLock}). A partial-load rule is applied via + * {@link #applyRule(String, Set)} after {@link #mount(StorageLocation)}. While a rule is applied + * ({@code ruleFingerprint != null}): + *

+ * {@link #clearRule()} releases the holds and clears the state. {@link #applyRule(String, Set)} with a different + * fingerprint / selection set diffs the current state and closes / acquires only the deltas, so overlapping bundle + * holds stay live across a rule swap. + */ + @GuardedBy("entryLock") + @Nullable + private String ruleFingerprint; + @GuardedBy("entryLock") + private Set ruleSelectedBundleNames = Set.of(); + @GuardedBy("entryLock") + private final Map> ruleBundleHolds = + new HashMap<>(); + @GuardedBy("entryLock") + @Nullable + private StorageLocation.ReservationHold metadataSelfHold; + public PartialSegmentMetadataCacheEntry( SegmentId segmentId, File localCacheDir, @@ -199,6 +230,326 @@ File getLocalCacheDir() return localCacheDir; } + /** + * The fingerprint of the partial-load rule currently applied to this entry, or {@code null} if no rule has been + * applied. Set by {@link #applyRule(String, Set)} and cleared by {@link #clearRule()}. {@link #doActualUnmount()} + * asserts that this field is {@code null} at unmount time, the {@code metadataSelfHold} taken by applyRule pins + * the entry's phaser and prevents eviction reclaim / doActualUnmount from firing while a rule is applied. A fresh + * mount therefore always starts with no rule applied. + */ + @Nullable + public String getRuleFingerprint() + { + entryLock.lock(); + try { + return ruleFingerprint; + } + finally { + entryLock.unlock(); + } + } + + /** + * The set of bundle names currently pinned by an applied partial-load rule, or an empty set if no rule is applied. + * Return value is immutable; safe to inspect without further synchronization. + */ + public Set getRuleSelectedBundleNames() + { + entryLock.lock(); + try { + return ruleSelectedBundleNames; + } + finally { + entryLock.unlock(); + } + } + + /** + * Whether a partial-load rule is currently applied to this entry. Equivalent to {@code getRuleFingerprint() != null}. + * While {@code true}, this entry holds a self-referential {@link StorageLocation.ReservationHold} that prevents cache + * from evicting it, so the metadata's V10 header stays resident until {@link #clearRule()} runs. + */ + public boolean isRuleHeld() + { + entryLock.lock(); + try { + return ruleFingerprint != null; + } + finally { + entryLock.unlock(); + } + } + + /** + * Apply (or replace) a partial-load rule on this entry. Must be called after {@link #mount(StorageLocation)}, the + * entry must be registered with its {@link StorageLocation} so it can acquire holds on itself and on its bundles. + *

+ * Concurrency contract. The caller MUST serialize {@code applyRule} and {@link #clearRule} for a given + * segment id through an external per-segment lock ({@code SegmentLocalCacheManager} uses its + * {@code ReferenceCountingLock segmentLocks}). This method releases {@link #entryLock} between Phase 2 (acquire + * holds outside the lock) and Phase 4 (release holds outside the lock); a concurrent invocation of + * {@code applyRule}/{@code clearRule} on the same entry can observe intermediate state during those windows and + * silently miss rule-hold installations or leak them. {@link #entryLock} alone is insufficient because it is + * released across the acquire/release phases by design (holding it across {@link StorageLocation} lock acquisition + * would invert the writeLock → entryLock ordering the storage layer already uses for eviction). + *

+ * On first application ({@code ruleFingerprint} transitions from {@code null}), a self-referential + * {@link StorageLocation.ReservationHold} is taken to pin the metadata entry in-cache. On repeat calls with the same + * {@code fingerprint} and matching {@code selectedBundleNames}, this is a no-op. On a genuine rule swap, only the + * delta between the previous and new selection is applied: bundle holds for names dropped from the selection are + * closed, and holds for names newly added are acquired for any bundle already registered with this entry (later + * arrivals get their hold via {@link #registerBundle}). + *

+ * Bundles whose {@link StorageLocation.ReservationHold} is currently zero-refcount (never mounted, or evicted but + * unregistered before this call) will not appear in {@link #linkedBundles} yet, so no hold is taken for them here; + * the caller is expected to drive an on-demand acquire per selected bundle name to trigger a fresh mount, which will + * call {@link #registerBundle} and pick up the rule-hold at that point. + * + * @throws DruidException if the entry is not currently mounted + */ + public void applyRule(String fingerprint, Set selectedBundleNames) + { + Objects.requireNonNull(fingerprint, "fingerprint"); + Objects.requireNonNull(selectedBundleNames, "selectedBundleNames"); + final Set newSelection = Set.copyOf(selectedBundleNames); + + // Phase 1: snapshot under entryLock and compute the diff. Do NOT call StorageLocation methods here — that would + // acquire the location's readLock while holding entryLock and could deadlock with a concurrent writeLock holder + // that needs entryLock + final StorageLocation loc; + final boolean needsSelfHold; + final List namesToRelease; + final List namesToAcquire; + entryLock.lock(); + try { + if (location == null) { + throw DruidException.defensive( + "applyRule on partial metadata entry[%s] requires the entry to be mounted", + id + ); + } + if (fingerprint.equals(ruleFingerprint) && newSelection.equals(ruleSelectedBundleNames)) { + return; + } + loc = location; + needsSelfHold = (metadataSelfHold == null); + namesToRelease = new ArrayList<>(); + for (String name : ruleBundleHolds.keySet()) { + if (!newSelection.contains(name)) { + namesToRelease.add(name); + } + } + namesToAcquire = new ArrayList<>(); + for (String name : newSelection) { + if (!ruleBundleHolds.containsKey(name) && findLinkedBundleByName(name) != null) { + namesToAcquire.add(name); + } + } + } + finally { + entryLock.unlock(); + } + + // Phase 2 + 3: acquire outside entryLock, then commit under entryLock. Track every hold this call acquires in + // `uncommittedHolds`; when a hold's ownership transfers into `metadataSelfHold` or `ruleBundleHolds`, it is + // removed from the list. Any throw between acquire and commit (or a race-lose at commit) leaves the hold in + // `uncommittedHolds`, and the finally at Phase 4 releases it. This is what makes applyRule all-or-nothing at + // the ReservationHold level: no leaks on partial failure, no stranded self-hold under a stale fingerprint. + final List> uncommittedHolds = new ArrayList<>(); + final List> displacedHolds = new ArrayList<>(); + try { + final StorageLocation.ReservationHold newSelfHold; + if (needsSelfHold) { + newSelfHold = loc.addWeakReservationHoldIfExists(id); + if (newSelfHold == null) { + throw DruidException.defensive( + "Failed to acquire self-referential rule-hold on partial metadata entry[%s]; entry is not weak-reserved", + id + ); + } + uncommittedHolds.add(newSelfHold); + } else { + newSelfHold = null; + } + final Map> acquired = new HashMap<>(); + for (String name : namesToAcquire) { + final StorageLocation.ReservationHold h = + loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, name)); + if (h != null) { + acquired.put(name, h); + uncommittedHolds.add(h); + } + } + + // Phase 3: commit under entryLock. Ownership transfers are done by removing from `uncommittedHolds` after + // installing into the field. Anything left in `uncommittedHolds` at the end lost a race and gets released + // in Phase 4 alongside `displacedHolds` (the pre-existing holds diffed out of ruleBundleHolds). + entryLock.lock(); + try { + if (newSelfHold != null) { + if (metadataSelfHold == null) { + metadataSelfHold = newSelfHold; + uncommittedHolds.remove(newSelfHold); + } + // else: a concurrent applyRule installed a self-hold (shouldn't happen under segmentLock, but defensive). + // newSelfHold stays in uncommittedHolds, gets released in Phase 4. + } + for (String name : namesToRelease) { + final StorageLocation.ReservationHold h = ruleBundleHolds.remove(name); + if (h != null) { + displacedHolds.add(h); + } + } + for (var e : acquired.entrySet()) { + if (!ruleBundleHolds.containsKey(e.getKey())) { + ruleBundleHolds.put(e.getKey(), e.getValue()); + uncommittedHolds.remove(e.getValue()); + } + // else: a concurrent registerBundle raced ahead and installed a hold for this name; ours stays in + // uncommittedHolds, gets released in Phase 4. + } + ruleFingerprint = fingerprint; + ruleSelectedBundleNames = newSelection; + } + finally { + entryLock.unlock(); + } + } + finally { + // Phase 4: release outside entryLock. `uncommittedHolds` contains every hold this call acquired but did not + // successfully transfer into a field (race-lose, mid-acquire failure, or commit failure). `displacedHolds` + // contains pre-existing holds that the diff removed from ruleBundleHolds after successful commit. + releaseHolds(uncommittedHolds); + releaseHolds(displacedHolds); + } + } + + /** + * Release the partial-load rule applied to this entry. Closes every {@link StorageLocation.ReservationHold} taken by + * {@link #applyRule(String, Set)} (the metadata self-hold and every bundle hold) and clears the rule state. Bundle + * entries that were kept resident only by the rule become eviction eligible; the metadata entry itself becomes + * eviction eligible as well (unless a concurrent query holds a reference). Safe to call when no rule is applied, no-op. + *

+ * Concurrency contract. Must be serialized against {@link #applyRule} for a given segment id via an external + * per-segment lock (see {@code applyRule}'s Concurrency contract). Hold releases run OUTSIDE {@link #entryLock}, so + * a concurrent {@code applyRule}/{@code clearRule} on the same entry can observe intermediate state. + *

+ * Snapshot-under-lock, release-outside-lock. Closing a {@link StorageLocation.ReservationHold} eventually acquires + * the location's writeLock (via {@code createWeakEntryReleaseRunnable}), so it must not run under {@link #entryLock} + * lest it invert the writeLock → entryLock ordering the storage layer already relies on. + */ + public void clearRule() + { + final List> toClose = new ArrayList<>(); + entryLock.lock(); + try { + toClose.addAll(ruleBundleHolds.values()); + ruleBundleHolds.clear(); + if (metadataSelfHold != null) { + toClose.add(metadataSelfHold); + metadataSelfHold = null; + } + ruleFingerprint = null; + ruleSelectedBundleNames = Set.of(); + } + finally { + entryLock.unlock(); + } + releaseHolds(toClose); + } + + private void releaseHolds(List> holds) + { + for (StorageLocation.ReservationHold h : holds) { + CloseableUtils.closeAndSuppressExceptions( + h, + t -> LOG.warn(t, "Failed releasing rule-hold on partial segment[%s]", segmentId) + ); + } + } + + @Nullable + private PartialSegmentBundleCacheEntry findLinkedBundleByName(String bundleName) + { + return linkedBundles.get(bundleName); + } + + /** + * Whether every container belonging to {@code bundleName} has been downloaded. + *

+ * Returns {@code false} when the entry is not mounted (no file mapper yet). + */ + public boolean isBundleFullyDownloaded(String bundleName) + { + final PartialSegmentFileMapperV10 mapper = getFileMapper(); + return mapper != null && mapper.isBundleFullyDownloaded(bundleName); + } + + /** + * Whether {@code bundleName} is currently held by an applied partial-load rule on this entry (i.e. present in + * {@link #ruleBundleHolds}). + *

+ * NOT the same as "is the bundle linked": a bundle can be linked ({@link #linkedBundles}) without being rule-held, + * because {@link #linkedBundles} is a lock-free {@link ConcurrentHashMap} keyset that + * {@link #registerBundle}/{@link #unregisterBundle} mutate outside {@link #entryLock} (to preserve lock ordering + * against the storage layer's writeLock → entryLock chain). + */ + public boolean isBundleRuleHeld(String bundleName) + { + entryLock.lock(); + try { + return ruleBundleHolds.containsKey(bundleName); + } + finally { + entryLock.unlock(); + } + } + + /** + * Mount and eagerly download the named bundle. Drives an acquire through the internal + * {@link #bundleAcquirer} (which mounts the bundle if needed, causing {@link #registerBundle} to fire), downloads + * every container in the bundle via {@link PartialSegmentFileMapperV10#ensureBundleDownloaded}, then releases the + * transient acquire hold. When a partial-load rule selects this bundle, the metadata's own rule-hold (installed by + * {@link #registerBundle} on the mount that this call drives) keeps the bundle resident after the transient hold + * releases. When no rule selects this bundle, the transient hold's release leaves the bundle eviction eligible; the + * caller should already have driven {@link #applyRule} first to install the rule state before calling this. + */ + public void ensureBundleResidentForRule(String bundleName) throws IOException + { + entryLock.lock(); + try { + if (location == null || fileMapper == null) { + throw DruidException.defensive( + "ensureBundleResidentForRule on unmounted partial metadata entry[%s]", + id + ); + } + } + finally { + entryLock.unlock(); + } + final Closeable acquired = bundleAcquirer.acquire(bundleName); + try { + final PartialSegmentFileMapperV10 mapper = getFileMapper(); + if (mapper == null) { + throw DruidException.defensive( + "Partial metadata entry[%s] lost its file mapper mid-acquire for bundle[%s]", + id, + bundleName + ); + } + mapper.ensureBundleDownloaded(bundleName); + } + finally { + CloseableUtils.closeAndSuppressExceptions(acquired, t -> LOG.warn( + t, + "Failed to release transient bundle-acquire hold for bundle[%s] on segment[%s]", + bundleName, + segmentId + )); + } + } + @Override public long getSize() { @@ -273,9 +624,12 @@ public void mount(StorageLocation mountLocation) throws IOException ours.set(null); } catch (Throwable t) { - // clear the future so the next caller gets a fresh attempt + // Clear the future so the next caller gets a fresh attempt. Signal awaiters with the exception BEFORE firing + // the onUnmount hook so any thread observing the failure via awaitMount(future) sees the exception before + // observing the hook's file-system side effect (info-file deletion). mountFuture.set(null); ours.setException(t); + runOnUnmountHookOnce(); Throwables.propagateIfInstanceOf(t, IOException.class); Throwables.propagateIfPossible(t); throw DruidException.defensive(t, "Failed to mount metadata entry[%s]", id); @@ -463,6 +817,17 @@ 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. + *

+ * Mid-mount race guard. A concurrent reclaim can invoke {@code unmount()} while {@link #doMount} is still + * running its slow deep-storage header fetch and has not yet installed {@link #references}. Firing the hook here + * would race with {@code doMount}. Detect this via {@link #mountFuture} (set by {@link #mount} before + * {@code doMount} runs; cleared by {@code doMount}'s failure path and by {@link #doActualUnmount}) and defer + * hook-firing to the mount path: on success {@code verifyStillReservedOrRollback} calls {@code unmount()} again + * with {@link #references} installed and the hook fires via {@link #doActualUnmount}; on failure {@link #mount}'s + * catch block fires the hook directly (see below). */ @Override public void unmount() @@ -470,6 +835,30 @@ public void unmount() final ReferenceCountingCloseableObject current = references.get(); if (current != null && !current.isClosed()) { current.close(); + return; + } + if (mountFuture.get() != null) { + // doMount is in flight and has not yet installed `references`; hook-firing is deferred to the mount path. + 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. + */ + 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); + } } } @@ -634,12 +1023,21 @@ public Optional acquireFullReference(@Nullable Closeable extraOnClose) */ private void doActualUnmount() { - final Runnable hook; entryLock.lock(); try { if (fileMapper == null) { return; } + // doActualUnmount must never fire while a partial-load rule is applied. + if (ruleFingerprint != null) { + throw DruidException.defensive( + "doActualUnmount fired for partial segment[%s] with rule state still applied (fingerprint=%s, " + + "selectedBundles=%s); the metadata self-hold invariant has been broken", + segmentId, + ruleFingerprint, + ruleSelectedBundleNames + ); + } try { fileMapper.close(); } @@ -654,21 +1052,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 +1194,7 @@ public Closeable acquire(String requestedBundleName) () -> PartialSegmentBundleCacheEntry.forBundle( PartialSegmentMetadataCacheEntry.this, bundleName, - inferParentBundles(bundleName) + inferBundleDependencies(bundleName) ) ); if (hold == null) { @@ -861,13 +1251,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 +1275,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); + * 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 +1302,34 @@ public List inferParentBundles(String ); } + /** + * Return {@code roots} plus every transitive dependency (as inferred by {@link #inferBundleDependencies}) in + * dependency-first order. + *

+ * 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. @@ -925,19 +1345,71 @@ private boolean hasBaseBundle() * Register a bundle entry as a current dependent of this metadata entry. Called by * {@link PartialSegmentBundleCacheEntry} after a successful mount; the drop path uses {@link #snapshotLinkedBundles} * to enumerate dependents for cascade-close. + *

+ * If a partial-load rule is currently applied and {@code bundle}'s name is in the rule's selection set, take a + * {@link StorageLocation.ReservationHold} on the bundle so it stays resident until the rule is cleared or swapped + * away. The hold acquisition uses a snapshot-under-lock, acquire-outside-lock, commit-under-lock pattern to avoid + * holding {@link #entryLock} while calling into {@link StorageLocation}'s read/write locks. */ void registerBundle(PartialSegmentBundleCacheEntry bundle) { - linkedBundles.add(bundle); + final String bundleName = bundle.getBundleName(); + linkedBundles.put(bundleName, bundle); + + // Phase 1: snapshot under entryLock. + final StorageLocation loc; + entryLock.lock(); + try { + if (location == null + || !ruleSelectedBundleNames.contains(bundleName) + || ruleBundleHolds.containsKey(bundleName)) { + return; + } + loc = location; + } + finally { + entryLock.unlock(); + } + + // Phase 2: acquire outside entryLock. Caller's mount path (via bundleAcquirer.acquire) already holds a transient + // ReservationHold on this bundle, so the weak entry is guaranteed present and this acquire cannot return null + // for a "just evicted" reason. + final StorageLocation.ReservationHold hold = + loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, bundleName)); + if (hold == null) { + return; + } + + // Phase 3: commit under entryLock, or discard if the rule state moved under us. + boolean stored = false; + entryLock.lock(); + try { + if (ruleSelectedBundleNames.contains(bundleName) && !ruleBundleHolds.containsKey(bundleName)) { + ruleBundleHolds.put(bundleName, hold); + stored = true; + } + } + finally { + entryLock.unlock(); + } + if (!stored) { + CloseableUtils.closeAndSuppressExceptions( + hold, + t -> LOG.warn(t, "Failed releasing discarded rule-hold for bundle[%s] on segment[%s]", bundleName, segmentId) + ); + } } /** * Reverse of {@link #registerBundle}. Called by {@link PartialSegmentBundleCacheEntry#unmount} so the metadata's - * view stays consistent with which bundles are actually mounted. + * view stays consistent with which bundles are actually mounted. Bookkeeping-only: never touches {@link + * #ruleBundleHolds}. */ void unregisterBundle(PartialSegmentBundleCacheEntry bundle) { - linkedBundles.remove(bundle); + // Remove only if the exact same instance is still registered under this name. A late-arriving unregister for a + // stale bundle instance (e.g. one that was replaced by a fresh remount) must not evict the fresh instance. + linkedBundles.remove(bundle.getBundleName(), bundle); } /** @@ -947,7 +1419,7 @@ void unregisterBundle(PartialSegmentBundleCacheEntry bundle) */ public Collection snapshotLinkedBundles() { - return new ArrayList<>(linkedBundles); + return new ArrayList<>(linkedBundles.values()); } /** 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..54c0375c9dc8 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 @@ -54,15 +54,21 @@ import java.io.InputStream; import java.nio.file.Files; import java.nio.file.StandardCopyOption; +import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.Objects; import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CancellationException; 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.RejectedExecutionException; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.Lock; @@ -746,6 +752,48 @@ private ReservedPartial findExistingPartialWithHold(SegmentId segmentId) return null; } + /** + * If any storage location has a non-{@link PartialSegmentMetadataCacheEntry} cache entry at {@code segmentId}, + * attempt to evict it via {@link StorageLocation#removeUnheldWeakEntry}. This handles the config-flip scenario + * where a segment with a {@link PartialLoadSpec} wrapper was previously queried while + * {@code virtualStoragePartialDownloadsEnabled=false} ({@link #acquireSegment} created a + * {@link CompleteSegmentCacheEntry} at this id, and a subsequent partial-load rule cannot reserve here without + * evicting the stale entry first). + *

+ * {@link StorageLocation#removeUnheldWeakEntry} is a no-op when the entry is held (in-flight query), so we + * detect via {@link StorageLocation#getCacheEntry} whether eviction actually happened. If not, throw a retryable + * {@link SegmentLoadingException}: the coordinator's load queue retries on the next sync, and by then the query + * should have released the hold and eviction will succeed. + */ + private void evictStaleNonPartialWeakEntry(SegmentId segmentId) throws SegmentLoadingException + { + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segmentId); + for (StorageLocation location : locations) { + final CacheEntry entry = location.getCacheEntry(id); + if (entry == null || entry instanceof PartialSegmentMetadataCacheEntry) { + continue; + } + location.removeUnheldWeakEntry(id); + final CacheEntry stillThere = location.getCacheEntry(id); + if (stillThere != null) { + throw new SegmentLoadingException( + "Stale non-partial cache entry[%s] at id[%s] on location[%s] blocks partial-load reservation and is " + + "currently held; the coordinator's load queue will retry", + stillThere.getClass().getSimpleName(), + id, + location.getPath() + ); + } + log.info( + "Evicted stale non-partial cache entry[%s] at location[%s] to make room for partial-load rule on " + + "segment[%s]", + entry.getClass().getSimpleName(), + location.getPath(), + segmentId + ); + } + } + /** * Try to open a range reader for the given segment's {@link LoadSpec}; returns {@code null} when the backend * doesn't support range reads (e.g. zipped storage), or when partial downloads are disabled via @@ -790,7 +838,14 @@ private ReservedPartial reservePartial(DataSegment dataSegment, SegmentRangeRead FileUtils.mkdirp(partialDir); } catch (IOException e) { - throw DruidException.defensive(e, "Failed to create partial cache dir for segment[%s]", dataSegment.getId()); + // Location is unwritable, fall through to next location rather than failing the whole reservation + log.warn( + e, + "Failed to create partial cache dir on location[%s] for segment[%s]; trying next location", + location.getPath(), + dataSegment.getId() + ); + continue; } final StorageLocation.ReservationHold hold = location.addWeakReservationHold( id, @@ -806,6 +861,7 @@ private ReservedPartial reservePartial(DataSegment dataSegment, SegmentRangeRead ) ); if (hold == null) { + atomicMoveAndDeleteCacheEntryDirectory(partialDir); continue; } try { @@ -817,19 +873,28 @@ private ReservedPartial reservePartial(DataSegment dataSegment, SegmentRangeRead location.getPath() ); } + // Unconditionally write the info file with the incoming DataSegment. An existing info file at this path + // may carry a STALE loadSpec. Skipping the write when the file exists would preserve that stale wrapper on + // disk, and a subsequent bootstrap would restore the segment via the wrong path (non-partial) or drift out of + // sync with the coordinator's applied rule. writeAtomically does the file-rename dance so a concurrent read + // sees either the old or new complete file, never a partial write. final File segmentInfoCacheFile = new File(getEffectiveInfoDir(), dataSegment.getId().toString()); - if (!segmentInfoCacheFile.exists()) { - FileUtils.mkdirp(getEffectiveInfoDir()); - FileUtils.writeAtomically(segmentInfoCacheFile, out -> { - jsonMapper.writeValue(out, dataSegment); - return null; - }); - } + FileUtils.mkdirp(getEffectiveInfoDir()); + FileUtils.writeAtomically(segmentInfoCacheFile, out -> { + jsonMapper.writeValue(out, dataSegment); + return null; + }); partial.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); return new ReservedPartial(partial, location, hold); } catch (Throwable t) { - throw CloseableUtils.closeAndWrapInCatch(t, hold); + // Close the hold (removing the never-mounted weak entry), then nuke the on-disk dir. + try { + throw CloseableUtils.closeAndWrapInCatch(t, hold); + } + finally { + atomicMoveAndDeleteCacheEntryDirectory(partialDir); + } } } throw DruidException.forPersona(DruidException.Persona.USER) @@ -872,6 +937,290 @@ private record ReservedPartial( { } + /** + * Apply (or reconcile) a partial-load rule for a segment coming through {@link #load}. The rule state lives on the + * {@link PartialSegmentMetadataCacheEntry} itself; this method just drives it: + *

    + *
  1. Materialize the wrapper's {@link PartialLoadSpec} and open a range reader over the backend. If the backend + * can't do range reads, clear any prior rule on the segment and return; the segment falls back to the + * weak-full-load path at query time.
    2. + *
  2. Acquire (or reserve) the metadata entry via the shared on-demand path ({@link #findOrReservePartial}), + * take a transient eviction-protective hold across the rest of the flow.
    3. + *
  3. Mount the metadata entry (idempotent).
    4. + *
  4. Compute the rule's selected bundle names from the wrapper + parsed segment metadata, and call + * {@link PartialSegmentMetadataCacheEntry#applyRule} which installs the metadata self-hold and takes + * holds on any already-registered selected bundles.
    5. + *
  5. Kick off async eager downloads on the loading pool for any selected bundle not yet linked. Each async task + * drives {@link PartialSegmentMetadataCacheEntry#ensureBundleResidentForRule}: its bundle mount will call + * {@link PartialSegmentMetadataCacheEntry#registerBundle}, which acquires the rule-hold on that bundle.
    6. + *
  6. Release the transient hold; the metadata's self-hold keeps the entry resident under the applied rule.
    7. + *
+ */ + private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException + { + final PartialLoadSpec wrapper = materializePartialLoadSpec(dataSegment); + final SegmentRangeReader rangeReader = openPartialRangeReader(dataSegment, wrapper); + final ReferenceCountingLock lock = lock(dataSegment); + synchronized (lock) { + try { + // If a stale non-partial cache entry sits at this segment id (a CompleteSegmentCacheEntry created by a prior + // acquireSegment while virtualStoragePartialDownloadsEnabled=false, for example), evict it before any + // partial-entry lookup or reservation; otherwise findExistingPartialWithHold's defensive type-check would + // throw, and reservePartial's addWeakReservationHold would land on the incompatible entry. If the stale + // entry is currently held (in-flight query), this throws a retryable SegmentLoadingException; the + // coordinator's load queue retries on next sync, and by then the query should have released. + evictStaleNonPartialWeakEntry(dataSegment.getId()); + + if (rangeReader == null) { + // Backend doesn't support range reads (e.g. zipped deep storage). The rule can't be honored as a partial + // load; clear any prior rule so the segment falls through to the ordinary weak-full-load path at query + // time. + final ReservedPartial existing = findExistingPartialWithHold(dataSegment.getId()); + if (existing != null) { + try { + existing.metadata().clearRule(); + log.warn( + "Backend for segment[%s] does not support range reads; released rule[fingerprint=%s], segment " + + "will fall back to weak full-load at query time", + dataSegment.getId(), + wrapper.getFingerprint() + ); + } + finally { + CloseableUtils.closeAndSuppressExceptions( + existing.hold(), + t -> log.warn(t, "Failed to release transient hold on partial metadata for segment[%s]", + dataSegment.getId()) + ); + } + } else { + // No prior entry AND range-reader is null on this fresh load: the coordinator asked for a partial rule + // this historical can't honor. Log so operators can debug "why isn't my rule applied on historical X". + log.warn( + "Backend for segment[%s] does not support range reads and no prior partial entry exists; rule" + + "[fingerprint=%s] cannot be applied on this historical", + dataSegment.getId(), + wrapper.getFingerprint() + ); + } + return; + } + + final ReservedPartial reserved = findOrReservePartial(dataSegment, rangeReader); + try { + final PartialSegmentMetadataCacheEntry metadata = reserved.metadata(); + try { + metadata.mount(reserved.location()); + } + catch (IOException e) { + throw new SegmentLoadingException( + e, + "Failed to mount partial metadata for segment[%s]", + dataSegment.getId() + ); + } + final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper(); + if (mapper == null) { + throw DruidException.defensive( + "Partial metadata for segment[%s] mounted without a file mapper", + dataSegment.getId() + ); + } + // Register the info-file cleanup hook BEFORE anything that could throw. Any throw between mount and + // awaitEagerDownloadsOrClearRule leaves the metadata entry weak-reserved with the hook attached; when + // cache eventually reclaims, the info file gets deleted along with the entry. + metadata.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); + final Set selected = Set.copyOf( + wrapper.getSelectedBundleNames(dataSegment, mapper.getSegmentFileMetadata()) + ); + final String priorFingerprint = metadata.getRuleFingerprint(); + metadata.applyRule(wrapper.getFingerprint(), selected); + if (priorFingerprint != null && !priorFingerprint.equals(wrapper.getFingerprint())) { + log.info( + "Reconciled partial-load rule for segment[%s]: fingerprint transitioned [%s] → [%s]", + dataSegment.getId(), + priorFingerprint, + wrapper.getFingerprint() + ); + } + // Block until every eager download completes so the announcement fingerprint reflects reality: any failure + // clears the rule state (releasing self-hold + all bundle rule-holds) and propagates as a load failure so + // the coordinator's load queue can retry on its next sync. The announced fingerprint == "rule fully + // realized" contract stays intact. + awaitEagerDownloadsOrClearRule(dataSegment, metadata, selected); + } + finally { + CloseableUtils.closeAndSuppressExceptions( + reserved.hold(), + t -> log.warn(t, "Failed to release transient hold on partial metadata for segment[%s]", + dataSegment.getId()) + ); + } + } + finally { + unlock(dataSegment, lock); + } + } + } + + /** + * Submit eager-download tasks to the loading pool for every rule-selected bundle not yet registered with + * {@code metadata}, then block until every task completes. On any failure the rule state is cleared + * and a {@link SegmentLoadingException} is thrown so the caller treats the load as failed and retries. + */ + private void awaitEagerDownloadsOrClearRule( + DataSegment dataSegment, + PartialSegmentMetadataCacheEntry metadata, + Set selected + ) throws SegmentLoadingException + { + final List> pending = new ArrayList<>(); + Throwable firstFailure = null; + try { + for (String bundleName : selected) { + // Skip only when the bundle is BOTH rule-held AND fully downloaded (every container's files present on disk). + if (metadata.isBundleRuleHeld(bundleName) && metadata.isBundleFullyDownloaded(bundleName)) { + continue; + } + pending.add(virtualStorageLoadingThreadPool.getExecutorService().submit(() -> { + metadata.ensureBundleResidentForRule(bundleName); + return null; + })); + } + } + catch (RejectedExecutionException e) { + // Pool shutting down or queue saturated. Any already-submitted futures may still run; track the rejection as + // firstFailure and fall through to the await/cancel + clearRule path so state is always cleaned up before we + // throw. + firstFailure = e; + } + + for (Future f : pending) { + if (firstFailure != null) { + // Cancel remaining futures on any failure so we don't wait for tasks whose result we no longer intend to + // commit. cancel(false) — NOT cancel(true) — because a running task is mid-NIO/FileChannel read via + // mapper.ensureBundleDownloaded, and Thread.interrupt() on an in-flight NIO op raises + // ClosedByInterruptException that closes the mapper's shared underlying FD, breaking still-mounted bundles + // for other queries reading the same segment. Not-yet-started tasks are marked CANCELLED; running tasks + // continue on their own timeline, observe ruleSelectedBundleNames == {} after clearRule, and release + // their transient bundleAcquirer holds without acquiring rule-holds. + f.cancel(false); + // Still drain the future so an ExecutionException from a task that completed-with-failure BEFORE we + // canceled is logged instead of silently swallowed. + try { + f.get(); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + catch (ExecutionException e) { + final Throwable cause = e.getCause() != null ? e.getCause() : e; + log.warn( + cause, + "Eager download of a rule-selected bundle for segment[%s] failed after the primary failure", + dataSegment.getId() + ); + } + catch (CancellationException e) { + // task never started; expected on the cancel path + } + continue; + } + try { + f.get(); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + firstFailure = e; + // Signal cancel WITHOUT interrupt — see cancel(false) comment above; interrupting mid-NIO closes the + // shared mapper FD. + f.cancel(false); + } + catch (ExecutionException e) { + firstFailure = e.getCause() != null ? e.getCause() : e; + } + catch (CancellationException e) { + // Defensive: shouldn't reach here on the first-failure branch (we only cancel after firstFailure is set), + // but if it does, treat it uniformly. + firstFailure = e; + } + } + + if (firstFailure != null) { + // Any late-completing pool task that still succeeds after we've cleared the rule will call registerBundle → + // observing ruleSelectedBundleNames == {} and skipping the rule-hold acquire. + metadata.clearRule(); + throw new SegmentLoadingException( + firstFailure, + "Failed eager download of rule-selected bundles for segment[%s]; cleared partial-load rule", + dataSegment.getId() + ); + } + } + + /** + * Materialize the segment's wrapped load spec to a {@link PartialLoadSpec}. + */ + 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 via the wrapper's {@link PartialLoadSpec#openRangeReader} + * (which delegates to the inner load spec), wrapping any {@link IOException} as a {@link SegmentLoadingException}. + * Returns {@code null} if the backend doesn't support range reads. + */ + @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()); + } + } + + /** + * The fingerprint of the currently-applied partial-load rule for {@code segmentId}, or {@code null} if no rule is + * currently applied (or the segment has no partial metadata entry (e.g. it was eager-loaded as a complete cache + * entry). Test-only. Delegates to {@link PartialSegmentMetadataCacheEntry#getRuleFingerprint} on the segment's + * cache entry. + */ + @VisibleForTesting + @Nullable + String getRuleFingerprintForSegment(SegmentId segmentId) + { + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segmentId); + for (StorageLocation location : locations) { + final CacheEntry entry = location.getCacheEntry(id); + if (entry instanceof PartialSegmentMetadataCacheEntry partial) { + return partial.getRuleFingerprint(); + } + } + return null; + } + @Nullable private AcquireSegmentAction acquireExistingSegment(SegmentCacheEntryIdentifier identifier) { @@ -927,6 +1276,15 @@ public void load(final DataSegment dataSegment) throws SegmentLoadingException "load() should not be called when virtualStorageIsEphemeral is true" ); } + // Partial-load-rule routing: a wrapped load spec carrying a fingerprint + per-spec selection (cluster groups, + // projections, ...) means the coordinator wants this segment loaded with rule holds pinning the selected + // bundles + metadata. Install the holds via loadPartial; other segments (no wrapper, or partials disabled) + // take the existing weak/no-op path below. + if (config.isVirtualStoragePartialDownloadsEnabled() + && PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec())) { + loadPartial(dataSegment); + return; + } // virtual storage doesn't do anything with loading immediately, but check to see if the segment is already cached // and if so, clear out the onUnmount action final ReferenceCountingLock lock = lock(dataSegment); @@ -999,6 +1357,17 @@ public void bootstrap( dataSegment.getId() ); } + // If the persisted DataSegment's loadSpec is a partial-load-rule wrapper, reapply the rule so the + // segment continues to be protected from eviction across the restart. + // + // Strict failure handling: reapplyRuleFromInfoFile throws on eager-download failure (and clears the + // rule state before throwing). The exception propagates up through SegmentManager.loadSegmentOnBootstrap, + // which calls cacheManager.drop and rethrows, so SegmentCacheBootstrapper marks the segment as failed + // and skips its announcement. + if (config.isVirtualStoragePartialDownloadsEnabled() + && PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec())) { + reapplyRuleFromInfoFile(dataSegment, partial); + } } else { throw DruidException.defensive( "Unexpected cache entry type[%s] for segment[%s] during bootstrap", @@ -1066,15 +1435,71 @@ public File getSegmentFiles(final DataSegment segment) @Override 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); + final ReferenceCountingLock lock = lock(segment); + synchronized (lock) { + try { + // partial-load-rule cleanup: if the segment's cache entry is a partial metadata entry with an applied rule, + // clear it + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segment.getId()); + for (StorageLocation location : locations) { + final CacheEntry entry = location.getCacheEntry(id); + if (entry instanceof PartialSegmentMetadataCacheEntry partial) { + partial.clearRule(); + // Force synchronous cleanup: clearRule releases the self-hold, but the info-file-deletion hook only + // fires on doActualUnmount which is triggered by the phaser hitting zero. removeUnheldWeakEntry + // unlinks the weak entry from cache and terminates the phaser now, firing the hook (and the mapper + // teardown) on the caller's thread. No-op if a concurrent query still holds the entry. + location.removeUnheldWeakEntry(id); + } + } + // full-segment drop path: release any {@link CompleteSegmentCacheEntry} still reserved. + for (StorageLocation location : locations) { + final CacheEntry entry = location.getCacheEntry(id); + if (entry != null) { + location.release(entry); + } + } + } + finally { + unlock(segment, lock); } } } + /** + * Reapply the persisted partial-load rule to a bootstrap-restored metadata entry. Reads the wrapper from the + * segment's info-file {@code loadSpec}, resolves the selected bundle names against the just-parsed on-disk + * metadata header, calls {@link PartialSegmentMetadataCacheEntry#applyRule}, then drives eager downloads for any + * selected bundle that wasn't restored from disk (via {@link #awaitEagerDownloadsOrClearRule}). On failure the + * exception marks the segment as failed → doesn't announce it → the coordinator's next sync re-issues load. + */ + private void reapplyRuleFromInfoFile(DataSegment dataSegment, PartialSegmentMetadataCacheEntry partial) + throws SegmentLoadingException + { + // evict any stale non-partial cache entry at this segment id before touching the partial state. + // Bootstrap ordering today produces a partial entry for the passed-in `partial`, so this is expected + // to be a no-op except during a config switch. + evictStaleNonPartialWeakEntry(dataSegment.getId()); + + final PartialLoadSpec wrapper = materializePartialLoadSpec(dataSegment); + final PartialSegmentFileMapperV10 mapper = partial.getFileMapper(); + if (mapper == null) { + throw DruidException.defensive( + "Bootstrap-restored partial metadata for segment[%s] has no file mapper", dataSegment.getId() + ); + } + // Register the info-file cleanup hook BEFORE anything that could throw. reserveFromDisk does NOT install an + // onUnmount hook on bootstrap-restored entries, so without this ordering a throw from getSelectedBundleNames + // or applyRule would leave the entry with no cleanup path. Setting the hook first ensures every teardown path + // deletes the info file consistently. + partial.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); + final Set selected = Set.copyOf( + wrapper.getSelectedBundleNames(dataSegment, mapper.getSegmentFileMetadata()) + ); + partial.applyRule(wrapper.getFingerprint(), selected); + awaitEagerDownloadsOrClearRule(dataSegment, partial, selected); + } + @Override public void shutdownBootstrap() { 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..5beeb750b4e9 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,142 @@ void testOnUnmountHookRunsAfterStorageLocationCleanup() throws IOException ); } + @Test + void testOnUnmountHookRunsOnReleaseBeforeMount() + { + // if a caller sets the onUnmount hook before mount() and then releases the reservation on a mount-failure path + // (or without ever calling mount), unmount() must still fire the hook — otherwise external cleanup + // (e.g. info-file deletion) would leak. + 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 testBundlesInMountOrderReturnsRootsInInputOrderWhenNoBase() + { + // With no __base bundle in the mapped segment (entry not mounted), inferBundleDependencies returns [] for every + // input, so the walker returns roots in their original iteration order. + final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE); + Assertions.assertEquals( + List.of("a", "b", "c"), + entry.bundlesInMountOrder(List.of("a", "b", "c")) + ); + } + + @Test + void testBundlesInMountOrderExpandsWithBaseWhenPresent() throws IOException + { + // A segment carrying __base: dependents' expansion places __base first, then the dependent itself. + final PartialSegmentMetadataCacheEntry entry = mountedEntryOver( + buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "some_projection") + ); + Assertions.assertEquals( + List.of(Projections.BASE_TABLE_PROJECTION_NAME, "some_projection"), + entry.bundlesInMountOrder(List.of("some_projection")) + ); + } + + @Test + void testBundlesInMountOrderDedupsBaseAcrossMultipleDependents() throws IOException + { + // When multiple dependents share a common dependency, mount order emits it exactly once at the front. + final PartialSegmentMetadataCacheEntry entry = mountedEntryOver( + buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a", "projection_b") + ); + Assertions.assertEquals( + List.of(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a", "projection_b"), + entry.bundlesInMountOrder(List.of("projection_a", "projection_b")) + ); + } + + @Test + void testRuleStateEmptyByDefault() + { + final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE); + Assertions.assertFalse(entry.isRuleHeld()); + Assertions.assertNull(entry.getRuleFingerprint()); + Assertions.assertEquals(Set.of(), entry.getRuleSelectedBundleNames()); + } + + @Test + void testApplyRuleBeforeMountThrows() + { + final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE); + Assertions.assertThrows( + DruidException.class, + () -> entry.applyRule("fp", Set.of("bundle")) + ); + } + + @Test + void testApplyRuleSetsFingerprintAndSelection() throws IOException + { + final PartialSegmentMetadataCacheEntry entry = mountedWeakEntryOver( + buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a") + ); + entry.applyRule("fp1", Set.of("projection_a")); + Assertions.assertTrue(entry.isRuleHeld()); + Assertions.assertEquals("fp1", entry.getRuleFingerprint()); + Assertions.assertEquals(Set.of("projection_a"), entry.getRuleSelectedBundleNames()); + } + + @Test + void testApplyRuleSameArgsIsIdempotent() throws IOException + { + final PartialSegmentMetadataCacheEntry entry = mountedWeakEntryOver( + buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a") + ); + entry.applyRule("fp1", Set.of("projection_a")); + entry.applyRule("fp1", Set.of("projection_a")); + Assertions.assertEquals("fp1", entry.getRuleFingerprint()); + Assertions.assertEquals(Set.of("projection_a"), entry.getRuleSelectedBundleNames()); + } + + @Test + void testApplyRuleSwapUpdatesFingerprintAndSelection() throws IOException + { + final PartialSegmentMetadataCacheEntry entry = mountedWeakEntryOver( + buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a", "projection_b") + ); + entry.applyRule("fp1", Set.of("projection_a")); + entry.applyRule("fp2", Set.of("projection_b")); + Assertions.assertEquals("fp2", entry.getRuleFingerprint()); + Assertions.assertEquals(Set.of("projection_b"), entry.getRuleSelectedBundleNames()); + } + + @Test + void testClearRuleReleasesState() throws IOException + { + final PartialSegmentMetadataCacheEntry entry = mountedWeakEntryOver( + buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a") + ); + entry.applyRule("fp1", Set.of("projection_a")); + Assertions.assertTrue(entry.isRuleHeld()); + entry.clearRule(); + Assertions.assertFalse(entry.isRuleHeld()); + Assertions.assertNull(entry.getRuleFingerprint()); + Assertions.assertEquals(Set.of(), entry.getRuleSelectedBundleNames()); + } + + @Test + void testClearRuleOnNoRuleIsNoop() throws IOException + { + final PartialSegmentMetadataCacheEntry entry = mountedWeakEntryOver( + buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME) + ); + entry.clearRule(); + Assertions.assertFalse(entry.isRuleHeld()); + } + @Test void testConstructorRejectsNonPositiveEstimate() { @@ -385,27 +522,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 +551,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 +574,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) @@ -472,8 +609,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 +630,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 { @@ -515,4 +652,28 @@ private PartialSegmentMetadataCacheEntry mountedEntryOver(File deepStorageDir) t return entry; } + /** + * Variant of {@link #mountedEntryOver} that uses {@link StorageLocation#reserveWeak} so the mounted entry is a real + * weak reservation — needed by the rule-holds state machine, which calls + * {@link StorageLocation#addWeakReservationHoldIfExists} on itself when {@code applyRule} runs. + */ + private PartialSegmentMetadataCacheEntry mountedWeakEntryOver(File deepStorageDir) throws IOException + { + final File cache = new File(tempDir, "cache_" + (fixtureSeq++)); + FileUtils.mkdirp(cache); + final StorageLocation location = new StorageLocation(cache, ESTIMATE * 4, null); + final PartialSegmentMetadataCacheEntry entry = new PartialSegmentMetadataCacheEntry( + SEGMENT_ID, + cache, + IndexIO.V10_FILE_NAME, + List.of(), + new DirectoryBackedRangeReader(deepStorageDir), + JSON_MAPPER, + null, + ESTIMATE + ); + Assertions.assertTrue(location.reserveWeak(entry)); + entry.mount(location); + return entry; + } } 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/SegmentLocalCacheManagerPartialRuleLoadTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java new file mode 100644 index 000000000000..ea833e59270c --- /dev/null +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java @@ -0,0 +1,666 @@ +/* + * 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.SegmentLazyLoadFailCallback; +import org.apache.druid.segment.TestHelper; +import org.apache.druid.segment.column.ColumnConfig; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.data.CompressionStrategy; +import org.apache.druid.segment.file.PartialSegmentFileMapperV10; +import org.apache.druid.segment.incremental.IncrementalIndexSchema; +import org.apache.druid.segment.projections.Projections; +import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory; +import org.apache.druid.server.metrics.NoopServiceEmitter; +import org.apache.druid.timeline.DataSegment; +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.timeline.partition.NoneShardSpec; +import org.joda.time.DateTime; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.File; +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ThreadLocalRandom; + +/** + * Manager-level tests for the partial-load-rule {@link SegmentLocalCacheManager#load} path: a segment whose + * {@code loadSpec} is a {@link PartialLoadSpec} wrapper drives + * {@link PartialSegmentMetadataCacheEntry#applyRule} on the metadata entry, which installs a self-referential + * {@link StorageLocation.ReservationHold} to keep the metadata resident and holds on each selected bundle already + * registered. Selected bundles not yet registered get eager-downloads submitted to the loading pool; {@code load} + * blocks until every eager download completes so the announced fingerprint reflects "rule fully realized". On any + * eager-download failure the rule state is cleared and a {@link SegmentLoadingException} propagates so the + * coordinator's load queue retries. The cache entries themselves live in {@code weakCacheEntries} exactly as they + * would from an on-demand acquire; the rule-holds keep them from being SIEVE-evicted for as long as the coordinator + * considers the segment loaded. + *

+ * {@code drop} calls {@link PartialSegmentMetadataCacheEntry#clearRule} — entries become unheld weak (subject to any + * remaining query holds) and SIEVE reclaims them naturally. + */ +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) { + // Drop the segment to release rule-holds and unmount mmap'd bundle files before @TempDir tries to clean up + // the cache directory. Without this, on some platforms the temp-dir cleanup fails on still-mapped files. + try { + manager.drop(partialWrapperSegment(List.of(AGG_BUNDLE))); + } + catch (Throwable ignored) { + // best-effort — some tests never installed a rule on this segment + } + manager.shutdown(); + } + } + + @Test + void testLoadInstallsRuleHoldsOnMetadataAndSelectedBundle() throws Exception + { + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + + // The metadata entry now carries the rule state. + final PartialSegmentMetadataCacheEntry metadata = weakReservedMetadata(location, SEGMENT_ID); + Assertions.assertTrue(metadata.isRuleHeld(), "rule must be applied to the metadata entry"); + Assertions.assertEquals(FINGERPRINT, metadata.getRuleFingerprint(), "fingerprint must be stored on the entry"); + // Metadata + selected bundle + its base parent are weak-reserved on the location; NOT static. + 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), "metadata entry should be weak-reserved"); + Assertions.assertTrue(location.isWeakReserved(aggId), "selected bundle should be weak-reserved"); + Assertions.assertTrue(location.isWeakReserved(baseId), "base dependency should be weak-reserved"); + Assertions.assertFalse(location.isReserved(metaId), "metadata entry should NOT be in staticCacheEntries"); + Assertions.assertFalse(location.isReserved(aggId), "selected bundle should NOT be in staticCacheEntries"); + } + + @Test + void testLoadDoesNotReserveNonSelectedBundles() throws Exception + { + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + + final PartialSegmentBundleCacheEntryIdentifier otherAggId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, OTHER_AGG_BUNDLE); + Assertions.assertFalse( + location.isWeakReserved(otherAggId), + "non-selected bundle must not be reserved by loadPartial" + ); + Assertions.assertNull( + location.getCacheEntry(otherAggId), + "non-selected bundle must not be in the cache at all after load" + ); + } + + @Test + void testLoadEagerlyDownloadsSelectedBundleContainers() throws Exception + { + // ensureBundleDownloaded should be a no-op after load — the eager-download pool task already downloaded the + // bundle. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + + final PartialSegmentMetadataCacheEntry metadata = weakReservedMetadata(location, SEGMENT_ID); + final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper(); + Assertions.assertNotNull(mapper, "metadata mount should produce a file mapper"); + final long downloadedBefore = mapper.getDownloadedBytes(); + mapper.ensureBundleDownloaded(AGG_BUNDLE); + Assertions.assertEquals( + downloadedBefore, + mapper.getDownloadedBytes(), + "ensureBundleDownloaded must be a no-op — the eager task already downloaded the bundle" + ); + } + + @Test + void testLoadWithoutVirtualStorageDoesNotInstallRuleHolds() throws Exception + { + // virtualStorage=false: partial-load rules aren't a concept here. load() takes the eager full-load path (via + // PartialLoadSpec.loadSegment → the materialized inner LocalLoadSpec), and no rule state is created. + manager = makeManager(false, false); + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + Assertions.assertNull( + manager.getRuleFingerprintForSegment(SEGMENT_ID), + "no rule state on the eager path" + ); + } + + @Test + void testLoadWithPartialsDisabledIsNoop() throws Exception + { + manager = makeManager(true, false); + final StorageLocation location = manager.getLocations().get(0); + + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + + Assertions.assertNull( + manager.getRuleFingerprintForSegment(SEGMENT_ID), + "with partial downloads disabled, loadPartial must not run" + ); + Assertions.assertFalse( + location.isWeakReserved(new SegmentCacheEntryIdentifier(SEGMENT_ID)), + "with partial downloads disabled, no reservation should be installed at load time" + ); + } + + @Test + void testSameFingerprintReIssueIsIdempotent() throws Exception + { + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + final PartialSegmentMetadataCacheEntry firstMeta = weakReservedMetadata(location, SEGMENT_ID); + + // Second load with the same fingerprint must be idempotent — same metadata entry, unchanged rule state. + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + final PartialSegmentMetadataCacheEntry secondMeta = weakReservedMetadata(location, SEGMENT_ID); + Assertions.assertSame(firstMeta, secondMeta, "same-fingerprint re-issue must not create a new metadata entry"); + Assertions.assertEquals(FINGERPRINT, secondMeta.getRuleFingerprint(), "fingerprint unchanged after re-issue"); + } + + @Test + void testRuleChangeSwapsHoldsWithoutDroppingSharedMetadata() throws Exception + { + // Same segment, different fingerprint: the rule's bundle selection changes but the underlying metadata entry is + // the same weak entry (its self-hold is never released during the swap; only bundle rule-holds diff). + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original")); + final PartialSegmentMetadataCacheEntry meta = weakReservedMetadata(location, SEGMENT_ID); + Assertions.assertEquals("v1:rule-original", meta.getRuleFingerprint()); + + manager.load(partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), "v2:rule-updated")); + + // Metadata entry is the SAME instance across the swap. + final PartialSegmentMetadataCacheEntry metaAfter = weakReservedMetadata(location, SEGMENT_ID); + Assertions.assertSame(meta, metaAfter, "rule change must reuse the underlying metadata entry"); + Assertions.assertEquals("v2:rule-updated", metaAfter.getRuleFingerprint(), "fingerprint must swap"); + // New rule's bundle is rule-held. + Assertions.assertTrue( + location.isWeakReserved(new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, OTHER_AGG_BUNDLE)), + "new rule's bundle should be reserved" + ); + } + + @Test + void testDropClearsRule() throws Exception + { + manager = makeManager(true, true); + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + Assertions.assertNotNull(manager.getRuleFingerprintForSegment(SEGMENT_ID)); + + manager.drop(partialWrapperSegment(List.of(AGG_BUNDLE))); + + Assertions.assertNull( + manager.getRuleFingerprintForSegment(SEGMENT_ID), + "drop must clear the rule state on the metadata entry" + ); + } + + @Test + void testRangeReaderNullDropsPriorRule() throws Exception + { + // Prior rule installed. Second load with a wrapper whose delegate can't produce a range reader releases the rule + // via clearRule so subsequent queries fall through to weak-full-load semantics. + manager = makeManager(true, true); + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original")); + Assertions.assertNotNull(manager.getRuleFingerprintForSegment(SEGMENT_ID)); + + manager.load(partialWrapperSegmentWithNullRangeReader(List.of(AGG_BUNDLE), "v2:rule-updated")); + + Assertions.assertNull( + manager.getRuleFingerprintForSegment(SEGMENT_ID), + "null range reader must clear the prior rule" + ); + } + + @Test + void testRangeReaderNullNoopWhenNoPriorRule() throws Exception + { + // No prior rule, no range reader → no rule installed, but also no failure. Segment simply doesn't get a rule; + // a later query would drive weak-full-load via the on-demand path. + manager = makeManager(true, true); + manager.load(partialWrapperSegmentWithNullRangeReader(List.of(AGG_BUNDLE), "v1:rule-noop")); + Assertions.assertNull(manager.getRuleFingerprintForSegment(SEGMENT_ID)); + } + + @Test + void testLoadFailureLeavesNoRuleApplied() throws Exception + { + // Wrapper referring to a non-existent projection — wrapper.getSelectedBundleNames throws before applyRule fires. + // In the v2 design there is no rollback-nuke of the partial dir on failure; the transient hold releases and the + // now-unheld weak metadata entry becomes SIEVE-eligible naturally (info-file cleanup fires via the entry's + // onUnmount hook when SIEVE eventually reclaims it). The important correctness invariant is: no stranded rule + // state, no partial rule fingerprint applied. + manager = makeManager(true, true); + final DataSegment segment = partialWrapperSegment(List.of("does_not_exist")); + + final Throwable thrown = Assertions.assertThrows( + Throwable.class, + () -> manager.load(segment) + ); + Assertions.assertTrue( + thrown.getMessage() != null && thrown.getMessage().contains("does not contain projection[does_not_exist]") + || thrown.getCause() != null + && thrown.getCause().getMessage().contains("does not contain projection[does_not_exist]"), + "unexpected throwable: " + thrown + ); + Assertions.assertNull( + manager.getRuleFingerprintForSegment(SEGMENT_ID), + "failed load must not leave a rule applied on the metadata entry" + ); + } + + @Test + void testBootstrapReinstallsRuleHoldsFromPersistedInfoFile() throws Exception + { + // Regression guard for the "brief post-bootstrap SIEVE-eligible window" concern: on historical restart, + // getCachedSegments+bootstrap must reapply the rule using the persisted DataSegment's loadSpec (the wire-form + // PartialLoadSpec is stored in the info file), so a rule-pinned segment is protected from eviction as soon as + // bootstrap completes — no waiting for the coordinator's next sync. + manager = makeManager(true, true); + final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE)); + + manager.load(segment); + Assertions.assertEquals(FINGERPRINT, manager.getRuleFingerprintForSegment(SEGMENT_ID)); + manager.shutdown(); + manager = null; + + // Fresh manager over the same cacheRoot — simulates a restart. + final SegmentLocalCacheManager restarted = makeManager(true, true); + try { + final List cached = restarted.getCachedSegments(); + Assertions.assertTrue( + cached.stream().anyMatch(s -> s.getId().equals(SEGMENT_ID)), + "restarted historical must rediscover the segment via its info file" + ); + + restarted.bootstrap(segment, SegmentLazyLoadFailCallback.NOOP); + + // Fingerprint round-trips through the info file → wrapper → applyRule. + Assertions.assertEquals( + FINGERPRINT, + restarted.getRuleFingerprintForSegment(SEGMENT_ID), + "bootstrap must reapply the persisted PartialLoadSpec wrapper's fingerprint" + ); + // Selected bundle + its base dependency are held after bootstrap restore (they were on disk and got + // restored by the metadata mount's PartialSegmentCacheBootstrap.restoreBundlesFromDisk, which register with + // the metadata; applyRule picks up their rule-holds from the linkedBundles state). + final StorageLocation loc = restarted.getLocations().get(0); + Assertions.assertTrue( + loc.isWeakReserved(new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE)), + "selected bundle must be reserved after bootstrap reapply" + ); + Assertions.assertTrue( + loc.isWeakReserved( + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME) + ), + "base dependency must be reserved after bootstrap reapply" + ); + } + finally { + restarted.shutdown(); + } + } + + @Test + void testReserveMkdirpFailureFallsThroughToNextLocation() throws Exception + { + // Setup: writable first (so info_dir defaults there), read-only second, dummy reservation on writable so + // LeastBytesUsed picks the read-only location first. mkdirp on read-only fails → release + continue to writable. + 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 usage so LeastBytesUsed picks readOnly first. + final SegmentId dummy = SegmentId.of("dummy", Intervals.of("2020/2021"), "v", 0); + Assertions.assertTrue( + manager.getLocations().get(0).reserveWeak(stubCacheEntry(new SegmentCacheEntryIdentifier(dummy), 4096L)) + ); + Assertions.assertTrue(readOnly.setReadOnly(), "test setup must be able to make readOnly location read-only"); + try { + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + + Assertions.assertFalse( + new File(readOnly, SEGMENT_ID.toString()).exists(), + "read-only location must NOT receive the segment" + ); + Assertions.assertTrue( + new File(writable, SEGMENT_ID.toString()).isDirectory(), + "load must fall through to the writable location" + ); + Assertions.assertNotNull(manager.getRuleFingerprintForSegment(SEGMENT_ID)); + } + finally { + Assertions.assertTrue(readOnly.setWritable(true), "test teardown must restore write permission"); + } + } + + @Test + void testLoadReusesExistingWeakEntryHeldByConcurrentQuery() throws Exception + { + // A concurrent query is holding a weak metadata entry from an on-demand acquire when the coordinator issues a + // partial-load rule for the same segment. loadPartial must reuse that entry — addWeakReservationHoldIfExists + // adds another party to the same phaser rather than installing a fresh entry alongside. + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(SEGMENT_ID); + + final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE)); + final File partialDir = new File(cacheRoot, SEGMENT_ID.toString()); + FileUtils.mkdirp(partialDir); + final PartialSegmentMetadataCacheEntry preExisting = new PartialSegmentMetadataCacheEntry( + SEGMENT_ID, + partialDir, + IndexIO.V10_FILE_NAME, + List.of(), + new DirectoryBackedRangeReader(DEEP_STORAGE_DIR), + jsonMapper, + null, + 16L * 1024L * 1024L + ); + final StorageLocation.ReservationHold queryHold = + location.addWeakReservationHold(id, () -> preExisting); + Assertions.assertNotNull(queryHold, "precondition: on-demand-style weak reserve must succeed"); + try { + Assertions.assertTrue(location.isWeakReserved(id)); + + manager.load(segment); + + // The rule hold reuses the pre-existing entry — same instance across the load. + final PartialSegmentMetadataCacheEntry ruleHeld = weakReservedMetadata(location, SEGMENT_ID); + Assertions.assertSame( + preExisting, + ruleHeld, + "loadPartial must add its hold to the pre-existing weak entry, not install a fresh one" + ); + Assertions.assertNotNull(manager.getRuleFingerprintForSegment(SEGMENT_ID)); + } + finally { + queryHold.close(); + } + } + + 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. + */ + 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, verifying it's in weak — not static. loadPartial's + * install is synchronous by the time load() returns; the entry is either fully mounted or the call threw. + */ + private static PartialSegmentMetadataCacheEntry weakReservedMetadata(StorageLocation location, SegmentId segmentId) + { + final CacheEntry entry = location.getCacheEntry(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 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() + { + } + }; + } +} From b65f9458a15a0f98140a26b92d353bbd66fee625 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Thu, 9 Jul 2026 15:01:48 -0700 Subject: [PATCH 2/6] add embedded test --- .../PartialProjectionLoadRuleQueryTest.java | 372 ++++++++++++++++++ 1 file changed, 372 insertions(+) create mode 100644 embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java new file mode 100644 index 000000000000..a3efef184dcc --- /dev/null +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java @@ -0,0 +1,372 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.testing.embedded.query; + +import org.apache.druid.catalog.guice.CatalogClientModule; +import org.apache.druid.catalog.guice.CatalogCoordinatorModule; +import org.apache.druid.catalog.model.Columns; +import org.apache.druid.catalog.model.TableMetadata; +import org.apache.druid.catalog.model.table.TableBuilder; +import org.apache.druid.common.utils.IdUtils; +import org.apache.druid.data.input.impl.AggregateProjectionSpec; +import org.apache.druid.data.input.impl.ClusteredValueGroupsBaseTableProjectionSpec; +import org.apache.druid.data.input.impl.LongDimensionSchema; +import org.apache.druid.data.input.impl.StringDimensionSchema; +import org.apache.druid.data.input.impl.TimestampSpec; +import org.apache.druid.indexer.granularity.SegmentGranularitySpec; +import org.apache.druid.indexing.common.task.TaskBuilder; +import org.apache.druid.indexing.common.task.batch.parallel.ParallelIndexSupervisorTask; +import org.apache.druid.java.util.common.HumanReadableBytes; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.granularity.Granularities; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.query.aggregation.LongSumAggregatorFactory; +import org.apache.druid.server.coordinator.rules.CannotMatchBehavior; +import org.apache.druid.server.coordinator.rules.ForeverPartialLoadRule; +import org.apache.druid.server.coordinator.rules.WildcardProjectionPartialLoadMatcher; +import org.apache.druid.server.metrics.LatchableEmitter; +import org.apache.druid.server.metrics.StorageMonitor; +import org.apache.druid.testing.embedded.EmbeddedBroker; +import org.apache.druid.testing.embedded.EmbeddedCoordinator; +import org.apache.druid.testing.embedded.EmbeddedDruidCluster; +import org.apache.druid.testing.embedded.EmbeddedHistorical; +import org.apache.druid.testing.embedded.EmbeddedIndexer; +import org.apache.druid.testing.embedded.EmbeddedOverlord; +import org.apache.druid.testing.embedded.EmbeddedRouter; +import org.apache.druid.testing.embedded.catalog.TestCatalogClient; +import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase; +import org.apache.druid.testing.embedded.minio.MinIOStorageResource; +import org.apache.druid.testing.embedded.msq.EmbeddedMSQApis; +import org.hamcrest.MatcherAssert; +import org.hamcrest.Matchers; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.io.File; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +/** + * End-to-end coverage for the partial-load-rule wiring: a segment with a clustered base + aggregate projection + * loaded under a {@link ForeverPartialLoadRule} whose matcher selects only the projection. The coordinator emits + * a {@code PartialProjectionLoadSpec} wrapper on the load request; the historical resolves the wrapper into the + * projection's bundle name, eagerly downloads that one bundle, and pins it via a rule-hold. Queries that the + * projection can serve are proven to run entirely off the pre-loaded bundle (zero on-demand loads); queries that + * miss the projection are proven to trigger on-demand loads of the base bundles. + *

+ * Exercises the full happy-path flow end-to-end: coordinator matcher → wire form → historical + * {@code loadPartial} → {@code applyRule} + eager download → announced fingerprint → query with + * projection hit vs miss → on-demand metric emission. + */ +class PartialProjectionLoadRuleQueryTest extends EmbeddedClusterTestBase +{ + private static final String PROJECTION_NAME = "country_delta"; + private static final long CACHE_SIZE = HumanReadableBytes.parse("1MiB"); + private static final long MAX_SIZE = HumanReadableBytes.parse("100MiB"); + private static final long ESTIMATE_SIZE = HumanReadableBytes.parse("2KiB"); + // Quiescence wait for the storage monitor: a few times its PT1s emission period so a single missed tick + // doesn't falsely read as "idle" while still bounding how long we wait once activity has actually stopped. + private static final long MONITOR_QUIESCE_TIMEOUT_MILLIS = 3_000L; + + private final EmbeddedBroker broker = new EmbeddedBroker(); + private final EmbeddedIndexer indexer = new EmbeddedIndexer(); + private final EmbeddedOverlord overlord = new EmbeddedOverlord(); + private final EmbeddedHistorical historical = new EmbeddedHistorical(); + private final EmbeddedCoordinator coordinator = new EmbeddedCoordinator(); + private final EmbeddedRouter router = new EmbeddedRouter(); + private final MinIOStorageResource storageResource = new MinIOStorageResource(); + + private EmbeddedMSQApis msqApis; + + @Override + public EmbeddedDruidCluster createCluster() + { + historical.setServerMemory(500_000_000) + .addProperty("druid.segmentCache.virtualStorage", "true") + .addProperty("druid.segmentCache.virtualStoragePartialDownloadsEnabled", "true") + .addProperty( + "druid.segmentCache.virtualStorageMetadataReservationEstimate", + String.valueOf(ESTIMATE_SIZE) + ) + .addProperty( + "druid.segmentCache.virtualStorageLoadThreads", + String.valueOf(Runtime.getRuntime().availableProcessors()) + ) + .addBeforeStartHook( + (cluster, self) -> self.addProperty( + "druid.segmentCache.locations", + StringUtils.format( + "[{\"path\":\"%s\",\"maxSize\":\"%s\"}]", + cluster.getTestFolder().newFolder().getAbsolutePath(), + CACHE_SIZE + ) + ) + ) + .addProperty("druid.server.maxSize", String.valueOf(MAX_SIZE)); + + broker.setServerMemory(200_000_000) + .addProperty("druid.sql.planner.enableSysQueriesTable", "true") + // Poll the catalog often so the broker's DruidSchema picks up the table definition quickly. + // Under virtualStorage the historical exposes only __time via VirtualPlaceholderSegment, so + // BrokerSegmentMetadataCache alone can't infer the real column schema — the catalog fills that gap. + .addProperty("druid.catalog.client.pollingPeriod", "100") + // Retain Dart reports long enough to fetch them after the query completes. + .addProperty("druid.msq.dart.controller.maxRetainedReportCount", "10"); + + coordinator.addProperty("druid.manager.segments.useIncrementalCache", "always"); + + overlord.addProperty("druid.manager.segments.useIncrementalCache", "always") + .addProperty("druid.manager.segments.pollDuration", "PT0.1s") + .addProperty("druid.catalog.client.maxSyncRetries", "0"); + + indexer.setServerMemory(300_000_000) + .addProperty("druid.worker.capacity", "2") + .addProperty("druid.processing.numThreads", "2") + .addProperty("druid.segment.handoff.pollDuration", "PT0.1s"); + + return EmbeddedDruidCluster + .withEmbeddedDerbyAndZookeeper() + .useLatchableEmitter() + .useDefaultTimeoutForLatchableEmitter(60) + .addResource(storageResource) + // Catalog: table schema must be published explicitly since historical exposes only __time under + // virtualStorage=true (see broker.druid.catalog.client.pollingPeriod above). + .addExtensions(CatalogClientModule.class, CatalogCoordinatorModule.class) + // Run queries via Dart so MSQ leaf processors take the PARTIAL acquire path (async cursor factory); + // native SQL takes the FULL acquire path which pins every bundle for the reference lifetime, defeating + // the memory savings of partial loads. + .addCommonProperty("druid.msq.dart.enabled", "true") + // Clustered base-table + projection segments require the V10 segment format. + .addCommonProperty("druid.indexer.task.buildV10", "true") + // Range reads over MinIO — required for partial loading to be functional. + .addCommonProperty("druid.storage.zip", "false") + .addCommonProperty("druid.monitoring.emissionPeriod", "PT1s") + .addServer(coordinator) + .addServer(overlord) + .addServer(indexer) + .addServer(historical) + .addServer(broker) + .addServer(router); + } + + @BeforeAll + void loadDataAndConfigureRule() throws IOException + { + msqApis = new EmbeddedMSQApis(cluster, overlord); + dataSource = "partial-projection-" + IdUtils.getRandomId(); + + // Under virtualStorage=true the historical announces segments via VirtualPlaceholderSegment, which only exposes + // __time. That means BrokerSegmentMetadataCache alone can't populate a full DruidSchema for SQL validation, so + // the queries below would fail at plan-time with "column not found". Publishing an explicit catalog table + // definition gives DruidSchema the real column list. Future work: derive a RowSignature from the partial-load + // header so this catalog registration becomes unnecessary. + final TestCatalogClient catalog = new TestCatalogClient(cluster); + final TableMetadata table = TableBuilder.datasource(dataSource, "HOUR") + .column(Columns.TIME_COLUMN, Columns.LONG) + .column("channel", Columns.SQL_VARCHAR) + .column("countryName", Columns.SQL_VARCHAR) + .column("delta", Columns.SQL_BIGINT) + .build(); + catalog.createTable(table, true); + + // Configure the partial-load rule BEFORE ingestion so the initial load applies the rule directly, rather than + // needing a rule change afterward. The matcher selects only the country_delta projection; base-table bundles + // (cluster groups) are NOT rule-loaded. + cluster.callApi().onLeaderCoordinator( + c -> c.updateRulesForDatasource( + dataSource, + List.of( + new ForeverPartialLoadRule( + Map.of("_default_tier", 1), + null, + new WildcardProjectionPartialLoadMatcher(List.of(PROJECTION_NAME), null), + CannotMatchBehavior.FALL_THROUGH + ) + ) + ) + ); + + ingestClusteredSegmentWithProjection(); + } + + @Override + protected void refreshDatasourceName() + { + // Fixed datasource across tests — rule + ingest are one-time setup. + } + + @Test + void testProjectionHitLoadsNothingOnDemand() + { + // The country_delta projection can serve group-by-countryName / sum-delta. The projection bundle was + // rule-loaded on ingest, so an MSQ leaf PARTIAL acquire should read it without triggering any downloads. + // MSQ input channel counters don't (yet) observe partial-load on-demand downloads (those happen inside the + // async cursor factory, downstream of the input channel), so we observe the historical's StorageMonitor + // directly — VSF_LOAD_* are incremented as each file is actually pulled from deep storage. + final LatchableEmitter emitter = quiesceAndFlushStorageMonitor(); + + runDartQuery( + "SELECT \"countryName\", SUM(\"delta\") FROM \"" + dataSource + "\" GROUP BY 1 ORDER BY 1", + "CA,3\nFR,7\nUS,17" + ); + + // Wait for post-query storage activity (if any) to settle before we sum the counters. + emitter.awaitMetricQuiescent(StorageMonitor.VSF_LOAD_BEGIN_COUNT, MONITOR_QUIESCE_TIMEOUT_MILLIS); + + Assertions.assertEquals( + 0L, + emitter.getMetricEventLongSum(StorageMonitor.VSF_LOAD_COUNT), + "expected no on-demand file loads for a projection-served Dart query, but the projection bundle was fetched at query time" + ); + Assertions.assertEquals( + 0L, + emitter.getMetricEventLongSum(StorageMonitor.VSF_LOAD_BYTES), + "expected zero on-demand load bytes for a projection-served Dart query" + ); + } + + @Test + void testProjectionMissLoadsBaseBundleOnDemand() + { + // Group by the clustering column — the projection can't serve this, so MSQ must fall back to the clustered + // base table, whose cluster-group bundles were NOT rule-loaded. Under PARTIAL acquire the async cursor + // factory triggers per-file downloads on demand, which show up on StorageMonitor's VSF_LOAD_* counters. + final LatchableEmitter emitter = quiesceAndFlushStorageMonitor(); + + runDartQuery( + "SELECT \"channel\", SUM(\"delta\") FROM \"" + dataSource + "\" GROUP BY 1 ORDER BY 1", + "#en,18\n#fr,9" + ); + + emitter.awaitMetricQuiescent(StorageMonitor.VSF_LOAD_BEGIN_COUNT, MONITOR_QUIESCE_TIMEOUT_MILLIS); + + MatcherAssert.assertThat( + "expected on-demand file loads for a projection-miss Dart query that needs the base-table bundles", + emitter.getMetricEventLongSum(StorageMonitor.VSF_LOAD_COUNT), + Matchers.greaterThan(0L) + ); + MatcherAssert.assertThat( + "expected non-zero on-demand load bytes for a projection-miss Dart query", + emitter.getMetricEventLongSum(StorageMonitor.VSF_LOAD_BYTES), + Matchers.greaterThan(0L) + ); + } + + /** + * Waits for any lingering storage-monitor activity from ingest / rule-application to settle, then flushes the + * emitter's event queue so subsequent {@code getMetricEventLongSum(...)} calls only reflect activity that + * happens after this point. + */ + private LatchableEmitter quiesceAndFlushStorageMonitor() + { + final LatchableEmitter emitter = historical.latchableEmitter(); + emitter.awaitMetricQuiescent(StorageMonitor.VSF_LOAD_BEGIN_COUNT, MONITOR_QUIESCE_TIMEOUT_MILLIS); + emitter.flush(); + return emitter; + } + + /** + * Submits {@code sql} via the Dart engine (MSQ leaves take the PARTIAL acquire path, so partial-load bundles + * only download the files the query actually reads) and asserts the CSV result matches {@code expectedCsv}. + */ + private void runDartQuery(String sql, String expectedCsv) + { + final String sqlQueryId = UUID.randomUUID().toString(); + final String result; + try { + result = msqApis.submitDartSqlAsync( + sql, + Map.of(QueryContexts.CTX_SQL_QUERY_ID, sqlQueryId), + broker + ).get(); + } + catch (Exception e) { + throw new RuntimeException("failed to run Dart query [" + sqlQueryId + "]: " + sql, e); + } + Assertions.assertEquals(expectedCsv, result.trim(), "unexpected result for Dart query [" + sqlQueryId + "]"); + } + + /** + * Ingests a single clustered base-table segment (clustered by {@code channel}) with a {@code country_delta} + * aggregate projection (group by {@code countryName}, sum {@code delta}). Copied from + * {@link ClusteredSegmentProjectionQueryTest#ingestClusteredSegmentWithProjection}; kept local so this test + * doesn't cross-depend on that class. + */ + private void ingestClusteredSegmentWithProjection() throws IOException + { + final File tmpDir = cluster.getTestFolder().newFolder(); + final File inputFile = new File(tmpDir, "clustered-input.json"); + final String inputData = + "{\"time\":\"2024-01-01T00:10:00Z\",\"channel\":\"#en\",\"countryName\":\"US\",\"delta\":10}\n" + + "{\"time\":\"2024-01-01T00:20:00Z\",\"channel\":\"#en\",\"countryName\":\"US\",\"delta\":5}\n" + + "{\"time\":\"2024-01-01T00:30:00Z\",\"channel\":\"#en\",\"countryName\":\"CA\",\"delta\":3}\n" + + "{\"time\":\"2024-01-01T00:40:00Z\",\"channel\":\"#fr\",\"countryName\":\"FR\",\"delta\":7}\n" + + "{\"time\":\"2024-01-01T00:50:00Z\",\"channel\":\"#fr\",\"countryName\":\"US\",\"delta\":2}\n"; + Files.write(inputFile.toPath(), inputData.getBytes(StandardCharsets.UTF_8)); + + final ClusteredValueGroupsBaseTableProjectionSpec clusterSpec = + ClusteredValueGroupsBaseTableProjectionSpec.builder() + .columns( + new StringDimensionSchema("channel"), + new StringDimensionSchema("countryName"), + new LongDimensionSchema("delta"), + new LongDimensionSchema("__time") + ) + .clusteringColumns("channel") + .build(); + + final AggregateProjectionSpec projection = + AggregateProjectionSpec.builder(PROJECTION_NAME) + .groupingColumns(new StringDimensionSchema("countryName")) + .aggregators(new LongSumAggregatorFactory("sumDelta", "delta")) + .build(); + + final SegmentGranularitySpec segmentGranularitySpec = new SegmentGranularitySpec( + Granularities.HOUR, + List.of(Intervals.of("2024-01-01/2024-01-02")) + ); + + final String taskId = IdUtils.getRandomId(); + final ParallelIndexSupervisorTask task = TaskBuilder + .ofTypeIndexParallel() + .jsonInputFormat() + .localInputSourceWithFiles(inputFile) + .dataSchema( + builder -> builder + .withDataSource(dataSource) + .withTimestamp(new TimestampSpec("time", "iso", null)) + .withSegmentGranularity(segmentGranularitySpec) + .withBaseTable(clusterSpec) + .withProjections(List.of(projection)) + ) + .tuningConfig(t -> t.withMaxNumConcurrentSubTasks(1)) + .withId(taskId); + + cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task)); + cluster.callApi().waitForTaskToSucceed(taskId, overlord); + cluster.callApi().waitForAllSegmentsToBeAvailable(dataSource, coordinator, broker); + } +} From eced2d66c8c80ecad5190eccb42e8957a7d22d4c Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Thu, 9 Jul 2026 17:32:51 -0700 Subject: [PATCH 3/6] clear up comments, fix to only eagerly evict unheld when was rule held --- .../PartialSegmentMetadataCacheEntry.java | 4 ---- .../loading/SegmentLocalCacheManager.java | 19 ++++++++----------- .../segment/loading/StorageLocation.java | 3 +-- ...tLocalCacheManagerPartialRuleLoadTest.java | 8 ++++---- 4 files changed, 13 insertions(+), 21 deletions(-) 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 9c6856fd8f10..2ee28fbe5d18 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 @@ -391,8 +391,6 @@ public void applyRule(String fingerprint, Set selectedBundleNames) metadataSelfHold = newSelfHold; uncommittedHolds.remove(newSelfHold); } - // else: a concurrent applyRule installed a self-hold (shouldn't happen under segmentLock, but defensive). - // newSelfHold stays in uncommittedHolds, gets released in Phase 4. } for (String name : namesToRelease) { final StorageLocation.ReservationHold h = ruleBundleHolds.remove(name); @@ -405,8 +403,6 @@ public void applyRule(String fingerprint, Set selectedBundleNames) ruleBundleHolds.put(e.getKey(), e.getValue()); uncommittedHolds.remove(e.getValue()); } - // else: a concurrent registerBundle raced ahead and installed a hold for this name; ours stays in - // uncommittedHolds, gets released in Phase 4. } ruleFingerprint = fingerprint; ruleSelectedBundleNames = newSelection; 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 816a65b1b4fe..193a94d6e9a4 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 @@ -1444,18 +1444,15 @@ public void drop(final DataSegment segment) for (StorageLocation location : locations) { final CacheEntry entry = location.getCacheEntry(id); if (entry instanceof PartialSegmentMetadataCacheEntry partial) { + final boolean partialLoaded = partial.isRuleHeld(); partial.clearRule(); - // Force synchronous cleanup: clearRule releases the self-hold, but the info-file-deletion hook only - // fires on doActualUnmount which is triggered by the phaser hitting zero. removeUnheldWeakEntry - // unlinks the weak entry from cache and terminates the phaser now, firing the hook (and the mapper - // teardown) on the caller's thread. No-op if a concurrent query still holds the entry. - location.removeUnheldWeakEntry(id); - } - } - // full-segment drop path: release any {@link CompleteSegmentCacheEntry} still reserved. - for (StorageLocation location : locations) { - final CacheEntry entry = location.getCacheEntry(id); - if (entry != null) { + if (partialLoaded) { + // if partially loaded, try to trigger an early eviction for the entry on the assumption it is unlikely + // to be used again. No-op if a concurrent query still holds the entry, leaving the entry to be lazily + // evicted later + location.removeUnheldWeakEntry(id); + } + } else if (entry != null) { location.release(entry); } } diff --git a/server/src/main/java/org/apache/druid/segment/loading/StorageLocation.java b/server/src/main/java/org/apache/druid/segment/loading/StorageLocation.java index 460a2fa9d2c6..49fc135a9a8e 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/StorageLocation.java +++ b/server/src/main/java/org/apache/druid/segment/loading/StorageLocation.java @@ -540,12 +540,11 @@ public void removeUnheldWeakEntry(CacheEntryIdentifier id) id, (cacheEntryIdentifier, weakCacheEntry) -> { if (weakCacheEntry.isHeld()) { - // a holder is responsible for cleanup when it releases; leave the entry in place return weakCacheEntry; } final boolean isMounted = weakCacheEntry.cacheEntry.isMounted(); unlinkWeakEntry(weakCacheEntry); - weakCacheEntry.unmount(); // terminate the phaser; fires cacheEntry.unmount() (idempotent) + weakCacheEntry.unmount(); if (isMounted) { weakStats.getAndUpdate(s -> s.evict(weakCacheEntry.cacheEntry.getSize())); } diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java index ea833e59270c..a8b1a3c9f927 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java @@ -340,6 +340,7 @@ void testRuleChangeSwapsHoldsWithoutDroppingSharedMetadata() throws Exception location.isWeakReserved(new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, OTHER_AGG_BUNDLE)), "new rule's bundle should be reserved" ); + Assertions.assertFalse(metaAfter.isBundleRuleHeld(AGG_BUNDLE)); } @Test @@ -414,10 +415,9 @@ void testLoadFailureLeavesNoRuleApplied() throws Exception @Test void testBootstrapReinstallsRuleHoldsFromPersistedInfoFile() throws Exception { - // Regression guard for the "brief post-bootstrap SIEVE-eligible window" concern: on historical restart, - // getCachedSegments+bootstrap must reapply the rule using the persisted DataSegment's loadSpec (the wire-form - // PartialLoadSpec is stored in the info file), so a rule-pinned segment is protected from eviction as soon as - // bootstrap completes — no waiting for the coordinator's next sync. + // on historical restart, getCachedSegments+bootstrap must reapply the rule using the persisted DataSegment's + // loadSpec (the wire-form PartialLoadSpec is stored in the info file), so a rule-pinned segment is protected from + // eviction as soon as bootstrap completes manager = makeManager(true, true); final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE)); From ce8e6a321ca2bf30ac7b16371f87bc34e566d422 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Thu, 9 Jul 2026 17:49:53 -0700 Subject: [PATCH 4/6] fix --- .../loading/SegmentLocalCacheManager.java | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java b/server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java index 193a94d6e9a4..508ee8c0ba37 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 @@ -1438,21 +1438,17 @@ public void drop(final DataSegment segment) final ReferenceCountingLock lock = lock(segment); synchronized (lock) { try { - // partial-load-rule cleanup: if the segment's cache entry is a partial metadata entry with an applied rule, - // clear it + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segment.getId()); for (StorageLocation location : locations) { final CacheEntry entry = location.getCacheEntry(id); if (entry instanceof PartialSegmentMetadataCacheEntry partial) { - final boolean partialLoaded = partial.isRuleHeld(); + // if the segment's cache entry is a partial metadata entry with an applied rule, clear it. Reclaim of + // on-disk state is left to eviction. partial.clearRule(); - if (partialLoaded) { - // if partially loaded, try to trigger an early eviction for the entry on the assumption it is unlikely - // to be used again. No-op if a concurrent query still holds the entry, leaving the entry to be lazily - // evicted later - location.removeUnheldWeakEntry(id); - } - } else if (entry != null) { + } + if (entry != null) { + // static entries are removed from both cache and disk immediately (no-op for weak held entries) location.release(entry); } } From d59bf60bae75a41cf242bb43e18ec77a297d0a44 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Fri, 10 Jul 2026 11:59:54 -0700 Subject: [PATCH 5/6] fixes --- .../PartialSegmentMetadataCacheEntry.java | 50 +++++++++++++++++++ .../loading/SegmentLocalCacheManager.java | 44 +++++++++++----- ...tLocalCacheManagerPartialRuleLoadTest.java | 45 +++++++++++++++++ 3 files changed, 127 insertions(+), 12 deletions(-) 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 2ee28fbe5d18..aa4c99b9eb63 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 @@ -410,6 +410,56 @@ public void applyRule(String fingerprint, Set selectedBundleNames) finally { entryLock.unlock(); } + + // Phase 3b: reconcile linkedBundles against the freshly-committed selection. A concurrent registerBundle + // whose linkedBundles.put landed AFTER our Phase 1 snapshot but whose Phase 1 selection-check ran BEFORE + // our Phase 3 commit would have observed the OLD selection and bailed without installing a hold, while + // Phase 1's namesToAcquire also missed it. Rescan + install to close that window. Loops in case a bundle + // arrives during reconciliation; termination is bounded by |newSelection| because each iteration either installs + // at least one hold or exits (a bundle that arrives after this loop terminates will see the committed + // newSelection in its own Phase 1 and install its hold via registerBundle Phase 3). + while (true) { + final List reconcileNames; + entryLock.lock(); + try { + reconcileNames = new ArrayList<>(); + for (String name : ruleSelectedBundleNames) { + if (!ruleBundleHolds.containsKey(name) && findLinkedBundleByName(name) != null) { + reconcileNames.add(name); + } + } + } + finally { + entryLock.unlock(); + } + if (reconcileNames.isEmpty()) { + break; + } + + final Map> reconcileHolds = + new HashMap<>(); + for (String name : reconcileNames) { + final StorageLocation.ReservationHold h = + loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, name)); + if (h != null) { + reconcileHolds.put(name, h); + uncommittedHolds.add(h); + } + } + + entryLock.lock(); + try { + for (var e : reconcileHolds.entrySet()) { + if (ruleSelectedBundleNames.contains(e.getKey()) && !ruleBundleHolds.containsKey(e.getKey())) { + ruleBundleHolds.put(e.getKey(), e.getValue()); + uncommittedHolds.remove(e.getValue()); + } + } + } + finally { + entryLock.unlock(); + } + } } finally { // Phase 4: release outside entryLock. `uncommittedHolds` contains every hold this call acquired but did not 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 508ee8c0ba37..3518cd7e7414 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 @@ -460,6 +460,23 @@ private void deleteSegmentInfoFile(DataSegment segment) } } + /** + * Write the info file for a partial-load segment, overwriting any existing content atomically. Distinct from + * {@link #storeInfoFile} which skips the write when the file already exists, for partial segments we must + * unconditionally rewrite so an incoming rule swap (new {@code fingerprint}/{@code delegate} inside the + * wrapped load spec) reaches disk. Otherwise bootstrap after a restart would restore the segment using the + * prior wrapper and re-announce the old rule until the coordinator resyncs. + */ + private void writePartialInfoFile(DataSegment segment) throws IOException + { + final File segmentInfoCacheFile = new File(getEffectiveInfoDir(), segment.getId().toString()); + FileUtils.mkdirp(getEffectiveInfoDir()); + FileUtils.writeAtomically(segmentInfoCacheFile, out -> { + jsonMapper.writeValue(out, segment); + return null; + }); + } + @Override public Optional acquireCachedSegment(final SegmentId segmentId, final AcquireMode acquireMode) { @@ -873,17 +890,7 @@ private ReservedPartial reservePartial(DataSegment dataSegment, SegmentRangeRead location.getPath() ); } - // Unconditionally write the info file with the incoming DataSegment. An existing info file at this path - // may carry a STALE loadSpec. Skipping the write when the file exists would preserve that stale wrapper on - // disk, and a subsequent bootstrap would restore the segment via the wrong path (non-partial) or drift out of - // sync with the coordinator's applied rule. writeAtomically does the file-rename dance so a concurrent read - // sees either the old or new complete file, never a partial write. - final File segmentInfoCacheFile = new File(getEffectiveInfoDir(), dataSegment.getId().toString()); - FileUtils.mkdirp(getEffectiveInfoDir()); - FileUtils.writeAtomically(segmentInfoCacheFile, out -> { - jsonMapper.writeValue(out, dataSegment); - return null; - }); + writePartialInfoFile(dataSegment); partial.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); return new ReservedPartial(partial, location, hold); } @@ -1009,6 +1016,19 @@ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException final ReservedPartial reserved = findOrReservePartial(dataSegment, rangeReader); try { final PartialSegmentMetadataCacheEntry metadata = reserved.metadata(); + // findOrReservePartial only invokes reservePartial (which writes the info file) on the fresh-reserve + // branch. On the find-existing branch the info file on disk still carries the PRIOR rule's wrapped + // load spec, so a rule swap here would apply in memory only. Rewrite unconditionally before mount. + try { + writePartialInfoFile(dataSegment); + } + catch (IOException e) { + throw new SegmentLoadingException( + e, + "Failed to write partial info file for segment[%s]", + dataSegment.getId() + ); + } try { metadata.mount(reserved.location()); } @@ -1077,7 +1097,7 @@ private void awaitEagerDownloadsOrClearRule( final List> pending = new ArrayList<>(); Throwable firstFailure = null; try { - for (String bundleName : selected) { + for (String bundleName : metadata.bundlesInMountOrder(selected)) { // Skip only when the bundle is BOTH rule-held AND fully downloaded (every container's files present on disk). if (metadata.isBundleRuleHeld(bundleName) && metadata.isBundleFullyDownloaded(bundleName)) { continue; diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java index a8b1a3c9f927..360fdd4cde9d 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java @@ -270,6 +270,25 @@ void testLoadEagerlyDownloadsSelectedBundleContainers() throws Exception ); } + @Test + void testLoadEagerlyDownloadsBaseDependencyOfSelectedProjection() throws Exception + { + // when the rule selects a projection bundle, __base is an implicit mount-time dependency (parent hold + sparse + // allocation) but its container bytes are NOT downloaded by ensureBundleDownloaded(projection). + manager = makeManager(true, true); + final StorageLocation location = manager.getLocations().get(0); + + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE))); + + final PartialSegmentMetadataCacheEntry metadata = weakReservedMetadata(location, SEGMENT_ID); + final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper(); + Assertions.assertNotNull(mapper, "metadata mount should produce a file mapper"); + Assertions.assertTrue( + mapper.isBundleFullyDownloaded(Projections.BASE_TABLE_PROJECTION_NAME), + "base dependency of the selected projection must be fully downloaded eagerly, not just sparse-mounted" + ); + } + @Test void testLoadWithoutVirtualStorageDoesNotInstallRuleHolds() throws Exception { @@ -343,6 +362,32 @@ void testRuleChangeSwapsHoldsWithoutDroppingSharedMetadata() throws Exception Assertions.assertFalse(metaAfter.isBundleRuleHeld(AGG_BUNDLE)); } + @Test + void testRuleSwapRewritesInfoFileOnDisk() throws Exception + { + // findOrReservePartial returns the existing entry on a rule swap, make sure we rewrite the + // info file. + manager = makeManager(true, true); + + manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original")); + manager.load(partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), "v2:rule-updated")); + + final File infoDir = new File(cacheRoot, "info_dir"); + final File infoFile = new File(infoDir, SEGMENT_ID.toString()); + Assertions.assertTrue(infoFile.exists(), "info file must exist after loadPartial"); + final DataSegment onDisk = jsonMapper.readValue(infoFile, DataSegment.class); + Assertions.assertEquals( + "v2:rule-updated", + onDisk.getLoadSpec().get("fingerprint"), + "info file on disk must carry the latest rule's fingerprint after a swap, not the prior rule's" + ); + Assertions.assertEquals( + List.of(OTHER_AGG_BUNDLE), + onDisk.getLoadSpec().get("projections"), + "info file on disk must carry the latest rule's projection selection" + ); + } + @Test void testDropClearsRule() throws Exception { From 5826302096843c53b1755c2bb99efcd5f9b9a484 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Fri, 10 Jul 2026 14:25:04 -0700 Subject: [PATCH 6/6] simplify a bit --- .../PartialSegmentMetadataCacheEntry.java | 110 ++++++++++-------- 1 file changed, 61 insertions(+), 49 deletions(-) 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 aa4c99b9eb63..94927bfb9efe 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 @@ -411,55 +411,8 @@ public void applyRule(String fingerprint, Set selectedBundleNames) entryLock.unlock(); } - // Phase 3b: reconcile linkedBundles against the freshly-committed selection. A concurrent registerBundle - // whose linkedBundles.put landed AFTER our Phase 1 snapshot but whose Phase 1 selection-check ran BEFORE - // our Phase 3 commit would have observed the OLD selection and bailed without installing a hold, while - // Phase 1's namesToAcquire also missed it. Rescan + install to close that window. Loops in case a bundle - // arrives during reconciliation; termination is bounded by |newSelection| because each iteration either installs - // at least one hold or exits (a bundle that arrives after this loop terminates will see the committed - // newSelection in its own Phase 1 and install its hold via registerBundle Phase 3). - while (true) { - final List reconcileNames; - entryLock.lock(); - try { - reconcileNames = new ArrayList<>(); - for (String name : ruleSelectedBundleNames) { - if (!ruleBundleHolds.containsKey(name) && findLinkedBundleByName(name) != null) { - reconcileNames.add(name); - } - } - } - finally { - entryLock.unlock(); - } - if (reconcileNames.isEmpty()) { - break; - } - - final Map> reconcileHolds = - new HashMap<>(); - for (String name : reconcileNames) { - final StorageLocation.ReservationHold h = - loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, name)); - if (h != null) { - reconcileHolds.put(name, h); - uncommittedHolds.add(h); - } - } - - entryLock.lock(); - try { - for (var e : reconcileHolds.entrySet()) { - if (ruleSelectedBundleNames.contains(e.getKey()) && !ruleBundleHolds.containsKey(e.getKey())) { - ruleBundleHolds.put(e.getKey(), e.getValue()); - uncommittedHolds.remove(e.getValue()); - } - } - } - finally { - entryLock.unlock(); - } - } + // Phase 3b: close the narrow race with concurrent registerBundle + reconcileLinkedBundlesWithSelection(loc, uncommittedHolds); } finally { // Phase 4: release outside entryLock. `uncommittedHolds` contains every hold this call acquired but did not @@ -470,6 +423,65 @@ public void applyRule(String fingerprint, Set selectedBundleNames) } } + /** + * Post-{@link #applyRule} reconciliation: scan {@link #linkedBundles} against the freshly-committed + * {@link #ruleSelectedBundleNames} and install a rule-hold for any selected bundle that's present in + * {@code linkedBundles} but missing from {@link #ruleBundleHolds}. + *

+ * Closes a race that can occur when bundle mounts triggered by concurrent queries run on the storage-loading pool + * without the per-segment lock, and can call {@link #registerBundle} concurrently with + * {@link #applyRule}. If the bundle lands in {@link #linkedBundles} after {@link #applyRule}'s Phase 1 + * snapshot but its {@code registerBundle} Phase 1 runs before Phase 3 has committed {@code newSelection}, + * {@code registerBundle} observes the OLD selection and bails without installing a hold, and Phase 1's + * {@code namesToAcquire} also missed it (it wasn't in {@code linkedBundles} at snapshot time). Without this + * reconciliation, the entry would incorrectly be evictable despite the rule. + */ + private void reconcileLinkedBundlesWithSelection( + StorageLocation loc, + List> uncommittedHolds + ) + { + final List reconcileNames; + entryLock.lock(); + try { + reconcileNames = new ArrayList<>(); + for (String name : ruleSelectedBundleNames) { + if (!ruleBundleHolds.containsKey(name) && findLinkedBundleByName(name) != null) { + reconcileNames.add(name); + } + } + } + finally { + entryLock.unlock(); + } + if (reconcileNames.isEmpty()) { + return; + } + + final Map> acquired = new HashMap<>(); + for (String name : reconcileNames) { + final StorageLocation.ReservationHold h = + loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, name)); + if (h != null) { + acquired.put(name, h); + uncommittedHolds.add(h); + } + } + + entryLock.lock(); + try { + for (var e : acquired.entrySet()) { + if (ruleSelectedBundleNames.contains(e.getKey()) && !ruleBundleHolds.containsKey(e.getKey())) { + ruleBundleHolds.put(e.getKey(), e.getValue()); + uncommittedHolds.remove(e.getValue()); + } + } + } + finally { + entryLock.unlock(); + } + } + /** * Release the partial-load rule applied to this entry. Closes every {@link StorageLocation.ReservationHold} taken by * {@link #applyRule(String, Set)} (the metadata self-hold and every bundle hold) and clears the rule state. Bundle