From 9f87049475bf8ad4b4c50605cdf4454c4828e006 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Tue, 25 Aug 2026 12:29:57 -0700 Subject: [PATCH 1/2] fix: separate storage location accounting for internal cache holds --- .../PartialSegmentBundleCacheEntry.java | 4 +- .../PartialSegmentMetadataCacheEntry.java | 10 +- .../segment/loading/StorageLocation.java | 135 ++++++++++++++++-- .../loading/VirtualStorageLocationStats.java | 11 ++ .../druid/server/metrics/StorageMonitor.java | 14 ++ .../PartialSegmentRestoreFromDiskTest.java | 2 +- .../segment/loading/StorageLocationTest.java | 35 +++++ 7 files changed, 188 insertions(+), 23 deletions(-) diff --git a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java index a5fe04998281..00d7f90170de 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java +++ b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java @@ -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]", 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 8c4d827fb693..1e11d950fe1a 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 @@ -362,7 +362,7 @@ public void applyRule(String fingerprint, Set selectedBundleNames) try { final StorageLocation.ReservationHold 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 selectedBundleNames) final Map> acquired = new HashMap<>(); for (String name : namesToAcquire) { final StorageLocation.ReservationHold 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> acquired = new HashMap<>(); for (String name : reconcileNames) { final StorageLocation.ReservationHold 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 hold = - loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, bundleName)); + loc.addInternalWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, bundleName)); if (hold == null) { return; } 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 0c3e4cc8d77c..ed350ee4e972 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 @@ -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 = new AtomicReference<>(); private final AtomicReference weakStats = new AtomicReference<>(); @@ -295,6 +303,25 @@ public boolean reserve(CacheEntry entry) */ @Nullable public ReservationHold 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 ReservationHold addInternalWeakReservationHoldIfExists(CacheEntryIdentifier entryId) + { + return addWeakReservationHoldIfExists(entryId, true); + } + + @Nullable + private ReservationHold addWeakReservationHoldIfExists( + CacheEntryIdentifier entryId, + boolean internal + ) { lock.readLock().lock(); try { @@ -304,12 +331,16 @@ public ReservationHold 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())); + trackWeakHold(existingEntry, internal); + if (!internal) { + weakStats.getAndUpdate(s -> s.hit(existingEntry.cacheEntry.getSize())); + } return new ReservationHold<>( (T) existingEntry.cacheEntry, - createWeakEntryReleaseRunnable(existingEntry, false) + createWeakEntryReleaseRunnable(existingEntry, false, internal) ); } return null; @@ -333,7 +364,30 @@ public ReservationHold addWeakReservationHold( Supplier entrySupplier ) { - final ReservationHold existingEntry = addWeakReservationHoldIfExists(entryId); + return addWeakReservationHold(entryId, entrySupplier, false); + } + + /** + * Internal version of {@link #addWeakReservationHold(CacheEntryIdentifier, Supplier)}, see + * {@link #addInternalWeakReservationHoldIfExists}. + */ + @Nullable + public ReservationHold addInternalWeakReservationHold( + CacheEntryIdentifier entryId, + Supplier entrySupplier + ) + { + return addWeakReservationHold(entryId, entrySupplier, true); + } + + @Nullable + private ReservationHold addWeakReservationHold( + CacheEntryIdentifier entryId, + Supplier entrySupplier, + boolean internal + ) + { + final ReservationHold existingEntry = addWeakReservationHoldIfExists(entryId, internal); if (existingEntry != null) { return existingEntry; } @@ -343,11 +397,13 @@ public ReservationHold addWeakReservationHold( WeakCacheEntry retryExistingEntry = weakCacheEntries.get(entryId); if (retryExistingEntry != null && retryExistingEntry.hold()) { retryExistingEntry.visited = true; - trackWeakHold(retryExistingEntry); - weakStats.getAndUpdate(s -> s.hit(retryExistingEntry.cacheEntry.getSize())); + trackWeakHold(retryExistingEntry, internal); + if (!internal) { + weakStats.getAndUpdate(s -> s.hit(retryExistingEntry.cacheEntry.getSize())); + } return new ReservationHold<>( (T) retryExistingEntry.cacheEntry, - createWeakEntryReleaseRunnable(retryExistingEntry, false) + createWeakEntryReleaseRunnable(retryExistingEntry, false, internal) ); } final CacheEntry newEntry = entrySupplier.get(); @@ -359,11 +415,11 @@ public ReservationHold addWeakReservationHold( newWeakEntry.hold(); linkNewWeakEntry(newWeakEntry); weakCacheEntries.put(newEntry.getId(), newWeakEntry); - trackWeakHold(newWeakEntry); + trackWeakHold(newWeakEntry, internal); weakStats.getAndUpdate(s -> s.loadBegin(newEntry.getSize())); hold = new ReservationHold<>( (T) newEntry, - createWeakEntryReleaseRunnable(newWeakEntry, true) + createWeakEntryReleaseRunnable(newWeakEntry, true, internal) ); } else { weakStats.getAndUpdate(WeakStats::reject); @@ -445,6 +501,12 @@ public void adjustReservation(CacheEntryIdentifier id, long newSize) final long holdDelta = delta * activeHolds; currHoldBytes.updateAndGet(v -> Math.max(0L, v - holdDelta)); } + // Same correction for internal holds + final long activeInternalHolds = weak.internalHolds.get(); + if (activeInternalHolds > 0) { + final long internalHoldDelta = delta * activeInternalHolds; + currInternalHoldBytes.updateAndGet(v -> Math.max(0L, v - internalHoldDelta)); + } } } finally { @@ -538,12 +600,13 @@ private void unmountEvictedWeakEntry(@Nullable WeakCacheEntry evicted) */ private Runnable createWeakEntryReleaseRunnable( final WeakCacheEntry weakEntry, - final boolean isNewEntry + final boolean isNewEntry, + final boolean internal ) { return () -> { weakEntry.release(); - trackWeakRelease(weakEntry); + trackWeakRelease(weakEntry, internal); if (!isNewEntry && !areWeakEntriesEphemeral) { // No need to consider removal from weakCacheEntries on hold release. @@ -783,16 +846,26 @@ public void trackWeakRangeRead(long bytes, long nanos) weakStats.getAndUpdate(s -> s.rangeRead(bytes, nanos)); } - private void trackWeakHold(WeakCacheEntry entry) + private void trackWeakHold(WeakCacheEntry entry, boolean internal) { currHoldCount.getAndIncrement(); currHoldBytes.getAndAdd(entry.cacheEntry.getSize()); + if (internal) { + entry.internalHolds.getAndIncrement(); + currInternalHoldCount.getAndIncrement(); + currInternalHoldBytes.getAndAdd(entry.cacheEntry.getSize()); + } } - private void trackWeakRelease(WeakCacheEntry entry) + private void trackWeakRelease(WeakCacheEntry entry, boolean internal) { currHoldCount.getAndDecrement(); currHoldBytes.getAndAdd(-entry.cacheEntry.getSize()); + if (internal) { + entry.internalHolds.getAndDecrement(); + currInternalHoldCount.getAndDecrement(); + currInternalHoldBytes.getAndAdd(-entry.cacheEntry.getSize()); + } } @VisibleForTesting @@ -839,6 +912,8 @@ public void reset() currStaticSizeBytes.set(0); currHoldCount.set(0); currHoldBytes.set(0); + currInternalHoldCount.set(0); + currInternalHoldBytes.set(0); resetStaticStats(); resetWeakStats(); } @@ -855,7 +930,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) + ); } /** @@ -1049,6 +1126,12 @@ protected boolean onAdvance(int phase, int registeredParties) */ private volatile boolean visited; + /** + * How many of this entry's holds are internal, so {@code adjustReservation} can correct the internal byte total + * by the same per-hold delta it applies to the overall one. + */ + private final AtomicLong internalHolds = new AtomicLong(0); + private WeakCacheEntry(CacheEntry cacheEntry) { this.cacheEntry = cacheEntry; @@ -1219,6 +1302,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); @@ -1233,11 +1318,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) @@ -1317,6 +1410,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() { diff --git a/server/src/main/java/org/apache/druid/segment/loading/VirtualStorageLocationStats.java b/server/src/main/java/org/apache/druid/segment/loading/VirtualStorageLocationStats.java index 035d30caffff..a68516e94f16 100644 --- a/server/src/main/java/org/apache/druid/segment/loading/VirtualStorageLocationStats.java +++ b/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 */ diff --git a/server/src/main/java/org/apache/druid/server/metrics/StorageMonitor.java b/server/src/main/java/org/apache/druid/server/metrics/StorageMonitor.java index 23a8bd4bbf4c..c44f93643559 100644 --- a/server/src/main/java/org/apache/druid/server/metrics/StorageMonitor.java +++ b/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())); diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentRestoreFromDiskTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentRestoreFromDiskTest.java index bb8b4e9f66f8..52611f30519f 100644 --- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentRestoreFromDiskTest.java +++ b/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)); 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 52e4c3ee4894..a8b31e8f9f11 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 @@ -584,6 +584,41 @@ public void testAdjustReservationWeakEntryShrinksHeldBytesWithMultipleHolds() th Assertions.assertEquals(0, location.getWeakStats().getHoldBytes()); } + @Test + public void testInternalHoldsAreNotCountedAsHitsButAreCountedAsPinned() + { + final StorageLocation location = new StorageLocation(tempDir, 100L, null); + final UnmountTrackingCacheEntry entry = new UnmountTrackingCacheEntry("a", 10); + final StorageLocation.ReservationHold reserver = + location.addWeakReservationHold(entry.getId(), () -> entry); + Assertions.assertNotNull(reserver); + + // A demand hold: somebody is waiting on this entry, so it is a cache hit. + final StorageLocation.ReservationHold query = location.addWeakReservationHoldIfExists(entry.getId()); + Assertions.assertNotNull(query); + Assertions.assertEquals(1, location.getWeakStats().getHitCount()); + Assertions.assertEquals(0, location.getWeakStats().getInternalHoldCount()); + + // A structural hold: another entry pinning this one. It pins against reclaim just the same, so it counts in the + // totals, but it is not demand and must not move the hit rate. + final StorageLocation.ReservationHold internal = + location.addInternalWeakReservationHoldIfExists(entry.getId()); + Assertions.assertNotNull(internal); + Assertions.assertEquals(1, location.getWeakStats().getHitCount(), "a structural hold is not a cache hit"); + Assertions.assertEquals(3, location.getWeakStats().getHoldCount(), "but it does pin the entry"); + Assertions.assertEquals(1, location.getWeakStats().getInternalHoldCount()); + Assertions.assertEquals(10, location.getWeakStats().getInternalHoldBytes()); + + internal.close(); + Assertions.assertEquals(0, location.getWeakStats().getInternalHoldCount()); + Assertions.assertEquals(0, location.getWeakStats().getInternalHoldBytes()); + Assertions.assertEquals(2, location.getWeakStats().getHoldCount()); + + query.close(); + reserver.close(); + Assertions.assertEquals(0, location.getWeakStats().getHoldCount()); + } + @Test public void testEphemeralWeakEntryUnmountCascadeDoesNotThrowConcurrentModification() { From fa98c923243e142ca0037f7364fe7163533bdea2 Mon Sep 17 00:00:00 2001 From: Clint Wylie Date: Wed, 26 Aug 2026 10:38:07 -0700 Subject: [PATCH 2/2] fix race --- .../segment/loading/StorageLocation.java | 67 ++++++++----------- .../segment/loading/StorageLocationTest.java | 66 ++++++++++++++++-- 2 files changed, 88 insertions(+), 45 deletions(-) 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 3c729cbaeb9a..758d957a0144 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 @@ -334,13 +334,14 @@ private ReservationHold addWeakReservationHoldIfExists // 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, internal); + final long heldBytes = existingEntry.cacheEntry.getSize(); + trackWeakHold(heldBytes, internal); if (!internal) { - weakStats.getAndUpdate(s -> s.hit(existingEntry.cacheEntry.getSize())); + weakStats.getAndUpdate(s -> s.hit(heldBytes)); } return new ReservationHold<>( (T) existingEntry.cacheEntry, - createWeakEntryReleaseRunnable(existingEntry, false, internal) + createWeakEntryReleaseRunnable(existingEntry, false, internal, heldBytes) ); } return null; @@ -397,13 +398,14 @@ private ReservationHold addWeakReservationHold( WeakCacheEntry retryExistingEntry = weakCacheEntries.get(entryId); if (retryExistingEntry != null && retryExistingEntry.hold()) { retryExistingEntry.visited = true; - trackWeakHold(retryExistingEntry, internal); + final long heldBytes = retryExistingEntry.cacheEntry.getSize(); + trackWeakHold(heldBytes, internal); if (!internal) { - weakStats.getAndUpdate(s -> s.hit(retryExistingEntry.cacheEntry.getSize())); + weakStats.getAndUpdate(s -> s.hit(heldBytes)); } return new ReservationHold<>( (T) retryExistingEntry.cacheEntry, - createWeakEntryReleaseRunnable(retryExistingEntry, false, internal) + createWeakEntryReleaseRunnable(retryExistingEntry, false, internal, heldBytes) ); } final CacheEntry newEntry = entrySupplier.get(); @@ -415,11 +417,12 @@ private ReservationHold addWeakReservationHold( newWeakEntry.hold(); linkNewWeakEntry(newWeakEntry); weakCacheEntries.put(newEntry.getId(), newWeakEntry); - trackWeakHold(newWeakEntry, internal); - 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, internal) + createWeakEntryReleaseRunnable(newWeakEntry, true, internal, heldBytes) ); } else { weakStats.getAndUpdate(WeakStats::reject); @@ -493,20 +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)); - } - // Same correction for internal holds - final long activeInternalHolds = weak.internalHolds.get(); - if (activeInternalHolds > 0) { - final long internalHoldDelta = delta * activeInternalHolds; - currInternalHoldBytes.updateAndGet(v -> Math.max(0L, v - internalHoldDelta)); - } } } finally { @@ -601,12 +590,13 @@ private void unmountEvictedWeakEntry(@Nullable WeakCacheEntry evicted) private Runnable createWeakEntryReleaseRunnable( final WeakCacheEntry weakEntry, final boolean isNewEntry, - final boolean internal + final boolean internal, + final long heldBytes ) { return () -> { weakEntry.release(); - trackWeakRelease(weakEntry, internal); + trackWeakRelease(heldBytes, internal); if (!isNewEntry && !areWeakEntriesEphemeral) { // No need to consider removal from weakCacheEntries on hold release. @@ -849,25 +839,30 @@ public void trackWeakRangeRead(long bytes, long nanos) weakStats.getAndUpdate(s -> s.rangeRead(bytes, nanos)); } - private void trackWeakHold(WeakCacheEntry entry, boolean internal) + /** + * {@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) { - entry.internalHolds.getAndIncrement(); currInternalHoldCount.getAndIncrement(); - currInternalHoldBytes.getAndAdd(entry.cacheEntry.getSize()); + currInternalHoldBytes.getAndAdd(heldBytes); } } - private void trackWeakRelease(WeakCacheEntry entry, boolean internal) + private void trackWeakRelease(long heldBytes, boolean internal) { currHoldCount.getAndDecrement(); - currHoldBytes.getAndAdd(-entry.cacheEntry.getSize()); + currHoldBytes.getAndAdd(-heldBytes); if (internal) { - entry.internalHolds.getAndDecrement(); currInternalHoldCount.getAndDecrement(); - currInternalHoldBytes.getAndAdd(-entry.cacheEntry.getSize()); + currInternalHoldBytes.getAndAdd(-heldBytes); } } @@ -1129,12 +1124,6 @@ protected boolean onAdvance(int phase, int registeredParties) */ private volatile boolean visited; - /** - * How many of this entry's holds are internal, so {@code adjustReservation} can correct the internal byte total - * by the same per-hold delta it applies to the overall one. - */ - private final AtomicLong internalHolds = new AtomicLong(0); - private WeakCacheEntry(CacheEntry cacheEntry) { this.cacheEntry = cacheEntry; 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 d5f184194794..23b89108e087 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 @@ -41,10 +41,12 @@ import java.util.Collections; import java.util.List; import java.util.Objects; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; class StorageLocationTest { @@ -549,12 +551,12 @@ public void testAdjustReservationWeakEntryShrinksHeldBytes() throws IOException Assertions.assertEquals(1, location.getWeakStats().getHoldCount()); Assertions.assertEquals(80, location.getWeakStats().getHoldBytes()); - // Shrink to 30: hold-bytes contribution from the active hold must shrink in lockstep so the eventual - // trackWeakRelease (which subtracts the new smaller size) leaves currHoldBytes at 0. + // Shrink to 30. The live hold keeps contributing the 80 it was taken at - it subtracts that same 80 on release, + // which is what keeps the total balanced without the resize having to reach into it. location.adjustReservation(entry.getId(), 30); Assertions.assertEquals(30, entry.getSize()); Assertions.assertEquals(30, location.currentWeakSizeBytes()); - Assertions.assertEquals(30, location.getWeakStats().getHoldBytes()); + Assertions.assertEquals(80, location.getWeakStats().getHoldBytes()); hold.close(); Assertions.assertEquals(0, location.getWeakStats().getHoldCount()); @@ -574,16 +576,68 @@ public void testAdjustReservationWeakEntryShrinksHeldBytesWithMultipleHolds() th Assertions.assertEquals(2, location.getWeakStats().getHoldCount()); Assertions.assertEquals(100, location.getWeakStats().getHoldBytes()); - // Shrink by 30 (50 → 20): each of the two active holds contributes -30, so currHoldBytes drops by 60. + // Shrink by 30 (50 -> 20). Both holds were taken at 50 and keep contributing it, so the total is untouched here + // and each release takes its own 50 back off. location.adjustReservation(entry.getId(), 20); - Assertions.assertEquals(40, location.getWeakStats().getHoldBytes()); + Assertions.assertEquals(100, location.getWeakStats().getHoldBytes()); hold1.close(); - Assertions.assertEquals(20, location.getWeakStats().getHoldBytes()); + Assertions.assertEquals(50, location.getWeakStats().getHoldBytes()); hold2.close(); Assertions.assertEquals(0, location.getWeakStats().getHoldBytes()); } + @Test + public void testConcurrentResizeAndReleaseLeavesHoldBytesBalanced() throws Exception + { + // A resize racing a release: the release runs outside the location lock, so it can land either side of the + // resize. Whichever way it interleaves, a hold subtracts exactly what it added, so the totals return to zero. + for (int i = 0; i < 500; i++) { + final StorageLocation location = new StorageLocation(tempDir, 1000L, null); + final TestResizableCacheEntry entry = new TestResizableCacheEntry("a" + i, 80); + final StorageLocation.ReservationHold reserver = + location.addWeakReservationHold(entry.getId(), () -> entry); + Assertions.assertNotNull(reserver); + final StorageLocation.ReservationHold queryHold = + location.addWeakReservationHoldIfExists(entry.getId()); + final StorageLocation.ReservationHold internalHold = + location.addInternalWeakReservationHoldIfExists(entry.getId()); + Assertions.assertNotNull(queryHold); + Assertions.assertNotNull(internalHold); + + final CountDownLatch start = new CountDownLatch(1); + final Future resizer = executorService.submit(() -> { + awaitUninterruptibly(start); + location.adjustReservation(entry.getId(), 30); + }); + final Future releaser = executorService.submit(() -> { + awaitUninterruptibly(start); + queryHold.close(); + internalHold.close(); + }); + start.countDown(); + resizer.get(); + releaser.get(); + reserver.close(); + + Assertions.assertEquals(0, location.getWeakStats().getHoldCount(), "iteration " + i); + Assertions.assertEquals(0, location.getWeakStats().getHoldBytes(), "iteration " + i); + Assertions.assertEquals(0, location.getWeakStats().getInternalHoldCount(), "iteration " + i); + Assertions.assertEquals(0, location.getWeakStats().getInternalHoldBytes(), "iteration " + i); + } + } + + private static void awaitUninterruptibly(CountDownLatch latch) + { + try { + Assertions.assertTrue(latch.await(30, TimeUnit.SECONDS)); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + @Test public void testInternalHoldsAreNotCountedAsHitsButAreCountedAsPinned() {