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 1c25f089871a..83e373080e63 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
@@ -36,6 +36,7 @@
import org.apache.druid.utils.CloseableUtils;
import javax.annotation.Nullable;
+import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
@@ -84,9 +85,10 @@
* instances. The mmap reflects writes through the shared page cache.
*
* State is persisted to disk so that the mapper can be restored after a process restart without re-fetching metadata
- * from deep storage. The raw V10 header bytes are written to a local file, and a compact bitmap file is appended to
- * the end of it to track which internal files have been downloaded (one bit per file, updated after each download). On
- * subsequent calls, the metadata is parsed from the local file instead of range-reading from deep storage.
+ * from deep storage. The local header file holds the raw V10 header bytes followed by a compact bitmap region that
+ * tracks which internal files have been downloaded (one bit per file, updated after each download). Both regions are
+ * written together, so the file's length is fixed for its lifetime at {@code headerSize + ceil(numFiles / 8)} bytes.
+ * On subsequent calls, the metadata is parsed from the local file instead of range-reading from deep storage.
*
* External segment files are supported via child {@link PartialSegmentFileMapperV10} instances, each targeting a
* different file in the segment's storage location.
@@ -120,6 +122,17 @@ public class PartialSegmentFileMapperV10 implements SegmentFileMapper
*/
public static final long DEFAULT_MAX_FETCH_RUN_BYTES = 64L * 1024 * 1024;
+ /**
+ * Detect if {@code localCacheDir} holds a partial-download header file for {@code targetFilename}.
+ */
+ public static boolean isPartialSegmentLayout(@Nullable File localCacheDir, String targetFilename)
+ {
+ if (localCacheDir == null || !localCacheDir.isDirectory()) {
+ return false;
+ }
+ return new File(localCacheDir, targetFilename + METADATA_HEADER_SUFFIX).exists();
+ }
+
/**
* Create (or restore) a lazy mapper for the main segment file with attached external file mappers. If persisted state
* exists locally from a previous session, metadata is read from disk. Otherwise, metadata is fetched from deep
@@ -194,15 +207,16 @@ private static PartialSegmentFileMapperV10 createForFile(
if (headerFile.exists()) {
try {
result = parseHeaderFile(headerFile, jsonMapper);
+ verifyPersistedHeaderLength(headerFile, result);
bitmapBuffer = mmapBitmap(headerFile, result);
}
catch (ClosedByInterruptException e) {
- // The header is fine, an interrupt aborted the mapping (see mapUninterruptibly). Treating this as corruption
- // would delete a valid local header and force a needless re-download, so leave the file alone and unwind.
+ // the header is fine, an interrupt aborted the mapping (see mapUninterruptibly), no need to necessarily delete
+ // let callers determine that
throw e;
}
catch (Exception e) {
- // corrupted file (partial write, truncated bitmap, bad JSON, etc.), delete and re-fetch
+ // corrupted file, delete
result = null;
if (!headerFile.delete()) {
LOG.warn(
@@ -216,10 +230,9 @@ private static PartialSegmentFileMapperV10 createForFile(
}
if (result == null) {
- fetchAndPersistHeader(rangeReader, targetFilename, headerFile);
- result = parseHeaderFile(headerFile, jsonMapper);
+ result = fetchAndPersistHeader(rangeReader, jsonMapper, targetFilename, headerFile);
bitmapBuffer = mmapBitmap(headerFile, result);
- downloadListener.onBytesDownloaded(headerFile.length());
+ downloadListener.onBytesDownloaded(headerFileSize(result));
}
final PartialSegmentFileMapperV10 mapper = new PartialSegmentFileMapperV10(
@@ -235,11 +248,10 @@ private static PartialSegmentFileMapperV10 createForFile(
);
try {
- // bitmap-vs-container repair pre-pass: if the bitmap claims a file is downloaded but its container file is
- // missing on disk, the bitmap is lying (e.g. partial-cache eviction that cleared containers but couldn't
- // atomically clear bits, or external file-system damage). Clear those bits before the restore loop so we don't
- // spuriously sparse-allocate empty containers in the restore loop's ensureContainerInitialized call and treat
- // their files as downloaded.
+ // if the bitmap claims a file is downloaded but its container file is missing on disk, the bitmap is lying
+ // (e.g. partial-cache eviction that cleared containers but couldn't atomically clear bits, or external
+ // file-system damage). Clear those bits before the restore loop so we don't spuriously sparse-allocate empty
+ // containers in the restore loop's ensureContainerInitialized call and treat their files as downloaded.
for (int i = 0; i < mapper.sortedFileNames.size(); i++) {
final int byteIndex = i / 8;
final int bitMask = 1 << (i % 8);
@@ -308,13 +320,11 @@ private static PartialSegmentFileMapperV10 createForFile(
// file names per container index (parallel to metadata.getContainers()) in ascending start-offset order, for
// whole-container bulk download and coalesced range planning (files tile back-to-back within a container, so offset
- // order is a total order). Built once from the immutable metadata.
+ // order is a total order).
private final List> containerFileNames;
- // external file mappers
private final Map externalMappers = new HashMap<>();
- // track which internal files have been downloaded
private final Set downloadedFiles = ConcurrentHashMap.newKeySet();
private final ConcurrentHashMap fileLocks = new ConcurrentHashMap<>();
private final ReentrantLock bitmapLock;
@@ -858,26 +868,18 @@ public void fetchRun(FetchRun run) throws IOException
}
/**
- * Total on-disk size of the header file(s) backing this mapper, summed across the main file and any external file
- * mappers. This is the actual reservation size that should be charged against the local cache once the metadata has
- * been fetched and persisted; callers can compare it against an up-front pessimistic estimate to decide whether to
- * shrink the reservation.
+ * Computed total on-disk size of the header file(s) backing this mapper, summed across the main file and any
+ * external file mappers.
*/
public long getOnDiskHeaderSize()
{
- long total = headerFileSize(localCacheDir, targetFilename);
+ long total = headerFileSize(headerSize, metadata);
for (PartialSegmentFileMapperV10 ext : externalMappers.values()) {
- total += headerFileSize(ext.localCacheDir, ext.targetFilename);
+ total += headerFileSize(ext.headerSize, ext.metadata);
}
return total;
}
- private static long headerFileSize(File dir, String filename)
- {
- final File header = new File(dir, filename + METADATA_HEADER_SUFFIX);
- return header.exists() ? header.length() : 0;
- }
-
/**
* Total bytes downloaded so far across all internal files, including external mappers.
*/
@@ -1204,12 +1206,32 @@ private void markDownloadedInBitmap(String name)
}
/**
- * Fetch the raw V10 header bytes from deep storage and write them to a local file. The bitmap region is not
- * included, it is created by {@link #mmapBitmap} after parsing. The file is parseable by
- * {@link SegmentFileMetadataReader#read(InputStream, ObjectMapper)}.
+ * On-disk footprint of a single header file: the raw V10 header followed by one bit per internal file.
*/
- private static void fetchAndPersistHeader(
+ private static long headerFileSize(long headerSize, SegmentFileMetadata metadata)
+ {
+ return headerSize + numBitmapBytes(metadata);
+ }
+
+ private static long headerFileSize(SegmentFileMetadataReader.Result result)
+ {
+ return headerFileSize(result.getHeaderSize(), result.getMetadata());
+ }
+
+ private static int numBitmapBytes(SegmentFileMetadata metadata)
+ {
+ return (metadata.getFiles().size() + 7) / 8;
+ }
+
+ /**
+ * Fetch the raw V10 header bytes from deep storage and persist them locally at the header file's final length: the
+ * raw header followed by the zeroed bitmap region.
+ *
+ * @return the parsed metadata of the header just persisted
+ */
+ private static SegmentFileMetadataReader.Result fetchAndPersistHeader(
SegmentRangeReader rangeReader,
+ ObjectMapper jsonMapper,
String targetFilename,
File headerFile
) throws IOException
@@ -1243,20 +1265,47 @@ private static void fetchAndPersistHeader(
actualHeaderSize = fixedHeader.length;
}
- // write fixed header + remaining metadata bytes to a local file atomically (write to temp, then rename)
- // to avoid leaving a partial file on disk if the process crashes mid-write
+ // Matches how SegmentFileMetadataReader reports malformed header bytes (bad version, bad lengths): these bytes
+ // cannot be a V10 header at all, so there is nothing to recover by re-reading them.
+ if (remainingBytes < 0) {
+ throw DruidException.defensive(
+ "Header of [%s] declares [%d] metadata bytes",
+ targetFilename,
+ remainingBytes
+ );
+ }
+
+ // Read the header in full: the fixed part already in hand, plus the metadata bytes that follow it.
+ final byte[] rawHeader = new byte[actualHeaderSize + (int) remainingBytes];
+ System.arraycopy(fixedHeader, 0, rawHeader, 0, actualHeaderSize);
+ try (InputStream remainingStream = rangeReader.readRange(targetFilename, actualHeaderSize, remainingBytes)) {
+ ByteStreams.readFully(remainingStream, rawHeader, actualHeaderSize, (int) remainingBytes);
+ }
+
+ final SegmentFileMetadataReader.Result result;
+ try (InputStream headerBytes = new ByteArrayInputStream(rawHeader)) {
+ result = SegmentFileMetadataReader.read(headerBytes, jsonMapper);
+ }
+ if (rawHeader.length != result.getHeaderSize()) {
+ throw DruidException.defensive(
+ "Read [%d] header bytes for [%s] but its metadata describes a [%d] byte header",
+ rawHeader.length,
+ targetFilename,
+ result.getHeaderSize()
+ );
+ }
+
+ // zeroed bitmap region (nothing is downloaded yet).
+ final byte[] emptyBitmap = new byte[numBitmapBytes(result.getMetadata())];
+
+ // writeAtomically fsyncs before the rename, so a crash leaves either no header file or a complete one.
FileUtils.mkdirp(headerFile.getParentFile());
FileUtils.writeAtomically(headerFile, out -> {
- out.write(fixedHeader, 0, actualHeaderSize);
- try (InputStream remainingStream = rangeReader.readRange(
- targetFilename,
- actualHeaderSize,
- remainingBytes
- )) {
- ByteStreams.limit(remainingStream, remainingBytes).transferTo(out);
- }
+ out.write(rawHeader);
+ out.write(emptyBitmap);
return null;
});
+ return result;
}
/**
@@ -1273,27 +1322,46 @@ private static SegmentFileMetadataReader.Result parseHeaderFile(
}
/**
- * Mmap the bitmap region of the header file as read-write. Extends the file if the bitmap region doesn't exist yet.
- * The channel is closed immediately after mapping.
+ * Mmap the bitmap region of the header file as read-write. The channel is closed immediately after mapping.
*/
private static MappedByteBuffer mmapBitmap(
File headerFile,
SegmentFileMetadataReader.Result result
) throws IOException
{
- final int numBitmapBytes = (result.getMetadata().getFiles().size() + 7) / 8;
- final long expectedSize = result.getHeaderSize() + numBitmapBytes;
+ final int numBitmapBytes = numBitmapBytes(result.getMetadata());
return mapUninterruptibly(() -> {
try (RandomAccessFile raf = new RandomAccessFile(headerFile, "rw");
FileChannel channel = raf.getChannel()) {
- if (raf.length() < expectedSize) {
- raf.setLength(expectedSize);
- }
return channel.map(FileChannel.MapMode.READ_WRITE, result.getHeaderSize(), numBitmapBytes);
}
});
}
+ /**
+ * Corruption check for a header file restored from a previous session: its length must be exactly the header plus
+ * one bit per internal file, since {@link #fetchAndPersistHeader} only ever publishes both regions together.
+ */
+ private static void verifyPersistedHeaderLength(File headerFile, SegmentFileMetadataReader.Result result)
+ throws IOException
+ {
+ final long expectedSize = headerFileSize(result);
+ final long actualSize = headerFile.length();
+ if (actualSize != expectedSize) {
+ throw new IOException(
+ StringUtils.format(
+ "Header file[%s] is [%d] bytes on disk but its metadata describes [%d] bytes (header[%d] plus "
+ + "bitmap[%d]); treating it as corrupt",
+ headerFile,
+ actualSize,
+ expectedSize,
+ result.getHeaderSize(),
+ numBitmapBytes(result.getMetadata())
+ )
+ );
+ }
+ }
+
/**
* Establish a memory mapping, shielding it from the calling thread's interrupt status.
*
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 3925976b88b9..88bb179cbc01 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
@@ -34,6 +34,8 @@
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import java.io.File;
import java.io.FileInputStream;
@@ -970,6 +972,141 @@ void testCreateWithExternals() throws IOException
}
}
+ @Test
+ void testIsPartialSegmentLayout() throws IOException
+ {
+ final File segmentFile = buildTestSegment(10, CompressionStrategy.NONE);
+ final File cacheDir = newCacheDir("layout");
+
+ // an empty (or absent, or non-directory) cache dir is not a partial layout
+ Assertions.assertFalse(PartialSegmentFileMapperV10.isPartialSegmentLayout(cacheDir, IndexIO.V10_FILE_NAME));
+ Assertions.assertFalse(PartialSegmentFileMapperV10.isPartialSegmentLayout(null, IndexIO.V10_FILE_NAME));
+ Assertions.assertFalse(
+ PartialSegmentFileMapperV10.isPartialSegmentLayout(new File(tempDir, "nonexistent"), IndexIO.V10_FILE_NAME)
+ );
+
+ // persisting the header is what makes the layout recognizable, container files or not
+ try (PartialSegmentFileMapperV10 mapper =
+ createMapper(new DirectoryBackedRangeReader(segmentFile.getParentFile()), cacheDir)) {
+ Assertions.assertTrue(PartialSegmentFileMapperV10.isPartialSegmentLayout(cacheDir, IndexIO.V10_FILE_NAME));
+ mapper.fetchFiles(Set.of("3"));
+ Assertions.assertTrue(PartialSegmentFileMapperV10.isPartialSegmentLayout(cacheDir, IndexIO.V10_FILE_NAME));
+ // ...and only for the file it was persisted for
+ Assertions.assertFalse(PartialSegmentFileMapperV10.isPartialSegmentLayout(cacheDir, "some-other-file"));
+ }
+ }
+
+ @Test
+ void testHeaderFileLengthMatchesComputedSize() throws IOException
+ {
+ final File segmentFile = buildTestSegment(10, CompressionStrategy.NONE);
+ final File cacheDir = newCacheDir("header_size");
+ final File headerFile = new File(
+ cacheDir,
+ IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX
+ );
+ final DirectoryBackedRangeReader rangeReader = new DirectoryBackedRangeReader(segmentFile.getParentFile());
+
+ final long freshSize;
+ // fresh fetch: the header is published at its final length, bitmap region included
+ try (PartialSegmentFileMapperV10 mapper = createMapper(rangeReader, cacheDir)) {
+ freshSize = mapper.getOnDiskHeaderSize();
+ Assertions.assertEquals(headerFile.length(), freshSize);
+ // the header is published in one write, so nothing intermediate is left behind to be cleaned up later
+ Assertions.assertArrayEquals(new String[]{headerFile.getName()}, cacheDir.list());
+ // downloading files must not change the header's footprint, only bits inside its bitmap region
+ mapper.fetchFiles(Set.of("3"));
+ Assertions.assertEquals(freshSize, mapper.getOnDiskHeaderSize());
+ Assertions.assertEquals(freshSize, headerFile.length());
+ }
+
+ // restore from the local header: same size, so a reservation made against one measurement is never grown by the
+ // other
+ try (PartialSegmentFileMapperV10 restored = createMapper(rangeReader, cacheDir)) {
+ Assertions.assertEquals(freshSize, restored.getOnDiskHeaderSize());
+ Assertions.assertEquals(freshSize, headerFile.length());
+ }
+ }
+
+ @Test
+ void testHeaderFileLengthMatchesComputedSizeWithExternals() throws IOException
+ {
+ final String externalName = "external.segment";
+ final File baseDir = new File(tempDir, "ext_size_base_" + ThreadLocalRandom.current().nextInt());
+ FileUtils.mkdirp(baseDir);
+
+ try (SegmentFileBuilderV10 builder = SegmentFileBuilderV10.create(JSON_MAPPER, baseDir)) {
+ for (int i = 0; i < 5; ++i) {
+ File tmpFile = new File(tempDir, StringUtils.format("ext-size-main-%s.bin", i));
+ Files.write(Ints.toByteArray(i), tmpFile);
+ builder.add(StringUtils.format("%d", i), tmpFile);
+ }
+ SegmentFileBuilder external = builder.getExternalBuilder(externalName);
+ for (int i = 5; i < 10; ++i) {
+ File tmpFile = new File(tempDir, StringUtils.format("ext-size-ext-%s.bin", i));
+ Files.write(Ints.toByteArray(i), tmpFile);
+ external.add(StringUtils.format("%d", i), tmpFile);
+ }
+ }
+
+ final File cacheDir = newCacheDir("ext_size");
+ final DirectoryBackedRangeReader rangeReader = new DirectoryBackedRangeReader(baseDir);
+ final File mainHeader = new File(
+ cacheDir,
+ IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX
+ );
+ final File externalHeader = new File(
+ cacheDir,
+ externalName + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX
+ );
+
+ // the computed size covers the entry point plus every attached external mapper's header
+ try (PartialSegmentFileMapperV10 mapper = createMapperWithExternal(rangeReader, cacheDir, externalName)) {
+ Assertions.assertEquals(mainHeader.length() + externalHeader.length(), mapper.getOnDiskHeaderSize());
+ }
+ }
+
+ @ParameterizedTest(name = "bytesMissing={0}")
+ @ValueSource(ints = {1, 2})
+ void testHeaderShorterThanItsMetadataIsRefetched(int bytesMissing) throws IOException
+ {
+ final File segmentFile = buildTestSegment(10, CompressionStrategy.NONE);
+ final File cacheDir = newCacheDir("short_header_" + bytesMissing);
+ final File headerFile = new File(
+ cacheDir,
+ IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX
+ );
+ final DirectoryBackedRangeReader rangeReader = new DirectoryBackedRangeReader(segmentFile.getParentFile());
+
+ final long expectedSize;
+ try (PartialSegmentFileMapperV10 mapper = createMapper(rangeReader, cacheDir)) {
+ expectedSize = mapper.getOnDiskHeaderSize();
+ // 10 files means a 2 byte bitmap region, and file "3" is bit 3 of the first of those bytes: dropping the last
+ // byte leaves that bit intact on disk, so a mapper that trusted a short file would report it as downloaded
+ Assertions.assertEquals(2, (mapper.getSegmentFileMetadata().getFiles().size() + 7) / 8);
+ mapper.fetchFiles(Set.of("3"));
+ Assertions.assertEquals(4, mapper.getDownloadedBytes());
+ }
+
+ // Shorten the header: the state a crash (or an interrupt) between persisting the header bytes and extending the
+ // file leaves behind, and the shape of a header written before the two were published together. The metadata
+ // still parses, so only the length tells us the file is incomplete. Growing it back instead would mean the same
+ // segment measures two different sizes depending on when it is looked at, and would resurrect whatever bits
+ // happened to survive.
+ try (RandomAccessFile raf = new RandomAccessFile(headerFile, "rw")) {
+ raf.setLength(expectedSize - bytesMissing);
+ }
+
+ try (PartialSegmentFileMapperV10 recovered = createMapper(rangeReader, cacheDir)) {
+ Assertions.assertEquals(expectedSize, recovered.getOnDiskHeaderSize());
+ Assertions.assertEquals(expectedSize, headerFile.length(), "the short header must have been re-fetched");
+ // treated as corrupt and re-fetched, so the download bitmap starts over rather than being partly trusted
+ Assertions.assertEquals(0, recovered.getDownloadedBytes());
+ recovered.fetchFiles(Set.of("3"));
+ Assertions.assertEquals(3, recovered.mapFile("3").getInt());
+ }
+ }
+
@Test
void testCorruptHeaderFileRecovery() 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
deleted file mode 100644
index 197c3154b6a2..000000000000
--- a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java
+++ /dev/null
@@ -1,352 +0,0 @@
-/*
- * 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.error.DruidException;
-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.projections.Projections;
-import org.apache.druid.timeline.SegmentId;
-
-import javax.annotation.Nullable;
-import java.io.File;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Comparator;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
-
-/**
- * Bootstraps partial-segment cache entries from existing on-disk state. Called by the cache manager on historical
- * startup for each segment directory that contains the partial-download layout (`{targetFilename}.header` plus one
- * or more `{targetFilename}.container.NNNNN` files).
- *
- * 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).
- *
- * Two-phase contract:
- *
- * - {@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.
- * - {@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).
- *
- * Dependency inference is delegated to {@link PartialSegmentMetadataCacheEntry#inferBundleDependencies}. A bundle
- * whose inferred dependency isn't itself present on disk is treated as orphaned: its on-disk container files
- * are deleted (via {@link PartialSegmentFileMapperV10#evictContainer}, which also clears the relevant bitmap bits) and
- * the bundle is not restored. The next access through the cache manager acquire path then triggers a clean cold
- * re-fetch, the same fall-back as when the cache manager finds a segment listed in the info directory but missing on
- * disk.
- */
-public final class PartialSegmentCacheBootstrap
-{
- private static final EmittingLogger LOG = new EmittingLogger(PartialSegmentCacheBootstrap.class);
-
- /**
- * 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 (parsing the header from local disk, no fetch) 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 rangeReader the segment's deep-storage range reader, retained for later on-demand fetches
- * @param jsonMapper used by the metadata entry's mount path to parse the header
- * @param storagePool thread pool the async cursor path submits on-demand column downloads to (which bounds
- * 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
- * @param coalesceGapBytes gap tolerance for coalesced range downloads, see
- * {@link PartialSegmentFileMapperV10#fetchFiles}
- * @param maxFetchRunBytes size cap for parallel-fetch range reads, see
- * {@link PartialSegmentFileMapperV10#planParallelFetch}
- * @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 PartialSegmentMetadataCacheEntry reserveFromDisk(
- SegmentId segmentId,
- File localCacheDir,
- String targetFilename,
- List externalFilenames,
- SegmentRangeReader rangeReader,
- ObjectMapper jsonMapper,
- @Nullable StorageLoadingThreadPool storagePool,
- StorageLocation location,
- long coalesceGapBytes,
- long maxFetchRunBytes
- )
- {
- final File headerFile = new File(localCacheDir, targetFilename + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX);
- if (!headerFile.exists()) {
- throw DruidException.defensive(
- "No on-disk header for partial segment[%s] at [%s]; nothing to restore",
- segmentId,
- headerFile
- );
- }
-
- // size the metadata reservation to the actual on-disk size so the location accounting is correct from the start
- final long actualMetadataSize = computeOnDiskHeaderSize(localCacheDir, targetFilename, externalFilenames);
- final PartialSegmentMetadataCacheEntry metadata = new PartialSegmentMetadataCacheEntry(
- segmentId,
- localCacheDir,
- targetFilename,
- externalFilenames,
- rangeReader,
- jsonMapper,
- storagePool,
- actualMetadataSize,
- coalesceGapBytes,
- maxFetchRunBytes
- );
-
- if (!location.reserveWeak(metadata)) {
- throw DruidException.defensive(
- "Failed to reserve metadata entry for partial segment[%s] at location[%s]",
- segmentId,
- location.getPath()
- );
- }
- return metadata;
- }
-
- /**
- * 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;
- }
-
- final SegmentId segmentId = metadata.getSegmentId();
- final File localCacheDir = metadata.getLocalCacheDir();
-
- // 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 dep : metadata.inferBundleDependencies(name)) {
- if (!presentBundleNames.contains(dep.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] (dependency unrestorable); next access "
- + "will trigger cold re-fetch",
- orphanName,
- segmentId
- );
- }
-
- // Mount the base bundle before any dependent bundle so its hold is available when dependents acquire deps.
- mountableBundleNames.sort(Comparator.comparing(name -> !Projections.BASE_TABLE_PROJECTION_NAME.equals(name)));
-
- final List mountedBundles = new ArrayList<>();
- boolean success = false;
- try {
- for (String bundleName : mountableBundleNames) {
- // Mountable bundles have all dependencies present by construction (orphans were filtered out above), so the
- // inferred dependency set is exactly what we want, no further filtering needed.
- final List parentIds = metadata.inferBundleDependencies(bundleName);
- final PartialSegmentBundleCacheEntry bundle = PartialSegmentBundleCacheEntry.forBundle(
- metadata,
- bundleName,
- parentIds
- );
- // weak-reserve with a temporary hold so the mount call's own parent-hold acquisition can succeed; release the
- // bootstrap hold immediately after, if the entry should remain alive for query-side access, the runtime
- // hold chain (transitive parents from aggregates, segment-level holds from acquire APIs) keeps it pinned.
- try (StorageLocation.ReservationHold bootstrapHold =
- location.addWeakReservationHold(bundle.getId(), () -> bundle)) {
- if (bootstrapHold == null) {
- throw DruidException.defensive(
- "Failed to reserve bundle entry[%s] in location[%s] during bootstrap",
- bundle.getId(),
- location.getPath()
- );
- }
- bundle.mount(location);
- }
- mountedBundles.add(bundle);
- }
- success = true;
- LOG.debug(
- "Restored bundles for partial segment[%s] from [%s]: bundles[%s], orphans[%s]",
- segmentId,
- localCacheDir,
- mountableBundleNames,
- orphanedBundleNames
- );
- }
- finally {
- if (!success) {
- // Reverse-dependency rollback for bundles only; the metadata entry's own rollback is the caller's
- // responsibility (it will fire when the propagated throw escapes doMount's try/catch).
- for (PartialSegmentBundleCacheEntry bundle : mountedBundles) {
- try {
- bundle.unmount();
- }
- catch (Throwable t) {
- LOG.warn(t, "Failed to roll back bundle[%s] during bootstrap failure for [%s]", bundle.getId(), segmentId);
- }
- }
- }
- }
- }
-
- /**
- * Check whether a directory looks like a partial-segment cache layout for the given target filename.
- */
- public static boolean isPartialSegmentLayout(File localCacheDir, String targetFilename)
- {
- if (localCacheDir == null || !localCacheDir.isDirectory()) {
- return false;
- }
- final File header = new File(
- localCacheDir,
- targetFilename + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX
- );
- return header.exists();
- }
-
- private static long computeOnDiskHeaderSize(File localCacheDir, String targetFilename, List externalFilenames)
- {
- long total = sizeOf(new File(
- localCacheDir,
- targetFilename + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX
- ));
- for (String external : externalFilenames) {
- total += sizeOf(new File(
- localCacheDir,
- external + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX
- ));
- }
- if (total <= 0) {
- // PartialSegmentMetadataCacheEntry requires a positive reservation; if all headers are zero-length the local
- // layout is degenerate and should not be restored
- throw DruidException.defensive(
- "Zero-sized header files in [%s]; refusing to restore",
- localCacheDir
- );
- }
- return total;
- }
-
- private static long sizeOf(File f)
- {
- return f.exists() ? f.length() : 0;
- }
-
- /**
- * 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
- * V10 filename for refs in the main mapper, or the external filename for refs in an external mapper.
- */
- private static List filterByContainerPresence(
- Set candidateBundleNames,
- PartialSegmentFileMapperV10 fileMapper,
- File localCacheDir
- )
- {
- final List restorable = new ArrayList<>();
- for (String bundleName : candidateBundleNames) {
- final List refs =
- PartialSegmentBundleCacheEntry.findContainersForBundle(fileMapper, bundleName);
- if (refs.isEmpty()) {
- continue;
- }
- boolean allPresent = true;
- for (PartialSegmentBundleCacheEntry.BundleContainerRef ref : refs) {
- final String mapperFilename = fileMapper.mapperForContainer(ref.externalFilename()).getTargetFilename();
- final File cf = new File(
- localCacheDir,
- StringUtils.format("%s.container.%05d", mapperFilename, ref.containerIndex())
- );
- if (!cf.exists()) {
- allPresent = false;
- break;
- }
- }
- if (allPresent) {
- restorable.add(bundleName);
- }
- }
- return restorable;
- }
-
- private PartialSegmentCacheBootstrap()
- {
- // utility class
- }
-
-}
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 ed9bfebc2890..8c4d827fb693 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
@@ -26,6 +26,7 @@
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.StringUtils;
import org.apache.druid.java.util.common.io.Closer;
import org.apache.druid.java.util.emitter.EmittingLogger;
import org.apache.druid.segment.PartialBundleAcquirer;
@@ -48,7 +49,9 @@
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
+import java.util.Comparator;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
@@ -229,15 +232,6 @@ 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;
- }
-
/**
* The fingerprint of the partial-load rule currently applied to this entry, or {@code null} if no rule has been
* applied. Set by {@link #applyRule(String, Set)} and cleared by {@link #clearRule()}. {@link #doActualUnmount()}
@@ -859,7 +853,7 @@ private void doMount(StorageLocation mountLocation) throws IOException
// 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 {
- PartialSegmentCacheBootstrap.restoreBundlesFromDisk(this, mountLocation);
+ restoreBundlesFromDisk(mountLocation);
}
catch (Throwable t) {
try {
@@ -874,15 +868,10 @@ private void doMount(StorageLocation mountLocation) throws IOException
}
}
catch (Throwable t) {
- // 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.
+ // Reclaim the reservation of an entry that is still registered here but no longer held, which the rollbacks
+ // above have just left with a closed mapper and no header on disk. No-op if anything holds this (including
+ // bundle entries). Runs outside entryLock (the inner blocks released it) so the writeLock -> entryLock order
+ // inside removeUnheldWeakEntry is respected.
try {
mountLocation.removeUnheldWeakEntry(id);
}
@@ -893,6 +882,165 @@ private void doMount(StorageLocation mountLocation) throws IOException
}
}
+ /**
+ * Discover, reserve, and mount any bundles whose container files survived on disk for this segment, so a restart (or
+ * an acquire that follows an eviction which left container files behind) doesn't re-fetch bytes that are already
+ * local. Invoked from {@link #doMount} once the file mapper is installed; safe to call unconditionally, on the
+ * fresh-acquire path there are no container files yet and this is a no-op.
+ *
+ * A bundle whose inferred dependency set includes a bundle that is not itself present on disk is treated as
+ * orphaned: restoring it would only produce a degenerate state where column reads that resolve into the
+ * missing parent fail at query time, so its container files are deleted (via
+ * {@link PartialSegmentFileMapperV10#evictContainer}, which also clears the matching bitmap bits) and the bundle is
+ * left unrestored. The next access through the acquire path then triggers a clean cold re-fetch, the same fall-back
+ * as when the cache manager finds a segment in the info directory but missing on disk.
+ *
+ * On any failure, bundles mounted so far are rolled back before the throw propagates. This entry's own rollback is
+ * {@link #doMount}'s responsibility, not this method's.
+ */
+ private void restoreBundlesFromDisk(StorageLocation location) throws IOException
+ {
+ final PartialSegmentFileMapperV10 mapper = getFileMapper();
+ if (mapper == null) {
+ // not mounted yet (or already 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 List presentBundleNames = filterByContainerPresence(
+ PartialSegmentBundleCacheEntry.bundleNames(mapper),
+ mapper
+ );
+ if (presentBundleNames.isEmpty()) {
+ // Fresh acquire path or a partial whose containers were all evicted: nothing to do.
+ return;
+ }
+
+ final List mountableBundleNames = new ArrayList<>();
+ final Set orphanedBundleNames = new HashSet<>();
+ for (String name : presentBundleNames) {
+ boolean orphaned = false;
+ for (PartialSegmentBundleCacheEntryIdentifier dep : inferBundleDependencies(name)) {
+ if (!presentBundleNames.contains(dep.bundleName())) {
+ orphaned = true;
+ break;
+ }
+ }
+ if (orphaned) {
+ orphanedBundleNames.add(name);
+ } else {
+ mountableBundleNames.add(name);
+ }
+ }
+
+ for (String orphanName : orphanedBundleNames) {
+ for (PartialSegmentBundleCacheEntry.BundleContainerRef ref :
+ PartialSegmentBundleCacheEntry.findContainersForBundle(mapper, orphanName)) {
+ mapper.mapperForContainer(ref.externalFilename()).evictContainer(ref.containerIndex());
+ }
+ LOG.debug(
+ "Deleted on-disk state of orphaned bundle[%s] for segment[%s] (dependency unrestorable); next access "
+ + "will trigger cold re-fetch",
+ orphanName,
+ segmentId
+ );
+ }
+
+ // Mount the base bundle before any dependent bundle so its hold is available when dependents acquire deps.
+ mountableBundleNames.sort(Comparator.comparing(name -> !Projections.BASE_TABLE_PROJECTION_NAME.equals(name)));
+
+ final List mountedBundles = new ArrayList<>();
+ boolean success = false;
+ try {
+ for (String bundleName : mountableBundleNames) {
+ // Mountable bundles have all dependencies present by construction (orphans were filtered out above), so the
+ // inferred dependency set is exactly what we want, no further filtering needed.
+ final PartialSegmentBundleCacheEntry bundle = PartialSegmentBundleCacheEntry.forBundle(
+ this,
+ bundleName,
+ inferBundleDependencies(bundleName)
+ );
+ // weak-reserve with a temporary hold so the mount call's own parent-hold acquisition can succeed; release the
+ // restore 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 restoreHold =
+ location.addWeakReservationHold(bundle.getId(), () -> bundle)) {
+ if (restoreHold == null) {
+ throw DruidException.defensive(
+ "Failed to reserve bundle entry[%s] in location[%s] while restoring from disk",
+ bundle.getId(),
+ location.getPath()
+ );
+ }
+ bundle.mount(location);
+ }
+ mountedBundles.add(bundle);
+ }
+ success = true;
+ LOG.debug(
+ "Restored bundles for partial segment[%s] from [%s]: bundles[%s], orphans[%s]",
+ segmentId,
+ localCacheDir,
+ mountableBundleNames,
+ orphanedBundleNames
+ );
+ }
+ finally {
+ if (!success) {
+ // Roll back the bundles, dependents before parents, so a parent is unheld (its dependent's holds released) by
+ // the time we try to remove it.
+ for (int i = mountedBundles.size() - 1; i >= 0; i--) {
+ final PartialSegmentBundleCacheEntry bundle = mountedBundles.get(i);
+ try {
+ bundle.unmount();
+ location.removeUnheldWeakEntry(bundle.getId());
+ }
+ catch (Throwable t) {
+ LOG.warn(t, "Failed to roll back bundle[%s] during restore failure for [%s]", bundle.getId(), segmentId);
+ }
+ }
+ }
+ }
+ }
+
+ /**
+ * The subset of {@code candidateBundleNames} whose every owned container file is present in the local cache
+ * directory. A bundle with no containers at all, or with only some of them on disk, is not restorable.
+ */
+ private List filterByContainerPresence(
+ Set candidateBundleNames,
+ PartialSegmentFileMapperV10 mapper
+ )
+ {
+ final List restorable = new ArrayList<>();
+ for (String bundleName : candidateBundleNames) {
+ final List refs =
+ PartialSegmentBundleCacheEntry.findContainersForBundle(mapper, bundleName);
+ if (refs.isEmpty()) {
+ continue;
+ }
+ boolean allPresent = true;
+ for (PartialSegmentBundleCacheEntry.BundleContainerRef ref : refs) {
+ final String mapperFilename = mapper.mapperForContainer(ref.externalFilename()).getTargetFilename();
+ final File containerFile = new File(
+ localCacheDir,
+ StringUtils.format("%s.container.%05d", mapperFilename, ref.containerIndex())
+ );
+ if (!containerFile.exists()) {
+ allPresent = false;
+ break;
+ }
+ }
+ if (allPresent) {
+ restorable.add(bundleName);
+ }
+ }
+ return restorable;
+ }
+
private static void awaitMount(SettableFuture future) throws IOException
{
try {
@@ -1335,8 +1483,8 @@ public Closeable acquire(String requestedBundleName)
}
/**
- * 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.
+ * Mount the bundle on the supplied location, closing the transient reservation 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,
@@ -1345,13 +1493,10 @@ private void mountBundleOrClose(
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 dependency holds/refs for its lifetime, so we hold these transient ones only across
- // the mount and release them immediately after a successful mount (acquire() is acyclic: base/root have no
- // dependencies, so the recursion terminates).
+ // Mount this bundle's parents 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 bundle keeps its own
+ // dependency holds/refs for its lifetime, so we hold these transient ones only across the mount and release
+ // them immediately after a successful mount.
final Closer parentHolds = Closer.create();
try {
for (PartialSegmentBundleCacheEntryIdentifier depId : inferBundleDependencies(bundleName)) {
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 a0948431e507..d8179e3ed45f 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
@@ -318,7 +318,7 @@ private void addFilesToCachedSegments(
// 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)) {
+ && PartialSegmentFileMapperV10.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
@@ -336,61 +336,19 @@ private void addFilesToCachedSegments(
atomicMoveAndDeleteCacheEntryDirectory(partialDir);
continue;
}
- SegmentRangeReader rangeReader;
- try {
- rangeReader = tryOpenRangeReader(segment);
- }
- catch (Exception e) {
- log.warn(e, "Failed to open a range reader for partial segment[%s] during bootstrap", segment.getId());
- rangeReader = null;
- }
- if (rangeReader == null) {
- // Anomalous: a layout on disk means range reads worked when it was written, so this should not happen (the
- // loadSpec is now non-range-capable, or no longer converts to a known type). Reclaim it and let the segment
- // re-load fresh on next access rather than failing bootstrap or reserving an entry that could never fetch.
- // Leave removeInfo true so it's treated as uncached.
- log.warn(
- "On-disk partial-load layout for segment[%s] in [%s] has no usable range reader (this should not "
- + "happen); deleting it so bootstrap can continue.",
- segment.getId(),
- partialDir
- );
- atomicMoveAndDeleteCacheEntryDirectory(partialDir);
- continue;
- }
+ // Nothing is reserved here: partial entries are reserved (and immediately mounted, which is what sizes them)
+ // one at a time in bootstrap().
removeInfo = false;
- try {
- PartialSegmentCacheBootstrap.reserveFromDisk(
- segment.getId(),
- partialDir,
- IndexIO.V10_FILE_NAME,
- List.of(),
- rangeReader,
- jsonMapper,
- virtualStorageLoadingThreadPool,
- location,
- config.getVirtualStorageCoalesceGapBytes(),
- config.getVirtualStorageMaxFetchRunBytes()
- );
- 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());
- }
+ cachedSegments.add(segment);
// do not fall through to 'complete' path since this was a partial
continue;
}
if (cacheEntry.checkExists(location.getPath())) {
removeInfo = false;
- final boolean reserveResult;
- if (config.isVirtualStorage()) {
- reserveResult = location.reserveWeak(cacheEntry);
- } else {
- reserveResult = location.reserve(cacheEntry);
- }
- if (!reserveResult) {
+ // Under virtual storage nothing is reserved here: as with the partial layout above, bootstrap() reserves this
+ // segment and mounts it under the resulting hold. The legacy path reserves it statically, up front.
+ if (!config.isVirtualStorage() && !location.reserve(cacheEntry)) {
log.makeAlert(
"storage[%s:%,d] has more segments than it is allowed. Currently loading Segment[%s:%,d]. Please increase druid.segmentCache.locations maxSize param",
location.getPath(),
@@ -870,72 +828,113 @@ private SegmentRangeReader tryOpenRangeReader(DataSegment dataSegment)
*/
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);
+ final ReservedPartial reserved = tryReservePartialAt(dataSegment, rangeReader, iterator.next(), true);
+ if (reserved != null) {
+ return reserved;
}
- catch (IOException e) {
- // Location is unwritable, fall through to next location rather than failing the whole reservation
+ }
+ 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()
+ );
+ }
+
+ /**
+ * Reserve a partial metadata entry for {@code dataSegment} on one specific location, returning {@code null} if that
+ * location can't take it (unwritable, or no capacity) so a caller walking locations can try the next one.
+ *
+ * @param writeInfoFile write the segment info file now. False for bootstrap restores, where it already exists and is
+ * what told us about this segment in the first place
+ */
+ @Nullable
+ private ReservedPartial tryReservePartialAt(
+ DataSegment dataSegment,
+ SegmentRangeReader rangeReader,
+ StorageLocation location,
+ boolean writeInfoFile
+ )
+ {
+ final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(dataSegment.getId());
+ final File partialDir = new File(location.getPath(), dataSegment.getId().toString());
+ // A pre-existing directory (e.g. bootstrap) will restore in place
+ final boolean createdPartialDir = !partialDir.exists();
+ try {
+ FileUtils.mkdirp(partialDir);
+ }
+ catch (IOException e) {
+ // Location is unwritable, let the caller fall through to the next location rather than failing the reservation
+ log.warn(
+ e,
+ "Failed to create partial cache dir on location[%s] for segment[%s]",
+ location.getPath(),
+ dataSegment.getId()
+ );
+ return null;
+ }
+ final StorageLocation.ReservationHold hold = location.addWeakReservationHold(
+ id,
+ () -> new PartialSegmentMetadataCacheEntry(
+ dataSegment.getId(),
+ partialDir,
+ IndexIO.V10_FILE_NAME,
+ List.of(),
+ rangeReader,
+ jsonMapper,
+ virtualStorageLoadingThreadPool,
+ config.getVirtualStorageMetadataReservationEstimate(),
+ config.getVirtualStorageCoalesceGapBytes(),
+ config.getVirtualStorageMaxFetchRunBytes()
+ )
+ );
+ if (hold == null) {
+ if (createdPartialDir) {
+ atomicMoveAndDeleteCacheEntryDirectory(partialDir);
+ } else {
+ // A layout is already on disk here but the location has no room for its metadata entry and nothing
+ // reclaimable, which most likely means the location's configured maxSize shrank since the layout was written.
+ // Worth saying out loud either way: bootstrap fails the segment on this, and an on-demand acquire moves on to
+ // another location, leaving the files here with no entry behind them.
log.warn(
- e,
- "Failed to create partial cache dir on location[%s] for segment[%s]; trying next location",
+ "Location[%s] with available bytes[%,d] cannot reserve metadata estimate[%,d] bytes for segment[%s], "
+ + "which already has partial cache state on disk there; check druid.segmentCache.locations maxSize.",
location.getPath(),
+ location.availableSizeBytes(),
+ config.getVirtualStorageMetadataReservationEstimate(),
dataSegment.getId()
);
- continue;
}
- final StorageLocation.ReservationHold hold = location.addWeakReservationHold(
- id,
- () -> new PartialSegmentMetadataCacheEntry(
- dataSegment.getId(),
- partialDir,
- IndexIO.V10_FILE_NAME,
- List.of(),
- rangeReader,
- jsonMapper,
- virtualStorageLoadingThreadPool,
- config.getVirtualStorageMetadataReservationEstimate(),
- config.getVirtualStorageCoalesceGapBytes(),
- config.getVirtualStorageMaxFetchRunBytes()
- )
- );
- if (hold == null) {
- atomicMoveAndDeleteCacheEntryDirectory(partialDir);
- continue;
+ return null;
+ }
+ 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()
+ );
}
- 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()
- );
- }
+ if (writeInfoFile) {
rewriteInfoFile(dataSegment);
- partial.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment));
- return new ReservedPartial(partial, location, hold);
}
- catch (Throwable t) {
- // Close the hold (removing the never-mounted weak entry), then nuke the on-disk dir.
- try {
- throw CloseableUtils.closeAndWrapInCatch(t, hold);
- }
- finally {
+ partial.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment));
+ return new ReservedPartial(partial, location, hold);
+ }
+ catch (Throwable t) {
+ // Close the hold (removing the never-mounted weak entry), then nuke the dir if we created it.
+ try {
+ throw CloseableUtils.closeAndWrapInCatch(t, hold);
+ }
+ finally {
+ if (createdPartialDir) {
atomicMoveAndDeleteCacheEntryDirectory(partialDir);
}
}
}
- 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()
- );
}
/**
@@ -1415,14 +1414,23 @@ public DataSegment 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
- // did the reserving for us
+ // During bootstrap, reserve whatever this segment left on disk and mount it. getCachedSegments only recognizes
+ // the layout; the reservation is made here, one segment at a time, so that a partial's pessimistic metadata
+ // estimate is only outstanding until its mount shrinks it, and so that every mount runs under a hold. That hold
+ // matters: reclaim passes over held entries only, so an unheld entry can be evicted by a parallel bootstrap
+ // thread's reservation while this one is still mounting it.
final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(dataSegment.getId());
// Assemble the loaded-profile inside the segment lock (null = no partial materialized)
PartialLoadProfile loadedProfile = null;
final ReferenceCountingLock lock = lock(dataSegment);
synchronized (lock) {
+ StorageLocation.ReservationHold bootstrapHold = null;
try {
+ // A segment has either a partial layout or a complete one, so at most one of these reserves anything
+ bootstrapHold = reservePartialForBootstrap(dataSegment, id);
+ if (bootstrapHold == null) {
+ bootstrapHold = reserveCompleteForBootstrap(dataSegment, id);
+ }
for (StorageLocation location : locations) {
final CacheEntry entry = location.getCacheEntry(id);
if (entry == null) {
@@ -1475,6 +1483,12 @@ public DataSegment bootstrap(
}
}
finally {
+ if (bootstrapHold != null) {
+ CloseableUtils.closeAndSuppressExceptions(
+ bootstrapHold,
+ t -> log.warn(t, "Failed to release bootstrap reservation hold for segment[%s]", dataSegment.getId())
+ );
+ }
unlock(dataSegment, lock);
}
}
@@ -1599,6 +1613,134 @@ private void releaseRuleForFullLoad(DataSegment dataSegment, PartialSegmentMetad
);
}
+ /**
+ * Whether any location already has a cache entry for {@code id}
+ */
+ private boolean isRegisteredAtAnyLocation(SegmentCacheEntryIdentifier id)
+ {
+ for (StorageLocation location : locations) {
+ if (location.getCacheEntry(id) != null) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /**
+ * Reserve a metadata entry for a partial-load layout that survived on disk, so {@link #bootstrap} has something to
+ * mount, and hand back the hold it was reserved under. Returns {@code null} when this segment has no on-disk partial
+ * layout, or already has an entry, in which case bootstrap proceeds as it does for complete segments.
+ *
+ * Restoring reads the local layout rather than deep storage: the mount parses the on-disk header and picks up
+ * whichever bundles still have all their container files (see
+ * {@link PartialSegmentMetadataCacheEntry#mount}). The exception is a header that turns out to be damaged, which is
+ * deleted and re-fetched, leaving the entry mounted with nothing marked as downloaded.
+ *
+ * @throws SegmentLoadingException if the layout exists but cannot be restored, so the segment is marked failed and
+ * the coordinator re-issues a load for it rather than it being announced with
+ * nothing behind it
+ */
+ @Nullable
+ private StorageLocation.ReservationHold reservePartialForBootstrap(
+ DataSegment dataSegment,
+ SegmentCacheEntryIdentifier id
+ ) throws SegmentLoadingException
+ {
+ if (isRegisteredAtAnyLocation(id)) {
+ return null;
+ }
+ for (StorageLocation location : locations) {
+ final File partialDir = new File(location.getPath(), dataSegment.getId().toString());
+ if (!PartialSegmentFileMapperV10.isPartialSegmentLayout(partialDir, IndexIO.V10_FILE_NAME)) {
+ continue;
+ }
+ SegmentRangeReader rangeReader;
+ try {
+ rangeReader = tryOpenRangeReader(dataSegment);
+ }
+ catch (Exception e) {
+ log.warn(e, "Failed to open a range reader for partial segment[%s] during bootstrap", dataSegment.getId());
+ rangeReader = null;
+ }
+ if (rangeReader == null) {
+ // Anomalous: a layout on disk means range reads worked when it was written, so this should not happen (the
+ // loadSpec is now non-range-capable, or no longer converts to a known type). Reclaim the layout so the next
+ // load fetches it fresh rather than reserving an entry that could never fetch anything. The info file goes
+ // with it: it asserts that this segment has local cache state, which stops being true here, and nothing else
+ // will remove it (no entry was reserved, so there is no unmount hook to fire).
+ log.warn(
+ "On-disk partial-load layout for segment[%s] in [%s] has no usable range reader (this should not "
+ + "happen); deleting it so the segment can be re-loaded.",
+ dataSegment.getId(),
+ partialDir
+ );
+ atomicMoveAndDeleteCacheEntryDirectory(partialDir);
+ deleteSegmentInfoFile(dataSegment);
+ throw new SegmentLoadingException(
+ "No usable range reader for partial segment[%s]; its local layout has been reclaimed",
+ dataSegment.getId()
+ );
+ }
+ final ReservedPartial reserved = tryReservePartialAt(dataSegment, rangeReader, location, false);
+ if (reserved == null) {
+ throw new SegmentLoadingException(
+ "Failed to reserve partial metadata for segment[%s] on location[%s] during bootstrap",
+ dataSegment.getId(),
+ location.getPath()
+ );
+ }
+ return reserved.hold;
+ }
+ return null;
+ }
+
+ /**
+ * Reserve a {@link CompleteSegmentCacheEntry} for an eagerly-downloaded segment whose files are on disk, so
+ * {@link #bootstrap} has something to mount, and hand back the hold it was reserved under. Returns {@code null} when
+ * this segment has no complete layout on disk, or already has an entry. Virtual storage only: the legacy path
+ * reserves these statically in {@link #getCachedSegments} and never evicts them.
+ *
+ * Reserving under a hold is what makes the mount safe: reclaim skips held entries only, so an unheld entry can be
+ * selected as a victim while a parallel bootstrap thread is still mounting it. Releasing the hold once mounted
+ * leaves the entry as evictable as any other weak entry.
+ *
+ * @throws SegmentLoadingException if the location cannot accept the reservation, which used to be an alert followed
+ * by mounting the segment unreserved, leaving the location under-counting the disk
+ * it was actually using
+ */
+ @Nullable
+ private StorageLocation.ReservationHold reserveCompleteForBootstrap(
+ DataSegment dataSegment,
+ SegmentCacheEntryIdentifier id
+ ) throws SegmentLoadingException
+ {
+ if (isRegisteredAtAnyLocation(id)) {
+ return null;
+ }
+ for (StorageLocation location : locations) {
+ final CacheEntry cacheEntry = new CompleteSegmentCacheEntry(dataSegment);
+ if (!((CompleteSegmentCacheEntry) cacheEntry).checkExists(location.getPath())) {
+ continue;
+ }
+ final StorageLocation.ReservationHold hold = location.addWeakReservationHold(
+ id,
+ () -> new CompleteSegmentCacheEntry(dataSegment)
+ );
+ if (hold == null) {
+ throw new SegmentLoadingException(
+ "Location[%s] with available bytes[%,d] cannot reserve segment[%s] of size[%,d] during bootstrap; check "
+ + "druid.segmentCache.locations maxSize",
+ location.getPath(),
+ location.availableSizeBytes(),
+ dataSegment.getId(),
+ dataSegment.getSize()
+ );
+ }
+ return hold;
+ }
+ return null;
+ }
+
/**
* Reapply the persisted partial-load rule to a bootstrap-restored metadata entry. Reads the wrapper from the
* segment's info-file {@code loadSpec}, resolves the selected bundle names against the just-parsed on-disk
@@ -1621,11 +1763,6 @@ private void reapplyRuleFromInfoFile(DataSegment dataSegment, PartialSegmentMeta
"Bootstrap-restored partial metadata for segment[%s] has no file mapper", dataSegment.getId()
);
}
- // Register the info-file cleanup hook BEFORE anything that could throw. reserveFromDisk does NOT install an
- // onUnmount hook on bootstrap-restored entries, so without this ordering a throw from getSelectedBundleNames
- // or applyRule would leave the entry with no cleanup path. Setting the hook first ensures every teardown path
- // deletes the info file consistently.
- partial.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment));
final Set selected = Set.copyOf(
wrapper.getSelectedBundleNames(dataSegment, mapper.getSegmentFileMetadata())
);
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 6330f20cb4c9..0c3e4cc8d77c 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
@@ -50,9 +50,10 @@
* with {@link #reserve(CacheEntry)}, where the space of the entry is accounted for and the storage space will not be
* recovered until {@link #release(CacheEntry)} is called. These entries are stored in {@link #staticCacheEntries}.
*
- * The second way is to store as a transient cache item with one of {@link #reserveWeak(CacheEntry)},
- * {@link #addWeakReservationHold(CacheEntryIdentifier, Supplier)}, or
- * {@link #addWeakReservationHoldIfExists(CacheEntryIdentifier)}. {@link CacheEntry} stored in this manner will exist on
+ * The second way is to store as a transient cache item with
+ * {@link #addWeakReservationHold(CacheEntryIdentifier, Supplier)} or
+ * {@link #addWeakReservationHoldIfExists(CacheEntryIdentifier)}. Both place a hold on the entry, so it cannot be
+ * reclaimed while the caller is still mounting or reading it. {@link CacheEntry} stored in this manner will exist on
* disk in this location until the point that another new reservation needs more space than remains available in the
* location, at which point {@link #reclaim(long)} will be called to try to call {@link CacheEntry#unmount()} on any
* eligible entries until enough space is available to store the new item.
@@ -284,54 +285,6 @@ public boolean reserve(CacheEntry entry)
}
}
- /**
- * Reserves space to store a 'weak' reservation for a given {@link CacheEntry}. Returns true if already reserved or
- * was able to be successfully reserved, or false if unable to be reserved. This method is intended for use during
- * 'bootstrapping'. To use weak cache entries in a query engine use
- * {@link #addWeakReservationHold(CacheEntryIdentifier, Supplier)} or
- * {@link #addWeakReservationHoldIfExists(CacheEntryIdentifier)}, which places a hold on cache entries to prevent
- * eviction until the hold is released.
- */
- public boolean reserveWeak(CacheEntry entry)
- {
- lock.readLock().lock();
- try {
- if (staticCacheEntries.containsKey(entry.getId())) {
- return true;
- }
- if (weakCacheEntries.containsKey(entry.getId())) {
- weakCacheEntries.get(entry.getId()).visited = true;
- return true;
- }
- }
- finally {
- lock.readLock().unlock();
- }
-
- lock.writeLock().lock();
- try {
- if (staticCacheEntries.containsKey(entry.getId())) {
- return true;
- }
- if (weakCacheEntries.containsKey(entry.getId())) {
- weakCacheEntries.get(entry.getId()).visited = true;
- return true;
- }
- final ReclaimResult reclaimResult = canHandleWeak(entry);
- unmountReclaimed(reclaimResult);
- if (reclaimResult.isSuccess()) {
- final WeakCacheEntry newEntry = new WeakCacheEntry(entry);
- linkNewWeakEntry(newEntry);
- weakCacheEntries.put(entry.getId(), newEntry);
- weakStats.getAndUpdate(s -> s.loadBegin(entry.getSize()));
- }
- return reclaimResult.isSuccess();
- }
- finally {
- lock.writeLock().unlock();
- }
- }
-
/**
* Returns a {@link ReservationHold} of a {@link CacheEntry} with a 'hold' placed on it, preventing it from being
* automatically removed by {@link #reclaim(long)} if the {@link CacheEntry} is one of {@link #weakCacheEntries} until
@@ -524,13 +477,12 @@ public void release(CacheEntry entry)
* 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.
+ * This is for reclaiming a reservation whose entry has nothing behind it any more, after a failed mount. It covers
+ * a hold released mid-mount, once the entry already reported {@link CacheEntry#isMounted()}. Every weak entry is
+ * created under a {@link ReservationHold} and releasing that hold removes an entry that never mounted, but the
+ * release runnable leaves a mounted one registered, so if the mount then fails anyway nothing else would remove it.
+ * The hold guard makes this safe to call unconditionally on any mount failure: an entry someone still holds is left
+ * to its holder's release runnable.
*/
public void removeUnheldWeakEntry(CacheEntryIdentifier id)
{
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 762230eed4eb..b3bcb026f050 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
@@ -661,7 +661,7 @@ private PartialSegmentMetadataCacheEntry mountedEntryOver(File deepStorageDir) t
}
/**
- * Variant of {@link #mountedEntryOver} that uses {@link StorageLocation#reserveWeak} so the mounted entry is a real
+ * Variant of {@link #mountedEntryOver} that leaves the entry registered but unheld, so the mounted entry is a real
* weak reservation — needed by the rule-holds state machine, which calls
* {@link StorageLocation#addWeakReservationHoldIfExists} on itself when {@code applyRule} runs.
*/
@@ -682,8 +682,13 @@ private PartialSegmentMetadataCacheEntry mountedWeakEntryOver(File deepStorageDi
PartialSegmentFileMapperV10.DEFAULT_COALESCE_GAP_BYTES,
PartialSegmentFileMapperV10.DEFAULT_MAX_FETCH_RUN_BYTES
);
- Assertions.assertTrue(location.reserveWeak(entry));
+ // Reserve under a hold, mount, then release: the entry stays registered but unheld, as it would after a
+ // bootstrap restore.
+ final StorageLocation.ReservationHold hold =
+ location.addWeakReservationHold(entry.getId(), () -> entry);
+ Assertions.assertNotNull(hold);
entry.mount(location);
+ hold.close();
return entry;
}
}
diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentRestoreFromDiskTest.java
similarity index 76%
rename from server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java
rename to server/src/test/java/org/apache/druid/segment/loading/PartialSegmentRestoreFromDiskTest.java
index c972f9eee77a..d3e651aff2e4 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java
+++ b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentRestoreFromDiskTest.java
@@ -26,7 +26,6 @@
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;
@@ -52,9 +51,12 @@
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
+import org.mockito.ArgumentMatchers;
+import org.mockito.Mockito;
import java.io.File;
import java.io.IOException;
+import java.io.RandomAccessFile;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -64,7 +66,7 @@
import java.util.concurrent.ThreadLocalRandom;
import java.util.stream.Collectors;
-class PartialSegmentCacheBootstrapTest
+class PartialSegmentRestoreFromDiskTest
{
private static final ObjectMapper JSON_MAPPER = TestHelper.makeJsonMapper();
private static final SegmentId SEGMENT_ID = SegmentId.of("test", Intervals.of("2025/2026"), "v1", 0);
@@ -296,13 +298,14 @@ void testRestoreRollsBackOnBundleReservationFailure() throws IOException
cacheDir,
IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX
);
- // size location exactly to the header size: metadata reservation fits, but the first bundle's weak reservation
- // has 0 bytes of remaining budget and no weak entries to reclaim, so addWeakReservationHold returns null
+ // Size the location and the reservation estimate exactly to the header size: the metadata reservation fits (and
+ // needs no shrink at mount), but the first bundle's weak reservation has 0 bytes of remaining budget and no weak
+ // entries to reclaim, so addWeakReservationHold returns null.
final StorageLocation location = new StorageLocation(cacheDir, headerFile.length(), null);
Assertions.assertThrows(
Throwable.class,
- () -> restoreFromDisk(location)
+ () -> restoreFromDisk(location, new DirectoryBackedRangeReader(deepStorageDir), headerFile.length())
);
// rollback must release the metadata reservation and leave no static/weak entries behind
@@ -317,6 +320,40 @@ void testRestoreRollsBackOnBundleReservationFailure() throws IOException
Assertions.assertFalse(headerFile.exists(), "bootstrap failure deletes the header via the unmount cleanup path");
}
+ /**
+ * The restore mounts {@code __base} before the bundles that depend on it, so a failure on a later bundle has to
+ * unwind one that already mounted. Unmounting it is not enough: its cache entry would stay registered with its
+ * reservation, while its containers are gone, and a later acquire reuses a registered entry by id, re-mounting a
+ * bundle bound to the metadata entry this failure tears down.
+ */
+ @Test
+ void testRollbackRemovesBundlesItAlreadyMounted() throws IOException
+ {
+ primeOnDiskState();
+
+ final StorageLocation location = Mockito.spy(new StorageLocation(cacheDir, ESTIMATE * 8, null));
+ final PartialSegmentBundleCacheEntryIdentifier aggId =
+ new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE);
+ // Refuse only the aggregate bundle's reservation, so the restore fails with __base already mounted.
+ Mockito.doReturn(null)
+ .when(location)
+ .addWeakReservationHold(ArgumentMatchers.eq(aggId), ArgumentMatchers.any());
+
+ Assertions.assertThrows(Throwable.class, () -> restoreFromDisk(location));
+
+ // Pins that __base really did mount and then got removed, rather than the rollback loop being empty because the
+ // restore failed before mounting anything.
+ Mockito.verify(location).removeUnheldWeakEntry(
+ new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME)
+ );
+ Assertions.assertEquals(
+ 0,
+ location.getWeakEntryCount(),
+ "rollback must remove the bundle entries it registered, not just unmount them"
+ );
+ Assertions.assertEquals(0, location.currentSizeBytes(), "rollback must release every reservation it took");
+ }
+
@Test
void testMountFailureRemovesLingeringWeakEntry() throws IOException
{
@@ -324,24 +361,10 @@ void testMountFailureRemovesLingeringWeakEntry() throws IOException
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), but with a range
- // reader that fails every fetch so we can force a mount failure below.
+ // A range reader that fails every fetch, so the mount below cannot rebuild the file mapper.
final SegmentRangeReader failingReader = (filename, offset, length) -> {
throw new IOException("simulated deep-storage fetch failure during mount");
};
- final PartialSegmentMetadataCacheEntry metadata = PartialSegmentCacheBootstrap.reserveFromDisk(
- SEGMENT_ID,
- cacheDir,
- IndexIO.V10_FILE_NAME,
- List.of(),
- failingReader,
- JSON_MAPPER,
- null,
- location,
- PartialSegmentFileMapperV10.DEFAULT_COALESCE_GAP_BYTES,
- PartialSegmentFileMapperV10.DEFAULT_MAX_FETCH_RUN_BYTES
- );
- Assertions.assertTrue(location.isWeakReserved(id));
// Delete the header so the mount's file-mapper build must re-fetch it from deep storage; the failing reader throws,
// failing the mount (the same poison that arises when create() detects header corruption and the deep-storage
@@ -352,7 +375,7 @@ void testMountFailureRemovesLingeringWeakEntry() throws IOException
);
Assertions.assertTrue(headerFile.delete());
- Assertions.assertThrows(Throwable.class, () -> metadata.mount(location));
+ Assertions.assertThrows(Throwable.class, () -> restoreFromDisk(location, failingReader, ESTIMATE));
// The failed mount must not leave the lingering weak entry behind: a later findExistingPartialWithHold would
// otherwise resurrect it and re-mount would fail forever (failing reader + deleted header). It must be gone so
@@ -365,9 +388,9 @@ void testMountFailureRemovesLingeringWeakEntry() throws IOException
@Test
void testRestoredEntryCanFetchUndownloadedFile() throws IOException
{
- // a bootstrap-restored entry must keep the segment's real deep-storage range reader so a later query can
- // fetch a bundle/column that wasn't on disk at startup. Previously the entry held a throwing disk-only reader, so
- // this fetch failed with "bootstrap should only read from local disk".
+ // a restored entry keeps the segment's real deep-storage range reader, so a later query can fetch a bundle or
+ // column that wasn't on disk at startup. The local layout is a head start, never the limit of what the entry can
+ // read.
primeOnDiskState();
final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null);
final PartialSegmentMetadataCacheEntry metadata = restoreFromDisk(location);
@@ -394,38 +417,71 @@ void testRestoredEntryCanFetchUndownloadedFile() throws IOException
}
@Test
- void testReserveFailsWhenHeaderMissing()
+ void testMountShrinksEstimateToHeaderFootprint() throws IOException
{
- // no priming: cacheDir is empty
+ primeOnDiskState();
+
final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null);
- Assertions.assertThrows(
- DruidException.class,
- () -> PartialSegmentCacheBootstrap.reserveFromDisk(
- SEGMENT_ID,
- cacheDir,
- IndexIO.V10_FILE_NAME,
- List.of(),
- new DirectoryBackedRangeReader(deepStorageDir),
- JSON_MAPPER,
- null,
- location,
- PartialSegmentFileMapperV10.DEFAULT_COALESCE_GAP_BYTES,
- PartialSegmentFileMapperV10.DEFAULT_MAX_FETCH_RUN_BYTES
- )
- );
+ final PartialSegmentMetadataCacheEntry metadata = restoreFromDisk(location);
+
+ // The entry is reserved with the pessimistic estimate and shrunk by mount, which is the only place an entry's
+ // size is computed. Restoring an existing layout is no exception: it goes through the same reserve-then-shrink,
+ // so there is no second, pre-mount size that the mount could disagree with.
+ final long headerFootprint = metadata.getFileMapper().getOnDiskHeaderSize();
+ Assertions.assertTrue(headerFootprint < ESTIMATE);
+ Assertions.assertEquals(headerFootprint, metadata.getSize());
+
+ // Re-mounting the same entry measures the same footprint, so the shrink is a no-op rather than a grow
+ metadata.unmount();
+ metadata.mount(location);
+ Assertions.assertEquals(headerFootprint, metadata.getFileMapper().getOnDiskHeaderSize());
+ Assertions.assertEquals(headerFootprint, metadata.getSize());
}
@Test
- void testIsPartialSegmentLayoutDetectsHeader() throws IOException
+ void testRestoreWithoutLocalHeaderColdFetches() throws IOException
+ {
+ // no priming: cacheDir is empty, so the mount fetches the header from deep storage like a fresh acquire
+ final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null);
+ final PartialSegmentMetadataCacheEntry metadata = restoreFromDisk(location);
+
+ Assertions.assertTrue(metadata.isMounted());
+ Assertions.assertEquals(metadata.getFileMapper().getOnDiskHeaderSize(), metadata.getSize());
+ }
+
+ /**
+ * A header file whose bitmap region never made it to disk (a crash, or an interrupt, between persisting the header
+ * bytes and extending the file) must be treated as corrupt and re-fetched. Before the header and its bitmap region
+ * were published together, this state was silently repaired by growing the file, which meant the same segment
+ * measured one size before the mount and a larger one after it, and a reservation sized from the short file was
+ * asked to grow, which storage locations do not support.
+ */
+ @Test
+ void testHeaderMissingBitmapRegionDoesNotGrowReservation() throws IOException
{
- Assertions.assertFalse(PartialSegmentCacheBootstrap.isPartialSegmentLayout(cacheDir, IndexIO.V10_FILE_NAME));
primeOnDiskState();
- Assertions.assertTrue(PartialSegmentCacheBootstrap.isPartialSegmentLayout(cacheDir, IndexIO.V10_FILE_NAME));
- Assertions.assertFalse(PartialSegmentCacheBootstrap.isPartialSegmentLayout(null, IndexIO.V10_FILE_NAME));
- Assertions.assertFalse(PartialSegmentCacheBootstrap.isPartialSegmentLayout(
- new File(perTestTempDir, "nonexistent"),
- IndexIO.V10_FILE_NAME
- ));
+
+ final File headerFile = new File(
+ cacheDir,
+ IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX
+ );
+ final int numFiles;
+ try (PartialSegmentFileMapperV10 introspect = createMapper(deepStorageDir, cacheDir)) {
+ numFiles = introspect.getSegmentFileMetadata().getFiles().size();
+ }
+ final long fullLength = headerFile.length();
+ final int bitmapBytes = (numFiles + 7) / 8;
+ Assertions.assertTrue(bitmapBytes > 0);
+ try (RandomAccessFile raf = new RandomAccessFile(headerFile, "rw")) {
+ raf.setLength(fullLength - bitmapBytes);
+ }
+
+ final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null);
+ final PartialSegmentMetadataCacheEntry metadata = restoreFromDisk(location);
+
+ Assertions.assertTrue(metadata.isMounted());
+ Assertions.assertEquals(fullLength, headerFile.length(), "the truncated header must have been re-fetched");
+ Assertions.assertEquals(fullLength, metadata.getSize());
}
@Test
@@ -529,26 +585,48 @@ private void primeOnDiskState() throws IOException
}
/**
- * 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.
+ * Two-step restore helper that mirrors what {@code SegmentLocalCacheManager} does in production: reserve the
+ * metadata entry with the up-front estimate under a hold (exactly as {@code bootstrap()} does when it finds a
+ * partial layout on disk), drive {@link PartialSegmentMetadataCacheEntry#mount} to trigger the file-mapper build,
+ * reservation shrink, and bundle restore, then release the hold. On a mount failure, releasing the hold removes the
+ * never-mounted weak entry, 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(
+ return restoreFromDisk(location, new DirectoryBackedRangeReader(deepStorageDir), ESTIMATE);
+ }
+
+ private PartialSegmentMetadataCacheEntry restoreFromDisk(
+ StorageLocation location,
+ SegmentRangeReader rangeReader,
+ long reservationEstimate
+ ) throws IOException
+ {
+ final PartialSegmentMetadataCacheEntry metadata = new PartialSegmentMetadataCacheEntry(
SEGMENT_ID,
cacheDir,
IndexIO.V10_FILE_NAME,
List.of(),
- new DirectoryBackedRangeReader(deepStorageDir),
+ rangeReader,
JSON_MAPPER,
null,
- location,
+ reservationEstimate,
PartialSegmentFileMapperV10.DEFAULT_COALESCE_GAP_BYTES,
PartialSegmentFileMapperV10.DEFAULT_MAX_FETCH_RUN_BYTES
);
- metadata.mount(location);
+ final StorageLocation.ReservationHold hold = location.addWeakReservationHold(
+ metadata.getId(),
+ () -> metadata
+ );
+ if (hold == null) {
+ throw new IOException("location has no capacity for the metadata reservation");
+ }
+ try {
+ metadata.mount(location);
+ }
+ finally {
+ hold.close();
+ }
return metadata;
}
diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerBootstrapReserveTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerBootstrapReserveTest.java
new file mode 100644
index 000000000000..e75f9c3d9137
--- /dev/null
+++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerBootstrapReserveTest.java
@@ -0,0 +1,140 @@
+/*
+ * 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.java.util.common.FileUtils;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.segment.SegmentLazyLoadFailCallback;
+import org.apache.druid.segment.TestHelper;
+import org.apache.druid.segment.column.ColumnConfig;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.partition.NoneShardSpec;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
+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.List;
+import java.util.Map;
+
+/**
+ * Covers what a virtual-storage bootstrap reserves, and what happens when it cannot. Both the partial and the complete
+ * layout are reserved by {@code bootstrap()} rather than by {@code getCachedSegments()}, so that every mount runs under
+ * a hold: reclaim passes over held entries only, so an unheld entry can be chosen as an eviction victim by a parallel
+ * bootstrap thread's reservation while this one is still mounting it.
+ */
+class SegmentLocalCacheManagerBootstrapReserveTest
+{
+ private static final SegmentId SEGMENT_ID = SegmentId.of("test", Intervals.of("2025/2026"), "v1", 0);
+ private static final long SEGMENT_SIZE = 4096L;
+
+ @TempDir
+ File tempDir;
+
+ private File cacheRoot;
+ private SegmentLocalCacheManager manager;
+ private DataSegment segment;
+
+ @BeforeEach
+ void setup() throws IOException
+ {
+ cacheRoot = new File(tempDir, "cache");
+ FileUtils.mkdirp(cacheRoot);
+ segment = DataSegment.builder(SEGMENT_ID)
+ .shardSpec(NoneShardSpec.instance())
+ .loadSpec(Map.of("type", "local", "path", new File(tempDir, "deep").getAbsolutePath()))
+ .size(SEGMENT_SIZE)
+ .build();
+ }
+
+ @AfterEach
+ void tearDown()
+ {
+ if (manager != null) {
+ manager.shutdown();
+ }
+ }
+
+ /**
+ * Build a virtual-storage manager over a location of the given size, with the on-disk shape of an eagerly
+ * downloaded segment: a directory named for the segment, plus its info file. The directory's contents don't matter
+ * to either test, since neither gets as far as a successful mount.
+ */
+ private void setUpManagerWithLocationSize(long locationSize) throws IOException
+ {
+ final ObjectMapper jsonMapper = TestHelper.makeJsonMapper();
+ final StorageLocationConfig locationConfig = new StorageLocationConfig(cacheRoot, locationSize, null);
+ final SegmentLoaderConfig loaderConfig = SegmentLoaderConfig.builder()
+ .locations(locationConfig)
+ .virtualStorage(true)
+ .build();
+ final List storageLocations = loaderConfig.toStorageLocations();
+ manager = new SegmentLocalCacheManager(
+ storageLocations,
+ loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
+ new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
+ TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
+ jsonMapper
+ );
+ FileUtils.mkdirp(new File(cacheRoot, SEGMENT_ID.toString()));
+ manager.storeInfoFile(segment);
+ }
+
+ @Test
+ void testCompleteLayoutIsReservedByBootstrapNotByGetCachedSegments() throws IOException
+ {
+ // Room to spare, so a reservation here would succeed if getCachedSegments still made one.
+ setUpManagerWithLocationSize(SEGMENT_SIZE * 4);
+
+ Assertions.assertEquals(List.of(segment), manager.getCachedSegments());
+ Assertions.assertNull(
+ manager.getLocations().get(0).getCacheEntry(new SegmentCacheEntryIdentifier(SEGMENT_ID)),
+ "getCachedSegments must only recognize the layout; reserving is bootstrap's job"
+ );
+ Assertions.assertEquals(0, manager.getLocations().get(0).currentSizeBytes());
+ }
+
+ @Test
+ void testBootstrapFailsSegmentWhenLocationCannotReserveIt() throws IOException
+ {
+ // Far too small for the segment, so the bootstrap reservation cannot be satisfied.
+ setUpManagerWithLocationSize(SEGMENT_SIZE / 4);
+ manager.getCachedSegments();
+
+ // Previously the too-small reservation was an alert, after which the segment was mounted anyway with nothing
+ // reserved for it, leaving the location under-counting the disk it was really using. Now the segment fails to
+ // bootstrap, so it is not announced and the coordinator re-issues a load for it.
+ Assertions.assertThrows(
+ SegmentLoadingException.class,
+ () -> manager.bootstrap(segment, SegmentLazyLoadFailCallback.NOOP)
+ );
+ Assertions.assertNull(
+ manager.getLocations().get(0).getCacheEntry(new SegmentCacheEntryIdentifier(SEGMENT_ID)),
+ "a segment that could not be reserved must not be left registered"
+ );
+ Assertions.assertEquals(0, manager.getLocations().get(0).currentSizeBytes());
+ }
+}
diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java
index 7a823904db63..e251395f4ecc 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java
+++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java
@@ -78,8 +78,10 @@
import java.io.File;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
+import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -754,9 +756,9 @@ 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
+ // and the segment info file BEFORE getCachedSegments runs, then verify the two-phase contract:
+ // - getCachedSegments() only recognizes the layout as cached; it reserves nothing
+ // - bootstrap(DataSegment) reserves the metadata entry and mounts it, which is also what sizes it
final File partialDir = new File(cacheRoot, SEGMENT_ID.toString());
FileUtils.mkdirp(partialDir);
primePartialOnDiskState(partialDir);
@@ -767,23 +769,29 @@ void testGetCachedSegmentsThenBootstrapMountsPartialEntry()
final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(SEGMENT_ID);
final StorageLocation location = manager.getLocations().get(0);
+ Assertions.assertNull(
+ location.getCacheEntry(id),
+ "getCachedSegments must not reserve a partial entry; reservation is deferred to bootstrap()"
+ );
+ Assertions.assertEquals(0, location.currentSizeBytes(), "nothing should be charged to the location yet");
+
+ manager.bootstrap(partialSegment, SegmentLazyLoadFailCallback.NOOP);
+
final CacheEntry reserved = location.getCacheEntry(id);
Assertions.assertInstanceOf(
PartialSegmentMetadataCacheEntry.class,
reserved,
- "getCachedSegments must reserve a partial metadata entry on the location"
+ "bootstrap must reserve a partial metadata entry on the location holding the layout"
);
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()");
+ // The estimate the entry was reserved with has been shrunk to the mounted footprint, and the bootstrap hold is
+ // released, so the restored entry is as evictable as any other
+ Assertions.assertEquals(partial.getFileMapper().getOnDiskHeaderSize(), partial.getSize());
+ Assertions.assertTrue(
+ partial.getSize() < manager.getConfig().getVirtualStorageMetadataReservationEstimate()
+ );
Assertions.assertFalse(
partial.snapshotLinkedBundles().isEmpty(),
"bundle restore must have linked at least one bundle to the metadata entry"
@@ -795,6 +803,26 @@ void testGetCachedSegmentsThenBootstrapMountsPartialEntry()
restoredBundles.contains(Projections.BASE_TABLE_PROJECTION_NAME),
"base bundle must be restored by bootstrap; got " + restoredBundles
);
+
+ // A bootstrap-restored entry owns its info file just like an on-demand one: the info file asserts that local
+ // cache state exists, and unmounting deletes that state (header files included), so the two go together.
+ // Without the hook, an evicted restored segment would leave the next startup trying to restore a layout that
+ // isn't there. (Note dropping the segment is not what triggers this: a weak entry survives drop, and reclaim is
+ // left to eviction, which is what the unmount below stands in for.)
+ final File infoFile = new File(new File(cacheRoot, "info_dir"), SEGMENT_ID.toString());
+ Assertions.assertTrue(infoFile.exists(), "bootstrap must leave the info file it restored from in place");
+
+ // Unmount dependents before parents so each parent hold is released before its bundle tears down, the order
+ // eviction's cascade produces.
+ final List bundles = new ArrayList<>(partial.snapshotLinkedBundles());
+ bundles.sort(Comparator.comparing(b -> Projections.BASE_TABLE_PROJECTION_NAME.equals(b.getBundleName())));
+ for (PartialSegmentBundleCacheEntry bundle : bundles) {
+ bundle.unmount();
+ }
+ partial.unmount();
+
+ Assertions.assertFalse(partial.isMounted(), "the restored entry should be unmounted");
+ Assertions.assertFalse(infoFile.exists(), "unmounting the restored entry must delete its info file");
}
@Test
@@ -805,7 +833,7 @@ void testGetCachedSegmentsDeletesPartialLayoutWhenPartialDisabled() throws IOExc
FileUtils.mkdirp(partialDir);
primePartialOnDiskState(partialDir);
manager.storeInfoFile(partialSegment);
- Assertions.assertTrue(PartialSegmentCacheBootstrap.isPartialSegmentLayout(partialDir, IndexIO.V10_FILE_NAME));
+ Assertions.assertTrue(PartialSegmentFileMapperV10.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.
@@ -846,7 +874,7 @@ void testGetCachedSegmentsDeletesPartialLayoutWhenPartialDisabled() throws IOExc
}
@Test
- void testGetCachedSegmentsDeletesPartialLayoutWhenRangeReaderUnavailable() throws IOException
+ void testBootstrapDeletesPartialLayoutWhenRangeReaderUnavailable() throws IOException
{
// Prime a valid partial on-disk layout (header written from the real deep-storage dir) as a previous run left it.
final File partialDir = new File(cacheRoot, SEGMENT_ID.toString());
@@ -855,8 +883,8 @@ void testGetCachedSegmentsDeletesPartialLayoutWhenRangeReaderUnavailable() throw
// ...but record the segment with a loadSpec whose storage can't produce a range reader: an existing directory
// that holds no V10 file, so LocalLoadSpec.openRangeReader returns null. This is the "shouldn't happen" case — a
- // partial layout on disk means range reads worked when it was written — so partial-enabled bootstrap must reclaim
- // the layout rather than reserve an entry that could never lazily fetch.
+ // partial layout on disk means range reads worked when it was written — so bootstrap must reclaim the layout and
+ // fail the segment rather than reserve an entry that could never lazily fetch.
final File noRangeReaderStorage = new File(perTestTempDir, "no_range_reader_storage");
FileUtils.mkdirp(noRangeReaderStorage);
final DataSegment unreadableSegment =
@@ -866,17 +894,27 @@ void testGetCachedSegmentsDeletesPartialLayoutWhenRangeReaderUnavailable() throw
.size(0)
.build();
manager.storeInfoFile(unreadableSegment);
- Assertions.assertTrue(PartialSegmentCacheBootstrap.isPartialSegmentLayout(partialDir, IndexIO.V10_FILE_NAME));
+ Assertions.assertTrue(PartialSegmentFileMapperV10.isPartialSegmentLayout(partialDir, IndexIO.V10_FILE_NAME));
+ final File infoFile = new File(new File(cacheRoot, "info_dir"), SEGMENT_ID.toString());
+ Assertions.assertTrue(infoFile.exists());
+ // The layout is on disk, so it still counts as cached; the reader is only needed once bootstrap tries to restore it
final List cached = manager.getCachedSegments();
- Assertions.assertFalse(
- cached.contains(unreadableSegment),
- "bootstrap must not return a partial segment whose deep storage can't produce a range reader"
+ Assertions.assertEquals(List.of(unreadableSegment), cached);
+
+ Assertions.assertThrows(
+ SegmentLoadingException.class,
+ () -> manager.bootstrap(unreadableSegment, SegmentLazyLoadFailCallback.NOOP),
+ "bootstrap must fail the segment so the coordinator re-issues a load for it"
);
Assertions.assertFalse(
partialDir.exists(),
"bootstrap must delete the unusable partial layout from disk"
);
+ Assertions.assertFalse(
+ infoFile.exists(),
+ "the info file must go with the layout; it claims local cache state that no longer exists"
+ );
Assertions.assertNull(
manager.getLocations().get(0).getCacheEntry(new SegmentCacheEntryIdentifier(SEGMENT_ID)),
"no cache entry should be reserved for the deleted partial layout"
@@ -884,7 +922,7 @@ void testGetCachedSegmentsDeletesPartialLayoutWhenRangeReaderUnavailable() throw
}
@Test
- void testGetCachedSegmentsDeletesPartialLayoutWhenLoadSpecUnconvertible() throws IOException
+ void testBootstrapDeletesPartialLayoutWhenLoadSpecUnconvertible() throws IOException
{
// Prime a valid partial on-disk layout as a previous run left it...
final File partialDir = new File(cacheRoot, SEGMENT_ID.toString());
@@ -892,8 +930,8 @@ void testGetCachedSegmentsDeletesPartialLayoutWhenLoadSpecUnconvertible() throws
primePartialOnDiskState(partialDir);
// ...but record the segment with a loadSpec whose type is no longer registered, so converting it to a LoadSpec
- // throws. Bootstrap must treat that broken segment like a null reader: delete the unusable layout and continue,
- // rather than aborting or reserving an entry that could never fetch.
+ // throws. Bootstrap must treat that broken segment like a null reader: delete the unusable layout and fail this
+ // segment, rather than aborting the whole bootstrap or reserving an entry that could never fetch.
final DataSegment unconvertibleSegment =
DataSegment.builder(SEGMENT_ID)
.shardSpec(NoneShardSpec.instance())
@@ -901,14 +939,23 @@ void testGetCachedSegmentsDeletesPartialLayoutWhenLoadSpecUnconvertible() throws
.size(0)
.build();
manager.storeInfoFile(unconvertibleSegment);
- Assertions.assertTrue(PartialSegmentCacheBootstrap.isPartialSegmentLayout(partialDir, IndexIO.V10_FILE_NAME));
+ Assertions.assertTrue(PartialSegmentFileMapperV10.isPartialSegmentLayout(partialDir, IndexIO.V10_FILE_NAME));
+ final File infoFile = new File(new File(cacheRoot, "info_dir"), SEGMENT_ID.toString());
+ Assertions.assertTrue(infoFile.exists());
final List cached = manager.getCachedSegments();
- Assertions.assertFalse(
- cached.contains(unconvertibleSegment),
- "bootstrap must not return a partial segment whose loadSpec can't be converted to a range reader"
+ Assertions.assertEquals(List.of(unconvertibleSegment), cached);
+
+ Assertions.assertThrows(
+ SegmentLoadingException.class,
+ () -> manager.bootstrap(unconvertibleSegment, SegmentLazyLoadFailCallback.NOOP),
+ "bootstrap must fail the segment so the coordinator re-issues a load for it"
);
Assertions.assertFalse(partialDir.exists(), "bootstrap must delete the unusable partial layout from disk");
+ Assertions.assertFalse(
+ infoFile.exists(),
+ "the info file must go with the layout; it claims local cache state that no longer exists"
+ );
Assertions.assertNull(
manager.getLocations().get(0).getCacheEntry(new SegmentCacheEntryIdentifier(SEGMENT_ID)),
"no cache entry should be reserved for the deleted partial layout"
diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
index a0ba9739b9c9..0daded9e8ea9 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
+++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
@@ -812,8 +812,8 @@ void testBootstrapReinstallsRuleHoldsFromPersistedInfoFile() throws Exception
"bootstrap must reapply the persisted PartialLoadSpec wrapper's fingerprint"
);
// Selected bundle + its base dependency are held after bootstrap restore (they were on disk and got
- // restored by the metadata mount's PartialSegmentCacheBootstrap.restoreBundlesFromDisk, which register with
- // the metadata; applyRule picks up their rule-holds from the linkedBundles state).
+ // restored by the metadata mount's own bundle restore, which registers them with the metadata; applyRule picks
+ // up their rule-holds from the linkedBundles state).
final StorageLocation loc = restarted.getLocations().get(0);
Assertions.assertTrue(
loc.isWeakReserved(new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE)),
@@ -844,9 +844,12 @@ void testReserveMkdirpFailureFallsThroughToNextLocation() throws Exception
// Bump writable's usage so LeastBytesUsed picks readOnly first.
final SegmentId dummy = SegmentId.of("dummy", Intervals.of("2020/2021"), "v", 0);
- Assertions.assertTrue(
- manager.getLocations().get(0).reserveWeak(stubCacheEntry(new SegmentCacheEntryIdentifier(dummy), 4096L))
- );
+ final CacheEntry filler = stubCacheEntry(new SegmentCacheEntryIdentifier(dummy), 4096L);
+ final StorageLocation.ReservationHold fillerHold =
+ manager.getLocations().get(0).addWeakReservationHold(filler.getId(), () -> filler);
+ Assertions.assertNotNull(fillerHold);
+ filler.mount(manager.getLocations().get(0));
+ fillerHold.close();
Assertions.assertTrue(readOnly.setReadOnly(), "test setup must be able to make readOnly location read-only");
try {
manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
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 59bd0d30de85..29c2e8f61626 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
@@ -1092,10 +1092,20 @@ public void testGetBootstrapSegmentVirtualStorageSegmentAlreadyCached() throws E
unzippedSegmentPathInLocation
);
- for (DataSegment segment : manager.getCachedSegments()) {
+ final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segmentToBootstrap.getId());
+ final List cached = manager.getCachedSegments();
+ Assertions.assertEquals(ImmutableList.of(segmentToBootstrap), cached);
+ // under virtual storage getCachedSegments only recognizes the layout; the reservation is bootstrap's job, so that
+ // the mount happens under a hold and cannot be evicted out from under itself
+ Assertions.assertNull(manager.getLocations().get(0).getCacheEntry(id));
+
+ for (DataSegment segment : cached) {
manager.bootstrap(segment, SegmentLazyLoadFailCallback.NOOP);
}
+ final CacheEntry entry = manager.getLocations().get(0).getCacheEntry(id);
+ Assertions.assertNotNull(entry, "bootstrap must reserve the complete entry it mounts");
+ Assertions.assertTrue(entry.isMounted());
// if bootstrapping a file that already exists it will be mounted by bootsrap
Assertions.assertNotNull(manager.getSegmentFiles(segmentToBootstrap));
}
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 b3b0f32d9140..52e4c3ee4894 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
@@ -72,16 +72,16 @@ public void testWeakReserveAndReclaim()
CacheEntry entry4 = new TestCacheEntry("4", 25);
CacheEntry entry5 = new TestCacheEntry("5", 25);
- location.reserveWeak(entry1);
- location.reserveWeak(entry2);
- location.reserveWeak(entry3);
- location.reserveWeak(entry4);
+ registerWeak(location, entry1);
+ registerWeak(location, entry2);
+ registerWeak(location, entry3);
+ registerWeak(location, entry4);
Assertions.assertEquals(100, location.currentWeakSizeBytes());
Assertions.assertTrue(location.isWeakReserved(entry1.getId()));
Assertions.assertTrue(location.isWeakReserved(entry2.getId()));
Assertions.assertTrue(location.isWeakReserved(entry3.getId()));
Assertions.assertTrue(location.isWeakReserved(entry4.getId()));
- location.reserveWeak(entry5);
+ registerWeak(location, entry5);
Assertions.assertEquals(100, location.currentWeakSizeBytes());
Assertions.assertFalse(location.isWeakReserved(entry1.getId()));
Assertions.assertTrue(location.isWeakReserved(entry2.getId()));
@@ -98,10 +98,10 @@ public void testRemoveFromHead()
CacheEntry entry3 = new TestCacheEntry("3", 25);
CacheEntry entry4 = new TestCacheEntry("4", 25);
- location.reserveWeak(entry1);
- location.reserveWeak(entry2);
- location.reserveWeak(entry3);
- location.reserveWeak(entry4);
+ registerWeak(location, entry1);
+ registerWeak(location, entry2);
+ registerWeak(location, entry3);
+ registerWeak(location, entry4);
Assertions.assertTrue(location.isWeakReserved(entry1.getId()));
Assertions.assertTrue(location.isWeakReserved(entry2.getId()));
Assertions.assertTrue(location.isWeakReserved(entry3.getId()));
@@ -126,10 +126,10 @@ public void testRemoveFromTail()
CacheEntry entry3 = new TestCacheEntry("3", 25);
CacheEntry entry4 = new TestCacheEntry("4", 25);
- location.reserveWeak(entry1);
- location.reserveWeak(entry2);
- location.reserveWeak(entry3);
- location.reserveWeak(entry4);
+ registerWeak(location, entry1);
+ registerWeak(location, entry2);
+ registerWeak(location, entry3);
+ registerWeak(location, entry4);
Assertions.assertTrue(location.isWeakReserved(entry1.getId()));
Assertions.assertTrue(location.isWeakReserved(entry2.getId()));
Assertions.assertTrue(location.isWeakReserved(entry3.getId()));
@@ -159,10 +159,10 @@ public void testRemoveRandom()
entries.add(entry3);
entries.add(entry4);
- location.reserveWeak(entry1);
- location.reserveWeak(entry2);
- location.reserveWeak(entry3);
- location.reserveWeak(entry4);
+ registerWeak(location, entry1);
+ registerWeak(location, entry2);
+ registerWeak(location, entry3);
+ registerWeak(location, entry4);
Assertions.assertTrue(location.isWeakReserved(entry1.getId()));
Assertions.assertTrue(location.isWeakReserved(entry2.getId()));
Assertions.assertTrue(location.isWeakReserved(entry3.getId()));
@@ -200,8 +200,8 @@ public void testBulkReservation()
Assertions.assertNotNull(closer.register(location.addWeakReservationHold(entry2.getId(), () -> entry2)));
Assertions.assertEquals(2, location.getWeakStats().getHoldCount());
Assertions.assertEquals(50, location.getWeakStats().getHoldBytes());
- Assertions.assertTrue(location.reserveWeak(entry3));
- Assertions.assertTrue(location.reserveWeak(entry4));
+ Assertions.assertTrue(registerWeak(location, entry3));
+ Assertions.assertTrue(registerWeak(location, entry4));
Assertions.assertEquals(100, location.currentWeakSizeBytes());
Assertions.assertEquals(2, location.getWeakStats().getHoldCount());
@@ -224,7 +224,7 @@ public void testBulkReservation()
Assertions.assertTrue(location.isWeakReserved(entry4.getId()));
Assertions.assertTrue(location.isWeakReserved(entry5.getId()));
- Assertions.assertTrue(location.reserveWeak(entry6));
+ Assertions.assertTrue(registerWeak(location, entry6));
Assertions.assertTrue(location.isWeakReserved(entry1.getId()));
Assertions.assertTrue(location.isWeakReserved(entry2.getId()));
@@ -249,11 +249,11 @@ public void testBulkReservation()
Assertions.assertTrue(location.isWeakReserved(entry7.getId()));
// all storage is held, cannot reserve
- Assertions.assertFalse(location.reserveWeak(entry8));
+ Assertions.assertFalse(registerWeak(location, entry8));
// release holds
CloseableUtils.closeAndWrapExceptions(closer);
- Assertions.assertTrue(location.reserveWeak(entry8));
+ Assertions.assertTrue(registerWeak(location, entry8));
Assertions.assertFalse(location.isWeakReserved(entry1.getId()));
Assertions.assertTrue(location.isWeakReserved(entry2.getId()));
@@ -361,7 +361,7 @@ public void testReserveWeakExistsConcurrency() throws ExecutionException, Interr
{
StorageLocation loc = new StorageLocation(tempDir, 1000L, null);
final TestSegmentCacheEntry entry = makeSegmentEntry("2024/2025", 10);
- loc.reserveWeak(entry);
+ registerWeak(loc, entry);
entry.mount(loc);
for (int i = 0; i < 1000; i++) {
@@ -399,7 +399,7 @@ public void testReclaimRestoreDoesNotCreateZombieEntries()
CacheEntry entry2 = new TestCacheEntry("2", 90);
CacheEntry entry3 = new TestCacheEntry("3", 20);
- location.reserveWeak(entry1);
+ registerWeak(location, entry1);
// hold entry2 so it cannot be evicted by reclaim
StorageLocation.ReservationHold hold2 = location.addWeakReservationHold(
entry2.getId(),
@@ -408,12 +408,12 @@ public void testReclaimRestoreDoesNotCreateZombieEntries()
// must free 20 bytes but can only evict entry1 (10). Fails and restores entry1
// where the bug was a mismatch caused by creating a new entry in the list but re-using the old entry for the map.
- Assertions.assertFalse(location.reserveWeak(entry3));
+ Assertions.assertFalse(registerWeak(location, entry3));
// the hand pointer reaches the new entry1, removes the old entry1 from the map which is a zombie, then wraps around
// to the same zombie entry1 again since its head — at which point the map no longer contains the ID and the defensive exception was
// thrown.
- Assertions.assertFalse(location.reserveWeak(entry3));
+ Assertions.assertFalse(registerWeak(location, entry3));
hold2.close();
}
@@ -425,7 +425,7 @@ public void testRemoveUnheldWeakEntry()
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(registerWeak(location, entry));
Assertions.assertTrue(location.isWeakReserved(entry.getId()));
Assertions.assertEquals(30, location.currentSizeBytes());
@@ -476,7 +476,7 @@ public void testAdjustReservationWeakEntry()
{
final StorageLocation location = new StorageLocation(tempDir, 100L, null);
final TestResizableCacheEntry entry = new TestResizableCacheEntry("a", 80);
- Assertions.assertTrue(location.reserveWeak(entry));
+ Assertions.assertTrue(registerWeak(location, entry));
Assertions.assertEquals(80, location.currentWeakSizeBytes());
location.adjustReservation(entry.getId(), 30);
@@ -541,7 +541,7 @@ public void testAdjustReservationWeakEntryShrinksHeldBytes() throws IOException
{
final StorageLocation location = new StorageLocation(tempDir, 100L, null);
final TestResizableCacheEntry entry = new TestResizableCacheEntry("a", 80);
- Assertions.assertTrue(location.reserveWeak(entry));
+ Assertions.assertTrue(registerWeak(location, entry));
// Acquire a hold BEFORE shrinking. trackWeakHold records 80 bytes against currHoldBytes.
final StorageLocation.ReservationHold hold = location.addWeakReservationHold(entry.getId(), () -> entry);
@@ -566,7 +566,7 @@ public void testAdjustReservationWeakEntryShrinksHeldBytesWithMultipleHolds() th
{
final StorageLocation location = new StorageLocation(tempDir, 100L, null);
final TestResizableCacheEntry entry = new TestResizableCacheEntry("a", 50);
- Assertions.assertTrue(location.reserveWeak(entry));
+ Assertions.assertTrue(registerWeak(location, entry));
// Two concurrent holds: trackWeakHold fires twice, so currHoldBytes = 2 * 50 = 100.
final StorageLocation.ReservationHold hold1 = location.addWeakReservationHold(entry.getId(), () -> entry);
@@ -629,9 +629,9 @@ public void testRemoveUnheldWeakEntryUnmountCascadeDoesNotThrowConcurrentModific
location.addWeakReservationHold(parent.getId(), () -> parent);
Assertions.assertNotNull(parentHold);
- // Child registered WITHOUT a hold (the bootstrap reserveWeak path that removeUnheldWeakEntry cleans up).
+ // Child left registered but unheld, the shape removeUnheldWeakEntry cleans up.
final CascadingUnmountCacheEntry child = new CascadingUnmountCacheEntry("child", 100L, parentHold);
- Assertions.assertTrue(location.reserveWeak(child));
+ Assertions.assertTrue(registerWeak(location, child));
child.mount(location);
Assertions.assertDoesNotThrow(() -> location.removeUnheldWeakEntry(child.getId()));
@@ -658,11 +658,11 @@ public void testReclaimUnmountCascadeDoesNotThrowConcurrentModification()
// Child registered unheld and mounted, sole holder of the parent's cache hold; it is the reclaim target.
final CascadingUnmountCacheEntry child = new CascadingUnmountCacheEntry("child", 40L, parentHold);
- Assertions.assertTrue(location.reserveWeak(child));
+ Assertions.assertTrue(registerWeak(location, child));
child.mount(location);
final CascadingUnmountCacheEntry filler = new CascadingUnmountCacheEntry("filler", 40L, null);
- Assertions.assertDoesNotThrow(() -> location.reserveWeak(filler));
+ Assertions.assertDoesNotThrow(() -> registerWeak(location, filler));
Assertions.assertTrue(child.wasUnmounted());
Assertions.assertTrue(parent.wasUnmounted());
@@ -757,6 +757,28 @@ public void unmount()
}
}
+ /**
+ * Register a weak entry the way production does — reserve under a hold, mount, release the hold — leaving it
+ * registered but unheld, which is the state reclaim and {@link StorageLocation#removeUnheldWeakEntry} act on.
+ * Returns false when the location could not accept the reservation at all.
+ */
+ private static boolean registerWeak(StorageLocation location, CacheEntry entry)
+ {
+ final StorageLocation.ReservationHold hold =
+ location.addWeakReservationHold(entry.getId(), () -> entry);
+ if (hold == null) {
+ return false;
+ }
+ try {
+ entry.mount(location);
+ }
+ catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ hold.close();
+ return true;
+ }
+
private static final class TestCacheEntry implements CacheEntry
{
private final StringCacheIdentifier id;