diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryVirtualStorageTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryVirtualStorageTest.java index 0e5a2a2baaf6..24e812a5f347 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryVirtualStorageTest.java +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryVirtualStorageTest.java @@ -73,6 +73,7 @@ class QueryVirtualStorageTest extends EmbeddedClusterTestBase private static final long SIZE_BYTES = 3778338L; private static final long CACHE_SIZE = HumanReadableBytes.parse("1MiB"); private static final long MAX_SIZE = HumanReadableBytes.parse("100MiB"); + private static final long ESTIMATE_SIZE = HumanReadableBytes.parse("2KiB"); private final EmbeddedBroker broker = new EmbeddedBroker(); private final EmbeddedIndexer indexer = new EmbeddedIndexer(); @@ -89,6 +90,11 @@ 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( @@ -155,8 +161,10 @@ void testQueryTooMuchData() RuntimeException.class, () -> cluster.runSql("select count(*) from \"%s\"", dataSource) ); - Assertions.assertTrue(t.getMessage().contains("Unable to load segment")); - Assertions.assertTrue(t.getMessage().contains("] on demand, ensure enough disk space has been allocated to load all segments involved in the query")); + Assertions.assertTrue(t.getMessage().contains("Unable to reserve bundle")); + Assertions.assertTrue(t.getMessage() + .contains( + "ensure enough disk space has been allocated to load all segments involved in the query")); } @Test @@ -296,11 +304,13 @@ void testQueryTooMuchDataButWithDart() MatcherAssert.assertThat(inputChannelSums.bytes(), Matchers.lessThanOrEqualTo(SIZE_BYTES)); // Verify stage 0 (segment read) VSF load counters - MatcherAssert.assertThat(inputChannelSums.loadFiles(), Matchers.greaterThan(0L)); + // partial loading is only partially metered at the moment, so depending on how stuff landed in and was evicted + // from the cache,there can be 0 loads (because loads are currently only counted when the metadata entry is mounted) + MatcherAssert.assertThat(inputChannelSums.loadFiles(), Matchers.greaterThanOrEqualTo(0L)); MatcherAssert.assertThat(inputChannelSums.loadFiles(), Matchers.lessThanOrEqualTo(24L)); - MatcherAssert.assertThat(inputChannelSums.loadTime(), Matchers.greaterThan(0L)); - MatcherAssert.assertThat(inputChannelSums.loadWait(), Matchers.greaterThan(0L)); - MatcherAssert.assertThat(inputChannelSums.loadBytes(), Matchers.greaterThan(0L)); + MatcherAssert.assertThat(inputChannelSums.loadTime(), Matchers.greaterThanOrEqualTo(0L)); + MatcherAssert.assertThat(inputChannelSums.loadWait(), Matchers.greaterThanOrEqualTo(0L)); + MatcherAssert.assertThat(inputChannelSums.loadBytes(), Matchers.greaterThanOrEqualTo(0L)); MatcherAssert.assertThat(inputChannelSums.loadBytes(), Matchers.lessThanOrEqualTo(SIZE_BYTES)); } @@ -322,7 +332,7 @@ private void assertQueryMetrics(int expectedEventCount, @Nullable Long expectedL long loadCount = getMetricLatestValue(emitter, DefaultQueryMetrics.QUERY_ON_DEMAND_LOAD_COUNT, expectedEventCount); if (expectedLoadCount != null) { - Assertions.assertEquals(expectedLoadCount, loadCount); + MatcherAssert.assertThat(loadCount, Matchers.lessThanOrEqualTo(expectedLoadCount)); } boolean hasLoads = loadCount > 0; diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/ServerManagerForQueryErrorTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/ServerManagerForQueryErrorTest.java index 6ff7c67b1eff..76f91f49edeb 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/ServerManagerForQueryErrorTest.java +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/ServerManagerForQueryErrorTest.java @@ -55,6 +55,7 @@ import org.apache.druid.segment.Segment; import org.apache.druid.segment.SegmentMapFunction; import org.apache.druid.segment.SegmentReference; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.server.SegmentManager; import org.apache.druid.server.ServerManager; import org.apache.druid.server.initialization.ServerConfig; @@ -225,7 +226,7 @@ protected LeafSegmentsBundle getLeafSegmentsBundle(Query query, SegmentMapFun missingSegments.add(segment.getDescriptor()); continue; } - Optional ref = segmentManager.acquireCachedSegment(dataSegment); + Optional ref = segmentManager.acquireCachedSegment(dataSegment, AcquireMode.FULL); if (ref.isPresent()) { segmentReferences.add( new SegmentReference( diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/task/CompactionTask.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/task/CompactionTask.java index 5804507f5d5e..7218f288255f 100644 --- a/indexing-service/src/main/java/org/apache/druid/indexing/common/task/CompactionTask.java +++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/task/CompactionTask.java @@ -95,6 +95,7 @@ import org.apache.druid.segment.indexing.CombinedDataSchema; import org.apache.druid.segment.indexing.DataSchema; import org.apache.druid.segment.indexing.TuningConfig; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.segment.loading.SegmentCacheManager; import org.apache.druid.segment.projections.AggregateProjectionSchema; @@ -890,7 +891,8 @@ private static ResourceHolder fetchSegmentInternal( { final Closer closer = Closer.create(); try { - final AcquireSegmentAction acquireAction = closer.register(segmentCacheManager.acquireSegment(dataSegment)); + final AcquireSegmentAction acquireAction = + closer.register(segmentCacheManager.acquireSegment(dataSegment, AcquireMode.FULL)); final ReferenceCountedObjectProvider segmentProvider = FutureUtils.getUnchecked(acquireAction.getSegmentFuture(), true).getReferenceProvider(); final Segment segment = segmentProvider.acquireReference().map(closer::register).get(); diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskTest.java index 6e3ae324b74c..48326727406c 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskTest.java @@ -132,6 +132,7 @@ import org.apache.druid.segment.indexing.DataSchema; import org.apache.druid.segment.indexing.TuningConfig; import org.apache.druid.segment.join.NoopJoinableFactory; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.segment.loading.AcquireSegmentResult; import org.apache.druid.segment.loading.NoopSegmentCacheManager; @@ -1978,7 +1979,7 @@ public void load(DataSegment segment) } @Override - public Optional acquireCachedSegment(SegmentId segmentId) + public Optional acquireCachedSegment(SegmentId segmentId, AcquireMode acquireMode) { for (Map.Entry entry : segments.entrySet()) { if (entry.getKey().getId().equals(segmentId)) { @@ -1991,7 +1992,7 @@ public Optional acquireCachedSegment(SegmentId segmentId) } @Override - public AcquireSegmentAction acquireSegment(DataSegment dataSegment) + public AcquireSegmentAction acquireSegment(DataSegment dataSegment, AcquireMode acquireMode) { final Segment segment = new QueryableIndexSegment(indexIO.loadIndex(segments.get(dataSegment)), dataSegment.getId()); diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractMultiPhaseParallelIndexingTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractMultiPhaseParallelIndexingTest.java index c4c2baa79997..abd74439d259 100644 --- a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractMultiPhaseParallelIndexingTest.java +++ b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractMultiPhaseParallelIndexingTest.java @@ -53,6 +53,7 @@ import org.apache.druid.segment.Segment; import org.apache.druid.segment.TestIndex; import org.apache.druid.segment.indexing.DataSchema; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.SegmentCacheManager; import org.apache.druid.segment.loading.SegmentLoadingException; import org.apache.druid.segment.loading.TombstoneLoadSpec; @@ -260,7 +261,7 @@ private Segment loadSegment(DataSegment dataSegment, File tempSegmentDir) .manufacturate(tempSegmentDir, null, false); try { cacheManager.load(dataSegment); - return cacheManager.acquireCachedSegment(dataSegment.getId()).orElseThrow(); + return cacheManager.acquireCachedSegment(dataSegment.getId(), AcquireMode.FULL).orElseThrow(); } catch (SegmentLoadingException e) { throw new RuntimeException(e); diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/AdaptedLoadableSegment.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/AdaptedLoadableSegment.java index 7000064be9b3..d66595946e1c 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/input/AdaptedLoadableSegment.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/AdaptedLoadableSegment.java @@ -31,6 +31,7 @@ import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.segment.ReferenceCountedSegmentProvider; import org.apache.druid.segment.Segment; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.segment.loading.AcquireSegmentResult; import org.apache.druid.timeline.DataSegment; @@ -129,16 +130,21 @@ public String description() } /** - * Adapted segments are not managed by SegmentManager, so they are never cached. + * Adapted segments are not managed by SegmentManager, so they are never cached. The {@code acquireMode} is ignored: + * an adapted segment is produced by its own async supplier, not the cache manager's full-vs-partial machinery. */ @Override - public Optional acquireIfCached() + public Optional acquireIfCached(AcquireMode acquireMode) { return Optional.empty(); } + /** + * The {@code acquireMode} is ignored: an adapted segment is produced by its own async supplier, not the cache + * manager's full-vs-partial machinery. + */ @Override - public AcquireSegmentAction acquire() + public AcquireSegmentAction acquire(AcquireMode acquireMode) { if (!acquired.compareAndSet(false, true)) { throw DruidException.defensive("Segment with descriptor[%s] is already acquired", descriptor); diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegment.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegment.java index 227c1ad2d492..3ecf39f1e1c9 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegment.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegment.java @@ -24,6 +24,7 @@ import org.apache.druid.error.DruidException; import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.segment.Segment; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.server.SegmentManager; import org.apache.druid.timeline.DataSegment; @@ -37,7 +38,8 @@ * {@link PhysicalInputSlice#getLoadableSegments()}, which are in turn generated by {@link InputSliceReader#attach}. * * These objects do not hold resources and do not have a lifecycle. They are just pointers to data that hasn't - * been acquired yet. Once {@link #acquire()} or {@link #acquireIfCached()} is called, the lifecycle begins. + * been acquired yet. Once {@link #acquire(AcquireMode)} or {@link #acquireIfCached(AcquireMode)} is called, the + * lifecycle begins. * * @see RegularLoadableSegment for segments loaded via {@link SegmentManager} * @see AdaptedLoadableSegment for segments adapted from other data sources @@ -61,23 +63,35 @@ public interface LoadableSegment * * If the returned {@link Optional} is present, the caller is responsible for closing the {@link Segment} * when finished with it. If the returned {@link Optional} is empty, the segment is not cached and the caller - * should use {@link #acquire()} to load it. + * should use {@link #acquire(AcquireMode)} to load it. + * + * The {@code acquireMode} selects what counts as a cached hit and how the segment is mounted; see {@link AcquireMode}. + * With {@link AcquireMode#PARTIAL} the returned {@link Segment} is mounted but, for a partial-download (virtual + * storage) segment, may not be fully downloaded yet: its data is fetched on demand at cursor-build time. Callers must + * therefore access it through the async API ({@link org.apache.druid.segment.CursorFactory#makeCursorHolderAsync}) + * rather than the synchronous {@code makeCursorHolder}, which refuses to download on a processing thread. * * @return an Optional containing the cached Segment if available, or empty if not cached */ - Optional acquireIfCached(); + Optional acquireIfCached(AcquireMode acquireMode); /** * Acquire the actual segment. Non-blocking operation. Once this is called, callers are responsible for closing the * {@link AcquireSegmentAction}. * + * The {@code acquireMode} selects how the segment is loaded; see {@link AcquireMode}. With {@link AcquireMode#PARTIAL} + * the returned {@link Segment} is mounted but, for a partial-download (virtual storage) segment, may not be fully + * downloaded yet: its data is fetched on demand at cursor-build time. Callers must therefore access it through the + * async API ({@link org.apache.druid.segment.CursorFactory#makeCursorHolderAsync}) rather than the + * synchronous {@code makeCursorHolder}, which refuses to download on a processing thread. + * * @throws DruidException if the segment has already been acquired */ - AcquireSegmentAction acquire(); + AcquireSegmentAction acquire(AcquireMode acquireMode); /** * Returns a future for the {@link DataSegment} object. For {@link RegularLoadableSegment}, the future is created - * lazily when this method is first called, or when a segment is acquired through {@link #acquire()}. + * lazily when this method is first called, or when a segment is acquired through {@link #acquire(AcquireMode)}. * For {@link AdaptedLoadableSegment}, this returns a failed future. */ ListenableFuture dataSegmentFuture(); diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java index b43eefdcf36a..54c9cf10a150 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java @@ -32,6 +32,7 @@ import org.apache.druid.query.SegmentDescriptor; import org.apache.druid.query.TableDataSource; import org.apache.druid.segment.Segment; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.server.SegmentManager; import org.apache.druid.timeline.DataSegment; @@ -127,13 +128,13 @@ public String description() } @Override - public synchronized Optional acquireIfCached() + public synchronized Optional acquireIfCached(AcquireMode acquireMode) { if (acquired) { throw DruidException.defensive("Segment with descriptor[%s] is already acquired", descriptor); } - final Optional cachedSegment = segmentManager.acquireCachedSegment(segmentId); + final Optional cachedSegment = segmentManager.acquireCachedSegment(segmentId, acquireMode); if (cachedSegment.isPresent()) { acquired = true; @@ -148,7 +149,7 @@ public synchronized Optional acquireIfCached() } @Override - public synchronized AcquireSegmentAction acquire() + public synchronized AcquireSegmentAction acquire(AcquireMode acquireMode) { if (acquired) { throw DruidException.defensive("Segment with descriptor[%s] is already acquired", descriptor); @@ -157,7 +158,7 @@ public synchronized AcquireSegmentAction acquire() acquired = true; if (cachedDataSegment != null) { - final AcquireSegmentAction action = segmentManager.acquireSegment(cachedDataSegment); + final AcquireSegmentAction action = segmentManager.acquireSegment(cachedDataSegment, acquireMode); return new AcquireSegmentAction( () -> LoadableSegmentUtils.countedLoad( action.getSegmentFuture(), @@ -177,7 +178,7 @@ public synchronized AcquireSegmentAction acquire() Suppliers.memoize(() -> FutureUtils.transformAsync( dataSegmentFutureSupplier.get(), dataSegment -> LoadableSegmentUtils.countedLoad( - closer.register(segmentManager.acquireSegment(dataSegment)).getSegmentFuture(), + closer.register(segmentManager.acquireSegment(dataSegment, acquireMode)).getSegmentFuture(), dataSegment.getSize(), inputCounters ) diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafStageProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafStageProcessor.java index f655010b4f15..608ea80fe21b 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafStageProcessor.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafStageProcessor.java @@ -55,6 +55,7 @@ import org.apache.druid.query.filter.SegmentPruner; import org.apache.druid.query.planning.ExecutionVertex; import org.apache.druid.segment.SegmentMapFunction; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.utils.CollectionUtils; import javax.annotation.Nullable; @@ -263,6 +264,23 @@ protected List filterBaseInput(final List slices, + final int loadahead + ) + { + return new ReadableInputQueue(partitionReader, slices, loadahead, AcquireMode.FULL); + } + /** * Read base inputs, where "base" is meant in the same sense as in {@link ExecutionVertex}: the primary datasource * that drives query processing. @@ -293,7 +311,7 @@ private ReadableInputQueue makeBaseInputQueue( } final List filteredSlices = filterBaseInput(physicalInputSlices); - return new ReadableInputQueue( + return makeReadableInputQueue( new StandardPartitionReader(context), filteredSlices, context.segmentLoadAheadCount() diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java index c2fe66ed0f48..6afe269e4371 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java @@ -34,6 +34,7 @@ import org.apache.druid.msq.input.stage.ReadablePartition; import org.apache.druid.segment.Segment; import org.apache.druid.segment.SegmentReference; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.utils.CloseableUtils; @@ -98,16 +99,25 @@ public class ReadableInputQueue implements Closeable private final StandardPartitionReader partitionReader; private final int loadahead; + + /** + * How segments in this queue are acquired (fully up front, or partially with on-demand column loading at query + * time). Threaded through to {@link LoadableSegment#acquire(AcquireMode)} and + * {@link LoadableSegment#acquireIfCached(AcquireMode)}. + */ + private final AcquireMode acquireMode; private final AtomicBoolean started = new AtomicBoolean(false); public ReadableInputQueue( final StandardPartitionReader partitionReader, final List slices, - final int loadahead + final int loadahead, + final AcquireMode acquireMode ) { this.partitionReader = partitionReader; this.loadahead = loadahead; + this.acquireMode = acquireMode; for (final PhysicalInputSlice slice : slices) { loadableSegments.addAll(slice.getLoadableSegments()); @@ -131,7 +141,7 @@ public void start() final List toLoad = new ArrayList<>(); // Temporarily store all non-cached segments LoadableSegment loadableSegment; while ((loadableSegment = loadableSegments.poll()) != null) { - final Optional cachedSegment = loadableSegment.acquireIfCached(); + final Optional cachedSegment = loadableSegment.acquireIfCached(acquireMode); if (cachedSegment.isPresent()) { final SegmentReferenceHolder holder = new SegmentReferenceHolder( new SegmentReference(loadableSegment.descriptor(), cachedSegment, null), @@ -288,7 +298,7 @@ private ListenableFuture loadNextSegment() return null; } - final AcquireSegmentAction acquireSegmentAction = nextLoadableSegment.acquire(); + final AcquireSegmentAction acquireSegmentAction = nextLoadableSegment.acquire(acquireMode); loadingSegments.add(acquireSegmentAction); return FutureUtils.transform( acquireSegmentAction.getSegmentFuture(), diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleStageProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleStageProcessor.java index 36eee4bef94c..208c57eefe40 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleStageProcessor.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleStageProcessor.java @@ -30,11 +30,13 @@ import org.apache.druid.frame.write.FrameWriterFactory; import org.apache.druid.msq.exec.ExecutionContext; import org.apache.druid.msq.exec.FrameContext; +import org.apache.druid.msq.exec.std.StandardPartitionReader; import org.apache.druid.msq.exec.std.StandardStageRunner; import org.apache.druid.msq.input.LoadableSegment; import org.apache.druid.msq.input.PhysicalInputSlice; import org.apache.druid.msq.querykit.BaseLeafStageProcessor; import org.apache.druid.msq.querykit.ReadableInput; +import org.apache.druid.msq.querykit.ReadableInputQueue; import org.apache.druid.msq.util.MultiStageQueryContext; import org.apache.druid.query.QueryToolChestWarehouse; import org.apache.druid.query.aggregation.AggregatorFactory; @@ -42,6 +44,7 @@ import org.apache.druid.query.groupby.GroupingEngine; import org.apache.druid.segment.SegmentMapFunction; import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.loading.AcquireMode; import org.joda.time.Interval; import javax.annotation.Nullable; @@ -169,6 +172,18 @@ protected List filterBaseInput(final List slices, + int loadahead + ) + { + // GroupBy reads segments through the async cursor API, so acquire them partially: only the columns the query + // touches are downloaded. + return new ReadableInputQueue(partitionReader, slices, loadahead, AcquireMode.PARTIAL); + } + @Override public boolean usesProcessingBuffers() { diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryStageProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryStageProcessor.java index 4de78269dd32..1e7fbe215c50 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryStageProcessor.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryStageProcessor.java @@ -28,8 +28,11 @@ import org.apache.druid.frame.processor.FrameProcessor; import org.apache.druid.frame.write.FrameWriterFactory; import org.apache.druid.msq.exec.FrameContext; +import org.apache.druid.msq.exec.std.StandardPartitionReader; +import org.apache.druid.msq.input.PhysicalInputSlice; import org.apache.druid.msq.querykit.BaseLeafStageProcessor; import org.apache.druid.msq.querykit.ReadableInput; +import org.apache.druid.msq.querykit.ReadableInputQueue; import org.apache.druid.query.DataSource; import org.apache.druid.query.Druids; import org.apache.druid.query.filter.DimFilter; @@ -38,9 +41,11 @@ import org.apache.druid.segment.SegmentMapFunction; import org.apache.druid.segment.VirtualColumns; import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.sql.calcite.filtration.Filtration; import javax.annotation.Nullable; +import java.util.List; import java.util.concurrent.atomic.AtomicLong; @JsonTypeName("scan") @@ -123,6 +128,18 @@ protected FrameProcessor makeProcessor( ); } + @Override + protected ReadableInputQueue makeReadableInputQueue( + StandardPartitionReader partitionReader, + List slices, + int loadahead + ) + { + // Scan reads segments through the async cursor API, so acquire them partially: only the columns the query + // touches are downloaded. + return new ReadableInputQueue(partitionReader, slices, loadahead, AcquireMode.PARTIAL); + } + @Override public boolean usesProcessingBuffers() { diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQCompactionTaskRunTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQCompactionTaskRunTest.java index e8c49f918d26..db881a76b5fd 100644 --- a/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQCompactionTaskRunTest.java +++ b/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQCompactionTaskRunTest.java @@ -214,7 +214,7 @@ public void setUpMSQ() ComplexMetrics.registerSerde(NestedDataComplexTypeSerde.TYPE_NAME, NestedDataComplexTypeSerde.INSTANCE); SegmentCacheManager segmentCacheManager = mock(SegmentCacheManager.class); - when(segmentCacheManager.acquireSegment(any())).thenAnswer(invocation -> { + when(segmentCacheManager.acquireSegment(any(), any())).thenAnswer(invocation -> { DataSegment segment = invocation.getArgument(0); QueryableIndexSegment index = new QueryableIndexSegment( new TestUtils().getTestIndexIO().loadIndex(new File((String) segment.getLoadSpec().get("path"))), @@ -225,7 +225,7 @@ public void setUpMSQ() null ); }); - when(segmentCacheManager.acquireCachedSegment(any())).thenReturn(Optional.empty()); + when(segmentCacheManager.acquireCachedSegment(any(), any())).thenReturn(Optional.empty()); GroupingEngine groupingEngine = GroupByQueryRunnerTest.makeQueryRunnerFactory( new GroupByQueryConfig(), TestGroupByBuffers.createDefault() diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/input/RegularLoadableSegmentTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/input/RegularLoadableSegmentTest.java index e68780677b70..003d8a6275cc 100644 --- a/multi-stage-query/src/test/java/org/apache/druid/msq/input/RegularLoadableSegmentTest.java +++ b/multi-stage-query/src/test/java/org/apache/druid/msq/input/RegularLoadableSegmentTest.java @@ -48,6 +48,7 @@ import org.apache.druid.segment.Segment; import org.apache.druid.segment.TestHelper; import org.apache.druid.segment.TestIndex; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.segment.loading.AcquireSegmentResult; import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStrategy; @@ -165,7 +166,9 @@ public void setUp() throws Exception cacheDir = tempDir.resolve("cache").toFile(); final SegmentLoaderConfig virtualLoaderConfig = new SegmentLoaderConfig() .setLocations(ImmutableList.of(new StorageLocationConfig(cacheDir, 10_000_000_000L, null))) - .setVirtualStorage(true).setVirtualStorageIsEphemeral(true); + .setVirtualStorage(true) + .setVirtualStorageIsEphemeral(true) + .setVirtualStoragePartialDownloadsEnabled(true); final List virtualLocations = virtualLoaderConfig.toStorageLocations(); segmentManagerDynamic = new SegmentManager( new SegmentLocalCacheManager( @@ -236,7 +239,7 @@ public void test_concurrency_dynamic() FutureUtils.transformAsync( f, ls -> { - final AcquireSegmentAction acquireAction = ls.acquire(); + final AcquireSegmentAction acquireAction = ls.acquire(AcquireMode.PARTIAL); return FutureUtils.transform(acquireAction.getSegmentFuture(), f2 -> Pair.of(acquireAction, f2)); } ), @@ -307,7 +310,7 @@ public void test_concurrency_preLoaded() throws SegmentLoadingException, IOExcep FutureUtils.transformAsync( f, ls -> { - final AcquireSegmentAction acquireAction = ls.acquire(); + final AcquireSegmentAction acquireAction = ls.acquire(AcquireMode.PARTIAL); return FutureUtils.transform(acquireAction.getSegmentFuture(), f2 -> Pair.of(acquireAction, f2)); } ), @@ -372,7 +375,7 @@ public void test_acquireIfCached_preLoaded() throws SegmentLoadingException, IOE ); // acquireIfCached should return a segment since it's loaded - final Optional cachedSegment = loadableSegment.acquireIfCached(); + final Optional cachedSegment = loadableSegment.acquireIfCached(AcquireMode.PARTIAL); Assertions.assertTrue(cachedSegment.isPresent()); try (final Segment acquiredSegment = cachedSegment.get()) { @@ -405,7 +408,7 @@ public void test_acquireIfCached_dynamic() ); // acquireIfCached should return empty since it's not loaded locally - final Optional cachedSegment = loadableSegment.acquireIfCached(); + final Optional cachedSegment = loadableSegment.acquireIfCached(AcquireMode.PARTIAL); Assertions.assertFalse(cachedSegment.isPresent()); } @@ -432,7 +435,7 @@ public void test_fetchSegment_dynamic() throws IOException Assertions.assertEquals(segment, fetchedDataSegment); // Verify segment acquisition works. - final AcquireSegmentAction acquireAction = loadableSegment.acquire(); + final AcquireSegmentAction acquireAction = loadableSegment.acquire(AcquireMode.PARTIAL); final AcquireSegmentResult acquireResult = FutureUtils.getUnchecked(acquireAction.getSegmentFuture(), false); final Optional acquiredSegmentOptional = acquireResult.getReferenceProvider().acquireReference(); Assertions.assertTrue(acquiredSegmentOptional.isPresent()); @@ -471,7 +474,7 @@ public void test_fetchSegment_preLoaded() throws IOException, SegmentLoadingExce Assertions.assertEquals(segment, fetchedDataSegment); // Verify segment acquisition works. - final AcquireSegmentAction acquireAction = loadableSegment.acquire(); + final AcquireSegmentAction acquireAction = loadableSegment.acquire(AcquireMode.PARTIAL); final AcquireSegmentResult acquireResult = FutureUtils.getUnchecked(acquireAction.getSegmentFuture(), false); final Optional acquiredSegmentOptional = acquireResult.getReferenceProvider().acquireReference(); Assertions.assertTrue(acquiredSegmentOptional.isPresent()); diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java index 651bcaf35795..54fc466c6180 100644 --- a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java +++ b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java @@ -156,6 +156,7 @@ import org.apache.druid.segment.column.ColumnHolder; import org.apache.druid.segment.column.RowSignature; import org.apache.druid.segment.incremental.IncrementalIndexSchema; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.DataSegmentPusher; import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStrategy; import org.apache.druid.segment.loading.LocalDataSegmentPusher; @@ -1424,7 +1425,10 @@ public void verifyResults() dataSegment.getDataSource() ); } - FutureUtils.getUnchecked(segmentCacheManager.acquireSegment(dataSegment).getSegmentFuture(), false); + FutureUtils.getUnchecked( + segmentCacheManager.acquireSegment(dataSegment, AcquireMode.FULL).getSegmentFuture(), + false + ); final QueryableIndex queryableIndex = indexIO.loadIndex(segmentCacheManager.getSegmentFiles(dataSegment)); final CursorFactory cursorFactory = new QueryableIndexCursorFactory(queryableIndex); diff --git a/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java b/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java index e22edd06750f..91d5a7f837b9 100644 --- a/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java +++ b/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java @@ -19,7 +19,12 @@ package org.apache.druid.common.asyncresource; +import com.google.common.util.concurrent.FutureCallback; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.MoreExecutors; import org.apache.druid.collections.ResourceHolder; +import org.apache.druid.utils.CloseableUtils; import java.io.Closeable; import java.util.List; @@ -55,6 +60,79 @@ public static AsyncResource ofCloseable(final T object) return retVal; } + /** + * Returns an {@link AsyncResource} backed by a {@link ListenableFuture}: it becomes ready when the future + * completes, exposing the future's (non-null) result via {@link AsyncResource#get()}. Closing the returned + * resource before the future completes cancels the future ({@code cancel(true)}). + *

+ * The result object must not have a lifecycle; closing the resource does not close the result. Use this for + * futures whose result is a plain value or a completion signal; use {@link #fromFutureCloseable} or + * {@link #ofCloseable} if the result has a lifecycle of its own. + */ + public static AsyncResource fromFutureUnmanaged(final ListenableFuture future) + { + final SettableAsyncResource retVal = new SettableAsyncResource<>(); + retVal.setCanceler(() -> future.cancel(true)); + Futures.addCallback( + future, + new FutureCallback<>() + { + @Override + public void onSuccess(T result) + { + retVal.set(result, null); + } + + @Override + public void onFailure(Throwable t) + { + retVal.setException(t); + } + }, + MoreExecutors.directExecutor() + ); + return retVal; + } + + /** + * Returns an {@link AsyncResource} backed by a {@link ListenableFuture} whose result owns a lifecycle: it + * becomes ready when the future completes, exposing the result via {@link AsyncResource#get()}, and the result is + * managed as a {@link Closeable}. Closing the resource closes the result, and a result that completes after + * the resource was already closed (a cancel/close-vs-completion race) is closed rather than leaked. Closing the + * returned resource before the future completes cancels the future ({@code cancel(true)}). + * + *

This is the managed counterpart of {@link #fromFutureUnmanaged}; use that for a future whose result is a plain + * value or a completion signal with no lifecycle. + */ + public static AsyncResource fromFutureCloseable(final ListenableFuture future) + { + final SettableAsyncResource retVal = new SettableAsyncResource<>(); + retVal.setCanceler(() -> future.cancel(true)); + Futures.addCallback( + future, + new FutureCallback<>() + { + @Override + public void onSuccess(T result) + { + // Lost the race with close()/cancel(): the resource is already closed, so set() returns false and we own + // closing the now-orphaned result. + if (!retVal.set(ResourceHolder.fromCloseable(result))) { + CloseableUtils.closeAndSuppressExceptions(result, ignored -> {}); + } + } + + @Override + public void onFailure(Throwable t) + { + retVal.setException(t); + } + }, + MoreExecutors.directExecutor() + ); + return retVal; + } + /** * Returns an {@link AsyncResource} that collects a list of underlying resources into a single lifecycle. * Calling {@link AsyncResource#close()} on the returned async resource causes the underlying async resource diff --git a/processing/src/main/java/org/apache/druid/segment/FilteredCursorFactory.java b/processing/src/main/java/org/apache/druid/segment/FilteredCursorFactory.java index 007543f76019..796ef59b8d20 100644 --- a/processing/src/main/java/org/apache/druid/segment/FilteredCursorFactory.java +++ b/processing/src/main/java/org/apache/druid/segment/FilteredCursorFactory.java @@ -46,6 +46,15 @@ public CursorHolder makeCursorHolder(CursorBuildSpec spec) return delegate.makeCursorHolder(CursorBuildSpec.builder(spec).andFilter(filter.toFilter()).build()); } + @Override + public AsyncCursorHolder makeCursorHolderAsync(CursorBuildSpec spec) + { + if (filter == null) { + return delegate.makeCursorHolderAsync(spec); + } + return delegate.makeCursorHolderAsync(CursorBuildSpec.builder(spec).andFilter(filter.toFilter()).build()); + } + @Override public RowSignature getRowSignature() { diff --git a/processing/src/main/java/org/apache/druid/segment/PartialBundleAcquirer.java b/processing/src/main/java/org/apache/druid/segment/PartialBundleAcquirer.java new file mode 100644 index 000000000000..dbd1632939d6 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/segment/PartialBundleAcquirer.java @@ -0,0 +1,51 @@ +/* + * 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; + +import org.apache.druid.common.asyncresource.AsyncResource; + +import java.io.Closeable; +import java.util.concurrent.Callable; + +/** + * Context object that bundles together pieces of cache-layer machinery needed to do on-demand partial segment loading: + *

    + *
  1. A hook for acquiring a hold on the cache-layer bundle.
    2. + *
  2. A hook for submitting bundle download tasks.
    3. + *
+ */ +public interface PartialBundleAcquirer +{ + /** + * Acquire a hold on the cache-layer bundle with the given name. The returned {@link Closeable} releases the hold + * when closed (the bundle itself stays in the cache for re-use by subsequent acquires). + *

+ * Implementations must be safe to call concurrently. + */ + Closeable acquire(String bundleName); + + /** + * Submit a bundle download task, returning an {@link AsyncResource} that becomes ready when the task completes. + * Closing the returned resource before completion cancels the task. The underlying executor bounds the concurrency + * of on-demand load work, so callers may submit as many tasks as needed. The task must return a non-null value + * (it may be a completion token if the work is purely a side effect), since {@link AsyncResource} cannot hold null. + */ + AsyncResource submitDownload(Callable task); +} diff --git a/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndex.java b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndex.java index 93c7e1fe2b56..ec2f3eb00171 100644 --- a/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndex.java +++ b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndex.java @@ -41,6 +41,7 @@ import org.apache.druid.segment.file.SegmentFileMetadata; import org.apache.druid.segment.projections.AggregateProjectionSchema; import org.apache.druid.segment.projections.BaseTableProjectionSchema; +import org.apache.druid.segment.projections.ClusteredValueGroupsBaseTableSchema; import org.apache.druid.segment.projections.ConstantTimeColumn; import org.apache.druid.segment.projections.ProjectionMetadata; import org.apache.druid.segment.projections.Projections; @@ -74,6 +75,7 @@ public class PartialQueryableIndex implements QueryableIndex { private final Interval dataInterval; private final int baseNumRows; + private final ProjectionMetadata baseProjectionMetadata; private final Indexed availableDimensions; private final List columnNames; private final BitmapFactory bitmapFactory; @@ -123,6 +125,18 @@ public PartialQueryableIndex( baseProjection.getSchema().getName() ); final BaseTableProjectionSchema baseSchema = (BaseTableProjectionSchema) baseProjection.getSchema(); + // Clustered V10 segments keep their data in per-cluster-group bundles, matched from metadata much like + // projections, but partial loading does not wire that up yet (and the write side isn't available to build one for + // testing). Fail loudly here rather than mis-treating the clustered base summary as a plain base table, whose + // top-level columns are empty. Both acquire modes route partial-eligible segments through this index, so this + // guards the full path too. Remove once partial loading supports clustered segments. + if (baseSchema instanceof ClusteredValueGroupsBaseTableSchema) { + throw DruidException.defensive( + "Clustered V10 segments are not yet supported for partial loading (interval[%s])", + metadata.getInterval() + ); + } + this.baseProjectionMetadata = baseProjection; this.baseNumRows = baseProjection.getNumRows(); this.baseProjectionPrefix = Projections.getProjectionSegmentInternalFilePrefix(baseSchema); this.dataInterval = Intervals.of(metadata.getInterval()); @@ -215,6 +229,24 @@ public Metadata getMetadata() return reconstructedMetadata; } + /** + * Returns the {@link ProjectionMetadata} for the base table projection. + */ + public ProjectionMetadata getBaseProjectionMetadata() + { + return baseProjectionMetadata; + } + + /** + * Whether every internal file referenced by this index's segment metadata (including any attached external file + * mappers) is already downloaded. Used by partial loaded segments to gate sync paths to confirm the segment + * is fully loaded before performing operations. Only async paths are allowed to perform on-demand downloads. + */ + public boolean isFullyDownloaded() + { + return fileMapper.isFullyDownloaded(); + } + @Override public Map getDimensionHandlers() { @@ -268,20 +300,21 @@ public QueryableProjection getProjection(CursorBuildSpec cursorB cursorBuildSpec, projections, dataInterval, - (projectionName, columnName) -> { - // check if the projection has this column using metadata column descriptors - final ProjectionMetadata projSpec = projectionSpecs.get(projectionName); - if (projSpec == null) { - return false; - } - final String smooshName = Projections.getProjectionSegmentInternalFileName(projSpec.getSchema(), columnName); - return metadata.getColumnDescriptors().containsKey(smooshName) - || getColumnCapabilities(columnName) == null; - }, + this::checkProjectionHasColumn, this::getProjectionQueryableIndex ); } + private boolean checkProjectionHasColumn(String projectionName, String columnName) + { + final ProjectionMetadata meta = projectionSpecs.get(projectionName); + if (meta == null) { + return false; + } + final String file = Projections.getProjectionSegmentInternalFileName(meta.getSchema(), columnName); + return metadata.getColumnDescriptors().containsKey(file) || getColumnCapabilities(columnName) == null; + } + @Nullable @Override public QueryableIndex getProjectionQueryableIndex(String name) diff --git a/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexCursorFactory.java b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexCursorFactory.java new file mode 100644 index 000000000000..31de0f9a3fb0 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexCursorFactory.java @@ -0,0 +1,364 @@ +/* + * 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; + +import org.apache.druid.common.asyncresource.AsyncResource; +import org.apache.druid.common.asyncresource.AsyncResources; +import org.apache.druid.error.DruidException; +import org.apache.druid.java.util.common.io.Closer; +import org.apache.druid.query.Order; +import org.apache.druid.query.OrderBy; +import org.apache.druid.query.aggregation.AggregatorFactory; +import org.apache.druid.segment.column.ColumnCapabilities; +import org.apache.druid.segment.column.ColumnHolder; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.projections.Projections; +import org.apache.druid.segment.projections.QueryableProjection; +import org.apache.druid.segment.vector.VectorCursor; +import org.apache.druid.utils.CloseableUtils; + +import javax.annotation.Nullable; +import java.io.Closeable; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Partial-aware {@link CursorFactory} for {@link PartialQueryableIndex}. + *

+ * Sync vs async contract. {@link #makeCursorHolder} requires the segment to already be fully downloaded, + * intended for callers that acquired the segment via the eager (download-everything-up-front) path, so by the time + * they ask for a cursor every internal file is on disk. If anything is missing it throws + * {@link DruidException#defensive} so that we never trigger downloads on the sync path, since processing threads must + * not block on deep-storage I/O. {@link #makeCursorHolderAsync} is the only path that performs downloads on demand; + * callers acknowledge that by opting into the async variant when they acquire a partial segment. + *

+ * Async download granularity. Pre-fetch is column-level. {@link #makeCursorHolderAsync} calls + * {@link QueryableIndex#getColumnHolder} on each required column; the memoized supplier on the underlying + * {@link PartialQueryableIndex} eagerly invokes + * {@link org.apache.druid.segment.file.PartialSegmentFileMapperV10#mapFile} inside that call, which is what triggers + * the deep-storage range read. The cursor holder constructed afterward sees the already-materialized holders via the + * same memoized suppliers, so no further downloads happen at cursor-read time. + *

+ * If a projection matches, the required columns are looked up against the projection's row selector and its rewritten + * {@link CursorBuildSpec} (which carries physical columns in the projection's namespace). When + * {@link CursorBuildSpec#getPhysicalColumns()} is {@code null}, every column on the chosen row selector is pre-fetched + * as required by the contract of {@link CursorBuildSpec}. + *

+ * Parallelism. Each column's materialization is submitted as a separate task to the supplied download executor. + * The cursor holder is constructed once every column task has completed. + */ +public class PartialQueryableIndexCursorFactory implements CursorFactory +{ + private final PartialQueryableIndex index; + private final QueryableIndexCursorFactory delegate; + private final PartialBundleAcquirer bundleAcquirer; + + public PartialQueryableIndexCursorFactory( + PartialQueryableIndex index, + TimeBoundaryInspector timeBoundaryInspector, + PartialBundleAcquirer bundleAcquirer + ) + { + this.index = index; + this.delegate = new QueryableIndexCursorFactory(index, timeBoundaryInspector); + this.bundleAcquirer = bundleAcquirer; + } + + @Override + public CursorHolder makeCursorHolder(CursorBuildSpec spec) + { + // refuse to download here so we never accidentally block a processing thread + if (!index.isFullyDownloaded()) { + throw DruidException.defensive( + "Sync makeCursorHolder requires the segment to be fully downloaded; use makeCursorHolderAsync for " + + "on-demand loading, or acquire the segment via the eager path so all files are loaded up front." + ); + } + return delegate.makeCursorHolder(spec); + } + + @Override + public AsyncCursorHolder makeCursorHolderAsync(CursorBuildSpec spec) + { + final QueryableProjection matched = index.getProjection(spec); + final QueryableIndex rowSelector = matched != null ? matched.getRowSelector() : index; + final Set requiredColumns = requiredColumns(rowSelector, matched, spec); + final String bundleName = matched != null ? matched.getName() : Projections.BASE_TABLE_PROJECTION_NAME; + + // Mount the cache-layer bundle before submitting downloads. Released on the cancel/failure paths (requestRelease) + // or handed to the produced cursor holder on success (releaser), exactly once, and never while an in-flight column + // download is mid-mapFile(). See BundleHoldRelease. + final BundleHoldRelease holdRelease = new BundleHoldRelease(bundleAcquirer.acquire(bundleName)); + + try { + // submit one materialization task per column so a multi-threaded download executor can fan them out + final List> columnDownloads = new ArrayList<>(requiredColumns.size()); + for (String column : requiredColumns) { + columnDownloads.add(submitColumnDownload(rowSelector, column, holdRelease)); + } + final AsyncResource> downloaded = AsyncResources.collect(columnDownloads); + + // Canceler runs if the awaiter closes this holder before it's ready (e.g. query cancel/timeout). Close the + // collected resource to cancel every column download that hasn't begun its deep-storage read yet (queued tasks + // are skipped; tasks parked on the download executor's permit are interrupted out of the interruptible wait + // before doing any I/O), then request the hold release through the handshake. Once the holder is produced and + // handed to set(), ownership transfers to the awaiter, which drains it via close() (cancel) or release() + // (success); the once-guard keeps the hold release safe across all of these paths. + final AsyncCursorHolder asyncHolder = new AsyncCursorHolder(() -> { + CloseableUtils.closeAndSuppressExceptions(downloaded, ignored -> {}); + holdRelease.requestRelease(); + }); + downloaded.addReadyCallback(() -> { + final CursorHolder holder; + try { + downloaded.get(); // surfaces any column-download failure (or cancellation) as the cause + final CursorHolder inner = delegate.makeCursorHolderForProjection(spec, matched); + // The produced holder takes ownership of both inner and the bundle hold; from here, closing it releases both + holder = wrapWithBundleRelease(inner, holdRelease.releaser()); + } + catch (Throwable t) { + // A column download failed or was canceled, this branch can fire while a sibling download is still + // mid-mapFile(), so the hold release must go through the handshake rather than dropping it here directly. + holdRelease.requestRelease(); + asyncHolder.setException(t); + return; + } + if (!asyncHolder.set(holder)) { + // wrapper was closed (awaiter canceled) while we were producing the holder; close it ourselves so the + // holder, its inner, and the bundle hold don't leak. + holder.close(); + } + }); + return asyncHolder; + } + catch (Throwable t) { + // Failure while submitting downloads / wiring up the holder (e.g. submitDownload shutdown rejection). A column + // download submitted before the failure may already be in flight, so release the bundle hold through the + // handshake (requestRelease defers to the last in-flight body) rather than dropping it directly mid-mapFile(). + throw CloseableUtils.closeAndWrapInCatch(t, holdRelease::requestRelease); + } + } + + @Override + public RowSignature getRowSignature() + { + return delegate.getRowSignature(); + } + + @Nullable + @Override + public ColumnCapabilities getColumnCapabilities(String column) + { + return delegate.getColumnCapabilities(column); + } + + /** + * Submit one column's materialization as a download task, running its body under the bundle-hold handshake (see + * {@link BundleHoldRelease#runDownloadBody}) so the hold stays alive until this body's + * {@link QueryableIndex#getColumnHolder} (and the + * {@link org.apache.druid.segment.file.PartialSegmentFileMapperV10#mapFile} it triggers) has finished. The returned + * token (the column name) is unused, the task's effect is the side effect of materializing the column into the file + * mapper. + */ + private AsyncResource submitColumnDownload( + QueryableIndex rowSelector, + String column, + BundleHoldRelease holdRelease + ) + { + return bundleAcquirer.submitDownload(() -> { + holdRelease.runDownloadBody(() -> rowSelector.getColumnHolder(column)); + return column; + }); + } + + /** + * Determine the set of physical column names to required from the chosen row selector given a {@link CursorBuildSpec} + */ + private static Set requiredColumns( + QueryableIndex rowSelector, + @Nullable QueryableProjection matched, + CursorBuildSpec originalSpec + ) + { + final CursorBuildSpec effective = matched != null ? matched.getCursorBuildSpec() : originalSpec; + if (effective.getPhysicalColumns() != null) { + final Set required = new LinkedHashSet<>(effective.getPhysicalColumns()); + // physicalColumns enumerates the selected columns, but QueryableIndexCursorHolder also reads __time while + // building the cursor, independent of physicalColumns: unconditionally for a time-ordered index (its + // interval-checking offset reads timestamps), and via a synthesized __time range filter for a non-time-ordered + // index whose data extends past the query interval. That read happens after the cursor holder is handed back, + // so __time must be pre-fetched on the async path or it becomes a lazy deep-storage download on a processing + // thread. Predicting exactly when the holder reads __time would mean replicating its internals (fragile, and it + // reads __time in the common cases anyway), so always include it: __time is cheap, and it resolves to a + // no-download constant column for projections that don't carry a real time column. + required.add(ColumnHolder.TIME_COLUMN_NAME); + return required; + } + // Conservative fallback when physicalColumns isn't declared, fetch every column on the chosen row selector + // plus __time (which is special-cased and not enumerated by getColumnNames()). + final Set all = new LinkedHashSet<>(rowSelector.getColumnNames()); + all.add(ColumnHolder.TIME_COLUMN_NAME); + return all; + } + + /** + * Wrap a {@link CursorHolder} so that closing it also closes {@code bundleHold}, releasing the cache-layer + * reservation hold on the bundle's containers. + */ + private static CursorHolder wrapWithBundleRelease(CursorHolder inner, Closeable bundleHold) + { + return new CursorHolder() + { + @Override + public Cursor asCursor() + { + return inner.asCursor(); + } + + @Override + public VectorCursor asVectorCursor() + { + return inner.asVectorCursor(); + } + + @Override + public boolean canVectorize() + { + return inner.canVectorize(); + } + + @Override + public boolean isPreAggregated() + { + return inner.isPreAggregated(); + } + + @Override + @Nullable + public List getAggregatorsForPreAggregated() + { + return inner.getAggregatorsForPreAggregated(); + } + + @Override + public List getOrdering() + { + return inner.getOrdering(); + } + + @Override + public Order getTimeOrder() + { + return inner.getTimeOrder(); + } + + @Override + public void close() + { + final Closer closer = Closer.create(); + closer.register(bundleHold); + closer.register(inner); + CloseableUtils.closeAndWrapExceptions(closer); + } + }; + } + + /** + * Owns a bundle hold's release for one {@link #makeCursorHolderAsync} build, guarding the file mapper's + * evict-vs-mapFile contract: releasing the hold makes the bundle entry eligible for eviction, which unmaps the + * bundle's containers, so a column download still inside {@code mapFile()} would then read (and slice a ByteBuffer + * from) an unmapped region resulting in a JVM SIGSEGV. The hold is released exactly once (CAS-guarded) and only once + * a release has been {@link #requestRelease() requested} (the awaiter canceled, or a download failed) AND no + * download body is executing, so an in-flight {@code mapFile()} always finishes under the hold. On the success path + * the produced cursor holder takes ownership of the hold via {@link #releaser()} and the request/in-flight handshake + * is unused. + */ + private static final class BundleHoldRelease + { + private final Closeable releaseHoldOnce; + private final AtomicInteger bodiesInFlight = new AtomicInteger(0); + private final AtomicBoolean releaseRequested = new AtomicBoolean(false); + + BundleHoldRelease(Closeable bundleHold) + { + // Release the bundle hold at most once: the canceler, the success-path holder close, the failure callback, and + // the last download body can all race to release it. + final AtomicBoolean released = new AtomicBoolean(false); + this.releaseHoldOnce = () -> { + if (released.compareAndSet(false, true)) { + bundleHold.close(); + } + }; + } + + /** + * Run one column-materialization body under the in-flight count. {@code incrementAndGet}-before-read pairs with + * {@code releaseRequested.set(true)}-before-read in {@link #requestRelease()}: whichever thread observes the + * other's write takes responsibility for the hold release, so it happens exactly once and never while a + * {@code mapFile()} is in flight. A body that observes a release request before starting skips the read entirely + * (its containers are about to be released and may be evicted); the last body to finish after a release request + * releases the hold. + */ + void runDownloadBody(Runnable read) + { + bodiesInFlight.incrementAndGet(); + try { + if (!releaseRequested.get()) { + read.run(); + } + } + finally { + if (bodiesInFlight.decrementAndGet() == 0 && releaseRequested.get()) { + CloseableUtils.closeAndSuppressExceptions(releaseHoldOnce, ignored -> {}); + } + } + } + + /** + * Request the hold release because the awaiter cancelled or a download failed. Releases now only if no download + * body is executing; otherwise the last body out (see {@link #runDownloadBody}) releases it once its + * {@code mapFile()} has finished. Idempotent. + */ + void requestRelease() + { + releaseRequested.set(true); + if (bodiesInFlight.get() == 0) { + CloseableUtils.closeAndSuppressExceptions(releaseHoldOnce, ignored -> {}); + } + } + + /** + * The raw, idempotent bundle-hold close. Handed to the produced cursor holder on the success path so that closing + * the holder releases the hold; it bypasses the in-flight handshake, which is safe there because the holder is only + * built after every download has completed (nothing is mid-{@code mapFile()}). Cancel/failure paths must release + * through {@link #requestRelease()} instead; the once-guard keeps the hold release exactly-once across whichever + * path fires. + */ + Closeable releaser() + { + return releaseHoldOnce; + } + } +} diff --git a/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexSegment.java b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexSegment.java new file mode 100644 index 000000000000..02b94b00f6cd --- /dev/null +++ b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexSegment.java @@ -0,0 +1,114 @@ +/* + * 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; + +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.utils.CloseableUtils; +import org.joda.time.Interval; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.io.Closeable; + +/** + * {@link Segment} wrapper around a {@link PartialQueryableIndex}. Mirrors {@link QueryableIndexSegment} but wires up + * the V10-specific {@link V10TimeBoundaryInspector} and {@link V10RowCountInspector} (which answer from + * {@link org.apache.druid.segment.projections.ProjectionMetadata} min/max-time and row-count fields without + * downloading any column data) and the partial-aware {@link PartialQueryableIndexCursorFactory} (which downloads + * required files on the supplied download executor before handing back a cursor). + *

+ * Lifecycle: this segment is intended to exist as a transient reference-hold scope over an externally-owned + * {@link PartialQueryableIndex}, it never closes the underlying index. The {@code onClose} hook is what + * {@link #close()} fires when the segment is closed (e.g. to 'release' reference tracking stuff in the cache layer). + */ +public class PartialQueryableIndexSegment implements ReferenceCountedSegmentProvider.LeafReference +{ + private final PartialQueryableIndex index; + private final PartialQueryableIndexCursorFactory cursorFactory; + private final TimeBoundaryInspector timeBoundaryInspector; + private final RowCountInspector rowCountInspector; + private final SegmentId segmentId; + private final Closeable onClose; + + public PartialQueryableIndexSegment( + PartialQueryableIndex index, + SegmentId segmentId, + Closeable onClose, + PartialBundleAcquirer bundleAcquirer + ) + { + this.index = index; + this.timeBoundaryInspector = V10TimeBoundaryInspector.forBaseProjection( + index.getBaseProjectionMetadata(), + index.getDataInterval() + ); + this.rowCountInspector = V10RowCountInspector.forBaseProjection(index.getBaseProjectionMetadata()); + this.cursorFactory = new PartialQueryableIndexCursorFactory( + index, + timeBoundaryInspector, + bundleAcquirer + ); + this.segmentId = segmentId; + this.onClose = onClose; + } + + @Override + public SegmentId getId() + { + return segmentId; + } + + @Override + public Interval getDataInterval() + { + return index.getDataInterval(); + } + + @Override + public void close() + { + CloseableUtils.closeAndWrapExceptions(onClose); + } + + @SuppressWarnings("unchecked") + @Nullable + @Override + public T as(@Nonnull Class clazz) + { + if (CursorFactory.class.equals(clazz)) { + return (T) cursorFactory; + } else if (QueryableIndex.class.equals(clazz)) { + return (T) index; + } else if (TimeBoundaryInspector.class.equals(clazz)) { + return (T) timeBoundaryInspector; + } else if (RowCountInspector.class.equals(clazz)) { + return (T) rowCountInspector; + } else if (Metadata.class.equals(clazz)) { + return (T) index.getMetadata(); + } + return null; + } + + @Override + public String getDebugString() + { + return getClass().getSimpleName() + ":" + segmentId; + } +} diff --git a/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorFactory.java b/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorFactory.java index 330984030585..51c1482cc89b 100644 --- a/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorFactory.java +++ b/processing/src/main/java/org/apache/druid/segment/QueryableIndexCursorFactory.java @@ -103,6 +103,25 @@ public CursorHolder makeCursorHolder(CursorBuildSpec spec) return new QueryableIndexCursorHolder(index, spec, timeBoundaryInspector); } + public CursorHolder makeCursorHolderForProjection( + CursorBuildSpec spec, + @Nullable QueryableProjection projection + ) + { + if (projection != null) { + return makeAggregateProjectionCursorHolder(projection); + } + + // Cluster-group dispatch runs after aggregate-projection match, before the regular base-table fallback + final ClusteredValueGroupsBaseTableSchema clusterSummary = index.getClusteredBaseSummary(); + if (clusterSummary != null) { + return makeClusteredCursorHolder(spec); + } + + // No projections, no clustering, regular full-segment cursor. + return new QueryableIndexCursorHolder(index, spec, timeBoundaryInspector); + } + private CursorHolder makeAggregateProjectionCursorHolder(QueryableProjection projection) { return new QueryableIndexCursorHolder( diff --git a/processing/src/main/java/org/apache/druid/segment/ReferenceCountedSegmentProvider.java b/processing/src/main/java/org/apache/druid/segment/ReferenceCountedSegmentProvider.java index f41912b32aa9..724e2c2cebb0 100644 --- a/processing/src/main/java/org/apache/druid/segment/ReferenceCountedSegmentProvider.java +++ b/processing/src/main/java/org/apache/druid/segment/ReferenceCountedSegmentProvider.java @@ -87,15 +87,27 @@ public Optional acquireReference() } /** - * Base type to mark {@link Segment} returned by {@link ReferenceCountedObjectProvider} as 'leaf' segments, - * to distinguish from other transformations which can be done on top of this segment, such as by - * {@link SegmentMapFunction} + * Marker interface for {@link Segment} types returned by {@link ReferenceCountedObjectProvider}, used to + * distinguish 'leaf' segments (segments held by a refcount, eligible to be wrapped by transformations like + * {@link SegmentMapFunction} or {@link RestrictedSegment}) from arbitrary {@link Segment} implementations. + *

+ * Most implementations should extend {@link BaseLeafReference} to inherit the standard delegating implementations + * of {@link Segment#getId}, {@link Segment#getDataInterval}, etc. A {@link Segment} that already provides its own + * delegation (e.g. {@link PartialQueryableIndexSegment}) can implement this interface directly as a pure marker. */ - public abstract static class LeafReference implements Segment + public interface LeafReference extends Segment + { + } + + /** + * Convenient base for {@link LeafReference} implementations that wrap a {@link Segment} and delegate the standard + * accessors to it. Subclasses implement {@link Segment#as} and {@link Segment#close}. + */ + public abstract static class BaseLeafReference implements LeafReference { protected final Segment baseSegment; - public LeafReference(Segment baseSegment) + public BaseLeafReference(Segment baseSegment) { this.baseSegment = baseSegment; } @@ -138,7 +150,7 @@ public void validateOrElseThrow(PolicyEnforcer policyEnforcer) * Wraps a {@link Segment} and decrements reference to this segment of {@link ReferenceCountedSegmentProvider} on * close (instead of closing the segment itself which is managed through the provider) */ - public final class ReferenceClosingSegment extends LeafReference + public final class ReferenceClosingSegment extends BaseLeafReference { private final Closeable referenceCloseable; private final AtomicBoolean isClosed = new AtomicBoolean(false); @@ -178,7 +190,7 @@ public ReferenceCountedSegmentProvider getProvider() } } - public static final class UnmanagedReference extends LeafReference + public static final class UnmanagedReference extends BaseLeafReference { public UnmanagedReference(Segment delegate) { @@ -199,7 +211,7 @@ public void close() } } - public static final class CloseableWrappedReference extends LeafReference + public static final class CloseableWrappedReference extends BaseLeafReference { private final Closeable closeable; diff --git a/processing/src/main/java/org/apache/druid/segment/RestrictedCursorFactory.java b/processing/src/main/java/org/apache/druid/segment/RestrictedCursorFactory.java index e033aac75a95..f2c56ecb78a2 100644 --- a/processing/src/main/java/org/apache/druid/segment/RestrictedCursorFactory.java +++ b/processing/src/main/java/org/apache/druid/segment/RestrictedCursorFactory.java @@ -59,6 +59,12 @@ public CursorHolder makeCursorHolder(CursorBuildSpec spec) return delegate.makeCursorHolder(policy.visit(spec)); } + @Override + public AsyncCursorHolder makeCursorHolderAsync(CursorBuildSpec spec) + { + return delegate.makeCursorHolderAsync(policy.visit(spec)); + } + @Override public RowSignature getRowSignature() { diff --git a/processing/src/main/java/org/apache/druid/segment/V10RowCountInspector.java b/processing/src/main/java/org/apache/druid/segment/V10RowCountInspector.java new file mode 100644 index 000000000000..415962858e75 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/segment/V10RowCountInspector.java @@ -0,0 +1,50 @@ +/* + * 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; + +import org.apache.druid.segment.projections.ProjectionMetadata; + +/** + * {@link RowCountInspector} for V10 segments. Reads the row count from the base projection's + * {@link ProjectionMetadata#getNumRows()} (persisted at write time), so it answers without reading any column data. + */ +public final class V10RowCountInspector implements RowCountInspector +{ + /** + * Build an inspector that reads the row count from the given base projection metadata. + */ + public static V10RowCountInspector forBaseProjection(ProjectionMetadata baseProjection) + { + return new V10RowCountInspector(baseProjection.getNumRows()); + } + + private final int numRows; + + private V10RowCountInspector(int numRows) + { + this.numRows = numRows; + } + + @Override + public int getNumRows() + { + return numRows; + } +} diff --git a/processing/src/main/java/org/apache/druid/segment/V10TimeBoundaryInspector.java b/processing/src/main/java/org/apache/druid/segment/V10TimeBoundaryInspector.java new file mode 100644 index 000000000000..e35227b4ba39 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/segment/V10TimeBoundaryInspector.java @@ -0,0 +1,82 @@ +/* + * 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; + +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.segment.column.ColumnHolder; +import org.apache.druid.segment.projections.ProjectionMetadata; +import org.joda.time.DateTime; +import org.joda.time.Interval; + +import javax.annotation.Nullable; + +/** + * {@link TimeBoundaryInspector} for V10 segments. Reads min/max {@link ColumnHolder#TIME_COLUMN_NAME} values from the + * base projection's {@link ProjectionMetadata} (persisted at write time by {@link IndexMergerBase}) so the inspector + * can answer without reading the {@code __time} column. Falls back to the segment's data interval (inexact) when the + * writer didn't supply min/max, for older V10 segments produced before the field existed. + *

+ * The exact-bounds path is independent of segment sort order: the writer records the actual minimum and maximum + * timestamps across all walked rows even when the segment is not time-sorted. This is a strict improvement over + * the eager-load {@link QueryableIndexTimeBoundaryInspector}, which reports inexact bounds for non-time-sorted + * segments because its only option is to read positions 0 and N-1 of the {@code __time} column. + */ +public final class V10TimeBoundaryInspector implements TimeBoundaryInspector +{ + /** + * Build an inspector that reads from the given base projection metadata, falling back to {@code fallbackInterval} + * when the writer didn't persist min/max times. + */ + public static V10TimeBoundaryInspector forBaseProjection(ProjectionMetadata baseProjection, Interval fallbackInterval) + { + return new V10TimeBoundaryInspector(baseProjection.getMinTime(), baseProjection.getMaxTime(), fallbackInterval); + } + + @Nullable + private final Long minTime; + @Nullable + private final Long maxTime; + private final Interval fallbackInterval; + + private V10TimeBoundaryInspector(@Nullable Long minTime, @Nullable Long maxTime, Interval fallbackInterval) + { + this.minTime = minTime; + this.maxTime = maxTime; + this.fallbackInterval = fallbackInterval; + } + + @Override + public DateTime getMinTime() + { + return minTime != null ? DateTimes.utc(minTime) : fallbackInterval.getStart(); + } + + @Override + public DateTime getMaxTime() + { + return maxTime != null ? DateTimes.utc(maxTime) : fallbackInterval.getEnd(); + } + + @Override + public boolean isMinMaxExact() + { + return minTime != null && maxTime != null; + } +} 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 5db4914a464a..c622584bea99 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 @@ -23,12 +23,14 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.io.ByteStreams; import com.google.common.primitives.Ints; +import org.apache.druid.collections.ResourceHolder; import org.apache.druid.error.DruidException; import org.apache.druid.java.util.common.ByteBufferUtils; import org.apache.druid.java.util.common.FileUtils; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.segment.CompressedPools; import org.apache.druid.segment.data.CompressionStrategy; import org.apache.druid.segment.loading.SegmentRangeReader; import org.apache.druid.utils.CloseableUtils; @@ -239,6 +241,15 @@ static PartialSegmentFileMapperV10 createForFile( private final File[] containerFiles; private final ReentrantLock[] containerLocks; + // bundle name -> indices (into metadata.getContainers()) of this single mapper's containers in that bundle. + // Computed once at construction from the immutable container metadata. Single-mapper scope only: stitching bundles + // across attached external mappers is the cache layer's concern (PartialSegmentBundleCacheEntry). + private final Map> bundleToContainerIndices; + + // file names per container index (parallel to metadata.getContainers()), for whole-container bulk download. Built + // once from the immutable metadata. + private final List> containerFileNames; + // external file mappers private final Map externalMappers = new HashMap<>(); @@ -273,13 +284,29 @@ private PartialSegmentFileMapperV10( fileNameToIndex.put(sortedFileNames.get(i), i); } - final int numContainers = metadata.getContainers().size(); + final List containerMetas = metadata.getContainers(); + final int numContainers = containerMetas.size(); this.containers = new MappedByteBuffer[numContainers]; this.containerFiles = new File[numContainers]; this.containerLocks = new ReentrantLock[numContainers]; + final Map> bundleIndices = new HashMap<>(); for (int i = 0; i < numContainers; i++) { this.containerLocks[i] = new ReentrantLock(); + bundleIndices.computeIfAbsent(containerMetas.get(i).getBundle(), name -> new ArrayList<>()).add(i); + } + bundleIndices.replaceAll((name, indices) -> List.copyOf(indices)); + this.bundleToContainerIndices = Map.copyOf(bundleIndices); + + final List> perContainer = new ArrayList<>(numContainers); + for (int i = 0; i < numContainers; i++) { + perContainer.add(new ArrayList<>()); } + for (Map.Entry entry : metadata.getFiles().entrySet()) { + perContainer.get(entry.getValue().getContainer()).add(entry.getKey()); + } + perContainer.replaceAll(List::copyOf); + this.containerFileNames = List.copyOf(perContainer); + this.bitmapLock = new ReentrantLock(); } @@ -288,6 +315,24 @@ public SegmentFileMetadata getSegmentFileMetadata() return metadata; } + /** + * The distinct bundle names present in this single mapper's containers (not including any attached external + * mappers). Computed once at construction. + */ + public Set getBundleNames() + { + return bundleToContainerIndices.keySet(); + } + + /** + * Indices, into this mapper's {@link SegmentFileMetadata#getContainers()} list, of the containers belonging to + * {@code bundleName}; empty if this mapper has no containers in that bundle. + */ + public List getContainerIndicesForBundle(String bundleName) + { + return bundleToContainerIndices.getOrDefault(bundleName, List.of()); + } + /** * Names of the external segment files attached to this mapper (each one is its own {@link PartialSegmentFileMapperV10} * accessible via {@link #getExternalMapper}). Empty for mappers with no externals. @@ -378,6 +423,76 @@ public ByteBuffer mapExternalFile(String filename, String name) throws IOExcepti return externalMapper.mapFile(name); } + /** + * Download every internal file referenced by this mapper's metadata (and recursively every attached external + * mapper's metadata) so that {@link #isFullyDownloaded} returns true afterward. Used by the eager + * {@code acquireSegment} path on partial-eligible segments to force the segment fully resident before returning. + * Containers already fully downloaded are skipped. + */ + public void ensureAllDownloaded() throws IOException + { + for (int containerIndex = 0; containerIndex < containers.length; containerIndex++) { + downloadContainer(containerIndex); + } + for (PartialSegmentFileMapperV10 external : externalMappers.values()) { + external.ensureAllDownloaded(); + } + } + + /** + * Download every container belonging to {@code bundleName} in this mapper, each in a single range read (see + * {@link #downloadContainer}). No-op for an unknown bundle. + */ + public void ensureBundleDownloaded(String bundleName) throws IOException + { + checkClosed(); + for (int containerIndex : getContainerIndicesForBundle(bundleName)) { + downloadContainer(containerIndex); + } + } + + /** + * Download an entire container in a single range read and mark every internal file it holds as downloaded. The whole + * region streams straight to the container's local sparse file at offset 0. + *

+ * No-op when every file in the container is already downloaded. A partially-downloaded container is re-fetched in + * full (already-present files are overwritten with byte-identical data). Holds the container lock for the whole + * fetch so a concurrent {@link #evictContainer} or container-init can't race the write; concurrent per-file + * {@link #ensureFileDownloaded} calls write byte-identical data so they remain safe, and the download-bookkeeping is + * gated on the atomic {@link #downloadedFiles} add so neither path double-counts. + */ + private void downloadContainer(int containerIndex) throws IOException + { + final List fileNames = containerFileNames.get(containerIndex); + if (fileNames.isEmpty() || downloadedFiles.containsAll(fileNames)) { + return; + } + containerLocks[containerIndex].lock(); + try { + checkClosed(); + if (downloadedFiles.containsAll(fileNames)) { + return; + } + ensureContainerInitialized(containerIndex); + + final SegmentFileContainerMetadata containerMeta = metadata.getContainers().get(containerIndex); + streamRangeIntoContainer( + containerIndex, + headerSize + containerMeta.getStartOffset(), + 0, + containerMeta.getSize(), + StringUtils.format("container[%d]", containerIndex) + ); + + for (String name : fileNames) { + markDownloaded(name, metadata.getFiles().get(name).getSize()); + } + } + finally { + containerLocks[containerIndex].unlock(); + } + } + /** * Pre-download a set of internal files so that subsequent {@link #mapFile(String)} calls for these files will not * trigger individual downloads. Files that are already downloaded are skipped. Useful for batch-downloading all @@ -436,6 +551,24 @@ public Set getDownloadedFiles() return Set.copyOf(downloadedFiles); } + /** + * Whether every internal file in this mapper's metadata (and recursively every attached external mapper's metadata) + * is downloaded. Used by the sync cursor-factory path to verify that the segment has been fully loaded before + * cursor construction; the async path is the only way to drive on-demand downloads. + */ + public boolean isFullyDownloaded() + { + if (!downloadedFiles.containsAll(metadata.getFiles().keySet())) { + return false; + } + for (PartialSegmentFileMapperV10 external : externalMappers.values()) { + if (!external.isFullyDownloaded()) { + return false; + } + } + return true; + } + @Override public void close() { @@ -485,9 +618,14 @@ private void ensureFileDownloaded(String name, SegmentInternalFileMetadata fileM } ensureContainerInitialized(fileMetadata.getContainer()); - downloadFileToContainer(name, fileMetadata); - downloadedFiles.add(name); - markDownloadedInBitmap(name); + streamRangeIntoContainer( + fileMetadata.getContainer(), + computeAbsoluteOffset(fileMetadata), + fileMetadata.getStartOffset(), + fileMetadata.getSize(), + StringUtils.format("file[%s]", name) + ); + markDownloaded(name, fileMetadata.getSize()); } finally { lock.unlock(); @@ -635,28 +773,35 @@ private void ensureContainerInitialized(int containerIndex) throws IOException } /** - * Download an internal file from deep storage and write it to the correct position in its local container file. - * Uses a short-lived {@link RandomAccessFile} for writing. The mmap sees the written data through the shared page - * cache. + * Stream {@code size} bytes starting at {@code absoluteOffset} in the segment file (deep storage) into container + * {@code containerIndex}'s local file at {@code localOffset}, via a short-lived {@link RandomAccessFile} (traditional + * I/O rather than an NIO channel, so a thread interrupt can't close the channel mid-write; the container's read-only + * mmap sees the bytes through the shared page cache). {@code what} names the unit being fetched (e.g. + * {@code "file[foo]"} or {@code "container[2]"}) for the end-of-stream error message. */ - private void downloadFileToContainer(String name, SegmentInternalFileMetadata fileMetadata) throws IOException + private void streamRangeIntoContainer( + int containerIndex, + long absoluteOffset, + long localOffset, + long size, + String what + ) throws IOException { - final long absoluteOffset = computeAbsoluteOffset(fileMetadata); - final long size = fileMetadata.getSize(); - - // stream directly from deep storage to the local container file to avoid holding the entire file in heap - try (InputStream is = rangeReader.readRange(targetFilename, absoluteOffset, size); - RandomAccessFile raf = new RandomAccessFile(containerFiles[fileMetadata.getContainer()], "rw")) { - raf.seek(fileMetadata.getStartOffset()); - final byte[] buf = new byte[8192]; + // stream straight from deep storage to the local container file to avoid heap-buffering the whole range + try (ResourceHolder bufHolder = CompressedPools.getOutputBytes(); + InputStream is = rangeReader.readRange(targetFilename, absoluteOffset, size); + RandomAccessFile raf = new RandomAccessFile(containerFiles[containerIndex], "rw")) { + final byte[] buf = bufHolder.get(); + raf.seek(localOffset); long remaining = size; while (remaining > 0) { final int toRead = (int) Math.min(buf.length, remaining); final int read = is.read(buf, 0, toRead); if (read < 0) { throw DruidException.defensive( - "unexpected end of stream for file[%s], expected[%s] more bytes", - name, + "unexpected end of stream for %s of [%s], expected[%s] more bytes", + what, + targetFilename, remaining ); } @@ -664,8 +809,20 @@ private void downloadFileToContainer(String name, SegmentInternalFileMetadata fi remaining -= read; } } + } - downloadedBytes.addAndGet(size); + /** + * Record an internal file as downloaded once its bytes are on disk: gate on the atomic {@link #downloadedFiles} add + * so a concurrent download of the same file via the other path (whole-container {@link #downloadContainer} vs + * per-file {@link #ensureFileDownloaded}) doesn't double-count its size (both write byte-identical data), then add + * its size and set its bitmap bit. + */ + private void markDownloaded(String name, long size) + { + if (downloadedFiles.add(name)) { + downloadedBytes.addAndGet(size); + markDownloadedInBitmap(name); + } } /** diff --git a/processing/src/test/java/org/apache/druid/segment/file/DirectoryBackedRangeReader.java b/processing/src/main/java/org/apache/druid/segment/loading/DirectoryBackedRangeReader.java similarity index 80% rename from processing/src/test/java/org/apache/druid/segment/file/DirectoryBackedRangeReader.java rename to processing/src/main/java/org/apache/druid/segment/loading/DirectoryBackedRangeReader.java index bc1df52946a6..42e2f591982f 100644 --- a/processing/src/test/java/org/apache/druid/segment/file/DirectoryBackedRangeReader.java +++ b/processing/src/main/java/org/apache/druid/segment/loading/DirectoryBackedRangeReader.java @@ -17,9 +17,7 @@ * under the License. */ -package org.apache.druid.segment.file; - -import org.apache.druid.segment.loading.SegmentRangeReader; +package org.apache.druid.segment.loading; import java.io.ByteArrayInputStream; import java.io.File; @@ -28,9 +26,9 @@ import java.io.RandomAccessFile; /** - * A {@link SegmentRangeReader} backed by a directory of files. Used across the partial-segment test suite (processing - * + server modules) to simulate deep-storage range reads against an on-disk layout produced by - * {@link SegmentFileBuilderV10} or {@link org.apache.druid.segment.IndexMergerV10}. + * A {@link SegmentRangeReader} backed by a directory of files: each {@code filename} passed to + * {@link #readRange} resolves to a child of the directory and is read at the requested byte range via + * {@link RandomAccessFile}. */ public class DirectoryBackedRangeReader implements SegmentRangeReader { diff --git a/processing/src/main/java/org/apache/druid/segment/projections/Projections.java b/processing/src/main/java/org/apache/druid/segment/projections/Projections.java index c4af08d8878f..8f03f7d06853 100644 --- a/processing/src/main/java/org/apache/druid/segment/projections/Projections.java +++ b/processing/src/main/java/org/apache/druid/segment/projections/Projections.java @@ -115,6 +115,7 @@ public static QueryableProjection findMatchingProjection( cursorBuildSpec.getQueryMetrics().projection(spec.getSchema().getName()); } return new QueryableProjection<>( + spec.getSchema().getName(), match.getCursorBuildSpec(), match.getRemapColumns(), getRowSelector.apply(spec.getSchema().getName()) diff --git a/processing/src/main/java/org/apache/druid/segment/projections/QueryableProjection.java b/processing/src/main/java/org/apache/druid/segment/projections/QueryableProjection.java index b32ea829f34d..d6a9b797b380 100644 --- a/processing/src/main/java/org/apache/druid/segment/projections/QueryableProjection.java +++ b/processing/src/main/java/org/apache/druid/segment/projections/QueryableProjection.java @@ -53,21 +53,32 @@ */ public class QueryableProjection { + private final String name; private final CursorBuildSpec cursorBuildSpec; private final Map remapColumns; private final T rowSelector; public QueryableProjection( + String name, CursorBuildSpec cursorBuildSpec, Map remapColumns, T rowSelector ) { + this.name = name; this.cursorBuildSpec = cursorBuildSpec; this.remapColumns = remapColumns; this.rowSelector = rowSelector; } + /** + * Name of the projection schema that matched the query. + */ + public String getName() + { + return name; + } + /** * The original {@link CursorBuildSpec} of a query can be modified if a projection matches the query, such as removing * virtual columns which have already been pre-computed. diff --git a/processing/src/test/java/org/apache/druid/common/asyncresource/AsyncResourcesTest.java b/processing/src/test/java/org/apache/druid/common/asyncresource/AsyncResourcesTest.java new file mode 100644 index 000000000000..dc25f7479c32 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/common/asyncresource/AsyncResourcesTest.java @@ -0,0 +1,177 @@ +/* + * 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.common.asyncresource; + +import com.google.common.util.concurrent.SettableFuture; +import org.apache.druid.error.DruidException; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.io.Closeable; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +public class AsyncResourcesTest +{ + @Test + public void testFromFutureSuccess() + { + final SettableFuture future = SettableFuture.create(); + final AsyncResource resource = AsyncResources.fromFutureUnmanaged(future); + + Assertions.assertFalse(resource.isReady()); + future.set("value"); + Assertions.assertTrue(resource.isReady()); + Assertions.assertEquals("value", resource.get()); + + // closing after a non-resource value completes is a no-op (nothing to release), and does not cancel the + // already-completed future + resource.close(); + Assertions.assertFalse(future.isCancelled()); + } + + @Test + public void testFromFutureFailureSurfacesAtGet() + { + final SettableFuture future = SettableFuture.create(); + final AsyncResource resource = AsyncResources.fromFutureUnmanaged(future); + + final RuntimeException boom = new IllegalStateException("boom"); + future.setException(boom); + Assertions.assertTrue(resource.isReady()); + final RuntimeException thrown = Assertions.assertThrows(RuntimeException.class, resource::get); + Assertions.assertSame(boom, thrown); + } + + @Test + public void testFromFutureCloseBeforeCompleteCancelsFuture() + { + final SettableFuture future = SettableFuture.create(); + final AsyncResource resource = AsyncResources.fromFutureUnmanaged(future); + + Assertions.assertFalse(resource.isReady()); + resource.close(); + Assertions.assertTrue(future.isCancelled(), "closing before completion must cancel the backing future"); + } + + @Test + public void testFromFutureReadyCallbackFiresOnCompletion() + { + final SettableFuture future = SettableFuture.create(); + final AsyncResource resource = AsyncResources.fromFutureUnmanaged(future); + + final boolean[] fired = {false}; + resource.addReadyCallback(() -> fired[0] = true); + Assertions.assertFalse(fired[0]); + + future.set("value"); + Assertions.assertTrue(fired[0], "ready callback should fire when the future completes"); + } + + @Test + public void testCollectOfFromFutureCancelsAllOnClose() + { + final SettableFuture a = SettableFuture.create(); + final SettableFuture b = SettableFuture.create(); + final AsyncResource> collected = + AsyncResources.collect(List.of(AsyncResources.fromFutureUnmanaged(a), AsyncResources.fromFutureUnmanaged(b))); + + Assertions.assertFalse(collected.isReady()); + // closing the collected resource before it is ready cancels every backing future + collected.close(); + Assertions.assertTrue(a.isCancelled()); + Assertions.assertTrue(b.isCancelled()); + } + + @Test + public void testCollectOfFromFutureReadyWhenAllComplete() + { + final SettableFuture a = SettableFuture.create(); + final SettableFuture b = SettableFuture.create(); + final AsyncResource> collected = + AsyncResources.collect(List.of(AsyncResources.fromFutureUnmanaged(a), AsyncResources.fromFutureUnmanaged(b))); + + a.set("a"); + Assertions.assertFalse(collected.isReady(), "not ready until every source completes"); + b.set("b"); + Assertions.assertTrue(collected.isReady()); + Assertions.assertEquals(List.of("a", "b"), collected.get()); + collected.close(); + } + + @Test + public void testFromFutureGetBeforeReadyThrows() + { + final AsyncResource resource = AsyncResources.fromFutureUnmanaged(SettableFuture.create()); + Assertions.assertThrows(DruidException.class, resource::get); + } + + @Test + public void testFromCloseableFutureClosesResultExactlyOnceOnClose() + { + final SettableFuture future = SettableFuture.create(); + final AsyncResource resource = AsyncResources.fromFutureCloseable(future); + final CloseableProbe probe = new CloseableProbe(); + + future.set(probe); + Assertions.assertTrue(resource.isReady()); + Assertions.assertSame(probe, resource.get()); + Assertions.assertEquals(0, probe.closeCount.get(), "result must not be closed while the resource is open"); + + resource.close(); + Assertions.assertEquals(1, probe.closeCount.get(), "closing the resource closes the managed result"); + } + + @Test + public void testFromCloseableFutureFailureSurfacesAtGet() + { + final SettableFuture future = SettableFuture.create(); + final AsyncResource resource = AsyncResources.fromFutureCloseable(future); + + final RuntimeException boom = new IllegalStateException("boom"); + future.setException(boom); + Assertions.assertSame(boom, Assertions.assertThrows(RuntimeException.class, resource::get)); + } + + @Test + public void testFromCloseableFutureCloseBeforeCompleteCancelsFuture() + { + final SettableFuture future = SettableFuture.create(); + final AsyncResource resource = AsyncResources.fromFutureCloseable(future); + + Assertions.assertFalse(resource.isReady()); + resource.close(); + Assertions.assertTrue(future.isCancelled(), "closing before completion cancels the backing future"); + } + + /** + * Closeable that records how many times it was closed. + */ + private static class CloseableProbe implements Closeable + { + private final AtomicInteger closeCount = new AtomicInteger(0); + + @Override + public void close() + { + closeCount.incrementAndGet(); + } + } +} diff --git a/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryTest.java b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryTest.java new file mode 100644 index 000000000000..74c1e4d4aa17 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryTest.java @@ -0,0 +1,757 @@ +/* + * 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; + +import com.google.common.collect.Sets; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.ListeningExecutorService; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.druid.common.asyncresource.AsyncResource; +import org.apache.druid.common.asyncresource.AsyncResources; +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.error.DruidException; +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.common.concurrent.Execs; +import org.apache.druid.java.util.common.granularity.Granularities; +import org.apache.druid.query.Order; +import org.apache.druid.query.aggregation.CountAggregatorFactory; +import org.apache.druid.query.aggregation.LongSumAggregatorFactory; +import org.apache.druid.query.filter.EqualityFilter; +import org.apache.druid.segment.column.ColumnConfig; +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.data.CompressionStrategy; +import org.apache.druid.segment.file.CountingRangeReader; +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.testing.InitializedNullHandlingTest; +import org.joda.time.DateTime; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.Closeable; +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.concurrent.Callable; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.ReentrantLock; + +class PartialQueryableIndexCursorFactoryTest extends InitializedNullHandlingTest +{ + private static final ColumnConfig COLUMN_CONFIG = ColumnConfig.DEFAULT; + private static final DateTime TIME = DateTimes.of("2025-01-01"); + private static final String PROJECTION_NAME = "dim1_metric1_sum"; + + private static final RowSignature ROW_SIGNATURE = RowSignature.builder() + .add("dim1", ColumnType.STRING) + .add("metric1", ColumnType.LONG) + .build(); + + private static final List PROJECTIONS = Collections.singletonList( + AggregateProjectionSpec.builder(PROJECTION_NAME) + .groupingColumns(new StringDimensionSchema("dim1")) + .aggregators( + new LongSumAggregatorFactory("_metric1_sum", "metric1"), + 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 sharedTempDir; + + private static File segmentDir; + private static ListeningExecutorService realExec; + + @TempDir + File perTestTempDir; + + @BeforeAll + static void buildSegment() + { + final File tmpDir = new File(sharedTempDir, "build_" + ThreadLocalRandom.current().nextInt()); + segmentDir = IndexBuilder.create() + .useV10() + .tmpDir(tmpDir) + .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(); + + realExec = MoreExecutors.listeningDecorator(Execs.singleThreaded("partial-cursor-test-%d")); + } + + @AfterAll + static void teardownExec() + { + realExec.shutdownNow(); + } + + + @Test + void testSyncMakeCursorHolderThrowsWhenSegmentNotFullyDownloaded() throws IOException + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + try (IndexAndMapper opened = openIndex(rangeReader, "sync_throws")) { + final PartialQueryableIndex index = opened.index(); + Assertions.assertFalse(index.isFullyDownloaded()); + final PartialQueryableIndexCursorFactory factory = new PartialQueryableIndexCursorFactory( + index, + QueryableIndexTimeBoundaryInspector.create(index), + noOpAcquirer(directExec()) + ); + + // Fresh mapper has only the header on disk; no internal files have been downloaded. The sync path must + // refuse rather than trigger downloads on the calling thread. + Assertions.assertThrows( + DruidException.class, + () -> factory.makeCursorHolder(CursorBuildSpec.FULL_SCAN) + ); + } + } + + @Test + void testSyncMakeCursorHolderSucceedsAfterFullDownload() throws Exception + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + try (IndexAndMapper opened = openIndex(rangeReader, "sync_after_async")) { + final PartialQueryableIndex index = opened.index(); + final PartialSegmentFileMapperV10 mapper = opened.mapper(); + final PartialQueryableIndexCursorFactory factory = new PartialQueryableIndexCursorFactory( + index, + QueryableIndexTimeBoundaryInspector.create(index), + noOpAcquirer(directExec()) + ); + + // Eagerly materialize the entire segment via the file mapper to simulate the eager-acquire path. + mapper.ensureFilesAvailable(mapper.getSegmentFileMetadata().getFiles().keySet()); + Assertions.assertTrue(mapper.isFullyDownloaded(), "test precondition: every internal file should be on disk"); + + try (CursorHolder holder = factory.makeCursorHolder(CursorBuildSpec.FULL_SCAN)) { + Assertions.assertNotNull(holder); + } + } + } + + @Test + void testAsyncMakeCursorHolderDefersDownloadUntilExecutorRuns() throws Exception + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + final CountDownLatch gate = new CountDownLatch(1); + final ExecutorService rawExec = Execs.singleThreaded("partial-cursor-defer-%d"); + final ListeningExecutorService gatedExec = MoreExecutors.listeningDecorator(rawExec); + try (IndexAndMapper opened = openIndex(rangeReader, "async_defer")) { + final PartialQueryableIndex index = opened.index(); + final PartialQueryableIndexCursorFactory factory = new PartialQueryableIndexCursorFactory( + index, + QueryableIndexTimeBoundaryInspector.create(index), + noOpAcquirer(gatedExec) + ); + rangeReader.resetCount(); + + // Hold the executor with a pre-queued blocker so the actual download task can't start. + // assign the blocker future to an (intentionally unused) local so errorprone's CheckReturnValue is satisfied + @SuppressWarnings("unused") + ListenableFuture unused = gatedExec.submit(() -> { + try { + gate.await(); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + + final AsyncCursorHolder asyncHolder = factory.makeCursorHolderAsync(CursorBuildSpec.FULL_SCAN); + Assertions.assertFalse( + asyncHolder.isReady(), + "async holder must not be ready while download executor is still blocked" + ); + Assertions.assertEquals(0, rangeReader.getReadCount(), "no download should have started yet"); + + final CountDownLatch ready = new CountDownLatch(1); + asyncHolder.addReadyCallback(ready::countDown); + + // Release the gate; download should now run. + gate.countDown(); + Assertions.assertTrue(ready.await(5, TimeUnit.SECONDS), "ready callback must fire after executor processes task"); + Assertions.assertTrue(asyncHolder.isReady()); + Assertions.assertTrue(rangeReader.getReadCount() > 0, "download must have happened by the time holder is ready"); + + try (CursorHolder holder = asyncHolder.release()) { + Assertions.assertNotNull(holder); + } + } + finally { + gatedExec.shutdownNow(); + rawExec.shutdownNow(); + } + } + + @Test + void testMatchedProjectionDownloadsOnlyRequestedColumns() throws IOException + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + try (IndexAndMapper opened = openIndex(rangeReader, "proj_columns")) { + final PartialQueryableIndex index = opened.index(); + final PartialSegmentFileMapperV10 mapper = opened.mapper(); + final PartialQueryableIndexCursorFactory factory = new PartialQueryableIndexCursorFactory( + index, + QueryableIndexTimeBoundaryInspector.create(index), + noOpAcquirer(directExec()) + ); + + // Query asks for group-by dim1 + sum(metric1), but NOT count. The projection has dim1 / _metric1_sum / _count; + // the projection match rewrites this to physical columns {dim1, _metric1_sum}. We should see those download, + // but NOT the _count file; proves column-level pruning, not just projection-level. + final CursorBuildSpec aggSpec = CursorBuildSpec.builder() + .setGroupingColumns(List.of("dim1")) + .setAggregators(List.of( + new LongSumAggregatorFactory("_metric1_sum", "metric1") + )) + .setPhysicalColumns(Set.of("dim1", "metric1")) + .build(); + + // Async path triggers the column-level pre-download. directExec makes the future complete synchronously. + try (AsyncCursorHolder asyncHolder = factory.makeCursorHolderAsync(aggSpec); + CursorHolder holder = asyncHolder.release()) { + Assertions.assertNotNull(holder); + final String projPrefix = PROJECTION_NAME + "/"; + final String basePrefix = Projections.BASE_TABLE_PROJECTION_NAME + "/"; + final Set downloaded = mapper.getDownloadedFiles(); + + // Requested projection columns materialized. + Assertions.assertTrue(downloaded.contains(projPrefix + "dim1"), "expected projection dim1; got: " + downloaded); + Assertions.assertTrue( + downloaded.contains(projPrefix + "_metric1_sum"), + "expected projection _metric1_sum; got: " + downloaded + ); + // The projection's _count file was NOT in the query; must not be downloaded + Assertions.assertFalse( + downloaded.contains(projPrefix + "_count"), + "expected projection _count NOT to be downloaded; got: " + downloaded + ); + // Base __time and base metric1 are NOT touched: projection dim1 may pull base dim1 as its parent column + // (legitimate dependency), but unrelated base columns must stay untouched. + Assertions.assertFalse( + downloaded.contains(basePrefix + ColumnHolder.TIME_COLUMN_NAME), + "expected base __time NOT to be downloaded; got: " + downloaded + ); + Assertions.assertFalse( + downloaded.contains(basePrefix + "metric1"), + "expected base metric1 NOT to be downloaded; got: " + downloaded + ); + } + } + } + + @Test + void testOpeningTimeOrderedProjectionCursorTriggersNoDownload() throws IOException + { + // A projection grouped by [__gran(HOUR), dim1] is time-ordered, so QueryableIndexCursorHolder reads the + // projection's (downloadable) time column when it builds the cursor (interval-checking offset + the projection's + // own time-boundary inspector). A GROUP BY dim1 query matches the projection by re-aggregating over __gran and + // does NOT list __time in physicalColumns. requiredColumns must nonetheless pre-fetch __time so that read happens + // during the async pre-fetch rather than as a lazy download when the cursor is opened on a processing thread. + final String projectionName = "hourly_dim1_metric1_sum"; + final File timeOrderedSegmentDir = buildTimeOrderedProjectionSegment(projectionName); + + final CountingRangeReader rangeReader = new CountingRangeReader(timeOrderedSegmentDir); + final File cacheDir = new File(perTestTempDir, "time_ordered_proj"); + FileUtils.mkdirp(cacheDir); + final PartialSegmentFileMapperV10 mapper = PartialSegmentFileMapperV10.create( + rangeReader, + TestHelper.makeJsonMapper(), + cacheDir, + IndexIO.V10_FILE_NAME, + Collections.emptyList() + ); + try { + final PartialQueryableIndex index = + new PartialQueryableIndex(mapper.getSegmentFileMetadata(), mapper, COLUMN_CONFIG); + + // Precondition: the matched projection's row selector really is time-ordered, so cursor build WOULD read its + // time column. Without this, the assertion below would be vacuously satisfied. + final QueryableIndex projectionIndex = index.getProjectionQueryableIndex(projectionName); + Assertions.assertNotNull(projectionIndex, "test setup: projection should be present"); + Assertions.assertEquals( + Order.ASCENDING, + Cursors.getTimeOrdering(projectionIndex.getOrdering()), + "test setup: projection must be time-ordered for this test to guard anything" + ); + + final PartialQueryableIndexCursorFactory factory = new PartialQueryableIndexCursorFactory( + index, + V10TimeBoundaryInspector.forBaseProjection(index.getBaseProjectionMetadata(), index.getDataInterval()), + noOpAcquirer(directExec()) + ); + + // ETERNITY interval so the offset is not clipped to a residual interval (a genuine reason to read __time); + // here the only __time read is the time-ordered offset/inspector during cursor build, which the pre-fetch + // must have already covered. physicalColumns deliberately omits __time. + final CursorBuildSpec aggSpec = CursorBuildSpec.builder() + .setInterval(Intervals.ETERNITY) + .setGroupingColumns(List.of("dim1")) + .setAggregators(List.of( + new LongSumAggregatorFactory("_metric1_sum", "metric1") + )) + .setPhysicalColumns(Set.of("dim1", "metric1")) + .build(); + + try (AsyncCursorHolder asyncHolder = factory.makeCursorHolderAsync(aggSpec); + CursorHolder holder = asyncHolder.release()) { + Assertions.assertNotNull(holder); + + // After the async pre-fetch resolves, the projection's required columns AND __time are on disk, even though + // __time was not declared in physicalColumns. + final String projPrefix = projectionName + "/"; + final Set beforeOpen = mapper.getDownloadedFiles(); + Assertions.assertTrue( + beforeOpen.contains(projPrefix + "_metric1_sum"), + "test setup: projection path should have run and pre-fetched _metric1_sum; got: " + beforeOpen + ); + // The projection's time column is stored under its granularity grouping name (__gran); it is exposed as + // __time but the on-disk smoosh file (and thus the downloaded-files entry) is __gran. + Assertions.assertTrue( + beforeOpen.contains(projPrefix + "__gran"), + "requiredColumns must pre-fetch the projection's time column for a time-ordered cursor; got: " + beforeOpen + ); + + // Opening the cursor reads __time (interval-checking offset + the projection's time-boundary inspector). It + // was pre-fetched above, so this must not trigger any further (lazy) download on the processing thread. + Assertions.assertNotNull(holder.asCursor()); + + final Set afterOpen = mapper.getDownloadedFiles(); + Assertions.assertEquals( + beforeOpen, + afterOpen, + "opening a projection-matched cursor must not trigger a lazy download; newly downloaded: " + + Sets.difference(afterOpen, beforeOpen) + ); + } + } + finally { + mapper.close(); + } + } + + @Test + void testOpeningTimeOrderedBaseCursorTriggersNoDownload() throws IOException + { + // The base table is time-ordered, so QueryableIndexCursorHolder reads __time when building the cursor (for its + // interval-checking offset) even though a raw scan does not list __time in physicalColumns. requiredColumns must + // pre-fetch __time so that read does not become a lazy download on the processing thread at cursor-open time. + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + try (IndexAndMapper opened = openIndex(rangeReader, "base_open_no_download")) { + final PartialQueryableIndex index = opened.index(); + final PartialSegmentFileMapperV10 mapper = opened.mapper(); + Assertions.assertEquals( + Order.ASCENDING, + Cursors.getTimeOrdering(index.getOrdering()), + "test setup: base table must be time-ordered" + ); + + final PartialQueryableIndexCursorFactory factory = new PartialQueryableIndexCursorFactory( + index, + V10TimeBoundaryInspector.forBaseProjection(index.getBaseProjectionMetadata(), index.getDataInterval()), + noOpAcquirer(directExec()) + ); + + // Raw scan (no grouping/aggregators) so the aggregate projection is not matched and we exercise the base path. + // __time is deliberately omitted from physicalColumns. + final CursorBuildSpec scanSpec = CursorBuildSpec.builder() + .setInterval(Intervals.ETERNITY) + .setPhysicalColumns(Set.of("dim1")) + .build(); + + try (AsyncCursorHolder asyncHolder = factory.makeCursorHolderAsync(scanSpec); + CursorHolder holder = asyncHolder.release()) { + Assertions.assertNotNull(holder); + + final String basePrefix = Projections.BASE_TABLE_PROJECTION_NAME + "/"; + final Set beforeOpen = mapper.getDownloadedFiles(); + // base path ran (dim1 pre-fetched) and __time was pre-fetched despite not being in physicalColumns + Assertions.assertTrue( + beforeOpen.contains(basePrefix + "dim1"), + "test setup: base path should have pre-fetched dim1; got: " + beforeOpen + ); + Assertions.assertTrue( + beforeOpen.contains(basePrefix + ColumnHolder.TIME_COLUMN_NAME), + "requiredColumns must pre-fetch __time for a time-ordered base cursor; got: " + beforeOpen + ); + + Assertions.assertNotNull(holder.asCursor()); + + final Set afterOpen = mapper.getDownloadedFiles(); + Assertions.assertEquals( + beforeOpen, + afterOpen, + "opening a time-ordered base cursor must not trigger a lazy download; newly downloaded: " + + Sets.difference(afterOpen, beforeOpen) + ); + } + } + } + + @Test + void testFilteredWrapperOverPartialUsesAsyncPath() throws IOException + { + // A FilteredCursorFactory wrapping a not-fully-downloaded partial cursor factory must route makeCursorHolderAsync + // through the delegate's async (download-on-demand) path. If FilteredCursorFactory fell back to the CursorFactory + // default makeCursorHolderAsync (which calls sync makeCursorHolder), the partial delegate would throw because the + // segment isn't fully downloaded. + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + try (IndexAndMapper opened = openIndex(rangeReader, "filtered_wrapper_async")) { + final PartialQueryableIndex index = opened.index(); + Assertions.assertFalse(index.isFullyDownloaded(), "test setup: segment must not be fully downloaded"); + + final PartialQueryableIndexCursorFactory partial = new PartialQueryableIndexCursorFactory( + index, + V10TimeBoundaryInspector.forBaseProjection(index.getBaseProjectionMetadata(), index.getDataInterval()), + noOpAcquirer(directExec()) + ); + final CursorFactory filtered = new FilteredCursorFactory( + partial, + new EqualityFilter("dim1", ColumnType.STRING, "a", null) + ); + + // directExec resolves the async download inline; release() returns the built holder. asCursor() exercises the + // cursor over the on-demand-loaded columns. Would throw here if the wrapper used the sync default. + try (AsyncCursorHolder asyncHolder = filtered.makeCursorHolderAsync(CursorBuildSpec.FULL_SCAN); + CursorHolder holder = asyncHolder.release()) { + Assertions.assertNotNull(holder); + Assertions.assertNotNull(holder.asCursor(), "async path should build a usable cursor without a sync download"); + } + } + } + + @Test + void testRowSignatureAndCapabilitiesDelegatedToUnderlyingFactory() throws IOException + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + try (IndexAndMapper opened = openIndex(rangeReader, "delegate")) { + final PartialQueryableIndex index = opened.index(); + final PartialQueryableIndexCursorFactory factory = new PartialQueryableIndexCursorFactory( + index, + QueryableIndexTimeBoundaryInspector.create(index), + noOpAcquirer(directExec()) + ); + + // delegating to QueryableIndexCursorFactory which returns the index's row signature + final RowSignature signature = factory.getRowSignature(); + Assertions.assertNotNull(signature); + Assertions.assertTrue(signature.size() > 0); + + Assertions.assertNotNull(factory.getColumnCapabilities("dim1")); + } + } + + @Test + void testCloseBeforeReadyCancelsColumnDownloadsAndReleasesHoldOnce() throws IOException + { + // When the awaiter closes the holder before it's ready (query cancel/timeout), the canceler must cancel every + // submitted column-download future and release the bundle hold exactly once (the canceler and the resulting + // failure callback both try to release it). A recording executor captures each submitted task as a future that + // never completes and never runs the task, so the holder stays not-ready and we can observe the cancellation. + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + try (IndexAndMapper opened = openIndex(rangeReader, "cancel_before_ready")) { + final PartialQueryableIndex index = opened.index(); + + // submitDownload records each task's never-completing future and returns it as an AsyncResource, so the holder + // stays not-ready and we can observe the cancellation. A bundle hold whose releases we count asserts + // exactly-once release across the canceler + failure callback. + final List> submitted = new ArrayList<>(); + final AtomicInteger holdReleases = new AtomicInteger(0); + final PartialBundleAcquirer acquirer = new PartialBundleAcquirer() + { + @Override + public Closeable acquire(String bundleName) + { + return holdReleases::incrementAndGet; + } + + @Override + public AsyncResource submitDownload(Callable task) + { + final SettableFuture future = SettableFuture.create(); + submitted.add(future); + return AsyncResources.fromFutureUnmanaged(future); + } + }; + + final PartialQueryableIndexCursorFactory factory = new PartialQueryableIndexCursorFactory( + index, + V10TimeBoundaryInspector.forBaseProjection(index.getBaseProjectionMetadata(), index.getDataInterval()), + acquirer + ); + + final AsyncCursorHolder asyncHolder = factory.makeCursorHolderAsync(CursorBuildSpec.FULL_SCAN); + + // Downloads were submitted but never complete, so the holder is not ready and the hold is still held. + Assertions.assertFalse(asyncHolder.isReady(), "downloads never complete, so the holder must not be ready"); + Assertions.assertFalse(submitted.isEmpty(), "at least one column download should have been submitted"); + submitted.forEach(f -> Assertions.assertFalse(f.isCancelled(), "futures must not be cancelled before close")); + Assertions.assertEquals(0, holdReleases.get(), "hold must still be held while the cursor is loading"); + + // Awaiter cancels (closes) before the holder is ready: every pending download is cancelled and the hold drops. + asyncHolder.close(); + + submitted.forEach(f -> Assertions.assertTrue(f.isCancelled(), "pending column downloads must be cancelled")); + Assertions.assertEquals( + 1, + holdReleases.get(), + "bundle hold must be released exactly once across the canceler and the failure callback" + ); + } + } + + @Test + void testCancelWhileDownloadMidMapFileDefersHoldReleaseUntilBodyFinishes() throws Exception + { + // Eviction-during-download safety: releasing the bundle hold makes the bundle entry evictable, and eviction unmaps + // its containers. A column download still inside mapFile() would then read an unmapped region (a JVM SIGSEGV). So + // when the awaiter cancels while a download is mid-mapFile(), the canceler must NOT release the hold; the last + // in-flight body releases it only after its read has finished. + final ReentrantLock gate = new ReentrantLock(); + final CountDownLatch readStarted = new CountDownLatch(1); + final AtomicBoolean blockedOnce = new AtomicBoolean(false); + // Park the first read issued from a download-executor thread, modeling a body parked inside mapFile() while holding + // the in-flight count. The mapper's create() reads (header + metadata) run on this thread during openIndex and + // must pass through untouched, so discriminate by thread identity rather than offset. gate.lock() is + // uninterruptible, modeling the mapper's non-interruptible traditional I/O: future.cancel(true) cannot abort it. + final Thread mainThread = Thread.currentThread(); + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir) + { + @Override + public InputStream readRange(String filename, long offset, long length) throws IOException + { + if (Thread.currentThread() != mainThread && blockedOnce.compareAndSet(false, true)) { + readStarted.countDown(); + gate.lock(); + gate.unlock(); + } + return super.readRange(filename, offset, length); + } + }; + + final ExecutorService rawExec = Execs.singleThreaded("partial-cursor-evict-race-%d"); + final ListeningExecutorService blockingExec = MoreExecutors.listeningDecorator(rawExec); + final AtomicInteger holdReleases = new AtomicInteger(0); + final CountDownLatch released = new CountDownLatch(1); + + gate.lock(); // hold the gate so the first column download parks inside its read + try (IndexAndMapper opened = openIndex(rangeReader, "evict_during_download")) { + final PartialQueryableIndex index = opened.index(); + final PartialBundleAcquirer acquirer = new PartialBundleAcquirer() + { + @Override + public Closeable acquire(String bundleName) + { + return () -> { + holdReleases.incrementAndGet(); + released.countDown(); + }; + } + + @Override + public AsyncResource submitDownload(Callable task) + { + return AsyncResources.fromFutureUnmanaged(blockingExec.submit(task)); + } + }; + + final PartialQueryableIndexCursorFactory factory = new PartialQueryableIndexCursorFactory( + index, + V10TimeBoundaryInspector.forBaseProjection(index.getBaseProjectionMetadata(), index.getDataInterval()), + acquirer + ); + + final AsyncCursorHolder asyncHolder = factory.makeCursorHolderAsync(CursorBuildSpec.FULL_SCAN); + + // A download body is now parked mid-mapFile(), holding the in-flight count; the hold is still held. + Assertions.assertTrue(readStarted.await(5, TimeUnit.SECONDS), "a column download must reach mapFile()"); + Assertions.assertFalse(asyncHolder.isReady(), "download is parked, so the holder must not be ready"); + Assertions.assertEquals(0, holdReleases.get(), "hold must still be held while a download is mid-mapFile"); + + // Awaiter cancels while the body is parked inside mapFile(). + asyncHolder.close(); + + // The hold must NOT be released yet: the in-flight body is still touching the bundle's container, so releasing + // now would let eviction unmap it mid-read. + Assertions.assertEquals( + 0, + holdReleases.get(), + "bundle hold must stay held until the in-flight mapFile() body finishes" + ); + + // Let the parked body finish its read; the last body out then releases the hold exactly once. + gate.unlock(); + Assertions.assertTrue(released.await(5, TimeUnit.SECONDS), "hold must be released once the in-flight body finishes"); + Assertions.assertEquals(1, holdReleases.get(), "bundle hold must be released exactly once"); + } + finally { + if (gate.isHeldByCurrentThread()) { + gate.unlock(); + } + blockingExec.shutdownNow(); + rawExec.shutdownNow(); + } + } + + private File buildTimeOrderedProjectionSegment(String projectionName) + { + final List projections = Collections.singletonList( + AggregateProjectionSpec.builder(projectionName) + .virtualColumns(Granularities.toVirtualColumn(Granularities.HOUR, "__gran")) + .groupingColumns( + new LongDimensionSchema("__gran"), + new StringDimensionSchema("dim1") + ) + .aggregators( + new LongSumAggregatorFactory("_metric1_sum", "metric1"), + new CountAggregatorFactory("_count") + ) + .build() + ); + final File tmpDir = new File(perTestTempDir, "build_time_ordered_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)); + return IndexBuilder.create() + .useV10() + .tmpDir(tmpDir) + .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(); + } + + private IndexAndMapper openIndex(CountingRangeReader rangeReader, String cacheName) throws IOException + { + final File cacheDir = new File(perTestTempDir, cacheName); + FileUtils.mkdirp(cacheDir); + final PartialSegmentFileMapperV10 mapper = PartialSegmentFileMapperV10.create( + rangeReader, + TestHelper.makeJsonMapper(), + cacheDir, + IndexIO.V10_FILE_NAME, + Collections.emptyList() + ); + return new IndexAndMapper( + new PartialQueryableIndex(mapper.getSegmentFileMetadata(), mapper, COLUMN_CONFIG), + mapper + ); + } + + private static ListeningExecutorService directExec() + { + return MoreExecutors.listeningDecorator(MoreExecutors.newDirectExecutorService()); + } + + private static PartialBundleAcquirer noOpAcquirer(ListeningExecutorService downloadExec) + { + return new PartialBundleAcquirer() + { + @Override + public Closeable acquire(String bundleName) + { + return () -> {}; + } + + @Override + public AsyncResource submitDownload(Callable task) + { + return AsyncResources.fromFutureUnmanaged(downloadExec.submit(task)); + } + }; + } + + private record IndexAndMapper(PartialQueryableIndex index, PartialSegmentFileMapperV10 mapper) + implements AutoCloseable + { + @Override + public void close() + { + mapper.close(); + } + } +} diff --git a/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexSegmentTest.java b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexSegmentTest.java new file mode 100644 index 000000000000..72156b740a1a --- /dev/null +++ b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexSegmentTest.java @@ -0,0 +1,272 @@ +/* + * 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; + +import com.google.common.util.concurrent.ListeningExecutorService; +import com.google.common.util.concurrent.MoreExecutors; +import org.apache.druid.common.asyncresource.AsyncResource; +import org.apache.druid.common.asyncresource.AsyncResources; +import org.apache.druid.data.input.InputRow; +import org.apache.druid.data.input.ListBasedInputRow; +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.java.util.common.DateTimes; +import org.apache.druid.java.util.common.FileUtils; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.segment.column.ColumnConfig; +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.data.CompressionStrategy; +import org.apache.druid.segment.file.CountingRangeReader; +import org.apache.druid.segment.file.PartialSegmentFileMapperV10; +import org.apache.druid.segment.incremental.IncrementalIndexSchema; +import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory; +import org.apache.druid.testing.InitializedNullHandlingTest; +import org.apache.druid.timeline.SegmentId; +import org.joda.time.DateTime; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.Closeable; +import java.io.File; +import java.io.IOException; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.ThreadLocalRandom; + +class PartialQueryableIndexSegmentTest extends InitializedNullHandlingTest +{ + private static final ColumnConfig COLUMN_CONFIG = ColumnConfig.DEFAULT; + private static final DateTime TIME = DateTimes.of("2025-01-01"); + + // expected exact bounds: first row at TIME, last row at TIME + 3 minutes + private static final DateTime EXPECTED_MIN_TIME = TIME; + private static final DateTime EXPECTED_MAX_TIME = TIME.plusMinutes(3); + + private static final SegmentId SEGMENT_ID = SegmentId.of("test", Intervals.of("2025/2026"), "v1", 0); + + // The async cursor factory's executor doesn't matter for the segment-shape tests; use direct-execution so async + // calls complete synchronously and tests stay simple. Cursor factory-specific async behavior is covered separately. + private static final ListeningExecutorService DIRECT_EXEC = + MoreExecutors.listeningDecorator(MoreExecutors.newDirectExecutorService()); + + private static final RowSignature ROW_SIGNATURE = RowSignature.builder() + .add("dim1", ColumnType.STRING) + .add("metric1", ColumnType.LONG) + .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 sharedTempDir; + + private static File segmentDir; + + @TempDir + File perTestTempDir; + + @BeforeAll + static void buildSegment() + { + final File tmpDir = new File(sharedTempDir, "build_" + ThreadLocalRandom.current().nextInt()); + segmentDir = IndexBuilder.create() + .useV10() + .tmpDir(tmpDir) + .segmentWriteOutMediumFactory(OffHeapMemorySegmentWriteOutMediumFactory.instance()) + .schema( + IncrementalIndexSchema.builder() + .withDimensionsSpec( + DimensionsSpec.builder() + .setDimensions( + List.of( + new StringDimensionSchema("dim1"), + new LongDimensionSchema("metric1") + ) + ) + .build() + ) + .withRollup(false) + .withMinTimestamp(TIME.getMillis()) + .build() + ) + .indexSpec(IndexSpec.builder().withMetadataCompression(CompressionStrategy.NONE).build()) + .rows(ROWS) + .buildMMappedIndexFile(); + } + + private PartialQueryableIndex openIndex(CountingRangeReader rangeReader, String cacheName) throws IOException + { + final File cacheDir = new File(perTestTempDir, cacheName); + FileUtils.mkdirp(cacheDir); + final PartialSegmentFileMapperV10 mapper = PartialSegmentFileMapperV10.create( + rangeReader, + TestHelper.makeJsonMapper(), + cacheDir, + IndexIO.V10_FILE_NAME, + Collections.emptyList() + ); + return new PartialQueryableIndex(mapper.getSegmentFileMetadata(), mapper, COLUMN_CONFIG); + } + + private static PartialQueryableIndexSegment makeSegment(PartialQueryableIndex index) + { + return new PartialQueryableIndexSegment(index, SEGMENT_ID, index::close, noOpAcquirer(DIRECT_EXEC)); + } + + /** + * No-op bundle acquirer for tests: every {@code acquire(...)} returns a noop {@link Closeable}; + * {@code submitDownload(...)} runs the task on the supplied executor and wraps the future as an + * {@link AsyncResource}. Production callers go through {@code PartialSegmentMetadataCacheEntry}'s bundle acquirer. + */ + private static PartialBundleAcquirer noOpAcquirer(ListeningExecutorService downloadExec) + { + return new PartialBundleAcquirer() + { + @Override + public Closeable acquire(String bundleName) + { + return () -> {}; + } + + @Override + public AsyncResource submitDownload(Callable task) + { + return AsyncResources.fromFutureUnmanaged(downloadExec.submit(task)); + } + }; + } + + @Test + void testAsTimeBoundaryInspectorReturnsV10Inspector() throws IOException + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + final PartialQueryableIndex index = openIndex(rangeReader, "boundary"); + + try (Segment segment = makeSegment(index)) { + final TimeBoundaryInspector inspector = segment.as(TimeBoundaryInspector.class); + Assertions.assertNotNull(inspector); + Assertions.assertInstanceOf(V10TimeBoundaryInspector.class, inspector); + } + } + + @Test + void testTimeBoundaryAnsweredFromMetadataWithoutColumnDownloads() throws IOException + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + final PartialQueryableIndex index = openIndex(rangeReader, "no_download"); + // header fetch happens during mapper creation; reset so we observe only what time-boundary lookups trigger + rangeReader.resetCount(); + + try (Segment segment = makeSegment(index)) { + final TimeBoundaryInspector inspector = segment.as(TimeBoundaryInspector.class); + Assertions.assertTrue(inspector.isMinMaxExact()); + Assertions.assertEquals(EXPECTED_MIN_TIME, inspector.getMinTime()); + Assertions.assertEquals(EXPECTED_MAX_TIME, inspector.getMaxTime()); + Assertions.assertEquals( + 0, + rangeReader.getReadCount(), + "metadata-only inspector must not trigger any column downloads" + ); + Assertions.assertFalse( + rangeReader.getReadFilenames().contains(ColumnHolder.TIME_COLUMN_NAME), + "must not download the __time column" + ); + } + } + + @Test + void testAsRowCountInspectorReturnsV10Inspector() throws IOException + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + final PartialQueryableIndex index = openIndex(rangeReader, "row_count"); + + try (Segment segment = makeSegment(index)) { + final RowCountInspector inspector = segment.as(RowCountInspector.class); + Assertions.assertNotNull(inspector); + Assertions.assertInstanceOf(V10RowCountInspector.class, inspector); + } + } + + @Test + void testRowCountAnsweredFromMetadataWithoutColumnDownloads() throws IOException + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + final PartialQueryableIndex index = openIndex(rangeReader, "row_count_no_download"); + // header fetch happens during mapper creation; reset so we observe only what the row-count lookup triggers + rangeReader.resetCount(); + + try (Segment segment = makeSegment(index)) { + final RowCountInspector inspector = segment.as(RowCountInspector.class); + Assertions.assertEquals(ROWS.size(), inspector.getNumRows()); + Assertions.assertEquals( + 0, + rangeReader.getReadCount(), + "metadata-only row count must not trigger any column downloads" + ); + } + } + + @Test + void testAsQueryableIndexReturnsUnderlyingIndex() throws IOException + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + final PartialQueryableIndex index = openIndex(rangeReader, "as_index"); + + try (Segment segment = makeSegment(index)) { + Assertions.assertSame(index, segment.as(QueryableIndex.class)); + } + } + + @Test + void testAsCursorFactoryReturnsNonNull() throws IOException + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + final PartialQueryableIndex index = openIndex(rangeReader, "cursor"); + + try (Segment segment = makeSegment(index)) { + final CursorFactory factory = segment.as(CursorFactory.class); + Assertions.assertNotNull(factory); + Assertions.assertInstanceOf(PartialQueryableIndexCursorFactory.class, factory); + } + } + + @Test + void testIdAndDataIntervalPassThrough() throws IOException + { + final CountingRangeReader rangeReader = new CountingRangeReader(segmentDir); + final PartialQueryableIndex index = openIndex(rangeReader, "id"); + + try (Segment segment = makeSegment(index)) { + Assertions.assertEquals(SEGMENT_ID, segment.getId()); + Assertions.assertEquals(index.getDataInterval(), segment.getDataInterval()); + } + } +} diff --git a/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexTest.java b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexTest.java index 6e82fb4a76a8..f40b0c5523b2 100644 --- a/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexTest.java +++ b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexTest.java @@ -38,9 +38,9 @@ import org.apache.druid.segment.column.ValueType; import org.apache.druid.segment.data.CompressionStrategy; import org.apache.druid.segment.file.CountingRangeReader; -import org.apache.druid.segment.file.DirectoryBackedRangeReader; import org.apache.druid.segment.file.PartialSegmentFileMapperV10; import org.apache.druid.segment.incremental.IncrementalIndexSchema; +import org.apache.druid.segment.loading.DirectoryBackedRangeReader; import org.apache.druid.segment.projections.QueryableProjection; import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory; import org.apache.druid.testing.InitializedNullHandlingTest; diff --git a/processing/src/test/java/org/apache/druid/segment/V10TimeBoundaryInspectorTest.java b/processing/src/test/java/org/apache/druid/segment/V10TimeBoundaryInspectorTest.java new file mode 100644 index 000000000000..321d76422b6c --- /dev/null +++ b/processing/src/test/java/org/apache/druid/segment/V10TimeBoundaryInspectorTest.java @@ -0,0 +1,109 @@ +/* + * 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; + +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.common.granularity.Granularities; +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.projections.AggregateProjectionSchema; +import org.apache.druid.segment.projections.ProjectionMetadata; +import org.apache.druid.segment.projections.ProjectionSchema; +import org.joda.time.Interval; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; + +class V10TimeBoundaryInspectorTest +{ + private static final Interval FALLBACK_INTERVAL = Intervals.of("2025-01-01/2025-01-02"); + + @Test + void testExactBoundsWhenMinMaxPresent() + { + final ProjectionMetadata projection = new ProjectionMetadata( + 100, + someSchema(), + DateTimes.of("2025-01-01T03:14:15").getMillis(), + DateTimes.of("2025-01-01T21:12:11").getMillis() + ); + + final V10TimeBoundaryInspector inspector = + V10TimeBoundaryInspector.forBaseProjection(projection, FALLBACK_INTERVAL); + + Assertions.assertTrue(inspector.isMinMaxExact(), "min/max from metadata are exact regardless of sort order"); + Assertions.assertEquals(DateTimes.of("2025-01-01T03:14:15"), inspector.getMinTime()); + Assertions.assertEquals(DateTimes.of("2025-01-01T21:12:11"), inspector.getMaxTime()); + } + + @Test + void testInexactFallbackWhenMinMaxMissing() + { + final ProjectionMetadata projection = new ProjectionMetadata(100, someSchema()); + + final V10TimeBoundaryInspector inspector = + V10TimeBoundaryInspector.forBaseProjection(projection, FALLBACK_INTERVAL); + + Assertions.assertFalse(inspector.isMinMaxExact(), "no metadata min/max → fall back to inexact interval bounds"); + Assertions.assertEquals(FALLBACK_INTERVAL.getStart(), inspector.getMinTime()); + Assertions.assertEquals(FALLBACK_INTERVAL.getEnd(), inspector.getMaxTime()); + } + + @Test + void testInexactFallbackWhenOnlyOneOfMinOrMaxMissing() + { + // both fields must be present to report exact; if either is null we fall back (this shouldn't happen) + final ProjectionMetadata onlyMin = new ProjectionMetadata( + 100, + someSchema(), + DateTimes.of("2025-01-01T03:00:00").getMillis(), + null + ); + Assertions.assertFalse( + V10TimeBoundaryInspector.forBaseProjection(onlyMin, FALLBACK_INTERVAL).isMinMaxExact() + ); + + final ProjectionMetadata onlyMax = new ProjectionMetadata( + 100, + someSchema(), + null, + DateTimes.of("2025-01-01T21:00:00").getMillis() + ); + Assertions.assertFalse( + V10TimeBoundaryInspector.forBaseProjection(onlyMax, FALLBACK_INTERVAL).isMinMaxExact() + ); + } + + private static ProjectionSchema someSchema() + { + return new AggregateProjectionSchema( + "p", + "time", + null, + VirtualColumns.create(Granularities.toVirtualColumn(Granularities.HOUR, "time")), + Arrays.asList("a", "time"), + new AggregatorFactory[]{new CountAggregatorFactory("count")}, + Arrays.asList(OrderBy.ascending("a"), OrderBy.ascending("time")) + ); + } +} diff --git a/processing/src/test/java/org/apache/druid/segment/file/CountingRangeReader.java b/processing/src/test/java/org/apache/druid/segment/file/CountingRangeReader.java index 37f50250abd0..4b47a55aab1f 100644 --- a/processing/src/test/java/org/apache/druid/segment/file/CountingRangeReader.java +++ b/processing/src/test/java/org/apache/druid/segment/file/CountingRangeReader.java @@ -19,6 +19,8 @@ package org.apache.druid.segment.file; +import org.apache.druid.segment.loading.DirectoryBackedRangeReader; + import java.io.File; import java.io.IOException; import java.io.InputStream; diff --git a/processing/src/test/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10Test.java b/processing/src/test/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10Test.java index 0294f65e20bf..a2403874757f 100644 --- a/processing/src/test/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10Test.java +++ b/processing/src/test/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10Test.java @@ -28,6 +28,7 @@ import org.apache.druid.segment.IndexIO; import org.apache.druid.segment.TestHelper; import org.apache.druid.segment.data.CompressionStrategy; +import org.apache.druid.segment.loading.DirectoryBackedRangeReader; import org.apache.druid.segment.loading.SegmentRangeReader; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -515,6 +516,84 @@ void testCorruptHeaderFileRecovery() throws IOException } } + @Test + void testEnsureAllDownloadedReadsOneRangePerContainer() throws IOException + { + final int numBundles = 3; + final int filesPerBundle = 4; + final File segmentFile = buildMultiBundleSegment(numBundles, filesPerBundle); + final File cacheDir = newCacheDir("ensure_all_bulk"); + final CountingRangeReader rangeReader = new CountingRangeReader(segmentFile.getParentFile()); + + try (PartialSegmentFileMapperV10 mapper = createMapper(rangeReader, cacheDir)) { + final int numContainers = mapper.getSegmentFileMetadata().getContainers().size(); + final int numFiles = mapper.getInternalFilenames().size(); + Assertions.assertEquals(numBundles, numContainers, "each tiny bundle should be its own container"); + Assertions.assertEquals(numBundles * filesPerBundle, numFiles); + + rangeReader.resetCount(); + mapper.ensureAllDownloaded(); + + // one range read per container, not one per internal file + Assertions.assertEquals(numContainers, rangeReader.getReadCount()); + Assertions.assertTrue(rangeReader.getReadCount() < numFiles, "bulk download must beat the per-file read count"); + Assertions.assertTrue(mapper.isFullyDownloaded()); + Assertions.assertEquals((long) numFiles * Integer.BYTES, mapper.getDownloadedBytes()); + + // a second pass is a no-op: every container is already fully downloaded + rangeReader.resetCount(); + mapper.ensureAllDownloaded(); + Assertions.assertEquals(0, rangeReader.getReadCount()); + } + + // bulk-downloaded contents match the eager mapper byte-for-byte + final DirectoryBackedRangeReader freshReader = new DirectoryBackedRangeReader(segmentFile.getParentFile()); + try (SegmentFileMapperV10 eager = SegmentFileMapperV10.create(segmentFile, JSON_MAPPER); + PartialSegmentFileMapperV10 lazy = createMapper(freshReader, newCacheDir("ensure_all_bulk_match"))) { + lazy.ensureAllDownloaded(); + for (String name : eager.getInternalFilenames()) { + Assertions.assertEquals(eager.mapFile(name), lazy.mapFile(name), "file[" + name + "]"); + } + } + } + + @Test + void testEnsureBundleDownloadedFetchesOnlyThatBundle() throws IOException + { + final File segmentFile = buildMultiBundleSegment(3, 4); + final File cacheDir = newCacheDir("ensure_bundle"); + final CountingRangeReader rangeReader = new CountingRangeReader(segmentFile.getParentFile()); + + try (PartialSegmentFileMapperV10 mapper = createMapper(rangeReader, cacheDir)) { + Assertions.assertTrue(mapper.getBundleNames().contains("b1")); + final int b1Containers = mapper.getContainerIndicesForBundle("b1").size(); + + rangeReader.resetCount(); + mapper.ensureBundleDownloaded("b1"); + + // one range read per container in the bundle, and nothing fetched from the other bundles + Assertions.assertEquals(b1Containers, rangeReader.getReadCount()); + Assertions.assertFalse(mapper.isFullyDownloaded(), "only one bundle was materialized"); + final Set downloaded = mapper.getDownloadedFiles(); + for (int i = 0; i < 4; i++) { + Assertions.assertTrue(downloaded.contains(StringUtils.format("b1/col_%d", i)), "b1 file should be downloaded"); + Assertions.assertFalse(downloaded.contains(StringUtils.format("b0/col_%d", i)), "b0 must stay undownloaded"); + Assertions.assertFalse(downloaded.contains(StringUtils.format("b2/col_%d", i)), "b2 must stay undownloaded"); + } + // the materialized bundle's files are mappable with the expected content + for (int i = 0; i < 4; i++) { + final ByteBuffer buf = mapper.mapFile(StringUtils.format("b1/col_%d", i)); + Assertions.assertNotNull(buf); + Assertions.assertEquals(100 + i, buf.getInt()); + } + + // an unknown bundle is a no-op + rangeReader.resetCount(); + mapper.ensureBundleDownloaded("does-not-exist"); + Assertions.assertEquals(0, rangeReader.getReadCount()); + } + } + private File newCacheDir(String name) throws IOException { final File dir = new File(tempDir, name + "_" + ThreadLocalRandom.current().nextInt()); @@ -522,6 +601,30 @@ private File newCacheDir(String name) throws IOException return dir; } + /** + * Build a segment whose files are spread across {@code numBundles} bundles. Each tiny file stays well under the max + * container size, so each bundle becomes its own container. Files in bundle {@code b} are named {@code "b{b}/col_{i}"} + * with content {@code b * 100 + i} (4 bytes). + */ + private File buildMultiBundleSegment(int numBundles, int filesPerBundle) throws IOException + { + final File baseDir = new File(tempDir, "multibundle_" + ThreadLocalRandom.current().nextInt()); + FileUtils.mkdirp(baseDir); + + try (SegmentFileBuilderV10 builder = SegmentFileBuilderV10.create(JSON_MAPPER, baseDir)) { + for (int b = 0; b < numBundles; b++) { + builder.startFileBundle("b" + b); + for (int i = 0; i < filesPerBundle; i++) { + final File tmpFile = new File(tempDir, StringUtils.format("mb-%d-%d.bin", b, i)); + Files.write(Ints.toByteArray(b * 100 + i), tmpFile); + builder.add(StringUtils.format("b%d/col_%d", b, i), tmpFile); + } + } + } + + return new File(baseDir, IndexIO.V10_FILE_NAME); + } + private File buildTestSegment(int numFiles, CompressionStrategy compression) throws IOException { final File baseDir = new File(tempDir, "base_" + ThreadLocalRandom.current().nextInt()); diff --git a/server/src/main/java/org/apache/druid/segment/loading/AcquireMode.java b/server/src/main/java/org/apache/druid/segment/loading/AcquireMode.java new file mode 100644 index 000000000000..28b58fb08496 --- /dev/null +++ b/server/src/main/java/org/apache/druid/segment/loading/AcquireMode.java @@ -0,0 +1,46 @@ +/* + * 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; + +/** + * How {@link SegmentCacheManager#acquireSegment} and {@link SegmentCacheManager#acquireCachedSegment} materialize a + * segment. Both modes return an open reference that callers must close when finished; they differ only in how much + * of the segment is downloaded before the reference is handed back. + */ +public enum AcquireMode +{ + /** + * Download the entire segment up front, so the returned {@link org.apache.druid.segment.Segment} is + * fully-materialized: every column is reachable synchronously by the time a cursor is built. This is the right + * choice for callers that use the synchronous {@link org.apache.druid.segment.CursorFactory#makeCursorHolder}. + */ + FULL, + + /** + * Mount only what is needed to serve the segment's schema (for a partial-eligible segment, just the metadata header). + * Callers must therefore use the async API such as + * {@link org.apache.druid.segment.CursorFactory#makeCursorHolderAsync} where column/bundle data is fetched on demand + * at cursor-build. + *

+ * Degrades to {@link #FULL} behavior when the segment's deep-storage layout cannot supply range reads (e.g. a + * non-V10 or zipped segment) or when partial downloads are disabled by configuration. + */ + PARTIAL +} diff --git a/server/src/main/java/org/apache/druid/segment/loading/LocalLoadSpec.java b/server/src/main/java/org/apache/druid/segment/loading/LocalLoadSpec.java index 4485792cc301..392f83b41eaf 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/LocalLoadSpec.java +++ b/server/src/main/java/org/apache/druid/segment/loading/LocalLoadSpec.java @@ -25,7 +25,9 @@ import com.fasterxml.jackson.annotation.JsonTypeName; import com.google.common.base.Preconditions; import org.apache.druid.guice.LocalDataStorageDruidModule; +import org.apache.druid.segment.IndexIO; +import javax.annotation.Nullable; import java.io.File; import java.nio.file.Files; import java.nio.file.Path; @@ -63,4 +65,18 @@ public LoadSpecResult loadSegment(final File outDir) throws SegmentLoadingExcept Preconditions.checkArgument(Files.exists(path), "[%s] does not exist", path); return new LoadSpecResult(puller.getSegmentFiles(path.toFile(), outDir).size()); } + + @Nullable + @Override + public SegmentRangeReader openRangeReader() + { + final File pathFile = path.toFile(); + if (!pathFile.isDirectory()) { + return null; + } + if (!new File(pathFile, IndexIO.V10_FILE_NAME).isFile()) { + return null; + } + return new DirectoryBackedRangeReader(pathFile); + } } diff --git a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java index ea7312cd26f6..c2dc29c75164 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java +++ b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java @@ -35,7 +35,9 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Collection; +import java.util.HashSet; import java.util.List; +import java.util.Set; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.ReentrantLock; @@ -140,14 +142,13 @@ public static List findContainersForBundle( ) { final List refs = new ArrayList<>(); - collectMatchingContainers(fileMapper.getSegmentFileMetadata(), bundleName, null, refs); + for (int containerIndex : fileMapper.getContainerIndicesForBundle(bundleName)) { + refs.add(new BundleContainerRef(null, containerIndex)); + } for (String externalFilename : fileMapper.getExternalFilenames()) { - collectMatchingContainers( - fileMapper.getExternalMapper(externalFilename).getSegmentFileMetadata(), - bundleName, - externalFilename, - refs - ); + for (int containerIndex : fileMapper.getExternalMapper(externalFilename).getContainerIndicesForBundle(bundleName)) { + refs.add(new BundleContainerRef(externalFilename, containerIndex)); + } } return List.copyOf(refs); } @@ -390,7 +391,7 @@ private void doMount(StorageLocation mountLocation) throws IOException // 2. References on metadata + parents (gates their deferred cleanup on this bundle's lifetime; matters for // statically-reserved dependencies where a drop fires `release()` directly without going through cache) - acquiredRefs.add(metadataEntry.acquireReference()); + acquiredRefs.add(metadataEntry.acquireMetadataReference()); for (PartialSegmentBundleCacheEntryIdentifier parentId : parentEntryIds) { final CacheEntry parentEntry = mountLocation.getCacheEntry(parentId); if (!(parentEntry instanceof PartialSegmentBundleCacheEntry)) { @@ -569,19 +570,50 @@ private void doActualUnmount() metadataEntry.unregisterBundle(this); } - private static void collectMatchingContainers( - SegmentFileMetadata fileMeta, - String bundleName, - @Nullable String externalFilename, - List out - ) + /** + * The distinct set of bundle names present across the segment's main file and every attached external file. A + * container written without an explicit {@link SegmentFileBuilder#startFileBundle} call (or from a segment that + * predates the bundle field) reports {@link SegmentFileBuilder#ROOT_BUNDLE_NAME}. Shared by the bootstrap + * discovery path and {@link #resolveBundleName}. + */ + public static Set bundleNames(PartialSegmentFileMapperV10 fileMapper) { - final List containers = fileMeta.getContainers(); - for (int ci = 0; ci < containers.size(); ci++) { - if (bundleName.equals(containers.get(ci).getBundle())) { - out.add(new BundleContainerRef(externalFilename, ci)); - } + // Union of each mapper's own bundle names (each computed once at mapper construction). + final Set names = new HashSet<>(fileMapper.getBundleNames()); + for (String externalFilename : fileMapper.getExternalFilenames()) { + names.addAll(fileMapper.getExternalMapper(externalFilename).getBundleNames()); + } + return names; + } + + /** + * Resolve the bundle name to actually acquire for a query that asked for {@code requestedBundleName}. Normally the + * requested name itself, but when no container exists with that name AND the segment's only bundle is + * {@link SegmentFileBuilder#ROOT_BUNDLE_NAME}, this method resolves to the root bundle so the whole segment becomes + * a single catch-all bundle and the query does not fail. This is to handle legacy v10 segments that might exist from + * before the bundle name was stored in the metadata. + *

+ * The fallback is deliberately gated on root being the segment's sole bundle: + *

    + *
  • A segment with named bundles present never silently reroutes a missing request to root, that stays a hard + * error (a genuinely missing bundle is a bug).
    • + *
  • A future writer that legitimately writes some files to the root bundle alongside named bundles is + * unaffected: its reader asks for {@code __root__} by name, which exists, so the fallback never fires.
    • + *
+ */ + public static String resolveBundleName(PartialSegmentFileMapperV10 fileMapper, String requestedBundleName) + { + if (!findContainersForBundle(fileMapper, requestedBundleName).isEmpty()) { + return requestedBundleName; + } + final Set present = bundleNames(fileMapper); + // legacy v10 segments from before bundle field was persisted in the segment + if (present.size() == 1 && present.contains(SegmentFileBuilder.ROOT_BUNDLE_NAME)) { + return SegmentFileBuilder.ROOT_BUNDLE_NAME; } + // Requested bundle is absent and there are named bundles present: leave the name as-is so the caller fails + // loudly (forBundle throws "no containers") rather than silently serving the wrong data. + return requestedBundleName; } private static void awaitMount(SettableFuture future) throws IOException 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 7793d2437206..88ea508a5326 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 @@ -24,12 +24,10 @@ import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.emitter.EmittingLogger; import org.apache.druid.segment.file.PartialSegmentFileMapperV10; -import org.apache.druid.segment.file.SegmentFileBuilder; -import org.apache.druid.segment.file.SegmentFileContainerMetadata; -import org.apache.druid.segment.file.SegmentFileMetadata; import org.apache.druid.segment.projections.Projections; import org.apache.druid.timeline.SegmentId; +import javax.annotation.Nullable; import java.io.File; import java.io.IOException; import java.io.InputStream; @@ -48,9 +46,18 @@ * The bootstrap is read-only with respect to deep storage; it never issues a range read. The on-disk header file is * parsed in-place by {@link PartialSegmentFileMapperV10#create} (which detects header corruption and, for that one * case, may delete the local copy; bootstrap callers should treat that as "no restorable state" and fall back to a - * cold start). Bundle entries are registered as weak entries on the storage location, mounted (which sparse-allocates - * any container files that weren't already present and re-establishes parent holds), and returned to the caller. + * cold start). *

+ * Two-phase contract: + *

    + *
  1. {@link #reserveFromDisk} (called from {@code getCachedSegments}, light): validates the header file is present, + * computes the actual on-disk reservation size, constructs the metadata cache entry, reserves it on the location. + * No mount, no range read.
    2. + *
  2. {@code metadata.mount(location)} (called from {@code bootstrap()} via the bootstrap executor, parallelizable): + * the metadata entry's own mount path builds the file mapper from the on-disk header, then internally invokes + * {@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).
    3. + *
* 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 @@ -62,29 +69,33 @@ public final class PartialSegmentCacheBootstrap private static final EmittingLogger LOG = new EmittingLogger(PartialSegmentCacheBootstrap.class); /** - * Restore a single partial segment's cache entries from its local on-disk layout. + * Reserve a partial segment's metadata cache entry on the supplied location from the on-disk header. Light-weight: + * no range read, no file mapper, no bundle work. The entry is registered as a weak cache entry (consistent with + * the runtime acquire path), so it is evictable once mounted; bootstrap-restored data is treated as a cache + * optimization, not a permanent fixture. The caller is expected to drive the actual mount through + * {@code metadata.mount(location)} later (typically via {@code SegmentCacheManager#bootstrap}), which builds the + * file mapper and cascades into {@link #restoreBundlesFromDisk}. * * @param segmentId the segment whose entries are being restored * @param localCacheDir the per-segment directory containing the header + container files * @param targetFilename the V10 entry-point filename - * @param externalFilenames any external segment file names that were registered as children of the entry-point - * file - * @param jsonMapper used to parse the header - * @param location the storage location these entries belong to; the metadata entry is registered as - * static and bundle entries are registered as weak - * @throws DruidException if the expected header file is missing, if a metadata or bundle reservation cannot be - * established on the location, or if the metadata entry mount produces a null file mapper - * @throws IOException propagated from {@link CacheEntry#mount} (metadata or bundle) or from on-disk header/bitmap - * I/O performed during restore + * @param externalFilenames any external segment file names that were registered as children of the entry-point file + * @param jsonMapper used by the metadata entry's mount path to parse the header + * @param storagePool thread pool the async cursor path submits on-demand column downloads to (which bounds + * load concurrency itself); may be null in tests that never invoke the cursor factory + * @param location the storage location to reserve the metadata entry on + * @return the reserved {@link PartialSegmentMetadataCacheEntry}; the caller is responsible for mounting it + * @throws DruidException if the expected header file is missing or the location cannot accept the reservation */ - public static RestoreResult restoreFromDisk( + public static PartialSegmentMetadataCacheEntry reserveFromDisk( SegmentId segmentId, File localCacheDir, String targetFilename, List externalFilenames, ObjectMapper jsonMapper, + @Nullable StorageLoadingThreadPool storagePool, StorageLocation location - ) throws IOException + ) { final File headerFile = new File(localCacheDir, targetFilename + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX); if (!headerFile.exists()) { @@ -104,80 +115,96 @@ public static RestoreResult restoreFromDisk( externalFilenames, BootstrapRangeReader.INSTANCE, jsonMapper, + storagePool, actualMetadataSize ); - if (!location.reserve(metadata)) { + if (!location.reserveWeak(metadata)) { throw DruidException.defensive( "Failed to reserve metadata entry for partial segment[%s] at location[%s]", segmentId, location.getPath() ); } + return metadata; + } - // From here, any throw must roll back: unmount any successfully-mounted bundles (releases their references on - // metadata) and then release the metadata entry from the location. Without this, a mid-bootstrap failure would - // leave the location partially reserved + mounted, which would confuse later restores and leak disk/memory. - final List mountedBundles = new ArrayList<>(); - boolean success = false; - try { - metadata.mount(location); + /** + * Discover, reserve, and mount any bundles whose container files survived on disk for the given partial segment. + * Invoked from {@link PartialSegmentMetadataCacheEntry#mount} after the file mapper has been built; safe to call + * unconditionally (on the fresh-acquire path there are no on-disk containers yet, so the call is a no-op). + *

+ * On any failure during bundle reservation or mount, attempts to roll back any partially-mounted bundles before + * propagating the throw. The metadata entry itself is NOT released here, the caller (typically + * {@link PartialSegmentMetadataCacheEntry#doMount}) handles metadata-level rollback on a propagated throw. + */ + static void restoreBundlesFromDisk(PartialSegmentMetadataCacheEntry metadata, StorageLocation location) + throws IOException + { + final PartialSegmentFileMapperV10 fileMapper = metadata.getFileMapper(); + if (fileMapper == null) { + // metadata isn't mounted yet (or has already been unmounted); nothing to restore + return; + } - // Discover bundle names across the main file and every external file, then keep only those whose owned - // container files actually exist on disk. Walks via the file mapper so the external mappers' SegmentFileMetadata - // are visited too — bundles can legitimately span the main file and one or more externals when the writer - // propagates startFileBundle across them. - final PartialSegmentFileMapperV10 fileMapper = metadata.getFileMapper(); - if (fileMapper == null) { - throw DruidException.defensive( - "Metadata entry mount produced null file mapper for segment[%s]", - segmentId - ); - } - final Set candidateBundleNames = discoverBundleNames(fileMapper); - final List presentBundleNames = filterByContainerPresence( - candidateBundleNames, - fileMapper, - localCacheDir - ); + final SegmentId segmentId = metadata.getSegmentId(); + final File localCacheDir = metadata.getLocalCacheDir(); - // Classify each present bundle as either mountable or orphaned. A bundle is orphaned when its inferred parent - // set includes a bundle that isn't itself present on disk; restoring it would only produce a degenerate state - // where column reads that resolve into the missing parent would fail at query time. Instead, delete the - // orphan's on-disk containers so the next access triggers a clean cold re-fetch from deep storage - final List mountableBundleNames = new ArrayList<>(); - final Set orphanedBundleNames = new HashSet<>(); - for (String name : presentBundleNames) { - boolean orphaned = false; - for (PartialSegmentBundleCacheEntryIdentifier parent : metadata.inferParentBundles(name)) { - if (!presentBundleNames.contains(parent.bundleName())) { - orphaned = true; - break; - } - } - if (orphaned) { - orphanedBundleNames.add(name); - } else { - mountableBundleNames.add(name); + // Discover bundle names across the main file and every external file, then keep only those whose owned + // container files actually exist on disk. Walks via the file mapper so the external mappers' SegmentFileMetadata + // are visited too; bundles can legitimately span the main file and one or more externals when the writer + // propagates startFileBundle across them. + final Set candidateBundleNames = PartialSegmentBundleCacheEntry.bundleNames(fileMapper); + final List presentBundleNames = filterByContainerPresence( + candidateBundleNames, + fileMapper, + localCacheDir + ); + if (presentBundleNames.isEmpty()) { + // Fresh acquire path or a partial whose containers were all evicted: nothing to do. + return; + } + + // Classify each present bundle as either mountable or orphaned. A bundle is orphaned when its inferred parent + // set includes a bundle that isn't itself present on disk; restoring it would only produce a degenerate state + // where column reads that resolve into the missing parent would fail at query time. Instead, delete the + // orphan's on-disk containers so the next access triggers a clean cold re-fetch from deep storage. + final List mountableBundleNames = new ArrayList<>(); + final Set orphanedBundleNames = new HashSet<>(); + for (String name : presentBundleNames) { + boolean orphaned = false; + for (PartialSegmentBundleCacheEntryIdentifier parent : metadata.inferParentBundles(name)) { + if (!presentBundleNames.contains(parent.bundleName())) { + orphaned = true; + break; } } + if (orphaned) { + orphanedBundleNames.add(name); + } else { + mountableBundleNames.add(name); + } + } - for (String orphanName : orphanedBundleNames) { - for (PartialSegmentBundleCacheEntry.BundleContainerRef ref : - PartialSegmentBundleCacheEntry.findContainersForBundle(fileMapper, orphanName)) { - fileMapper.mapperForContainer(ref.externalFilename()).evictContainer(ref.containerIndex()); - } - LOG.debug( - "Deleted on-disk state of orphaned bundle[%s] for segment[%s] (parent unrestorable); next access " - + "will trigger cold re-fetch", - orphanName, - segmentId - ); + for (String orphanName : orphanedBundleNames) { + for (PartialSegmentBundleCacheEntry.BundleContainerRef ref : + PartialSegmentBundleCacheEntry.findContainersForBundle(fileMapper, orphanName)) { + fileMapper.mapperForContainer(ref.externalFilename()).evictContainer(ref.containerIndex()); } + LOG.debug( + "Deleted on-disk state of orphaned bundle[%s] for segment[%s] (parent 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 - mountableBundleNames.sort(Comparator.comparing(name -> !Projections.BASE_TABLE_PROJECTION_NAME.equals(name))); + // mount base bundle before any dependent bundle so its hold is available when dependents acquire parent holds + 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. @@ -188,8 +215,8 @@ public static RestoreResult restoreFromDisk( parentIds ); // weak-reserve with a temporary hold so the mount call's own parent-hold acquisition can succeed; release the - // bootstrap hold immediately after, if the entry should remain alive for query-side access, the runtime hold - // chain (transitive parents from aggregates, segment-level holds from acquire APIs) keeps it pinned. + // bootstrap hold immediately after, if the entry should remain alive for query-side access, the runtime + // hold chain (transitive parents from aggregates, segment-level holds from acquire APIs) keeps it pinned. try (StorageLocation.ReservationHold bootstrapHold = location.addWeakReservationHold(bundle.getId(), () -> bundle)) { if (bootstrapHold == null) { @@ -203,23 +230,19 @@ public static RestoreResult restoreFromDisk( } mountedBundles.add(bundle); } - + success = true; LOG.debug( - "Restored partial segment[%s] from [%s]: metadata size[%d], bundles[%s], orphans[%s]", + "Restored bundles for partial segment[%s] from [%s]: bundles[%s], orphans[%s]", segmentId, localCacheDir, - actualMetadataSize, mountableBundleNames, orphanedBundleNames ); - success = true; - return new RestoreResult(metadata, mountedBundles); } finally { if (!success) { - // Roll back in reverse-dependency order: bundles first (so they release references on metadata + parents) - // then the metadata entry itself. The bundle/metadata cleanup is best-effort, log and continue rather than - // shadow the original failure. + // Reverse-dependency rollback for bundles only; the metadata entry's own rollback is the caller's + // responsibility (it will fire when the propagated throw escapes doMount's try/catch). for (PartialSegmentBundleCacheEntry bundle : mountedBundles) { try { bundle.unmount(); @@ -228,12 +251,6 @@ public static RestoreResult restoreFromDisk( LOG.warn(t, "Failed to roll back bundle[%s] during bootstrap failure for [%s]", bundle.getId(), segmentId); } } - try { - location.release(metadata); - } - catch (Throwable t) { - LOG.warn(t, "Failed to roll back metadata entry for partial segment[%s]", segmentId); - } } } } @@ -281,29 +298,6 @@ private static long sizeOf(File f) return f.exists() ? f.length() : 0; } - /** - * Discover the bundle names present in a segment by walking each container's - * {@link SegmentFileContainerMetadata#getBundle bundle} across the main file and every attached external file. - * The bundle field is always non-null — containers written without an explicit {@code startFileBundle} call - * (including those from older segments) default to {@link SegmentFileBuilder#ROOT_BUNDLE_NAME}. - */ - private static Set discoverBundleNames(PartialSegmentFileMapperV10 fileMapper) - { - final Set names = new HashSet<>(); - collectBundleNames(fileMapper.getSegmentFileMetadata(), names); - for (String externalFilename : fileMapper.getExternalFilenames()) { - collectBundleNames(fileMapper.getExternalMapper(externalFilename).getSegmentFileMetadata(), names); - } - return names; - } - - private static void collectBundleNames(SegmentFileMetadata fileMeta, Set out) - { - for (SegmentFileContainerMetadata container : fileMeta.getContainers()) { - out.add(container.getBundle()); - } - } - /** * Keep only bundles whose every owned container file exists on disk. The on-disk path for a container is * {@code {mapperTargetFilename}.container.{containerIndex:05d}} where {@code mapperTargetFilename} is the main @@ -367,29 +361,4 @@ public InputStream readRange(String filename, long offset, long length) } } - /** - * Hold-acquire result of a partial-segment restore: the always-static metadata entry plus the list of bundle - * entries (already mounted, registered as weak entries) discovered on disk for this segment. - */ - public static final class RestoreResult - { - private final PartialSegmentMetadataCacheEntry metadata; - private final List bundles; - - RestoreResult(PartialSegmentMetadataCacheEntry metadata, List bundles) - { - this.metadata = metadata; - this.bundles = List.copyOf(bundles); - } - - public PartialSegmentMetadataCacheEntry getMetadata() - { - return metadata; - } - - public List getBundles() - { - return bundles; - } - } } 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 c301123dd38a..244dd680fd39 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 @@ -20,16 +20,26 @@ package org.apache.druid.segment.loading; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Throwables; import com.google.common.util.concurrent.SettableFuture; import com.google.errorprone.annotations.concurrent.GuardedBy; +import org.apache.druid.common.asyncresource.AsyncResource; import org.apache.druid.error.DruidException; +import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.emitter.EmittingLogger; +import org.apache.druid.segment.PartialBundleAcquirer; +import org.apache.druid.segment.PartialQueryableIndex; +import org.apache.druid.segment.PartialQueryableIndexSegment; import org.apache.druid.segment.ReferenceCountingCloseableObject; +import org.apache.druid.segment.Segment; +import org.apache.druid.segment.column.ColumnConfig; import org.apache.druid.segment.file.PartialSegmentFileMapperV10; import org.apache.druid.segment.file.SegmentFileBuilder; import org.apache.druid.segment.file.SegmentFileMetadata; import org.apache.druid.segment.projections.Projections; import org.apache.druid.timeline.SegmentId; +import org.apache.druid.utils.CloseableUtils; import javax.annotation.Nullable; import java.io.Closeable; @@ -38,7 +48,9 @@ import java.util.ArrayList; import java.util.Collection; import java.util.List; +import java.util.Optional; import java.util.Set; +import java.util.concurrent.Callable; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicReference; @@ -60,7 +72,7 @@ * Reference-counted deferred cleanup. {@link #unmount()} does not necessarily release resources synchronously. * Callers that need the file mapper to stay alive across an intervening drop (e.g. a query reading column data * through {@link PartialSegmentBundleCacheEntry}, or another component that needs the parsed - * {@link SegmentFileMetadata}) acquire a reference via {@link #acquireReference()}; while any references are + * {@link SegmentFileMetadata}) acquire a reference via {@link #acquireMetadataReference()}; while any references are * outstanding, the actual close-file-mapper work is deferred. When the last reference releases the cleanup fires on * that thread. Bundle entries hold one such reference per active mount, so the typical pattern is: mount metadata, * mount bundle (which acquires a reference on metadata), use the bundle, unmount bundle (releases its reference and @@ -72,7 +84,7 @@ * the mapper is closed in {@link #doActualUnmount}. This is the right place to schedule work that should run only when * the entry is truly purged. */ -public class PartialSegmentMetadataCacheEntry implements ResizableCacheEntry +public class PartialSegmentMetadataCacheEntry implements SegmentCacheEntry, ResizableCacheEntry { private static final EmittingLogger LOG = new EmittingLogger(PartialSegmentMetadataCacheEntry.class); @@ -85,6 +97,9 @@ public class PartialSegmentMetadataCacheEntry implements ResizableCacheEntry private final ObjectMapper jsonMapper; private final long reservationEstimate; + @Nullable + private final StorageLoadingThreadPool storagePool; + // ReentrantLock instead of synchronized to avoid pinning virtual threads pre-JEP 491 private final ReentrantLock entryLock = new ReentrantLock(); @@ -99,6 +114,17 @@ public class PartialSegmentMetadataCacheEntry implements ResizableCacheEntry @GuardedBy("entryLock") @Nullable private PartialSegmentFileMapperV10 fileMapper; + // Cached PartialQueryableIndex for the mounted file mapper. Built lazily on first {@code acquireReference} call + // so concurrent acquireReference calls share the same memoized column-holder suppliers. first read of a column + // deserializes once and the ColumnHolder stays cached for subsequent queries against the same entry. Cleared + // during doActualUnmount so a subsequent mount rebuilds it against the fresh mapper. + @GuardedBy("entryLock") + @Nullable + private PartialQueryableIndex queryableIndex; + + // Bundle-acquirer view of this entry, cached at construction. Stateless (delegates to the entry's mutable fields + // under entryLock at call time), so safe to share across all acquireReference() calls. + private final PartialBundleAcquirer bundleAcquirer; // Optional deferred-cleanup hook invoked by doActualUnmount after the mapper is closed. private final AtomicReference onUnmount = new AtomicReference<>(); @@ -110,8 +136,8 @@ public class PartialSegmentMetadataCacheEntry implements ResizableCacheEntry // 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 - // (acquired via acquireReference()) are released. Re-created on mount-after-cleanup-completion. Null when the entry - // has never been mounted. + // (acquired via acquireMetadataReference()) are released. Re-created on mount-after-cleanup-completion. Null when + // the entry has never been mounted. private final AtomicReference> references = new AtomicReference<>(); // CAS+SettableFuture mount-dedup gate, mirroring the bundle entry's pattern. Without this, mount()'s slow range-read @@ -127,6 +153,7 @@ public PartialSegmentMetadataCacheEntry( List externalFilenames, SegmentRangeReader rangeReader, ObjectMapper jsonMapper, + @Nullable StorageLoadingThreadPool storagePool, long reservationEstimate ) { @@ -144,8 +171,10 @@ public PartialSegmentMetadataCacheEntry( this.externalFilenames = List.copyOf(externalFilenames); this.rangeReader = rangeReader; this.jsonMapper = jsonMapper; + this.storagePool = storagePool; this.reservationEstimate = reservationEstimate; this.currentSize = reservationEstimate; + this.bundleAcquirer = createBundleAcquirer(); } @Override @@ -154,11 +183,21 @@ public SegmentCacheEntryIdentifier getId() return id; } + @Override public SegmentId getSegmentId() { return segmentId; } + /** + * The per-segment cache directory this entry reads from and writes to. Exposed for the bundle-restore helper in + * {@link PartialSegmentCacheBootstrap#restoreBundlesFromDisk} so it can locate the on-disk container files. + */ + File getLocalCacheDir() + { + return localCacheDir; + } + @Override public long getSize() { @@ -236,15 +275,8 @@ public void mount(StorageLocation mountLocation) throws IOException // clear the future so the next caller gets a fresh attempt mountFuture.set(null); ours.setException(t); - if (t instanceof IOException) { - throw (IOException) t; - } - if (t instanceof RuntimeException) { - throw (RuntimeException) t; - } - if (t instanceof Error) { - throw (Error) t; - } + Throwables.propagateIfInstanceOf(t, IOException.class); + Throwables.propagateIfPossible(t); throw DruidException.defensive(t, "Failed to mount metadata entry[%s]", id); } verifyStillReservedOrRollback(mountLocation); @@ -280,87 +312,129 @@ private void doMount(StorageLocation mountLocation) throws IOException // blocked on it. adjustReservation also runs outside entryLock: StorageLocation.release goes // writeLock -> entryLock (via release -> unmount), so entryLock -> writeLock here would be a deadlock-prone // lock-order inversion. - entryLock.lock(); try { - if (location != null && fileMapper != null) { - if (!location.equals(mountLocation)) { - throw DruidException.defensive( - "Already mounted[%s] in location[%s] which differs from requested[%s]", - id, - location.getPath(), - mountLocation.getPath() - ); + entryLock.lock(); + try { + if (location != null && fileMapper != null) { + if (!location.equals(mountLocation)) { + throw DruidException.defensive( + "Already mounted[%s] in location[%s] which differs from requested[%s]", + id, + location.getPath(), + mountLocation.getPath() + ); + } + return; } - return; } - } - finally { - entryLock.unlock(); - } + finally { + entryLock.unlock(); + } - final PartialSegmentFileMapperV10 mapper = PartialSegmentFileMapperV10.create( - rangeReader, - jsonMapper, - localCacheDir, - targetFilename, - externalFilenames - ); + final PartialSegmentFileMapperV10 mapper = PartialSegmentFileMapperV10.create( + rangeReader, + jsonMapper, + localCacheDir, + targetFilename, + externalFilenames + ); - final long sizeToAdjust; - try { - final long actualSize = mapper.getOnDiskHeaderSize(); - if (actualSize > reservationEstimate) { - throw DruidException.forPersona(DruidException.Persona.OPERATOR) - .ofCategory(DruidException.Category.RUNTIME_FAILURE) - .build( - "Partial segment metadata for [%s] is [%d] bytes on disk, exceeding the " - + "configured reservation estimate of [%d] bytes. Increase " - + "druid.segmentCache.virtualStorageMetadataReservationEstimate.", - segmentId, - actualSize, - reservationEstimate - ); + final long sizeToAdjust; + try { + final long actualSize = mapper.getOnDiskHeaderSize(); + if (actualSize > reservationEstimate) { + throw DruidException.forPersona(DruidException.Persona.OPERATOR) + .ofCategory(DruidException.Category.RUNTIME_FAILURE) + .build( + "Partial segment metadata for [%s] is [%d] bytes on disk, exceeding the " + + "configured reservation estimate of [%d] bytes. Increase " + + "druid.segmentCache.virtualStorageMetadataReservationEstimate.", + segmentId, + actualSize, + reservationEstimate + ); + } + sizeToAdjust = actualSize < reservationEstimate ? actualSize : -1; + + entryLock.lock(); + try { + location = mountLocation; + fileMapper = mapper; + // Install (or re-install, after a previous mount/unmount cycle terminated the prior Phaser) the + // reference-counted gate over cleanup. Future acquireMetadataReference() / unmount() calls operate on this + // instance. + references.set(new ReferenceCountingCloseableObject(this::doActualUnmount) {}); + } + finally { + entryLock.unlock(); + } + } + catch (Throwable t) { + // mount failed; close mmaps and delete the on-disk header files so a retry starts clean. Mirrors the eager + // SegmentCacheEntry behavior: simpler to redo a small header range-read than to reason about whatever partial + // on-disk state the failure left. Crash-mid-mount across JVM restarts is still handled by the mapper's own + // corruption recovery when bootstrap runs at next startup; this path covers the in-process retry case. + try { + mapper.close(); + } + catch (Throwable closeError) { + t.addSuppressed(closeError); + } + try { + deleteHeaderFiles(); + } + catch (Throwable deleteError) { + t.addSuppressed(deleteError); + } + throw t; } - sizeToAdjust = actualSize < reservationEstimate ? actualSize : -1; - entryLock.lock(); + // Only shrink the reservation if the entry is still registered with the location. If we lost the reservation + // mid-mount (concurrent canceler / drop), adjustReservation would throw; defer to the post-mount check in + // mount() to roll back cleanly instead. + if (sizeToAdjust >= 0 && (mountLocation.isReserved(id) || mountLocation.isWeakReserved(id))) { + mountLocation.adjustReservation(id, sizeToAdjust); + } + + // Restore any bundles whose container files survived on disk for this segment. No-op on the fresh-acquire path + // (the partialDir was just created and contains no container files yet). The bundle restore needs metadata's + // file mapper, which is now installed, so this runs after the commit above. On any failure, fire our own + // deferred-cleanup gate via unmount() to undo the file-mapper install + delete header files; synchronous since + // we just installed the gate and no external caller has had a chance to acquire a reference yet. The location + // reservation release stays the caller's responsibility (matches mount's overall contract). try { - location = mountLocation; - fileMapper = mapper; - // Install (or re-install, after a previous mount/unmount cycle terminated the prior Phaser) the - // reference-counted gate over cleanup. Future acquireReference() / unmount() calls operate on this instance. - references.set(new ReferenceCountingCloseableObject(this::doActualUnmount) {}); + PartialSegmentCacheBootstrap.restoreBundlesFromDisk(this, mountLocation); } - finally { - entryLock.unlock(); + catch (Throwable t) { + try { + unmount(); + } + catch (Throwable cleanupError) { + t.addSuppressed(cleanupError); + } + Throwables.propagateIfInstanceOf(t, IOException.class); + Throwables.propagateIfPossible(t); + throw DruidException.defensive(t, "Failed to restore bundles for partial segment[%s]", segmentId); } } catch (Throwable t) { - // mount failed; close mmaps and delete the on-disk header files so a retry starts clean. Mirrors the eager - // SegmentCacheEntry behavior: simpler to redo a small header range-read than to reason about whatever partial - // on-disk state the failure left. Crash-mid-mount across JVM restarts is still handled by the mapper's own - // corruption recovery when bootstrap runs at next startup; this path covers the in-process retry case. + // A failed mount must not leave a lingering, un-re-mountable weak entry in the location. The inner rollbacks + // above close the mapper and delete the on-disk header, so any weak entry left behind is poison: a later + // findExistingPartialWithHold would resurrect it and re-mount would fail again (the header is gone and the + // bootstrap reserve path's entry uses a disk-only range reader). Remove it here so the next acquire rebuilds a + // fresh, deep-storage-capable entry via reservePartial. This is keyed on the entry being unheld: the runtime + // acquire path holds the entry via the AcquireSegmentAction's loadCleanup, so removeUnheldWeakEntry is a no-op + // there and the holder's release runnable performs cleanup instead. The bootstrap reserve path + // (StorageLocation.reserveWeak) places no hold, so this is what cleans it up. Runs outside entryLock (the + // inner blocks released it) so the writeLock -> entryLock order inside removeUnheldWeakEntry is respected. try { - mapper.close(); + mountLocation.removeUnheldWeakEntry(id); } - catch (Throwable closeError) { - t.addSuppressed(closeError); - } - try { - deleteHeaderFiles(); - } - catch (Throwable deleteError) { - t.addSuppressed(deleteError); + catch (Throwable removeError) { + t.addSuppressed(removeError); } throw t; } - - // Only shrink the reservation if the entry is still registered with the location. If we lost the reservation - // mid-mount (concurrent canceler / drop), adjustReservation would throw; defer to the post-mount check in - // mount() to roll back cleanly instead. - if (sizeToAdjust >= 0 && (mountLocation.isReserved(id) || mountLocation.isWeakReserved(id))) { - mountLocation.adjustReservation(id, sizeToAdjust); - } } private static void awaitMount(SettableFuture future) throws IOException @@ -384,10 +458,10 @@ private static void awaitMount(SettableFuture future) throws IOException } /** - * Triggers cleanup of this entry. If any references acquired via {@link #acquireReference()} are still 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. + * Triggers cleanup of this entry. If any references acquired via {@link #acquireMetadataReference()} are still + * 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. */ @Override public void unmount() @@ -403,10 +477,15 @@ public void unmount() * intervening {@link #unmount} call. The returned {@link Closeable} must be closed when the caller is done; at * that point if {@code unmount()} has already been called and no other references remain, the deferred cleanup * fires on the closing thread. + *

+ * This is the metadata-layer refcount used by {@link PartialSegmentBundleCacheEntry} (to keep the file mapper alive + * across an intervening drop) and by {@link #acquireReference} (which acquires one of these refs for each + * {@link PartialQueryableIndexSegment} it hands out). Callers that just want a {@link Segment} should use + * {@link #acquireReference} instead. * * @throws DruidException if the entry has never been mounted, or has already been cleaned up */ - public Closeable acquireReference() + public Closeable acquireMetadataReference() { final ReferenceCountingCloseableObject current = references.get(); if (current == null) { @@ -422,6 +501,130 @@ public Closeable acquireReference() )); } + /** + * Build a {@link Segment} backed by this entry's mounted file mapper. Returns {@link Optional#empty()} when the + * entry isn't mounted. The segment internally acquires one {@link #acquireMetadataReference} so that closing the + * segment is what releases the metadata-layer reference. + * + * @throws DruidException if no {@code storagePool} was supplied at construction time (this entry is + * metadata-only and cannot produce a queryable segment) + */ + @Override + public Optional acquireReference() + { + return buildSegmentWithExtraClose(null); + } + + /** + * Partial entries compose extra closeables natively into the segment's single {@code onClose} hook via a + * {@link Closer}. + */ + @Override + public Optional acquireReference(Closeable extraOnClose) + { + return buildSegmentWithExtraClose(extraOnClose); + } + + private Optional buildSegmentWithExtraClose(@Nullable Closeable extraOnClose) + { + final PartialQueryableIndex index; + entryLock.lock(); + try { + if (fileMapper == null) { + CloseableUtils.closeAndSuppressExceptions(extraOnClose, ignored -> {}); + return Optional.empty(); + } + if (queryableIndex == null) { + queryableIndex = new PartialQueryableIndex( + fileMapper.getSegmentFileMetadata(), + fileMapper, + ColumnConfig.DEFAULT + ); + } + index = queryableIndex; + } + finally { + entryLock.unlock(); + } + if (!hasStoragePool()) { + CloseableUtils.closeAndSuppressExceptions(extraOnClose, ignored -> {}); + throw DruidException.defensive( + "Cannot build segment for partial entry[%s]; no storage loading thread pool was supplied at construction", + id + ); + } + final Closeable metadataRef; + try { + metadataRef = acquireMetadataReference(); + } + catch (DruidException raceLost) { + // Defensive: acquireMetadataReference only fails if the entry is mid-unmount, but every caller of + // acquireReference holds a reservation hold on this entry across the call, and a held weak entry can't be + // concurrently unmounted, so this branch is unreachable today + CloseableUtils.closeAndSuppressExceptions(extraOnClose, ignored -> {}); + return Optional.empty(); + } + // Compose metadata-ref + optional extra into a single Closer so the segment's single-Closeable onClose contract + // stays simple. The Closer suppresses subsequent failures into the first, so every hook gets a chance to run. + final Closer onClose = Closer.create(); + if (extraOnClose != null) { + onClose.register(extraOnClose); + } + onClose.register(metadataRef); + return Optional.of( + new PartialQueryableIndexSegment( + index, + segmentId, + onClose, + bundleAcquirer + ) + ); + } + + /** + * Acquire a fully-materialized reference for {@link AcquireMode#FULL} with the metadata reference plus a + * hold on every bundle, all folded into the returned segment's close so the segment stays fully resident for its + * lifetime. + *

+ * Returns {@link Optional#empty()} when the segment is not fully resident in the cache. + */ + public Optional acquireFullReference(@Nullable Closeable extraOnClose) + { + final PartialSegmentFileMapperV10 mapper = getFileMapper(); + if (mapper == null || !isFullyDownloaded()) { + CloseableUtils.closeAndSuppressExceptions(extraOnClose, ignored -> {}); + return Optional.empty(); + } + final Closer bundleHolds = Closer.create(); + try { + // acquire() mounts each bundle's parents first, so iteration order doesn't matter. When the entry is resident + // every bundle is already mounted, so each acquire is a cheap refcount; only the rare eviction race below does + // any work (a discarded empty sparse-mount). + for (String bundleName : PartialSegmentBundleCacheEntry.bundleNames(mapper)) { + bundleHolds.register(bundleAcquirer.acquire(bundleName)); + } + // Re-check under the holds: a bundle could have been evicted between the check above and taking its hold, + // in which case acquire() re-mounted it as an empty container and isFullyDownloaded() is now false. Bail so the + // caller downloads via the executor path rather than handing back a segment that fails in makeCursorHolder. + if (!isFullyDownloaded()) { + CloseableUtils.closeAndSuppressExceptions(bundleHolds, ignored -> {}); + CloseableUtils.closeAndSuppressExceptions(extraOnClose, ignored -> {}); + return Optional.empty(); + } + } + catch (Throwable t) { + CloseableUtils.closeAndSuppressExceptions(bundleHolds, ignored -> {}); + CloseableUtils.closeAndSuppressExceptions(extraOnClose, ignored -> {}); + throw t; + } + final Closer onClose = Closer.create(); + if (extraOnClose != null) { + onClose.register(extraOnClose); + } + onClose.register(bundleHolds); + return acquireReference(onClose); + } + /** * The actual unmount work, invoked by the reference-counted gate's {@code onAdvance} once every outstanding * reference (plus the wrapper's own initial party) has been released. Closes the file mapper, deletes the on-disk @@ -443,6 +646,9 @@ private void doActualUnmount() LOG.warn(t, "Failed to close partial segment file mapper for [%s]", segmentId); } fileMapper = null; + // Drop the cached PartialQueryableIndex; the file mapper it referenced is now closed. A subsequent mount + // rebuilds it against a fresh mapper. + queryableIndex = null; location = null; // Clear the mount-dedup gate so a subsequent mount() on this same instance starts a fresh attempt. mountFuture.set(null); @@ -464,6 +670,23 @@ private void doActualUnmount() } } + /** + * Whether every internal file referenced by this segment's metadata (including any attached external mappers) has + * been downloaded. Delegates to {@link PartialSegmentFileMapperV10#isFullyDownloaded}. Returns false when the + * entry is not mounted. + */ + @Override + public boolean isFullyDownloaded() + { + entryLock.lock(); + try { + return fileMapper != null && fileMapper.isFullyDownloaded(); + } + finally { + entryLock.unlock(); + } + } + /** * Returns the file mapper held by this entry while mounted, or null if the entry has not been mounted. */ @@ -489,6 +712,177 @@ public SegmentFileMetadata getSegmentFileMetadata() return mapper == null ? null : mapper.getSegmentFileMetadata(); } + /** + * Build the {@link PartialBundleAcquirer} cached at construction time on {@link #bundleAcquirer}. The acquirer the + * partial-aware cursor factory uses after projection matching looks up an existing bundle entry by name on this + * metadata's storage location, creates + reserves a fresh one when none exists, mounts it (idempotent), and + * returns a per-cursor handle holding two protections: + *

    + *
  • a {@link StorageLocation.ReservationHold} so the cache doesn't evict the bundle while the cursor is reading, + * and
    • + *
  • a {@link PartialSegmentBundleCacheEntry#acquireReference reference} on the bundle entry so an explicit + * drop / unmount defers its actual cleanup until the cursor finishes.
    • + *
+ * Closing the returned {@link Closeable} releases both. Once the last hold drops the bundle becomes eviction-eligible + * and is unlinked from the location naturally on the next reclaim. + */ + @VisibleForTesting + PartialBundleAcquirer getBundleAcquirer() + { + return bundleAcquirer; + } + + /** + * Whether a usable {@link StorageLoadingThreadPool} was supplied at construction. When false, this entry can serve + * schema/metadata but cannot mount bundles or submit column-load tasks (e.g. a metadata-only stub entry). + */ + private boolean hasStoragePool() + { + return storagePool != null && storagePool.isAvailable(); + } + + private PartialBundleAcquirer createBundleAcquirer() + { + return new PartialBundleAcquirer() + { + @Override + public AsyncResource submitDownload(Callable task) + { + if (!hasStoragePool()) { + throw DruidException.defensive( + "No storage loading thread pool was supplied at construction for partial entry[%s]; cannot serve " + + "column-load tasks", + id + ); + } + return storagePool.submitUnmanagedAsyncResource(task); + } + + @Override + public Closeable acquire(String requestedBundleName) + { + final StorageLocation loc; + final PartialSegmentFileMapperV10 mapper; + entryLock.lock(); + try { + loc = location; + mapper = fileMapper; + } + finally { + entryLock.unlock(); + } + if (loc == null || mapper == null) { + // Programming error: acquire shouldn't be called on an unmounted metadata entry. The cache manager only + // hands out segments produced by acquireReference(), which only succeeds after mount. + throw DruidException.defensive( + "Metadata entry for [%s] is not mounted; cannot acquire bundle[%s]", + segmentId, + requestedBundleName + ); + } + + // Map the requested projection/base bundle name onto an actual on-disk bundle. Normally identity, but a + // legacy/untagged segment (sole bundle is the root bundle) resolves any request to the root catch-all. + final String bundleName = PartialSegmentBundleCacheEntry.resolveBundleName(mapper, requestedBundleName); + final PartialSegmentBundleCacheEntryIdentifier bundleId = + new PartialSegmentBundleCacheEntryIdentifier(segmentId, bundleName); + + // Take the hold atomically: reuse an existing entry if present, otherwise build + reserve a fresh one. Both + // addWeakReservationHold* variants take the hold under the location lock, so there is no getCacheEntry-then- + // hold race window. addWeakReservationHold is atomic create-or-get, so a concurrent acquire that raced in + // returns the same entry rather than a duplicate. And because a held entry can't be concurrently unmounted + // (SIEVE reclaim and removeUnheldWeakEntry both skip held entries, and the metadata cascade is gated by the + // bundle's reference on the metadata), the subsequent mount + acquireReference can't lose a drop race. + final StorageLocation.ReservationHold existingHold = + loc.addWeakReservationHoldIfExists(bundleId); + final StorageLocation.ReservationHold hold = + existingHold != null + ? existingHold + : loc.addWeakReservationHold( + bundleId, + () -> PartialSegmentBundleCacheEntry.forBundle( + PartialSegmentMetadataCacheEntry.this, + bundleName, + inferParentBundles(bundleName) + ) + ); + if (hold == null) { + throw DruidException.forPersona(DruidException.Persona.USER) + .ofCategory(DruidException.Category.CAPACITY_EXCEEDED) + .build( + "Unable to reserve bundle[%s] for segment[%s]; ensure enough disk space has been " + + "allocated to load all segments involved in the query", + bundleName, + segmentId + ); + } + + final PartialSegmentBundleCacheEntry bundle = hold.getEntry(); + if (!bundle.isMounted()) { + // Freshly built, or a stale entry from a prior cycle whose cleanup ran: (re-)mount it. mount() is idempotent + // via its mount-future dedup, so a concurrent acquirer that already mounted it makes this a cheap no-op. + mountBundleOrClose(bundle, loc, hold, bundleName); + } + + // Acquire the bundle's own reference, which gates its deferred cleanup on explicit unmount calls. The hold and + // the reference must close together when the cursor closes; they protect against different eviction triggers + // (SIEVE eviction vs. explicit drop cascade). + final Closeable bundleRef; + try { + bundleRef = bundle.acquireReference(); + } + catch (Throwable t) { + throw CloseableUtils.closeAndWrapInCatch(t, hold); + } + + return () -> { + try { + bundleRef.close(); + } + finally { + hold.close(); + } + }; + } + + /** + * Mount the bundle on the supplied location, closing the bootstrap hold and propagating a {@link DruidException} + * if the mount fails. Used by both the fresh-build and re-mount branches. + */ + private void mountBundleOrClose( + PartialSegmentBundleCacheEntry bundle, + StorageLocation loc, + StorageLocation.ReservationHold hold, + String bundleName + ) + { + // Mount this bundle's parents (e.g. __base for a projection bundle) FIRST: the bundle's own mount() takes + // 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). + final Closer parentHolds = Closer.create(); + try { + for (PartialSegmentBundleCacheEntryIdentifier parentId : inferParentBundles(bundleName)) { + parentHolds.register(acquire(parentId.bundleName())); + } + bundle.mount(loc); + } + catch (Throwable t) { + CloseableUtils.closeAndSuppressExceptions(hold, t::addSuppressed); + CloseableUtils.closeAndSuppressExceptions(parentHolds, t::addSuppressed); + throw DruidException.defensive(t, "Failed to mount bundle[%s] for segment[%s]", bundleName, segmentId); + } + CloseableUtils.closeAndSuppressExceptions( + parentHolds, + e -> LOG.warn(e, "Failed to release transient parent holds after mounting bundle[%s]", bundleName) + ); + } + }; + } + /** * Structural inference of the parent bundles that the given {@code bundleName} depends on within this segment. * Single source of truth for both bootstrap (which post-filters by what's actually restorable on disk) and the @@ -549,6 +943,7 @@ public Collection snapshotLinkedBundles() * hook after closing the file mapper and deleting the entry's storage-location files, outside the entry lock. * Replaces any previously-set hook. Pass {@code null} to clear. */ + @Override public void setOnUnmount(@Nullable Runnable hook) { onUnmount.set(hook); diff --git a/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheEntry.java b/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheEntry.java new file mode 100644 index 000000000000..40a8a4f22567 --- /dev/null +++ b/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheEntry.java @@ -0,0 +1,100 @@ +/* + * 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 org.apache.druid.error.DruidException; +import org.apache.druid.segment.ReferenceCountedSegmentProvider; +import org.apache.druid.segment.Segment; +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.utils.CloseableUtils; + +import javax.annotation.Nullable; +import java.io.Closeable; +import java.util.Optional; + +/** + * A {@link CacheEntry} that manages a {@link Segment} on local disk. + */ +public interface SegmentCacheEntry extends CacheEntry +{ + /** + * Identity of the segment this entry manages. Distinct from {@link #getId()}, which returns the + * {@link CacheEntryIdentifier} used by {@link StorageLocation} as a map key. + */ + SegmentId getSegmentId(); + + /** + * Acquire a closeable {@link Segment} reference, if one is available. Returns {@link Optional#empty()} when the + * entry is not mounted. + *

+ * The returned {@link Segment} must be closed by the caller. The entry tracks its own outstanding references and + * defers any cleanup work until they're all released. + */ + Optional acquireReference(); + + /** + * Acquire a segment reference and compose {@code extraOnClose} into the segment's close lifecycle. The + * implementation takes ownership of {@code extraOnClose}: on a successful acquire it becomes part of the returned + * segment's close; on a miss (entry not mounted, or a race lost to a concurrent drop) it is closed immediately + * before returning {@link Optional#empty()}. + */ + default Optional acquireReference(Closeable extraOnClose) + { + final Optional segment; + try { + segment = acquireReference(); + } + catch (Throwable t) { + CloseableUtils.closeAndSuppressExceptions(extraOnClose, t::addSuppressed); + throw t; + } + if (segment.isEmpty()) { + CloseableUtils.closeAndSuppressExceptions(extraOnClose, ignored -> {}); + return Optional.empty(); + } + if (segment.get() instanceof ReferenceCountedSegmentProvider.LeafReference leaf) { + return ReferenceCountedSegmentProvider.wrapCloseable(leaf, extraOnClose); + } + // Defensive: the returned segment isn't composable via wrapCloseable. Close both to avoid leaks and throw. + CloseableUtils.closeAndSuppressExceptions(segment.get(), ignored -> {}); + CloseableUtils.closeAndSuppressExceptions(extraOnClose, ignored -> {}); + throw DruidException.defensive( + "Segment[%s] of type[%s] from acquireReference() is not a LeafReference; cannot compose extra close", + segment.get().getDebugString(), + segment.get().getClass().getSimpleName() + ); + } + + /** + * Attach (or clear) a hook that runs when this entry is finally purged from the cache. Replaces any previously-set + * hook; pass {@code null} to clear. Implementations must invoke the hook exactly once, after all on-disk and + * in-memory state owned by the entry has been released, never while a reference acquired via + * {@link #acquireReference} is outstanding. + */ + void setOnUnmount(@Nullable Runnable hook); + + /** + * Whether every byte the segment needs is already on local disk. Complete entries always return {@code true} once + * mounted; partial entries return {@code true} only once every internal file has been pulled from deep storage. + * Callers can use this to gate synchronous segment operations on this so they never block a processing thread to + * load data on-demand. + */ + boolean isFullyDownloaded(); +} diff --git a/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java b/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java index d3fa1b84f2d4..c81c1d256a84 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java +++ b/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java @@ -19,11 +19,10 @@ package org.apache.druid.segment.loading; -import org.apache.druid.query.SegmentDescriptor; +import org.apache.druid.segment.CursorFactory; import org.apache.druid.segment.ReferenceCountedObjectProvider; import org.apache.druid.segment.Segment; import org.apache.druid.segment.SegmentLazyLoadFailCallback; -import org.apache.druid.segment.SegmentMapFunction; import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.SegmentId; @@ -71,7 +70,7 @@ public interface SegmentCacheManager /** * Given a {@link DataSegment}, which contains the instructions for where and how to fetch a {@link Segment} from * deep storage, this method tries to load and subsequently serve it to callers via - * {@link #acquireCachedSegment(SegmentId)} or {@link #acquireSegment(DataSegment)}. If the segment + * {@link #acquireCachedSegment(SegmentId, AcquireMode)} or {@link #acquireSegment(DataSegment, AcquireMode)}. If the segment * cannot be loaded either due to error or insufficient storage space, this method throws a * {@link SegmentLoadingException}. * @@ -103,31 +102,42 @@ public interface SegmentCacheManager void drop(DataSegment segment); /** - * Applies a {@link SegmentMapFunction} to a {@link Segment} if it is available in the cache. If not present in any - * storage location, this method will not attempt to download the {@link DataSegment} from deep storage. The - * {@link Segment} returned by this method is considered an open reference, cache implementations must not allow it - * to be dropped until it has been closed. As such, the returned {@link Segment} must be closed when the caller is - * finished doing segment things. + * Returns a {@link Segment} if it is available in the cache. If not present in any storage location, this method will + * not attempt to download the {@link DataSegment} from deep storage. The {@link Segment} returned by this method is + * considered an open reference, cache implementations must not allow it to be dropped until it has been closed. As + * such, the returned {@link Segment} must be closed when the caller is finished doing segment things. + *

+ * The {@code acquireMode} controls whether a partially-loaded cached entry is acceptable. With + * {@link AcquireMode#FULL} a cached entry is returned only if it is behaviorally fully-materialized. With + * {@link AcquireMode#PARTIAL} a mounted (but possibly not-yet-fully-downloaded) entry is sufficient, and callers + * acknowledge that synchronous methods may refuse if the segment is not fully downloaded; they must use asynchronous + * methods like {@link CursorFactory#makeCursorHolderAsync} instead. See {@link AcquireMode} for details. */ - Optional acquireCachedSegment(SegmentId segmentId); + Optional acquireCachedSegment(SegmentId segmentId, AcquireMode acquireMode); /** - * Returns a {@link AcquireSegmentAction} for a given {@link DataSegment} and {@link SegmentDescriptor}, which returns - * a reference provider for the {@link Segment} if already present in the cache, or tries to fetch from deep storage - * and map if not. The {@link Segment} returned by the provider returned by this method are considered an open - * reference, cache implementations must not allow the segment to be dropped until it has been closed. As such, the - * returned {@link Segment} from {@link ReferenceCountedObjectProvider#acquireReference()} must be closed when the - * caller is finished doing segment things. + * Returns a {@link AcquireSegmentAction} for a given {@link DataSegment}, which returns a reference provider for the + * {@link Segment} if already present in the cache, or tries to fetch from deep storage and map if not. The + * {@link Segment} returned by the provider returned by this method are considered an open reference, cache + * implementations must not allow the segment to be dropped until it has been closed. As such, the returned + * {@link Segment} from {@link ReferenceCountedObjectProvider#acquireReference()} must be closed when the caller is + * finished doing segment things. + *

+ * The {@code acquireMode} controls how much of the segment is downloaded before the action's future resolves. With + * {@link AcquireMode#FULL} the segment is fully materialized up front. With {@link AcquireMode#PARTIAL}, when the + * segment's {@link LoadSpec} supports range reads, only a minimal amount of metadata is downloaded and additional + * loading is deferred to the async methods callers use to interact with the returned segment; partial-ineligible + * segments fall back to {@link AcquireMode#FULL} behavior. See {@link AcquireMode} for details. */ - AcquireSegmentAction acquireSegment(DataSegment dataSegment); + AcquireSegmentAction acquireSegment(DataSegment dataSegment, AcquireMode acquireMode); /** - * Alternative to {@link #acquireCachedSegment(SegmentId)}, to return the {@link File} location of the segment files - * stored in the cache, instead of a {@link Optional}. Unlike {@link #acquireCachedSegment(SegmentId)} and - * {@link #acquireSegment(DataSegment)}, this method does not provide any protections for callers, - * and should only be used by callers that are in control of when {@link #drop(DataSegment)} is called. This method - * will not download the segment files from deep storage if they do not already exist in the cache, callers should use - * {@link #load(DataSegment)} before calling this method. + * Alternative to {@link #acquireCachedSegment(SegmentId, AcquireMode)}, to return the {@link File} location of the + * segment files stored in the cache, instead of a {@link Optional}. Unlike + * {@link #acquireCachedSegment(SegmentId, AcquireMode)} and {@link #acquireSegment(DataSegment, AcquireMode)}, this + * method does not provide any protections for callers, and should only be used by callers that are in control of + * when {@link #drop(DataSegment)} is called. This method will not download the segment files from deep storage if + * they do not already exist in the cache, callers should use {@link #load(DataSegment)} before calling this method. */ @Nullable File getSegmentFiles(DataSegment segment) throws SegmentLoadingException; diff --git a/server/src/main/java/org/apache/druid/segment/loading/SegmentLoaderConfig.java b/server/src/main/java/org/apache/druid/segment/loading/SegmentLoaderConfig.java index 1fbe3679ad03..63e87c4b1008 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/SegmentLoaderConfig.java +++ b/server/src/main/java/org/apache/druid/segment/loading/SegmentLoaderConfig.java @@ -105,6 +105,15 @@ public class SegmentLoaderConfig @JsonProperty("virtualStorageMetadataReservationEstimate") private long virtualStorageMetadataReservationEstimate = 16L * 1024L * 1024L; + /** + * When true, partial-eligible V10 segments are mounted via the partial machinery and + * {@link SegmentCacheManager#acquireSegment} with {@link AcquireMode#PARTIAL} returns a metadata-anchored segment + * whose columns are downloaded on demand. When false (the default), {@link AcquireMode#PARTIAL} falls back to + * {@link AcquireMode#FULL} so the entire segment is downloaded up front (matching pre-partial-download behavior). + */ + @JsonProperty("virtualStoragePartialDownloadsEnabled") + private boolean virtualStoragePartialDownloadsEnabled = false; + private long combinedMaxSize = 0; public List getLocations() @@ -197,12 +206,23 @@ public long getVirtualStorageMetadataReservationEstimate() return virtualStorageMetadataReservationEstimate; } + public boolean isVirtualStoragePartialDownloadsEnabled() + { + return virtualStorage && virtualStoragePartialDownloadsEnabled; + } + public SegmentLoaderConfig setLocations(List locations) { this.locations = Lists.newArrayList(locations); return this; } + public SegmentLoaderConfig setVirtualStoragePartialDownloadsEnabled(boolean enabled) + { + this.virtualStoragePartialDownloadsEnabled = enabled; + return this; + } + /** * Sets {@link #virtualStorage}. */ @@ -266,6 +286,7 @@ public String toString() ", virtualStorageUseVirtualThreads=" + virtualStorageUseVirtualThreads + ", virtualStorageIsEphemeral=" + virtualStorageIsEphemeral + ", virtualStorageMetadataReservationEstimate=" + virtualStorageMetadataReservationEstimate + + ", virtualStoragePartialDownloadsEnabled=" + virtualStoragePartialDownloadsEnabled + ", combinedMaxSize=" + combinedMaxSize + '}'; } 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 a43400efcb1a..55f808e2c096 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 @@ -41,12 +41,14 @@ import org.apache.druid.segment.ReferenceCountedSegmentProvider; import org.apache.druid.segment.Segment; import org.apache.druid.segment.SegmentLazyLoadFailCallback; +import org.apache.druid.segment.file.PartialSegmentFileMapperV10; import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.SegmentId; import org.apache.druid.utils.CloseableUtils; import javax.annotation.Nonnull; import javax.annotation.Nullable; +import java.io.Closeable; import java.io.File; import java.io.IOException; import java.io.InputStream; @@ -61,7 +63,6 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.Semaphore; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.Lock; @@ -92,7 +93,8 @@ public class SegmentLocalCacheManager implements SegmentCacheManager * segments simultaneously, one of them creates a lock first using {@link #lock(DataSegment)}. And then, all threads * compete with each other to get the lock. Finally, the lock should be released using * {@link #unlock(DataSegment, ReferenceCountingLock)}. A lock must be acquired any time a {@link SegmentCacheEntry} - * needs to be assigned to a {@link StorageLocation}. + * (either {@link CompleteSegmentCacheEntry} or {@link PartialSegmentMetadataCacheEntry}) needs to be assigned to a + * {@link StorageLocation}. *

* An example usage is: *

@@ -288,7 +290,7 @@ private void addFilesToCachedSegments( boolean removeInfo = true; - final SegmentCacheEntry cacheEntry = new SegmentCacheEntry(segment); + final CompleteSegmentCacheEntry cacheEntry = new CompleteSegmentCacheEntry(segment); for (StorageLocation location : locations) { // check for migrate from old nested local storage path format final File legacyPath = new File(location.getPath(), DataSegmentPusher.getDefaultStorageDir(segment, false)); @@ -304,6 +306,48 @@ private void addFilesToCachedSegments( cleanupLegacyCacheLocation(location.getPath(), legacyPath); } + // Partial-segment layout is signaled by a {targetFilename}.header file in the segment dir + final File partialDir = cacheEntry.toPotentialLocation(location.getPath()); + if (partialDir.exists() + && PartialSegmentCacheBootstrap.isPartialSegmentLayout(partialDir, IndexIO.V10_FILE_NAME)) { + if (!config.isVirtualStoragePartialDownloadsEnabled()) { + // Partial downloads are disabled but a partial-load layout (header + sparse containers) is on disk, e.g. the + // operator toggled druid.segmentCache.virtualStoragePartialDownloadsEnabled off. The eager path can't serve + // this layout, so reclaim it now and let the segment be re-downloaded in full on next load, rather than + // reserving it and failing when a query later tries to acquire it (see acquireExistingSegment). Leave + // removeInfo true so the segment is treated as uncached and re-loaded. Use the same move-then-delete as the + // other cache-dir removals; safe without a segment lock here because bootstrap runs before the node serves, + // so nothing concurrently mounts or drops this entry. + log.info( + "Deleting on-disk partial-load layout for segment[%s] in [%s] because partial downloads are disabled; " + + "it will be re-loaded via full download on next access.", + segment.getId(), + partialDir + ); + atomicMoveAndDeleteCacheEntryDirectory(partialDir); + continue; + } + removeInfo = false; + try { + PartialSegmentCacheBootstrap.reserveFromDisk( + segment.getId(), + partialDir, + IndexIO.V10_FILE_NAME, + List.of(), + jsonMapper, + virtualStorageLoadingThreadPool, + location + ); + cachedSegments.add(segment); + } + catch (Throwable t) { + // Reservation failed (header missing, location full, etc.) + log.warn(t, "Failed to reserve partial segment[%s] from disk; cold fetch on next access", segment.getId()); + } + // do not fall through to 'complete' path since this was a partial + continue; + } + if (cacheEntry.checkExists(location.getPath())) { removeInfo = false; final boolean reserveResult; @@ -364,9 +408,12 @@ public void removeInfoFile(final DataSegment segment) // defer deleting until the unmount operation of the cache entry, if possible, so that if the process stops before // the segment files are deleted, they can be properly managed on startup (since the info entry still exists) for (StorageLocation location : locations) { - final SegmentCacheEntry cacheEntry = location.getCacheEntry(entryId); - if (cacheEntry != null) { - isCached = isCached || cacheEntry.setDeleteInfoFileOnUnmount(); + final SegmentCacheEntry entry = location.getCacheEntry(entryId); + // gate isCached on isMounted() because an unmounted but reserved entry has no on-disk state worth deferring for; + // we want isCached = true exactly when there is real cached state that the unmount hook should clean up. + if (entry != null && entry.isMounted()) { + entry.setOnUnmount(() -> deleteSegmentInfoFile(segment)); + isCached = true; } } @@ -385,42 +432,26 @@ private void deleteSegmentInfoFile(DataSegment segment) } @Override - public Optional acquireCachedSegment(final SegmentId segmentId) + public Optional acquireCachedSegment(final SegmentId segmentId, final AcquireMode acquireMode) { - final SegmentCacheEntryIdentifier cacheEntryIdentifier = new SegmentCacheEntryIdentifier(segmentId); - for (StorageLocation location : locations) { - final SegmentCacheEntry cacheEntry = location.getStaticCacheEntry(cacheEntryIdentifier); - if (cacheEntry != null) { - return cacheEntry.acquireReference(); - } - final StorageLocation.ReservationHold hold = - location.addWeakReservationHoldIfExists(cacheEntryIdentifier); - try { - if (hold != null) { - if (hold.getEntry().isMounted()) { - Optional segment = hold.getEntry().acquireReference(); - if (segment.isPresent()) { - return ReferenceCountedSegmentProvider.wrapCloseable( - (ReferenceCountedSegmentProvider.LeafReference) segment.get(), - hold - ); - } - } - - hold.close(); - } - } - catch (Throwable e) { - hold.close(); - throw e; - } - } - return Optional.empty(); + // PARTIAL accepts a mounted-but-not-fully-downloaded entry; FULL (and PARTIAL when partial downloads are disabled) + // requires a behaviorally fully-materialized entry. + final boolean requireFullyDownloaded = + acquireMode == AcquireMode.FULL || !config.isVirtualStoragePartialDownloadsEnabled(); + return acquireCachedInternal(segmentId, requireFullyDownloaded); } @Override - public AcquireSegmentAction acquireSegment(final DataSegment dataSegment) + public AcquireSegmentAction acquireSegment(final DataSegment dataSegment, final AcquireMode acquireMode) { + // Partial-eligible segments route through the partial machinery: PARTIAL mounts lazily (header only, bundles load + // on demand), FULL force-downloads everything up front. Partial-ineligible segments fall through to the eager path + // below for both modes. + final SegmentRangeReader rangeReader = tryOpenRangeReader(dataSegment); + if (rangeReader != null) { + return acquirePartialInternal(dataSegment, rangeReader, acquireMode == AcquireMode.FULL); + } + final SegmentCacheEntryIdentifier identifier = new SegmentCacheEntryIdentifier(dataSegment.getId()); final AcquireSegmentAction acquireExisting = acquireExistingSegment(identifier); if (acquireExisting != null) { @@ -442,9 +473,9 @@ public AcquireSegmentAction acquireSegment(final DataSegment dataSegment) final Iterator iterator = strategy.getLocations(); while (iterator.hasNext()) { final StorageLocation location = iterator.next(); - final StorageLocation.ReservationHold hold = location.addWeakReservationHold( + final StorageLocation.ReservationHold hold = location.addWeakReservationHold( identifier, - () -> new SegmentCacheEntry(dataSegment) + () -> new CompleteSegmentCacheEntry(dataSegment) ); try { if (hold != null) { @@ -456,7 +487,7 @@ public AcquireSegmentAction acquireSegment(final DataSegment dataSegment) jsonMapper.writeValue(out, dataSegment); return null; }); - hold.getEntry().setDeleteInfoFileOnUnmount(); + hold.getEntry().setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); } return new AcquireSegmentAction( @@ -482,6 +513,357 @@ public AcquireSegmentAction acquireSegment(final DataSegment dataSegment) } } + /** + * Shared implementation for {@link #acquireCachedSegment}. Walks the storage locations checking for existing entries. + *

+ * When {@code requireFullyDownloaded} is {@code true} the entry must also report + * {@link SegmentCacheEntry#isFullyDownloaded()} ({@link AcquireMode#FULL}), which never hands back an entry that + * isn't behaviorally fully-materialized. When {@code false} a mounted entry is sufficient + * ({@link AcquireMode#PARTIAL}'s lazy-mount contract), even if some bundles are still not yet on disk. In either case + * {@link SegmentCacheEntry#acquireReference(Closeable)} composes the weak-entry hold into the returned segment's close + * lifecycle. + */ + private Optional acquireCachedInternal(SegmentId segmentId, boolean requireFullyDownloaded) + { + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segmentId); + for (StorageLocation location : locations) { + final SegmentCacheEntry staticEntry = location.getStaticCacheEntry(id); + if (staticEntry != null) { + if (staticEntry.isMounted() && (!requireFullyDownloaded || staticEntry.isFullyDownloaded())) { + return staticEntry.acquireReference(); + } + return Optional.empty(); + } + if (!config.isVirtualStorage()) { + continue; + } + final StorageLocation.ReservationHold hold = location.addWeakReservationHoldIfExists(id); + if (hold != null) { + final SegmentCacheEntry entry = hold.getEntry(); + if (entry.isMounted() && (!requireFullyDownloaded || entry.isFullyDownloaded())) { + if (requireFullyDownloaded && entry instanceof PartialSegmentMetadataCacheEntry partial) { + // A FULL reference on a partial segment must pin every bundle for the segment's lifetime (the sync cursor + // factory requires isFullyDownloaded(), which a metadata-only hold can't guarantee under SIEVE eviction). + // acquireFullReference folds `hold` into the segment close on success, or closes it and returns empty when + // a bundle is no longer resident — in which case we fall through to the downloading acquireSegment path. + return partial.acquireFullReference(hold); + } + return entry.acquireReference(hold); + } + hold.close(); + return Optional.empty(); + } + } + return Optional.empty(); + } + + /** + * Shared scaffolding for the partial-eligible branch of {@link #acquireSegment}. {@code fullDownload=false} powers + * {@link AcquireMode#PARTIAL} (lazy mount, header bytes only; bundles mount on demand at query time); + * {@code fullDownload=true} powers {@link AcquireMode#FULL} (mount + + * {@link PartialSegmentFileMapperV10#ensureAllDownloaded} so the returned segment is fully-materialized). + *

+ * Fast path: {@link #findExistingPartialWithHold} locates an existing entry across locations under a hold. If the + * entry is already usable (mounted, and fully-downloaded when required), return an immediate-future action whose + * {@code loadCleanup} is the hold (the supplier mints fresh segments per call, each with its own metadata + * reference, so no separate cleanup is needed). + *

+ * Slow path: under the per-segment lock, {@link #findOrReservePartial} reuses an existing not-yet-usable entry or + * reserves a fresh weak one, and the action submits mount (+ optional ensureAllDownloaded) to + * {@link #virtualStorageLoadingThreadPool} so callers that yield on the future never block a processing thread on + * deep-storage I/O. + */ + private AcquireSegmentAction acquirePartialInternal( + DataSegment dataSegment, + SegmentRangeReader rangeReader, + boolean fullDownload + ) + { + final ReservedPartial existing = findExistingPartialWithHold(dataSegment.getId()); + if (existing != null) { + final PartialSegmentMetadataCacheEntry partial = existing.metadata; + if (!fullDownload && partial.isMounted()) { + // Lazy/PARTIAL contract: hand back a metadata-anchored segment; bundles mount on demand at query time via the + // async cursor factory, so a metadata-only hold is correct here. The full-download fast path is intentionally + // omitted: a FULL segment must hold every bundle for its lifetime (the sync cursor factory requires + // isFullyDownloaded(), which a metadata-only hold can't guarantee under SIEVE eviction), so it goes through + // the slow path's bundle-holding + ensureAllDownloaded dance. The resident full case is already served without + // an executor hop upstream by acquireCachedSegment -> acquireCachedInternal -> acquireFullReference. + return new AcquireSegmentAction( + () -> Futures.immediateFuture(AcquireSegmentResult.cached(partial::acquireReference)), + existing.hold + ); + } + // Entry exists but isn't usable on the fast path (not mounted, or a full download is required). Release the + // fast-path hold and let the slow path re-find with a fresh hold and drive mount on the executor. + CloseableUtils.closeAndSuppressExceptions(existing.hold, ignored -> {}); + } + + final ReferenceCountingLock lock = lock(dataSegment); + synchronized (lock) { + try { + // Find an existing entry (in any state: mounted, not-yet-mounted, partially-downloaded) under a cache hold, + // or reserve a fresh weak entry with a hold if none exists. The metadata entry's mount-future dedup makes a + // no-op cheap when the entry is already mounted (the typical re-find case after the fast path closed its hold). + final ReservedPartial reserved = findOrReservePartial(dataSegment, rangeReader); + // holdHolder holds the metadata reservation hold immediately; the full-download path adds a hold for every + // bundle it mounts, all released together when the AcquireSegmentAction closes. + final HoldHolder holdHolder = new HoldHolder(reserved.hold); + return new AcquireSegmentAction( + // Memoized so repeat getSegmentFuture() calls return the same future rather than scheduling duplicate + // executor tasks. The entry's own mount-future dedup would prevent the actual work from being duplicated, + // but the executor scheduling and timing capture would still be wasted. + Suppliers.memoize(() -> { + // Capture submit time on first invocation of getSegmentFuture(), so waitTime measures the queue delay + // until the executor picks up the task. loadTime then covers mount (+ ensureAllDownloaded for the + // full-download path). + final long submitNanos = System.nanoTime(); + return virtualStorageLoadingThreadPool.getExecutorService().submit(() -> { + // The executor bounds concurrency itself (permit acquired inside the task on the worker thread), so + // waitNanos measures both the queue delay and any permit wait until this task actually starts. + final long taskStartNanos = System.nanoTime(); + final long waitNanos = taskStartNanos - submitNanos; + final boolean wasMounted = reserved.metadata.isMounted(); + // mount() is idempotent via PartialSegmentMetadataCacheEntry's mount-future dedup; already-mounted + // returns immediately, a concurrent mount is awaited, a fresh entry is mounted. The weak entry's + // hold-release runnable removes a never-mounted entry from the cache when our loadCleanup hold + // closes, so no explicit rollback is needed on failure. + try { + reserved.metadata.mount(reserved.location); + } + catch (IOException e) { + throw DruidException.defensive( + e, + "Failed to mount partial metadata for segment[%s]", + dataSegment.getId() + ); + } + // Pin the metadata across the rest of the task + final Closeable taskMetadataRef; + try { + taskMetadataRef = reserved.metadata.acquireMetadataReference(); + } + catch (DruidException raceLost) { + throw DruidException.defensive( + raceLost, + "Partial metadata for segment[%s] was dropped before %s task could complete", + dataSegment.getId(), + fullDownload ? "full-download" : "lazy mount" + ); + } + try { + final PartialSegmentFileMapperV10 mapper = reserved.metadata.getFileMapper(); + final long loadSizeBytes; + if (fullDownload) { + // Delta of internal-file bytes downloaded by this task. The small header range-read isn't + // counted in getDownloadedBytes; it's the negligible mount cost. + final long downloadedBefore = mapper.getDownloadedBytes(); + // Mount every bundle so the containers it owns are reserved on the location, SIEVE-evictable, and + // deleted on drop. Writing through the file mapper alone (ensureAllDownloaded) would leave those + // containers unmanaged by any cache entry. Each acquire returns a hold added to loadCleanup, so + // the fully-materialized segment stays resident for the eager (sync-cursor) caller until the + // action closes; on later eviction acquireCachedSegment returns empty and the caller re-downloads. + // acquire() mounts each bundle's parents first, so iteration order doesn't matter. + for (String bundleName : PartialSegmentBundleCacheEntry.bundleNames(mapper)) { + holdHolder.add(reserved.metadata.getBundleAcquirer().acquire(bundleName)); + } + mapper.ensureAllDownloaded(); + loadSizeBytes = mapper.getDownloadedBytes() - downloadedBefore; + } else { + // Lazy mount: report the header bytes when this task caused the mount; 0 when the entry was + // already mounted (a concurrent acquirer or earlier query did the load). + loadSizeBytes = wasMounted ? 0L : mapper.getOnDiskHeaderSize(); + } + // Record the actual-load bytes on the location (VSF_LOAD_*): the header bytes for a lazy mount, or the + // full downloaded delta for a full download. The partial metadata entry is always weak (reservePartial + // uses addWeakReservationHold). Skip when nothing was pulled (a cache hit / already-mounted re-acquire + // reports 0) so we don't log a spurious 0-byte load. On-demand per-column downloads on the lazy path + // happen later at query time and are captured by the query-layer load metrics (the + // AcquireSegmentResult load size), not by this per-location stat. + if (loadSizeBytes > 0) { + reserved.location.trackWeakLoad(loadSizeBytes); + } + final long loadNanos = System.nanoTime() - taskStartNanos; + return new AcquireSegmentResult( + reserved.metadata::acquireReference, + loadSizeBytes, + waitNanos, + loadNanos + ); + } + finally { + CloseableUtils.closeAndSuppressExceptions(taskMetadataRef, ignored -> {}); + } + }); + }), + holdHolder + ); + } + finally { + unlock(dataSegment, lock); + } + } + } + + /** + * Locate an existing partial metadata entry across storage locations and attach a eviction-protective hold to it. + * Returns {@code null} when no entry exists at any location. Race-safe (the hold prevents eviction from picking + * the entry between this lookup and the caller's subsequent use). Caller must close the returned hold (typically + * via {@link AcquireSegmentAction#loadCleanup} or by closing it directly when falling through to a reserve path). + */ + @Nullable + private ReservedPartial findExistingPartialWithHold(SegmentId segmentId) + { + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segmentId); + for (StorageLocation location : locations) { + final StorageLocation.ReservationHold hold = location.addWeakReservationHoldIfExists(id); + if (hold == null) { + continue; + } + try { + if (hold.getEntry() instanceof PartialSegmentMetadataCacheEntry partial) { + return new ReservedPartial(partial, location, hold); + } + throw DruidException.defensive( + "Unexpected non-partial cache entry[%s] at id[%s] on location[%s]", + hold.getEntry().getClass().getSimpleName(), + id, + location.getPath() + ); + } + catch (Throwable t) { + throw CloseableUtils.closeAndWrapInCatch(t, hold); + } + } + return null; + } + + /** + * 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 + * {@link SegmentLoaderConfig#isVirtualStoragePartialDownloadsEnabled}. Used to gate the partial-eligible branch of + * {@link #acquireSegment}; a null result causes that call site to fall through to the eager extraction path. + */ + @Nullable + private SegmentRangeReader tryOpenRangeReader(DataSegment dataSegment) + { + if (!config.isVirtualStoragePartialDownloadsEnabled()) { + return null; + } + try { + final LoadSpec loadSpec = jsonMapper.convertValue(dataSegment.getLoadSpec(), LoadSpec.class); + return loadSpec.openRangeReader(); + } + catch (IOException e) { + throw DruidException.forPersona(DruidException.Persona.OPERATOR) + .ofCategory(DruidException.Category.RUNTIME_FAILURE) + .build(e, "Failed to open range reader for segment[%s]", dataSegment.getId()); + } + } + + /** + * Synchronously reserve a partial metadata cache entry as a weak entry on the first storage location with capacity, + * placing a eviction-protective hold so it can't be evicted before the caller's {@link AcquireSegmentAction} is + * closed. Writes the segment info file so bootstrap can see this segment on restart. + *

+ * Uses {@link StorageLocation#addWeakReservationHold} which is atomic: it returns an existing entry's hold if one + * raced in (rare since callers hold the per-segment lock), or installs the freshly-built entry and returns its + * hold. On info-file write failure, releases the hold (which removes the never-mounted weak entry) and propagates. + * Throws CAPACITY_EXCEEDED if no location accepts the reservation. + */ + private ReservedPartial reservePartial(DataSegment dataSegment, SegmentRangeReader rangeReader) + { + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(dataSegment.getId()); + final Iterator iterator = strategy.getLocations(); + while (iterator.hasNext()) { + final StorageLocation location = iterator.next(); + final File partialDir = new File(location.getPath(), dataSegment.getId().toString()); + try { + FileUtils.mkdirp(partialDir); + } + catch (IOException e) { + throw DruidException.defensive(e, "Failed to create partial cache dir for segment[%s]", dataSegment.getId()); + } + final StorageLocation.ReservationHold hold = location.addWeakReservationHold( + id, + () -> new PartialSegmentMetadataCacheEntry( + dataSegment.getId(), + partialDir, + IndexIO.V10_FILE_NAME, + List.of(), + rangeReader, + jsonMapper, + virtualStorageLoadingThreadPool, + config.getVirtualStorageMetadataReservationEstimate() + ) + ); + if (hold == null) { + continue; + } + try { + if (!(hold.getEntry() instanceof PartialSegmentMetadataCacheEntry partial)) { + throw DruidException.defensive( + "Unexpected non-partial cache entry[%s] at id[%s] on location[%s]", + hold.getEntry().getClass().getSimpleName(), + id, + location.getPath() + ); + } + 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; + }); + } + partial.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); + return new ReservedPartial(partial, location, hold); + } + catch (Throwable t) { + throw CloseableUtils.closeAndWrapInCatch(t, hold); + } + } + throw DruidException.forPersona(DruidException.Persona.USER) + .ofCategory(DruidException.Category.CAPACITY_EXCEEDED) + .build( + "Unable to reserve partial metadata for segment[%s]; ensure enough disk space has been allocated", + dataSegment.getId() + ); + } + + /** + * Return a handle to a partial metadata entry for the given segment, paired with a eviction-protective hold: either + * an existing entry (in any mount state, located by {@link #findExistingPartialWithHold}), or a fresh weak + * reservation from {@link #reservePartial}. The caller is expected to drive + * {@link PartialSegmentMetadataCacheEntry#mount} on the returned handle inside the + * {@link AcquireSegmentAction}'s future; mount is idempotent via its mount-future dedup, so an already-mounted + * entry's mount call is cheap. The hold rides in the action's {@code loadCleanup} and is released when the + * action closes. + */ + private ReservedPartial findOrReservePartial(DataSegment dataSegment, SegmentRangeReader rangeReader) + { + final ReservedPartial existing = findExistingPartialWithHold(dataSegment.getId()); + if (existing != null) { + return existing; + } + return reservePartial(dataSegment, rangeReader); + } + + /** + * Pairing of a partial metadata entry (either pre-existing, discovered via {@link #findExistingPartialWithHold}, or + * freshly reserved by {@link #reservePartial}) with the storage location it lives on plus a eviction-protective hold. + * The hold rides in the {@link AcquireSegmentAction#loadCleanup} so the entry is protected from eviction across + * the action's lifetime; closing the action releases the hold. + */ + private record ReservedPartial( + PartialSegmentMetadataCacheEntry metadata, + StorageLocation location, + StorageLocation.ReservationHold hold + ) + { + } + @Nullable private AcquireSegmentAction acquireExistingSegment(SegmentCacheEntryIdentifier identifier) { @@ -492,16 +874,30 @@ private AcquireSegmentAction acquireExistingSegment(SegmentCacheEntryIdentifier location.addWeakReservationHoldIfExists(identifier) ); if (hold != null) { - if (hold.getEntry().isMounted()) { + if (!(hold.getEntry() instanceof CompleteSegmentCacheEntry complete)) { + // The eager (complete) acquire path found a non-complete entry under this id. Defensive backstop: when + // partial downloads are disabled, getCachedSegments now deletes any on-disk partial layout at bootstrap + // rather than reserving it, so a partial entry should not exist on this path. If one somehow does (e.g. a + // bootstrap delete failed), surface a clear operator error rather than a ClassCastException. + throw DruidException.forPersona(DruidException.Persona.OPERATOR) + .ofCategory(DruidException.Category.RUNTIME_FAILURE) + .build( + "Segment[%s] has partial-load cache state on disk but partial downloads are " + + "disabled; clear the segment cache directory or re-enable " + + "druid.segmentCache.virtualStoragePartialDownloadsEnabled", + identifier + ); + } + if (complete.isMounted()) { return new AcquireSegmentAction( - () -> Futures.immediateFuture(AcquireSegmentResult.cached(hold.getEntry().referenceProvider)), + () -> Futures.immediateFuture(AcquireSegmentResult.cached(complete.referenceProvider)), hold ); } else { // go ahead and mount it, someone else is probably trying this as well, but mount is done under a segment // lock and is a no-op if already mounted, and if we win we need it to be mounted return new AcquireSegmentAction( - makeOnDemandLoadSupplier(hold.getEntry(), location), + makeOnDemandLoadSupplier(complete, location), hold ); } @@ -532,7 +928,7 @@ public void load(final DataSegment dataSegment) throws SegmentLoadingException for (StorageLocation location : locations) { final SegmentCacheEntry cacheEntry = location.getCacheEntry(cacheEntryIdentifier); if (cacheEntry != null) { - cacheEntry.clearOnUnmount(); + cacheEntry.setOnUnmount(null); } } } @@ -542,11 +938,11 @@ public void load(final DataSegment dataSegment) throws SegmentLoadingException } return; } - final SegmentCacheEntry cacheEntry = new SegmentCacheEntry(dataSegment); + final CompleteSegmentCacheEntry cacheEntry = new CompleteSegmentCacheEntry(dataSegment); final ReferenceCountingLock lock = lock(dataSegment); synchronized (lock) { try { - final SegmentCacheEntry entry = assignLocationAndMount(cacheEntry, SegmentLazyLoadFailCallback.NOOP); + final CompleteSegmentCacheEntry entry = assignLocationAndMount(cacheEntry, SegmentLazyLoadFailCallback.NOOP); if (loadOnDownloadExec != null) { loadOnDownloadExec.submit(entry::loadIntoPageCache); } @@ -569,18 +965,38 @@ public void bootstrap( "bootstrap() should not be called when virtualStorageIsEphemeral is true" ); } - // during bootstrap, check if the segment exists in a location and mount it, getCachedSegments already + // during bootstrap, check if the segment exists in a location and mount it; getCachedSegments already // did the reserving for us final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(dataSegment.getId()); final ReferenceCountingLock lock = lock(dataSegment); synchronized (lock) { try { for (StorageLocation location : locations) { - final SegmentCacheEntry entry = location.getCacheEntry(id); - if (entry != null) { - entry.lazyLoadCallback = loadFailed; - entry.clearOnUnmount(); - entry.mount(location); + final CacheEntry entry = location.getCacheEntry(id); + if (entry == null) { + continue; + } + if (entry instanceof CompleteSegmentCacheEntry complete) { + complete.lazyLoadCallback = loadFailed; + complete.setOnUnmount(null); + complete.mount(location); + } else if (entry instanceof PartialSegmentMetadataCacheEntry partial) { + try { + partial.mount(location); + } + catch (IOException e) { + throw new SegmentLoadingException( + e, + "Failed to mount partial metadata for segment[%s]", + dataSegment.getId() + ); + } + } else { + throw DruidException.defensive( + "Unexpected cache entry type[%s] for segment[%s] during bootstrap", + entry.getClass().getSimpleName(), + dataSegment.getId() + ); } } } @@ -593,7 +1009,7 @@ public void bootstrap( final ReferenceCountingLock lock = lock(dataSegment); synchronized (lock) { try { - final SegmentCacheEntry entry = assignLocationAndMount(new SegmentCacheEntry(dataSegment), loadFailed); + final CompleteSegmentCacheEntry entry = assignLocationAndMount(new CompleteSegmentCacheEntry(dataSegment), loadFailed); if (loadOnBootstrapExec != null) { loadOnBootstrapExec.submit(entry::loadIntoPageCache); } @@ -608,15 +1024,28 @@ public void bootstrap( @Override public File getSegmentFiles(final DataSegment segment) { - final SegmentCacheEntry cacheEntry = new SegmentCacheEntry(segment); + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segment.getId()); final ReferenceCountingLock lock = lock(segment); synchronized (lock) { try { for (StorageLocation location : locations) { - final SegmentCacheEntry entry = location.getCacheEntry(cacheEntry.id); - if (entry != null) { - return entry.storageDir; + final CacheEntry entry = location.getCacheEntry(id); + if (entry == null) { + continue; + } + if (entry instanceof CompleteSegmentCacheEntry complete) { + return complete.storageDir; } + // The only caller of SegmentCacheManager#getSegmentFiles is DruidSegmentInputEntity (native batch ingest), + // which constructs its cache manager with virtualStorage=false, so a partial entry showing up here is a + // programming error, not a runtime configuration issue. Fail loudly rather than silently returning null, + // which would lead to a confusing "missing segment file" error downstream. + throw DruidException.defensive( + "getSegmentFiles[%s] called on a partial-segment cache entry (type[%s]); this API is only supported " + + "for complete cache entries (callers should not enable virtual storage on their cache manager)", + segment.getId(), + entry.getClass().getSimpleName() + ); } } finally { @@ -631,7 +1060,7 @@ public void drop(final DataSegment segment) { final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segment.getId()); for (StorageLocation location : locations) { - final SegmentCacheEntry entry = location.getCacheEntry(id); + final CacheEntry entry = location.getCacheEntry(id); if (entry != null) { location.release(entry); } @@ -681,7 +1110,7 @@ public StorageLoadingThreadPool getLoadingThreadPool() @VisibleForTesting boolean isSegmentCached(final DataSegment segment) { - final SegmentCacheEntry cacheEntry = new SegmentCacheEntry(segment); + final CompleteSegmentCacheEntry cacheEntry = new CompleteSegmentCacheEntry(segment); for (StorageLocation location : locations) { if (cacheEntry.checkExists(location.getPath())) { return true; @@ -699,9 +1128,9 @@ boolean isSegmentCached(final DataSegment segment) @Nullable public ReferenceCountedSegmentProvider getSegmentReferenceProvider(DataSegment segment) { - final SegmentCacheEntry cacheEntry = new SegmentCacheEntry(segment); + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segment.getId()); for (StorageLocation location : locations) { - final SegmentCacheEntry entry = location.getCacheEntry(cacheEntry.id); + final CompleteSegmentCacheEntry entry = checkComplete(location.getCacheEntry(id), id); if (entry != null) { return entry.referenceProvider; } @@ -740,7 +1169,7 @@ private File getEffectiveInfoDir() } private Supplier> makeOnDemandLoadSupplier( - final SegmentCacheEntry entry, + final CompleteSegmentCacheEntry entry, final StorageLocation location ) { @@ -749,28 +1178,15 @@ private Supplier> makeOnDemandLoadSupplie final long startTime = System.nanoTime(); return virtualStorageLoadingThreadPool.getExecutorService().submit( () -> { - // Acquire a permit inside the task so that waiting parks the (virtual) thread instead of blocking the - // submitter. In fixed-pool mode permits is null and the pool size itself bounds concurrency. - final Semaphore loadingPermits = virtualStorageLoadingThreadPool.getPermits(); - if (loadingPermits != null) { - loadingPermits.acquireUninterruptibly(); - } - try { - final long execStartTime = System.nanoTime(); - final long waitTime = execStartTime - startTime; - entry.mount(location); - return new AcquireSegmentResult( - entry.referenceProvider, - entry.dataSegment.getSize(), - waitTime, - System.nanoTime() - execStartTime - ); - } - finally { - if (loadingPermits != null) { - loadingPermits.release(); - } - } + final long execStartTime = System.nanoTime(); + final long waitTime = execStartTime - startTime; + entry.mount(location); + return new AcquireSegmentResult( + entry.referenceProvider, + entry.dataSegment.getSize(), + waitTime, + System.nanoTime() - execStartTime + ); } ); } @@ -815,8 +1231,8 @@ private void unlock(final DataSegment dataSegment, final ReferenceCountingLock l ); } - private SegmentCacheEntry assignLocationAndMount( - final SegmentCacheEntry cacheEntry, + private CompleteSegmentCacheEntry assignLocationAndMount( + final CompleteSegmentCacheEntry cacheEntry, final SegmentLazyLoadFailCallback segmentLoadFailCallback ) throws SegmentLoadingException { @@ -824,10 +1240,10 @@ private SegmentCacheEntry assignLocationAndMount( for (StorageLocation location : locations) { if (cacheEntry.checkExists(location.getPath())) { if (location.isReserved(cacheEntry.id) || location.reserve(cacheEntry)) { - final SegmentCacheEntry entry = location.getCacheEntry(cacheEntry.id); + final CompleteSegmentCacheEntry entry = checkComplete(location.getCacheEntry(cacheEntry.id), cacheEntry.id); if (entry != null) { entry.lazyLoadCallback = segmentLoadFailCallback; - entry.clearOnUnmount(); + entry.setOnUnmount(null); entry.mount(location); return entry; } @@ -846,10 +1262,10 @@ private SegmentCacheEntry assignLocationAndMount( final StorageLocation location = locationsIterator.next(); if (location.reserve(cacheEntry)) { try { - final SegmentCacheEntry entry = location.getCacheEntry(cacheEntry.id); + final CompleteSegmentCacheEntry entry = checkComplete(location.getCacheEntry(cacheEntry.id), cacheEntry.id); if (entry != null) { entry.lazyLoadCallback = segmentLoadFailCallback; - entry.clearOnUnmount(); + entry.setOnUnmount(null); entry.mount(location); return entry; } @@ -862,6 +1278,28 @@ private SegmentCacheEntry assignLocationAndMount( throw new SegmentLoadingException("Failed to load segment[%s] in all locations.", cacheEntry.id); } + /** + * Narrow a {@link CacheEntry} to {@link CompleteSegmentCacheEntry} with a defensive check. Returns null when the + * entry is missing. Throws when a non-complete entry (e.g. {@link PartialSegmentMetadataCacheEntry}) is registered + * under the same identifier. Callers of {@link #assignLocationAndMount} only operate on the non-virtual-storage + * eager path, so encountering a partial entry here is a programming error. + */ + @Nullable + private static CompleteSegmentCacheEntry checkComplete(@Nullable CacheEntry entry, SegmentCacheEntryIdentifier id) + { + if (entry == null) { + return null; + } + if (entry instanceof CompleteSegmentCacheEntry complete) { + return complete; + } + throw DruidException.defensive( + "Expected a complete cache entry for [%s], got [%s].", + id, + entry.getClass().getSimpleName() + ); + } + /** * Performs an atomic move to a sibling {@link #DROP_PATH} directory, and then deletes the directory and logs about * it. This method should only be called under the lock of a {@link #segmentLocks}. @@ -948,7 +1386,7 @@ private void decrement() } } - private final class SegmentCacheEntry implements CacheEntry + private final class CompleteSegmentCacheEntry implements SegmentCacheEntry { private final SegmentCacheEntryIdentifier id; private final DataSegment dataSegment; @@ -962,7 +1400,7 @@ private final class SegmentCacheEntry implements CacheEntry // https://openjdk.org/jeps/491, we could consider switching back after java 24+ is the minimum version private final ReentrantLock entryLock = new ReentrantLock(); - private SegmentCacheEntry(final DataSegment dataSegment) + private CompleteSegmentCacheEntry(final DataSegment dataSegment) { this.dataSegment = dataSegment; this.id = new SegmentCacheEntryIdentifier(dataSegment.getId()); @@ -1079,7 +1517,7 @@ public void mount(StorageLocation mountLocation) throws SegmentLoadingException } if (config.isVirtualStorageEphemeral()) { - setDeleteInfoFileOnUnmount(); + setOnUnmount(() -> deleteSegmentInfoFile(dataSegment)); } } catch (SegmentLoadingException e) { @@ -1154,6 +1592,13 @@ public void unmount() } } + @Override + public SegmentId getSegmentId() + { + return dataSegment.getId(); + } + + @Override public Optional acquireReference() { entryLock.lock(); @@ -1168,30 +1613,23 @@ public Optional acquireReference() } } - public boolean setDeleteInfoFileOnUnmount() + @Override + public void setOnUnmount(@Nullable Runnable hook) { entryLock.lock(); try { - if (location == null) { - return false; - } - onUnmount.set(() -> deleteSegmentInfoFile(dataSegment)); - return true; + onUnmount.set(hook); } finally { entryLock.unlock(); } } - public void clearOnUnmount() + @Override + public boolean isFullyDownloaded() { - entryLock.lock(); - try { - onUnmount.set(null); - } - finally { - entryLock.unlock(); - } + // Complete entries extract every byte at mount time, so by definition a mounted entry is fully downloaded. + return isMounted(); } public void loadIntoPageCache() @@ -1296,7 +1734,7 @@ public boolean equals(Object o) if (o == null || getClass() != o.getClass()) { return false; } - SegmentCacheEntry that = (SegmentCacheEntry) o; + CompleteSegmentCacheEntry that = (CompleteSegmentCacheEntry) o; return Objects.equals(dataSegment, that.dataSegment); } @@ -1306,4 +1744,49 @@ public int hashCode() return Objects.hashCode(dataSegment); } } + + /** + * The {@link AcquireSegmentAction#close()} cleanup for a partial acquire: a set of cache holds that keep the + * acquired segment's entries resident for the action's lifetime. Seeded with the metadata reservation hold at + * construction; the full-download path {@link #add adds} a hold per bundle it mounts while the load future runs. + * Closing releases everything (LIFO). + *

+ * Thread-safe because the future that {@link #add}s bundle holds runs on the load executor while a different thread + * may close the action (query cancel / timeout racing a blocked {@code getSegmentFuture().get()}). A hold added + * after the action has already been closed is closed immediately rather than leaked. + */ + private static final class HoldHolder implements Closeable + { + @GuardedBy("this") + private final Closer holds = Closer.create(); + @GuardedBy("this") + private boolean closed = false; + + private HoldHolder(Closeable initialHold) + { + holds.register(initialHold); + } + + private void add(Closeable hold) + { + final boolean alreadyClosed; + synchronized (this) { + alreadyClosed = closed; + if (!alreadyClosed) { + holds.register(hold); + } + } + if (alreadyClosed) { + // the action was closed while the load future was still running; release the late hold rather than leak it + CloseableUtils.closeAndSuppressExceptions(hold, ignored -> {}); + } + } + + @Override + public synchronized void close() throws IOException + { + closed = true; + holds.close(); + } + } } diff --git a/server/src/main/java/org/apache/druid/segment/loading/StorageLoadingThreadPool.java b/server/src/main/java/org/apache/druid/segment/loading/StorageLoadingThreadPool.java index 9c8de953b9d3..18668a51e2c6 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/StorageLoadingThreadPool.java +++ b/server/src/main/java/org/apache/druid/segment/loading/StorageLoadingThreadPool.java @@ -19,18 +19,30 @@ package org.apache.druid.segment.loading; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.util.concurrent.ForwardingListeningExecutorService; +import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListeningExecutorService; import com.google.common.util.concurrent.MoreExecutors; +import org.apache.druid.common.asyncresource.AsyncResource; +import org.apache.druid.common.asyncresource.AsyncResources; import org.apache.druid.error.DruidException; import org.apache.druid.guice.StorageNodeModule; import org.apache.druid.java.util.common.concurrent.Execs; import org.apache.druid.java.util.common.lifecycle.LifecycleStop; import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.segment.PartialBundleAcquirer; import org.apache.druid.segment.loading.external.VirtualStorageManager; import javax.annotation.Nullable; +import java.io.Closeable; +import java.util.Collection; +import java.util.List; +import java.util.concurrent.Callable; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; /** * Holds the thread pool used for background loading by {@link SegmentLocalCacheManager} and @@ -41,21 +53,17 @@ public class StorageLoadingThreadPool private static final Logger log = new Logger(StorageNodeModule.class); private final ListeningExecutorService exec; - private final Semaphore permits; public StorageLoadingThreadPool( - @Nullable final ListeningExecutorService exec, - @Nullable final Semaphore permits + @Nullable final ListeningExecutorService exec ) { this.exec = exec; - this.permits = permits; } public static StorageLoadingThreadPool createFromConfig(final SegmentLoaderConfig config) { final ListeningExecutorService exec; - final Semaphore permits; if (config.isVirtualStorage()) { if (config.getVirtualStorageLoadThreads() <= 0) { @@ -71,20 +79,21 @@ public static StorageLoadingThreadPool createFromConfig(final SegmentLoaderConfi "Using virtual storage mode with virtual threads - max concurrent on demand loads: [%d].", config.getVirtualStorageLoadThreads() ); - permits = new Semaphore(config.getVirtualStorageLoadThreads()); - exec = MoreExecutors.listeningDecorator( - Executors.newThreadPerTaskExecutor( - Thread.ofVirtual() - .name("VirtualStorageOnDemandLoadingThread-", 0) - .factory() - ) + exec = new PermitBoundedListeningExecutorService( + MoreExecutors.listeningDecorator( + Executors.newThreadPerTaskExecutor( + Thread.ofVirtual() + .name("VirtualStorageOnDemandLoadingThread-", 0) + .factory() + ) + ), + new Semaphore(config.getVirtualStorageLoadThreads()) ); } else { log.info( "Using virtual storage mode with fixed platform thread pool - on demand load threads: [%d].", config.getVirtualStorageLoadThreads() ); - permits = null; exec = MoreExecutors.listeningDecorator( Executors.newFixedThreadPool( config.getVirtualStorageLoadThreads(), @@ -94,10 +103,9 @@ public static StorageLoadingThreadPool createFromConfig(final SegmentLoaderConfi } } else { exec = null; - permits = null; } - return new StorageLoadingThreadPool(exec, permits); + return new StorageLoadingThreadPool(exec); } /** @@ -105,7 +113,7 @@ public static StorageLoadingThreadPool createFromConfig(final SegmentLoaderConfi */ public static StorageLoadingThreadPool none() { - return new StorageLoadingThreadPool(null, null); + return new StorageLoadingThreadPool(null); } public boolean isAvailable() @@ -113,6 +121,13 @@ public boolean isAvailable() return exec != null; } + /** + * Executor for on-demand load work. Concurrency is bounded by the executor itself: in the virtual-thread path it is + * a {@link PermitBoundedListeningExecutorService} wrapping an unbounded thread-per-virtual-thread executor (so the + * permit count, not the thread count, is the bound); in the fixed-pool path the pool size is the bound. Every + * on-demand-load submission, including those routed through {@link #submitUnmanagedAsyncResource}, is therefore + * bounded without callers having to acquire a permit themselves. + */ public ListeningExecutorService getExecutorService() { if (exec == null) { @@ -122,15 +137,33 @@ public ListeningExecutorService getExecutorService() } /** - * Bounds the number of in-flight on-demand loads when running with virtual threads. If null, no permits are used; - * the natural limit of the {@link #getExecutorService()} is sufficient. If nonnull, callers must acquire a permit - * in the tasks submitted to {@link #getExecutorService()}, before doing real work. Permits must be released before - * the task completes. + * Submit a task to the pool and hand back an {@link AsyncResource} that becomes ready when the task completes, + * exposing the task's (non-null) result. The result is treated as a plain value with no lifecycle: closing the + * returned resource does not close the result; closing it before completion cancels the task. + * + *

This is the unmanaged counterpart of {@link #submitCloseableAsyncResource}; use that when the task's result + * owns a lifecycle. + * + * @see AsyncResources#fromFutureUnmanaged(ListenableFuture) + */ + public AsyncResource submitUnmanagedAsyncResource(Callable task) + { + return AsyncResources.fromFutureUnmanaged(getExecutorService().submit(task)); + } + + /** + * Submit a task whose result owns a lifecycle and hand back an {@link AsyncResource} that manages it: closing + * the returned resource closes the result, a result produced after a cancel/close race is closed rather than leaked, + * and closing it before completion cancels the task. + * + *

This is the managed counterpart of {@link #submitUnmanagedAsyncResource}; use that when the task's result is a + * plain value with no lifecycle. + * + * @see AsyncResources#fromFutureCloseable(ListenableFuture) */ - @Nullable - public Semaphore getPermits() + public AsyncResource submitCloseableAsyncResource(Callable task) { - return permits; + return AsyncResources.fromFutureCloseable(getExecutorService().submit(task)); } @LifecycleStop @@ -140,4 +173,127 @@ public void stop() exec.shutdownNow(); } } + + /** + * A {@link ListeningExecutorService} that caps the number of concurrently-running submitted tasks at a semaphore's + * permit count, acquiring a permit (on the worker thread) before each task body runs and releasing it after. Used to + * bound concurrent virtual-storage on-demand loads when the backing executor is an unbounded thread-per-virtual-thread + * executor: the permit count is the concurrency bound, not the thread count, and the wait for a permit parks a virtual + * thread rather than blocking the submitter. + *

+ * The permit wait is interruptible: a task whose future is canceled with {@code mayInterruptIfRunning} (or a + * {@code shutdownNow}) while it is parked on the permit is interrupted out of the wait and aborts before running its + * body. Canceling a query stops not only its queued column downloads but also those blocked on the permit, + * before any deep-storage I/O begins. A task that has already passed the permit and started its body runs to + * completion (aborting in-flight reads is handled separately by the task itself, not here). + *

+ * Only {@code execute}/{@code submit} are bounded; those are the only submission paths on-demand load work uses + * (including the {@link PartialBundleAcquirer#submitDownload} column-download path). {@code invokeAll}/ + * {@code invokeAny} throw {@link UnsupportedOperationException} so the concurrency bound can never be silently + * bypassed by a future caller. + *

+ * Callers must not submit a task that itself blocks on another task submitted to this executor while holding a + * permit, or all permits could be exhausted by waiters; the on-demand load tasks here never nest submissions. + */ + @VisibleForTesting + static final class PermitBoundedListeningExecutorService extends ForwardingListeningExecutorService + { + private final ListeningExecutorService delegate; + private final Semaphore permits; + + PermitBoundedListeningExecutorService(ListeningExecutorService delegate, Semaphore permits) + { + this.delegate = delegate; + this.permits = permits; + } + + @Override + protected ListeningExecutorService delegate() + { + return delegate; + } + + @Override + public void execute(Runnable command) + { + delegate.execute(withPermit(command)); + } + + @Override + public ListenableFuture submit(Callable task) + { + return delegate.submit(withPermit(task)); + } + + @Override + public ListenableFuture submit(Runnable task) + { + return delegate.submit(withPermit(task)); + } + + @Override + public ListenableFuture submit(Runnable task, T result) + { + return delegate.submit(withPermit(task), result); + } + + @Override + public List> invokeAll(Collection> tasks) + { + throw new UnsupportedOperationException("invokeAll is not permit-bounded; use submit"); + } + + @Override + public List> invokeAll(Collection> tasks, long timeout, TimeUnit unit) + { + throw new UnsupportedOperationException("invokeAll is not permit-bounded; use submit"); + } + + @Override + public T invokeAny(Collection> tasks) + { + throw new UnsupportedOperationException("invokeAny is not permit-bounded; use submit"); + } + + @Override + public T invokeAny(Collection> tasks, long timeout, TimeUnit unit) + { + throw new UnsupportedOperationException("invokeAny is not permit-bounded; use submit"); + } + + private Runnable withPermit(Runnable task) + { + return () -> { + try { + permits.acquire(); + } + catch (InterruptedException e) { + // Interrupted while waiting for a permit (e.g. the task's future was canceled with mayInterruptIfRunning). + // Abort before running the body. Restore the flag and surface as a failure rather than a silent success; + // if the future was canceled it already reports as such, so this only matters for a stray interrupt. + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + try { + task.run(); + } + finally { + permits.release(); + } + }; + } + + private Callable withPermit(Callable task) + { + return () -> { + permits.acquire(); + try { + return task.call(); + } + finally { + permits.release(); + } + }; + } + } } 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 67f23c2545b3..fbc756f8b0b9 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 @@ -478,8 +478,12 @@ public void adjustReservation(CacheEntryIdentifier id, long newSize) currSizeBytes.getAndAdd(-delta); if (weak == null) { currStaticSizeBytes.getAndAdd(-delta); + // The reservation (loadBegin) was recorded at the pre-shrink size; correct its byte total to match. + staticStats.getAndUpdate(s -> s.shrinkLoadBegin(delta)); } else { currWeakSizeBytes.getAndAdd(-delta); + // The reservation (loadBegin) was recorded at the pre-shrink size; correct its byte total to match. + weakStats.getAndUpdate(s -> s.shrinkLoadBegin(delta)); // Each active hold contributed entry.getSize() to currHoldBytes via trackWeakHold; shrink each hold's // contribution by the same delta so a future trackWeakRelease (which subtracts the new smaller size) lands // on the correct total. Clamp at 0 defensively against any pre-existing drift. @@ -515,6 +519,45 @@ public void release(CacheEntry entry) } } + /** + * Remove a {@link #weakCacheEntries} entry that currently has no outstanding holds, unlinking it from the SIEVE + * queue and terminating its phaser (which fires the underlying {@link CacheEntry#unmount}). No-op when the entry + * is absent, is a {@link #staticCacheEntries} entry, or still has outstanding holds. + *

+ * This exists for callers that register a weak entry without a {@link ReservationHold} (the bootstrap + * reserve path uses {@link #reserveWeak}) and need to clean it up after a failed mount. The normal runtime path + * registers weak entries via {@link #addWeakReservationHold} and relies on the hold's release runnable to evict a + * never-mounted entry on close; an entry registered without a hold has no such cleanup, so a failed mount would + * otherwise leave it lingering (and un-re-mountable if its on-disk state was deleted by the mount rollback). The + * hold guard makes this safe to call unconditionally on any mount failure: a held entry (runtime path) is left to + * its holder's release runnable. + */ + public void removeUnheldWeakEntry(CacheEntryIdentifier id) + { + lock.writeLock().lock(); + try { + weakCacheEntries.computeIfPresent( + 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) + if (isMounted) { + weakStats.getAndUpdate(s -> s.evict(weakCacheEntry.cacheEntry.getSize())); + } + return null; + } + ); + } + finally { + lock.writeLock().unlock(); + } + } + /** * Creates a release runnable for a {@link WeakCacheEntry} that handles immediate eviction when configured. * If {@link #areWeakEntriesEphemeral} is true and there are no more holds after releasing, the entry is immediately @@ -1109,6 +1152,16 @@ public StaticStats loadBegin(long size) return this; } + /** + * Correct the reserved (loadBegin) byte total downward by {@code delta} when a reservation is shrunk via + * {@link StorageLocation#adjustReservation}. The load-begin count is unchanged; only the byte total is corrected. + */ + public StaticStats shrinkLoadBegin(long delta) + { + loadBeginBytes.getAndAdd(-delta); + return this; + } + public StaticStats load(long size) { loadCount.getAndIncrement(); @@ -1203,6 +1256,17 @@ public WeakStats loadBegin(long size) return this; } + /** + * Correct the reserved (loadBegin) byte total downward by {@code delta} when a reservation is shrunk via + * {@link StorageLocation#adjustReservation} (e.g. a partial metadata entry's pessimistic estimate reduced to the + * actual header size after mount). The load-begin count is unchanged; only the byte total is corrected. + */ + public WeakStats shrinkLoadBegin(long delta) + { + loadBeginBytes.getAndAdd(-delta); + return this; + } + public WeakStats load(long size) { loadCount.getAndIncrement(); diff --git a/server/src/main/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManager.java b/server/src/main/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManager.java index d8a351cbc826..c9c22591cffd 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManager.java +++ b/server/src/main/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManager.java @@ -19,7 +19,6 @@ package org.apache.druid.segment.loading.external; -import com.google.common.util.concurrent.ListenableFuture; import com.google.inject.Inject; import org.apache.druid.collections.ResourceHolder; import org.apache.druid.common.asyncresource.AsyncResource; @@ -38,7 +37,6 @@ import java.util.Iterator; import java.util.List; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.Semaphore; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.LongSupplier; @@ -233,34 +231,9 @@ public AsyncResource reserveAndPopulateAsync( } if (loadingThreadPool.isAvailable()) { - final SettableAsyncResource resource = new SettableAsyncResource<>(); - final ListenableFuture future = loadingThreadPool.getExecutorService().submit( - () -> { - try { - final Semaphore loadingPermits = loadingThreadPool.getPermits(); - if (loadingPermits != null) { - loadingPermits.acquire(); - } - try { - final CachedFile theCachedFile = reserveAndPopulate(identifier, sizeSupplier, populator); - if (!resource.set(ResourceHolder.fromCloseable(theCachedFile))) { - theCachedFile.close(); - } - } - finally { - if (loadingPermits != null) { - loadingPermits.release(); - } - } - } - catch (Throwable e) { - resource.setException(e); - } - } + return loadingThreadPool.submitCloseableAsyncResource( + () -> reserveAndPopulate(identifier, sizeSupplier, populator) ); - - resource.setCanceler(() -> future.cancel(true)); - return resource; } else { final SettableAsyncResource resource = new SettableAsyncResource<>(); try { diff --git a/server/src/main/java/org/apache/druid/server/SegmentManager.java b/server/src/main/java/org/apache/druid/server/SegmentManager.java index 7c9618044927..a3d58e102510 100644 --- a/server/src/main/java/org/apache/druid/server/SegmentManager.java +++ b/server/src/main/java/org/apache/druid/server/SegmentManager.java @@ -38,6 +38,7 @@ import org.apache.druid.segment.SegmentReference; import org.apache.druid.segment.join.table.IndexedTable; import org.apache.druid.segment.join.table.ReferenceCountedIndexedTableProvider; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.segment.loading.SegmentCacheManager; import org.apache.druid.segment.loading.SegmentLoadingException; @@ -137,12 +138,12 @@ public Optional> getTimeline(Tabl * present in the cache on demand. *

* What this means mechanically, is that for each {@link DataSegmentAndDescriptor} we check if it is already cached - * with {@link #acquireCachedSegment(DataSegment)} to add to {@link LeafSegmentsBundle#cachedSegments}, else if - * {@link #canLoadSegmentOnDemand(DataSegment)} is true it is added to {@link LeafSegmentsBundle#loadableSegments} or - * {@link LeafSegmentsBundle#missingSegments} if not. + * with {@link #acquireCachedSegment(DataSegment, AcquireMode)} to add to {@link LeafSegmentsBundle#cachedSegments}, + * else if {@link #canLoadSegmentOnDemand(DataSegment)} is true it is added to + * {@link LeafSegmentsBundle#loadableSegments} or {@link LeafSegmentsBundle#missingSegments} if not. *

* The segments in {@link LeafSegmentsBundle#loadableSegments} can be retrieved with - * {@link #acquireSegment(DataSegment)} to ensure they are loaded from deep storage. + * {@link #acquireSegment(DataSegment, AcquireMode)} to ensure they are loaded from deep storage. */ public LeafSegmentsBundle getSegmentsBundle( List segments, @@ -162,7 +163,7 @@ public LeafSegmentsBundle getSegmentsBundle( missingSegments.add(segment.getDescriptor()); continue; } - final Optional ref = acquireCachedSegment(dataSegment); + final Optional ref = acquireCachedSegment(dataSegment, AcquireMode.FULL); if (ref.isPresent()) { try { final Optional mapped = segmentMapFunction.apply(ref).map(safetyNet::register); @@ -193,36 +194,44 @@ public LeafSegmentsBundle getSegmentsBundle( } /** - * Returns a {@link Segment} transformed with a {@link SegmentMapFunction}, if it is available in the cache. The - * returned {@link Segment} must be closed when the caller is finished doing segment things. This method will not - * download a {@link DataSegment} if it is not already present in {@link #cacheManager}, use - * {@link #acquireSegment(DataSegment)} instead. + * Returns a {@link Segment}, if it is available in the cache. The returned {@link Segment} must be closed when the + * caller is finished doing segment things. This method will not download a {@link DataSegment} if it is not already + * present in {@link #cacheManager}, use {@link #acquireSegment(DataSegment, AcquireMode)} instead. + *

+ * With {@link AcquireMode#PARTIAL} the returned segment may not be fully loaded, and callers must use async methods + * like {@link org.apache.druid.segment.CursorFactory#makeCursorHolderAsync} to download data on-demand; the + * synchronous {@code makeCursorHolder} will fail if anything is still missing. See {@link AcquireMode}. */ - public Optional acquireCachedSegment(SegmentId segmentId) + public Optional acquireCachedSegment(SegmentId segmentId, AcquireMode acquireMode) { - return cacheManager.acquireCachedSegment(segmentId); + return cacheManager.acquireCachedSegment(segmentId, acquireMode); } /** - * Convenience overload of {@link #acquireCachedSegment(SegmentId)} that accepts a {@link DataSegment}. + * Convenience overload of {@link #acquireCachedSegment(SegmentId, AcquireMode)} that accepts a {@link DataSegment}. */ - public Optional acquireCachedSegment(DataSegment dataSegment) + public Optional acquireCachedSegment(DataSegment dataSegment, AcquireMode acquireMode) { - return acquireCachedSegment(dataSegment.getId()); + return acquireCachedSegment(dataSegment.getId(), acquireMode); } /** - * Returns a {@link AcquireSegmentAction}, where calling {@link AcquireSegmentAction#getSegmentFuture()} will either return - * immediately if the {@link Segment} is in the cache, or possibly try to fetch the segment from deep storage if not. - * The returned {@link Segment}, if present, must be closed when the caller is finished doing segment things. + * Returns a {@link AcquireSegmentAction}, where calling {@link AcquireSegmentAction#getSegmentFuture()} will either + * return immediately if the {@link Segment} is in the cache, or possibly try to fetch the segment from deep storage + * if not. The returned {@link Segment}, if present, must be closed when the caller is finished doing segment things. *

* Calling this method is treated as an intent to acquire and use the segment via resolving the future, and cache * manager implementations will place a hold on this segment until the 'loadCleanup' closer is closed - typically * after resolving the future to acquire the reference to the actual {@link Segment} object. + *

+ * With {@link AcquireMode#PARTIAL} the action resolves to a partial-load capable segment (when the segment supports + * range reads), and callers must use async methods like + * {@link org.apache.druid.segment.CursorFactory#makeCursorHolderAsync} to download data on-demand. See + * {@link AcquireMode}. */ - public AcquireSegmentAction acquireSegment(DataSegment dataSegment) + public AcquireSegmentAction acquireSegment(DataSegment dataSegment, AcquireMode acquireMode) { - return cacheManager.acquireSegment(dataSegment); + return cacheManager.acquireSegment(dataSegment, acquireMode); } /** @@ -335,7 +344,8 @@ private void loadSegmentInternal( ); long numOfRows = 0; - final Optional loadedSegment = cacheManager.acquireCachedSegment(dataSegment.getId()); + final Optional loadedSegment = + cacheManager.acquireCachedSegment(dataSegment.getId(), AcquireMode.FULL); if (loadedSegment.isPresent()) { final Segment segment = loadedSegment.get(); final IndexedTable table = segment.as(IndexedTable.class); @@ -402,7 +412,8 @@ public void dropSegment(final DataSegment dataSegment) if (oldSegmentRef != null) { try (final Closer closer = Closer.create()) { - final Optional oldSegment = cacheManager.acquireCachedSegment(oldSegmentRef.getId()); + final Optional oldSegment = + cacheManager.acquireCachedSegment(oldSegmentRef.getId(), AcquireMode.FULL); long numberOfRows = oldSegment.map(segment -> { closer.register(segment); final RowCountInspector countInspector = segment.as(RowCountInspector.class); diff --git a/server/src/main/java/org/apache/druid/server/ServerManager.java b/server/src/main/java/org/apache/druid/server/ServerManager.java index d8daf88cd341..60e854f0f6c3 100644 --- a/server/src/main/java/org/apache/druid/server/ServerManager.java +++ b/server/src/main/java/org/apache/druid/server/ServerManager.java @@ -70,6 +70,7 @@ import org.apache.druid.segment.SegmentMapFunction; import org.apache.druid.segment.SegmentReference; import org.apache.druid.segment.TimeBoundaryInspector; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.segment.loading.AcquireSegmentResult; import org.apache.druid.segment.loading.VirtualPlaceholderSegment; @@ -86,6 +87,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Optional; +import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicLong; @@ -280,8 +282,8 @@ protected ArrayList getOrLoadBundleSegments( } /** - * Given a list of {@link DataSegmentAndDescriptor}, uses {@link SegmentManager#acquireSegment(DataSegment)} for each - * to obtain a 'reference' to segments in the cache (or load from deep storage if necessary/supported by the + * Given a list of {@link DataSegmentAndDescriptor}, uses {@link SegmentManager#acquireSegment(DataSegment, AcquireMode)} + * for each to obtain a 'reference' to segments in the cache (or load from deep storage if necessary/supported by the * storage layer). *

* For each of these segments, we then apply a {@link SegmentMapFunction} to prepare for processing. The returned @@ -307,7 +309,7 @@ protected LoadSegmentsResult getOrLoadSegmentReferences( if (segment.getDataSegment() == null) { actions.add(safetyNet.register(AcquireSegmentAction.missingSegment())); } else { - actions.add(safetyNet.register(segmentManager.acquireSegment(segment.getDataSegment()))); + actions.add(safetyNet.register(segmentManager.acquireSegment(segment.getDataSegment(), AcquireMode.FULL))); } } } @@ -416,6 +418,8 @@ protected LoadSegmentsResult getOrLoadSegmentReferences( final DruidException toThrow; if (failure instanceof DruidException) { toThrow = (DruidException) failure; + } else if (failure instanceof ExecutionException ee && ee.getCause() instanceof DruidException druidException) { + toThrow = druidException; } else if (timedOut) { toThrow = DruidException.forPersona(DruidException.Persona.USER) .ofCategory(DruidException.Category.TIMEOUT) diff --git a/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java b/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java index 1fed47d38ea6..1db7a1618eb3 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java +++ b/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java @@ -95,13 +95,13 @@ public void drop(DataSegment segment) } @Override - public Optional acquireCachedSegment(SegmentId segmentId) + public Optional acquireCachedSegment(SegmentId segmentId, AcquireMode acquireMode) { throw new UnsupportedOperationException(); } @Override - public AcquireSegmentAction acquireSegment(DataSegment dataSegment) + public AcquireSegmentAction acquireSegment(DataSegment dataSegment, AcquireMode acquireMode) { throw new UnsupportedOperationException(); } diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java index e794c3abcd65..1d94cab8d2f0 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java @@ -43,7 +43,6 @@ 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.DirectoryBackedRangeReader; import org.apache.druid.segment.file.PartialSegmentFileMapperV10; import org.apache.druid.segment.file.SegmentFileBuilder; import org.apache.druid.segment.file.SegmentFileBuilderV10; @@ -64,6 +63,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Set; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.ThreadLocalRandom; @@ -260,6 +260,40 @@ void testMountAcquiresParentHoldsForAggregateBundle() throws IOException aggEntry.unmount(); } + @Test + void testRuntimeAcquireOfProjectionBundleMountsParentBaseFirst() throws IOException + { + // Runtime acquire path (the cursor factory's PartialBundleAcquirer): acquiring a projection bundle must first + // mount its parent __base bundle, since the projection bundle's mount() takes holds/references on its parents and + // fails with "parent entry not registered" otherwise. Unlike the bootstrap restore path, the runtime path reaches + // the projection bundle directly with no __base mounted, so the acquirer is responsible for the ordering. + final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); + final PartialSegmentMetadataCacheEntry metadata = newMetadataEntry(); + Assertions.assertTrue(location.reserve(metadata)); + metadata.mount(location); + + final PartialSegmentBundleCacheEntryIdentifier baseId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME); + final PartialSegmentBundleCacheEntryIdentifier aggId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); + + // Precondition: nothing has mounted __base yet; the acquirer must do it as the projection bundle's parent. + Assertions.assertNull(location.getCacheEntry(baseId), "test setup: __base must not be mounted yet"); + + final Closeable handle = metadata.getBundleAcquirer().acquire(AGG_BUNDLE); + try { + final PartialSegmentBundleCacheEntry agg = location.getCacheEntry(aggId); + final PartialSegmentBundleCacheEntry base = location.getCacheEntry(baseId); + Assertions.assertNotNull(agg, "projection bundle should be registered"); + Assertions.assertTrue(agg.isMounted(), "projection bundle should be mounted"); + Assertions.assertNotNull(base, "parent __base bundle should be auto-registered by the acquirer"); + Assertions.assertTrue(base.isMounted(), "parent __base bundle should be auto-mounted by the acquirer"); + } + finally { + handle.close(); + } + } + @Test void testMountFailsIfMetadataNotRegisteredWithLocation() throws IOException { @@ -487,6 +521,7 @@ void testForBundleAcceptsBundleNameContainingSlash() throws IOException List.of(), new DirectoryBackedRangeReader(deepDir), JSON_MAPPER, + null, ESTIMATE ); Assertions.assertTrue(location.reserve(metadata)); @@ -534,6 +569,7 @@ void testForBundleSpansMainAndExternalContainers() throws IOException List.of(externalName), new DirectoryBackedRangeReader(deepDir), JSON_MAPPER, + null, ESTIMATE ); Assertions.assertTrue(location.reserve(metadata)); @@ -576,6 +612,7 @@ void testForBundleRootOwnsAllUngroupedContainers() throws IOException List.of(), new DirectoryBackedRangeReader(deepDir), JSON_MAPPER, + null, ESTIMATE ); Assertions.assertTrue(location.reserve(metadata)); @@ -593,6 +630,75 @@ void testForBundleRootOwnsAllUngroupedContainers() throws IOException ); } + @Test + void testResolveBundleNameFallsBackToRootWhenRootIsSoleBundle() throws IOException + { + // A segment whose every container defaulted to the root bundle (legacy / never called startFileBundle). + final File deepDir = writeRootOnlySegment(); + final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); + final PartialSegmentMetadataCacheEntry metadata = new PartialSegmentMetadataCacheEntry( + SEGMENT_ID, + cacheDir, + IndexIO.V10_FILE_NAME, + List.of(), + new DirectoryBackedRangeReader(deepDir), + JSON_MAPPER, + null, + ESTIMATE + ); + Assertions.assertTrue(location.reserve(metadata)); + metadata.mount(location); + final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper(); + + Assertions.assertEquals( + Set.of(SegmentFileBuilder.ROOT_BUNDLE_NAME), + PartialSegmentBundleCacheEntry.bundleNames(mapper), + "an untagged segment has only the root bundle" + ); + // A request for the base bundle (or any projection) resolves to the root catch-all, since root is the sole bundle. + Assertions.assertEquals( + SegmentFileBuilder.ROOT_BUNDLE_NAME, + PartialSegmentBundleCacheEntry.resolveBundleName(mapper, Projections.BASE_TABLE_PROJECTION_NAME) + ); + Assertions.assertEquals( + SegmentFileBuilder.ROOT_BUNDLE_NAME, + PartialSegmentBundleCacheEntry.resolveBundleName(mapper, "some_projection") + ); + // A request for the root bundle by name resolves to itself (it exists; no fallback logic needed). + Assertions.assertEquals( + SegmentFileBuilder.ROOT_BUNDLE_NAME, + PartialSegmentBundleCacheEntry.resolveBundleName(mapper, SegmentFileBuilder.ROOT_BUNDLE_NAME) + ); + } + + @Test + void testResolveBundleNameNoFallbackWhenNamedBundlesPresent() throws IOException + { + final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); + final PartialSegmentMetadataCacheEntry metadata = newMetadataEntry(); + Assertions.assertTrue(location.reserve(metadata)); + metadata.mount(location); + final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper(); + + final Set present = PartialSegmentBundleCacheEntry.bundleNames(mapper); + Assertions.assertTrue(present.contains(Projections.BASE_TABLE_PROJECTION_NAME)); + Assertions.assertTrue(present.contains(AGG_BUNDLE)); + Assertions.assertFalse(present.contains(SegmentFileBuilder.ROOT_BUNDLE_NAME), "a tagged segment has no root bundle"); + + // Present bundles resolve to themselves. + Assertions.assertEquals( + Projections.BASE_TABLE_PROJECTION_NAME, + PartialSegmentBundleCacheEntry.resolveBundleName(mapper, Projections.BASE_TABLE_PROJECTION_NAME) + ); + Assertions.assertEquals(AGG_BUNDLE, PartialSegmentBundleCacheEntry.resolveBundleName(mapper, AGG_BUNDLE)); + // An absent request does NOT fall back to root when named bundles are present: the name passes through unchanged + // so the downstream forBundle fails loudly rather than silently serving the wrong data. + Assertions.assertEquals( + "does_not_exist", + PartialSegmentBundleCacheEntry.resolveBundleName(mapper, "does_not_exist") + ); + } + @Test void testMountRollsBackIfEntryNoLongerWeakReservedAtLocation() throws Exception { @@ -783,6 +889,7 @@ private PartialSegmentMetadataCacheEntry newMetadataEntry() List.of(), new DirectoryBackedRangeReader(deepStorageDir), JSON_MAPPER, + null, ESTIMATE ); } diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java index 93776d6a6cd5..d79809f7d7ac 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java @@ -40,7 +40,6 @@ 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.DirectoryBackedRangeReader; import org.apache.druid.segment.file.PartialSegmentFileMapperV10; import org.apache.druid.segment.incremental.IncrementalIndexSchema; import org.apache.druid.segment.projections.Projections; @@ -147,32 +146,25 @@ void testRestoreRebuildsBothEntriesFromDisk() throws IOException primeOnDiskState(); final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); - final PartialSegmentCacheBootstrap.RestoreResult result = PartialSegmentCacheBootstrap.restoreFromDisk( - SEGMENT_ID, - cacheDir, - IndexIO.V10_FILE_NAME, - List.of(), - JSON_MAPPER, - location - ); + final PartialSegmentMetadataCacheEntry metadata = restoreFromDisk(location); - Assertions.assertNotNull(result.getMetadata()); - Assertions.assertTrue(result.getMetadata().isMounted()); + Assertions.assertTrue(metadata.isMounted()); // metadata entry size matches on-disk header size, NOT a pessimistic estimate final long headerSize = new File( cacheDir, IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX ).length(); - Assertions.assertEquals(headerSize, result.getMetadata().getSize()); + Assertions.assertEquals(headerSize, metadata.getSize()); // bundles: should have at least __base; aggregate may or may not exist on disk depending on what we primed - Assertions.assertFalse(result.getBundles().isEmpty()); - final Set bundleNames = result.getBundles().stream() - .map(PartialSegmentBundleCacheEntry::getBundleName) - .collect(Collectors.toSet()); + final List bundles = List.copyOf(metadata.snapshotLinkedBundles()); + Assertions.assertFalse(bundles.isEmpty()); + final Set bundleNames = bundles.stream() + .map(PartialSegmentBundleCacheEntry::getBundleName) + .collect(Collectors.toSet()); Assertions.assertTrue(bundleNames.contains(Projections.BASE_TABLE_PROJECTION_NAME)); Assertions.assertTrue(bundleNames.contains(AGG_BUNDLE)); - for (PartialSegmentBundleCacheEntry bundle : result.getBundles()) { + for (PartialSegmentBundleCacheEntry bundle : bundles) { Assertions.assertTrue(bundle.isMounted(), "bundle " + bundle.getBundleName() + " should be mounted"); } } @@ -183,17 +175,10 @@ void testRestoreEstablishesParentHoldOnBase() throws IOException primeOnDiskState(); final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); - final PartialSegmentCacheBootstrap.RestoreResult result = PartialSegmentCacheBootstrap.restoreFromDisk( - SEGMENT_ID, - cacheDir, - IndexIO.V10_FILE_NAME, - List.of(), - JSON_MAPPER, - location - ); + final PartialSegmentMetadataCacheEntry metadata = restoreFromDisk(location); // aggregate bundle should declare __base as its parent - final PartialSegmentBundleCacheEntry agg = result.getBundles().stream() + final PartialSegmentBundleCacheEntry agg = metadata.snapshotLinkedBundles().stream() .filter(b -> AGG_BUNDLE.equals(b.getBundleName())) .findFirst().orElseThrow(); Assertions.assertEquals(1, agg.getParentEntryIds().size()); @@ -233,19 +218,12 @@ void testRestoreSkipsBundlesWithMissingContainers() throws IOException } final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); - final PartialSegmentCacheBootstrap.RestoreResult result = PartialSegmentCacheBootstrap.restoreFromDisk( - SEGMENT_ID, - cacheDir, - IndexIO.V10_FILE_NAME, - List.of(), - JSON_MAPPER, - location - ); + final PartialSegmentMetadataCacheEntry metadata = restoreFromDisk(location); // base should still be restored; aggregate is skipped because its containers were removed - final Set bundleNames = result.getBundles().stream() - .map(PartialSegmentBundleCacheEntry::getBundleName) - .collect(Collectors.toSet()); + final Set bundleNames = metadata.snapshotLinkedBundles().stream() + .map(PartialSegmentBundleCacheEntry::getBundleName) + .collect(Collectors.toSet()); Assertions.assertTrue(bundleNames.contains(Projections.BASE_TABLE_PROJECTION_NAME)); // Note: base and aggregate may legitimately share a container (small test segment with sub-cap data), in that // case both end up restored. Don't assert absence of aggregate; assert presence of base. @@ -302,19 +280,12 @@ void testRestoreDeletesOrphanedBundleAndSkipsIt() throws IOException } final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); - final PartialSegmentCacheBootstrap.RestoreResult result = PartialSegmentCacheBootstrap.restoreFromDisk( - SEGMENT_ID, - cacheDir, - IndexIO.V10_FILE_NAME, - List.of(), - JSON_MAPPER, - location - ); + final PartialSegmentMetadataCacheEntry metadata = restoreFromDisk(location); // Neither base (containers missing) nor aggregate (orphaned, parent missing) should be restored. - final Set restoredNames = result.getBundles().stream() - .map(PartialSegmentBundleCacheEntry::getBundleName) - .collect(Collectors.toSet()); + final Set restoredNames = metadata.snapshotLinkedBundles().stream() + .map(PartialSegmentBundleCacheEntry::getBundleName) + .collect(Collectors.toSet()); Assertions.assertFalse(restoredNames.contains(AGG_BUNDLE), "orphan must not be restored"); // Aggregate's container files (those exclusively owned by it) must have been deleted by the orphan cleanup. @@ -341,15 +312,8 @@ void testRestoreRollsBackOnBundleReservationFailure() throws IOException final StorageLocation location = new StorageLocation(cacheDir, headerFile.length(), null); Assertions.assertThrows( - DruidException.class, - () -> PartialSegmentCacheBootstrap.restoreFromDisk( - SEGMENT_ID, - cacheDir, - IndexIO.V10_FILE_NAME, - List.of(), - JSON_MAPPER, - location - ) + Throwable.class, + () -> restoreFromDisk(location) ); // rollback must release the metadata reservation and leave no static/weak entries behind @@ -365,18 +329,58 @@ void testRestoreRollsBackOnBundleReservationFailure() throws IOException } @Test - void testRestoreFailsWhenHeaderMissing() + void testMountFailureRemovesLingeringWeakEntry() throws IOException + { + primeOnDiskState(); + final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(SEGMENT_ID); + + // Reserve the metadata entry weakly, exactly as the bootstrap path does (no protecting hold). + final PartialSegmentMetadataCacheEntry metadata = PartialSegmentCacheBootstrap.reserveFromDisk( + SEGMENT_ID, + cacheDir, + IndexIO.V10_FILE_NAME, + List.of(), + JSON_MAPPER, + null, + location + ); + Assertions.assertTrue(location.isWeakReserved(id)); + + // Delete the header so the mount's file-mapper build must re-fetch it from deep storage; the bootstrap + // disk-only range reader throws, failing the mount (the same poison that arises when create() detects header + // corruption and tries to re-fetch). The mount rollback deletes the header, so the entry is now both unmounted + // and un-re-mountable. + final File headerFile = new File( + cacheDir, + IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX + ); + Assertions.assertTrue(headerFile.delete()); + + Assertions.assertThrows(Throwable.class, () -> metadata.mount(location)); + + // The failed mount must not leave the lingering weak entry behind: a later findExistingPartialWithHold would + // otherwise resurrect it and re-mount would fail forever (disk-only reader + deleted header). It must be gone so + // the cold acquire path rebuilds a fresh, deep-storage-capable entry via reservePartial. + Assertions.assertEquals(0, location.getWeakEntryCount(), "failed mount must remove the lingering weak entry"); + Assertions.assertFalse(location.isWeakReserved(id)); + Assertions.assertEquals(0, location.currentSizeBytes(), "failed mount must release the reservation"); + } + + @Test + void testReserveFailsWhenHeaderMissing() { // no priming: cacheDir is empty final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null); Assertions.assertThrows( DruidException.class, - () -> PartialSegmentCacheBootstrap.restoreFromDisk( + () -> PartialSegmentCacheBootstrap.reserveFromDisk( SEGMENT_ID, cacheDir, IndexIO.V10_FILE_NAME, List.of(), JSON_MAPPER, + null, location ) ); @@ -471,6 +475,7 @@ private void primeOnDiskState() throws IOException List.of(), new DirectoryBackedRangeReader(deepStorageDir), JSON_MAPPER, + null, ESTIMATE ); Assertions.assertTrue(seedLocation.reserve(seedMeta)); @@ -503,4 +508,25 @@ private void primeOnDiskState() throws IOException baseHold.close(); } + /** + * Two-step restore helper that mirrors what {@code SegmentLocalCacheManager} does in production: reserve the metadata + * entry via {@link PartialSegmentCacheBootstrap#reserveFromDisk}, then drive {@link PartialSegmentMetadataCacheEntry#mount} + * to trigger the file-mapper build + bundle restore. On a mount failure, {@code mount}'s own rollback removes the + * (unheld) weak entry from the location, so tests can assert on a clean location without any extra cleanup here. + */ + private PartialSegmentMetadataCacheEntry restoreFromDisk(StorageLocation location) throws IOException + { + final PartialSegmentMetadataCacheEntry metadata = PartialSegmentCacheBootstrap.reserveFromDisk( + SEGMENT_ID, + cacheDir, + IndexIO.V10_FILE_NAME, + List.of(), + JSON_MAPPER, + null, + location + ); + metadata.mount(location); + return metadata; + } + } 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 4beaf9ba4ffa..620234e562b0 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 @@ -31,7 +31,6 @@ import org.apache.druid.segment.TestHelper; import org.apache.druid.segment.data.CompressionStrategy; import org.apache.druid.segment.file.CountingRangeReader; -import org.apache.druid.segment.file.DirectoryBackedRangeReader; import org.apache.druid.segment.file.PartialSegmentFileMapperV10; import org.apache.druid.segment.file.SegmentFileBuilderV10; import org.apache.druid.segment.projections.Projections; @@ -45,6 +44,7 @@ import java.io.File; import java.io.IOException; import java.util.List; +import java.util.Optional; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; @@ -229,6 +229,7 @@ void testConstructorRejectsNonPositiveEstimate() List.of(), new DirectoryBackedRangeReader(segmentFile.getParentFile()), JSON_MAPPER, + null, 0 ) ); @@ -257,7 +258,7 @@ void testUnmountDefersHeaderDeleteWhileReferenceHeld() throws Exception ); Assertions.assertTrue(headerFile.exists()); - final Closeable ref = entry.acquireReference(); + final Closeable ref = entry.acquireMetadataReference(); Assertions.assertTrue(entry.isMounted()); entry.unmount(); @@ -283,6 +284,7 @@ void testConcurrentMountIsDeduplicated() throws Exception List.of(), rangeReader, JSON_MAPPER, + null, ESTIMATE ); Assertions.assertTrue(location.reserve(entry)); @@ -352,14 +354,14 @@ void testMountRollsBackIfEntryNoLongerReservedAtLocation() throws Exception } @Test - void testAcquireReferenceBeforeMountThrows() + void testAcquireMetadataReferenceBeforeMountThrows() { final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE); - Assertions.assertThrows(DruidException.class, entry::acquireReference); + Assertions.assertThrows(DruidException.class, entry::acquireMetadataReference); } @Test - void testAcquireReferenceAfterCleanupCompletesThrows() throws Exception + void testAcquireMetadataReferenceAfterCleanupCompletesThrows() throws Exception { final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 4, null); final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE); @@ -367,7 +369,17 @@ void testAcquireReferenceAfterCleanupCompletesThrows() throws Exception entry.mount(location); entry.unmount(); // no references; cleanup runs synchronously Assertions.assertFalse(entry.isMounted()); - Assertions.assertThrows(DruidException.class, entry::acquireReference); + Assertions.assertThrows(DruidException.class, entry::acquireMetadataReference); + } + + @Test + void testAcquireReferenceBeforeMountReturnsEmpty() + { + // Segment-level acquireReference returns empty (not throw) when the entry isn't mounted, matches the + // SegmentCacheEntry contract used by SegmentLocalCacheManager.acquireCachedSegment to skip locations that don't + // have the segment. + final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE); + Assertions.assertEquals(Optional.empty(), entry.acquireReference()); } @Test @@ -402,6 +414,7 @@ private PartialSegmentMetadataCacheEntry newEntry(long estimate) List.of(), new DirectoryBackedRangeReader(segmentFile.getParentFile()), JSON_MAPPER, + null, estimate ); } diff --git a/server/src/test/java/org/apache/druid/segment/loading/PermitBoundedListeningExecutorServiceTest.java b/server/src/test/java/org/apache/druid/segment/loading/PermitBoundedListeningExecutorServiceTest.java new file mode 100644 index 000000000000..bf338e2017f8 --- /dev/null +++ b/server/src/test/java/org/apache/druid/segment/loading/PermitBoundedListeningExecutorServiceTest.java @@ -0,0 +1,106 @@ +/* + * 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.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.ListeningExecutorService; +import com.google.common.util.concurrent.MoreExecutors; +import org.apache.druid.java.util.common.concurrent.Execs; +import org.apache.druid.segment.loading.StorageLoadingThreadPool.PermitBoundedListeningExecutorService; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.util.concurrent.CancellationException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Tests the cancellation behavior of {@link PermitBoundedListeningExecutorService}: a task parked on the permit (but + * not yet running its body) must be abortable via {@code cancel(true)}. This is the "mid-tier" download cancellation + * that lets the on-demand load path stop queued/permit-blocked column downloads before they touch deep storage. + */ +class PermitBoundedListeningExecutorServiceTest +{ + private ListeningExecutorService backing; + + @AfterEach + void tearDown() + { + if (backing != null) { + backing.shutdownNow(); + } + } + + @Test + @Timeout(30) + void testCancelInterruptsTaskWaitingOnPermitBeforeItRuns() throws Exception + { + backing = MoreExecutors.listeningDecorator(Execs.multiThreaded(2, "permit-bounded-test-%d")); + final Semaphore permits = new Semaphore(1); + final PermitBoundedListeningExecutorService exec = + new PermitBoundedListeningExecutorService(backing, permits); + + final CountDownLatch holderRunning = new CountDownLatch(1); + final CountDownLatch releaseHolder = new CountDownLatch(1); + + // Task A takes the only permit and parks in its body, so any other task must wait on the permit to start. + final ListenableFuture holder = exec.submit(() -> { + holderRunning.countDown(); + try { + releaseHolder.await(); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + Assertions.assertTrue(holderRunning.await(10, TimeUnit.SECONDS), "holder task should start and take the permit"); + + // Task B starts on another thread but blocks acquiring the permit (0 available) before its body can run. + final AtomicBoolean bodyRan = new AtomicBoolean(false); + final ListenableFuture blocked = exec.submit(() -> { + bodyRan.set(true); + return null; + }); + + // Wait until B is genuinely parked on the semaphore, not merely submitted. Guarded by @Timeout against a hang. + for (int i = 0; i < 2000 && !permits.hasQueuedThreads(); i++) { + Thread.sleep(5); + } + Assertions.assertTrue(permits.hasQueuedThreads(), "task B should be parked waiting on the permit"); + Assertions.assertEquals(0, permits.availablePermits(), "the holder still owns the only permit"); + + // Cancelling with interruption must abort the permit wait before B's body runs. + Assertions.assertTrue(blocked.cancel(true), "cancel(true) should interrupt the permit-blocked task"); + Assertions.assertTrue(blocked.isCancelled()); + Assertions.assertThrows(CancellationException.class, blocked::get); + Assertions.assertFalse(bodyRan.get(), "permit-blocked task body must not run after cancellation"); + + // An interrupted acquire takes no permit, so cancelling B must not have consumed or leaked one: the holder still + // owns the only permit, and releasing the holder returns it. + Assertions.assertEquals(0, permits.availablePermits()); + releaseHolder.countDown(); + holder.get(10, TimeUnit.SECONDS); + Assertions.assertEquals(1, permits.availablePermits(), "the holder must return its permit on completion"); + } +} diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerConcurrencyTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerConcurrencyTest.java index 5073ddf9b16e..afbcc0bc7507 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerConcurrencyTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerConcurrencyTest.java @@ -940,7 +940,7 @@ public Integer call() throws SegmentLoadingException { final Closer closer = Closer.create(); final AcquireSegmentAction action = closer.register( - segmentManager.acquireSegment(segment) + segmentManager.acquireSegment(segment, AcquireMode.FULL) ); try { final AcquireSegmentResult result = @@ -995,7 +995,7 @@ public Integer call() if (maxDelayBefore > 0) { Thread.sleep(ThreadLocalRandom.current().nextInt(maxDelayBefore)); } - final Optional segmentReference = segmentManager.acquireCachedSegment(segment.getId()).map(closer::register); + final Optional segmentReference = segmentManager.acquireCachedSegment(segment.getId(), AcquireMode.FULL).map(closer::register); if (segmentReference.isPresent()) { RowCountInspector gadget = segmentReference.get().as(RowCountInspector.class); if (maxDelayAfter > 0) { @@ -1033,7 +1033,7 @@ public Integer call() throws SegmentLoadingException { final Closer closer = Closer.create(); final AcquireSegmentAction action = closer.register( - segmentManager.acquireSegment(segment) + segmentManager.acquireSegment(segment, AcquireMode.FULL) ); try { final Future result = action.getSegmentFuture(); 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 new file mode 100644 index 000000000000..d8d37381381c --- /dev/null +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java @@ -0,0 +1,694 @@ +/* + * 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.CursorBuildSpec; +import org.apache.druid.segment.CursorFactory; +import org.apache.druid.segment.CursorHolder; +import org.apache.druid.segment.IndexBuilder; +import org.apache.druid.segment.IndexIO; +import org.apache.druid.segment.IndexSpec; +import org.apache.druid.segment.PartialQueryableIndexSegment; +import org.apache.druid.segment.Segment; +import org.apache.druid.segment.SegmentLazyLoadFailCallback; +import org.apache.druid.segment.TestHelper; +import org.apache.druid.segment.TimeBoundaryInspector; +import org.apache.druid.segment.V10TimeBoundaryInspector; +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.apache.druid.utils.CloseableUtils; +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.Collections; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; + +class SegmentLocalCacheManagerPartialAcquireTest +{ + 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 RowSignature ROW_SIGNATURE = RowSignature.builder() + .add("dim1", ColumnType.STRING) + .add("metric1", ColumnType.LONG) + .build(); + + private static final List PROJECTIONS = Collections.singletonList( + AggregateProjectionSpec.builder(AGG_BUNDLE) + .groupingColumns(new StringDimensionSchema("dim1")) + .aggregators( + new LongSumAggregatorFactory("_metric1_sum", "metric1"), + 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; + private DataSegment partialSegment; + + @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.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); + + final StorageLocationConfig locConfig = new StorageLocationConfig(cacheRoot, 1024L * 1024L * 1024L, null); + final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig() + .setLocations(List.of(locConfig)) + .setVirtualStorage(true) + .setVirtualStoragePartialDownloadsEnabled(true); + final List storageLocations = loaderConfig.toStorageLocations(); + manager = new SegmentLocalCacheManager( + storageLocations, + loaderConfig, + StorageLoadingThreadPool.createFromConfig(loaderConfig), + new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations), + TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT), + jsonMapper + ) + { + // Tripwire: the partial-aware acquire paths should be self-contained; they must not fall through to the + // eager-extract cached lookup (AcquireMode.FULL, the non-partial branch of acquireCachedSegment). + @Override + public Optional acquireCachedSegment(SegmentId segmentId, AcquireMode acquireMode) + { + if (acquireMode == AcquireMode.FULL) { + Assertions.fail("should not fall back to FULL acquireCachedSegment"); + } + return super.acquireCachedSegment(segmentId, acquireMode); + } + }; + + // DataSegment with a LocalLoadSpec pointing at the deep storage directory (unzipped V10 layout). + partialSegment = DataSegment.builder() + .dataSource(SEGMENT_ID.getDataSource()) + .interval(SEGMENT_ID.getInterval()) + .version(SEGMENT_ID.getVersion()) + .shardSpec(NoneShardSpec.instance()) + .loadSpec(Map.of("type", "local", "path", DEEP_STORAGE_DIR.getAbsolutePath())) + .size(0) + .build(); + } + + @AfterEach + void tearDown() + { + // Drop the segment to release any partial cache entries this test mounted (closes file mappers, unmaps containers + // and the header bitmap, deletes on-disk artifacts). Safe to call even when the test never acquired the segment: + // drop is a no-op for missing entries. Then shut down the manager's executors so threads don't leak across tests. + if (manager != null) { + manager.drop(partialSegment); + manager.shutdown(); + } + } + + @Test + void testAcquirePartialSegmentReturnsPartialAwareSegment() throws ExecutionException, InterruptedException, IOException + { + try (AcquireSegmentAction action = manager.acquireSegment(partialSegment, AcquireMode.PARTIAL)) { + final AcquireSegmentResult result = action.getSegmentFuture().get(); + try (Segment segment = result.getReferenceProvider().acquireReference().orElseThrow()) { + Assertions.assertEquals(SEGMENT_ID, segment.getId()); + Assertions.assertInstanceOf(PartialQueryableIndexSegment.class, segment); + + // The time boundary inspector reads from metadata, so requesting bounds should not download any columns. + final TimeBoundaryInspector inspector = segment.as(TimeBoundaryInspector.class); + Assertions.assertNotNull(inspector); + Assertions.assertInstanceOf(V10TimeBoundaryInspector.class, inspector); + Assertions.assertTrue(inspector.isMinMaxExact()); + Assertions.assertEquals(TIME, inspector.getMinTime()); + Assertions.assertEquals(TIME.plusMinutes(3), inspector.getMaxTime()); + } + } + } + + @Test + void testAcquirePartialAsyncCursorHolderProducesWorkingCursor() + throws ExecutionException, InterruptedException, IOException + { + try (AcquireSegmentAction action = manager.acquireSegment(partialSegment, AcquireMode.PARTIAL)) { + final AcquireSegmentResult result = action.getSegmentFuture().get(); + try (Segment segment = result.getReferenceProvider().acquireReference().orElseThrow()) { + final CursorFactory factory = segment.as(CursorFactory.class); + Assertions.assertNotNull(factory); + // Drive the async path; with the manager's executor the future will complete in the background. + try (var asyncHolder = factory.makeCursorHolderAsync(CursorBuildSpec.FULL_SCAN)) { + // Block in the test for readiness via a busy callback latch. + final CountDownLatch ready = new CountDownLatch(1); + asyncHolder.addReadyCallback(ready::countDown); + Assertions.assertTrue( + ready.await(15, TimeUnit.SECONDS), + "async cursor must materialize within the test timeout" + ); + try (var cursorHolder = asyncHolder.release()) { + Assertions.assertNotNull(cursorHolder); + } + } + } + } + } + + @Test + void testPartialAcquireRecordsLoadAndCorrectsReservationOvercount() + throws ExecutionException, InterruptedException, IOException + { + final StorageLocation loc = manager.getLocations().get(0); + final long estimate = new SegmentLoaderConfig().getVirtualStorageMetadataReservationEstimate(); + + // A lazy (PARTIAL) acquire mounts only the metadata entry (downloads the header); bundles stay lazy until a cursor + // is built, so exactly one weak reservation + one weak load is expected. + try (AcquireSegmentAction action = manager.acquireSegment(partialSegment, AcquireMode.PARTIAL)) { + try (Segment ignored = action.getSegmentFuture().get().getReferenceProvider().acquireReference().orElseThrow()) { + final StorageLocation.WeakStats stats = loc.getWeakStats(); + + // The header download is recorded as an actual load (previously partial mounts recorded no load at all). + Assertions.assertEquals(1, stats.getLoadCount(), "metadata header load should be recorded once"); + Assertions.assertTrue(stats.getLoadBytes() > 0, "header load bytes should be > 0"); + Assertions.assertTrue( + stats.getLoadBytes() < estimate, + "the header is far smaller than the pessimistic reservation estimate" + ); + + // The reservation (loadBegin) was corrected from the pessimistic estimate down to the actual header size, so it + // matches the loaded bytes rather than overcounting at ~16MB. + Assertions.assertEquals(1, stats.getLoadBeginCount()); + Assertions.assertEquals( + stats.getLoadBytes(), + stats.getLoadBeginBytes(), + "loadBegin bytes must be corrected to the actual header size, not the reservation estimate" + ); + } + } + } + + @Test + void testSecondAcquireReturnsCachedSegment() throws ExecutionException, InterruptedException, IOException + { + try (AcquireSegmentAction first = manager.acquireSegment(partialSegment, AcquireMode.PARTIAL)) { + try (Segment ignored = first.getSegmentFuture().get().getReferenceProvider().acquireReference().orElseThrow()) { + // entry is registered + mounted + } + } + + // Second acquire should find the existing mounted entry (cached fast path). + final Optional cached = manager.acquireCachedSegment(SEGMENT_ID, AcquireMode.PARTIAL); + try { + Assertions.assertTrue(cached.isPresent(), "second acquire should hit the cached fast path"); + } + finally { + cached.ifPresent(CloseableUtils::closeAndWrapExceptions); + } + } + + @Test + void testAsyncCursorHoldsBundleHoldUntilCursorClose() + throws ExecutionException, InterruptedException, IOException + { + final StorageLocation loc = manager.getLocations().get(0); + final PartialSegmentBundleCacheEntryIdentifier baseBundleId = new PartialSegmentBundleCacheEntryIdentifier( + SEGMENT_ID, + Projections.BASE_TABLE_PROJECTION_NAME + ); + final long initialHoldBytes = loc.getWeakStats().getHoldBytes(); + + try (AcquireSegmentAction action = manager.acquireSegment(partialSegment, AcquireMode.PARTIAL)) { + final AcquireSegmentResult result = action.getSegmentFuture().get(); + // The action holds a SIEVE-protective hold on the metadata entry for its entire lifetime. + Assertions.assertTrue( + loc.getWeakStats().getHoldBytes() > initialHoldBytes, + "metadata storage-location hold must be active for the acquire action's lifetime" + ); + try (Segment segment = result.getReferenceProvider().acquireReference().orElseThrow()) { + try (var asyncHolder = segment.as(CursorFactory.class).makeCursorHolderAsync(CursorBuildSpec.FULL_SCAN)) { + final CountDownLatch ready = new CountDownLatch(1); + asyncHolder.addReadyCallback(ready::countDown); + Assertions.assertTrue(ready.await(15, TimeUnit.SECONDS)); + final long duringCursorHoldBytes; + try (var cursorHolder = asyncHolder.release()) { + duringCursorHoldBytes = loc.getWeakStats().getHoldBytes(); + // Bundle's own SIEVE hold (from bundleAcquirer.acquire) adds to hold-bytes for the cursor's lifetime, + // on top of the metadata hold and the bundle's persistent SIEVE hold on the metadata acquired at mount. + Assertions.assertTrue( + duringCursorHoldBytes > initialHoldBytes, + "cursor lifecycle must add storage-location holds beyond the initial baseline" + ); + Assertions.assertNotNull(cursorHolder); + } + // After cursor close the cursor-owned bundle hold is released. The bundle entry's persistent SIEVE hold on + // the metadata (acquired during bundle.mount) is still alive until the bundle itself is unmounted by SIEVE, + // so hold-bytes drop but don't return all the way to the pre-cursor baseline. + Assertions.assertTrue( + loc.getWeakStats().getHoldBytes() < duringCursorHoldBytes, + "cursor close must release the cursor-owned bundle hold" + ); + } + } + Assertions.assertTrue(loc.isWeakReserved(baseBundleId), "bundle entry remains registered for re-use"); + } + } + + @Test + void testAsyncCursorBuildMountsMatchedProjectionBundle() + throws ExecutionException, InterruptedException, IOException + { + try (AcquireSegmentAction action = manager.acquireSegment(partialSegment, AcquireMode.PARTIAL)) { + final AcquireSegmentResult result = action.getSegmentFuture().get(); + try (Segment segment = result.getReferenceProvider().acquireReference().orElseThrow()) { + // base-table cursor: drives the base bundle to mount via the cache layer + final CursorBuildSpec scanSpec = CursorBuildSpec.FULL_SCAN; + try (var asyncHolder = segment.as(CursorFactory.class).makeCursorHolderAsync(scanSpec)) { + final CountDownLatch ready = new CountDownLatch(1); + asyncHolder.addReadyCallback(ready::countDown); + Assertions.assertTrue(ready.await(15, TimeUnit.SECONDS)); + try (var cursorHolder = asyncHolder.release()) { + Assertions.assertNotNull(cursorHolder); + } + } + } + } + + // Confirm a bundle entry was registered on the storage location for the base projection + final StorageLocation loc = manager.getLocations().get(0); + final PartialSegmentBundleCacheEntryIdentifier baseBundleId = new PartialSegmentBundleCacheEntryIdentifier( + SEGMENT_ID, + Projections.BASE_TABLE_PROJECTION_NAME + ); + Assertions.assertTrue( + loc.isWeakReserved(baseBundleId), + "base bundle should be registered with the storage location after the async cursor build" + ); + } + + @Test + void testAcquireSegmentForcesFullDownloadOnPartialEligible() + throws ExecutionException, InterruptedException, IOException + { + // acquireSegment (eager API) on a partial-eligible segment should route through the partial machinery and + // force-download every internal file so the returned segment supports sync makeCursorHolder. + final StorageLocation loc = manager.getLocations().get(0); + final PartialSegmentBundleCacheEntryIdentifier baseBundleId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME); + final PartialSegmentBundleCacheEntryIdentifier aggBundleId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); + try (AcquireSegmentAction action = manager.acquireSegment(partialSegment, AcquireMode.FULL)) { + final AcquireSegmentResult result = action.getSegmentFuture().get(); + try (Segment segment = result.getReferenceProvider().acquireReference().orElseThrow()) { + Assertions.assertInstanceOf(PartialQueryableIndexSegment.class, segment); + + // Sync makeCursorHolder must succeed, everything was force-downloaded during acquire. + final CursorFactory factory = segment.as(CursorFactory.class); + Assertions.assertNotNull(factory); + try (CursorHolder holder = factory.makeCursorHolder(CursorBuildSpec.FULL_SCAN)) { + Assertions.assertNotNull(holder); + } + } + + // The eager path must download into proper bundle cache entries (reserved + evictable + cleaned-up on drop), + // not write the containers unmanaged through the file mapper. While the action is held, every bundle the + // segment owns is registered + held on the location. + Assertions.assertTrue(loc.isWeakReserved(baseBundleId), "eager download must create the base bundle entry"); + Assertions.assertTrue(loc.isWeakReserved(aggBundleId), "eager download must create the projection bundle entry"); + Assertions.assertTrue(loc.getWeakStats().getHoldBytes() > 0, "the action must hold the bundles it mounted"); + } + + // Confirm the metadata entry exists on the location and reports fully downloaded. + final CacheEntry entry = loc.getCacheEntry(new SegmentCacheEntryIdentifier(SEGMENT_ID)); + Assertions.assertInstanceOf(PartialSegmentMetadataCacheEntry.class, entry); + Assertions.assertTrue( + ((PartialSegmentMetadataCacheEntry) entry).isFullyDownloaded(), + "force-download path must leave the segment fully downloaded after acquire returns" + ); + } + + @Test + void testCachedFullAcquirePinsBundlesAgainstEviction() + throws ExecutionException, InterruptedException, IOException + { + // Regression: a FULL acquire of a partial segment must pin every bundle for the segment's lifetime. The cached + // fast path (acquireCachedSegment FULL -> acquireCachedInternal) used to hand back a metadata-only hold, so a + // bundle SIEVE-evicted under cache pressure mid-query cleared the mapper's downloaded-file set and the sync + // makeCursorHolder then failed with "requires the segment to be fully downloaded". Uses a plain manager so we can + // call acquireCachedSegment(FULL) directly (the shared fixture installs a tripwire that forbids it). + final File plainCacheRoot = new File(perTestTempDir, "plain_cache_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)); + FileUtils.mkdirp(plainCacheRoot); + final StorageLocationConfig locConfig = new StorageLocationConfig(plainCacheRoot, 1024L * 1024L * 1024L, null); + final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig() + .setLocations(List.of(locConfig)) + .setVirtualStorage(true) + .setVirtualStoragePartialDownloadsEnabled(true); + final List storageLocations = loaderConfig.toStorageLocations(); + final SegmentLocalCacheManager plain = new SegmentLocalCacheManager( + storageLocations, + loaderConfig, + StorageLoadingThreadPool.createFromConfig(loaderConfig), + new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations), + TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT), + jsonMapper + ); + final StorageLocation loc = plain.getLocations().get(0); + final PartialSegmentBundleCacheEntryIdentifier baseBundleId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME); + final PartialSegmentBundleCacheEntryIdentifier aggBundleId = + new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE); + try { + // Fully download + mount, then release the eager acquire so the bundles become unheld (resident but evictable) — + // the state a previously-queried segment is in when the next query's acquireCachedSegment(FULL) finds it. + try (AcquireSegmentAction action = plain.acquireSegment(partialSegment, AcquireMode.FULL)) { + action.getSegmentFuture().get(); + } + Assertions.assertTrue(loc.isWeakReserved(baseBundleId), "base bundle resident after eager acquire"); + Assertions.assertTrue(loc.isWeakReserved(aggBundleId), "agg bundle resident after eager acquire"); + + final Segment cached = plain.acquireCachedSegment(SEGMENT_ID, AcquireMode.FULL).orElseThrow( + () -> new AssertionError("cached FULL acquire of a resident partial segment must return a segment") + ); + try { + Assertions.assertInstanceOf(PartialQueryableIndexSegment.class, cached); + // Simulate SIEVE eviction under cache pressure: a bundle held by the FULL reference must be left in place. + loc.removeUnheldWeakEntry(baseBundleId); + loc.removeUnheldWeakEntry(aggBundleId); + Assertions.assertTrue(loc.isWeakReserved(baseBundleId), "FULL acquire must pin the base bundle against eviction"); + Assertions.assertTrue(loc.isWeakReserved(aggBundleId), "FULL acquire must pin the agg bundle against eviction"); + // Still fully resident, so the sync cursor factory succeeds (this is the call that throws without the fix). + try (CursorHolder holder = cached.as(CursorFactory.class).makeCursorHolder(CursorBuildSpec.FULL_SCAN)) { + Assertions.assertNotNull(holder); + } + } + finally { + CloseableUtils.closeAndWrapExceptions(cached); + } + + // Once the FULL reference closes the bundles are unheld again and free to reclaim. Evict child-before-parent: + // the agg bundle holds the base bundle (its transitive parent), so the base stays reserved until the agg is gone. + loc.removeUnheldWeakEntry(aggBundleId); + Assertions.assertFalse(loc.isWeakReserved(aggBundleId), "agg bundle evictable once the FULL reference closes"); + loc.removeUnheldWeakEntry(baseBundleId); + Assertions.assertFalse(loc.isWeakReserved(baseBundleId), "base bundle evictable once the agg bundle is gone"); + } + finally { + plain.drop(partialSegment); + plain.shutdown(); + } + } + + @Test + void testPartialDownloadsDisabledFallsBackToEager() + throws ExecutionException, InterruptedException, IOException + { + // Rebuild the manager with partial downloads disabled. Both partial-aware acquire APIs must fall through to their + // eager counterparts: the returned segment is an eager QueryableIndexSegment (NOT PartialQueryableIndexSegment), + // and no PartialSegmentMetadataCacheEntry is registered on the location. + final StorageLocationConfig locConfig = new StorageLocationConfig(cacheRoot, 1024L * 1024L * 1024L, null); + final SegmentLoaderConfig disabledConfig = new SegmentLoaderConfig() + .setLocations(List.of(locConfig)) + .setVirtualStorage(true) + .setVirtualStoragePartialDownloadsEnabled(false); + final List storageLocations = disabledConfig.toStorageLocations(); + final SegmentLocalCacheManager disabledManager = new SegmentLocalCacheManager( + storageLocations, + disabledConfig, + StorageLoadingThreadPool.createFromConfig(disabledConfig), + new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations), + TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT), + jsonMapper + ); + + try { + try (AcquireSegmentAction action = disabledManager.acquireSegment(partialSegment, AcquireMode.PARTIAL)) { + final AcquireSegmentResult result = action.getSegmentFuture().get(); + try (Segment segment = result.getReferenceProvider().acquireReference().orElseThrow()) { + Assertions.assertEquals(SEGMENT_ID, segment.getId()); + Assertions.assertFalse( + segment instanceof PartialQueryableIndexSegment, + "partial downloads disabled, acquireSegment(PARTIAL) must fall back to FULL and return a non-partial segment" + ); + } + } + + // No partial metadata entry should have been created on the location. + Assertions.assertFalse( + disabledManager.getLocations().get(0).getCacheEntry(new SegmentCacheEntryIdentifier(SEGMENT_ID)) + instanceof PartialSegmentMetadataCacheEntry, + "partial downloads disabled, no PartialSegmentMetadataCacheEntry should be registered" + ); + + // Cached lookup via the partial API also degrades to the eager cached path. + try (Segment cached = disabledManager.acquireCachedSegment(SEGMENT_ID, AcquireMode.PARTIAL).orElseThrow()) { + Assertions.assertEquals(SEGMENT_ID, cached.getId()); + Assertions.assertFalse( + cached instanceof PartialQueryableIndexSegment, + "partial downloads disabled, acquireCachedSegment(PARTIAL) must return the eager cached segment" + ); + } + } + finally { + // Local manager, not picked up by tearDown(); drop + shut down here to release reservations and stop threads. + disabledManager.drop(partialSegment); + disabledManager.shutdown(); + } + } + + @Test + void testGetCachedSegmentsThenBootstrapMountsPartialEntry() + throws IOException, SegmentLoadingException + { + // Simulate process-restart state: prime the partial on-disk layout (header file + sparse-allocated containers) + // and the segment info file BEFORE getCachedSegments runs, then verify the new two-phase contract: + // - getCachedSegments() reserves the metadata entry on the location but doesn't mount it + // - bootstrap(DataSegment) is what triggers the actual mount + bundle restore via polymorphic dispatch + final File partialDir = new File(cacheRoot, SEGMENT_ID.toString()); + FileUtils.mkdirp(partialDir); + primePartialOnDiskState(partialDir); + manager.storeInfoFile(partialSegment); + + final List cached = manager.getCachedSegments(); + Assertions.assertEquals(List.of(partialSegment), cached); + + final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(SEGMENT_ID); + final StorageLocation location = manager.getLocations().get(0); + final CacheEntry reserved = location.getCacheEntry(id); + Assertions.assertInstanceOf( + PartialSegmentMetadataCacheEntry.class, + reserved, + "getCachedSegments must reserve a partial metadata entry on the location" + ); + final PartialSegmentMetadataCacheEntry partial = (PartialSegmentMetadataCacheEntry) reserved; + Assertions.assertFalse( + partial.isMounted(), + "metadata entry should NOT be mounted after getCachedSegments; mount is deferred to bootstrap()" + ); + + // bootstrap() dispatches polymorphically: partial entry -> metadata.mount(location), which cascades into bundle + // restore via PartialSegmentCacheBootstrap.restoreBundlesFromDisk. + manager.bootstrap(partialSegment, SegmentLazyLoadFailCallback.NOOP); + + Assertions.assertTrue(partial.isMounted(), "metadata entry must be mounted after bootstrap()"); + Assertions.assertFalse( + partial.snapshotLinkedBundles().isEmpty(), + "bundle restore must have linked at least one bundle to the metadata entry" + ); + final Set restoredBundles = partial.snapshotLinkedBundles().stream() + .map(PartialSegmentBundleCacheEntry::getBundleName) + .collect(Collectors.toSet()); + Assertions.assertTrue( + restoredBundles.contains(Projections.BASE_TABLE_PROJECTION_NAME), + "base bundle must be restored by bootstrap; got " + restoredBundles + ); + } + + @Test + void testGetCachedSegmentsDeletesPartialLayoutWhenPartialDisabled() throws IOException + { + // Prime a partial on-disk layout + info file as a previous (partial-enabled) run would have left behind. + final File partialDir = new File(cacheRoot, SEGMENT_ID.toString()); + FileUtils.mkdirp(partialDir); + primePartialOnDiskState(partialDir); + manager.storeInfoFile(partialSegment); + Assertions.assertTrue(PartialSegmentCacheBootstrap.isPartialSegmentLayout(partialDir, IndexIO.V10_FILE_NAME)); + + // A manager with partial downloads disabled (operator toggled the flag off) over the same cache dir must reclaim + // the now-unusable partial layout at bootstrap rather than reserving it and failing at query time. + final StorageLocationConfig locConfig = new StorageLocationConfig(cacheRoot, 1024L * 1024L * 1024L, null); + final SegmentLoaderConfig disabledConfig = new SegmentLoaderConfig() + .setLocations(List.of(locConfig)) + .setVirtualStorage(true) + .setVirtualStoragePartialDownloadsEnabled(false); + final List storageLocations = disabledConfig.toStorageLocations(); + final SegmentLocalCacheManager disabledManager = new SegmentLocalCacheManager( + storageLocations, + disabledConfig, + StorageLoadingThreadPool.createFromConfig(disabledConfig), + new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations), + TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT), + jsonMapper + ); + + try { + final List cached = disabledManager.getCachedSegments(); + Assertions.assertFalse( + cached.contains(partialSegment), + "partial-disabled bootstrap must not return a segment whose only cache state is a partial layout" + ); + Assertions.assertFalse( + partialDir.exists(), + "partial-disabled bootstrap must delete the unusable partial layout from disk" + ); + Assertions.assertNull( + disabledManager.getLocations().get(0).getCacheEntry(new SegmentCacheEntryIdentifier(SEGMENT_ID)), + "no cache entry should be reserved for the deleted partial layout" + ); + } + finally { + disabledManager.shutdown(); + } + } + + /** + * Lay down the on-disk artifacts a previous process run would have left behind in the given partial directory: + * a V10 header file and sparse-allocated container files for every container the segment metadata declares. The + * temporary file mapper used to write the header is closed immediately afterward, it does not delete its files + * on close, so the artifacts persist for the test to bootstrap-restore. + */ + private void primePartialOnDiskState(File partialDir) throws IOException + { + try (PartialSegmentFileMapperV10 seed = + PartialSegmentFileMapperV10.create( + new DirectoryBackedRangeReader(DEEP_STORAGE_DIR), + jsonMapper, + partialDir, + IndexIO.V10_FILE_NAME, + List.of() + )) { + final int numContainers = seed.getSegmentFileMetadata().getContainers().size(); + for (int i = 0; i < numContainers; i++) { + seed.initializeContainer(i); + } + } + } +} diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialDropTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialDropTest.java new file mode 100644 index 000000000000..6613b49dc1b4 --- /dev/null +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialDropTest.java @@ -0,0 +1,315 @@ +/* + * 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.ObjectMapper; +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.java.util.common.DateTimes; +import org.apache.druid.java.util.common.FileUtils; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.query.aggregation.CountAggregatorFactory; +import org.apache.druid.query.aggregation.LongSumAggregatorFactory; +import org.apache.druid.segment.IndexBuilder; +import org.apache.druid.segment.IndexIO; +import org.apache.druid.segment.IndexSpec; +import org.apache.druid.segment.TestHelper; +import org.apache.druid.segment.column.ColumnConfig; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.data.CompressionStrategy; +import org.apache.druid.segment.file.PartialSegmentFileMapperV10; +import org.apache.druid.segment.incremental.IncrementalIndexSchema; +import org.apache.druid.segment.projections.Projections; +import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory; +import org.apache.druid.testing.InitializedNullHandlingTest; +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.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.Collections; +import java.util.List; +import java.util.concurrent.ThreadLocalRandom; + +/** + * Manager-level tests for {@code drop} / {@code removeInfoFile} on partial segments. Partial entries (metadata + + * bundles) are weak cache entries; their lifecycle is owned by {@link StorageLocation}'s SIEVE + hold-phaser + * machinery. Coordinator {@code drop} is effectively a no-op for the partial branch — SIEVE drives the actual + * eviction once holds drop and capacity is needed. {@code removeInfoFile} still wires info-file deletion via the + * metadata entry's {@link PartialSegmentMetadataCacheEntry#setOnUnmount} hook so the info entry survives until the + * (now SIEVE-driven) cleanup actually fires. + *

+ * These tests use {@link StorageLocation#setAreWeakEntriesEphemeral} so hold-release triggers immediate eviction, + * a deterministic substitute for SIEVE's capacity-pressure-driven reclaim. + */ +class SegmentLocalCacheManagerPartialDropTest extends InitializedNullHandlingTest +{ + private static final ObjectMapper JSON_MAPPER = TestHelper.makeJsonMapper(); + private static final SegmentId SEGMENT_ID = SegmentId.of("test", Intervals.of("2025/2026"), "v1", 0); + private static final long ESTIMATE = 16 * 1024 * 1024L; + private static final String AGG_BUNDLE = "dim1_metric1_sum"; + + private static final DateTime TIME = DateTimes.of("2025-01-01"); + private static final RowSignature ROW_SIGNATURE = RowSignature.builder() + .add("dim1", ColumnType.STRING) + .add("metric1", ColumnType.LONG) + .build(); + + private static final List PROJECTIONS = Collections.singletonList( + AggregateProjectionSpec.builder(AGG_BUNDLE) + .groupingColumns(new StringDimensionSchema("dim1")) + .aggregators( + new LongSumAggregatorFactory("_metric1_sum", "metric1"), + 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 sharedTempDir; + + private static File deepStorageDir; + + @TempDir + File perTestTempDir; + + private File cacheDir; + private File infoDir; + private SegmentLocalCacheManager manager; + private StorageLocation location; + private DataSegment dataSegment; + + @BeforeAll + static void buildSegment() + { + final File tmp = new File(sharedTempDir, "build_" + ThreadLocalRandom.current().nextInt()); + deepStorageDir = 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(); + } + + @BeforeEach + void setup() throws IOException + { + cacheDir = new File(perTestTempDir, "cache_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)); + FileUtils.mkdirp(cacheDir); + // SegmentLocalCacheManager defaults the info dir to /info_dir when not configured. + infoDir = new File(cacheDir, "info_dir"); + FileUtils.mkdirp(infoDir); + + final StorageLocationConfig locConfig = new StorageLocationConfig(cacheDir, ESTIMATE * 16, null); + final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig().setLocations(List.of(locConfig)); + final List locations = loaderConfig.toStorageLocations(); + manager = new SegmentLocalCacheManager( + locations, + loaderConfig, + StorageLoadingThreadPool.createFromConfig(loaderConfig), + new LeastBytesUsedStorageLocationSelectorStrategy(locations), + TestHelper.getTestIndexIO(JSON_MAPPER, ColumnConfig.DEFAULT), + JSON_MAPPER + ); + location = manager.getLocations().get(0); + // Drive eviction off hold-release so tests don't need to provoke SIEVE via capacity pressure. + location.setAreWeakEntriesEphemeral(true); + dataSegment = DataSegment.builder() + .dataSource(SEGMENT_ID.getDataSource()) + .interval(SEGMENT_ID.getInterval()) + .version(SEGMENT_ID.getVersion()) + .shardSpec(NoneShardSpec.instance()) + .size(0) + .build(); + } + + /** + * Pair of a mounted partial metadata entry and the {@link StorageLocation.ReservationHold} the test holds against + * it; closing the hold (under {@link StorageLocation#setAreWeakEntriesEphemeral} mode) is what triggers SIEVE + * cleanup. + */ + private record HeldMetadata( + PartialSegmentMetadataCacheEntry metadata, + StorageLocation.ReservationHold hold + ) + { + } + + private HeldMetadata registerMountedMetadata() + { + final PartialSegmentMetadataCacheEntry metadata = new PartialSegmentMetadataCacheEntry( + SEGMENT_ID, + cacheDir, + IndexIO.V10_FILE_NAME, + List.of(), + new DirectoryBackedRangeReader(deepStorageDir), + JSON_MAPPER, + null, + ESTIMATE + ); + final StorageLocation.ReservationHold hold = location.addWeakReservationHold( + new SegmentCacheEntryIdentifier(SEGMENT_ID), + () -> metadata + ); + Assertions.assertNotNull(hold); + try { + metadata.mount(location); + } + catch (IOException e) { + throw new RuntimeException(e); + } + return new HeldMetadata(metadata, hold); + } + + private record HeldBundle( + PartialSegmentBundleCacheEntry bundle, + StorageLocation.ReservationHold hold + ) + { + } + + private HeldBundle registerMountedBundle( + PartialSegmentMetadataCacheEntry metadata, + String bundleName, + List parents + ) throws IOException + { + final PartialSegmentBundleCacheEntry bundle = PartialSegmentBundleCacheEntry.forBundle(metadata, bundleName, parents); + final StorageLocation.ReservationHold hold = + location.addWeakReservationHold(bundle.getId(), () -> bundle); + Assertions.assertNotNull(hold); + bundle.mount(location); + return new HeldBundle(bundle, hold); + } + + @Test + void testSieveEvictionCascadesThroughBundlesAndMetadata() throws IOException + { + final HeldMetadata held = registerMountedMetadata(); + final HeldBundle base = registerMountedBundle( + held.metadata, + Projections.BASE_TABLE_PROJECTION_NAME, + List.of() + ); + final HeldBundle agg = registerMountedBundle(held.metadata, AGG_BUNDLE, List.of(base.bundle.getId())); + Assertions.assertEquals(2, held.metadata.snapshotLinkedBundles().size()); + + final File headerFile = new File(cacheDir, IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX); + Assertions.assertTrue(headerFile.exists(), "header file must exist before eviction"); + + // Coordinator drop is a no-op for the partial path — entries linger until SIEVE picks them. + manager.drop(dataSegment); + Assertions.assertTrue(held.metadata.isMounted(), "drop alone must not unmount weak partial metadata"); + Assertions.assertTrue(headerFile.exists(), "drop alone must not delete the header file"); + + // Release dependent (agg → base → metadata) holds in order. Each release triggers ephemeral-mode eviction of + // the corresponding entry, which releases that entry's parent holds and metadata reference; once the metadata + // hold itself is closed and the last reference drops, the file mapper closes and the header is deleted. + agg.hold.close(); + Assertions.assertFalse(agg.bundle.isMounted(), "agg bundle should unmount once its hold releases"); + base.hold.close(); + Assertions.assertFalse(base.bundle.isMounted(), "base bundle should unmount once its hold releases (agg dropped its parent hold)"); + held.hold.close(); + Assertions.assertFalse(held.metadata.isMounted(), "metadata should unmount once its hold releases and all bundle references are gone"); + Assertions.assertFalse(headerFile.exists(), "header file should be deleted as part of metadata cleanup"); + Assertions.assertTrue(held.metadata.snapshotLinkedBundles().isEmpty(), "metadata should have no linked bundles"); + Assertions.assertFalse(location.isWeakReserved(new SegmentCacheEntryIdentifier(SEGMENT_ID))); + } + + @Test + void testRemoveInfoFileDefersInfoDeleteUntilSieveEviction() throws IOException + { + final HeldMetadata held = registerMountedMetadata(); + final HeldBundle base = registerMountedBundle(held.metadata, Projections.BASE_TABLE_PROJECTION_NAME, List.of()); + + // Simulate the info file that the manager writes when a segment is loaded. + final File infoFile = new File(infoDir, SEGMENT_ID.toString()); + Assertions.assertTrue(infoFile.createNewFile()); + + manager.removeInfoFile(dataSegment); + Assertions.assertTrue( + infoFile.exists(), + "removeInfoFile must defer the actual delete until the partial entry's cleanup runs" + ); + + // Cascading release triggers ephemeral-mode eviction; the onUnmount hook wired by removeInfoFile fires as part + // of metadata cleanup. + base.hold.close(); + held.hold.close(); + Assertions.assertFalse(held.metadata.isMounted()); + Assertions.assertFalse(infoFile.exists(), "info file should be deleted as part of metadata cleanup"); + } + + @Test + void testRemoveInfoFileWithNoCachedEntryDeletesImmediately() throws IOException + { + // No registered entry, manager's removeInfoFile fallback should delete on the spot. + final File infoFile = new File(infoDir, SEGMENT_ID.toString()); + Assertions.assertTrue(infoFile.createNewFile()); + + manager.removeInfoFile(dataSegment); + Assertions.assertFalse(infoFile.exists(), "with no cache entry, removeInfoFile should delete immediately"); + } + + @Test + void testDropWithNoEntryIsNoOp() + { + // Nothing registered for this segment; drop should not throw. + manager.drop(dataSegment); + } +} diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerTest.java index c6b7e61f12a0..d6f96739fc20 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerTest.java @@ -849,7 +849,7 @@ public void testGetBootstrapSegment() throws SegmentLoadingException final DataSegment dataSegment = TestSegmentUtils.makeSegment("foo", "v1", Intervals.of("2020/2021")); manager.bootstrap(dataSegment, SegmentLazyLoadFailCallback.NOOP); - Segment actualBootstrapSegment = manager.acquireCachedSegment(dataSegment.getId()).orElse(null); + Segment actualBootstrapSegment = manager.acquireCachedSegment(dataSegment.getId(), AcquireMode.FULL).orElse(null); Assert.assertNotNull(actualBootstrapSegment); Assert.assertEquals(dataSegment.getId(), actualBootstrapSegment.getId()); Assert.assertEquals(dataSegment.getInterval(), actualBootstrapSegment.getDataInterval()); @@ -887,7 +887,7 @@ public List getLocations() final DataSegment dataSegment = TestSegmentUtils.makeSegment("foo", "v1", Intervals.of("2020/2021")); manager.bootstrap(dataSegment, () -> {}); - Segment actualBootstrapSegment = manager.acquireCachedSegment(dataSegment.getId()).orElse(null); + Segment actualBootstrapSegment = manager.acquireCachedSegment(dataSegment.getId(), AcquireMode.FULL).orElse(null); Assert.assertNotNull(actualBootstrapSegment); Assert.assertEquals(dataSegment.getId(), actualBootstrapSegment.getId()); Assert.assertEquals(dataSegment.getInterval(), actualBootstrapSegment.getDataInterval()); @@ -937,8 +937,8 @@ public File getInfoDir() manager.load(segmentToLoad); Assert.assertNull(manager.getSegmentFiles(segmentToLoad)); - Assert.assertFalse(manager.acquireCachedSegment(segmentToLoad.getId()).isPresent()); - AcquireSegmentAction segmentAction = manager.acquireSegment(segmentToLoad); + Assert.assertFalse(manager.acquireCachedSegment(segmentToLoad.getId(), AcquireMode.FULL).isPresent()); + AcquireSegmentAction segmentAction = manager.acquireSegment(segmentToLoad, AcquireMode.FULL); AcquireSegmentResult result = segmentAction.getSegmentFuture().get(); Optional theSegment = result.getReferenceProvider().acquireReference(); Assert.assertTrue(theSegment.isPresent()); @@ -953,7 +953,7 @@ public File getInfoDir() Assert.assertNotNull(manager.getSegmentFiles(segmentToLoad)); // can actually load them again because load doesn't really do anything - AcquireSegmentAction segmentActionAfterDrop = manager.acquireSegment(segmentToLoad); + AcquireSegmentAction segmentActionAfterDrop = manager.acquireSegment(segmentToLoad, AcquireMode.FULL); AcquireSegmentResult resultAfterDrop = segmentActionAfterDrop.getSegmentFuture().get(); Optional theSegmentAfterDrop = resultAfterDrop.getReferenceProvider().acquireReference(); Assert.assertTrue(theSegmentAfterDrop.isPresent()); @@ -1053,8 +1053,8 @@ public File getInfoDir() manager.bootstrap(segmentToBootstrap, SegmentLazyLoadFailCallback.NOOP); Assert.assertNull(manager.getSegmentFiles(segmentToBootstrap)); - Assert.assertFalse(manager.acquireCachedSegment(segmentToBootstrap.getId()).isPresent()); - AcquireSegmentAction segmentAction = manager.acquireSegment(segmentToBootstrap); + Assert.assertFalse(manager.acquireCachedSegment(segmentToBootstrap.getId(), AcquireMode.FULL).isPresent()); + AcquireSegmentAction segmentAction = manager.acquireSegment(segmentToBootstrap, AcquireMode.FULL); AcquireSegmentResult result = segmentAction.getSegmentFuture().get(); Optional theSegment = result.getReferenceProvider().acquireReference(); Assert.assertTrue(theSegment.isPresent()); @@ -1071,7 +1071,7 @@ public File getInfoDir() // can actually load them again because bootstrap doesn't really do anything unless the segment is already // present in the cache - AcquireSegmentAction segmentActionAfterDrop = manager.acquireSegment(segmentToBootstrap); + AcquireSegmentAction segmentActionAfterDrop = manager.acquireSegment(segmentToBootstrap, AcquireMode.FULL); AcquireSegmentResult resultAfterDrop = segmentActionAfterDrop.getSegmentFuture().get(); Optional theSegmentAfterDrop = resultAfterDrop.getReferenceProvider().acquireReference(); Assert.assertTrue(theSegmentAfterDrop.isPresent()); @@ -1152,7 +1152,7 @@ public void testGetSegmentAfterDroppedWithNoVirtualStorageEnabled() throws Excep // ensure that if virtual storage is not enabled, we do not download the segment (callers might have a DataSegment // reference which was originally cached and then dropped before attempting to acquire a segment. if virtual storage // is not enabled, this should return a missing segment instead of downloading - AcquireSegmentAction segmentAction = manager.acquireSegment(segmentToLoad); + AcquireSegmentAction segmentAction = manager.acquireSegment(segmentToLoad, AcquireMode.FULL); AcquireSegmentResult result = segmentAction.getSegmentFuture().get(); Optional theSegment = result.getReferenceProvider().acquireReference(); Assert.assertFalse(theSegment.isPresent()); @@ -1205,8 +1205,8 @@ public File getInfoDir() manager.load(segmentToLoad); Assert.assertNull(manager.getSegmentFiles(segmentToLoad)); - Assert.assertFalse(manager.acquireCachedSegment(segmentToLoad.getId()).isPresent()); - AcquireSegmentAction segmentAction = manager.acquireSegment(segmentToLoad); + Assert.assertFalse(manager.acquireCachedSegment(segmentToLoad.getId(), AcquireMode.FULL).isPresent()); + AcquireSegmentAction segmentAction = manager.acquireSegment(segmentToLoad, AcquireMode.FULL); // now drop it before we actually load it, but dropping a weakly held reference does not remove the entry from the // cache, deferring it until eviction @@ -1214,7 +1214,7 @@ public File getInfoDir() // however, we also have a hold, so it will not be evicted final DataSegment cannotLoad = makeTestDataSegment(segmentDeepStorageDir, 1, TEST_DATA_RELATIVE_PATH_2); - Assert.assertThrows(DruidException.class, () -> manager.acquireSegment(cannotLoad)); + Assert.assertThrows(DruidException.class, () -> manager.acquireSegment(cannotLoad, AcquireMode.FULL)); // and we can still mount and use the segment we are holding AcquireSegmentResult result = segmentAction.getSegmentFuture().get(); @@ -1231,7 +1231,7 @@ public File getInfoDir() // now that the hold has been released, we can load the other segment and evict the one that was held createSegmentZipInLocation(segmentDeepStorageDir, TEST_DATA_RELATIVE_PATH_2); manager.load(cannotLoad); - AcquireSegmentAction segmentActionAfterDrop = manager.acquireSegment(cannotLoad); + AcquireSegmentAction segmentActionAfterDrop = manager.acquireSegment(cannotLoad, AcquireMode.FULL); AcquireSegmentResult resultDrop = segmentActionAfterDrop.getSegmentFuture().get(); Optional theSegmentAfterDrop = resultDrop.getReferenceProvider().acquireReference(); Assert.assertTrue(theSegmentAfterDrop.isPresent()); @@ -1289,7 +1289,7 @@ public File getInfoDir() createSegmentZipInLocation(segmentDeepStorageDir, TEST_DATA_RELATIVE_PATH); // Acquire the segment (load() is not allowed with evictImmediately) - AcquireSegmentAction segmentAction = manager.acquireSegment(segmentToLoad); + AcquireSegmentAction segmentAction = manager.acquireSegment(segmentToLoad, AcquireMode.FULL); AcquireSegmentResult result = segmentAction.getSegmentFuture().get(); Optional theSegment = result.getReferenceProvider().acquireReference(); Assert.assertTrue(theSegment.isPresent()); @@ -1316,7 +1316,7 @@ public File getInfoDir() Assert.assertFalse(infoFile.exists()); // Verify the segment can be loaded again if needed - AcquireSegmentAction segmentActionAfterEvict = manager.acquireSegment(segmentToLoad); + AcquireSegmentAction segmentActionAfterEvict = manager.acquireSegment(segmentToLoad, AcquireMode.FULL); AcquireSegmentResult resultAfterEvict = segmentActionAfterEvict.getSegmentFuture().get(); Optional theSegmentAfterEvict = resultAfterEvict.getReferenceProvider().acquireReference(); Assert.assertTrue(theSegmentAfterEvict.isPresent()); @@ -1374,7 +1374,7 @@ public void testGetTombstoneSegment() throws SegmentLoadingException .build(); manager.load(tombstone); - Segment segment = manager.acquireCachedSegment(tombstone.getId()).orElse(null); + Segment segment = manager.acquireCachedSegment(tombstone.getId(), AcquireMode.FULL).orElse(null); Assert.assertEquals(tombstone.getId(), segment.getId()); Assert.assertEquals(interval, segment.getDataInterval()); diff --git a/server/src/test/java/org/apache/druid/segment/loading/StorageLocationTest.java b/server/src/test/java/org/apache/druid/segment/loading/StorageLocationTest.java index 2a0d8d858478..433a5b414bde 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/StorageLocationTest.java +++ b/server/src/test/java/org/apache/druid/segment/loading/StorageLocationTest.java @@ -416,6 +416,36 @@ public void testReclaimRestoreDoesNotCreateZombieEntries() hold2.close(); } + @Test + public void testRemoveUnheldWeakEntry() + { + final StorageLocation location = new StorageLocation(tempDir, 100L, null); + final UnmountTrackingCacheEntry entry = new UnmountTrackingCacheEntry("a", 30); + + // an unheld weak entry is removed: unlinked from the queue, unmounted, and its size reclaimed + Assertions.assertTrue(location.reserveWeak(entry)); + Assertions.assertTrue(location.isWeakReserved(entry.getId())); + Assertions.assertEquals(30, location.currentSizeBytes()); + + location.removeUnheldWeakEntry(entry.getId()); + Assertions.assertFalse(location.isWeakReserved(entry.getId())); + Assertions.assertEquals(0, location.getWeakEntryCount()); + Assertions.assertEquals(0, location.currentSizeBytes()); + Assertions.assertTrue(entry.unmountCalled, "removeUnheldWeakEntry must unmount the entry"); + + // absent id is a no-op + location.removeUnheldWeakEntry(new StringCacheIdentifier("missing")); + + // a held weak entry is left in place: its holder's release runnable is responsible for cleanup + final UnmountTrackingCacheEntry held = new UnmountTrackingCacheEntry("b", 40); + final StorageLocation.ReservationHold hold = location.addWeakReservationHold(held.getId(), () -> held); + Assertions.assertNotNull(hold); + location.removeUnheldWeakEntry(held.getId()); + Assertions.assertTrue(location.isWeakReserved(held.getId()), "a held weak entry must not be removed"); + Assertions.assertFalse(held.unmountCalled, "a held weak entry must not be unmounted"); + hold.close(); + } + @Test public void testAdjustReservationStaticEntry() { @@ -681,6 +711,54 @@ public void unmount() } } + /** + * A {@link CacheEntry} that tracks mount/unmount so tests can assert that lifecycle hooks fired. + */ + private static final class UnmountTrackingCacheEntry implements CacheEntry + { + private final StringCacheIdentifier id; + private final long size; + private boolean mounted = false; + private boolean unmountCalled = false; + + private UnmountTrackingCacheEntry(String id, long size) + { + this.id = new StringCacheIdentifier(id); + this.size = size; + } + + @Override + public StringCacheIdentifier getId() + { + return id; + } + + @Override + public long getSize() + { + return size; + } + + @Override + public boolean isMounted() + { + return mounted; + } + + @Override + public void mount(StorageLocation location) + { + mounted = true; + } + + @Override + public void unmount() + { + unmountCalled = true; + mounted = false; + } + } + private static final class TestResizableCacheEntry implements ResizableCacheEntry { private final StringCacheIdentifier id; diff --git a/server/src/test/java/org/apache/druid/server/SegmentManagerTest.java b/server/src/test/java/org/apache/druid/server/SegmentManagerTest.java index 3b0d6c6e9db6..ea0412709a95 100644 --- a/server/src/test/java/org/apache/druid/server/SegmentManagerTest.java +++ b/server/src/test/java/org/apache/druid/server/SegmentManagerTest.java @@ -41,6 +41,7 @@ import org.apache.druid.segment.TestHelper; import org.apache.druid.segment.TestIndex; import org.apache.druid.segment.TestSegmentUtils; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.segment.loading.AcquireSegmentResult; import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStrategy; @@ -489,7 +490,7 @@ public void testGetSegmentsBundleVirtual() throws IOException, ExecutionExceptio ) ); - final AcquireSegmentAction action = virtualSegmentManager.acquireSegment(toLoad); + final AcquireSegmentAction action = virtualSegmentManager.acquireSegment(toLoad, AcquireMode.FULL); AcquireSegmentResult result = action.getSegmentFuture().get(); Assert.assertNotNull(result); Assert.assertEquals(1L, result.getLoadSizeBytes()); diff --git a/server/src/test/java/org/apache/druid/test/utils/TestSegmentCacheManager.java b/server/src/test/java/org/apache/druid/test/utils/TestSegmentCacheManager.java index fe135005d19d..06059e6cd8d5 100644 --- a/server/src/test/java/org/apache/druid/test/utils/TestSegmentCacheManager.java +++ b/server/src/test/java/org/apache/druid/test/utils/TestSegmentCacheManager.java @@ -27,6 +27,7 @@ import org.apache.druid.segment.Segment; import org.apache.druid.segment.SegmentLazyLoadFailCallback; import org.apache.druid.segment.TestSegmentUtils; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.AcquireSegmentAction; import org.apache.druid.segment.loading.AcquireSegmentResult; import org.apache.druid.segment.loading.NoopSegmentCacheManager; @@ -141,7 +142,7 @@ private ReferenceCountedSegmentProvider getSegmentInternal(final DataSegment seg } @Override - public Optional acquireCachedSegment(SegmentId segmentId) + public Optional acquireCachedSegment(SegmentId segmentId, AcquireMode acquireMode) { if (observedSegmentsRemovedFromCache.contains(segmentId)) { return Optional.empty(); @@ -154,7 +155,7 @@ public Optional acquireCachedSegment(SegmentId segmentId) } @Override - public AcquireSegmentAction acquireSegment(DataSegment dataSegment) + public AcquireSegmentAction acquireSegment(DataSegment dataSegment, AcquireMode acquireMode) { if (observedSegmentsRemovedFromCache.contains(dataSegment.getId())) { return AcquireSegmentAction.missingSegment(); diff --git a/server/src/test/java/org/apache/druid/test/utils/TestSegmentManager.java b/server/src/test/java/org/apache/druid/test/utils/TestSegmentManager.java index 771e4bb31b61..c0d0a9b63cbb 100644 --- a/server/src/test/java/org/apache/druid/test/utils/TestSegmentManager.java +++ b/server/src/test/java/org/apache/druid/test/utils/TestSegmentManager.java @@ -21,6 +21,7 @@ import org.apache.druid.java.util.common.ISE; import org.apache.druid.segment.Segment; +import org.apache.druid.segment.loading.AcquireMode; import org.apache.druid.segment.loading.SegmentLoadingException; import org.apache.druid.server.SegmentManager; import org.apache.druid.timeline.DataSegment; @@ -104,7 +105,7 @@ public Segment getSegment(final SegmentId segmentId) if (dataSegment == null) { return null; } - return segmentManager.acquireCachedSegment(dataSegment).orElse(null); + return segmentManager.acquireCachedSegment(dataSegment, AcquireMode.FULL).orElse(null); } /**