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 9e17b1c48a1 fix: cache entries use self hold to avoid eviction while
mounting, and other fixes (#20182)
9e17b1c48a1 is described below
commit 9e17b1c48a11ffee625eea19e2fe88c483352713
Author: Clint Wylie <[email protected]>
AuthorDate: Fri Aug 28 15:03:01 2026 -0700
fix: cache entries use self hold to avoid eviction while mounting, and
other fixes (#20182)
---
.../segment/file/PartialSegmentFileMapperV10.java | 194 +++++++++++++-------
.../file/PartialSegmentFileMapperV10Test.java | 56 ++++++
.../loading/PartialSegmentBundleCacheEntry.java | 98 +++++++---
.../loading/PartialSegmentMetadataCacheEntry.java | 15 +-
.../segment/loading/SegmentLocalCacheManager.java | 25 ++-
.../PartialSegmentBundleCacheEntryTest.java | 149 +++++++++++++++
.../PartialSegmentMetadataCacheEntryTest.java | 99 ++++++++++
.../SegmentLocalCacheManagerConcurrencyTest.java | 199 ++++++++++++---------
...SegmentLocalCacheManagerPartialAcquireTest.java | 136 ++++++++++++++
9 files changed, 791 insertions(+), 180 deletions(-)
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 83e373080e6..5f2af4699e8 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
@@ -312,6 +312,11 @@ public class PartialSegmentFileMapperV10 implements
SegmentFileMapper
private final ReentrantLock[] containerLocks;
// per-container eviction generation; see getBundleGeneration
private final AtomicLongArray containerGenerations;
+ /**
+ * Number of fetches currently writing into each container, and whether an
eviction is waiting on them.
+ */
+ private final int[] containerFetchesInFlight;
+ private final boolean[] containerEvictionPending;
// bundle name -> indices (into metadata.getContainers()) of this single
mapper's containers in that bundle.
// Computed once at construction from the immutable container metadata.
Single-mapper scope only: stitching bundles
@@ -372,6 +377,8 @@ public class PartialSegmentFileMapperV10 implements
SegmentFileMapper
this.containerFiles = new File[numContainers];
this.containerLocks = new ReentrantLock[numContainers];
this.containerGenerations = new AtomicLongArray(numContainers);
+ this.containerFetchesInFlight = new int[numContainers];
+ this.containerEvictionPending = new boolean[numContainers];
final Map<String, List<Integer>> bundleIndices = new HashMap<>();
for (int i = 0; i < numContainers; i++) {
this.containerLocks[i] = new ReentrantLock();
@@ -531,12 +538,24 @@ public class PartialSegmentFileMapperV10 implements
SegmentFileMapper
* read (its files tile back-to-back), and a partially-downloaded container
skips its already-resident spans.
*/
public void ensureAllDownloaded() throws IOException
+ {
+ fetchAllContainers();
+ if (!isFullyDownloaded()) {
+ throw DruidException.defensive(
+ "Failed to download every file of [%s]; residency was cleared
mid-download, which means a container was "
+ + "evicted while a fetch was writing into it",
+ targetFilename
+ );
+ }
+ }
+
+ private void fetchAllContainers() throws IOException
{
for (int containerIndex = 0; containerIndex < containers.length;
containerIndex++) {
fetchFiles(containerFileNames.get(containerIndex));
}
for (PartialSegmentFileMapperV10 external : externalMappers.values()) {
- external.ensureAllDownloaded();
+ external.fetchAllContainers();
}
}
@@ -823,40 +842,46 @@ public class PartialSegmentFileMapperV10 implements
SegmentFileMapper
}
checkClosed();
- int from = 0;
- int to = runFiles.size();
- while (from < to && downloadedFiles.contains(runFiles.get(from))) {
- from++;
- }
- while (to > from && downloadedFiles.contains(runFiles.get(to - 1))) {
- to--;
- }
- if (from == to) {
- // the whole run became resident while we were waiting on the locks
- return;
- }
- final List<String> remaining = runFiles.subList(from, to);
- final SegmentInternalFileMetadata first =
metadata.getFiles().get(remaining.get(0));
- final long startOffset = first.getStartOffset();
- // scan for the span end rather than assuming the last file ends last:
shrinking the read below any covered
- // file's end would mark that file downloaded without its bytes on disk
(zero-length files share start offsets)
- long endOffset = startOffset;
- for (String name : remaining) {
- final SegmentInternalFileMetadata fileMeta =
metadata.getFiles().get(name);
- endOffset = Math.max(endOffset, fileMeta.getStartOffset() +
fileMeta.getSize());
+ beginContainerFetch(containerIndex);
+ try {
+ int from = 0;
+ int to = runFiles.size();
+ while (from < to && downloadedFiles.contains(runFiles.get(from))) {
+ from++;
+ }
+ while (to > from && downloadedFiles.contains(runFiles.get(to - 1))) {
+ to--;
+ }
+ if (from == to) {
+ // the whole run became resident while we were waiting on the locks
+ return;
+ }
+ final List<String> remaining = runFiles.subList(from, to);
+ final SegmentInternalFileMetadata first =
metadata.getFiles().get(remaining.get(0));
+ final long startOffset = first.getStartOffset();
+ // scan for the span end rather than assuming the last file ends last:
shrinking the read below any covered
+ // file's end would mark that file downloaded without its bytes on
disk (zero-length files share start offsets)
+ long endOffset = startOffset;
+ for (String name : remaining) {
+ final SegmentInternalFileMetadata fileMeta =
metadata.getFiles().get(name);
+ endOffset = Math.max(endOffset, fileMeta.getStartOffset() +
fileMeta.getSize());
+ }
+ final long length = endOffset - startOffset;
+
+ ensureContainerInitialized(containerIndex);
+ streamRangeIntoContainer(
+ containerIndex,
+ computeAbsoluteOffset(first),
+ startOffset,
+ length,
+ StringUtils.format("files[%d] in container[%d]", remaining.size(),
containerIndex)
+ );
+ for (String name : remaining) {
+ markDownloaded(name, metadata.getFiles().get(name).getSize());
+ }
}
- final long length = endOffset - startOffset;
-
- ensureContainerInitialized(containerIndex);
- streamRangeIntoContainer(
- containerIndex,
- computeAbsoluteOffset(first),
- startOffset,
- length,
- StringUtils.format("files[%d] in container[%d]", remaining.size(),
containerIndex)
- );
- for (String name : remaining) {
- markDownloaded(name, metadata.getFiles().get(name).getSize());
+ finally {
+ endContainerFetch(containerIndex);
}
}
finally {
@@ -964,22 +989,63 @@ public class PartialSegmentFileMapperV10 implements
SegmentFileMapper
ensureContainerInitialized(containerIndex);
}
+ /**
+ * Register that a fetch is about to write into this container, so a
concurrent {@link #evictContainer} defers
+ * rather than deleting the file out from under it. Paired with {@link
#endContainerFetch} in a finally.
+ */
+ private void beginContainerFetch(int containerIndex)
+ {
+ containerLocks[containerIndex].lock();
+ try {
+ containerFetchesInFlight[containerIndex]++;
+ }
+ finally {
+ containerLocks[containerIndex].unlock();
+ }
+ }
+
+ /**
+ * Drop this fetch's claim on the container and, if it was the last one and
an eviction was deferred while it ran,
+ * carry that eviction out now.
+ */
+ private void endContainerFetch(int containerIndex)
+ {
+ boolean evictNow = false;
+ containerLocks[containerIndex].lock();
+ try {
+ containerFetchesInFlight[containerIndex]--;
+ if (containerFetchesInFlight[containerIndex] == 0 &&
containerEvictionPending[containerIndex]) {
+ containerEvictionPending[containerIndex] = false;
+ evictNow = true;
+ }
+ }
+ finally {
+ containerLocks[containerIndex].unlock();
+ }
+ if (evictNow) {
+ // Runs from fetchRun's finally, so it must not throw over whatever
brought us here.
+ try {
+ evictContainer(containerIndex);
+ }
+ catch (Throwable t) {
+ LOG.warn(t, "Failed to run deferred eviction of container[%d] for
[%s]", containerIndex, targetFilename);
+ }
+ }
+ }
+
/**
* Reverse of {@link #initializeContainer(int)}: unmap the in-memory view of
the container, delete the local
* container file, and clear the bitmap bits + {@link #downloadedFiles}
entries for every internal file that lived
* in this container.
* <p>
- * Used by per-bundle cache entries on unmount/eviction to release the disk
and memory footprint of one bundle
- * without affecting other bundles sharing the same {@link
PartialSegmentFileMapperV10}. After eviction, the
- * container's files are non-resident again: {@link #mapFile} throws for
them until a subsequent fetch re-downloads
- * them (re-initializing the container and repopulating the bitmap
incrementally).
- * <p>
- * <b>Concurrency contract.</b> The caller is responsible for ensuring no
concurrent {@link #mapFile} (or
- * {@link #fetchFiles}/{@link #fetchRun}) call is in flight for any file in
this container. This is enforced one layer up
- * by the cache-entry refcount: {@code PartialSegmentBundleCacheEntry} only
invokes {@code evictContainer} from its
- * {@code doActualUnmount} callback, which fires only after every reference
acquired via {@code acquireReference()}
- * has been closed. Bypassing that gate is dangerous, {@link
ByteBufferUtils#unmap} frees the off-heap mapping, so a
- * {@link ByteBuffer#slice} from a concurrent reader is a JVM SIGSEGV, not a
recoverable error.
+ * <b>Concurrency contract.</b> No concurrent {@link #mapFile} may be in
flight for any file in this container;
+ * that is enforced one layer up by the cache-entry refcount, since {@code
PartialSegmentBundleCacheEntry} evicts
+ * from its {@code doActualUnmount} callback, which fires only after every
reference acquired via
+ * {@code acquireReference()} has been closed. Bypassing that gate is
dangerous: {@link ByteBufferUtils#unmap} frees
+ * the off-heap mapping, so a {@link ByteBuffer#slice} from a concurrent
reader is a JVM SIGSEGV, not a recoverable
+ * error. An in-flight {@link #fetchFiles}/{@link #fetchRun} is handled here
instead of by the caller: the eviction
+ * is deferred to whichever fetch finishes last, because callers can hold
the storage location's write lock and must
+ * not block on a deep-storage read.
* <p>
* No-op if the container has not been initialized.
*/
@@ -988,6 +1054,12 @@ public class PartialSegmentFileMapperV10 implements
SegmentFileMapper
checkClosed();
containerLocks[containerIndex].lock();
try {
+ if (containerFetchesInFlight[containerIndex] > 0) {
+ // A fetch is writing into this container. Deleting the file now would
leave it writing to a null File, so
+ // hand the eviction to whichever fetch finishes last rather than
blocking here.
+ containerEvictionPending[containerIndex] = true;
+ return;
+ }
final MappedByteBuffer existing = containers[containerIndex];
if (existing != null) {
ByteBufferUtils.unmap(existing);
@@ -1012,27 +1084,25 @@ public class PartialSegmentFileMapperV10 implements
SegmentFileMapper
);
}
containerFiles[containerIndex] = null;
+
+ // clear bitmap bits + downloadedFiles entries for files that lived in
this container.
+ for (Map.Entry<String, SegmentInternalFileMetadata> entry :
metadata.getFiles().entrySet()) {
+ if (entry.getValue().getContainer() != containerIndex) {
+ continue;
+ }
+ final String fileName = entry.getKey();
+ if (downloadedFiles.remove(fileName)) {
+ downloadedBytes.addAndGet(-entry.getValue().getSize());
+ }
+ clearBitmapBit(fileName);
+ }
+
+ // last: readers that observe the bumped generation must also observe
the cleared residency above
+ containerGenerations.incrementAndGet(containerIndex);
}
finally {
containerLocks[containerIndex].unlock();
}
-
- // clear bitmap bits + downloadedFiles entries for files that lived in
this container. Iterates
- // metadata.getFiles() without external synchronization:
SegmentFileMetadata is constructed once at mapper
- // creation and its file map is effectively immutable for the mapper's
lifetime, so concurrent iteration is safe.
- for (Map.Entry<String, SegmentInternalFileMetadata> entry :
metadata.getFiles().entrySet()) {
- if (entry.getValue().getContainer() != containerIndex) {
- continue;
- }
- final String fileName = entry.getKey();
- if (downloadedFiles.remove(fileName)) {
- downloadedBytes.addAndGet(-entry.getValue().getSize());
- }
- clearBitmapBit(fileName);
- }
-
- // last: readers that observe the bumped generation must also observe the
cleared residency above
- containerGenerations.incrementAndGet(containerIndex);
}
/**
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 6060d887aac..b3d4ba182b7 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
@@ -1317,6 +1317,62 @@ class PartialSegmentFileMapperV10Test
return new File(baseDir, IndexIO.V10_FILE_NAME);
}
+ @Test
+ void testEvictContainerDefersWhileAFetchIsWritingIntoIt() throws Exception
+ {
+ final File segmentFile = buildTestSegment(20, CompressionStrategy.NONE);
+ final File cacheDir = newCacheDir("evict-defer");
+ final CountDownLatch fetching = new CountDownLatch(1);
+ final CountDownLatch release = new CountDownLatch(1);
+ final AtomicBoolean armed = new AtomicBoolean(false);
+ final AtomicBoolean gated = new AtomicBoolean(false);
+ final SegmentRangeReader delegate = new
CountingRangeReader(segmentFile.getParentFile());
+ // park the first read once armed, so a fetch is provably mid-flight when
the eviction arrives. Armed only after
+ // the mapper exists, since create() does its own header read on this
thread.
+ final SegmentRangeReader gatedReader = (filename, offset, length) -> {
+ if (armed.get() && gated.compareAndSet(false, true)) {
+ fetching.countDown();
+ try {
+ Assertions.assertTrue(release.await(30, TimeUnit.SECONDS));
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException(e);
+ }
+ }
+ return delegate.readRange(filename, offset, length);
+ };
+
+ final ExecutorService exec = Execs.multiThreaded(1, "evict-defer-test-%d");
+ try (PartialSegmentFileMapperV10 mapper = createMapper(gatedReader,
cacheDir)) {
+ final String fileName =
mapper.getSegmentFileMetadata().getFiles().keySet().iterator().next();
+ final int containerIndex =
mapper.getSegmentFileMetadata().getFiles().get(fileName).getContainer();
+ armed.set(true);
+
+ final Future<?> fetch = exec.submit(() -> {
+ mapper.fetchFiles(List.of(fileName));
+ return null;
+ });
+ Assertions.assertTrue(fetching.await(30, TimeUnit.SECONDS), "fetch must
reach the gated read");
+
+ // Evicting now would delete the file the fetch is streaming into. It
must defer instead of tearing it down.
+ mapper.evictContainer(containerIndex);
+
+ release.countDown();
+ fetch.get(30, TimeUnit.SECONDS);
+
+ // The fetch completed, and the deferred eviction then ran: the file is
no longer resident.
+ Assertions.assertFalse(
+ mapper.getDownloadedFiles().contains(fileName),
+ "the deferred eviction should have run once the fetch finished"
+ );
+ }
+ finally {
+ release.countDown();
+ exec.shutdownNow();
+ }
+ }
+
private static PartialSegmentFileMapperV10 createMapper(
SegmentRangeReader rangeReader,
File localCacheDir
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 00d7f90170d..ccf2acab3f9 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
@@ -281,7 +281,12 @@ public class PartialSegmentBundleCacheEntry implements
CacheEntry
if (!mountFuture.compareAndSet(null, ours)) {
continue;
}
- try {
+ // Hold this entry against reclaim while the mount establishes state;
reclaim passes over held entries only,
+ // and an entry evicted mid-mount finishes into a location that no
longer knows about it. Internal, since a
+ // mount is not somebody waiting on the bundle, and null when the entry
is already gone - nothing to protect.
+ final StorageLocation.ReservationHold<PartialSegmentBundleCacheEntry>
selfHold =
+ mountLocation.addInternalWeakReservationHoldIfExists(id);
+ try (selfHold) {
doMount(mountLocation);
ours.set(null);
}
@@ -289,6 +294,10 @@ public class PartialSegmentBundleCacheEntry implements
CacheEntry
// clear the gate so the next caller gets a fresh attempt
mountFuture.set(null);
ours.setException(t);
+ // If it is already gone, no unmount() will follow to sweep them up,
so reap them here.
+ if (!isStillRegistered(mountLocation)) {
+ reapContainersOfUncommittedMount();
+ }
switch (t) {
case IOException ioException -> throw ioException;
case RuntimeException runtimeException -> throw runtimeException;
@@ -311,7 +320,7 @@ public class PartialSegmentBundleCacheEntry implements
CacheEntry
*/
private void verifyStillReservedOrRollback(StorageLocation mountLocation)
{
- if (!mountLocation.isReserved(id) && !mountLocation.isWeakReserved(id)) {
+ if (!isStillRegistered(mountLocation)) {
LOG.debug(
"Aborting mount of bundle[%s] in location[%s]; entry was evicted
while mounting",
id,
@@ -321,6 +330,65 @@ public class PartialSegmentBundleCacheEntry implements
CacheEntry
}
}
+ /**
+ * Whether the cache still knows about this entry.
+ */
+ private boolean isStillRegistered(StorageLocation mountLocation)
+ {
+ return mountLocation.isReserved(id) || mountLocation.isWeakReserved(id);
+ }
+
+ /**
+ * Evict containers left behind by a mount that never committed. A failed
mount deliberately keeps the containers it
+ * had already initialized so a retry can reuse them (see {@link #doMount});
this entry being torn down is that retry
+ * never coming, and nothing else will clean them up since {@link
#doActualUnmount} runs only for a mount that
+ * committed.
+ */
+ private void reapContainersOfUncommittedMount()
+ {
+ final SettableFuture<Void> ours = SettableFuture.create();
+ if (!mountFuture.compareAndSet(null, ours)) {
+ // something else is mounting, bail out
+ return;
+ }
+ try {
+ // Null once the metadata entry has unmounted, which closes the file
mapper; evictContainer would throw on it.
+ final PartialSegmentFileMapperV10 fileMapper =
metadataEntry.getFileMapper();
+ if (fileMapper != null) {
+ evictOwnedContainers(fileMapper);
+ }
+ }
+ finally {
+ // Clear the gate before completing it, so a mount that joined while we
held it finds a fresh gate to claim
+ // rather than this one. It fails instead of believing a mount happened,
which is accurate: it asked to mount an
+ // entry that was being torn down.
+ mountFuture.set(null);
+ ours.setException(DruidException.defensive("Bundle[%s] was torn down
while mounting", id));
+ }
+ }
+
+ /**
+ * Evict every container this bundle owns, routing each ref to the mapper
that owns it. Best-effort: a container that
+ * fails to evict is logged rather than aborting the rest, since every
caller is already tearing down.
+ */
+ private void evictOwnedContainers(PartialSegmentFileMapperV10 fileMapper)
+ {
+ for (BundleContainerRef ref : containerRefs) {
+ try {
+
fileMapper.mapperForContainer(ref.externalFilename()).evictContainer(ref.containerIndex());
+ }
+ catch (Throwable t) {
+ LOG.warn(
+ t,
+ "Failed to evict container[%s/%d] for bundle[%s]",
+ ref.externalFilename(),
+ ref.containerIndex(),
+ id
+ );
+ }
+ }
+ }
+
private void doMount(StorageLocation mountLocation) throws IOException
{
// Pre-check inside entryLock; after this we release entryLock so the
hold-acquisition + container-init work below
@@ -432,19 +500,6 @@ public class PartialSegmentBundleCacheEntry implements
CacheEntry
}
finally {
if (!committed) {
- // Evict any containers that were successfully initialized before the
failure. Mirrors the eager
- // SegmentCacheEntry behavior: retry from a clean slate is simpler
than reasoning about partial on-disk state.
- // evictContainer is a no-op for containers that were never
initialized, so we can iterate the full set
- // without tracking how far the initialization loop got.
- for (BundleContainerRef ref : containerRefs) {
- try {
-
fileMapper.mapperForContainer(ref.externalFilename()).evictContainer(ref.containerIndex());
- }
- catch (Throwable t) {
- LOG.warn(t, "Failed to evict container[%s/%d] for bundle[%s]
during mount rollback",
- ref.externalFilename(), ref.containerIndex(), id);
- }
- }
if (registered) {
try {
metadataEntry.unregisterBundle(this);
@@ -485,7 +540,11 @@ public class PartialSegmentBundleCacheEntry implements
CacheEntry
final ReferenceCountingCloseableObject<Closeable> current =
references.get();
if (current != null && !current.isClosed()) {
current.close();
+ return;
}
+ // Nothing mounted to tear down, so doActualUnmount will not run for this
entry; sweep up after any mount that
+ // failed partway and left containers behind.
+ reapContainersOfUncommittedMount();
}
/**
@@ -534,14 +593,7 @@ public class PartialSegmentBundleCacheEntry implements
CacheEntry
final PartialSegmentFileMapperV10 fileMapper =
metadataEntry.getFileMapper();
// file mapper may be null if metadata was already unmounted
(out-of-order shutdown); evictContainer would NPE
if (fileMapper != null) {
- for (BundleContainerRef ref : containerRefs) {
- try {
-
fileMapper.mapperForContainer(ref.externalFilename()).evictContainer(ref.containerIndex());
- }
- catch (Throwable t) {
- LOG.warn(t, "Failed to evict container[%s/%d] for bundle[%s]",
ref.externalFilename(), ref.containerIndex(), id);
- }
- }
+ evictOwnedContainers(fileMapper);
}
refsToRelease = new ArrayList<>(dependencyReferences);
dependencyReferences.clear();
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 1e11d950fe1..96dde6cba21 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
@@ -757,6 +757,13 @@ public class PartialSegmentMetadataCacheEntry implements
SegmentCacheEntry, Resi
// blocked on it. adjustReservation also runs outside entryLock:
StorageLocation.release goes
// writeLock -> entryLock (via release -> unmount), so entryLock ->
writeLock here would be a deadlock-prone
// lock-order inversion.
+ // Hold this entry against reclaim for as long as the mount is
establishing state. reclaim passes over held
+ // entries only, and an entry evicted mid-mount finishes into a location
that no longer knows about it: the
+ // post-mount check then rolls the whole thing back, or worse commits it
if a fresh entry has taken the id in the
+ // meantime, leaving a mapper nothing will ever unmount. Released before
the cleanup in the catch below, which only
+ // acts on an entry no one holds.
+ final StorageLocation.ReservationHold<SegmentCacheEntry> selfHold =
+ mountLocation.addInternalWeakReservationHoldIfExists(id);
try {
entryLock.lock();
try {
@@ -870,8 +877,9 @@ public class PartialSegmentMetadataCacheEntry implements
SegmentCacheEntry, Resi
catch (Throwable t) {
// 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.
+ // bundle entries), so the mount's own hold has to go first. Runs
outside entryLock (the inner blocks released
+ // it) so the writeLock -> entryLock order inside removeUnheldWeakEntry
is respected.
+ CloseableUtils.closeAndSuppressExceptions(selfHold, e -> LOG.warn(e,
"Failed to release mount hold[%s]", id));
try {
mountLocation.removeUnheldWeakEntry(id);
}
@@ -880,6 +888,9 @@ public class PartialSegmentMetadataCacheEntry implements
SegmentCacheEntry, Resi
}
throw t;
}
+ finally {
+ CloseableUtils.closeAndSuppressExceptions(selfHold, e -> LOG.warn(e,
"Failed to release mount hold[%s]", id));
+ }
}
/**
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 f441515e1e6..91fd583d4db 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
@@ -690,11 +690,20 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
if (fullDownload) {
// Delta of internal-file bytes downloaded by this task
final long downloadedBefore = mapper.getDownloadedBytes();
- // Mount every bundle so the containers it owns are
reserved on the location
- for (String bundleName :
PartialSegmentBundleCacheEntry.bundleNames(mapper)) {
-
holdHolder.add(reserved.metadata.getBundleAcquirer().acquire(bundleName));
+ // Mount every bundle so the containers it owns are
reserved on the location, and keep those
+ // references here for the duration of the download
instead of handing them straight to
+ // holdHolder so that the caller abandoning a load doesn't
release the holds until load is
+ // finished.
+ final List<Closeable> bundleRefs = new ArrayList<>();
+ try {
+ for (String bundleName :
PartialSegmentBundleCacheEntry.bundleNames(mapper)) {
+
bundleRefs.add(reserved.metadata.getBundleAcquirer().acquire(bundleName));
+ }
+ mapper.ensureAllDownloaded();
+ }
+ finally {
+ bundleRefs.forEach(holdHolder::add);
}
- mapper.ensureAllDownloaded();
loadSizeBytes = mapper.getDownloadedBytes() -
downloadedBefore;
} else {
// Lazy mount: the header bytes when this task caused the
mount; 0 when the entry was already
@@ -2162,6 +2171,11 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
return;
}
+ // Hold this entry against reclaim for the duration of the mount,
including the post-mount reservation check
+ // below. reclaim passes over held entries only, and evicting one
mid-mount unmounts a storage directory this
+ // mount is filling, or when the directory was already resident, one it
is about to serve.
+ final StorageLocation.ReservationHold<SegmentCacheEntry> selfHold =
+ mountLocation.addInternalWeakReservationHoldIfExists(this.id);
try {
entryLock.lock();
try {
@@ -2262,6 +2276,9 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
unmount();
throw t;
}
+ finally {
+ CloseableUtils.closeAndSuppressExceptions(selfHold, e -> log.warn(e,
"Failed to release mount hold[%s]", id));
+ }
}
@Override
diff --git
a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java
b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java
index 27974a919dc..3d08c548cc8 100644
---
a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java
+++
b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntryTest.java
@@ -30,6 +30,7 @@ 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;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.concurrent.Execs;
@@ -75,6 +76,8 @@ class PartialSegmentBundleCacheEntryTest
private static final ObjectMapper JSON_MAPPER = TestHelper.makeJsonMapper();
private static final SegmentId SEGMENT_ID = SegmentId.of("test",
Intervals.of("2025/2026"), "v1", 0);
private static final String AGG_BUNDLE = "dim1_metric1_sum";
+ private static final String TWO_CONTAINER_BUNDLE = "two_container";
+ private static final String TWO_CONTAINER_EXTERNAL =
"two_container_ext.segment";
private static final long ESTIMATE = 16 * 1024 * 1024L;
private static final DateTime TIME = DateTimes.of("2025-01-01");
@@ -745,6 +748,152 @@ class PartialSegmentBundleCacheEntryTest
);
}
+ @Test
+ void testFailedMountKeepsInitializedContainersForRetry() throws IOException
+ {
+ final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE *
8, null);
+ final PartialSegmentMetadataCacheEntry metadata =
newTwoContainerBundleMetadata();
+ Assertions.assertTrue(location.reserve(metadata));
+ metadata.mount(location);
+
+ final PartialSegmentBundleCacheEntry bundle =
+ PartialSegmentBundleCacheEntry.forBundle(metadata,
TWO_CONTAINER_BUNDLE, List.of());
+ Assertions.assertNotNull(location.addWeakReservationHold(bundle.getId(),
() -> bundle));
+
+ final List<PartialSegmentBundleCacheEntry.BundleContainerRef> refs =
bundle.getContainerRefs();
+ Assertions.assertEquals(2, refs.size());
+ final File firstContainerFile = containerFileFor(refs.getFirst());
+ final File blocked = containerFileFor(refs.getLast());
+ // Fail the container-initialization loop on its last container by parking
a directory where that container's file
+ // needs to go, so RandomAccessFile(..., "rw") can't open it. The one
before it initializes normally.
+ FileUtils.mkdirp(blocked);
+
+ Assertions.assertThrows(IOException.class, () -> bundle.mount(location));
+ Assertions.assertFalse(bundle.isMounted());
+
+ // The entry is still in the cache, so the container that did initialize
stays put for a retry to reuse rather
+ // than being torn down and re-created.
+ Assertions.assertTrue(
+ firstContainerFile.exists(),
+ "a failed mount should leave already-initialized containers in place"
+ );
+
+ // ...and the retry succeeds against that reused state once the failure
clears.
+ Assertions.assertTrue(blocked.delete());
+ bundle.mount(location);
+ Assertions.assertTrue(bundle.isMounted());
+ for (PartialSegmentBundleCacheEntry.BundleContainerRef ref : refs) {
+ Assertions.assertTrue(containerFileFor(ref).exists(), "container " + ref
+ " should be allocated after retry");
+ }
+ }
+
+ @Test
+ void testFailedMountEvictsContainersWhenEntryIsNoLongerReserved() throws
IOException
+ {
+ final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE *
8, null);
+ // ephemeral mode: releasing the last hold drops the weak entry
immediately, so the mount below runs for an entry
+ // the cache no longer knows about.
+ location.setAreWeakEntriesEphemeral(true);
+ final PartialSegmentMetadataCacheEntry metadata =
newTwoContainerBundleMetadata();
+ Assertions.assertTrue(location.reserve(metadata));
+ metadata.mount(location);
+
+ final PartialSegmentBundleCacheEntry bundle =
+ PartialSegmentBundleCacheEntry.forBundle(metadata,
TWO_CONTAINER_BUNDLE, List.of());
+ try (StorageLocation.ReservationHold<?> hold =
location.addWeakReservationHold(bundle.getId(), () -> bundle)) {
+ Assertions.assertNotNull(hold);
+ }
+ Assertions.assertFalse(location.isWeakReserved(bundle.getId()), "ephemeral
release should have evicted");
+
+ final List<PartialSegmentBundleCacheEntry.BundleContainerRef> refs =
bundle.getContainerRefs();
+ Assertions.assertEquals(2, refs.size());
+ final File firstContainerFile = containerFileFor(refs.getFirst());
+ FileUtils.mkdirp(containerFileFor(refs.getLast()));
+
+ Assertions.assertThrows(IOException.class, () -> bundle.mount(location));
+ Assertions.assertFalse(bundle.isMounted());
+
+ // Nothing will ever mount this entry again, and unmount only cleans up a
mount that committed, so the container
+ // this attempt initialized would have no owner. It has to be reaped by
the failed mount itself.
+ Assertions.assertFalse(
+ firstContainerFile.exists(),
+ "containers of an entry the cache has dropped should be evicted by the
failed mount"
+ );
+ }
+
+ @Test
+ void testUnmountReapsContainersLeftBehindByAFailedMount() throws IOException
+ {
+ final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE *
8, null);
+ final PartialSegmentMetadataCacheEntry metadata =
newTwoContainerBundleMetadata();
+ Assertions.assertTrue(location.reserve(metadata));
+ metadata.mount(location);
+
+ final PartialSegmentBundleCacheEntry bundle =
+ PartialSegmentBundleCacheEntry.forBundle(metadata,
TWO_CONTAINER_BUNDLE, List.of());
+ Assertions.assertNotNull(location.addWeakReservationHold(bundle.getId(),
() -> bundle));
+
+ final List<PartialSegmentBundleCacheEntry.BundleContainerRef> refs =
bundle.getContainerRefs();
+ Assertions.assertEquals(2, refs.size());
+ final File firstContainerFile = containerFileFor(refs.getFirst());
+ FileUtils.mkdirp(containerFileFor(refs.getLast()));
+
+ Assertions.assertThrows(IOException.class, () -> bundle.mount(location));
+ Assertions.assertFalse(bundle.isMounted());
+ Assertions.assertTrue(firstContainerFile.exists(), "the failed mount
should have kept its container");
+
+ // The entry is dropped without anyone retrying the mount. doActualUnmount
never runs for a mount that did not
+ // commit, and nothing else deletes container files, so unmount() has to
reap them or they outlive the entry with
+ // their reservation already released.
+ bundle.unmount();
+ Assertions.assertFalse(
+ firstContainerFile.exists(),
+ "unmount should evict containers left behind by a mount that never
committed"
+ );
+ }
+
+ /**
+ * A metadata entry for a segment whose {@link #TWO_CONTAINER_BUNDLE} bundle
spans two containers, one in the main
+ * file and one in an external file. Two containers is what lets a test fail
a mount partway through the
+ * container-initialization loop, with one container already initialized
behind it.
+ */
+ private PartialSegmentMetadataCacheEntry newTwoContainerBundleMetadata()
throws IOException
+ {
+ final int salt = ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE);
+ final File deepDir = temporaryFolder.newFolder("two_container_deep_" +
salt);
+ try (SegmentFileBuilderV10 builder =
SegmentFileBuilderV10.create(JSON_MAPPER, deepDir)) {
+ // Attach the external builder BEFORE startFileBundle so the bundle
propagates to it.
+ final SegmentFileBuilder external =
builder.getExternalBuilder(TWO_CONTAINER_EXTERNAL);
+ builder.startFileBundle(TWO_CONTAINER_BUNDLE);
+
+ final File mainTmp = temporaryFolder.newFile("two-container-main-" +
salt + ".bin");
+ Files.write(Ints.toByteArray(1), mainTmp);
+ builder.add(TWO_CONTAINER_BUNDLE + "/main_col", mainTmp);
+
+ final File extTmp = temporaryFolder.newFile("two-container-ext-" + salt
+ ".bin");
+ Files.write(Ints.toByteArray(2), extTmp);
+ external.add(TWO_CONTAINER_BUNDLE + "/ext_col", extTmp);
+ }
+ return new PartialSegmentMetadataCacheEntry(
+ SEGMENT_ID,
+ cacheDir,
+ IndexIO.V10_FILE_NAME,
+ List.of(TWO_CONTAINER_EXTERNAL),
+ new DirectoryBackedRangeReader(deepDir),
+ JSON_MAPPER,
+ null,
+ ESTIMATE,
+ PartialSegmentFileMapperV10.DEFAULT_COALESCE_GAP_BYTES,
+ PartialSegmentFileMapperV10.DEFAULT_MAX_FETCH_RUN_BYTES
+ );
+ }
+
+ private File
containerFileFor(PartialSegmentBundleCacheEntry.BundleContainerRef ref)
+ {
+ final String mapperFilename = ref.externalFilename() != null ?
ref.externalFilename() : IndexIO.V10_FILE_NAME;
+ return new File(cacheDir, StringUtils.format("%s.container.%05d",
mapperFilename, ref.containerIndex()));
+ }
+
@Test
void testAcquireReferenceBeforeMountThrows() throws IOException
{
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 36f7162cc9a..0d0b588747f 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
@@ -49,7 +49,9 @@ import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
@@ -57,6 +59,7 @@ class PartialSegmentMetadataCacheEntryTest
{
private static final ObjectMapper JSON_MAPPER = TestHelper.makeJsonMapper();
private static final SegmentId SEGMENT_ID = SegmentId.of("test",
Intervals.of("2025/2026"), "v1", 0);
+ private static final SegmentId OTHER_SEGMENT_ID = SegmentId.of("other",
Intervals.of("2025/2026"), "v1", 0);
private static final long ESTIMATE = 16 * 1024 * 1024L;
@RegisterExtension
@@ -122,6 +125,66 @@ class PartialSegmentMetadataCacheEntryTest
Assertions.assertFalse(headerFile.exists(), "mount failure must delete the
on-disk header file");
}
+ @Test
+ void testReclaimCannotEvictAnEntryWhileItIsMounting() throws Exception
+ {
+ // Room for one entry's reservation and no more, so a second reservation
can only succeed by reclaiming the first.
+ final StorageLocation location = new StorageLocation(cacheDir, ESTIMATE,
null);
+ final CountDownLatch entered = new CountDownLatch(1);
+ final CountDownLatch release = new CountDownLatch(1);
+ final PartialSegmentMetadataCacheEntry entry = newGatedEntry(ESTIMATE,
entered, release);
+
+ final StorageLocation.ReservationHold<SegmentCacheEntry> reserver =
+ location.addWeakReservationHold(entry.getId(), () -> entry);
+ Assertions.assertNotNull(reserver);
+
+ final ExecutorService exec = Execs.multiThreaded(1,
"mount-self-hold-test-%d");
+ try {
+ final Future<?> mounting = exec.submit(() -> {
+ entry.mount(location);
+ return null;
+ });
+ Assertions.assertTrue(entered.await(30, TimeUnit.SECONDS), "mount must
reach the gated read");
+
+ // The acquire that started this mount gives up. Nothing outside the
mount holds the entry now, which is
+ // exactly when reclaim would previously have been free to take it.
+ reserver.close();
+
+ // A reservation that only fits if this entry is evicted. It must fail
rather than pull the entry out from
+ // under the mount that is still filling it.
+ final PartialSegmentMetadataCacheEntry other = new
PartialSegmentMetadataCacheEntry(
+ OTHER_SEGMENT_ID,
+ cacheDir,
+ IndexIO.V10_FILE_NAME,
+ List.of(),
+ new DirectoryBackedRangeReader(segmentFile.getParentFile()),
+ JSON_MAPPER,
+ null,
+ ESTIMATE,
+ PartialSegmentFileMapperV10.DEFAULT_COALESCE_GAP_BYTES,
+ PartialSegmentFileMapperV10.DEFAULT_MAX_FETCH_RUN_BYTES
+ );
+ final StorageLocation.ReservationHold<SegmentCacheEntry> contender =
+ location.addWeakReservationHold(other.getId(), () -> other);
+ Assertions.assertNull(contender, "reclaim must not evict an entry whose
mount is still in flight");
+ Assertions.assertNotNull(
+ location.getCacheEntry(entry.getId()),
+ "the mounting entry must still be registered"
+ );
+
+ release.countDown();
+ mounting.get(30, TimeUnit.SECONDS);
+ }
+ finally {
+ release.countDown();
+ exec.shutdownNow();
+ }
+
+ // The mount ran to completion against an entry the location still knows
about.
+ Assertions.assertTrue(entry.isMounted());
+ Assertions.assertNotNull(location.getCacheEntry(entry.getId()));
+ }
+
@Test
void testMountIsIdempotentInSameLocation() throws IOException
{
@@ -581,6 +644,42 @@ class PartialSegmentMetadataCacheEntryTest
Assertions.assertEquals(List.of(),
entry.inferBundleDependencies("some_projection"));
}
+ /**
+ * A range reader that parks the first read until the test lets it through,
so a mount can be held mid-flight.
+ */
+ private PartialSegmentMetadataCacheEntry newGatedEntry(long estimate,
CountDownLatch entered, CountDownLatch release)
+ {
+ final SegmentRangeReader delegate = new
DirectoryBackedRangeReader(segmentFile.getParentFile());
+ final AtomicBoolean gated = new AtomicBoolean(false);
+ final SegmentRangeReader gatedReader = (filename, offset, length) -> {
+ if (gated.compareAndSet(false, true)) {
+ entered.countDown();
+ try {
+ if (!release.await(30, TimeUnit.SECONDS)) {
+ throw new IOException("timed out waiting for the test to release
the mount");
+ }
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException(e);
+ }
+ }
+ return delegate.readRange(filename, offset, length);
+ };
+ return new PartialSegmentMetadataCacheEntry(
+ SEGMENT_ID,
+ cacheDir,
+ IndexIO.V10_FILE_NAME,
+ List.of(),
+ gatedReader,
+ JSON_MAPPER,
+ null,
+ estimate,
+ PartialSegmentFileMapperV10.DEFAULT_COALESCE_GAP_BYTES,
+ PartialSegmentFileMapperV10.DEFAULT_MAX_FETCH_RUN_BYTES
+ );
+ }
+
private PartialSegmentMetadataCacheEntry newEntry(long estimate)
{
return new PartialSegmentMetadataCacheEntry(
diff --git
a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerConcurrencyTest.java
b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerConcurrencyTest.java
index 921d1597cc2..18f6a086626 100644
---
a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerConcurrencyTest.java
+++
b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerConcurrencyTest.java
@@ -434,15 +434,15 @@ class SegmentLocalCacheManagerConcurrencyTest
Assertions.assertTrue(t instanceof TimeoutException || t instanceof
ExecutionException, t.toString());
}
Thread.sleep(20);
- Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
- Assertions.assertEquals(0, location2.getWeakStats().getHoldCount());
+ awaitNoHolds(location);
+ awaitNoHolds(location2);
currentBatch.clear();
}
}
- Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
- Assertions.assertEquals(0, location2.getWeakStats().getHoldCount());
+ awaitNoHolds(location);
+ awaitNoHolds(location2);
Assertions.assertTrue(4 >= location.getWeakEntryCount());
Assertions.assertTrue(4 >= location2.getWeakEntryCount());
// 5 because __drop path
@@ -500,8 +500,8 @@ class SegmentLocalCacheManagerConcurrencyTest
}
Thread.sleep(5);
}
- Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
- Assertions.assertEquals(0, location2.getWeakStats().getHoldCount());
+ awaitNoHolds(location);
+ awaitNoHolds(location2);
currentBatch.clear();
}
}
@@ -569,8 +569,8 @@ class SegmentLocalCacheManagerConcurrencyTest
}
}
- Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
- Assertions.assertEquals(0, location2.getWeakStats().getHoldCount());
+ awaitNoHolds(location);
+ awaitNoHolds(location2);
totalSuccess += success;
totalEmpty += empty;
totalRows += rows;
@@ -592,6 +592,85 @@ class SegmentLocalCacheManagerConcurrencyTest
assertNoLooseEnds();
}
+ @Test
+ public void testAcquireSegmentOnDemandRandomSegmentWithInterrupt() throws
IOException, InterruptedException
+ {
+ final int segmentCount = 8;
+ final int iterations = 2000;
+ final int concurrentReads = 10;
+ final File localStorageFolder = new File(tempDir, "local_storage_folder");
+
+ final Interval interval = Intervals.of("2019-01-01/P1D");
+
+ makeSegmentsToLoad(segmentCount, localStorageFolder, interval,
segmentsToWeakLoad);
+
+ final List<DataSegment> currentBatch = new ArrayList<>();
+ for (int i = 0; i < iterations; i++) {
+
currentBatch.add(segmentsToWeakLoad.get(ThreadLocalRandom.current().nextInt(segmentCount)));
+ // process batches of 10 requests at a time
+ if (currentBatch.size() == concurrentReads) {
+ final List<InterruptedLoad> weakLoads = currentBatch
+ .stream()
+ .map(segment -> new InterruptedLoad(virtualStorageManager,
segment))
+ .collect(Collectors.toList());
+ final List<Future<Integer>> futures = new ArrayList<>();
+ for (InterruptedLoad weakLoad : weakLoads) {
+ futures.add(executorService.submit(weakLoad));
+ }
+ for (Future<Integer> future : futures) {
+ try {
+ future.get(20L, TimeUnit.MILLISECONDS);
+ }
+ catch (Throwable t) {
+ }
+ }
+ while (true) {
+ boolean allDone = true;
+ for (Future<?> f : futures) {
+ allDone = allDone && f.isDone();
+ }
+ if (allDone) {
+ break;
+ }
+ Thread.sleep(5);
+ }
+ awaitNoHolds(location);
+ awaitNoHolds(location2);
+ currentBatch.clear();
+ }
+ }
+
+ Assertions.assertTrue(location.getWeakStats().getHitCount() >= 0);
+ Assertions.assertTrue(location.getWeakStats().getHitBytes() >= 0);
+ Assertions.assertTrue(location2.getWeakStats().getHitCount() >= 0);
+ Assertions.assertTrue(location2.getWeakStats().getHitBytes() >= 0);
+
+ // now ensure that we can successfully do stuff after all those interrupts
+ int totalSuccess = 0;
+ int totalFailures = 0;
+ for (int i = 0; i < iterations; i++) {
+ int segment = ThreadLocalRandom.current().nextInt(segmentCount);
+ currentBatch.add(segmentsToWeakLoad.get(segment));
+ // process batches of 10 requests at a time
+ if (currentBatch.size() == concurrentReads) {
+
+ BatchResult result = testWeakBatch(i, currentBatch, false);
+ totalSuccess += result.success;
+ totalFailures += result.exceptions.size();
+ currentBatch.clear();
+ }
+ }
+ Assertions.assertEquals(iterations, totalSuccess);
+ Assertions.assertEquals(0, totalFailures);
+ awaitNoHolds(location);
+ awaitNoHolds(location2);
+ Assertions.assertTrue(4 >= location.getWeakEntryCount());
+ Assertions.assertTrue(4 >= location2.getWeakEntryCount());
+ // 5 because __drop path
+ Assertions.assertTrue(5 >= location.getPath().listFiles().length);
+ Assertions.assertTrue(5 >= location2.getPath().listFiles().length);
+ }
+
private void testWeakLoad(
int iterations,
int segmentCount,
@@ -735,89 +814,10 @@ class SegmentLocalCacheManagerConcurrencyTest
return new BatchResult(exceptions, success, rows);
}
- @Test
- public void testAcquireSegmentOnDemandRandomSegmentWithInterrupt() throws
IOException, InterruptedException
- {
- final int segmentCount = 8;
- final int iterations = 2000;
- final int concurrentReads = 10;
- final File localStorageFolder = new File(tempDir, "local_storage_folder");
-
- final Interval interval = Intervals.of("2019-01-01/P1D");
-
- makeSegmentsToLoad(segmentCount, localStorageFolder, interval,
segmentsToWeakLoad);
-
- final List<DataSegment> currentBatch = new ArrayList<>();
- for (int i = 0; i < iterations; i++) {
-
currentBatch.add(segmentsToWeakLoad.get(ThreadLocalRandom.current().nextInt(segmentCount)));
- // process batches of 10 requests at a time
- if (currentBatch.size() == concurrentReads) {
- final List<InterruptedLoad> weakLoads = currentBatch
- .stream()
- .map(segment -> new InterruptedLoad(virtualStorageManager,
segment))
- .collect(Collectors.toList());
- final List<Future<Integer>> futures = new ArrayList<>();
- for (InterruptedLoad weakLoad : weakLoads) {
- futures.add(executorService.submit(weakLoad));
- }
- for (Future<Integer> future : futures) {
- try {
- future.get(20L, TimeUnit.MILLISECONDS);
- }
- catch (Throwable t) {
- }
- }
- while (true) {
- boolean allDone = true;
- for (Future<?> f : futures) {
- allDone = allDone && f.isDone();
- }
- if (allDone) {
- break;
- }
- Thread.sleep(5);
- }
- Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
- Assertions.assertEquals(0, location2.getWeakStats().getHoldCount());
- currentBatch.clear();
- }
- }
-
- Assertions.assertTrue(location.getWeakStats().getHitCount() >= 0);
- Assertions.assertTrue(location.getWeakStats().getHitBytes() >= 0);
- Assertions.assertTrue(location2.getWeakStats().getHitCount() >= 0);
- Assertions.assertTrue(location2.getWeakStats().getHitBytes() >= 0);
-
- // now ensure that we can successfully do stuff after all those interrupts
- int totalSuccess = 0;
- int totalFailures = 0;
- for (int i = 0; i < iterations; i++) {
- int segment = ThreadLocalRandom.current().nextInt(segmentCount);
- currentBatch.add(segmentsToWeakLoad.get(segment));
- // process batches of 10 requests at a time
- if (currentBatch.size() == concurrentReads) {
-
- BatchResult result = testWeakBatch(i, currentBatch, false);
- totalSuccess += result.success;
- totalFailures += result.exceptions.size();
- currentBatch.clear();
- }
- }
- Assertions.assertEquals(iterations, totalSuccess);
- Assertions.assertEquals(0, totalFailures);
- Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
- Assertions.assertEquals(0, location2.getWeakStats().getHoldCount());
- Assertions.assertTrue(4 >= location.getWeakEntryCount());
- Assertions.assertTrue(4 >= location2.getWeakEntryCount());
- // 5 because __drop path
- Assertions.assertTrue(5 >= location.getPath().listFiles().length);
- Assertions.assertTrue(5 >= location2.getPath().listFiles().length);
- }
-
private void assertNoLooseEnds()
{
- Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
- Assertions.assertEquals(0, location2.getWeakStats().getHoldCount());
+ awaitNoHolds(location);
+ awaitNoHolds(location2);
Assertions.assertTrue(4 >= location.getWeakEntryCount());
Assertions.assertTrue(4 >= location2.getWeakEntryCount());
// 5 because __drop path
@@ -922,6 +922,27 @@ class SegmentLocalCacheManagerConcurrencyTest
.build();
}
+ /**
+ * Waits for every hold on a location to be released, then asserts there are
none.
+ * <p>
+ * A mount holds its own entry until it has finished establishing state, and
abandoning the acquire that triggered
+ * it does not stop that mount, so a hold can briefly outlive the caller
that asked for it. Anything that leaks a
+ * hold still fails here, it just takes the timeout to say so.
+ */
+ private static void awaitNoHolds(StorageLocation location)
+ {
+ for (int i = 0; i < 300 && location.getWeakStats().getHoldCount() > 0;
i++) {
+ try {
+ Thread.sleep(10);
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ break;
+ }
+ }
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ }
+
private static class BatchResult
{
public final List<Throwable> exceptions;
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 4a5be4714b3..63c42d1e1c6 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
@@ -19,6 +19,9 @@
package org.apache.druid.segment.loading;
+import com.fasterxml.jackson.annotation.JacksonInject;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.InjectableValues;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.jsontype.NamedType;
@@ -79,6 +82,7 @@ import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
+import javax.annotation.Nullable;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
@@ -95,6 +99,7 @@ import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
class SegmentLocalCacheManagerPartialAcquireTest
@@ -305,6 +310,59 @@ class SegmentLocalCacheManagerPartialAcquireTest
}
}
+ @Test
+ void testCancellingFullAcquireMidDownloadDoesNotEvictContainersUnderIt()
throws Exception
+ {
+ // Ephemeral weak entries, as an MSQ/Dart worker runs: releasing the last
hold on a bundle drops it from the cache
+ // then and there, which is what turns an abandoned acquire into an
unmount of the bundles it was downloading.
+ manager.getLocations().get(0).setAreWeakEntriesEphemeral(true);
+
+ jsonMapper.registerSubtypes(new NamedType(GatedLocalLoadSpec.class,
GatedLocalLoadSpec.TYPE));
+ final DataSegment gatedSegment = DataSegment.builder(SEGMENT_ID)
+
.shardSpec(NoneShardSpec.instance())
+ .loadSpec(Map.of(
+ "type",
GatedLocalLoadSpec.TYPE,
+ "path",
DEEP_STORAGE_DIR.getAbsolutePath()
+ ))
+ .size(0)
+ .build();
+
+ // Warm the metadata entry first, and keep it held for the rest of the
test. Mounting metadata range-reads the
+ // header through the same reader the download uses, so warming it means
the gate below can only be reached from
+ // inside ensureAllDownloaded. A partial acquire downloads nothing else,
so no container is resident yet.
+ try (AcquireSegmentAction warm = manager.acquireSegment(gatedSegment,
AcquireMode.PARTIAL)) {
+ warm.getSegmentFuture().get();
+
+ final DownloadGate gate = new DownloadGate();
+ DOWNLOAD_GATE.set(gate);
+ try {
+ final AcquireSegmentAction full = manager.acquireSegment(gatedSegment,
AcquireMode.FULL);
+ // Held across the close() below: getSegmentFuture() refuses to hand
out the future once the action is closed,
+ // and the load task keeps running either way - closing an action
cancels the caller's interest, not the task.
+ final ListenableFuture<AcquireSegmentResult> future =
full.getSegmentFuture();
+
+ Assertions.assertTrue(
+ gate.entered.await(60, TimeUnit.SECONDS),
+ "the full download should have reached a container fetch"
+ );
+
+ // Abandon the acquire while a fetch is parked mid-write. This closes
the action's HoldHolder on THIS thread,
+ // which is what used to unmount the bundles - and evict the
containers - out from under the running download.
+ full.close();
+ gate.release.countDown();
+
+ // The download must still finish against containers that are all
still there. Before the download owned its
+ // own bundle references, this failed: either an NPE from a fetch
whose container file had been deleted, or
+ // ensureAllDownloaded's post-condition once the eviction cleared
residency behind it.
+ final AcquireSegmentResult result = future.get(60, TimeUnit.SECONDS);
+ Assertions.assertNotNull(result);
+ }
+ finally {
+ DOWNLOAD_GATE.set(null);
+ }
+ }
+ }
+
@Test
void testAcquirePartialSegmentReturnsPartialAwareSegment() throws
ExecutionException, InterruptedException, IOException
{
@@ -1172,4 +1230,82 @@ class SegmentLocalCacheManagerPartialAcquireTest
}
}
}
+
+ /**
+ * Parks the first range read that reaches it until the test lets it go, so
a download can be caught mid-write.
+ * Consulted by every reader {@link GatedLocalLoadSpec} hands out, since the
reader doing the download is the one the
+ * file mapper was created with, not the one the acquire being tested opened.
+ */
+ private static final AtomicReference<DownloadGate> DOWNLOAD_GATE = new
AtomicReference<>();
+
+ private static final class DownloadGate
+ {
+ private final CountDownLatch entered = new CountDownLatch(1);
+ private final CountDownLatch release = new CountDownLatch(1);
+
+ private void arrive() throws IOException
+ {
+ entered.countDown();
+ try {
+ if (!release.await(60, TimeUnit.SECONDS)) {
+ throw new IOException("Timed out waiting for the test to release the
download gate");
+ }
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException(e);
+ }
+ }
+ }
+
+ /**
+ * A {@link LoadSpec} that loads exactly like {@link LocalLoadSpec} but
hands out range readers that consult
+ * {@link #DOWNLOAD_GATE} first. Registered as a Jackson subtype by the test
that needs it.
+ */
+ public static class GatedLocalLoadSpec implements LoadSpec
+ {
+ static final String TYPE = "gated-local";
+
+ private final LocalLoadSpec delegate;
+ private final String path;
+
+ @JsonCreator
+ public GatedLocalLoadSpec(
+ @JacksonInject LocalDataSegmentPuller puller,
+ @JsonProperty(value = "path", required = true) String path
+ )
+ {
+ this.delegate = new LocalLoadSpec(puller, path);
+ this.path = path;
+ }
+
+ @JsonProperty
+ public String getPath()
+ {
+ return path;
+ }
+
+ @Override
+ public LoadSpecResult loadSegment(File destDir) throws
SegmentLoadingException
+ {
+ return delegate.loadSegment(destDir);
+ }
+
+ @Nullable
+ @Override
+ public SegmentRangeReader openRangeReader() throws IOException
+ {
+ final SegmentRangeReader reader = delegate.openRangeReader();
+ if (reader == null) {
+ return null;
+ }
+ return (filename, offset, length) -> {
+ final DownloadGate gate = DOWNLOAD_GATE.get();
+ if (gate != null) {
+ gate.arrive();
+ }
+ return reader.readRange(filename, offset, length);
+ };
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
