fix: partial segment bootstrap

server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java

@@ -30,12 +30,10 @@
 import javax.annotation.Nullable;
 import java.io.File;
 import java.io.IOException;
-import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Objects;
 import java.util.Set;
 
 /**
@@ -74,12 +72,13 @@ public final class PartialSegmentCacheBootstrap
    * the runtime acquire path), so it is evictable once mounted; bootstrap-restored data is treated as a cache
    * optimization, not a permanent fixture. The caller is expected to drive the actual mount through
    * {@code metadata.mount(location)} later (typically via {@code SegmentCacheManager#bootstrap}), which builds the
-   * file mapper and cascades into {@link #restoreBundlesFromDisk}.
+   * 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
@@ -92,6 +91,7 @@ public static PartialSegmentMetadataCacheEntry reserveFromDisk(
       File localCacheDir,
       String targetFilename,
       List<String> externalFilenames,
+      SegmentRangeReader rangeReader,
       ObjectMapper jsonMapper,
       @Nullable StorageLoadingThreadPool storagePool,
       StorageLocation location
@@ -113,7 +113,7 @@ public static PartialSegmentMetadataCacheEntry reserveFromDisk(
         localCacheDir,
         targetFilename,
         externalFilenames,
-        BootstrapRangeReader.INSTANCE,
+        rangeReader,
         jsonMapper,
         storagePool,
         actualMetadataSize
@@ -340,25 +340,4 @@ private PartialSegmentCacheBootstrap()
     // utility class
   }
 
-  /**
-   * Stub range reader used during bootstrap: the on-disk header is expected to exist and parse, so no fetch is needed.
-   * If for any reason {@link PartialSegmentFileMapperV10#create} decides to re-fetch (e.g. header corruption), this
-   * reader throws so we fail loudly rather than silently re-downloading without the operator's knowledge.
-   */
-  private static final class BootstrapRangeReader implements SegmentRangeReader
-  {
-    static final BootstrapRangeReader INSTANCE = new BootstrapRangeReader();
-
-    @Override
-    public InputStream readRange(String filename, long offset, long length)
-    {
-      throw DruidException.defensive(
-          "BootstrapRangeReader was asked to fetch [%s] @[%d:%d]; bootstrap should only read from local disk",
-          Objects.toString(filename),
-          offset,
-          length
-      );
-    }
-  }
-
 }

server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java

@@ -327,13 +327,36 @@ 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;
+        }
         removeInfo = false;
         try {
           PartialSegmentCacheBootstrap.reserveFromDisk(
               segment.getId(),
               partialDir,
               IndexIO.V10_FILE_NAME,
               List.of(),
+              rangeReader,
               jsonMapper,
               virtualStorageLoadingThreadPool,
               location
@@ -1121,8 +1144,8 @@ boolean isSegmentCached(final DataSegment segment)
 
   /**
    * Testing use only please, any callers that want to do stuff with segments should use
-   * {@link #acquireCachedSegment(SegmentId)} or {@link #acquireSegment(DataSegment)} instead. Does not hold locks
-   * and so is not really safe to use while the cache manager is active
+   * {@link #acquireCachedSegment(SegmentId, AcquireMode)} or {@link #acquireSegment(DataSegment, AcquireMode)} instead.
+   * Does not hold locks and so is not really safe to use while the cache manager is active
    */
   @VisibleForTesting
   @Nullable

server/src/test/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrapTest.java

@@ -335,22 +335,26 @@ 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).
+    // 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.
+    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
     );
     Assertions.assertTrue(location.isWeakReserved(id));
 
-    // Delete the header so the mount's file-mapper build must re-fetch it from deep storage; the bootstrap
-    // disk-only range reader throws, failing the mount (the same poison that arises when create() detects header
-    // corruption and tries to re-fetch). The mount rollback deletes the header, so the entry is now both unmounted
-    // and un-re-mountable.
+    // 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
+    // fetch fails). The mount rollback deletes the header, so the entry is now both unmounted and un-re-mountable.
     final File headerFile = new File(
         cacheDir,
         IndexIO.V10_FILE_NAME + PartialSegmentFileMapperV10.METADATA_HEADER_SUFFIX
@@ -360,13 +364,43 @@ void testMountFailureRemovesLingeringWeakEntry() throws IOException
     Assertions.assertThrows(Throwable.class, () -> metadata.mount(location));
 
     // The failed mount must not leave the lingering weak entry behind: a later findExistingPartialWithHold would
-    // otherwise resurrect it and re-mount would fail forever (disk-only reader + deleted header). It must be gone so
+    // otherwise resurrect it and re-mount would fail forever (failing reader + deleted header). It must be gone so
     // the cold acquire path rebuilds a fresh, deep-storage-capable entry via reservePartial.
     Assertions.assertEquals(0, location.getWeakEntryCount(), "failed mount must remove the lingering weak entry");
     Assertions.assertFalse(location.isWeakReserved(id));
     Assertions.assertEquals(0, location.currentSizeBytes(), "failed mount must release the reservation");
   }
 
+  @Test
+  void testRestoredEntryCanLazilyFetchUndownloadedFile() throws IOException
+  {
+    // a bootstrap-restored entry must keep the segment's real deep-storage range reader so a later query can lazily
+    // 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".
+    primeOnDiskState();
+    final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE * 8, null);
+    final PartialSegmentMetadataCacheEntry metadata = restoreFromDisk(location);
+
+    final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper();
+    Assertions.assertNotNull(mapper, "restored entry must be mounted");
+
+    // Priming sparse-allocated the containers but downloaded no column data, so every internal file is un-downloaded.
+    final String fileToFetch = mapper.getSegmentFileMetadata().getFiles().keySet().stream()
+                                     .filter(f -> f.startsWith(Projections.BASE_TABLE_PROJECTION_NAME + "/"))
+                                     .findFirst()
+                                     .orElseThrow();
+    Assertions.assertFalse(
+        mapper.getDownloadedFiles().contains(fileToFetch),
+        "precondition: " + fileToFetch + " should not be downloaded yet"
+    );
+
+    Assertions.assertNotNull(
+        mapper.mapFile(fileToFetch),
+        "restored entry must lazily fetch an un-downloaded file from deep storage"
+    );
+    Assertions.assertTrue(mapper.getDownloadedFiles().contains(fileToFetch));
+  }
+
   @Test
   void testReserveFailsWhenHeaderMissing()
   {
@@ -379,6 +413,7 @@ void testReserveFailsWhenHeaderMissing()
             cacheDir,
             IndexIO.V10_FILE_NAME,
             List.of(),
+            new DirectoryBackedRangeReader(deepStorageDir),
             JSON_MAPPER,
             null,
             location
@@ -521,6 +556,7 @@ private PartialSegmentMetadataCacheEntry restoreFromDisk(StorageLocation locatio
         cacheDir,
         IndexIO.V10_FILE_NAME,
         List.of(),
+        new DirectoryBackedRangeReader(deepStorageDir),
         JSON_MAPPER,
         null,
         location

server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java

@@ -792,6 +792,76 @@ void testGetCachedSegmentsDeletesPartialLayoutWhenPartialDisabled() throws IOExc
     }
   }
 
+  @Test
+  void testGetCachedSegmentsDeletesPartialLayoutWhenRangeReaderUnavailable() 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());
+    FileUtils.mkdirp(partialDir);
+    primePartialOnDiskState(partialDir);
+
+    // ...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.
+    final File noRangeReaderStorage = new File(perTestTempDir, "no_range_reader_storage");
+    FileUtils.mkdirp(noRangeReaderStorage);
+    final DataSegment unreadableSegment =
+        DataSegment.builder(SEGMENT_ID)
+                   .shardSpec(NoneShardSpec.instance())
+                   .loadSpec(Map.of("type", "local", "path", noRangeReaderStorage.getAbsolutePath()))
+                   .size(0)
+                   .build();
+    manager.storeInfoFile(unreadableSegment);
+    Assertions.assertTrue(PartialSegmentCacheBootstrap.isPartialSegmentLayout(partialDir, IndexIO.V10_FILE_NAME));
+
+    final List<DataSegment> 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.assertFalse(
+        partialDir.exists(),
+        "bootstrap must delete the unusable partial layout from disk"
+    );
+    Assertions.assertNull(
+        manager.getLocations().get(0).getCacheEntry(new SegmentCacheEntryIdentifier(SEGMENT_ID)),
+        "no cache entry should be reserved for the deleted partial layout"
+    );
+  }
+
+  @Test
+  void testGetCachedSegmentsDeletesPartialLayoutWhenLoadSpecUnconvertible() throws IOException
+  {
+    // Prime a valid partial on-disk layout as a previous run left it...
+    final File partialDir = new File(cacheRoot, SEGMENT_ID.toString());
+    FileUtils.mkdirp(partialDir);
+    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.
+    final DataSegment unconvertibleSegment =
+        DataSegment.builder(SEGMENT_ID)
+                   .shardSpec(NoneShardSpec.instance())
+                   .loadSpec(Map.of("type", "no-such-loadspec-type"))
+                   .size(0)
+                   .build();
+    manager.storeInfoFile(unconvertibleSegment);
+    Assertions.assertTrue(PartialSegmentCacheBootstrap.isPartialSegmentLayout(partialDir, IndexIO.V10_FILE_NAME));
+
+    final List<DataSegment> 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.assertFalse(partialDir.exists(), "bootstrap must delete the unusable partial layout from disk");
+    Assertions.assertNull(
+        manager.getLocations().get(0).getCacheEntry(new SegmentCacheEntryIdentifier(SEGMENT_ID)),
+        "no cache entry should be reserved for the deleted partial layout"
+    );
+  }
+
   /**
    * Lay down the on-disk artifacts a previous process run would have left behind in the given partial directory:
    * a V10 header file and sparse-allocated container files for every container the segment metadata declares. The