This is an automated email from the ASF dual-hosted git repository.
kfaraz 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 543256bd422 fix: loading pool callers decide when to permit (#19717)
543256bd422 is described below
commit 543256bd422f6f17a1ab3952d0aab538fc22f30b
Author: Clint Wylie <[email protected]>
AuthorDate: Wed Jul 22 00:45:37 2026 -0700
fix: loading pool callers decide when to permit (#19717)
Reworks `StorageLoadingThreadPool` to no longer bake the semaphore permits
into the executor,
and instead acquire the permits as close as possible to the actual deep
storage reads, combining
them was a mistake and it could at least in some scenarios lead to
deadlocks.
`SegmentRangeReader` is done with a wrapper that bakes the permit process
into the stream for
range reads; FilePopulator similarly wraps populate in an acquire, and the
complete segment load
path acquires directly when calling `loadInLocation`.
---
.../loading/PermitLimitedSegmentRangeReader.java | 86 ++++++
.../segment/loading/SegmentLocalCacheManager.java | 35 ++-
.../segment/loading/StorageLoadingThreadPool.java | 302 ++++++++-------------
.../StorageLocationVirtualStorageManager.java | 10 +-
.../PermitBoundedListeningExecutorServiceTest.java | 106 --------
.../PermitLimitedSegmentRangeReaderTest.java | 88 ++++++
.../loading/StorageLoadingThreadPoolTest.java | 85 ++++++
7 files changed, 413 insertions(+), 299 deletions(-)
diff --git
a/server/src/main/java/org/apache/druid/segment/loading/PermitLimitedSegmentRangeReader.java
b/server/src/main/java/org/apache/druid/segment/loading/PermitLimitedSegmentRangeReader.java
new file mode 100644
index 00000000000..f6edc4e8dde
--- /dev/null
+++
b/server/src/main/java/org/apache/druid/segment/loading/PermitLimitedSegmentRangeReader.java
@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.segment.loading;
+
+import java.io.Closeable;
+import java.io.FilterInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+
+/**
+ * A {@link SegmentRangeReader} decorator that bounds concurrent deep-storage
range reads: each {@link #readRange} call
+ * acquires an on-demand-load permit from {@link
StorageLoadingThreadPool#acquireLoadPermit()} and holds it for the
+ * lifetime of the returned stream (released on close), so the permit spans
exactly the wire transfer.
+ */
+public class PermitLimitedSegmentRangeReader implements SegmentRangeReader
+{
+ private final SegmentRangeReader delegate;
+ private final StorageLoadingThreadPool pool;
+
+ public PermitLimitedSegmentRangeReader(SegmentRangeReader delegate,
StorageLoadingThreadPool pool)
+ {
+ this.delegate = delegate;
+ this.pool = pool;
+ }
+
+ @Override
+ public InputStream readRange(String filename, long offset, long length)
throws IOException
+ {
+ final Closeable permit = pool.acquireLoadPermit();
+ try {
+ return new PermitReleasingInputStream(delegate.readRange(filename,
offset, length), permit);
+ }
+ catch (Throwable t) {
+ try {
+ permit.close();
+ }
+ catch (Exception e) {
+ t.addSuppressed(e);
+ }
+ throw t;
+ }
+ }
+
+ /**
+ * Releases the permit when the stream is closed (after the transfer
completes or the reader aborts). The permit's
+ * {@code close()} is idempotent, so a double-close of the stream
over-releases nothing.
+ */
+ private static final class PermitReleasingInputStream extends
FilterInputStream
+ {
+ private final Closeable permit;
+
+ PermitReleasingInputStream(InputStream in, Closeable permit)
+ {
+ super(in);
+ this.permit = permit;
+ }
+
+ @Override
+ public void close() throws IOException
+ {
+ try {
+ super.close();
+ }
+ finally {
+ permit.close();
+ }
+ }
+ }
+}
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 5059c932a03..ec99cd4c1c3 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
@@ -664,13 +664,15 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
// executor tasks. The entry's own mount-future dedup would
prevent the actual work from being duplicated,
// but the executor scheduling and timing capture would still be
wasted.
Suppliers.memoize(() -> {
- // Capture submit time on first invocation of
getSegmentFuture(), so waitTime measures the queue delay
- // until the executor picks up the task. loadTime then covers
mount (+ ensureAllDownloaded for the
- // full-download path).
+ // Capture submit time on first invocation of
getSegmentFuture(). loadTime then covers mount
+ // (+ ensureAllDownloaded for the full-download path).
final long submitNanos = System.nanoTime();
return
virtualStorageLoadingThreadPool.getExecutorService().submit(() -> {
- // The executor bounds concurrency itself (permit acquired
inside the task on the worker thread), so
- // waitNanos measures both the queue delay and any permit wait
until this task actually starts.
+ // waitNanos is only the executor scheduling delay until this
task body starts; it no longer reflects
+ // load-slot contention when using virtual threads. Load
permits are acquired inside the deep-storage
+ // reads now, so that wait is folded into loadTime instead,
and the query-time bundle/column fetches
+ // (separate permit-bounded tasks) are not reflected here at
all. A meaningful load-wait metric would
+ // have to time the permit acquire at the read sites and
aggregate it across those fetches.
final long taskStartNanos = System.nanoTime();
final long waitNanos = taskStartNanos - submitNanos;
final boolean wasMounted = reserved.metadata.isMounted();
@@ -829,7 +831,12 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
}
try {
final LoadSpec loadSpec =
jsonMapper.convertValue(dataSegment.getLoadSpec(), LoadSpec.class);
- return loadSpec.openRangeReader();
+ final SegmentRangeReader rangeReader = loadSpec.openRangeReader();
+ if (rangeReader == null) {
+ return null;
+ }
+ // Bound concurrent deep-storage reads at the actual range-read (see
PermitLimitedSegmentRangeReader).
+ return new PermitLimitedSegmentRangeReader(rangeReader,
virtualStorageLoadingThreadPool);
}
catch (IOException e) {
throw DruidException.forPersona(DruidException.Persona.OPERATOR)
@@ -1235,7 +1242,12 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
throws SegmentLoadingException
{
try {
- return wrapper.openRangeReader();
+ final SegmentRangeReader rangeReader = wrapper.openRangeReader();
+ if (rangeReader == null) {
+ return null;
+ }
+ // Bound concurrent deep-storage reads at the actual range-read (see
PermitLimitedSegmentRangeReader).
+ return new PermitLimitedSegmentRangeReader(rangeReader,
virtualStorageLoadingThreadPool);
}
catch (IOException e) {
throw new SegmentLoadingException(e, "Failed to open range reader for
segment[%s]", dataSegment.getId());
@@ -1659,6 +1671,9 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
final long startTime = System.nanoTime();
return virtualStorageLoadingThreadPool.getExecutorService().submit(
() -> {
+ // waitTime is only the executor scheduling delay; when using
virtual threads for the pool, load-slot
+ // contention is folded into loadTime now, since mount
acquires the load permit around the deep-storage
+ // read.
final long execStartTime = System.nanoTime();
final long waitTime = execStartTime - startTime;
entry.mount(location);
@@ -1952,7 +1967,11 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
}
}
if (needsLoad) {
- loadInLocationWithStartMarker(dataSegment, storageDir);
+ // Hold a load permit only around the actual deep-storage read,
not the surrounding entryLock or the
+ // factorize/deserialize below. Acquired after entryLock, so a
permit is never held while blocking on it.
+ try (StorageLoadingThreadPool.LoadPermit ignored =
virtualStorageLoadingThreadPool.acquireLoadPermit()) {
+ loadInLocationWithStartMarker(dataSegment, storageDir);
+ }
}
final SegmentizerFactory factory = getSegmentFactory(storageDir);
diff --git
a/server/src/main/java/org/apache/druid/segment/loading/StorageLoadingThreadPool.java
b/server/src/main/java/org/apache/druid/segment/loading/StorageLoadingThreadPool.java
index b76f347cb14..866e5eeeb82 100644
---
a/server/src/main/java/org/apache/druid/segment/loading/StorageLoadingThreadPool.java
+++
b/server/src/main/java/org/apache/druid/segment/loading/StorageLoadingThreadPool.java
@@ -19,9 +19,6 @@
package org.apache.druid.segment.loading;
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.util.concurrent.ForwardingListeningExecutorService;
-import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import org.apache.druid.common.asyncresource.AsyncResource;
@@ -30,81 +27,75 @@ import org.apache.druid.error.DruidException;
import org.apache.druid.java.util.common.concurrent.Execs;
import org.apache.druid.java.util.common.lifecycle.LifecycleStop;
import org.apache.druid.java.util.common.logger.Logger;
-import org.apache.druid.segment.PartialBundleAcquirer;
import org.apache.druid.segment.loading.external.VirtualStorageManager;
import javax.annotation.Nullable;
import java.io.Closeable;
-import java.util.Collection;
-import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
import java.util.concurrent.Semaphore;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
/**
* Holds the thread pool used for background loading by {@link
SegmentLocalCacheManager} and
* {@link VirtualStorageManager}.
+ * <p>
+ * <b>Submissions are not automatically concurrency-bounded.</b> The executor
returned by {@link #getExecutorService()}
+ * runs submitted tasks as fast as it can (in the virtual-thread mode, one
virtual thread per task). Instead, the
+ * number of concurrent <em>deep-storage reads</em> is bounded by a permit
that callers acquire via
+ * {@link #acquireLoadPermit()} around <em>only</em> the actual I/O, never
around lock acquisition, reservation, or
+ * deserialization. Any new load path that reads from deep storage must
acquire a permit around that read, or it will
+ * be unbounded.
*/
public class StorageLoadingThreadPool
{
- private static final Logger log = new Logger(StorageLoadingThreadPool.class);
-
- private final ListeningExecutorService exec;
-
- public StorageLoadingThreadPool(
- @Nullable final ListeningExecutorService exec
- )
- {
- this.exec = exec;
- }
-
public static StorageLoadingThreadPool createFromConfig(final
SegmentLoaderConfig config)
{
- final ListeningExecutorService exec;
+ if (!config.isVirtualStorage()) {
+ return new StorageLoadingThreadPool(null, null);
+ }
- if (config.isVirtualStorage()) {
- if (config.getVirtualStorageLoadThreads() <= 0) {
- throw DruidException.forPersona(DruidException.Persona.OPERATOR)
- .ofCategory(DruidException.Category.INVALID_INPUT)
- .build(
- "virtualStorageLoadThreads must be greater
than 0, got [%d]",
- config.getVirtualStorageLoadThreads()
- );
- }
- if (config.isVirtualStorageUseVirtualThreads()) {
- log.info(
- "Using virtual storage mode with virtual threads - max concurrent
on demand loads: [%d].",
- config.getVirtualStorageLoadThreads()
- );
- exec = new PermitBoundedListeningExecutorService(
- MoreExecutors.listeningDecorator(
- Executors.newThreadPerTaskExecutor(
- Thread.ofVirtual()
- .name("VirtualStorageOnDemandLoadingThread-", 0)
- .factory()
- )
- ),
- new Semaphore(config.getVirtualStorageLoadThreads())
- );
- } else {
- log.info(
- "Using virtual storage mode with fixed platform thread pool - on
demand load threads: [%d].",
- config.getVirtualStorageLoadThreads()
- );
- exec = MoreExecutors.listeningDecorator(
- Executors.newFixedThreadPool(
- config.getVirtualStorageLoadThreads(),
-
Execs.makeThreadFactory("VirtualStorageOnDemandLoadingThread-%s")
- )
- );
- }
- } else {
- exec = null;
+ if (config.getVirtualStorageLoadThreads() <= 0) {
+ throw DruidException.forPersona(DruidException.Persona.OPERATOR)
+ .ofCategory(DruidException.Category.INVALID_INPUT)
+ .build(
+ "virtualStorageLoadThreads must be greater than
0, got [%d]",
+ config.getVirtualStorageLoadThreads()
+ );
}
- return new StorageLoadingThreadPool(exec);
+ final ListeningExecutorService exec;
+ final Semaphore permits;
+ if (config.isVirtualStorageUseVirtualThreads()) {
+ log.info(
+ "Using virtual storage mode with virtual threads - max concurrent on
demand loads: [%d].",
+ config.getVirtualStorageLoadThreads()
+ );
+ // Unbounded thread-per-virtual-thread executor; concurrency is bounded
by the permit count, acquired by callers
+ // via acquireLoadPermit() around the actual deep-storage reads.
+ exec = MoreExecutors.listeningDecorator(
+ Executors.newThreadPerTaskExecutor(
+ Thread.ofVirtual()
+ .name("VirtualStorageOnDemandLoadingThread-", 0)
+ .factory()
+ )
+ );
+ permits = new Semaphore(config.getVirtualStorageLoadThreads());
+ } else {
+ log.info(
+ "Using virtual storage mode with fixed platform thread pool - on
demand load threads: [%d].",
+ config.getVirtualStorageLoadThreads()
+ );
+ // Fixed pool: the thread count is the concurrency bound, so no separate
permit is needed.
+ exec = MoreExecutors.listeningDecorator(
+ Executors.newFixedThreadPool(
+ config.getVirtualStorageLoadThreads(),
+ Execs.makeThreadFactory("VirtualStorageOnDemandLoadingThread-%s")
+ )
+ );
+ permits = null;
+ }
+ return new StorageLoadingThreadPool(exec, permits);
}
/**
@@ -112,7 +103,31 @@ public class StorageLoadingThreadPool
*/
public static StorageLoadingThreadPool none()
{
- return new StorageLoadingThreadPool(null);
+ return new StorageLoadingThreadPool(null, null);
+ }
+
+ private static final Logger log = new Logger(StorageLoadingThreadPool.class);
+
+ /**
+ * A permit handle whose {@code close()} releases nothing. Returned by
{@link #acquireLoadPermit()} when there is no
+ * semaphore (the fixed-thread-pool mode, where the thread count is the
bound, and the "no pool" instance).
+ */
+ private static final LoadPermit NOOP_PERMIT = () -> {};
+
+ @Nullable
+ private final ListeningExecutorService exec;
+ /**
+ * Bounds concurrent on-demand deep-storage reads in the virtual-thread
mode, where the executor is otherwise
+ * unbounded (one virtual thread per task). Null in the fixed-thread-pool
mode (the pool size is the bound) and in
+ * the "no pool" instance.
+ */
+ @Nullable
+ private final Semaphore permits;
+
+ private StorageLoadingThreadPool(@Nullable final ListeningExecutorService
exec, @Nullable final Semaphore permits)
+ {
+ this.exec = exec;
+ this.permits = permits;
}
public boolean isAvailable()
@@ -121,11 +136,8 @@ public class StorageLoadingThreadPool
}
/**
- * Executor for on-demand load work. Concurrency is bounded by the executor
itself: in the virtual-thread path it is
- * a {@link PermitBoundedListeningExecutorService} wrapping an unbounded
thread-per-virtual-thread executor (so the
- * permit count, not the thread count, is the bound); in the fixed-pool path
the pool size is the bound. Every
- * on-demand-load submission, including those routed through {@link
#submitUnmanagedAsyncResource}, is therefore
- * bounded without callers having to acquire a permit themselves.
+ * Executor for on-demand load work. Not concurrency-bounded by itself.
Callers doing deep-storage reads with this
+ * pool must bound themselves via {@link #acquireLoadPermit()} around the
read.
*/
public ListeningExecutorService getExecutorService()
{
@@ -135,6 +147,37 @@ public class StorageLoadingThreadPool
return exec;
}
+ /**
+ * Acquire one on-demand-load permit, returning a {@link LoadPermit} that
releases it (idempotently) on close. Callers
+ * must hold the permit only around the actual deep-storage read, not around
lock acquisition, reservation, or
+ * deserialization.
+ * <p>
+ * In the fixed-thread-pool mode (or the "no pool" instance) there is no
semaphore and this returns a no-op handle
+ * (the thread count, or nothing, is the bound). The acquire is
interruptible: if the thread is interrupted while
+ * waiting for a permit (e.g. a canceled load), the interrupt flag is
restored and a {@link RuntimeException} is
+ * thrown so the load aborts before any I/O begins.
+ */
+ public LoadPermit acquireLoadPermit()
+ {
+ final Semaphore p = permits;
+ if (p == null) {
+ return NOOP_PERMIT;
+ }
+ try {
+ p.acquire();
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ final AtomicBoolean released = new AtomicBoolean(false);
+ return () -> {
+ if (released.compareAndSet(false, true)) {
+ p.release();
+ }
+ };
+ }
+
/**
* Submit a task to the pool and hand back an {@link AsyncResource} that
becomes ready when the task completes,
* exposing the task's (non-null) result. The result is treated as a plain
value with no lifecycle: closing the
@@ -143,7 +186,7 @@ public class StorageLoadingThreadPool
* <p>This is the unmanaged counterpart of {@link
#submitCloseableAsyncResource}; use that when the task's result
* owns a lifecycle.
*
- * @see AsyncResources#fromFutureUnmanaged(ListenableFuture)
+ * @see AsyncResources#fromFutureUnmanaged
*/
public <T> AsyncResource<T> submitUnmanagedAsyncResource(Callable<T> task)
{
@@ -158,7 +201,7 @@ public class StorageLoadingThreadPool
* <p>This is the managed counterpart of {@link
#submitUnmanagedAsyncResource}; use that when the task's result is a
* plain value with no lifecycle.
*
- * @see AsyncResources#fromFutureCloseable(ListenableFuture)
+ * @see AsyncResources#fromFutureCloseable
*/
public <T extends Closeable> AsyncResource<T>
submitCloseableAsyncResource(Callable<T> task)
{
@@ -174,125 +217,16 @@ public class StorageLoadingThreadPool
}
/**
- * A {@link ListeningExecutorService} that caps the number of
concurrently-running submitted tasks at a semaphore's
- * permit count, acquiring a permit (on the worker thread) before each task
body runs and releasing it after. Used to
- * bound concurrent virtual-storage on-demand loads when the backing
executor is an unbounded thread-per-virtual-thread
- * executor: the permit count is the concurrency bound, not the thread
count, and the wait for a permit parks a virtual
- * thread rather than blocking the submitter.
- * <p>
- * The permit wait is <b>interruptible</b>: a task whose future is canceled
with {@code mayInterruptIfRunning} (or a
- * {@code shutdownNow}) while it is parked on the permit is interrupted out
of the wait and aborts before running its
- * body. Canceling a query stops not only its queued column downloads but
also those blocked on the permit,
- * before any deep-storage I/O begins. A task that has already passed the
permit and started its body runs to
- * completion (aborting in-flight reads is handled separately by the task
itself, not here).
- * <p>
- * Only {@code execute}/{@code submit} are bounded; those are the only
submission paths on-demand load work uses
- * (including the {@link PartialBundleAcquirer#submitDownload}
column-download path). {@code invokeAll}/
- * {@code invokeAny} throw {@link UnsupportedOperationException} so the
concurrency bound can never be silently
- * bypassed by a future caller.
- * <p>
- * Callers must not submit a task that itself blocks on another task
submitted to this executor while holding a
- * permit, or all permits could be exhausted by waiters; the on-demand load
tasks here never nest submissions.
+ * A handle to a held on-demand-load permit. {@link #close()} releases the
permit.
+ *
+ * @see #acquireLoadPermit()
*/
- @VisibleForTesting
- static final class PermitBoundedListeningExecutorService extends
ForwardingListeningExecutorService
+ public interface LoadPermit extends Closeable
{
- private final ListeningExecutorService delegate;
- private final Semaphore permits;
-
- PermitBoundedListeningExecutorService(ListeningExecutorService delegate,
Semaphore permits)
- {
- this.delegate = delegate;
- this.permits = permits;
- }
-
+ /**
+ * Releases the permit.
+ */
@Override
- protected ListeningExecutorService delegate()
- {
- return delegate;
- }
-
- @Override
- public void execute(Runnable command)
- {
- delegate.execute(withPermit(command));
- }
-
- @Override
- public <T> ListenableFuture<T> submit(Callable<T> task)
- {
- return delegate.submit(withPermit(task));
- }
-
- @Override
- public ListenableFuture<?> submit(Runnable task)
- {
- return delegate.submit(withPermit(task));
- }
-
- @Override
- public <T> ListenableFuture<T> submit(Runnable task, T result)
- {
- return delegate.submit(withPermit(task), result);
- }
-
- @Override
- public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>>
tasks)
- {
- throw new UnsupportedOperationException("invokeAll is not
permit-bounded; use submit");
- }
-
- @Override
- public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>>
tasks, long timeout, TimeUnit unit)
- {
- throw new UnsupportedOperationException("invokeAll is not
permit-bounded; use submit");
- }
-
- @Override
- public <T> T invokeAny(Collection<? extends Callable<T>> tasks)
- {
- throw new UnsupportedOperationException("invokeAny is not
permit-bounded; use submit");
- }
-
- @Override
- public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long
timeout, TimeUnit unit)
- {
- throw new UnsupportedOperationException("invokeAny is not
permit-bounded; use submit");
- }
-
- private Runnable withPermit(Runnable task)
- {
- return () -> {
- try {
- permits.acquire();
- }
- catch (InterruptedException e) {
- // Interrupted while waiting for a permit (e.g. the task's future
was canceled with mayInterruptIfRunning).
- // Abort before running the body. Restore the flag and surface as a
failure rather than a silent success;
- // if the future was canceled it already reports as such, so this
only matters for a stray interrupt.
- Thread.currentThread().interrupt();
- throw new RuntimeException(e);
- }
- try {
- task.run();
- }
- finally {
- permits.release();
- }
- };
- }
-
- private <T> Callable<T> withPermit(Callable<T> task)
- {
- return () -> {
- permits.acquire();
- try {
- return task.call();
- }
- finally {
- permits.release();
- }
- };
- }
+ void close();
}
}
diff --git
a/server/src/main/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManager.java
b/server/src/main/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManager.java
index c9c22591cff..48243b607e5 100644
---
a/server/src/main/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManager.java
+++
b/server/src/main/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManager.java
@@ -100,6 +100,14 @@ public class StorageLocationVirtualStorageManager
implements VirtualStorageManag
FilePopulator populator
)
{
+ // Hold a load permit only around the actual populate (the deep-storage
read), not the reservation/hold or the
+ // per-identifier population lock below (the permit is acquired inside
populate, which runs on mount).
+ final FilePopulator permittedPopulator = file -> {
+ try (StorageLoadingThreadPool.LoadPermit ignored =
loadingThreadPool.acquireLoadPermit()) {
+ populator.populate(file);
+ }
+ };
+
// Get or create lock for this identifier
final PopulationLock lock = populationLocks.computeIfAbsent(identifier,
ignored -> new PopulationLock());
@@ -125,7 +133,7 @@ public class StorageLocationVirtualStorageManager
implements VirtualStorageManag
// Reserve space and acquire a hold, using a cache entry that will
call the populator on mount.
final StorageLocation.ReservationHold<CacheEntry> hold =
location.addWeakReservationHold(
cacheId,
- () -> new DownloadableCacheEntry(cacheId, sizeBytes,
populator, locationFile)
+ () -> new DownloadableCacheEntry(cacheId, sizeBytes,
permittedPopulator, locationFile)
{
final AtomicBoolean mounted = new AtomicBoolean(false);
diff --git
a/server/src/test/java/org/apache/druid/segment/loading/PermitBoundedListeningExecutorServiceTest.java
b/server/src/test/java/org/apache/druid/segment/loading/PermitBoundedListeningExecutorServiceTest.java
deleted file mode 100644
index bf338e2017f..00000000000
---
a/server/src/test/java/org/apache/druid/segment/loading/PermitBoundedListeningExecutorServiceTest.java
+++ /dev/null
@@ -1,106 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-package org.apache.druid.segment.loading;
-
-import com.google.common.util.concurrent.ListenableFuture;
-import com.google.common.util.concurrent.ListeningExecutorService;
-import com.google.common.util.concurrent.MoreExecutors;
-import org.apache.druid.java.util.common.concurrent.Execs;
-import
org.apache.druid.segment.loading.StorageLoadingThreadPool.PermitBoundedListeningExecutorService;
-import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.Timeout;
-
-import java.util.concurrent.CancellationException;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.Semaphore;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
-
-/**
- * Tests the cancellation behavior of {@link
PermitBoundedListeningExecutorService}: a task parked on the permit (but
- * not yet running its body) must be abortable via {@code cancel(true)}. This
is the "mid-tier" download cancellation
- * that lets the on-demand load path stop queued/permit-blocked column
downloads before they touch deep storage.
- */
-class PermitBoundedListeningExecutorServiceTest
-{
- private ListeningExecutorService backing;
-
- @AfterEach
- void tearDown()
- {
- if (backing != null) {
- backing.shutdownNow();
- }
- }
-
- @Test
- @Timeout(30)
- void testCancelInterruptsTaskWaitingOnPermitBeforeItRuns() throws Exception
- {
- backing = MoreExecutors.listeningDecorator(Execs.multiThreaded(2,
"permit-bounded-test-%d"));
- final Semaphore permits = new Semaphore(1);
- final PermitBoundedListeningExecutorService exec =
- new PermitBoundedListeningExecutorService(backing, permits);
-
- final CountDownLatch holderRunning = new CountDownLatch(1);
- final CountDownLatch releaseHolder = new CountDownLatch(1);
-
- // Task A takes the only permit and parks in its body, so any other task
must wait on the permit to start.
- final ListenableFuture<?> holder = exec.submit(() -> {
- holderRunning.countDown();
- try {
- releaseHolder.await();
- }
- catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- }
- });
- Assertions.assertTrue(holderRunning.await(10, TimeUnit.SECONDS), "holder
task should start and take the permit");
-
- // Task B starts on another thread but blocks acquiring the permit (0
available) before its body can run.
- final AtomicBoolean bodyRan = new AtomicBoolean(false);
- final ListenableFuture<?> blocked = exec.submit(() -> {
- bodyRan.set(true);
- return null;
- });
-
- // Wait until B is genuinely parked on the semaphore, not merely
submitted. Guarded by @Timeout against a hang.
- for (int i = 0; i < 2000 && !permits.hasQueuedThreads(); i++) {
- Thread.sleep(5);
- }
- Assertions.assertTrue(permits.hasQueuedThreads(), "task B should be parked
waiting on the permit");
- Assertions.assertEquals(0, permits.availablePermits(), "the holder still
owns the only permit");
-
- // Cancelling with interruption must abort the permit wait before B's body
runs.
- Assertions.assertTrue(blocked.cancel(true), "cancel(true) should interrupt
the permit-blocked task");
- Assertions.assertTrue(blocked.isCancelled());
- Assertions.assertThrows(CancellationException.class, blocked::get);
- Assertions.assertFalse(bodyRan.get(), "permit-blocked task body must not
run after cancellation");
-
- // An interrupted acquire takes no permit, so cancelling B must not have
consumed or leaked one: the holder still
- // owns the only permit, and releasing the holder returns it.
- Assertions.assertEquals(0, permits.availablePermits());
- releaseHolder.countDown();
- holder.get(10, TimeUnit.SECONDS);
- Assertions.assertEquals(1, permits.availablePermits(), "the holder must
return its permit on completion");
- }
-}
diff --git
a/server/src/test/java/org/apache/druid/segment/loading/PermitLimitedSegmentRangeReaderTest.java
b/server/src/test/java/org/apache/druid/segment/loading/PermitLimitedSegmentRangeReaderTest.java
new file mode 100644
index 00000000000..193d8474b1f
--- /dev/null
+++
b/server/src/test/java/org/apache/druid/segment/loading/PermitLimitedSegmentRangeReaderTest.java
@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.segment.loading;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+
+class PermitLimitedSegmentRangeReaderTest
+{
+ private static StorageLoadingThreadPool oneThreadPool()
+ {
+ return StorageLoadingThreadPool.createFromConfig(
+ new SegmentLoaderConfig()
+ {
+ @Override
+ public int getVirtualStorageLoadThreads()
+ {
+ return 1;
+ }
+ }.setVirtualStorage(true)
+ );
+ }
+
+ @Test
+ @Timeout(30)
+ void testPermitIsReleasedWhenTheReturnedStreamIsClosed() throws IOException
+ {
+ final StorageLoadingThreadPool pool = oneThreadPool();
+ try {
+ final SegmentRangeReader delegate = (filename, offset, length) -> new
ByteArrayInputStream(new byte[0]);
+ final PermitLimitedSegmentRangeReader reader = new
PermitLimitedSegmentRangeReader(delegate, pool);
+
+ // Acquires the single permit and releases it on close. The second read
would hang under @Timeout if the permit
+ // were held for longer than the stream's lifetime.
+ reader.readRange("file", 0, 1).close();
+ reader.readRange("file", 0, 1).close();
+ }
+ finally {
+ pool.stop();
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ void testPermitIsReleasedWhenTheDelegateReadFails() throws IOException
+ {
+ final StorageLoadingThreadPool pool = oneThreadPool();
+ try {
+ final boolean[] fail = {true};
+ final SegmentRangeReader delegate = (filename, offset, length) -> {
+ if (fail[0]) {
+ throw new IOException("boom");
+ }
+ return new ByteArrayInputStream(new byte[0]);
+ };
+ final PermitLimitedSegmentRangeReader reader = new
PermitLimitedSegmentRangeReader(delegate, pool);
+
+ // A failed read must release the permit rather than leak it, so a later
read still acquires it.
+ Assertions.assertThrows(IOException.class, () ->
reader.readRange("file", 0, 1));
+ fail[0] = false;
+ reader.readRange("file", 0, 1).close();
+ }
+ finally {
+ pool.stop();
+ }
+ }
+}
diff --git
a/server/src/test/java/org/apache/druid/segment/loading/StorageLoadingThreadPoolTest.java
b/server/src/test/java/org/apache/druid/segment/loading/StorageLoadingThreadPoolTest.java
index 78c5fde240d..196b1d96437 100644
---
a/server/src/test/java/org/apache/druid/segment/loading/StorageLoadingThreadPoolTest.java
+++
b/server/src/test/java/org/apache/druid/segment/loading/StorageLoadingThreadPoolTest.java
@@ -22,9 +22,40 @@ package org.apache.druid.segment.loading;
import org.apache.druid.error.DruidException;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
class StorageLoadingThreadPoolTest
{
+ private static SegmentLoaderConfig oneVirtualThreadConfig()
+ {
+ return new SegmentLoaderConfig()
+ {
+ @Override
+ public int getVirtualStorageLoadThreads()
+ {
+ return 1;
+ }
+ }.setVirtualStorage(true);
+ }
+
+ private static SegmentLoaderConfig fixedThreadConfig()
+ {
+ return new SegmentLoaderConfig()
+ {
+ @Override
+ public int getVirtualStorageLoadThreads()
+ {
+ return 2;
+ }
+
+ @Override
+ public boolean isVirtualStorageUseVirtualThreads()
+ {
+ return false;
+ }
+ }.setVirtualStorage(true);
+ }
+
@Test
void testCreateFromConfigIsUnavailableWhenNotVirtualStorage()
{
@@ -59,4 +90,58 @@ class StorageLoadingThreadPoolTest
}.setVirtualStorage(true);
Assertions.assertThrows(DruidException.class, () ->
StorageLoadingThreadPool.createFromConfig(config));
}
+
+ @Test
+ @Timeout(30)
+ void testAcquireLoadPermitReleasesOnCloseForReuse()
+ {
+ final StorageLoadingThreadPool pool =
StorageLoadingThreadPool.createFromConfig(oneVirtualThreadConfig());
+ try {
+ final StorageLoadingThreadPool.LoadPermit first =
pool.acquireLoadPermit();
+ first.close();
+ first.close(); // idempotent: a double close must not over-release the
single permit
+ // The permit is available again, so this returns immediately (it would
hang under @Timeout if close() had not
+ // released it, or the extra release above had inflated the count).
+ pool.acquireLoadPermit().close();
+ }
+ finally {
+ pool.stop();
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ void testAcquireLoadPermitIsNoOpWithoutSemaphore()
+ {
+ // Fixed-thread mode has no semaphore (the thread count is the bound), so
acquiring repeatedly without releasing
+ // must never block.
+ final StorageLoadingThreadPool pool =
StorageLoadingThreadPool.createFromConfig(fixedThreadConfig());
+ try {
+ pool.acquireLoadPermit();
+ pool.acquireLoadPermit();
+ pool.acquireLoadPermit().close();
+ }
+ finally {
+ pool.stop();
+ }
+ }
+
+ @Test
+ @Timeout(30)
+ void testAcquireLoadPermitIsInterruptibleAndDoesNotConsumeAPermit()
+ {
+ // A load interrupted while acquiring must abort (throw) and take no
permit - this is the cancel-before-I/O
+ // guarantee that lets a cancelled query stop permit-blocked downloads
before they touch deep storage.
+ final StorageLoadingThreadPool pool =
StorageLoadingThreadPool.createFromConfig(oneVirtualThreadConfig());
+ try {
+ Thread.currentThread().interrupt();
+ Assertions.assertThrows(RuntimeException.class, pool::acquireLoadPermit);
+ Assertions.assertTrue(Thread.interrupted(), "interrupt flag should be
restored (and is cleared here)");
+ // The aborted acquire consumed no permit, so this succeeds (would hang
under @Timeout had it leaked one).
+ pool.acquireLoadPermit().close();
+ }
+ finally {
+ pool.stop();
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
