fix: separate storage location accounting for internal cache holds

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

@@ -362,7 +362,7 @@ private void doMount(StorageLocation mountLocation) throws IOException
     try {
       // 1. Cache holds on metadata + parents (prevents cache eviction of weak dependencies)
       final StorageLocation.ReservationHold<?> metadataHold =
-          mountLocation.addWeakReservationHoldIfExists(metadataEntry.getId());
+          mountLocation.addInternalWeakReservationHoldIfExists(metadataEntry.getId());
       if (metadataHold == null) {
         throw DruidException.defensive(
             "Cannot acquire metadata hold for [%s]; metadata entry not registered with location[%s]",
@@ -374,7 +374,7 @@ private void doMount(StorageLocation mountLocation) throws IOException
 
       for (PartialSegmentBundleCacheEntryIdentifier parentId : parentEntryIds) {
         final StorageLocation.ReservationHold<?> parentHold =
-            mountLocation.addWeakReservationHoldIfExists(parentId);
+            mountLocation.addInternalWeakReservationHoldIfExists(parentId);
         if (parentHold == null) {
           throw DruidException.defensive(
               "Cannot acquire parent hold for [%s]; parent entry not registered with location[%s]",

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

@@ -362,7 +362,7 @@ public void applyRule(String fingerprint, Set<String> selectedBundleNames)
     try {
       final StorageLocation.ReservationHold<PartialSegmentMetadataCacheEntry> newSelfHold;
       if (needsSelfHold) {
-        newSelfHold = loc.addWeakReservationHoldIfExists(id);
+        newSelfHold = loc.addInternalWeakReservationHoldIfExists(id);
         if (newSelfHold == null) {
           throw DruidException.defensive(
               "Failed to acquire self-referential rule-hold on partial metadata entry[%s]; entry is not weak-reserved",
@@ -376,7 +376,7 @@ public void applyRule(String fingerprint, Set<String> selectedBundleNames)
       final Map<String, StorageLocation.ReservationHold<PartialSegmentBundleCacheEntry>> acquired = new HashMap<>();
       for (String name : namesToAcquire) {
         final StorageLocation.ReservationHold<PartialSegmentBundleCacheEntry> h =
-            loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, name));
+            loc.addInternalWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, name));
         if (h != null) {
           acquired.put(name, h);
           uncommittedHolds.add(h);
@@ -463,7 +463,7 @@ private void reconcileLinkedBundlesWithSelection(
     final Map<String, StorageLocation.ReservationHold<PartialSegmentBundleCacheEntry>> acquired = new HashMap<>();
     for (String name : reconcileNames) {
       final StorageLocation.ReservationHold<PartialSegmentBundleCacheEntry> h =
-          loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, name));
+          loc.addInternalWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, name));
       if (h != null) {
         acquired.put(name, h);
         uncommittedHolds.add(h);
@@ -967,7 +967,7 @@ private void restoreBundlesFromDisk(StorageLocation location) throws IOException
         // 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)) {
+                 location.addInternalWeakReservationHold(bundle.getId(), () -> bundle)) {
           if (restoreHold == null) {
             throw DruidException.defensive(
                 "Failed to reserve bundle entry[%s] in location[%s] while restoring from disk",
@@ -1618,7 +1618,7 @@ void registerBundle(PartialSegmentBundleCacheEntry bundle)
     // ReservationHold on this bundle, so the weak entry is guaranteed present and this acquire cannot return null
     // for a "just evicted" reason.
     final StorageLocation.ReservationHold<PartialSegmentBundleCacheEntry> hold =
-        loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, bundleName));
+        loc.addInternalWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, bundleName));
     if (hold == null) {
       return;
     }

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

@@ -123,6 +123,14 @@ public class StorageLocation
   private final AtomicLong currWeakSizeBytes = new AtomicLong(0);
   private final AtomicLong currHoldCount = new AtomicLong(0);
   private final AtomicLong currHoldBytes = new AtomicLong(0);
+  /**
+   * The subset of {@link #currHoldCount}/{@link #currHoldBytes} that is structural rather than demand (a bundle
+   * pinning its metadata entry or a parent bundle, a partial-load rule pinning what it selected, bootstrap restoring
+   * a bundle). These pin an entry against {@link #reclaim} exactly as a query hold does, so they belong in the totals,
+   * but they are not somebody waiting on an entry, broken out so the two can be told apart.
+   */
+  private final AtomicLong currInternalHoldCount = new AtomicLong(0);
+  private final AtomicLong currInternalHoldBytes = new AtomicLong(0);
 
   private final AtomicReference<StaticStats> staticStats = new AtomicReference<>();
   private final AtomicReference<WeakStats> weakStats = new AtomicReference<>();
@@ -295,6 +303,25 @@ public boolean reserve(CacheEntry entry)
    */
   @Nullable
   public <T extends CacheEntry> ReservationHold<T> addWeakReservationHoldIfExists(CacheEntryIdentifier entryId)
+  {
+    return addWeakReservationHoldIfExists(entryId, false);
+  }
+
+  /**
+   * Effectively the same as {@link #addWeakReservationHoldIfExists(CacheEntryIdentifier)} only accounted differently,
+   * for internal use to protect an entry from {@link #reclaim}.
+   */
+  @Nullable
+  public <T extends CacheEntry> ReservationHold<T> addInternalWeakReservationHoldIfExists(CacheEntryIdentifier entryId)
+  {
+    return addWeakReservationHoldIfExists(entryId, true);
+  }
+
+  @Nullable
+  private <T extends CacheEntry> ReservationHold<T> addWeakReservationHoldIfExists(
+      CacheEntryIdentifier entryId,
+      boolean internal
+  )
   {
     lock.readLock().lock();
     try {
@@ -304,12 +331,17 @@ public <T extends CacheEntry> ReservationHold<T> addWeakReservationHoldIfExists(
 
       WeakCacheEntry existingEntry = weakCacheEntries.get(entryId);
       if (existingEntry != null && existingEntry.hold()) {
+        // visited is set for internal holds too: an entry a live dependency is pinning genuinely is in use, and that
+        // is what this flag tells the reclaim scan.
         existingEntry.visited = true;
-        trackWeakHold(existingEntry);
-        weakStats.getAndUpdate(s -> s.hit(existingEntry.cacheEntry.getSize()));
+        final long heldBytes = existingEntry.cacheEntry.getSize();
+        trackWeakHold(heldBytes, internal);
+        if (!internal) {
+          weakStats.getAndUpdate(s -> s.hit(heldBytes));
+        }
         return new ReservationHold<>(
             (T) existingEntry.cacheEntry,
-            createWeakEntryReleaseRunnable(existingEntry, false)
+            createWeakEntryReleaseRunnable(existingEntry, false, internal, heldBytes)
         );
       }
       return null;
@@ -333,7 +365,30 @@ public <T extends CacheEntry> ReservationHold<T> addWeakReservationHold(
       Supplier<? extends CacheEntry> entrySupplier
   )
   {
-    final ReservationHold<T> existingEntry = addWeakReservationHoldIfExists(entryId);
+    return addWeakReservationHold(entryId, entrySupplier, false);
+  }
+
+  /**
+   * Internal version of {@link #addWeakReservationHold(CacheEntryIdentifier, Supplier)}, see
+   * {@link #addInternalWeakReservationHoldIfExists}.
+   */
+  @Nullable
+  public <T extends CacheEntry> ReservationHold<T> addInternalWeakReservationHold(
+      CacheEntryIdentifier entryId,
+      Supplier<? extends CacheEntry> entrySupplier
+  )
+  {
+    return addWeakReservationHold(entryId, entrySupplier, true);
+  }
+
+  @Nullable
+  private <T extends CacheEntry> ReservationHold<T> addWeakReservationHold(
+      CacheEntryIdentifier entryId,
+      Supplier<? extends CacheEntry> entrySupplier,
+      boolean internal
+  )
+  {
+    final ReservationHold<T> existingEntry = addWeakReservationHoldIfExists(entryId, internal);
     if (existingEntry != null) {
       return existingEntry;
     }
@@ -343,11 +398,14 @@ public <T extends CacheEntry> ReservationHold<T> addWeakReservationHold(
       WeakCacheEntry retryExistingEntry = weakCacheEntries.get(entryId);
       if (retryExistingEntry != null && retryExistingEntry.hold()) {
         retryExistingEntry.visited = true;
-        trackWeakHold(retryExistingEntry);
-        weakStats.getAndUpdate(s -> s.hit(retryExistingEntry.cacheEntry.getSize()));
+        final long heldBytes = retryExistingEntry.cacheEntry.getSize();
+        trackWeakHold(heldBytes, internal);
+        if (!internal) {
+          weakStats.getAndUpdate(s -> s.hit(heldBytes));
+        }
         return new ReservationHold<>(
             (T) retryExistingEntry.cacheEntry,
-            createWeakEntryReleaseRunnable(retryExistingEntry, false)
+            createWeakEntryReleaseRunnable(retryExistingEntry, false, internal, heldBytes)
         );
       }
       final CacheEntry newEntry = entrySupplier.get();
@@ -359,11 +417,12 @@ public <T extends CacheEntry> ReservationHold<T> addWeakReservationHold(
         newWeakEntry.hold();
         linkNewWeakEntry(newWeakEntry);
         weakCacheEntries.put(newEntry.getId(), newWeakEntry);
-        trackWeakHold(newWeakEntry);
-        weakStats.getAndUpdate(s -> s.loadBegin(newEntry.getSize()));
+        final long heldBytes = newEntry.getSize();
+        trackWeakHold(heldBytes, internal);
+        weakStats.getAndUpdate(s -> s.loadBegin(heldBytes));
         hold = new ReservationHold<>(
             (T) newEntry,
-            createWeakEntryReleaseRunnable(newWeakEntry, true)
+            createWeakEntryReleaseRunnable(newWeakEntry, true, internal, heldBytes)
         );
       } else {
         weakStats.getAndUpdate(WeakStats::reject);
@@ -437,14 +496,6 @@ public void adjustReservation(CacheEntryIdentifier id, long newSize)
         currWeakSizeBytes.getAndAdd(-delta);
         // The reservation (loadBegin) was recorded at the pre-shrink size; correct its byte total to match.
         weakStats.getAndUpdate(s -> s.shrinkLoadBegin(delta));
-        // Each active hold contributed entry.getSize() to currHoldBytes via trackWeakHold; shrink each hold's
-        // contribution by the same delta so a future trackWeakRelease (which subtracts the new smaller size) lands
-        // on the correct total. Clamp at 0 defensively against any pre-existing drift.
-        final long activeHolds = weak.holdReferents.getRegisteredParties() - 1L;
-        if (activeHolds > 0) {
-          final long holdDelta = delta * activeHolds;
-          currHoldBytes.updateAndGet(v -> Math.max(0L, v - holdDelta));
-        }
       }
     }
     finally {
@@ -538,12 +589,14 @@ private void unmountEvictedWeakEntry(@Nullable WeakCacheEntry evicted)
    */
   private Runnable createWeakEntryReleaseRunnable(
       final WeakCacheEntry weakEntry,
-      final boolean isNewEntry
+      final boolean isNewEntry,
+      final boolean internal,
+      final long heldBytes
   )
   {
     return () -> {
       weakEntry.release();
-      trackWeakRelease(weakEntry);
+      trackWeakRelease(heldBytes, internal);
 
       if (!isNewEntry && !areWeakEntriesEphemeral) {
         // No need to consider removal from weakCacheEntries on hold release.
@@ -786,16 +839,31 @@ public void trackWeakRangeRead(long bytes, long nanos)
     weakStats.getAndUpdate(s -> s.rangeRead(bytes, nanos));
   }
 
-  private void trackWeakHold(WeakCacheEntry entry)
+  /**
+   * {@code heldBytes} is the entry's size as of when the hold was taken, and the matching
+   * {@link #trackWeakRelease} subtracts that same number rather than re-reading the entry. The size can change under
+   * a live hold ({@link #adjustReservation} shrinks a partial segment's pessimistic estimate once its real size is
+   * known), and re-reading it would leave the totals permanently skewed by the difference. Pairing each add with an
+   * identical subtract keeps them balanced without either side taking a lock.
+   */
+  private void trackWeakHold(long heldBytes, boolean internal)
   {
     currHoldCount.getAndIncrement();
-    currHoldBytes.getAndAdd(entry.cacheEntry.getSize());
+    currHoldBytes.getAndAdd(heldBytes);
+    if (internal) {
+      currInternalHoldCount.getAndIncrement();
+      currInternalHoldBytes.getAndAdd(heldBytes);
+    }
   }
 
-  private void trackWeakRelease(WeakCacheEntry entry)
+  private void trackWeakRelease(long heldBytes, boolean internal)
   {
     currHoldCount.getAndDecrement();
-    currHoldBytes.getAndAdd(-entry.cacheEntry.getSize());
+    currHoldBytes.getAndAdd(-heldBytes);
+    if (internal) {
+      currInternalHoldCount.getAndDecrement();
+      currInternalHoldBytes.getAndAdd(-heldBytes);
+    }
   }
 
   @VisibleForTesting
@@ -842,6 +910,8 @@ public void reset()
     currStaticSizeBytes.set(0);
     currHoldCount.set(0);
     currHoldBytes.set(0);
+    currInternalHoldCount.set(0);
+    currInternalHoldBytes.set(0);
     resetStaticStats();
     resetWeakStats();
   }
@@ -858,7 +928,9 @@ public StaticStats resetStaticStats()
 
   public WeakStats resetWeakStats()
   {
-    return weakStats.getAndSet(new WeakStats(currWeakSizeBytes, currHoldCount, currHoldBytes));
+    return weakStats.getAndSet(
+        new WeakStats(currWeakSizeBytes, currHoldCount, currHoldBytes, currInternalHoldCount, currInternalHoldBytes)
+    );
   }
 
   /**
@@ -1222,6 +1294,8 @@ public static final class WeakStats implements VirtualStorageLocationStats
     private final AtomicLong sizeUsed;
     private final AtomicLong holdCount;
     private final AtomicLong holdBytes;
+    private final AtomicLong internalHoldCount;
+    private final AtomicLong internalHoldBytes;
     private final AtomicLong loadBeginCount = new AtomicLong(0);
     private final AtomicLong loadBeginBytes = new AtomicLong(0);
     private final AtomicLong loadCount = new AtomicLong(0);
@@ -1236,11 +1310,19 @@ public static final class WeakStats implements VirtualStorageLocationStats
     private final AtomicLong readBytes = new AtomicLong(0);
     private final AtomicLong readTimeNanos = new AtomicLong(0);
 
-    public WeakStats(AtomicLong sizeUsed, AtomicLong holdCount, AtomicLong holdBytes)
+    public WeakStats(
+        AtomicLong sizeUsed,
+        AtomicLong holdCount,
+        AtomicLong holdBytes,
+        AtomicLong internalHoldCount,
+        AtomicLong internalHoldBytes
+    )
     {
       this.sizeUsed = sizeUsed;
       this.holdCount = holdCount;
       this.holdBytes = holdBytes;
+      this.internalHoldCount = internalHoldCount;
+      this.internalHoldBytes = internalHoldBytes;
     }
 
     public WeakStats hit(long size)
@@ -1320,6 +1402,18 @@ public long getHoldBytes()
       return holdBytes.get();
     }
 
+    @Override
+    public long getInternalHoldCount()
+    {
+      return internalHoldCount.get();
+    }
+
+    @Override
+    public long getInternalHoldBytes()
+    {
+      return internalHoldBytes.get();
+    }
+
     @Override
     public long getHitCount()
     {

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

@@ -39,6 +39,17 @@ public interface VirtualStorageLocationStats
    */
   long getHoldBytes();
 
+  /**
+   * internal holds, such as those which one cache entry places on another it depends on, or that a partial-load rule
+   * places on what it selected, rather than a caller waiting on the entry.
+   */
+  long getInternalHoldCount();
+
+  /**
+   * Total bytes from the holds counted by {@link #getInternalHoldCount()}.
+   */
+  long getInternalHoldBytes();
+
   /**
    * Number of operations for which an entry was already present during the measurement period
    */

server/src/main/java/org/apache/druid/server/metrics/StorageMonitor.java

@@ -100,6 +100,18 @@ public class StorageMonitor extends AbstractMonitor
    */
   public static final String VSF_HOLD_BYTES = "storage/virtual/hold/bytes";
 
+  /**
+   * internal holds, such as one cache entry places on another it depends on, or that a partial-load rule places on
+   * what it selected, rather than a caller waiting on the entry. Counted in the totals above because they pin an
+   * entry against reclaim identically.
+   */
+  public static final String VSF_INTERNAL_HOLD_COUNT = "storage/virtual/hold/internal/count";
+
+  /**
+   * Total bytes from the holds counted by {@link #VSF_INTERNAL_HOLD_COUNT}.
+   */
+  public static final String VSF_INTERNAL_HOLD_BYTES = "storage/virtual/hold/internal/bytes";
+
   /**
    * Number of acquire operations during the measurement period that found an existing weakly-held entry already in
    * virtual storage.
@@ -215,6 +227,8 @@ public boolean doMonitor(ServiceEmitter emitter)
         emitter.emit(builder.setMetric(VSF_USED_BYTES, weakStats.getUsedBytes()));
         emitter.emit(builder.setMetric(VSF_HOLD_COUNT, weakStats.getHoldCount()));
         emitter.emit(builder.setMetric(VSF_HOLD_BYTES, weakStats.getHoldBytes()));
+        emitter.emit(builder.setMetric(VSF_INTERNAL_HOLD_COUNT, weakStats.getInternalHoldCount()));
+        emitter.emit(builder.setMetric(VSF_INTERNAL_HOLD_BYTES, weakStats.getInternalHoldBytes()));
         emitter.emit(builder.setMetric(VSF_HIT_COUNT, weakStats.getHitCount()));
         emitter.emit(builder.setMetric(VSF_HIT_BYTES, weakStats.getHitBytes()));
         emitter.emit(builder.setMetric(VSF_LOAD_BEGIN_COUNT, weakStats.getLoadBeginCount()));

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

@@ -336,7 +336,7 @@ void testRollbackRemovesBundlesItAlreadyMounted() throws IOException
     // 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());
+           .addInternalWeakReservationHold(ArgumentMatchers.eq(aggId), ArgumentMatchers.any());
 
     Assertions.assertThrows(Throwable.class, () -> restoreFromDisk(location));