This is an automated email from the ASF dual-hosted git repository.
clintropolis pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 1ed7baae5cb fix: more chill abandoned partial loads and safer deferred
eviction of weak held entries (#20073)
1ed7baae5cb is described below
commit 1ed7baae5cb858810f1e197ea1ec46680322874e
Author: Clint Wylie <[email protected]>
AuthorDate: Tue Aug 25 20:02:17 2026 -0700
fix: more chill abandoned partial loads and safer deferred eviction of weak
held entries (#20073)
---
.../segment/loading/SegmentLocalCacheManager.java | 95 +++++++----
.../druid/segment/loading/StorageLocation.java | 7 +-
...SegmentLocalCacheManagerPartialAcquireTest.java | 185 +++++++++++++++++++++
.../druid/segment/loading/StorageLocationTest.java | 43 ++++-
4 files changed, 294 insertions(+), 36 deletions(-)
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 d8179e3ed45..f441515e1e6 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
@@ -383,18 +383,30 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
@Override
public void storeInfoFile(final DataSegment segment) throws IOException
+ {
+ writeInfoFile(segment, false);
+ }
+
+ /**
+ * Internal info file writer, where {@code overwrite} can be set to force
writing the file to disk unconditionally.
+ *
+ * @return whether this call wrote the file
+ */
+ private boolean writeInfoFile(final DataSegment segment, final boolean
overwrite) throws IOException
{
final File segmentInfoCacheFile = getSegmentInfoFile(segment);
- if (!segmentInfoCacheFile.exists()) {
- FileUtils.mkdirp(segmentInfoCacheFile.getParentFile());
- FileUtils.writeAtomically(
- segmentInfoCacheFile,
- out -> {
- jsonMapper.writeValue(out, segment);
- return null;
- }
- );
+ if (!overwrite && segmentInfoCacheFile.exists()) {
+ return false;
}
+ FileUtils.mkdirp(getEffectiveInfoDir());
+ FileUtils.writeAtomically(
+ segmentInfoCacheFile,
+ out -> {
+ jsonMapper.writeValue(out, segment);
+ return null;
+ }
+ );
+ return true;
}
@Override
@@ -444,12 +456,7 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
*/
private void rewriteInfoFile(DataSegment segment) throws IOException
{
- final File segmentInfoCacheFile = getSegmentInfoFile(segment);
- FileUtils.mkdirp(getEffectiveInfoDir());
- FileUtils.writeAtomically(segmentInfoCacheFile, out -> {
- jsonMapper.writeValue(out, segment);
- return null;
- });
+ writeInfoFile(segment, true);
}
@Override
@@ -501,13 +508,8 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
try {
if (hold != null) {
// write the segment info file if it doesn't exist. this can
happen if we are loading after a drop
- final File segmentInfoCacheFile =
getSegmentInfoFile(dataSegment);
- if (!segmentInfoCacheFile.exists()) {
- FileUtils.mkdirp(getEffectiveInfoDir());
- FileUtils.writeAtomically(segmentInfoCacheFile, out -> {
- jsonMapper.writeValue(out, dataSegment);
- return null;
- });
+ if (writeInfoFile(dataSegment, false)) {
+ // if we wrote it, set the hook to clean it up too
hold.getEntry().setOnUnmount(() ->
deleteSegmentInfoFile(dataSegment));
}
@@ -654,7 +656,12 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
try {
reserved.metadata.mount(reserved.location);
}
- catch (IOException e) {
+ catch (IOException | RuntimeException e) {
+ // only fail if caller was expecting this to finish
+ if (holdHolder.isClosed()) {
+ log.debug(e, "Mount of segment[%s] failed after its
acquire was abandoned", dataSegment.getId());
+ return AcquireSegmentResult.empty();
+ }
throw DruidException.defensive(
e,
"Failed to mount partial metadata for segment[%s]",
@@ -667,12 +674,15 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
taskMetadataRef =
reserved.metadata.acquireMetadataReference();
}
catch (DruidException raceLost) {
- throw DruidException.defensive(
- raceLost,
- "Partial metadata for segment[%s] was dropped before %s
task could complete",
- dataSegment.getId(),
- fullDownload ? "full-download" : "lazy mount"
- );
+ // The entry was evicted between mounting it and pinning it.
What keeps it resident across that
+ // window is the hold in holdHolder, so an acquire that
still holds one is entitled to its segment
+ // and getting here means that guarantee has been broken so
rethrow
+ if (!holdHolder.isClosed()) {
+ throw raceLost;
+ }
+ // hold is release so we have an abandoned acquire; nothing
is waiting on this result, report
+ // the segment as unavailable rather than failing a query
that is already on its way down
+ return AcquireSegmentResult.empty();
}
try {
final PartialSegmentFileMapperV10 mapper =
reserved.metadata.getFileMapper();
@@ -950,6 +960,17 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
{
final ReservedPartial existing =
findExistingPartialWithHold(dataSegment.getId());
if (existing != null) {
+ if (!existing.metadata().isMounted()) {
+ // Restore the info file if it is missing
+ try {
+ if (writeInfoFile(dataSegment, false)) {
+ existing.metadata().setOnUnmount(() ->
deleteSegmentInfoFile(dataSegment));
+ }
+ }
+ catch (IOException e) {
+ log.warn(e, "Failed to restore info file for cached segment[%s]",
dataSegment.getId());
+ }
+ }
return existing;
}
return reservePartial(dataSegment, rangeReader);
@@ -2458,14 +2479,26 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
{
@GuardedBy("this")
private final Closer holds = Closer.create();
- @GuardedBy("this")
- private boolean closed = false;
+ /**
+ * Volatile rather than guarded, so {@link #isClosed} can be answered
without queueing behind a {@link #close}
+ * that is working through a hold-release cascade (location write lock,
unmount, unmap, file deletion).
+ */
+ private volatile boolean closed = false;
private HoldHolder(Closeable initialHold)
{
holds.register(initialHold);
}
+ /**
+ * Whether the {@link AcquireSegmentAction} these holds belong to has been
closed, i.e. whoever wanted the segment
+ * has stopped waiting for it and every hold taken so far is being (or has
been) released.
+ */
+ private boolean isClosed()
+ {
+ return closed;
+ }
+
private void add(Closeable hold)
{
final boolean alreadyClosed;
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 0c3e4cc8d77..5236aa19689 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
@@ -556,11 +556,14 @@ public class StorageLocation
weakCacheEntries.computeIfPresent(
weakEntry.cacheEntry.getId(),
(cacheEntryIdentifier, weakCacheEntry) -> {
+ if (weakCacheEntry != weakEntry || weakCacheEntry.isHeld()) {
+ // Someone else's entry, or someone else is still using ours;
either way, theirs to clean up.
+ return weakCacheEntry;
+ }
// If we never successfully mounted, go ahead and remove so we
don't have a dead entry.
// Furthermore, if evictImmediatelyOnHoldRelease is set, evict
on release if all holds are gone.
final boolean isMounted = weakCacheEntry.cacheEntry.isMounted();
- if ((isNewEntry && !isMounted)
- || (areWeakEntriesEphemeral && !weakCacheEntry.isHeld())) {
+ if ((isNewEntry && !isMounted) || areWeakEntriesEphemeral) {
unlinkWeakEntry(weakCacheEntry);
if (isMounted) {
weakStats.getAndUpdate(s ->
s.evict(weakCacheEntry.cacheEntry.getSize()));
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 ccc47cfb6d3..4a5be4714b3 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
@@ -22,6 +22,7 @@ package org.apache.druid.segment.loading;
import com.fasterxml.jackson.databind.InjectableValues;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.jsontype.NamedType;
+import com.google.common.util.concurrent.ListenableFuture;
import org.apache.druid.data.input.InputRow;
import org.apache.druid.data.input.ListBasedInputRow;
import org.apache.druid.data.input.MapBasedInputRow;
@@ -37,6 +38,7 @@ import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.FileUtils;
import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.java.util.common.granularity.Granularities;
+import org.apache.druid.java.util.common.io.Closer;
import org.apache.druid.java.util.emitter.EmittingLogger;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.query.aggregation.CountAggregatorFactory;
@@ -79,6 +81,7 @@ import org.junit.jupiter.api.extension.RegisterExtension;
import java.io.File;
import java.io.IOException;
+import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -962,6 +965,188 @@ class SegmentLocalCacheManagerPartialAcquireTest
);
}
+ @Test
+ void testAbandonedAcquireLeavesAConcurrentAcquiresEntryAlone()
+ throws ExecutionException, InterruptedException, IOException
+ {
+ final StorageLocation location = manager.getLocations().get(0);
+ final SegmentCacheEntryIdentifier id = new
SegmentCacheEntryIdentifier(SEGMENT_ID);
+
+ // Two acquires holding the same freshly reserved entry, neither of which
resolves its future, so nothing mounts
+ // it. The acquire that created the entry gives up first - a canceled or
timed out query, say.
+ final PartialSegmentMetadataCacheEntry entry;
+ final Closer acquires = Closer.create();
+ try {
+ // The first acquire reserves the entry; the second joins it. Both are
registered with the closer so a failed
+ // assertion cannot leak a hold into the shared manager fixture;
AcquireSegmentAction.close() is idempotent, so
+ // closing the first one early below is safe.
+ final AcquireSegmentAction creator = acquires.register(
+ manager.acquireSegment(partialSegment, AcquireMode.PARTIAL)
+ );
+ acquires.register(manager.acquireSegment(partialSegment,
AcquireMode.PARTIAL));
+ entry =
Assertions.assertInstanceOf(PartialSegmentMetadataCacheEntry.class,
location.getCacheEntry(id));
+ Assertions.assertFalse(entry.isMounted());
+
+ creator.close();
+ Assertions.assertSame(
+ entry,
+ location.getCacheEntry(id),
+ "the entry the second acquire is holding must survive the first
giving up"
+ );
+ }
+ finally {
+ acquires.close();
+ }
+ // The second acquire letting go does not remove it either - removal is
the creating hold's job - so it is left
+ // registered and unmounted, reclaimable, and reusable by the next acquire.
+ Assertions.assertSame(entry, location.getCacheEntry(id));
+ Assertions.assertFalse(entry.isMounted());
+
+ // A later acquire mounts that same entry and serves the segment from it.
+ try (AcquireSegmentAction action = manager.acquireSegment(partialSegment,
AcquireMode.PARTIAL)) {
+ final AcquireSegmentResult result = action.getSegmentFuture().get();
+ try (Segment segment =
result.getReferenceProvider().acquireReference().orElseThrow()) {
+ Assertions.assertEquals(SEGMENT_ID, segment.getId());
+ final TimeBoundaryInspector inspector =
segment.as(TimeBoundaryInspector.class);
+ Assertions.assertNotNull(inspector);
+ Assertions.assertEquals(TIME, inspector.getMinTime());
+ Assertions.assertEquals(TIME.plusMinutes(3), inspector.getMaxTime());
+ }
+ Assertions.assertSame(entry, location.getCacheEntry(id));
+ Assertions.assertTrue(entry.isMounted());
+ }
+ }
+
+ @Test
+ void testAbandonedAcquireWhoseMountFailsResolvesToAnUnavailableSegmentToo()
throws Exception
+ {
+ // Deep storage that can produce a range reader now but not serve a read
later: openRangeReader() checks the V10
+ // file at acquire time, and the test deletes it before the load task gets
to run.
+ final File vanishingStorage =
temporaryFolder.newFolder("vanishing_storage");
+ final File v10File = new File(vanishingStorage, IndexIO.V10_FILE_NAME);
+ Files.copy(new File(DEEP_STORAGE_DIR, IndexIO.V10_FILE_NAME).toPath(),
v10File.toPath());
+ final DataSegment vanishingSegment = DataSegment.builder(SEGMENT_ID)
+
.shardSpec(NoneShardSpec.instance())
+ .loadSpec(Map.of("type",
"local", "path", vanishingStorage.getAbsolutePath()))
+ .size(0)
+ .build();
+
+ final File gatedCacheRoot =
temporaryFolder.newFolder("gated_cache_mount_failure");
+ final SegmentLoaderConfig gatedConfig = SegmentLoaderConfig.builder()
+ .locations(new StorageLocationConfig(gatedCacheRoot, 1024L * 1024L *
1024L, null))
+ .virtualStorage(true)
+ .virtualStoragePartialDownloadsEnabled(true)
+ .virtualStorageUseVirtualThreads(false)
+ .virtualStorageLoadThreads(1)
+ .build();
+ final List<StorageLocation> storageLocations =
gatedConfig.toStorageLocations();
+ final StorageLoadingThreadPool gatedPool =
StorageLoadingThreadPool.createFromConfig(gatedConfig);
+ final SegmentLocalCacheManager gatedManager = new SegmentLocalCacheManager(
+ storageLocations,
+ gatedConfig,
+ gatedPool,
+ new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
+ TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
+ jsonMapper
+ );
+
+ final CountDownLatch atGate = new CountDownLatch(1);
+ final CountDownLatch openGate = new CountDownLatch(1);
+ try {
+ // (intentionally unused) local so errorprone's CheckReturnValue is
satisfied
+ @SuppressWarnings("unused")
+ ListenableFuture<?> unused = gatedPool.getExecutorService().submit(() ->
{
+ atGate.countDown();
+ return openGate.await(30, TimeUnit.SECONDS);
+ });
+ Assertions.assertTrue(atGate.await(30, TimeUnit.SECONDS), "loading
thread must reach the gate");
+
+ final ListenableFuture<AcquireSegmentResult> future;
+ try (AcquireSegmentAction action =
gatedManager.acquireSegment(vanishingSegment, AcquireMode.PARTIAL)) {
+ future = action.getSegmentFuture();
+ }
+ // The mount will now fail on its header read rather than rolling back
cleanly, so the loss surfaces from
+ // mount() instead of from the pin - which is no reason to fail a query
that has already given up.
+ Assertions.assertTrue(v10File.delete(), "test needs the deep-storage
file gone before the mount runs");
+ openGate.countDown();
+
+ final AcquireSegmentResult result = future.get(30, TimeUnit.SECONDS);
+ Assertions.assertTrue(
+ result.getReferenceProvider().acquireReference().isEmpty(),
+ "an abandoned acquire whose mount failed must resolve to an
unavailable segment"
+ );
+ }
+ finally {
+ openGate.countDown();
+ gatedManager.drop(vanishingSegment);
+ gatedManager.shutdown();
+ gatedPool.stop();
+ }
+ }
+
+ @Test
+ void testAbandonedAcquireResolvesToAnUnavailableSegmentRatherThanFailing()
throws Exception
+ {
+ // One fixed loading thread, so the test can hold the load task at a gate
while it abandons the acquire.
+ final File gatedCacheRoot = temporaryFolder.newFolder("gated_cache");
+ final SegmentLoaderConfig gatedConfig = SegmentLoaderConfig.builder()
+ .locations(new StorageLocationConfig(gatedCacheRoot, 1024L * 1024L *
1024L, null))
+ .virtualStorage(true)
+ .virtualStoragePartialDownloadsEnabled(true)
+ .virtualStorageUseVirtualThreads(false)
+ .virtualStorageLoadThreads(1)
+ .build();
+ final List<StorageLocation> storageLocations =
gatedConfig.toStorageLocations();
+ final StorageLoadingThreadPool gatedPool =
StorageLoadingThreadPool.createFromConfig(gatedConfig);
+ final SegmentLocalCacheManager gatedManager = new SegmentLocalCacheManager(
+ storageLocations,
+ gatedConfig,
+ gatedPool,
+ new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
+ TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
+ jsonMapper
+ );
+
+ final CountDownLatch atGate = new CountDownLatch(1);
+ final CountDownLatch openGate = new CountDownLatch(1);
+ try {
+ // (intentionally unused) local so errorprone's CheckReturnValue is
satisfied
+ @SuppressWarnings("unused")
+ ListenableFuture<?> unused = gatedPool.getExecutorService().submit(() ->
{
+ atGate.countDown();
+ return openGate.await(30, TimeUnit.SECONDS);
+ });
+ Assertions.assertTrue(atGate.await(30, TimeUnit.SECONDS), "loading
thread must reach the gate");
+
+ // Start the acquire (its load task queues behind the gate) and then
abandon it (like a canceled or timed
+ // out query does); closing the action releases the hold that was
keeping the reserved entry resident.
+ final ListenableFuture<AcquireSegmentResult> future;
+ try (AcquireSegmentAction action =
gatedManager.acquireSegment(partialSegment, AcquireMode.PARTIAL)) {
+ future = action.getSegmentFuture();
+ }
+ openGate.countDown();
+
+ // The task now mounts an entry the location no longer knows about, so
it can never pin it. Nothing is waiting
+ // on the result, and a segment that is not there is not a reason to
fail a query.
+ final AcquireSegmentResult result = future.get(30, TimeUnit.SECONDS);
+ Assertions.assertTrue(
+ result.getReferenceProvider().acquireReference().isEmpty(),
+ "an abandoned acquire must resolve to an unavailable segment"
+ );
+ Assertions.assertNull(
+ gatedManager.getLocations().get(0).getCacheEntry(new
SegmentCacheEntryIdentifier(SEGMENT_ID)),
+ "the abandoned entry must not be left behind in the cache"
+ );
+ }
+ finally {
+ openGate.countDown();
+ gatedManager.drop(partialSegment);
+ gatedManager.shutdown();
+ // shutdown() only stops the manager's own executor; the loading pool is
stopped through its lifecycle hook
+ gatedPool.stop();
+ }
+ }
+
/**
* 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
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 52e4c3ee489..b4b01f6dd21 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,46 @@ class StorageLocationTest
Assertions.assertEquals(0, location.getWeakStats().getHoldBytes());
}
+ @Test
+ public void
testReleasingReserversHoldDoesNotEvictAnEntryAnotherHolderIsStillUsing()
+ {
+ final StorageLocation location = new StorageLocation(tempDir, 100L, null);
+ final UnmountTrackingCacheEntry entry = new UnmountTrackingCacheEntry("a",
10);
+
+ // One acquirer reserves the entry, a second takes its own hold on the
same entry, and then the first gives up
+ // before anything has managed to mount it
+ final StorageLocation.ReservationHold<?> reserver =
location.addWeakReservationHold(entry.getId(), () -> entry);
+ final StorageLocation.ReservationHold<?> other =
location.addWeakReservationHold(entry.getId(), () -> entry);
+ Assertions.assertNotNull(reserver);
+ Assertions.assertNotNull(other);
+ Assertions.assertFalse(entry.isMounted());
+
+ reserver.close();
+
+ // The second holder is still using the entry (e.g. mid-mount) and
evicting it out from under a hold
+ // is what the hold exists to prevent: a mount that completes to find its
entry unregistered rolls itself back,
+ // leaving the holder with a reference to an entry it was promised would
stay alive.
+ Assertions.assertSame(entry, location.getCacheEntry(entry.getId()));
+ Assertions.assertFalse(entry.unmountCalled);
+
+ // A later acquirer for this id reuses the registered entry rather than
registering a second one alongside it,
+ // which is what keeps the on-disk state two entries for the same id would
share owned by exactly one of them.
+ final UnmountTrackingCacheEntry replacement = new
UnmountTrackingCacheEntry("a", 10);
+ final StorageLocation.ReservationHold<?> reacquire =
+ location.addWeakReservationHold(entry.getId(), () -> replacement);
+ Assertions.assertNotNull(reacquire);
+ Assertions.assertSame(entry, reacquire.getEntry());
+ Assertions.assertFalse(replacement.unmountCalled);
+ reacquire.close();
+ Assertions.assertSame(entry, location.getCacheEntry(entry.getId()));
+
+ // Releasing that hold does not remove it either - removal is the creating
hold's job - so it is left registered
+ // and unmounted, for reclaim to take when the space is wanted.
+ other.close();
+ Assertions.assertSame(entry, location.getCacheEntry(entry.getId()));
+ Assertions.assertFalse(entry.unmountCalled);
+ }
+
@Test
public void
testEphemeralWeakEntryUnmountCascadeDoesNotThrowConcurrentModification()
{
@@ -822,9 +862,6 @@ class StorageLocationTest
}
}
- /**
- * A {@link CacheEntry} that tracks mount/unmount so tests can assert that
lifecycle hooks fired.
- */
private static final class UnmountTrackingCacheEntry implements CacheEntry
{
private final StringCacheIdentifier id;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
