You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@druid.apache.org by gi...@apache.org on 2026/06/03 21:09:54 UTC
(druid) branch master updated: feat: Background fetching of files for SQL EXTERN. (#19539)
This is an automated email from the ASF dual-hosted git repository.
gianm 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 03ab4e2b1e3 feat: Background fetching of files for SQL EXTERN. (#19539)
03ab4e2b1e3 is described below
commit 03ab4e2b1e345d5ee2dc4df78a9a3232b0aa8927
Author: Gian Merlino <gi...@gmail.com>
AuthorDate: Wed Jun 3 14:09:47 2026 -0700
feat: Background fetching of files for SQL EXTERN. (#19539)
This patch adds the setting "backgroundFetchExternalFiles". When set,
cloud storage files referenced by ExternalSegment (EXTERN) are fetched
asynchronously into the task's storage locations. The setting defaults
to true.
To support this, new infrastructure is added:
1) VirtualStorageManager, a layer on top of StorageLocation that
provides a simple "fetch and cache a file" API.
2) StorageLoadingThreadPool, an extraction of the thread pool from
SegmentLocalCacheManager so it can be shared with
VirtualStorageManager.
3) AsyncResource, a Future-like utility that provides better tools for
managing Closeable resources. It is used by VirtualStorageManager to
provide the asynchronously-fetched file handles.
---
docs/multi-stage-query/reference.md | 1 +
.../testing/embedded/msq/EmbeddedMSQApis.java | 94 +-
.../embedded/msq/EmbeddedMSQRealtimeQueryTest.java | 30 +-
.../testing/embedded/msq/S3ExternQueryTest.java | 215 +++++
.../embedded/query/QueryVirtualStorageTest.java | 47 +-
.../common/SegmentCacheManagerFactory.java | 26 +-
.../druid/indexing/common/TaskToolboxFactory.java | 2 +-
.../druid/indexing/input/DruidInputSource.java | 3 +-
.../indexing/input/DruidSegmentInputEntity.java | 2 +-
.../seekablestream/SettableByteEntity.java | 2 +-
.../common/task/CompactionTaskRunBase.java | 10 +-
.../druid/indexing/common/task/IndexTaskTest.java | 2 +
.../indexing/common/task/IngestionTestBase.java | 2 +-
.../AbstractMultiPhaseParallelIndexingTest.java | 2 +-
.../apache/druid/msq/counters/ChannelCounters.java | 13 +-
.../druid/msq/dart/worker/DartFrameContext.java | 10 +
.../druid/msq/dart/worker/DartWorkerContext.java | 5 +
.../dart/worker/DartWorkerContextFactoryImpl.java | 9 +-
.../msq/exec/DataServerQueryHandlerUtils.java | 12 +-
.../org/apache/druid/msq/exec/FrameContext.java | 6 +
.../org/apache/druid/msq/exec/RunWorkOrder.java | 9 +-
.../druid/msq/indexing/IndexerFrameContext.java | 10 +
.../druid/msq/indexing/IndexerWorkerContext.java | 24 +-
.../druid/msq/input/AdaptedLoadableSegment.java | 115 ++-
.../apache/druid/msq/input/LoadableSegment.java | 8 -
.../druid/msq/input/LoadableSegmentUtils.java | 83 ++
.../druid/msq/input/RegularLoadableSegment.java | 32 +-
.../input/external/ExternalInputSliceReader.java | 175 +++-
.../druid/msq/input/external/ExternalSegment.java | 52 +-
.../msq/input/inline/InlineInputSliceReader.java | 5 +-
.../msq/input/lookup/LookupInputSliceReader.java | 5 +-
.../druid/msq/querykit/BaseLeafFrameProcessor.java | 29 -
.../druid/msq/querykit/ReadableInputQueue.java | 7 -
.../druid/msq/querykit/SegmentReferenceHolder.java | 14 -
.../druid/msq/util/MultiStageQueryContext.java | 8 +
.../msq/dart/worker/DartFrameContextTest.java | 1 +
.../druid/msq/exec/MSQCompactionTaskRunTest.java | 2 +
.../msq/indexing/IndexerFrameContextTest.java | 1 +
.../msq/input/RegularLoadableSegmentTest.java | 5 +-
.../querykit/scan/ScanQueryFrameProcessorTest.java | 2 -
.../druid/msq/sql/MSQTaskQueryMakerTest.java | 2 +
.../test/AbstractMSQComponentSupplierDelegate.java | 7 +-
.../druid/msq/test/CalciteMSQTestsHelper.java | 2 +-
.../org/apache/druid/msq/test/MSQTestBase.java | 27 +-
.../druid/msq/test/MSQTestWorkerContext.java | 7 +
.../druid/common/asyncresource/AsyncResource.java | 139 +++
.../druid/common/asyncresource/AsyncResources.java | 114 +++
.../common/asyncresource/CollectAsyncResource.java | 98 ++
.../common/asyncresource/RecoverAsyncResource.java | 124 +++
.../asyncresource/SettableAsyncResource.java | 335 +++++++
.../asyncresource/TransformAsyncResource.java | 86 ++
.../druid/data/input/BytesCountingInputEntity.java | 6 +
.../org/apache/druid/data/input/InputEntity.java | 13 +-
.../apache/druid/data/input/InputFilePointer.java | 34 +-
.../org/apache/druid/data/input/InputSource.java | 10 +
.../druid/data/input/RetryingInputEntity.java | 12 +-
.../apache/druid/data/input/impl/ByteEntity.java | 2 +-
.../data/input/impl/CloudObjectInputSource.java | 63 ++
.../apache/druid/data/input/impl/FileEntity.java | 14 +-
.../druid/data/input/impl/LocalInputSource.java | 10 +
.../java/org/apache/druid/io/FilePopulator.java | 36 +-
.../apache/druid/segment/AsyncCursorHolder.java | 254 +----
.../incremental/NoopRowIngestionMeters.java | 5 +-
.../asyncresource/CollectAsyncResourceTest.java | 131 +++
.../asyncresource/RecoverAsyncResourceTest.java | 184 ++++
.../asyncresource/SettableAsyncResourceTest.java | 276 ++++++
.../asyncresource/TransformAsyncResourceTest.java | 156 +++
.../druid/data/input/ResourceInputSource.java | 9 +-
.../input/impl/CloudObjectInputSourceTest.java | 129 ++-
.../druid/data/input/impl/FileEntityTest.java | 125 +++
.../query/aggregation/AggregationTestHelper.java | 2 +-
.../druid/guice/LocalDataStorageDruidModule.java | 4 -
.../org/apache/druid/guice/StorageNodeModule.java | 14 +
.../org/apache/druid/metadata/input/SqlEntity.java | 2 +-
.../druid/segment/loading/SegmentCacheManager.java | 5 +
.../druid/segment/loading/SegmentLoaderConfig.java | 31 +-
.../segment/loading/SegmentLocalCacheManager.java | 82 +-
.../segment/loading/StorageLoadingThreadPool.java | 143 +++
.../druid/segment/loading/external/CachedFile.java | 87 ++
.../loading/external/DownloadableCacheEntry.java | 314 ++++++
.../StorageLocationVirtualStorageManager.java | 299 ++++++
.../external/StringCacheEntryIdentifier.java} | 27 +-
.../loading/external/VirtualStorageManager.java | 97 ++
.../segment/loading/NoopSegmentCacheManager.java | 8 +-
.../segment/loading/SegmentLoaderConfigTest.java | 2 +-
.../SegmentLocalCacheManagerConcurrencyTest.java | 2 +
.../loading/SegmentLocalCacheManagerTest.java | 39 +-
.../external/DownloadableCacheEntryTest.java | 225 +++++
.../StorageLocationVirtualStorageManagerTest.java | 1013 ++++++++++++++++++++
.../apache/druid/server/SegmentManagerTest.java | 3 +
.../SegmentCacheBootstrapperCacheTest.java | 3 +
...egmentManagerBroadcastJoinIndexedTableTest.java | 2 +
.../SegmentManagerThreadSafetyTest.java | 2 +
93 files changed, 5261 insertions(+), 635 deletions(-)
diff --git a/docs/multi-stage-query/reference.md b/docs/multi-stage-query/reference.md
index 6fb63924ee7..82d85b52eb7 100644
--- a/docs/multi-stage-query/reference.md
+++ b/docs/multi-stage-query/reference.md
@@ -417,6 +417,7 @@ The following table lists the context parameters for the MSQ task engine:
| `includeAllCounters` | SELECT, INSERT or REPLACE<br /><br />Whether to include counters that were added in Druid 31 or later. This is a backwards compatibility option that must be set to `false` during a rolling update from versions prior to Druid 31. | `true` |
| `maxFrameSize` | SELECT, INSERT or REPLACE<br /><br />Size of frames used for data transfer within the MSQ engine. You generally do not need to change this unless you have very large rows. | `1000000` (1 MB) |
| `maxInputFilesPerWorker` | SELECT, INSERT, REPLACE<br /><br />Maximum number of input files or segments per worker. If a single worker would need to read more than this number of files, the query fails with a `TooManyInputFiles` error. In this case, you should either increase this limit if your tasks have enough memory to handle more files, add more workers by increasing `maxNumTasks`, or split your query into smaller queries that process fewer files. | 10,000 |
+| `backgroundFetchExternalFiles` | SELECT, INSERT, REPLACE<br /><br />Controls how workers read cloud storage files (e.g. `s3`, `gs`, `azure`) referenced by [`EXTERN`](#extern-function). When `true` (the default), each worker fetches its input files asynchronously when possible before reading them. This overlaps downloading with processing, which generally improves throughput when reading larger numbers of files. Otherwise, workers stream each file directly from cloud storage while proce [...]
| `maxPartitions` | SELECT, INSERT, REPLACE<br /><br />Maximum number of output partitions for any single stage. For INSERT or REPLACE queries, this controls the maximum number of segments that can be generated. If the query would exceed this limit, it fails with a `TooManyPartitions` error. You can increase this limit if needed, break your query into smaller queries, or use a larger target segment size (via `rowsPerSegment`). | 25,000 |
| `maxThreads` | SELECT, INSERT or REPLACE<br /><br />Maximum number of threads to use for processing. This only has an effect if it is greater than zero and less than the default thread count based on system configuration. Otherwise, it is ignored, and workers use the default thread count. | Not set (use default thread count) |
diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQApis.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQApis.java
index e58ad81b465..8154ac988a3 100644
--- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQApis.java
+++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQApis.java
@@ -243,44 +243,61 @@ public class EmbeddedMSQApis
}
/**
- * Returns the sum of completed queries across all input channel snapshots from all stages and workers.
+ * Returns the sums of all input channel counters across all workers.
*/
- public long getQueriesSum(final MSQTaskReportPayload payload)
+ public ChannelSums getInputChannelSums(final MSQTaskReportPayload payload, final int stageNumber)
{
- return getAllInputChannelCounters(payload)
- .stream()
- .filter(snapshot -> snapshot.getQueries() != null)
- .mapToLong(snapshot -> Arrays.stream(snapshot.getQueries()).sum())
- .sum();
+ long rows = 0;
+ long bytes = 0;
+ long files = 0;
+ long totalFiles = 0;
+ long queries = 0;
+ long totalQueries = 0;
+ long loadBytes = 0;
+ long loadFiles = 0;
+ long loadTime = 0;
+ long loadWait = 0;
+
+ for (final ChannelCounters.Snapshot snapshot : getAllInputChannelCounters(payload, stageNumber)) {
+ rows += sum(snapshot.getRows());
+ bytes += sum(snapshot.getBytes());
+ files += sum(snapshot.getFiles());
+ totalFiles += sum(snapshot.getTotalFiles());
+ queries += sum(snapshot.getQueries());
+ totalQueries += sum(snapshot.getTotalQueries());
+ loadBytes += sum(snapshot.getLoadBytes());
+ loadFiles += sum(snapshot.getLoadFiles());
+ loadTime += sum(snapshot.getLoadTime());
+ loadWait += sum(snapshot.getLoadWait());
+ }
+
+ return new ChannelSums(rows, bytes, files, totalFiles, queries, totalQueries, loadBytes, loadFiles, loadTime, loadWait);
}
/**
- * Returns the sum of files read across all input channel snapshots from all stages and workers.
+ * Sums the values of a nullable channel counter array, treating {@code null} as empty.
*/
- public long getFilesSum(final MSQTaskReportPayload payload)
+ private static long sum(@Nullable final long[] values)
{
- return getAllInputChannelCounters(payload)
- .stream()
- .filter(snapshot -> snapshot.getFiles() != null)
- .mapToLong(snapshot -> Arrays.stream(snapshot.getFiles()).sum())
- .sum();
+ return values == null ? 0 : Arrays.stream(values).sum();
}
/**
- * Returns all {@link ChannelCounters.Snapshot} from input channels across all stages and workers.
+ * Returns all {@link ChannelCounters.Snapshot} from input channels across all workers for a stage.
*/
- private List<ChannelCounters.Snapshot> getAllInputChannelCounters(final MSQTaskReportPayload payload)
+ private List<ChannelCounters.Snapshot> getAllInputChannelCounters(
+ final MSQTaskReportPayload payload,
+ final int stageNumber
+ )
{
- final Map<Integer, Map<Integer, CounterSnapshots>> countersMap = payload.getCounters().copyMap();
final List<ChannelCounters.Snapshot> snapshots = new ArrayList<>();
+ final Map<Integer, CounterSnapshots> stageMap = payload.getCounters().snapshotForStage(stageNumber);
- for (final Map.Entry<Integer, Map<Integer, CounterSnapshots>> stageEntry : countersMap.entrySet()) {
- for (final Map.Entry<Integer, CounterSnapshots> workerEntry : stageEntry.getValue().entrySet()) {
- for (final Map.Entry<String, QueryCounterSnapshot> counterEntry : workerEntry.getValue().getMap().entrySet()) {
- if (counterEntry.getKey().startsWith("input")
- && counterEntry.getValue() instanceof ChannelCounters.Snapshot) {
- snapshots.add((ChannelCounters.Snapshot) counterEntry.getValue());
- }
+ for (final Map.Entry<Integer, CounterSnapshots> workerEntry : stageMap.entrySet()) {
+ for (final Map.Entry<String, QueryCounterSnapshot> counterEntry : workerEntry.getValue().getMap().entrySet()) {
+ if (counterEntry.getKey().startsWith("input")
+ && counterEntry.getValue() instanceof ChannelCounters.Snapshot counterSnapshot) {
+ snapshots.add(counterSnapshot);
}
}
}
@@ -297,4 +314,33 @@ public class EmbeddedMSQApis
throw DruidException.defensive(e, "Failed to parse query report response[%s]", responseJson);
}
}
+
+ /**
+ * Sums of input channel counters computed by {@link #getInputChannelSums(MSQTaskReportPayload, int)}.
+ *
+ * @param rows total rows read
+ * @param bytes total bytes read
+ * @param files total files read
+ * @param totalFiles total number of files to read
+ * @param queries total queries completed
+ * @param totalQueries total number of queries to run
+ * @param loadBytes total bytes loaded into the virtual storage file cache (VSF)
+ * @param loadFiles total files loaded into the VSF
+ * @param loadTime total time (in milliseconds) spent loading files into the VSF
+ * @param loadWait total time (in milliseconds) spent waiting to load files into the VSF
+ */
+ public record ChannelSums(
+ long rows,
+ long bytes,
+ long files,
+ long totalFiles,
+ long queries,
+ long totalQueries,
+ long loadBytes,
+ long loadFiles,
+ long loadTime,
+ long loadWait
+ )
+ {
+ }
}
diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeQueryTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeQueryTest.java
index a651ccf03ad..2d3e53b2a03 100644
--- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeQueryTest.java
+++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQRealtimeQueryTest.java
@@ -38,7 +38,6 @@ import org.apache.druid.testing.embedded.EmbeddedOverlord;
import org.apache.druid.testing.embedded.EmbeddedRouter;
import org.hamcrest.CoreMatchers;
import org.hamcrest.MatcherAssert;
-import org.hamcrest.Matchers;
import org.junit.internal.matchers.ThrowableMessageMatcher;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
@@ -189,9 +188,18 @@ public class EmbeddedMSQRealtimeQueryTest extends BaseRealtimeQueryTest
payload.getResults().getResults()
);
- // Verify that realtime queries were issued and no files were read (all data is realtime).
- MatcherAssert.assertThat(msqApis.getQueriesSum(payload), Matchers.greaterThan(0L));
- Assertions.assertEquals(0, msqApis.getFilesSum(payload));
+ // Verify that for the first stage, realtime queries were issued and no files were read (all data is realtime).
+ final EmbeddedMSQApis.ChannelSums channelSums = msqApis.getInputChannelSums(payload, 0);
+
+ // 2 realtime tasks running on an Indexer => 1 query that gets data from both.
+ Assertions.assertEquals(1, channelSums.queries());
+ Assertions.assertEquals(1, channelSums.totalQueries());
+ Assertions.assertEquals(0, channelSums.files());
+ Assertions.assertEquals(0, channelSums.totalFiles());
+
+ // We get 1 row back with the COUNT from both tasks.
+ Assertions.assertEquals(1, channelSums.rows());
+ Assertions.assertEquals(0, channelSums.bytes()); // Realtime queries do not report bytes
}
@Test
@@ -432,6 +440,19 @@ public class EmbeddedMSQRealtimeQueryTest extends BaseRealtimeQueryTest
),
payload.getResults().getResults()
);
+
+ // Verify that for the first stage, realtime queries were issued and no files were read (all data is realtime).
+ final EmbeddedMSQApis.ChannelSums channelSums = msqApis.getInputChannelSums(payload, 0);
+
+ // 2 realtime tasks running on an Indexer => 1 query that gets data from both.
+ Assertions.assertEquals(1, channelSums.queries());
+ Assertions.assertEquals(1, channelSums.totalQueries());
+ Assertions.assertEquals(0, channelSums.files());
+ Assertions.assertEquals(0, channelSums.totalFiles());
+
+ // We get 3 rows back.
+ Assertions.assertEquals(3, channelSums.rows());
+ Assertions.assertEquals(0, channelSums.bytes()); // Realtime queries do not report bytes
}
@Test
@@ -488,7 +509,6 @@ public class EmbeddedMSQRealtimeQueryTest extends BaseRealtimeQueryTest
);
}
-
@Test
@Timeout(60)
public void test_scanWithFilter_dart()
diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/S3ExternQueryTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/S3ExternQueryTest.java
new file mode 100644
index 00000000000..8eb06e76a39
--- /dev/null
+++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/S3ExternQueryTest.java
@@ -0,0 +1,215 @@
+/*
+ * 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.testing.embedded.msq;
+
+import org.apache.druid.data.input.s3.S3InputSourceDruidModule;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.msq.indexing.report.MSQTaskReportPayload;
+import org.apache.druid.sql.calcite.BaseCalciteQueryTest;
+import org.apache.druid.testing.embedded.EmbeddedBroker;
+import org.apache.druid.testing.embedded.EmbeddedCoordinator;
+import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
+import org.apache.druid.testing.embedded.EmbeddedIndexer;
+import org.apache.druid.testing.embedded.EmbeddedOverlord;
+import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase;
+import org.apache.druid.testing.embedded.minio.MinIOStorageResource;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.core.sync.RequestBody;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.zip.GZIPOutputStream;
+
+/**
+ * Tests MSQ task-based SELECT queries that read external data via {@code EXTERN} from S3 (backed by MinIO).
+ */
+public class S3ExternQueryTest extends EmbeddedClusterTestBase
+{
+ /**
+ * Base key under the bucket where the external input files are uploaded.
+ */
+ private static final String DATA_PATH = "extern-input";
+ private static final String PLAIN_FILE = "data1.json";
+ private static final String GZ_FILE = "data2.json.gz";
+
+ private static final String PLAIN_FILE_CONTENT = """
+ {"timestamp":"2020-01-01T00:00:00Z","page":"A","added":10}
+ {"timestamp":"2020-01-01T01:00:00Z","page":"B","added":20}
+ """;
+
+ private static final String GZ_FILE_CONTENT = """
+ {"timestamp":"2020-01-02T00:00:00Z","page":"C","added":30}
+ {"timestamp":"2020-01-02T01:00:00Z","page":"D","added":40}
+ {"timestamp":"2020-01-02T02:00:00Z","page":"E","added":50}
+ """;
+
+ /**
+ * Total number of bytes uploaded across {@link #PLAIN_FILE} and {@link #GZ_FILE}. Populated by
+ * {@link #uploadExternalFiles()} and used to verify the VSF {@code loadBytes} counter.
+ */
+ private long totalUploadedBytes;
+
+ private final MinIOStorageResource storageResource = new MinIOStorageResource();
+ private final EmbeddedOverlord overlord = new EmbeddedOverlord();
+ private final EmbeddedCoordinator coordinator = new EmbeddedCoordinator();
+ private final EmbeddedIndexer indexer = new EmbeddedIndexer()
+ .setServerMemory(300_000_000L)
+ .addProperty("druid.worker.capacity", "2");
+ private final EmbeddedBroker broker = new EmbeddedBroker().setServerMemory(200_000_000);
+
+ private EmbeddedMSQApis msqApis;
+
+ @Override
+ protected EmbeddedDruidCluster createCluster()
+ {
+ return EmbeddedDruidCluster
+ .withEmbeddedDerbyAndZookeeper()
+ .useLatchableEmitter()
+ .addResource(storageResource)
+ .addExtension(S3InputSourceDruidModule.class)
+ .addServer(overlord)
+ .addServer(coordinator)
+ .addServer(indexer)
+ .addServer(broker);
+ }
+
+ @BeforeAll
+ public void setupCluster() throws IOException
+ {
+ msqApis = new EmbeddedMSQApis(cluster, overlord);
+ uploadExternalFiles();
+ }
+
+ @Test
+ public void test_extern_backgroundFetchEnabled()
+ {
+ runQueryAndVerify(true);
+ }
+
+ @Test
+ public void test_extern_backgroundFetchDisabled()
+ {
+ runQueryAndVerify(false);
+ }
+
+ private void runQueryAndVerify(final boolean backgroundFetchExternalFiles)
+ {
+ final String inputSourceJson = StringUtils.format(
+ "{\"type\":\"s3\",\"uris\":[\"s3://%s/%s/%s\",\"s3://%s/%s/%s\"]}",
+ storageResource.getBucket(), DATA_PATH, PLAIN_FILE,
+ storageResource.getBucket(), DATA_PATH, GZ_FILE
+ );
+
+ final String sql = StringUtils.format(
+ """
+ SET backgroundFetchExternalFiles = %s;
+ SELECT page, added
+ FROM TABLE(
+ EXTERN(
+ '%s',
+ '{"type":"json"}'
+ )
+ ) EXTEND ("page" VARCHAR, "added" BIGINT)
+ ORDER BY page
+ """,
+ backgroundFetchExternalFiles,
+ inputSourceJson
+ );
+
+ final MSQTaskReportPayload report = msqApis.runTaskSqlAndGetReport(sql);
+
+ BaseCalciteQueryTest.assertResultsEquals(
+ sql,
+ List.of(
+ new Object[]{"A", 10},
+ new Object[]{"B", 20},
+ new Object[]{"C", 30},
+ new Object[]{"D", 40},
+ new Object[]{"E", 50}
+ ),
+ report.getResults().getResults()
+ );
+
+ // Verify input counters.
+ final EmbeddedMSQApis.ChannelSums channelSums = msqApis.getInputChannelSums(report, 0);
+ Assertions.assertEquals(2, channelSums.files(), "files");
+ Assertions.assertEquals(2, channelSums.totalFiles(), "totalFiles");
+ Assertions.assertEquals(5, channelSums.rows(), "rows");
+ Assertions.assertEquals(PLAIN_FILE_CONTENT.length() + GZ_FILE_CONTENT.length(), channelSums.bytes(), "bytes");
+ Assertions.assertEquals(0, channelSums.queries(), "queries");
+ Assertions.assertEquals(0, channelSums.queries(), "totalQueries");
+
+ // Verify load counters.
+ if (backgroundFetchExternalFiles) {
+ Assertions.assertEquals(totalUploadedBytes, channelSums.loadBytes(), "VSF bytes loaded");
+ Assertions.assertEquals(2, channelSums.loadFiles(), "VSF files loaded");
+ Assertions.assertTrue(channelSums.loadTime() >= 0, "VSF load time");
+ Assertions.assertEquals(0, channelSums.loadWait(), "VSF load wait time");
+ } else {
+ Assertions.assertEquals(0, channelSums.loadBytes(), "VSF bytes loaded");
+ Assertions.assertEquals(0, channelSums.loadFiles(), "VSF files loaded");
+ Assertions.assertEquals(0, channelSums.loadTime(), "VSF load time");
+ Assertions.assertEquals(0, channelSums.loadWait(), "VSF load wait time");
+ }
+ }
+
+ /**
+ * Uploads {@link #PLAIN_FILE} (plain JSON) and {@link #GZ_FILE} (gzipped JSON) to the MinIO bucket.
+ */
+ private void uploadExternalFiles() throws IOException
+ {
+ final S3Client s3Client = storageResource.getS3Client();
+
+ final byte[] plainBytes = PLAIN_FILE_CONTENT.getBytes(StandardCharsets.UTF_8);
+ final byte[] gzBytes = gzip(GZ_FILE_CONTENT);
+ totalUploadedBytes += plainBytes.length + gzBytes.length;
+
+ s3Client.putObject(
+ PutObjectRequest.builder()
+ .bucket(storageResource.getBucket())
+ .key(DATA_PATH + "/" + PLAIN_FILE)
+ .build(),
+ RequestBody.fromBytes(plainBytes)
+ );
+
+ s3Client.putObject(
+ PutObjectRequest.builder()
+ .bucket(storageResource.getBucket())
+ .key(DATA_PATH + "/" + GZ_FILE)
+ .build(),
+ RequestBody.fromBytes(gzBytes)
+ );
+ }
+
+ private static byte[] gzip(final String content) throws IOException
+ {
+ final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ try (final GZIPOutputStream gzipStream = new GZIPOutputStream(baos)) {
+ gzipStream.write(StringUtils.toUtf8(content));
+ }
+ return baos.toByteArray();
+ }
+}
diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryVirtualStorageTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryVirtualStorageTest.java
index df55055751d..0e5a2a2baaf 100644
--- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryVirtualStorageTest.java
+++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/QueryVirtualStorageTest.java
@@ -25,7 +25,7 @@ import org.apache.druid.indexer.TaskState;
import org.apache.druid.java.util.common.HumanReadableBytes;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
-import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.msq.indexing.report.MSQTaskReport;
import org.apache.druid.msq.indexing.report.MSQTaskReportPayload;
import org.apache.druid.query.DefaultQueryMetrics;
import org.apache.druid.query.DruidProcessingConfigTest;
@@ -47,6 +47,8 @@ import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase;
import org.apache.druid.testing.embedded.minio.MinIOStorageResource;
import org.apache.druid.testing.embedded.msq.EmbeddedDurableShuffleStorageTest;
import org.apache.druid.testing.embedded.msq.EmbeddedMSQApis;
+import org.hamcrest.MatcherAssert;
+import org.hamcrest.Matchers;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
@@ -278,29 +280,28 @@ class QueryVirtualStorageTest extends EmbeddedClusterTestBase
// Now fetch the report using the SQL query ID
final GetQueryReportResponse reportResponse = msqApis.getDartQueryReport(sqlQueryId, broker);
-
- // Verify the report response
Assertions.assertNotNull(reportResponse, "Report response should not be null");
- ChannelCounters.Snapshot segmentChannelCounters =
- (ChannelCounters.Snapshot) reportResponse.getReportMap()
- .findReport("multiStageQuery")
- .map(r ->
- ((MSQTaskReportPayload) r.getPayload()).getCounters()
- .snapshotForStage(0)
- .get(0)
- .getMap()
- .get("input0")
- ).orElse(null);
-
- Assertions.assertNotNull(segmentChannelCounters);
- Assertions.assertArrayEquals(new long[]{24L}, segmentChannelCounters.getFiles());
- Assertions.assertTrue(segmentChannelCounters.getLoadFiles()[0] > 0 && segmentChannelCounters.getLoadFiles()[0] <= segmentChannelCounters.getFiles()[0]);
- // size of all segments at time of writing, possibly we have to load all of them, but possibly less depending on
- // test order
- Assertions.assertTrue(segmentChannelCounters.getLoadBytes()[0] > 0);
- Assertions.assertTrue(segmentChannelCounters.getLoadBytes()[0] <= SIZE_BYTES);
- Assertions.assertTrue(segmentChannelCounters.getLoadTime()[0] > 0);
- Assertions.assertTrue(segmentChannelCounters.getLoadWait()[0] > 0);
+
+ final MSQTaskReportPayload reportPayload =
+ ((MSQTaskReport) reportResponse.getReportMap().get(MSQTaskReport.REPORT_KEY)).getPayload();
+
+ // Verify stage 0 (segment read) input counters
+ final EmbeddedMSQApis.ChannelSums inputChannelSums = msqApis.getInputChannelSums(reportPayload, 0);
+ Assertions.assertEquals(24L, inputChannelSums.files());
+ Assertions.assertEquals(24L, inputChannelSums.totalFiles());
+ Assertions.assertEquals(0L, inputChannelSums.queries());
+ Assertions.assertEquals(0L, inputChannelSums.totalQueries());
+ Assertions.assertEquals(39244L, inputChannelSums.rows());
+ MatcherAssert.assertThat(inputChannelSums.bytes(), Matchers.greaterThan(0L));
+ MatcherAssert.assertThat(inputChannelSums.bytes(), Matchers.lessThanOrEqualTo(SIZE_BYTES));
+
+ // Verify stage 0 (segment read) VSF load counters
+ MatcherAssert.assertThat(inputChannelSums.loadFiles(), Matchers.greaterThan(0L));
+ MatcherAssert.assertThat(inputChannelSums.loadFiles(), Matchers.lessThanOrEqualTo(24L));
+ MatcherAssert.assertThat(inputChannelSums.loadTime(), Matchers.greaterThan(0L));
+ MatcherAssert.assertThat(inputChannelSums.loadWait(), Matchers.greaterThan(0L));
+ MatcherAssert.assertThat(inputChannelSums.loadBytes(), Matchers.greaterThan(0L));
+ MatcherAssert.assertThat(inputChannelSums.loadBytes(), Matchers.lessThanOrEqualTo(SIZE_BYTES));
}
@Test
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/SegmentCacheManagerFactory.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/SegmentCacheManagerFactory.java
index 63ee7e8c495..cf5dfc7d0e9 100644
--- a/indexing-service/src/main/java/org/apache/druid/indexing/common/SegmentCacheManagerFactory.java
+++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/SegmentCacheManagerFactory.java
@@ -27,8 +27,10 @@ import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStr
import org.apache.druid.segment.loading.SegmentCacheManager;
import org.apache.druid.segment.loading.SegmentLoaderConfig;
import org.apache.druid.segment.loading.SegmentLocalCacheManager;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
import org.apache.druid.segment.loading.StorageLocation;
import org.apache.druid.segment.loading.StorageLocationConfig;
+import org.apache.druid.timeline.DataSegment;
import java.io.File;
import java.util.Collections;
@@ -52,16 +54,34 @@ public class SegmentCacheManagerFactory
this.jsonMapper = mapper;
}
- public SegmentCacheManager manufacturate(File storageDir, boolean virtualStorage)
+ /**
+ * Creates a new {@link SegmentCacheManager} backed by a new storage location in {@code storageDir}, and a new
+ * loading thread pool of default size.
+ *
+ * @param storageDir storage location
+ * @param maxSize size limit, or null for no limit
+ * @param virtualStorage whether to configure the cache manager in ephemeral virtual storage mode. In this mode,
+ * loading is triggered by {@link SegmentCacheManager#acquireSegment(DataSegment)}, and
+ * segment files are deleted as soon as all holds are closed.
+ */
+ public SegmentCacheManager manufacturate(File storageDir, Long maxSize, boolean virtualStorage)
{
+ final StorageLocationConfig locationConfig = new StorageLocationConfig(
+ storageDir,
+ maxSize != null ? maxSize : Long.MAX_VALUE,
+ null
+ );
final SegmentLoaderConfig loaderConfig =
new SegmentLoaderConfig()
- .setLocations(Collections.singletonList(new StorageLocationConfig(storageDir, null, null)))
- .setVirtualStorage(virtualStorage, virtualStorage);
+ .setLocations(Collections.singletonList(locationConfig))
+ .setVirtualStorage(virtualStorage)
+ .setVirtualStorageIsEphemeral(virtualStorage);
final List<StorageLocation> storageLocations = loaderConfig.toStorageLocations();
+ final StorageLoadingThreadPool loadingThreadPool = StorageLoadingThreadPool.createFromConfig(loaderConfig);
return new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ loadingThreadPool,
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
indexIO,
jsonMapper
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/common/TaskToolboxFactory.java b/indexing-service/src/main/java/org/apache/druid/indexing/common/TaskToolboxFactory.java
index 2cd5b8dd609..6ee32438d49 100644
--- a/indexing-service/src/main/java/org/apache/druid/indexing/common/TaskToolboxFactory.java
+++ b/indexing-service/src/main/java/org/apache/druid/indexing/common/TaskToolboxFactory.java
@@ -247,7 +247,7 @@ public class TaskToolboxFactory
.queryProcessingPool(queryProcessingPool)
.joinableFactory(joinableFactory)
.monitorSchedulerProvider(monitorSchedulerProvider)
- .segmentCacheManager(segmentCacheManagerFactory.manufacturate(taskWorkDir, true))
+ .segmentCacheManager(segmentCacheManagerFactory.manufacturate(taskWorkDir, null, true))
.jsonMapper(jsonMapper)
.taskWorkDir(taskWorkDir)
.indexIO(indexIO)
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/input/DruidInputSource.java b/indexing-service/src/main/java/org/apache/druid/indexing/input/DruidInputSource.java
index ab7ea72730a..84fead9916b 100644
--- a/indexing-service/src/main/java/org/apache/druid/indexing/input/DruidInputSource.java
+++ b/indexing-service/src/main/java/org/apache/druid/indexing/input/DruidInputSource.java
@@ -313,7 +313,8 @@ public class DruidInputSource extends AbstractInputSource implements SplittableI
@Override
protected InputSourceReader fixedFormatReader(InputRowSchema inputRowSchema, @Nullable File temporaryDirectory)
{
- final SegmentCacheManager segmentCacheManager = segmentCacheManagerFactory.manufacturate(temporaryDirectory, false);
+ final SegmentCacheManager segmentCacheManager =
+ segmentCacheManagerFactory.manufacturate(temporaryDirectory, null, false);
final List<TimelineObjectHolder<String, DataSegment>> timeline = createTimeline();
final Iterator<DruidSegmentInputEntity> entityIterator = FluentIterable
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/input/DruidSegmentInputEntity.java b/indexing-service/src/main/java/org/apache/druid/indexing/input/DruidSegmentInputEntity.java
index 0f7f5a6b941..85ab75c3787 100644
--- a/indexing-service/src/main/java/org/apache/druid/indexing/input/DruidSegmentInputEntity.java
+++ b/indexing-service/src/main/java/org/apache/druid/indexing/input/DruidSegmentInputEntity.java
@@ -59,7 +59,7 @@ public class DruidSegmentInputEntity implements InputEntity
}
@Override
- public InputStream open()
+ public InputStream openRaw()
{
throw new UnsupportedOperationException("Don't call this");
}
diff --git a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SettableByteEntity.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SettableByteEntity.java
index b4eb42a6b48..3faf436f59c 100644
--- a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SettableByteEntity.java
+++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SettableByteEntity.java
@@ -62,7 +62,7 @@ public class SettableByteEntity<T extends ByteEntity> implements InputEntity
}
@Override
- public InputStream open()
+ public InputStream openRaw()
{
// Duplicate the entity buffer, because the stream will update its position.
final SettableByteBufferInputStream stream = new SettableByteBufferInputStream();
diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskRunBase.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskRunBase.java
index 2269da4e331..99d478229a8 100644
--- a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskRunBase.java
+++ b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskRunBase.java
@@ -101,6 +101,7 @@ import org.apache.druid.segment.loading.NoopDataSegmentKiller;
import org.apache.druid.segment.loading.SegmentCacheManager;
import org.apache.druid.segment.loading.SegmentLoaderConfig;
import org.apache.druid.segment.loading.SegmentLocalCacheManager;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
import org.apache.druid.segment.loading.StorageLocation;
import org.apache.druid.segment.loading.StorageLocationConfig;
import org.apache.druid.segment.loading.TombstoneLoadSpec;
@@ -1263,7 +1264,7 @@ public abstract class CompactionTaskRunBase
Assert.assertEquals(new NumberedShardSpec(0, 1), segments.get(0).getShardSpec());
final File cacheDir = temporaryFolder.newFolder();
- final SegmentCacheManager segmentCacheManager = segmentCacheManagerFactory.manufacturate(cacheDir, false);
+ final SegmentCacheManager segmentCacheManager = segmentCacheManagerFactory.manufacturate(cacheDir, null, false);
List<String> rowsFromSegment = new ArrayList<>();
for (DataSegment segment : segments) {
@@ -1375,7 +1376,7 @@ public abstract class CompactionTaskRunBase
Assert.assertEquals(new NumberedShardSpec(0, 1), segments.get(0).getShardSpec());
final File cacheDir = temporaryFolder.newFolder();
- final SegmentCacheManager segmentCacheManager = segmentCacheManagerFactory.manufacturate(cacheDir, false);
+ final SegmentCacheManager segmentCacheManager = segmentCacheManagerFactory.manufacturate(cacheDir, null, false);
List<String> rowsFromSegment = new ArrayList<>();
for (DataSegment segment : segments) {
@@ -1492,7 +1493,7 @@ public abstract class CompactionTaskRunBase
Assert.assertEquals(new NumberedShardSpec(0, 1), compactSegment.getShardSpec());
final File cacheDir = temporaryFolder.newFolder();
- final SegmentCacheManager segmentCacheManager = segmentCacheManagerFactory.manufacturate(cacheDir, false);
+ final SegmentCacheManager segmentCacheManager = segmentCacheManagerFactory.manufacturate(cacheDir, null, false);
List<String> rowsFromSegment = new ArrayList<>();
segmentCacheManager.load(compactSegment);
@@ -1714,6 +1715,7 @@ public abstract class CompactionTaskRunBase
final SegmentCacheManager cacheManager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestIndex.INDEX_IO,
objectMapper
@@ -1748,7 +1750,7 @@ public abstract class CompactionTaskRunBase
protected List<String> getCSVFormatRowsFromSegments(List<DataSegment> segments) throws Exception
{
final File cacheDir = temporaryFolder.newFolder();
- final SegmentCacheManager segmentCacheManager = segmentCacheManagerFactory.manufacturate(cacheDir, false);
+ final SegmentCacheManager segmentCacheManager = segmentCacheManagerFactory.manufacturate(cacheDir, null, false);
List<String> rowsFromSegment = new ArrayList<>();
for (DataSegment segment : segments) {
diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/IndexTaskTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/IndexTaskTest.java
index a2b0b225553..399cb9a513e 100644
--- a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/IndexTaskTest.java
+++ b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/IndexTaskTest.java
@@ -88,6 +88,7 @@ import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStr
import org.apache.druid.segment.loading.SegmentCacheManager;
import org.apache.druid.segment.loading.SegmentLoaderConfig;
import org.apache.druid.segment.loading.SegmentLocalCacheManager;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
import org.apache.druid.segment.loading.StorageLocation;
import org.apache.druid.segment.loading.StorageLocationConfig;
import org.apache.druid.segment.realtime.WindowedCursorFactory;
@@ -219,6 +220,7 @@ public class IndexTaskTest extends IngestionTestBase
segmentCacheManager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestIndex.INDEX_IO,
jsonMapper
diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/IngestionTestBase.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/IngestionTestBase.java
index 47a5ee76cd4..7ca28f7f3b7 100644
--- a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/IngestionTestBase.java
+++ b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/IngestionTestBase.java
@@ -227,7 +227,7 @@ public abstract class IngestionTestBase extends InitializedNullHandlingTest
public SegmentCacheManager newSegmentLoader(File storageDir)
{
- return segmentCacheManagerFactory.manufacturate(storageDir, true);
+ return segmentCacheManagerFactory.manufacturate(storageDir, null, true);
}
public ObjectMapper getObjectMapper()
diff --git a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractMultiPhaseParallelIndexingTest.java b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractMultiPhaseParallelIndexingTest.java
index 1790c0d36bb..c4c2baa7999 100644
--- a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractMultiPhaseParallelIndexingTest.java
+++ b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/batch/parallel/AbstractMultiPhaseParallelIndexingTest.java
@@ -257,7 +257,7 @@ abstract class AbstractMultiPhaseParallelIndexingTest extends AbstractParallelIn
private Segment loadSegment(DataSegment dataSegment, File tempSegmentDir)
{
final SegmentCacheManager cacheManager = new SegmentCacheManagerFactory(TestIndex.INDEX_IO, getObjectMapper())
- .manufacturate(tempSegmentDir, false);
+ .manufacturate(tempSegmentDir, null, false);
try {
cacheManager.load(dataSegment);
return cacheManager.acquireCachedSegment(dataSegment.getId()).orElseThrow();
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/counters/ChannelCounters.java b/multi-stage-query/src/main/java/org/apache/druid/msq/counters/ChannelCounters.java
index 975cd50aaf0..2df4ea1b5d6 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/counters/ChannelCounters.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/counters/ChannelCounters.java
@@ -40,7 +40,7 @@ import java.util.concurrent.TimeUnit;
*/
public class ChannelCounters implements QueryCounter
{
- private static final int NO_PARTITION = 0;
+ public static final int NO_PARTITION = 0;
@GuardedBy("this")
private final LongList rows = new LongArrayList();
@@ -162,7 +162,16 @@ public class ChannelCounters implements QueryCounter
}
}
- private void add(
+ /**
+ * Most-generic "add" method.
+ *
+ * @param partitionNumber add counters to this partition
+ * @param nRows add this number of rows
+ * @param nBytes add this number of bytes
+ * @param nFrames add this number of frames
+ * @param nFiles add this number of files
+ */
+ public void add(
final int partitionNumber,
final long nRows,
final long nBytes,
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartFrameContext.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartFrameContext.java
index fca50d51efe..2c679e9f348 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartFrameContext.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartFrameContext.java
@@ -40,6 +40,7 @@ import org.apache.druid.segment.SegmentWrangler;
import org.apache.druid.segment.incremental.NoopRowIngestionMeters;
import org.apache.druid.segment.incremental.RowIngestionMeters;
import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.server.SegmentManager;
import javax.annotation.Nullable;
@@ -54,6 +55,7 @@ public class DartFrameContext implements FrameContext
private final FrameWriterSpec frameWriterSpec;
private final SegmentWrangler segmentWrangler;
private final SegmentManager segmentManager;
+ private final VirtualStorageManager virtualStorageManager;
private final CoordinatorClient coordinatorClient;
private final WorkerContext workerContext;
@@ -78,6 +80,7 @@ public class DartFrameContext implements FrameContext
final FrameWriterSpec frameWriterSpec,
final SegmentWrangler segmentWrangler,
final SegmentManager segmentManager,
+ final VirtualStorageManager virtualStorageManager,
final CoordinatorClient coordinatorClient,
@Nullable final ProcessingBuffersSet processingBuffersSet,
final WorkerMemoryParameters memoryParameters,
@@ -89,6 +92,7 @@ public class DartFrameContext implements FrameContext
this.segmentWrangler = segmentWrangler;
this.frameWriterSpec = frameWriterSpec;
this.segmentManager = segmentManager;
+ this.virtualStorageManager = virtualStorageManager;
this.coordinatorClient = coordinatorClient;
this.workerContext = workerContext;
this.processingBuffersSet = processingBuffersSet;
@@ -121,6 +125,12 @@ public class DartFrameContext implements FrameContext
return segmentManager;
}
+ @Override
+ public VirtualStorageManager virtualStorageManager()
+ {
+ return virtualStorageManager;
+ }
+
@Override
public CoordinatorClient coordinatorClient()
{
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContext.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContext.java
index 39888b75717..e43d79d29f7 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContext.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContext.java
@@ -51,6 +51,7 @@ import org.apache.druid.query.QueryContext;
import org.apache.druid.query.QueryContexts;
import org.apache.druid.query.policy.PolicyEnforcer;
import org.apache.druid.segment.SegmentWrangler;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.server.DruidNode;
import org.apache.druid.server.SegmentManager;
import org.apache.druid.utils.CloseableUtils;
@@ -81,6 +82,7 @@ public class DartWorkerContext implements WorkerContext
private final DartWorkerClient workerClient;
private final SegmentWrangler segmentWrangler;
private final SegmentManager segmentManager;
+ private final VirtualStorageManager virtualStorageManager;
private final CoordinatorClient coordinatorClient;
private final MemoryIntrospector memoryIntrospector;
private final ProcessingBuffersProvider processingBuffersProvider;
@@ -109,6 +111,7 @@ public class DartWorkerContext implements WorkerContext
final DruidProcessingConfig processingConfig,
final SegmentWrangler segmentWrangler,
final SegmentManager segmentManager,
+ final VirtualStorageManager virtualStorageManager,
final CoordinatorClient coordinatorClient,
final MemoryIntrospector memoryIntrospector,
final ProcessingBuffersProvider processingBuffersProvider,
@@ -131,6 +134,7 @@ public class DartWorkerContext implements WorkerContext
this.workerClient = workerClient;
this.segmentWrangler = segmentWrangler;
this.segmentManager = segmentManager;
+ this.virtualStorageManager = virtualStorageManager;
this.coordinatorClient = coordinatorClient;
this.memoryIntrospector = memoryIntrospector;
this.processingBuffersProvider = processingBuffersProvider;
@@ -258,6 +262,7 @@ public class DartWorkerContext implements WorkerContext
FrameWriterSpec.fromContext(workOrder.getWorkerContext()),
segmentWrangler,
segmentManager,
+ virtualStorageManager,
coordinatorClient,
workOrder.getStageDefinition().getProcessor().usesProcessingBuffers() ? processingBuffersSet.get() : null,
memoryParameters,
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContextFactoryImpl.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContextFactoryImpl.java
index dc74cea3513..16e8e187b98 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContextFactoryImpl.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerContextFactoryImpl.java
@@ -37,10 +37,10 @@ import org.apache.druid.msq.exec.WorkerContext;
import org.apache.druid.msq.input.InputSliceReaderProvider;
import org.apache.druid.query.DruidProcessingConfig;
import org.apache.druid.query.QueryContext;
-import org.apache.druid.query.groupby.GroupingEngine;
import org.apache.druid.query.policy.PolicyEnforcer;
import org.apache.druid.rpc.ServiceClientFactory;
import org.apache.druid.segment.SegmentWrangler;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.server.DruidNode;
import org.apache.druid.server.SegmentManager;
@@ -61,8 +61,8 @@ public class DartWorkerContextFactoryImpl implements DartWorkerContextFactory
private final ServiceClientFactory serviceClientFactory;
private final DruidProcessingConfig processingConfig;
private final SegmentWrangler segmentWrangler;
- private final GroupingEngine groupingEngine;
private final SegmentManager segmentManager;
+ private final VirtualStorageManager virtualStorageManager;
private final CoordinatorClient coordinatorClient;
private final MemoryIntrospector memoryIntrospector;
private final ProcessingBuffersProvider processingBuffersProvider;
@@ -81,8 +81,8 @@ public class DartWorkerContextFactoryImpl implements DartWorkerContextFactory
@EscalatedGlobal ServiceClientFactory serviceClientFactory,
DruidProcessingConfig processingConfig,
SegmentWrangler segmentWrangler,
- GroupingEngine groupingEngine,
SegmentManager segmentManager,
+ VirtualStorageManager virtualStorageManager,
CoordinatorClient coordinatorClient,
MemoryIntrospector memoryIntrospector,
@Dart ProcessingBuffersProvider processingBuffersProvider,
@@ -101,8 +101,8 @@ public class DartWorkerContextFactoryImpl implements DartWorkerContextFactory
this.processingConfig = processingConfig;
this.segmentWrangler = segmentWrangler;
this.coordinatorClient = coordinatorClient;
- this.groupingEngine = groupingEngine;
this.segmentManager = segmentManager;
+ this.virtualStorageManager = virtualStorageManager;
this.memoryIntrospector = memoryIntrospector;
this.processingBuffersProvider = processingBuffersProvider;
this.outbox = outbox;
@@ -130,6 +130,7 @@ public class DartWorkerContextFactoryImpl implements DartWorkerContextFactory
processingConfig,
segmentWrangler,
segmentManager,
+ virtualStorageManager,
coordinatorClient,
memoryIntrospector,
processingBuffersProvider,
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/DataServerQueryHandlerUtils.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/DataServerQueryHandlerUtils.java
index 392e6ee378e..9a034a152e8 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/DataServerQueryHandlerUtils.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/DataServerQueryHandlerUtils.java
@@ -35,6 +35,8 @@ import org.apache.druid.query.RestrictedDataSource;
import org.apache.druid.query.SegmentDescriptor;
import org.apache.druid.query.TableDataSource;
import org.apache.druid.query.context.ResponseContext;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.scan.ScanResultValue;
import org.apache.druid.sql.calcite.planner.PlannerContext;
import java.util.Collections;
@@ -130,7 +132,15 @@ public class DataServerQueryHandlerUtils
return Yielders.each(
mappingFunction.apply(sequence)
.map(row -> {
- channelCounters.incrementRowCount();
+ if (row instanceof RowsAndColumns rac) {
+ channelCounters.addRAC(rac, ChannelCounters.NO_PARTITION);
+ } else if (row instanceof ScanResultValue scanResult
+ && scanResult.getEvents() instanceof List<?> eventsList) {
+ // Special handling for ScanQuery
+ channelCounters.add(ChannelCounters.NO_PARTITION, eventsList.size(), 0, 0, 0);
+ } else {
+ channelCounters.incrementRowCount();
+ }
return row;
})
);
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/FrameContext.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/FrameContext.java
index 1a3b99538d1..8cc2e9a66f3 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/FrameContext.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/FrameContext.java
@@ -30,6 +30,7 @@ import org.apache.druid.segment.IndexMerger;
import org.apache.druid.segment.SegmentWrangler;
import org.apache.druid.segment.incremental.RowIngestionMeters;
import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.server.SegmentManager;
import javax.annotation.Nullable;
@@ -55,6 +56,11 @@ public interface FrameContext extends Closeable
*/
SegmentManager segmentManager();
+ /**
+ * Returns the virtual storage manager for caching files.
+ */
+ VirtualStorageManager virtualStorageManager();
+
/**
* Returns the coordinator client for fetching DataSegment metadata when not available locally.
* May be null if no coordinator client is available (e.g., in Dart workers).
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java
index 20593f26621..e064d78640e 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java
@@ -403,7 +403,14 @@ public class RunWorkOrder
LinkedHashMap<Class<? extends InputSlice>, InputSliceReader> readers = new LinkedHashMap<>();
readers.put(NilInputSlice.class, NilInputSliceReader.INSTANCE);
readers.put(StageInputSlice.class, StageInputSliceReader.INSTANCE);
- readers.put(ExternalInputSlice.class, new ExternalInputSliceReader(frameContext.tempDir("external")));
+ readers.put(
+ ExternalInputSlice.class,
+ new ExternalInputSliceReader(
+ frameContext.virtualStorageManager(),
+ frameContext.tempDir("external"),
+ MultiStageQueryContext.isBackgroundFetchExternalFiles(workOrder.getWorkerContext())
+ )
+ );
readers.put(InlineInputSlice.class, new InlineInputSliceReader(frameContext.segmentWrangler()));
readers.put(LookupInputSlice.class, new LookupInputSliceReader(frameContext.segmentWrangler()));
readers.put(SegmentsInputSlice.class, new SegmentsInputSliceReader(frameContext, reindex));
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerFrameContext.java b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerFrameContext.java
index 13da1393c71..2684c25a6d6 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerFrameContext.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerFrameContext.java
@@ -39,6 +39,7 @@ import org.apache.druid.segment.IndexMerger;
import org.apache.druid.segment.SegmentWrangler;
import org.apache.druid.segment.incremental.RowIngestionMeters;
import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.server.SegmentManager;
import javax.annotation.Nullable;
@@ -51,6 +52,7 @@ public class IndexerFrameContext implements FrameContext
private final FrameWriterSpec frameWriterSpec;
private final IndexIO indexIO;
private final SegmentManager segmentManager;
+ private final VirtualStorageManager virtualStorageManager;
@Nullable
private final CoordinatorClient coordinatorClient;
@@ -75,6 +77,7 @@ public class IndexerFrameContext implements FrameContext
FrameWriterSpec frameWriterSpec,
IndexIO indexIO,
SegmentManager segmentManager,
+ VirtualStorageManager virtualStorageManager,
@Nullable CoordinatorClient coordinatorClient,
@Nullable ProcessingBuffersSet processingBuffersSet,
IndexerDataServerQueryHandlerFactory dataServerQueryHandlerFactory,
@@ -87,6 +90,7 @@ public class IndexerFrameContext implements FrameContext
this.frameWriterSpec = frameWriterSpec;
this.indexIO = indexIO;
this.segmentManager = segmentManager;
+ this.virtualStorageManager = virtualStorageManager;
this.coordinatorClient = coordinatorClient;
this.processingBuffersSet = processingBuffersSet;
this.memoryParameters = memoryParameters;
@@ -118,6 +122,12 @@ public class IndexerFrameContext implements FrameContext
return segmentManager;
}
+ @Override
+ public VirtualStorageManager virtualStorageManager()
+ {
+ return virtualStorageManager;
+ }
+
@Override
@Nullable
public CoordinatorClient coordinatorClient()
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerWorkerContext.java b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerWorkerContext.java
index b37340b4a42..3e1c151e5f2 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerWorkerContext.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerWorkerContext.java
@@ -30,6 +30,7 @@ import org.apache.druid.guice.annotations.EscalatedGlobal;
import org.apache.druid.guice.annotations.Smile;
import org.apache.druid.indexing.common.SegmentCacheManagerFactory;
import org.apache.druid.indexing.common.TaskToolbox;
+import org.apache.druid.indexing.common.config.TaskConfig;
import org.apache.druid.java.util.common.io.Closer;
import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.msq.exec.ControllerClient;
@@ -60,7 +61,10 @@ import org.apache.druid.rpc.indexing.OverlordClient;
import org.apache.druid.rpc.indexing.SpecificTaskRetryPolicy;
import org.apache.druid.rpc.indexing.SpecificTaskServiceLocator;
import org.apache.druid.segment.IndexIO;
+import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStrategy;
import org.apache.druid.segment.loading.SegmentCacheManager;
+import org.apache.druid.segment.loading.external.StorageLocationVirtualStorageManager;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.server.DruidNode;
import org.apache.druid.server.SegmentManager;
import org.apache.druid.server.metrics.StorageMonitor;
@@ -88,6 +92,7 @@ public class IndexerWorkerContext implements WorkerContext
private final ServiceLocator controllerLocator;
private final IndexIO indexIO;
private final SegmentManager segmentManager;
+ private final VirtualStorageManager virtualStorageManager;
private final StorageMonitor storageMonitor;
@Nullable
private final CoordinatorClient coordinatorClient;
@@ -113,6 +118,7 @@ public class IndexerWorkerContext implements WorkerContext
final ServiceLocator controllerLocator,
final IndexIO indexIO,
final SegmentManager segmentManager,
+ final VirtualStorageManager virtualStorageManager,
final StorageMonitor storageMonitor,
@Nullable final CoordinatorClient coordinatorClient,
final ServiceClientFactory clientFactory,
@@ -128,6 +134,7 @@ public class IndexerWorkerContext implements WorkerContext
this.controllerLocator = controllerLocator;
this.indexIO = indexIO;
this.segmentManager = segmentManager;
+ this.virtualStorageManager = virtualStorageManager;
this.storageMonitor = storageMonitor;
this.coordinatorClient = coordinatorClient;
this.clientFactory = clientFactory;
@@ -166,10 +173,23 @@ public class IndexerWorkerContext implements WorkerContext
)
{
final IndexIO indexIO = injector.getInstance(IndexIO.class);
+ final TaskConfig taskConfig = injector.getInstance(TaskConfig.class);
final SegmentCacheManager cacheManager =
injector.getInstance(SegmentCacheManagerFactory.class)
- .manufacturate(new File(toolbox.getIndexingTmpDir(), "segment-fetch"), true);
+ .manufacturate(
+ new File(toolbox.getIndexingTmpDir(), "segment-fetch"),
+ // Divide tmpStorageBytesPerTask by 3 so the local cache never takes up the majority of space.
+ // In a typical leaf stage run, we may need some disk space for inputs and some for outputs.
+ taskConfig.getTmpStorageBytesPerTask() > 0 ? taskConfig.getTmpStorageBytesPerTask() / 3 : null,
+ true
+ );
final SegmentManager segmentManager = new SegmentManager(cacheManager);
+ final VirtualStorageManager virtualStorageManager =
+ new StorageLocationVirtualStorageManager(
+ cacheManager.getLocations(),
+ new LeastBytesUsedStorageLocationSelectorStrategy(cacheManager.getLocations()),
+ cacheManager.getLoadingThreadPool()
+ );
final StorageMonitor storageMonitor = new StorageMonitor(cacheManager.getLocations(), task::getMetricBuilder);
toolbox.addMonitor(storageMonitor);
final ServiceClientFactory serviceClientFactory =
@@ -190,6 +210,7 @@ public class IndexerWorkerContext implements WorkerContext
new SpecificTaskServiceLocator(task.getControllerTaskId(), overlordClient),
indexIO,
segmentManager,
+ virtualStorageManager,
storageMonitor,
toolbox.getCoordinatorClient(),
serviceClientFactory,
@@ -320,6 +341,7 @@ public class IndexerWorkerContext implements WorkerContext
FrameWriterSpec.fromContext(workOrder.getWorkerContext()),
indexIO,
segmentManager,
+ virtualStorageManager,
coordinatorClient,
workOrder.getStageDefinition().getProcessor().usesProcessingBuffers() ? processingBuffersSet.get() : null,
dataServerQueryHandlerFactory,
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/AdaptedLoadableSegment.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/AdaptedLoadableSegment.java
index 8897170c5e6..7000064be9b 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/input/AdaptedLoadableSegment.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/AdaptedLoadableSegment.java
@@ -21,7 +21,12 @@ package org.apache.druid.msq.input;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
+import org.apache.druid.common.asyncresource.AsyncResource;
+import org.apache.druid.common.asyncresource.AsyncResources;
import org.apache.druid.error.DruidException;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.io.Closer;
import org.apache.druid.msq.counters.ChannelCounters;
import org.apache.druid.query.SegmentDescriptor;
import org.apache.druid.segment.ReferenceCountedSegmentProvider;
@@ -29,11 +34,11 @@ import org.apache.druid.segment.Segment;
import org.apache.druid.segment.loading.AcquireSegmentAction;
import org.apache.druid.segment.loading.AcquireSegmentResult;
import org.apache.druid.timeline.DataSegment;
-import org.joda.time.Interval;
import javax.annotation.Nullable;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
/**
* Implementation of {@link LoadableSegment} for segments adapted from non-regular sources such as inline data,
@@ -44,64 +49,70 @@ import java.util.concurrent.atomic.AtomicBoolean;
public class AdaptedLoadableSegment implements LoadableSegment
{
private final AtomicBoolean acquired = new AtomicBoolean(false);
+ private final Supplier<AsyncResource<AcquireSegmentResult>> asyncSegmentSupplier;
private final SegmentDescriptor descriptor;
@Nullable
- private final ChannelCounters inputCounters;
- @Nullable
private final String description;
- private final ListenableFuture<DataSegment> dataSegmentFuture;
- private final AcquireSegmentAction acquireSegmentAction;
+ @Nullable
+ private final ChannelCounters inputCounters;
- private AdaptedLoadableSegment(
- final Segment segment,
+ /**
+ * Creates a wrapper around a supplier of an {@link AcquireSegmentResult}. The lifecycle of the supplied
+ * {@link AsyncResource} is tied to the {@link AcquireSegmentAction} returned from {@link #acquire()}.
+ *
+ * @param asyncSegmentSupplier the supplier to wrap
+ * @param descriptor descriptor containing the interval to use for filtering
+ * @param description user-oriented description for error messages
+ * @param inputCounters counters for tracking input via {@link LoadableSegmentUtils#countedLoad}.
+ */
+ public AdaptedLoadableSegment(
+ final Supplier<AsyncResource<AcquireSegmentResult>> asyncSegmentSupplier,
final SegmentDescriptor descriptor,
@Nullable final String description,
@Nullable final ChannelCounters inputCounters
)
{
+ this.asyncSegmentSupplier = asyncSegmentSupplier;
this.descriptor = descriptor;
this.description = description;
this.inputCounters = inputCounters;
-
- // These segments don't have an associated DataSegment
- this.dataSegmentFuture = Futures.immediateFailedFuture(
- DruidException.defensive("DataSegment not available for adapted segments")
- );
-
- // Pre-create the acquire action since the segment is already available
- final ListenableFuture<AcquireSegmentResult> segmentFuture =
- Futures.immediateFuture(AcquireSegmentResult.cached(ReferenceCountedSegmentProvider.of(segment)));
- this.acquireSegmentAction = new AcquireSegmentAction(() -> segmentFuture, null);
}
/**
* Creates an AdaptedLoadableSegment wrapper around a Segment object which is not a regular Druid segment,
* has no associated {@link DataSegment}, and whose lifecycle is not managed by the LoadableSegment instance.
*
- * @param segment the segment to wrap
- * @param queryInterval the query interval to use for filtering
- * @param description user-oriented description for error messages
- * @param channelCounters counters for tracking input
+ * @param segment the segment to wrap
+ * @param descriptor descriptor containing the interval to use for filtering
+ * @param description user-oriented description for error messages
+ * @param inputCounters counters for tracking input
*/
- public static AdaptedLoadableSegment create(
+ public static AdaptedLoadableSegment fromUnmanagedSegment(
final Segment segment,
- final Interval queryInterval,
+ final SegmentDescriptor descriptor,
@Nullable final String description,
- @Nullable final ChannelCounters channelCounters
+ @Nullable final ChannelCounters inputCounters
)
{
+ // Pre-create the acquire result since the segment is already available
+ final AcquireSegmentResult acquireSegmentResult =
+ AcquireSegmentResult.cached(ReferenceCountedSegmentProvider.of(segment));
+
+ final AsyncResource<AcquireSegmentResult> resource = AsyncResources.unmanaged(acquireSegmentResult);
return new AdaptedLoadableSegment(
- segment,
- new SegmentDescriptor(queryInterval, "0", 0),
+ () -> resource,
+ descriptor,
description,
- channelCounters
+ inputCounters
);
}
@Override
public ListenableFuture<DataSegment> dataSegmentFuture()
{
- return dataSegmentFuture;
+ return Futures.immediateFailedFuture(
+ DruidException.defensive("DataSegment not available for adapted segments")
+ );
}
@Override
@@ -110,13 +121,6 @@ public class AdaptedLoadableSegment implements LoadableSegment
return descriptor;
}
- @Override
- @Nullable
- public ChannelCounters inputCounters()
- {
- return inputCounters;
- }
-
@Override
@Nullable
public String description()
@@ -139,6 +143,45 @@ public class AdaptedLoadableSegment implements LoadableSegment
if (!acquired.compareAndSet(false, true)) {
throw DruidException.defensive("Segment with descriptor[%s] is already acquired", descriptor);
}
- return acquireSegmentAction;
+
+ // Synchronized by itself; Closer is not thread-safe.
+ final Closer closer = Closer.create();
+ final AtomicBoolean closed = new AtomicBoolean(false);
+
+ return new AcquireSegmentAction(
+ () -> {
+ final AsyncResource<AcquireSegmentResult> asyncSegment = asyncSegmentSupplier.get();
+ synchronized (closer) {
+ if (closed.get()) {
+ asyncSegment.close();
+ return Futures.immediateFailedFuture(new ISE("Already closed"));
+ } else {
+ closer.register(asyncSegment);
+ }
+ }
+
+ final SettableFuture<AcquireSegmentResult> retVal = SettableFuture.create();
+
+ asyncSegment.addReadyCallback(() -> {
+ try {
+ retVal.set(asyncSegment.get());
+ }
+ catch (Throwable e) {
+ retVal.setException(e);
+ }
+ });
+
+ // Use byteCount = 0 for adapted segments; we can't really tell what it is from the AcquireSegmentResult
+ // (the "load size" may not be the entire size if the segment was fully or partially cached). Implementations
+ // call ChannelCounters#incrementBytes if they have something useful to put there.
+ return LoadableSegmentUtils.countedLoad(retVal, 0, inputCounters);
+ },
+ () -> {
+ synchronized (closer) {
+ closed.set(true);
+ closer.close();
+ }
+ }
+ );
}
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegment.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegment.java
index 5c04d78383c..227c1ad2d49 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegment.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegment.java
@@ -22,7 +22,6 @@ package org.apache.druid.msq.input;
import com.google.common.util.concurrent.ListenableFuture;
import org.apache.druid.common.guava.FutureUtils;
import org.apache.druid.error.DruidException;
-import org.apache.druid.msq.counters.ChannelCounters;
import org.apache.druid.query.SegmentDescriptor;
import org.apache.druid.segment.Segment;
import org.apache.druid.segment.loading.AcquireSegmentAction;
@@ -51,13 +50,6 @@ public interface LoadableSegment
*/
SegmentDescriptor descriptor();
- /**
- * Returns input counters that should be updated as this segment is queried, or null if it is not necessary
- * to update input counters.
- */
- @Nullable
- ChannelCounters inputCounters();
-
/**
* User-oriented description, suitable for inclusion in log or error messages.
*/
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegmentUtils.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegmentUtils.java
new file mode 100644
index 00000000000..a092777eeef
--- /dev/null
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/LoadableSegmentUtils.java
@@ -0,0 +1,83 @@
+/*
+ * 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.msq.input;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.segment.PhysicalSegmentInspector;
+import org.apache.druid.segment.Segment;
+import org.apache.druid.segment.loading.AcquireSegmentAction;
+import org.apache.druid.segment.loading.AcquireSegmentResult;
+
+import javax.annotation.Nullable;
+import java.util.Optional;
+
+public class LoadableSegmentUtils
+{
+ /**
+ * Given a future from {@link AcquireSegmentAction#getSegmentFuture()}, wraps it with logic to increment
+ * counters as follows:
+ *
+ * <ul>
+ * <li>{@link ChannelCounters#addLoad(AcquireSegmentResult)} when the load completes</li>
+ * <li>{@link ChannelCounters#addFile(long, long)} each time a reference is acquired from
+ * {@link AcquireSegmentResult#getReferenceProvider()}. The row count is taken from
+ * {@link #getSegmentRowCount(Segment)}, and byte count is taken from {@code byteCount}.</li>
+ * </ul>
+ */
+ public static ListenableFuture<AcquireSegmentResult> countedLoad(
+ final ListenableFuture<AcquireSegmentResult> segmentFuture,
+ final long byteCount,
+ @Nullable final ChannelCounters channelCounters
+ )
+ {
+ if (channelCounters == null) {
+ return segmentFuture;
+ } else {
+ return FutureUtils.transform(
+ segmentFuture,
+ result -> {
+ channelCounters.addLoad(result);
+ return new AcquireSegmentResult(
+ () -> {
+ final Optional<Segment> segment = result.getReferenceProvider().acquireReference();
+ final int rowCount = segment.map(LoadableSegmentUtils::getSegmentRowCount).orElse(0);
+ channelCounters.addFile(rowCount, byteCount);
+ return segment;
+ },
+ result.getLoadSizeBytes(),
+ result.getWaitTimeNanos(),
+ result.getLoadTimeNanos()
+ );
+ }
+ );
+ }
+ }
+
+ /**
+ * Gets the number of rows for a segment, using a {@link PhysicalSegmentInspector}. Returns 0 when unknown.
+ */
+ public static int getSegmentRowCount(final Segment segment)
+ {
+ final PhysicalSegmentInspector inspector = segment.as(PhysicalSegmentInspector.class);
+ return inspector != null ? inspector.getNumRows() : 0;
+ }
+}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java
index 7582cd2514c..b43eefdcf36 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java
@@ -77,7 +77,7 @@ public class RegularLoadableSegment implements LoadableSegment
* @param segmentManager segment manager for loading and caching segments
* @param segmentId the segment ID to load
* @param descriptor segment descriptor for querying
- * @param inputCounters optional counters for tracking input
+ * @param inputCounters optional counters for tracking input via {@link LoadableSegmentUtils#countedLoad}
* @param coordinatorClient optional client for fetching DataSegment from Coordinator when not available locally
* @param isReindex true if this is a DML command writing to the same table it's reading from
*/
@@ -119,13 +119,6 @@ public class RegularLoadableSegment implements LoadableSegment
return descriptor;
}
- @Override
- @Nullable
- public ChannelCounters inputCounters()
- {
- return inputCounters;
- }
-
@Override
@Nullable
public String description()
@@ -143,6 +136,13 @@ public class RegularLoadableSegment implements LoadableSegment
final Optional<Segment> cachedSegment = segmentManager.acquireCachedSegment(segmentId);
if (cachedSegment.isPresent()) {
acquired = true;
+
+ // Update counters in the manner of LoadableSegmentUtils#countedLoad (which we aren't using here).
+ if (inputCounters != null) {
+ final int rowCount = LoadableSegmentUtils.getSegmentRowCount(cachedSegment.get());
+ final long byteCount = cachedDataSegment != null ? cachedDataSegment.getSize() : 0;
+ inputCounters.addFile(rowCount, byteCount);
+ }
}
return cachedSegment;
}
@@ -157,7 +157,15 @@ public class RegularLoadableSegment implements LoadableSegment
acquired = true;
if (cachedDataSegment != null) {
- return segmentManager.acquireSegment(cachedDataSegment);
+ final AcquireSegmentAction action = segmentManager.acquireSegment(cachedDataSegment);
+ return new AcquireSegmentAction(
+ () -> LoadableSegmentUtils.countedLoad(
+ action.getSegmentFuture(),
+ cachedDataSegment.getSize(),
+ inputCounters
+ ),
+ action
+ );
} else {
// Create a shim AcquireSegmentAction that doesn't acquire a hold (yet). We can't make a real
// AcquireSegmentAction yet because we don't have the DataSegment object. It needs to be fetched
@@ -168,7 +176,11 @@ public class RegularLoadableSegment implements LoadableSegment
return new AcquireSegmentAction(
Suppliers.memoize(() -> FutureUtils.transformAsync(
dataSegmentFutureSupplier.get(),
- dataSegment -> closer.register(segmentManager.acquireSegment(dataSegment)).getSegmentFuture()
+ dataSegment -> LoadableSegmentUtils.countedLoad(
+ closer.register(segmentManager.acquireSegment(dataSegment)).getSegmentFuture(),
+ dataSegment.getSize(),
+ inputCounters
+ )
)),
closer
);
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalInputSliceReader.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalInputSliceReader.java
index 310175b0ef8..5a95c53f4c4 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalInputSliceReader.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalInputSliceReader.java
@@ -19,18 +19,22 @@
package org.apache.druid.msq.input.external;
+import org.apache.druid.common.asyncresource.AsyncResource;
+import org.apache.druid.common.asyncresource.AsyncResources;
import org.apache.druid.data.input.ColumnsFilter;
+import org.apache.druid.data.input.InputFilePointer;
import org.apache.druid.data.input.InputFormat;
import org.apache.druid.data.input.InputRowSchema;
import org.apache.druid.data.input.InputSource;
import org.apache.druid.data.input.InputSourceReader;
-import org.apache.druid.data.input.InputStats;
import org.apache.druid.data.input.impl.DimensionsSpec;
import org.apache.druid.data.input.impl.InlineInputSource;
+import org.apache.druid.data.input.impl.LocalInputSource;
import org.apache.druid.data.input.impl.TimestampSpec;
import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.msq.counters.ChannelCounters;
import org.apache.druid.msq.counters.CounterNames;
import org.apache.druid.msq.counters.CounterTracker;
@@ -45,12 +49,14 @@ import org.apache.druid.msq.input.PhysicalInputSlice;
import org.apache.druid.msq.input.stage.ReadablePartitions;
import org.apache.druid.msq.util.DimensionSchemaUtils;
import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.segment.ReferenceCountedSegmentProvider;
import org.apache.druid.segment.RowBasedSegment;
-import org.apache.druid.segment.Segment;
import org.apache.druid.segment.column.ColumnHolder;
import org.apache.druid.segment.column.RowSignature;
-import org.apache.druid.segment.incremental.SimpleRowIngestionMeters;
-import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.segment.loading.AcquireSegmentResult;
+import org.apache.druid.segment.loading.external.CachedFile;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
+import org.apache.druid.utils.CloseableUtils;
import java.io.File;
import java.util.ArrayList;
@@ -64,13 +70,21 @@ import java.util.stream.Collectors;
*/
public class ExternalInputSliceReader implements InputSliceReader
{
- public static final String SEGMENT_ID = "__external";
- public static final SegmentDescriptor SEGMENT_DESCRIPTOR = SegmentId.dummy(SEGMENT_ID).toDescriptor();
+ private static final Logger log = new Logger(ExternalInputSliceReader.class);
+
+ private final VirtualStorageManager virtualStorageManager;
private final File temporaryDirectory;
+ private final boolean backgroundFetchExternalFiles;
- public ExternalInputSliceReader(final File temporaryDirectory)
+ public ExternalInputSliceReader(
+ final VirtualStorageManager virtualStorageManager,
+ final File temporaryDirectory,
+ final boolean backgroundFetchExternalFiles
+ )
{
+ this.virtualStorageManager = virtualStorageManager;
this.temporaryDirectory = temporaryDirectory;
+ this.backgroundFetchExternalFiles = backgroundFetchExternalFiles;
}
public static boolean isFileBasedInputSource(final InputSource inputSource)
@@ -92,26 +106,15 @@ public class ExternalInputSliceReader implements InputSliceReader
final List<LoadableSegment> loadableSegments = new ArrayList<>();
for (final InputSource inputSource : externalInputSlice.getInputSources()) {
- // The LoadableSegment generated here does not acquire a real hold, and ends up loading the external data in a
- // processing thread (when the cursor is created). Ideally, this would be better integrated with the virtual
- // storage system, giving us storage holds and the ability to load data outside of a processing thread.
- final Segment segment = makeExternalSegment(
+ final LoadableSegment segment = makeExternalSegment(
inputSource,
externalInputSlice.getInputFormat(),
externalInputSlice.getSignature(),
- new File(temporaryDirectory, String.valueOf(inputNumber)),
inputCounters,
counters.warnings(),
warningPublisher
);
- loadableSegments.add(
- AdaptedLoadableSegment.create(
- segment,
- Intervals.ETERNITY,
- StringUtils.format("external[%s]", inputSource.toString()),
- null
- )
- );
+ loadableSegments.add(segment);
}
return new PhysicalInputSlice(ReadablePartitions.empty(), loadableSegments, Collections.emptyList());
@@ -120,11 +123,10 @@ public class ExternalInputSliceReader implements InputSliceReader
/**
* Creates a lazy segment that fetches external data when a cursor is created.
*/
- private Segment makeExternalSegment(
+ private LoadableSegment makeExternalSegment(
final InputSource inputSource,
final InputFormat inputFormat,
final RowSignature signature,
- final File temporaryDirectory,
final ChannelCounters channelCounters,
final WarningCounters warningCounters,
final Consumer<Throwable> warningPublisher
@@ -144,27 +146,128 @@ public class ExternalInputSliceReader implements InputSliceReader
ColumnsFilter.all()
);
- final InputSourceReader reader;
+ final String description = StringUtils.format("external[%s]", inputSource.toString());
+ final SegmentDescriptor descriptor = new SegmentDescriptor(Intervals.ETERNITY, "0", 0);
+ final List<InputFilePointer> filePointers = inputSource.asFilePointers();
+
+ if (filePointers != null && backgroundFetchExternalFiles) {
+ // The LoadableSegment generated here will cause files to be fetched in VSF loading threads
+ // when acquire() is called.
+ return new AdaptedLoadableSegment(
+ () -> {
+ final long startTime = System.nanoTime();
+ final List<AsyncResource<CachedFile>> cachedFileResources = new ArrayList<>(filePointers.size());
+ try {
+ for (final InputFilePointer filePointer : filePointers) {
+ cachedFileResources.add(
+ virtualStorageManager.reserveAndPopulateAsync(
+ filePointer.uri().toString(),
+ filePointer.sizeSupplier(),
+ filePointer.populator()
+ )
+ );
+ }
+
+ final AsyncResource<AcquireSegmentResult> fetched = AsyncResources.transform(
+ AsyncResources.collect(cachedFileResources),
+ cachedFiles -> {
+ long totalSize = 0;
+ final List<File> files = new ArrayList<>(cachedFiles.size());
+
+ for (final CachedFile cachedFile : cachedFiles) {
+ files.add(cachedFile.getFile());
+ totalSize += cachedFile.getFile().length();
+ }
+
+ final InputSource localInputSource = new LocalInputSource(null, null, files, null);
+ final ExternalSegment segment = new ExternalSegment(
+ localInputSource,
+ makeReader(schema, localInputSource, inputFormat, channelCounters),
+ warningCounters,
+ warningPublisher,
+ channelCounters,
+ signature
+ );
+
+ return new AcquireSegmentResult(
+ ReferenceCountedSegmentProvider.of(segment),
+ totalSize,
+ 0L,
+ System.nanoTime() - startTime
+ );
+ }
+ );
+
+ // If the fetch fails because of insufficient storage space, release the fetched files and fall
+ // back to creating an ExternalSegment that streams data directly from the original input source.
+ return AsyncResources.recover(
+ fetched,
+ e -> {
+ if (VirtualStorageManager.isInsufficientStorage(e)) {
+ log.noStackTrace()
+ .info(e, "Insufficient storage space to prefetch[%s]; streaming instead.", description);
+ return AcquireSegmentResult.cached(
+ ReferenceCountedSegmentProvider.of(
+ new ExternalSegment(
+ inputSource,
+ makeReader(schema, inputSource, inputFormat, channelCounters),
+ warningCounters,
+ warningPublisher,
+ channelCounters,
+ signature
+ )
+ )
+ );
+ } else {
+ return null;
+ }
+ }
+ );
+ }
+ catch (Throwable e) {
+ // Close any cachedFileResources that have been created prior to this exception being thrown.
+ throw CloseableUtils.closeAndWrapInCatch(e, CloseableUtils.forIterable(cachedFileResources));
+ }
+ },
+ descriptor,
+ description,
+ channelCounters
+ );
+ } else {
+ // The LoadableSegment generated here does not acquire a real hold, and ends up loading the external data in a
+ // processing thread (when the cursor is created).
+ return AdaptedLoadableSegment.fromUnmanagedSegment(
+ new ExternalSegment(
+ inputSource,
+ makeReader(schema, inputSource, inputFormat, channelCounters),
+ warningCounters,
+ warningPublisher,
+ channelCounters,
+ signature
+ ),
+ descriptor,
+ description,
+ channelCounters
+ );
+ }
+ }
+
+ private InputSourceReader makeReader(
+ final InputRowSchema schema,
+ final InputSource inputSource,
+ final InputFormat inputFormat,
+ final ChannelCounters channelCounters
+ )
+ {
final boolean incrementCounters = isFileBasedInputSource(inputSource);
- final InputStats inputStats = new SimpleRowIngestionMeters();
if (incrementCounters) {
- reader = new CountableInputSourceReader(
+ return new CountableInputSourceReader(
inputSource.reader(schema, inputFormat, temporaryDirectory),
channelCounters
);
} else {
- reader = inputSource.reader(schema, inputFormat, temporaryDirectory);
+ return inputSource.reader(schema, inputFormat, temporaryDirectory);
}
-
- return new ExternalSegment(
- inputSource,
- reader,
- inputStats,
- warningCounters,
- warningPublisher,
- channelCounters,
- signature
- );
}
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalSegment.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalSegment.java
index 4c21616eb3f..48e22f1f298 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalSegment.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalSegment.java
@@ -31,9 +31,12 @@ import org.apache.druid.msq.counters.WarningCounters;
import org.apache.druid.msq.indexing.error.CannotParseExternalDataFault;
import org.apache.druid.segment.RowBasedSegment;
import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.segment.incremental.NoopRowIngestionMeters;
+import org.apache.druid.utils.CloseableUtils;
import java.io.IOException;
import java.util.NoSuchElementException;
+import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
/**
@@ -46,16 +49,14 @@ public class ExternalSegment extends RowBasedSegment<InputRow>
/**
* @param inputSource {@link InputSource} that the segment is a representation of
* @param reader reader to read the external input source
- * @param inputStats input stats
* @param warningCounters warning counters tracking the warnings generated while reading the external source
* @param warningPublisher publisher to report the warnings generated
- * @param channelCounters channel counters to increment as we read through the files/units of the external source
+ * @param channelCounters channel counters, will be used for {@link ChannelCounters#incrementBytes(long)}
* @param signature signature of the external source
*/
public ExternalSegment(
final InputSource inputSource,
final InputSourceReader reader,
- final InputStats inputStats,
final WarningCounters warningCounters,
final Consumer<Throwable> warningPublisher,
final ChannelCounters channelCounters,
@@ -70,7 +71,7 @@ public class ExternalSegment extends RowBasedSegment<InputRow>
public CloseableIterator<InputRow> make()
{
try {
- CloseableIterator<InputRow> baseIterator = reader.read(inputStats);
+ CloseableIterator<InputRow> baseIterator = reader.read(makeInputStats(inputSource, channelCounters));
return new CloseableIterator<>()
{
private InputRow next = null;
@@ -119,20 +120,7 @@ public class ExternalSegment extends RowBasedSegment<InputRow>
@Override
public void cleanup(CloseableIterator<InputRow> iterFromMake)
{
- try {
- iterFromMake.close();
- // We increment the file count whenever the caller calls clean up. So we can double count here
- // if the callers are not careful.
- // This logic only works because we are using FilePerSplitHintSpec. Each input source only
- // has one file.
- if (ExternalInputSliceReader.isFileBasedInputSource(inputSource)) {
- channelCounters.incrementFileCount();
- channelCounters.incrementBytes(inputStats.getProcessedBytes());
- }
- }
- catch (IOException e) {
- throw new RuntimeException(e);
- }
+ CloseableUtils.closeAndWrapExceptions(iterFromMake);
}
}
),
@@ -149,4 +137,32 @@ public class ExternalSegment extends RowBasedSegment<InputRow>
{
return signature;
}
+
+ /**
+ * Create {@link InputStats} that calls {@link ChannelCounters#incrementBytes(long)} as data is read from files.
+ */
+ private static InputStats makeInputStats(final InputSource inputSource, final ChannelCounters channelCounters)
+ {
+ if (ExternalInputSliceReader.isFileBasedInputSource(inputSource)) {
+ return new InputStats()
+ {
+ private final AtomicLong processedBytes = new AtomicLong();
+
+ @Override
+ public void incrementProcessedBytes(final long n)
+ {
+ processedBytes.addAndGet(n);
+ channelCounters.incrementBytes(n);
+ }
+
+ @Override
+ public long getProcessedBytes()
+ {
+ return processedBytes.get();
+ }
+ };
+ } else {
+ return new NoopRowIngestionMeters();
+ }
+ }
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/inline/InlineInputSliceReader.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/inline/InlineInputSliceReader.java
index 7e6c8daa7eb..0a1dc9305b2 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/input/inline/InlineInputSliceReader.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/inline/InlineInputSliceReader.java
@@ -30,6 +30,7 @@ import org.apache.druid.msq.input.PhysicalInputSlice;
import org.apache.druid.msq.input.stage.ReadablePartitions;
import org.apache.druid.msq.input.table.RichSegmentDescriptor;
import org.apache.druid.query.InlineDataSource;
+import org.apache.druid.query.SegmentDescriptor;
import org.apache.druid.segment.InlineSegmentWrangler;
import org.apache.druid.segment.Segment;
import org.apache.druid.segment.SegmentWrangler;
@@ -70,9 +71,9 @@ public class InlineInputSliceReader implements InputSliceReader
for (final Segment segment : segmentWrangler.getSegmentsForIntervals(dataSource, Intervals.ONLY_ETERNITY)) {
segments.add(
- AdaptedLoadableSegment.create(
+ AdaptedLoadableSegment.fromUnmanagedSegment(
segment,
- Intervals.ETERNITY,
+ new SegmentDescriptor(Intervals.ETERNITY, "0", 0),
"inline data",
counters.channel(CounterNames.inputChannel(inputNumber))
)
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/input/lookup/LookupInputSliceReader.java b/multi-stage-query/src/main/java/org/apache/druid/msq/input/lookup/LookupInputSliceReader.java
index a9ee5db9e70..0eeea6f6ff0 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/input/lookup/LookupInputSliceReader.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/input/lookup/LookupInputSliceReader.java
@@ -32,6 +32,7 @@ import org.apache.druid.msq.input.LoadableSegment;
import org.apache.druid.msq.input.PhysicalInputSlice;
import org.apache.druid.msq.input.stage.ReadablePartitions;
import org.apache.druid.query.LookupDataSource;
+import org.apache.druid.query.SegmentDescriptor;
import org.apache.druid.segment.Segment;
import org.apache.druid.segment.SegmentWrangler;
import org.apache.druid.utils.CloseableUtils;
@@ -87,9 +88,9 @@ public class LookupInputSliceReader implements InputSliceReader
throw new ISE("Lookup[%s] has multiple segments; cannot read", lookupName);
}
- final LoadableSegment loadableSegment = AdaptedLoadableSegment.create(
+ final LoadableSegment loadableSegment = AdaptedLoadableSegment.fromUnmanagedSegment(
segment,
- Intervals.ETERNITY,
+ new SegmentDescriptor(Intervals.ETERNITY, "0", 0),
StringUtils.format("lookup[%s]", lookupName),
counters.channel(CounterNames.inputChannel(inputNumber))
);
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java
index 788da66072d..742ef6a974c 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java
@@ -32,15 +32,12 @@ import org.apache.druid.frame.write.FrameWriterFactory;
import org.apache.druid.java.util.common.Unit;
import org.apache.druid.java.util.common.io.Closer;
import org.apache.druid.java.util.common.logger.Logger;
-import org.apache.druid.msq.counters.ChannelCounters;
import org.apache.druid.msq.exec.DataServerQueryHandler;
import org.apache.druid.msq.input.table.SegmentsInputSlice;
-import org.apache.druid.segment.PhysicalSegmentInspector;
import org.apache.druid.segment.ReferenceCountedSegmentProvider;
import org.apache.druid.segment.Segment;
import org.apache.druid.segment.SegmentMapFunction;
import org.apache.druid.segment.SegmentReference;
-import org.apache.druid.utils.CloseableUtils;
import java.io.IOException;
import java.util.Collections;
@@ -154,19 +151,6 @@ public abstract class BaseLeafFrameProcessor implements FrameProcessor<Object>
throw DruidException.defensive("Missing segmentReference[%s]", segmentHolder.getDescriptor());
}
- try {
- final ChannelCounters counters = segmentHolder.getInputCounters();
- if (counters != null) {
- // Attach a counters.addFile call to the closer, to ensure input metrics are updated.
- // Get row count prior to mapping, because mapped segments often do not provide PhysicalSegmentInspector.
- final int rowCount = getSegmentRowCount(segmentReference);
- closer.register(() -> counters.addFile(rowCount, 0));
- }
- }
- catch (Throwable e) {
- throw CloseableUtils.closeAndWrapInCatch(e, segmentReference);
- }
-
final Segment segment = closer.register(segmentReference.map(segmentMapFn)).getSegmentReference().orElse(null);
if (segment == null) {
throw DruidException.defensive("Missing segment[%s]", segmentHolder.getDescriptor());
@@ -187,17 +171,4 @@ public abstract class BaseLeafFrameProcessor implements FrameProcessor<Object>
segment.getId()
));
}
-
- /**
- * Helper to get the number of rows for a segment, using a {@link PhysicalSegmentInspector}. Returns 0 when the
- * number is unknown.
- */
- private int getSegmentRowCount(final SegmentReference segmentReference)
- {
- return segmentReference
- .getSegmentReference()
- .flatMap(segment -> Optional.ofNullable(segment.as(PhysicalSegmentInspector.class)))
- .map(PhysicalSegmentInspector::getNumRows)
- .orElse(0);
- }
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java
index b965e8babc2..c2fe66ed0f4 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java
@@ -27,7 +27,6 @@ import org.apache.druid.common.guava.FutureUtils;
import org.apache.druid.error.DruidException;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.java.util.common.logger.Logger;
-import org.apache.druid.msq.counters.ChannelCounters;
import org.apache.druid.msq.exec.DataServerQueryHandler;
import org.apache.druid.msq.exec.std.StandardPartitionReader;
import org.apache.druid.msq.input.LoadableSegment;
@@ -136,7 +135,6 @@ public class ReadableInputQueue implements Closeable
if (cachedSegment.isPresent()) {
final SegmentReferenceHolder holder = new SegmentReferenceHolder(
new SegmentReference(loadableSegment.descriptor(), cachedSegment, null),
- loadableSegment.inputCounters(),
loadableSegment.description()
);
loadedSegments.add(holder);
@@ -299,17 +297,12 @@ public class ReadableInputQueue implements Closeable
// Transfer segment from "loadingSegments" to "loadedSegments" and return a reference to it.
if (loadingSegments.remove(acquireSegmentAction)) {
try {
- final ChannelCounters inputCounters = nextLoadableSegment.inputCounters();
- if (inputCounters != null) {
- inputCounters.addLoad(segment);
- }
final SegmentReferenceHolder referenceHolder = new SegmentReferenceHolder(
new SegmentReference(
nextLoadableSegment.descriptor(),
segment.getReferenceProvider().acquireReference(),
acquireSegmentAction // Release the hold when the SegmentReference is closed.
),
- nextLoadableSegment.inputCounters(),
nextLoadableSegment.description()
);
loadedSegments.add(referenceHolder);
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/SegmentReferenceHolder.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/SegmentReferenceHolder.java
index 71470eb620d..b6f38704240 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/SegmentReferenceHolder.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/SegmentReferenceHolder.java
@@ -20,7 +20,6 @@
package org.apache.druid.msq.querykit;
import com.google.common.base.Preconditions;
-import org.apache.druid.msq.counters.ChannelCounters;
import org.apache.druid.query.SegmentDescriptor;
import org.apache.druid.segment.SegmentReference;
@@ -35,20 +34,16 @@ public class SegmentReferenceHolder
{
private final AtomicReference<SegmentReference> segmentReference = new AtomicReference<>();
@Nullable
- private final ChannelCounters inputCounters;
- @Nullable
private final String description;
private final SegmentDescriptor descriptor;
public SegmentReferenceHolder(
SegmentReference segmentReference,
- @Nullable ChannelCounters inputCounters,
@Nullable String description
)
{
this.segmentReference.set(Preconditions.checkNotNull(segmentReference, "segmentReference"));
this.descriptor = Preconditions.checkNotNull(segmentReference, "segmentReference").getSegmentDescriptor();
- this.inputCounters = inputCounters;
this.description = description;
}
@@ -71,15 +66,6 @@ public class SegmentReferenceHolder
return descriptor;
}
- /**
- * Input counters that should be incremented as we read, or null if none should be incremented.
- */
- @Nullable
- public ChannelCounters getInputCounters()
- {
- return inputCounters;
- }
-
/**
* User-oriented description, suitable for inclusion in error messages.
*/
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/util/MultiStageQueryContext.java b/multi-stage-query/src/main/java/org/apache/druid/msq/util/MultiStageQueryContext.java
index 79b1c5ca341..d731ca074ae 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/util/MultiStageQueryContext.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/util/MultiStageQueryContext.java
@@ -171,6 +171,9 @@ public class MultiStageQueryContext
public static final String CTX_REMOVE_NULL_BYTES = "removeNullBytes";
public static final boolean DEFAULT_REMOVE_NULL_BYTES = false;
+ public static final String CTX_BACKGROUND_FETCH_EXTERNAL_FILES = "backgroundFetchExternalFiles";
+ public static final boolean DEFAULT_BACKGROUND_FETCH_EXTERNAL_FILES = true;
+
/**
* Hint to {@link StageProcessor} implementations about whether they should attempt to use
* {@link FrameCombiner} when doing sort-based aggregations.
@@ -473,6 +476,11 @@ public class MultiStageQueryContext
return queryContext.getBoolean(CTX_REMOVE_NULL_BYTES, DEFAULT_REMOVE_NULL_BYTES);
}
+ public static boolean isBackgroundFetchExternalFiles(final QueryContext queryContext)
+ {
+ return queryContext.getBoolean(CTX_BACKGROUND_FETCH_EXTERNAL_FILES, DEFAULT_BACKGROUND_FETCH_EXTERNAL_FILES);
+ }
+
public static boolean isUseCombiner(final QueryContext queryContext)
{
return queryContext.getBoolean(CTX_USE_COMBINER, DEFAULT_USE_COMBINER);
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartFrameContextTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartFrameContextTest.java
index 056115d87a2..d687be8e92b 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartFrameContextTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartFrameContextTest.java
@@ -135,6 +135,7 @@ public class DartFrameContextTest
null,
null,
null,
+ null,
processingBuffersSet,
null,
null,
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQCompactionTaskRunTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQCompactionTaskRunTest.java
index 895de0319da..e8c49f918d2 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQCompactionTaskRunTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQCompactionTaskRunTest.java
@@ -93,6 +93,7 @@ import org.apache.druid.segment.loading.DataSegmentPusher;
import org.apache.druid.segment.loading.LocalDataSegmentPusher;
import org.apache.druid.segment.loading.LocalDataSegmentPusherConfig;
import org.apache.druid.segment.loading.SegmentCacheManager;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.segment.nested.NestedDataComplexTypeSerde;
import org.apache.druid.segment.serde.ComplexMetrics;
import org.apache.druid.server.security.Escalator;
@@ -249,6 +250,7 @@ public class MSQCompactionTaskRunTest extends CompactionTaskRunBase
.toInstance(new ForwardingQueryProcessingPool(Execs.singleThreaded("Test-runner-processing-pool"))),
binder -> binder.bind(ObjectMapper.class).annotatedWith(Json.class).toInstance(objectMapper),
binder -> binder.bind(SegmentCacheManager.class).toInstance(segmentCacheManager),
+ binder -> binder.bind(VirtualStorageManager.class).toInstance(MSQTestBase.makeNilVirtualStorageManager()),
binder -> binder.bind(GroupingEngine.class).toInstance(groupingEngine)
);
injector = Guice.createInjector(modules);
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/IndexerFrameContextTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/IndexerFrameContextTest.java
index 6832198173b..560d10b11f6 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/IndexerFrameContextTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/IndexerFrameContextTest.java
@@ -135,6 +135,7 @@ public class IndexerFrameContextTest
null,
null,
null,
+ null,
processingBuffersSet,
null,
null,
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/input/RegularLoadableSegmentTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/input/RegularLoadableSegmentTest.java
index 4bf712811de..c24a1506ec2 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/input/RegularLoadableSegmentTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/input/RegularLoadableSegmentTest.java
@@ -55,6 +55,7 @@ import org.apache.druid.segment.loading.LoadSpec;
import org.apache.druid.segment.loading.SegmentLoaderConfig;
import org.apache.druid.segment.loading.SegmentLoadingException;
import org.apache.druid.segment.loading.SegmentLocalCacheManager;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
import org.apache.druid.segment.loading.StorageLocation;
import org.apache.druid.segment.loading.StorageLocationConfig;
import org.apache.druid.server.SegmentManager;
@@ -164,12 +165,13 @@ class RegularLoadableSegmentTest extends InitializedNullHandlingTest
cacheDir = tempDir.resolve("cache").toFile();
final SegmentLoaderConfig virtualLoaderConfig = new SegmentLoaderConfig()
.setLocations(ImmutableList.of(new StorageLocationConfig(cacheDir, 10_000_000_000L, null)))
- .setVirtualStorage(true, true);
+ .setVirtualStorage(true).setVirtualStorageIsEphemeral(true);
final List<StorageLocation> virtualLocations = virtualLoaderConfig.toStorageLocations();
segmentManagerDynamic = new SegmentManager(
new SegmentLocalCacheManager(
virtualLocations,
virtualLoaderConfig,
+ StorageLoadingThreadPool.createFromConfig(virtualLoaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(virtualLocations),
TestIndex.INDEX_IO,
jsonMapper
@@ -185,6 +187,7 @@ class RegularLoadableSegmentTest extends InitializedNullHandlingTest
new SegmentLocalCacheManager(
localLocations,
localLoaderConfig,
+ StorageLoadingThreadPool.createFromConfig(localLoaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(localLocations),
TestIndex.INDEX_IO,
jsonMapper
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessorTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessorTest.java
index 9922be0a6c2..8b617059082 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessorTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessorTest.java
@@ -133,7 +133,6 @@ public class ScanQueryFrameProcessorTest extends FrameProcessorTestBase
segmentReferenceProvider.acquireReference(),
null
),
- null,
null
)
),
@@ -396,7 +395,6 @@ public class ScanQueryFrameProcessorTest extends FrameProcessorTestBase
segmentReferenceProvider.acquireReference(),
null
),
- null,
null
)
),
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java
index f7947676ed7..c72efa72aa1 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java
@@ -93,6 +93,7 @@ import org.apache.druid.segment.column.RowSignature;
import org.apache.druid.segment.join.JoinConditionAnalysis;
import org.apache.druid.segment.join.JoinType;
import org.apache.druid.segment.join.JoinableFactoryWrapper;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.server.QueryResponse;
import org.apache.druid.server.QueryStackTests;
import org.apache.druid.server.SegmentManager;
@@ -222,6 +223,7 @@ public class MSQTaskQueryMakerTest
binder -> {
binder.bind(WireTransferableContext.class).toInstance(FrameTestUtil.WT_CONTEXT_LEGACY);
binder.bind(CoordinatorClient.class).to(NoopCoordinatorClient.class);
+ binder.bind(VirtualStorageManager.class).toInstance(MSQTestBase.makeNilVirtualStorageManager());
}
);
Injector injector = Guice.createInjector(defaultModule, BoundFieldModule.of(this));
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java b/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java
index 4d1fa384e2e..ede5f61662c 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/test/AbstractMSQComponentSupplierDelegate.java
@@ -49,7 +49,6 @@ public class AbstractMSQComponentSupplierDelegate extends QueryComponentSupplier
{
return DruidModuleCollection.of(
super.getCoreModule(),
- new MSQTestModule(),
new IndexingServiceTuningConfigModule(),
new JoinableFactoryModule(),
new MSQExternalDataSourceModule(),
@@ -59,6 +58,12 @@ public class AbstractMSQComponentSupplierDelegate extends QueryComponentSupplier
);
}
+ @Override
+ public DruidModule getOverrideModule()
+ {
+ return DruidModuleCollection.of(super.getOverrideModule(), new MSQTestModule());
+ }
+
@Override
public Class<? extends SqlEngine> getSqlEngineClass()
{
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteMSQTestsHelper.java b/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteMSQTestsHelper.java
index bef94919dcf..837a545db1d 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteMSQTestsHelper.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteMSQTestsHelper.java
@@ -95,7 +95,7 @@ public class CalciteMSQTestsHelper
public SegmentCacheManager provideSegmentCacheManager(ObjectMapper testMapper, TempDirProducer tempDirProducer)
{
return new SegmentCacheManagerFactory(TestIndex.INDEX_IO, testMapper)
- .manufacturate(tempDirProducer.newTempFolder("test"), true);
+ .manufacturate(tempDirProducer.newTempFolder("test"), null, true);
}
@Provides
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
index 10c3e8cbec3..65a9fb5b0e8 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
@@ -155,10 +155,14 @@ import org.apache.druid.segment.column.ColumnHolder;
import org.apache.druid.segment.column.RowSignature;
import org.apache.druid.segment.incremental.IncrementalIndexSchema;
import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStrategy;
import org.apache.druid.segment.loading.LocalDataSegmentPusher;
import org.apache.druid.segment.loading.LocalDataSegmentPusherConfig;
import org.apache.druid.segment.loading.LocalLoadSpec;
import org.apache.druid.segment.loading.SegmentCacheManager;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
+import org.apache.druid.segment.loading.external.StorageLocationVirtualStorageManager;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.segment.realtime.appenderator.AppenderatorsManager;
import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory;
import org.apache.druid.segment.writeout.SegmentWriteOutMediumFactory;
@@ -466,7 +470,7 @@ public class MSQTestBase extends BaseCalciteQueryTest
indexIO = new IndexIO(objectMapper, ColumnConfig.DEFAULT);
segmentCacheManager =
- new SegmentCacheManagerFactory(indexIO, objectMapper).manufacturate(newTempFolder("cacheManager"), true);
+ new SegmentCacheManagerFactory(indexIO, objectMapper).manufacturate(newTempFolder("cacheManager"), null, true);
testSegmentManager = new TestSegmentManager();
@@ -556,6 +560,13 @@ public class MSQTestBase extends BaseCalciteQueryTest
// Requirement of WorkerMemoryParameters.createProductionInstanceForWorker(injector)
binder -> binder.bind(AppenderatorsManager.class).toProvider(() -> null),
binder -> binder.bind(SegmentManager.class).toInstance(testSegmentManager.getSegmentManager()),
+ binder -> binder.bind(VirtualStorageManager.class).toInstance(
+ new StorageLocationVirtualStorageManager(
+ segmentCacheManager.getLocations(),
+ new LeastBytesUsedStorageLocationSelectorStrategy(segmentCacheManager.getLocations()),
+ segmentCacheManager.getLoadingThreadPool()
+ )
+ ),
new JoinableFactoryModule(),
new IndexingServiceTuningConfigModule(),
Modules.override(new MSQSqlModule()).with(
@@ -797,7 +808,7 @@ public class MSQTestBase extends BaseCalciteQueryTest
testSegmentManager.addSegment(dataSegment, segment);
acquiredSegment = testSegmentManager.getSegment(segmentId);
}
- return AdaptedLoadableSegment.create(acquiredSegment, descriptor.getInterval(), null, counters);
+ return AdaptedLoadableSegment.fromUnmanagedSegment(acquiredSegment, descriptor, null, counters);
}
public SelectTester testSelectQuery()
@@ -847,6 +858,18 @@ public class MSQTestBase extends BaseCalciteQueryTest
);
}
+ /**
+ * Creates an non-functional {@link VirtualStorageManager} suitable for tests.
+ */
+ public static VirtualStorageManager makeNilVirtualStorageManager()
+ {
+ return new StorageLocationVirtualStorageManager(
+ ImmutableList.of(),
+ new LeastBytesUsedStorageLocationSelectorStrategy(ImmutableList.of()),
+ StorageLoadingThreadPool.none()
+ );
+ }
+
private String runMultiStageQuery(
String query,
Map<String, Object> context,
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java
index 02ae7c56e1a..ef519108834 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java
@@ -54,6 +54,7 @@ import org.apache.druid.segment.column.ColumnConfig;
import org.apache.druid.segment.incremental.NoopRowIngestionMeters;
import org.apache.druid.segment.incremental.RowIngestionMeters;
import org.apache.druid.segment.loading.DataSegmentPusher;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory;
import org.apache.druid.server.DruidNode;
import org.apache.druid.server.SegmentManager;
@@ -258,6 +259,12 @@ public class MSQTestWorkerContext implements WorkerContext
return injector.getInstance(SegmentManager.class);
}
+ @Override
+ public VirtualStorageManager virtualStorageManager()
+ {
+ return injector.getInstance(VirtualStorageManager.class);
+ }
+
@Override
public CoordinatorClient coordinatorClient()
{
diff --git a/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResource.java b/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResource.java
new file mode 100644
index 00000000000..a92b9047784
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResource.java
@@ -0,0 +1,139 @@
+/*
+ * 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.common.asyncresource;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.collections.ResourceHolder;
+import org.apache.druid.error.DruidException;
+
+import java.io.Closeable;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+/**
+ * Represents a resource that requires some cleanup and that is acquired asynchronously. This class can in principle
+ * be used for resources that do not need cleanup or are not acquired asynchronously, but it is most useful when
+ * both are true. The wrapper generally owns the resource lifecycle; see "to consume a resource" below for details.
+ *
+ * <p>To produce a resource, generally you should create and populate {@link SettableAsyncResource}.
+ *
+ * <p>To consume a resource, use {@link #addReadyCallback(Runnable)}, {@link #await()}, or {@link #await(long)}
+ * to wait for the resource to become ready. Then use {@link #get()} to retrieve the resource. When you are done
+ * with the resource, call {@link #close()} on the {@code AsyncResource} object to close it. Do *not* close the
+ * resource {@code T} itself, even if it is {@link Closeable}, as this may lead to a double-close. The one
+ * exception to this rule is if you call {@link SettableAsyncResource#release()}: in this case you actually
+ * *must* call close on the resource {@code T} itself.
+ *
+ * <h3>Why not use Futures?</h3>
+ * Often {@link Future} or {@link ListenableFuture} are used for objects that are created asynchronously.
+ * These are, however, problematic when the object is a resource that requires cleanup. The biggest issue is
+ * handling cancellation. As soon as a caller gets a {@code Future<Closeable>}, it becomes responsible for
+ * closing the resource once it is created. If the caller is a query that could itself be canceled, it must
+ * still arrange for the resource to be closed.
+ *
+ * <p>The caller can do something like this to deal with it:
+ *
+ * <pre>
+ * // Upon query cancellation, attach a callback to the future that closes the resource once it becomes available.
+ * Futures.addCallback(
+ * resourceFuture,
+ * new FutureCallback<>() {
+ * void onSuccess(Closeable resource) { resource.close(); }
+ * void onFailure(Throwable t) { }
+ * }
+ * );
+ * </pre>
+ *
+ * But this is awkward, and doesn't allow resource acquisition to actually be canceled. Canceling the future isn't
+ * reliable, because it can lead to an orphaned resource: the asynchronous acquisition can complete in a race with
+ * cancellation, and in this case, the resource becomes eligible for GC without completing the future and therefore
+ * without being closed.
+ *
+ * <p>AsyncResource handles this problem by automatically closing the resource in
+ * {@link SettableAsyncResource#set(ResourceHolder)} when the {@link SettableAsyncResource} has been canceled.
+ */
+public interface AsyncResource<T> extends Closeable
+{
+ /**
+ * Whether resource acquisition has completed (successfully or with failure). To wait for this to become true
+ * asynchronously, use {@link #addReadyCallback(Runnable)}. To block until readiness, use {@link #await()}
+ * or {@link #await(long)}.
+ */
+ boolean isReady();
+
+ /**
+ * Register a callback to fire when {@link #isReady()} becomes true (whether the load succeeded or failed). If the
+ * holder is already ready, the callback fires immediately in the calling thread. Callbacks are not fired if
+ * {@link #close()} is called prior to the resource becoming available.
+ *
+ * <p>Throws {@link DruidException} if {@link #close()} has been called prior to this method.
+ */
+ void addReadyCallback(Runnable callback);
+
+ /**
+ * Retrieve the underlying object. May be called any number of times, and the same object will be returned.
+ *
+ * <p>Throws {@link DruidException} if the underlying object is not ready or if {@link #close()} has been called.
+ * Also throws an exception if the resource acquisition failed.
+ */
+ T get();
+
+ /**
+ * Block until {@link #isReady()} returns true. Does not close the resource if interrupted; callers must still
+ * call {@link #close()}.
+ *
+ * <p>Throws {@link DruidException} if {@link #close()} has been called prior to this method.
+ */
+ default T await() throws InterruptedException
+ {
+ final CountDownLatch latch = new CountDownLatch(1);
+ addReadyCallback(latch::countDown);
+ latch.await();
+ return get();
+ }
+
+ /**
+ * Block until {@link #isReady()} returns true, up to some timeout. Does not close the resource if interrupted
+ * or if waiting times out; callers must still call {@link #close()}.
+ *
+ * <p>Throws {@link DruidException} if {@link #close()} has been called prior to this method.
+ */
+ default T await(long timeoutMillis) throws InterruptedException, TimeoutException
+ {
+ final CountDownLatch latch = new CountDownLatch(1);
+ addReadyCallback(latch::countDown);
+ if (!latch.await(timeoutMillis, TimeUnit.MILLISECONDS)) {
+ throw new TimeoutException();
+ }
+ return get();
+ }
+
+ /**
+ * Closes the resource if it is ready, and has not been released by {@link SettableAsyncResource#release()}.
+ * If acquisition is still in progress, it is canceled if possible.
+ *
+ * <p>Despite {@link Closeable} requiring this method to be idempotent, it is not necessarily
+ * going to be idempotent. Do not close more than once.
+ */
+ @Override
+ void close();
+}
diff --git a/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java b/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java
new file mode 100644
index 00000000000..e22edd06750
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/common/asyncresource/AsyncResources.java
@@ -0,0 +1,114 @@
+/*
+ * 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.common.asyncresource;
+
+import org.apache.druid.collections.ResourceHolder;
+
+import java.io.Closeable;
+import java.util.List;
+import java.util.function.Function;
+
+/**
+ * Utility functions for {@link AsyncResource}.
+ */
+public class AsyncResources
+{
+ private AsyncResources()
+ {
+ // No instantiation.
+ }
+
+ /**
+ * Returns an {@link AsyncResource} that wraps the underlying object, which does not have a lifecycle.
+ */
+ public static <T> AsyncResource<T> unmanaged(final T object)
+ {
+ final SettableAsyncResource<T> retVal = new SettableAsyncResource<>();
+ retVal.set(object, null);
+ return retVal;
+ }
+
+ /**
+ * Returns an {@link AsyncResource} that wraps the underlying closeable object.
+ */
+ public static <T extends Closeable> AsyncResource<T> ofCloseable(final T object)
+ {
+ final SettableAsyncResource<T> retVal = new SettableAsyncResource<>();
+ retVal.set(ResourceHolder.fromCloseable(object));
+ return retVal;
+ }
+
+ /**
+ * Returns an {@link AsyncResource} that collects a list of underlying resources into a single lifecycle.
+ * Calling {@link AsyncResource#close()} on the returned async resource causes the underlying async resource
+ * to be closed.
+ *
+ * <p>The transformation generally happens eagerly in the thread that provides the source resource, so it is
+ * important that it run quickly.
+ *
+ * <p>The target of {@code function} need not be {@link Closeable}, and even if it is {@link Closeable}, it
+ * is not closed (only the source is closed). This transform utility is meant for transformations that do
+ * not introduce new resource lifecycles.
+ */
+ public static <T, R> AsyncResource<R> transform(
+ final AsyncResource<T> sourceResource,
+ final Function<T, R> function
+ )
+ {
+ return new TransformAsyncResource<>(sourceResource, function);
+ }
+
+ /**
+ * Returns an {@link AsyncResource} that collects a list of underlying resources into a single lifecycle.
+ * Calling {@link AsyncResource#close()} on the returned async resource causes the underlying async resources
+ * to also be closed.
+ */
+ public static <T> AsyncResource<List<T>> collect(final List<AsyncResource<T>> asyncResources)
+ {
+ return new CollectAsyncResource<>(asyncResources);
+ }
+
+ /**
+ * Returns an {@link AsyncResource} that recovers from an exception in {@code sourceResource}.
+ *
+ * <p>If the source resource suceeds, the recoverFn is not called and the underlying source resource is returned.
+ * On the other hand, if the source resource fails, the {@code recoverFn} is called with the exception and is
+ * given a chance to substitute a fallback value. Recovery generally happens eagerly in the thread that provides
+ * the source resource, so it is important that it run quickly.
+ *
+ * <p>When recovery happens, the {@code sourceResource} is closed immediately. Otherwise, the {@code sourceResoruce}
+ * is closed when the resource returned by this function is closed.
+ *
+ * <p>The target of {@code function} need not be {@link Closeable}, and even if it is {@link Closeable}, it
+ * is not closed (only the source is closed). This transform utility is meant for transformations that do
+ * not introduce new resource lifecycles.
+ *
+ * @param sourceResource the source resource
+ * @param recoverFn called when the source resource fails. Returns nonnull to recover, or null to keep the
+ * error condition as-is.
+ */
+ public static <T> AsyncResource<T> recover(
+ final AsyncResource<T> sourceResource,
+ final Function<Throwable, T> recoverFn
+ )
+ {
+ return new RecoverAsyncResource<>(sourceResource, recoverFn);
+ }
+}
diff --git a/processing/src/main/java/org/apache/druid/common/asyncresource/CollectAsyncResource.java b/processing/src/main/java/org/apache/druid/common/asyncresource/CollectAsyncResource.java
new file mode 100644
index 00000000000..3955e0b25a4
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/common/asyncresource/CollectAsyncResource.java
@@ -0,0 +1,98 @@
+/*
+ * 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.common.asyncresource;
+
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.utils.CloseableUtils;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Resource that collects a list of underlying resources into a single lifecycle.
+ *
+ * @see AsyncResources#collect(List) for more details
+ */
+public class CollectAsyncResource<T> implements AsyncResource<List<T>>
+{
+ private final List<AsyncResource<T>> sourceResources;
+ private final AtomicInteger readyCount = new AtomicInteger(0);
+ private final SettableAsyncResource<List<T>> targetResource = new SettableAsyncResource<>();
+
+ /**
+ * Constructor. Can also be created with {@link AsyncResources#collect(List)}.
+ */
+ CollectAsyncResource(final List<AsyncResource<T>> sourceResources)
+ {
+ this.sourceResources = sourceResources;
+
+ if (sourceResources.isEmpty()) {
+ targetResource.set(List.of(), null);
+ } else {
+ for (final AsyncResource<T> asyncResource : sourceResources) {
+ asyncResource.addReadyCallback(this::onOneSourceReady);
+ }
+ }
+ }
+
+ @Override
+ public boolean isReady()
+ {
+ return targetResource.isReady();
+ }
+
+ @Override
+ public void addReadyCallback(Runnable callback)
+ {
+ targetResource.addReadyCallback(callback);
+ }
+
+ @Override
+ public List<T> get()
+ {
+ return targetResource.get();
+ }
+
+ @Override
+ public void close()
+ {
+ final Closer closer = Closer.create();
+ closer.registerAll(sourceResources);
+ closer.register(targetResource);
+ CloseableUtils.closeAndWrapExceptions(closer);
+ }
+
+ private void onOneSourceReady()
+ {
+ if (readyCount.incrementAndGet() == sourceResources.size()) {
+ try {
+ final List<T> resources = new ArrayList<>(sourceResources.size());
+ for (final AsyncResource<T> asyncResource : sourceResources) {
+ resources.add(asyncResource.get());
+ }
+ targetResource.set(resources, null); // no Closeable here, since we own asyncResources ourselves
+ }
+ catch (Throwable e) {
+ targetResource.setException(e);
+ }
+ }
+ }
+}
diff --git a/processing/src/main/java/org/apache/druid/common/asyncresource/RecoverAsyncResource.java b/processing/src/main/java/org/apache/druid/common/asyncresource/RecoverAsyncResource.java
new file mode 100644
index 00000000000..a83a70fd74f
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/common/asyncresource/RecoverAsyncResource.java
@@ -0,0 +1,124 @@
+/*
+ * 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.common.asyncresource;
+
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.utils.CloseableUtils;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
+
+/**
+ * Resource that passes through an underlying resource, but can substitute a fallback on failure.
+ *
+ * @see AsyncResources#recover(AsyncResource, Function) for more details
+ */
+public class RecoverAsyncResource<T> implements AsyncResource<T>
+{
+ private static final Logger log = new Logger(RecoverAsyncResource.class);
+
+ private final AsyncResource<T> sourceResource;
+ private final Function<Throwable, T> recovery;
+ private final SettableAsyncResource<T> targetResource = new SettableAsyncResource<>();
+
+ /**
+ * Whether {@link #sourceResource} has been closed. The source is closed either after recovery (so its resources
+ * are released before handing the recovered object to the caller), or otherwise in {@link #close()}.
+ */
+ private final AtomicBoolean sourceClosed = new AtomicBoolean(false);
+
+ /**
+ * Constructor. Can also be created with {@link AsyncResources#recover(AsyncResource, Function)}.
+ */
+ RecoverAsyncResource(
+ final AsyncResource<T> sourceResource,
+ final Function<Throwable, T> recovery
+ )
+ {
+ this.sourceResource = sourceResource;
+ this.recovery = recovery;
+ sourceResource.addReadyCallback(this::onSourceReady);
+ }
+
+ @Override
+ public boolean isReady()
+ {
+ return targetResource.isReady();
+ }
+
+ @Override
+ public T get()
+ {
+ return targetResource.get();
+ }
+
+ @Override
+ public void addReadyCallback(Runnable callback)
+ {
+ targetResource.addReadyCallback(callback);
+ }
+
+ @Override
+ public void close()
+ {
+ final Closer closer = Closer.create();
+ if (sourceClosed.compareAndSet(false, true)) {
+ closer.register(sourceResource);
+ }
+ closer.register(targetResource);
+ CloseableUtils.closeAndWrapExceptions(closer);
+ }
+
+ private void onSourceReady()
+ {
+ final T value;
+ try {
+ value = sourceResource.get();
+ }
+ catch (Throwable e) {
+ final T recovered;
+ try {
+ recovered = recovery.apply(e);
+ }
+ catch (Throwable e2) {
+ e.addSuppressed(e2);
+ targetResource.setException(e);
+ return;
+ }
+
+ if (recovered != null) {
+ // Release the source's resources before handing back the recovered object.
+ if (sourceClosed.compareAndSet(false, true)) {
+ CloseableUtils.closeAndSuppressExceptions(
+ sourceResource,
+ ex -> log.warn(ex, "Failed to close source resource during recovery")
+ );
+ }
+ targetResource.set(recovered, null);
+ } else {
+ targetResource.setException(e);
+ }
+ return;
+ }
+
+ targetResource.set(value, null);
+ }
+}
diff --git a/processing/src/main/java/org/apache/druid/common/asyncresource/SettableAsyncResource.java b/processing/src/main/java/org/apache/druid/common/asyncresource/SettableAsyncResource.java
new file mode 100644
index 00000000000..9894d16a311
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/common/asyncresource/SettableAsyncResource.java
@@ -0,0 +1,335 @@
+/*
+ * 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.common.asyncresource;
+
+import com.google.common.base.Throwables;
+import com.google.common.util.concurrent.SettableFuture;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import org.apache.druid.collections.ResourceHolder;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.java.util.common.Either;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.segment.AsyncCursorHolder;
+import org.apache.druid.utils.CloseableUtils;
+
+import javax.annotation.Nullable;
+import java.io.Closeable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
+
+/**
+ * Basic utility that allows creating {@link AsyncResource} wrappers. Analogous to JDK {@link CompletableFuture}
+ * or Guava {@link SettableFuture}. See {@link AsyncResource} for details about why you would want to use this
+ * instead of {@link Future}.
+ *
+ * <p>In addition to {@link #get()}, there is also {@link #release()}. Releasing is not allowed by instances
+ * of this class, but may be allowed by subclasses. When releasing is allowed, it enables an ownership-transfer
+ * model. See {@link #release()} for more information on this model.
+ */
+public class SettableAsyncResource<T> implements AsyncResource<T>
+{
+ private static final Logger LOG = new Logger(SettableAsyncResource.class);
+
+ /**
+ * Whether {@link #release()} is allowed.
+ */
+ private final boolean releasable;
+
+ /**
+ * Callbacks provided by {@link #addReadyCallback(Runnable)}.
+ */
+ @GuardedBy("this")
+ private final List<Runnable> readyCallbacks = new ArrayList<>();
+
+ /**
+ * Canceler provided by {@link #setCanceler(Runnable)}. Ideally would be {@code GuardeBy("this")}, but isn't
+ * annotated because errorprone doesn't like the reference to canceler::run escaping.
+ */
+ @Nullable
+ private Runnable canceler;
+
+ /**
+ * Result set by {@link #setInternal(Either)}.
+ */
+ @Nullable
+ @GuardedBy("this")
+ private ResourceHolder<T> result = null;
+
+ /**
+ * Error set by {@link #setInternal(Either)}.
+ */
+ @Nullable
+ @GuardedBy("this")
+ private Throwable error = null;
+
+ @GuardedBy("this")
+ private State state = State.NEW;
+
+ /**
+ * Constructor.
+ */
+ public SettableAsyncResource()
+ {
+ this(false);
+ }
+
+ /**
+ * Constructor for subclasses that allow {@link #release()}. They should provide "true" here and then
+ * also override {@link #release()} to change its visibility from protected to public.
+ */
+ protected SettableAsyncResource(boolean releasable)
+ {
+ this.releasable = releasable;
+ }
+
+ /**
+ * Set a canceler that will be called from {@link #close()} if this {@link AsyncResource} is closed prior
+ * to the resource being available. Calling this method has no effect if {@link #close()} has already been called
+ * or if the resource is already available.
+ */
+ public synchronized void setCanceler(Runnable newCanceler)
+ {
+ if (canceler != null) {
+ throw DruidException.defensive("canceler already set, cannot call setCanceler()");
+ }
+
+ if (state == State.NEW) {
+ this.canceler = newCanceler;
+ }
+ }
+
+ /**
+ * Provides a resource and closer for the resource. Transitions the {@link AsyncResource} into a "ready" state,
+ * where {@link #isReady()} returns true, if it has not yet been closed. Returns true if this transition happened
+ * successfully, false otherwise.
+ *
+ * <p>If this method returns true, it also fires all the callbacks that were registered via
+ * {@link #addReadyCallback(Runnable)}. Once this method returns true, {@link #close()} will no longer call
+ * the canceler from {@link #setCanceler(Runnable)}.
+ *
+ * <p>If this method returns false, the producer is responsible for closing the resource itself.
+ *
+ * <p>Throws {@link DruidException} if this resource was already completed from a prior call to this method or
+ * {@link #setException}).
+ */
+ public boolean set(T object, @Nullable Closeable closer)
+ {
+ if (object == null) {
+ throw DruidException.defensive("object cannot be null");
+ }
+ final ResourceHolder<T> resourceHolder = new ResourceHolder<>()
+ {
+ @Override
+ public T get()
+ {
+ return object;
+ }
+
+ @Override
+ public void close()
+ {
+ CloseableUtils.closeAndWrapExceptions(closer);
+ }
+ };
+ return setInternal(Either.value(resourceHolder));
+ }
+
+ /**
+ * Variant of {@link #set(Object, Closeable)} for callers that have a {@link ResourceHolder}.
+ */
+ public boolean set(ResourceHolder<T> holder)
+ {
+ return set(holder.get(), holder);
+ }
+
+ /**
+ * Provides an exception for a resource that failed to load. Transitions the {@link AsyncResource} into a "ready"
+ * state, where {@link #isReady()} returns true, if it has not yet been closed.
+ *
+ * <p>If this method successfully transitions to "ready", it also fires all the callbacks that were registered via
+ * {@link #addReadyCallback(Runnable)}. Afterwards, {@link #close()} will no longer call the canceler from
+ * {@link #setCanceler(Runnable)}.
+ *
+ * <p>Throws {@link DruidException} if this resource was already completed from a prior call to this method or
+ * {@link #set}).
+ */
+ public void setException(Throwable t)
+ {
+ setInternal(Either.error(t));
+ }
+
+ @Override
+ public synchronized boolean isReady()
+ {
+ return state == State.READY;
+ }
+
+ @Override
+ public synchronized T get()
+ {
+ return switch (state) {
+ case NEW -> throw DruidException.defensive("Not ready yet");
+ case READY -> {
+ if (error != null) {
+ Throwables.throwIfUnchecked(error);
+ throw DruidException.forPersona(DruidException.Persona.DEVELOPER)
+ .ofCategory(DruidException.Category.UNCATEGORIZED)
+ .build(error, error.getMessage());
+ } else {
+ yield result.get();
+ }
+ }
+ case RELEASED -> throw DruidException.defensive("Resource has been released");
+ case CLOSED -> throw DruidException.defensive("Closed");
+ };
+ }
+
+ /**
+ * Take ownership of the underlying object. After this returns, {@link #close()} on this
+ * {@link AsyncResource} is a no-op; the caller is responsible for closing the returned {@code T}. Useful when
+ * passing the resource to something else that prefers to take full ownership of it.
+ *
+ * <p>This method enables a resource-transfer model. It is exposed by certain subclasses, such as
+ * {@link AsyncCursorHolder}, where the resource is itself {@link Closeable} and where that close method
+ * encapsulates all necessary resource releasing logic. It is provided for convenience of certain callers,
+ * although note that you must avoid releasing if you intend to use combinators such as
+ * {@link AsyncResources#collect}, {@link AsyncResources#transform}, etc. These combinators will fail to properly
+ * encapsulate resource lifecycle if resources have been released.
+ *
+ * <p>Throws {@link DruidException} if the holder is not yet ready, has already been released, or if
+ * {@link #close()} has been called.
+ */
+ protected synchronized T release()
+ {
+ if (!releasable) {
+ throw DruidException.defensive("Not releasable");
+ }
+
+ final T object = get();
+ // Clear result to allow GC.
+ result = null;
+ state = State.RELEASED;
+ return object;
+ }
+
+ @Override
+ public void addReadyCallback(Runnable callback)
+ {
+ final boolean fireImmediately;
+ synchronized (this) {
+ switch (state) {
+ case NEW -> {
+ readyCallbacks.add(callback);
+ fireImmediately = false;
+ }
+ case READY -> fireImmediately = true;
+ default -> throw DruidException.defensive("Cannot addReadyCallback in state[%s]", state);
+ }
+ }
+ if (fireImmediately) {
+ callback.run();
+ }
+ }
+
+ @Override
+ public void close()
+ {
+ final Closeable deferredCloseable;
+
+ synchronized (this) {
+ deferredCloseable = switch (state) {
+ case NEW -> canceler != null ? canceler::run : null;
+ case READY -> result;
+ case RELEASED -> null;
+ default -> throw DruidException.defensive("Already closed");
+ };
+
+ // Clear result and canceler to allow GC.
+ result = null;
+ canceler = null;
+ state = State.CLOSED;
+ }
+
+ CloseableUtils.closeAndSuppressExceptions(
+ deferredCloseable,
+ e -> LOG.warn(e, "Failed to call cleaner of class[%s]", deferredCloseable.getClass())
+ );
+ }
+
+ @GuardedBy("this")
+ private List<Runnable> drainCallbacks()
+ {
+ final List<Runnable> snapshot = List.copyOf(readyCallbacks);
+ readyCallbacks.clear();
+ return snapshot;
+ }
+
+ private boolean setInternal(Either<Throwable, ResourceHolder<T>> value)
+ {
+ final boolean didSet;
+ final List<Runnable> callbacksToFire;
+
+ synchronized (this) {
+ didSet = switch (state) {
+ case NEW -> {
+ if (value.isError()) {
+ error = value.error();
+ } else {
+ result = value.valueOrThrow();
+ }
+
+ state = State.READY;
+ yield true;
+ }
+ case READY, RELEASED -> throw DruidException.defensive("Already complete, cannot call set/setException again");
+ case CLOSED -> false;
+ };
+
+ // Clear canceler to allow GC.
+ canceler = null;
+ callbacksToFire = drainCallbacks();
+ }
+ fireCallbacks(callbacksToFire);
+ return didSet;
+ }
+
+ private static void fireCallbacks(List<Runnable> callbacks)
+ {
+ for (final Runnable callback : callbacks) {
+ try {
+ callback.run();
+ }
+ catch (Throwable t) {
+ // Best-effort; one bad callback shouldn't break others.
+ LOG.warn(t, "callback exception");
+ }
+ }
+ }
+
+ enum State
+ {
+ NEW,
+ READY,
+ RELEASED,
+ CLOSED
+ }
+}
diff --git a/processing/src/main/java/org/apache/druid/common/asyncresource/TransformAsyncResource.java b/processing/src/main/java/org/apache/druid/common/asyncresource/TransformAsyncResource.java
new file mode 100644
index 00000000000..8f6bcf6bcee
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/common/asyncresource/TransformAsyncResource.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.common.asyncresource;
+
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.utils.CloseableUtils;
+
+import java.util.function.Function;
+
+/**
+ * Resource that is the result of calling a function on an underlying resource.
+ *
+ * @see AsyncResources#transform(AsyncResource, Function) for more details
+ */
+public class TransformAsyncResource<T, R> implements AsyncResource<R>
+{
+ private final AsyncResource<T> sourceResource;
+ private final Function<T, R> function;
+ private final SettableAsyncResource<R> targetResource;
+
+ /**
+ * Constructor. Can also be created with {@link AsyncResources#transform(AsyncResource, Function)}.
+ */
+ TransformAsyncResource(final AsyncResource<T> sourceResource, final Function<T, R> function)
+ {
+ this.sourceResource = sourceResource;
+ this.function = function;
+ this.targetResource = new SettableAsyncResource<>();
+ sourceResource.addReadyCallback(this::onSourceReady);
+ }
+
+ @Override
+ public boolean isReady()
+ {
+ return targetResource.isReady();
+ }
+
+ @Override
+ public R get()
+ {
+ return targetResource.get();
+ }
+
+ @Override
+ public void addReadyCallback(Runnable callback)
+ {
+ targetResource.addReadyCallback(callback);
+ }
+
+ @Override
+ public void close()
+ {
+ final Closer closer = Closer.create();
+ closer.register(sourceResource);
+ closer.register(targetResource);
+ CloseableUtils.closeAndWrapExceptions(closer);
+ }
+
+ private void onSourceReady()
+ {
+ try {
+ final R target = function.apply(sourceResource.get());
+ targetResource.set(target, null); // no Closeable here, since we own sourceResource ourselves
+ }
+ catch (Throwable e) {
+ targetResource.setException(e);
+ }
+ }
+}
diff --git a/processing/src/main/java/org/apache/druid/data/input/BytesCountingInputEntity.java b/processing/src/main/java/org/apache/druid/data/input/BytesCountingInputEntity.java
index 21485bfa02f..c1b4739bb81 100644
--- a/processing/src/main/java/org/apache/druid/data/input/BytesCountingInputEntity.java
+++ b/processing/src/main/java/org/apache/druid/data/input/BytesCountingInputEntity.java
@@ -57,6 +57,12 @@ public class BytesCountingInputEntity implements InputEntity
return new BytesCountingInputStream(baseInputEntity.open(), inputStats);
}
+ @Override
+ public InputStream openRaw() throws IOException
+ {
+ return new BytesCountingInputStream(baseInputEntity.openRaw(), inputStats);
+ }
+
public InputEntity getBaseInputEntity()
{
return baseInputEntity;
diff --git a/processing/src/main/java/org/apache/druid/data/input/InputEntity.java b/processing/src/main/java/org/apache/druid/data/input/InputEntity.java
index d6ea7211573..f990a94de8f 100644
--- a/processing/src/main/java/org/apache/druid/data/input/InputEntity.java
+++ b/processing/src/main/java/org/apache/druid/data/input/InputEntity.java
@@ -63,13 +63,22 @@ public interface InputEntity
URI getUri();
/**
- * Opens an {@link InputStream} on the input entity directly.
+ * Opens an {@link InputStream} on the input entity, decompressing the raw bytes if the entity is compressed.
* This is the basic way to read the given entity.
* This method may be called multiple times to re-read the data from the entity.
*
* @see #fetch
*/
- InputStream open() throws IOException;
+ default InputStream open() throws IOException
+ {
+ // Implementations that decompress must override this method.
+ return openRaw();
+ }
+
+ /**
+ * Opens an {@link InputStream} on the raw bytes of the input entity, without any decompression.
+ */
+ InputStream openRaw() throws IOException;
/**
* Fetches the input entity into the local storage.
diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java b/processing/src/main/java/org/apache/druid/data/input/InputFilePointer.java
similarity index 57%
copy from server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java
copy to processing/src/main/java/org/apache/druid/data/input/InputFilePointer.java
index bd9c8c65def..651eee96849 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java
+++ b/processing/src/main/java/org/apache/druid/data/input/InputFilePointer.java
@@ -17,27 +17,21 @@
* under the License.
*/
-package org.apache.druid.segment.loading;
+package org.apache.druid.data.input;
-import org.junit.Assert;
-import org.junit.Test;
+import org.apache.druid.io.FilePopulator;
-public class SegmentLoaderConfigTest
-{
- @Test
- public void testSetVirtualStorage()
- {
- final SegmentLoaderConfig config = new SegmentLoaderConfig();
-
- // Verify default values
- Assert.assertFalse(config.isVirtualStorage());
- Assert.assertFalse(config.isVirtualStorageEphemeral());
-
- // Set both to true
- config.setVirtualStorage(true, true);
+import java.net.URI;
+import java.util.function.LongSupplier;
- // Verify both fields are set
- Assert.assertTrue(config.isVirtualStorage());
- Assert.assertTrue(config.isVirtualStorageEphemeral());
- }
+/**
+ * Pointer to a file that may be on remote storage.
+ *
+ * @param uri URI of the file
+ * @param sizeSupplier supplier for the size of the file. If the file is compressed, this is the compressed size
+ * @param populator populator that knows how to fetch the file. If the file is compressed, this fetches the
+ * compressed file
+ */
+public record InputFilePointer(URI uri, LongSupplier sizeSupplier, FilePopulator populator)
+{
}
diff --git a/processing/src/main/java/org/apache/druid/data/input/InputSource.java b/processing/src/main/java/org/apache/druid/data/input/InputSource.java
index be815742be1..23a16409482 100644
--- a/processing/src/main/java/org/apache/druid/data/input/InputSource.java
+++ b/processing/src/main/java/org/apache/druid/data/input/InputSource.java
@@ -33,6 +33,7 @@ import org.apache.druid.java.util.common.UOE;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.File;
+import java.util.List;
import java.util.Set;
/**
@@ -78,6 +79,15 @@ public interface InputSource
*/
boolean needsFormat();
+ /**
+ * If this input source is representable as a set of files, return pointers to those files.
+ */
+ @Nullable
+ default List<InputFilePointer> asFilePointers()
+ {
+ return null;
+ }
+
/**
* Creates an {@link InputSourceReader}.
*
diff --git a/processing/src/main/java/org/apache/druid/data/input/RetryingInputEntity.java b/processing/src/main/java/org/apache/druid/data/input/RetryingInputEntity.java
index ac2481e989a..c91abb07324 100644
--- a/processing/src/main/java/org/apache/druid/data/input/RetryingInputEntity.java
+++ b/processing/src/main/java/org/apache/druid/data/input/RetryingInputEntity.java
@@ -36,13 +36,21 @@ public abstract class RetryingInputEntity implements InputEntity
@Override
public InputStream open() throws IOException
{
- RetryingInputStream<?> retryingInputStream = new RetryingInputStream<>(
+ return CompressionUtils.decompress(openRaw(), getPath());
+ }
+
+ /**
+ * Opens a raw {@link InputStream} on the entity, without decompression.
+ */
+ @Override
+ public InputStream openRaw() throws IOException
+ {
+ return new RetryingInputStream<>(
this,
new RetryingInputEntityOpenFunction(),
getRetryCondition(),
getMaxRetries()
);
- return CompressionUtils.decompress(retryingInputStream, getPath());
}
// override this in sub-classes to customize retries
diff --git a/processing/src/main/java/org/apache/druid/data/input/impl/ByteEntity.java b/processing/src/main/java/org/apache/druid/data/input/impl/ByteEntity.java
index db4f3396ae4..3e7307d0f14 100644
--- a/processing/src/main/java/org/apache/druid/data/input/impl/ByteEntity.java
+++ b/processing/src/main/java/org/apache/druid/data/input/impl/ByteEntity.java
@@ -63,7 +63,7 @@ public class ByteEntity implements InputEntity
}
@Override
- public InputStream open()
+ public InputStream openRaw()
{
return new ByteBufferInputStream(buffer.duplicate());
}
diff --git a/processing/src/main/java/org/apache/druid/data/input/impl/CloudObjectInputSource.java b/processing/src/main/java/org/apache/druid/data/input/impl/CloudObjectInputSource.java
index 6e9587ca0d7..b047f812810 100644
--- a/processing/src/main/java/org/apache/druid/data/input/impl/CloudObjectInputSource.java
+++ b/processing/src/main/java/org/apache/druid/data/input/impl/CloudObjectInputSource.java
@@ -29,6 +29,7 @@ import org.apache.druid.data.input.AbstractInputSource;
import org.apache.druid.data.input.FilePerSplitHintSpec;
import org.apache.druid.data.input.InputEntity;
import org.apache.druid.data.input.InputFileAttribute;
+import org.apache.druid.data.input.InputFilePointer;
import org.apache.druid.data.input.InputFormat;
import org.apache.druid.data.input.InputRowSchema;
import org.apache.druid.data.input.InputSourceReader;
@@ -39,6 +40,7 @@ import org.apache.druid.data.input.impl.systemfield.SystemFieldDecoratorFactory;
import org.apache.druid.data.input.impl.systemfield.SystemFieldInputSource;
import org.apache.druid.data.input.impl.systemfield.SystemFields;
import org.apache.druid.java.util.common.CloseableIterators;
+import org.apache.druid.java.util.common.FileUtils;
import org.apache.druid.utils.CollectionUtils;
import org.apache.druid.utils.Streams;
@@ -49,6 +51,7 @@ import java.net.URI;
import java.nio.file.FileSystems;
import java.nio.file.PathMatcher;
import java.nio.file.Paths;
+import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
@@ -211,6 +214,66 @@ public abstract class CloudObjectInputSource
.iterator();
}
+ @Nullable
+ @Override
+ public List<InputFilePointer> asFilePointers()
+ {
+ if (!systemFields.getFields().isEmpty()) {
+ // System fields cannot be added to file pointers.
+ return null;
+ }
+
+ final List<CloudObjectLocation> locations;
+ if (!CollectionUtils.isNullOrEmpty(objects)) {
+ locations = objects;
+ } else if (!CollectionUtils.isNullOrEmpty(uris)) {
+ locations = Lists.transform(uris, CloudObjectLocation::new);
+ } else {
+ // Only uris + objects can be expressed as file pointers. Prefixes could be if we listed them here, but
+ // we'd rather avoid doing that, since this method is likely to be called in a thread that should not
+ // be doing blocking I/O.
+ return null;
+ }
+
+ final PathMatcher globMatcher =
+ StringUtils.isNotBlank(objectGlob)
+ ? FileSystems.getDefault().getPathMatcher("glob:" + objectGlob)
+ : null;
+
+ final CloudObjectSplitWidget splitWidget = getSplitWidget();
+ final List<InputFilePointer> pointers = new ArrayList<>();
+ for (final CloudObjectLocation location : locations) {
+ if (globMatcher != null && !globMatcher.matches(Paths.get(location.getPath()))) {
+ continue;
+ }
+
+ final InputEntity entity = createEntity(location);
+ final URI uri = location.toUri(scheme);
+ pointers.add(
+ new InputFilePointer(
+ uri,
+ () -> {
+ try {
+ return splitWidget.getObjectSize(location);
+ }
+ catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ },
+ dstFile -> FileUtils.copyLarge(
+ entity::openRaw,
+ dstFile,
+ new byte[InputEntity.DEFAULT_FETCH_BUFFER_SIZE],
+ entity.getRetryCondition(),
+ InputEntity.DEFAULT_MAX_NUM_FETCH_TRIES,
+ "Failed to fetch [" + uri + "]"
+ )
+ )
+ );
+ }
+ return pointers;
+ }
+
@Override
public boolean equals(Object o)
{
diff --git a/processing/src/main/java/org/apache/druid/data/input/impl/FileEntity.java b/processing/src/main/java/org/apache/druid/data/input/impl/FileEntity.java
index 01762a9e879..617fd895563 100644
--- a/processing/src/main/java/org/apache/druid/data/input/impl/FileEntity.java
+++ b/processing/src/main/java/org/apache/druid/data/input/impl/FileEntity.java
@@ -38,8 +38,14 @@ public class FileEntity implements InputEntity
}
@Override
- public CleanableFile fetch(File temporaryDirectory, byte[] fetchBuffer)
+ public CleanableFile fetch(File temporaryDirectory, byte[] fetchBuffer) throws IOException
{
+ if (CompressionUtils.Format.fromFileName(file.getName()) != null) {
+ // The file appears to be compressed. Decompress it into temporaryDirectory.
+ return InputEntity.super.fetch(temporaryDirectory, fetchBuffer);
+ }
+
+ // The file is not compressed, so there is no need to copy it.
return new CleanableFile()
{
@Override
@@ -72,4 +78,10 @@ public class FileEntity implements InputEntity
{
return CompressionUtils.decompress(new FileInputStream(file), file.getName());
}
+
+ @Override
+ public InputStream openRaw() throws IOException
+ {
+ return new FileInputStream(file);
+ }
}
diff --git a/processing/src/main/java/org/apache/druid/data/input/impl/LocalInputSource.java b/processing/src/main/java/org/apache/druid/data/input/impl/LocalInputSource.java
index 365a59afecf..981e23766f1 100644
--- a/processing/src/main/java/org/apache/druid/data/input/impl/LocalInputSource.java
+++ b/processing/src/main/java/org/apache/druid/data/input/impl/LocalInputSource.java
@@ -37,6 +37,7 @@ import org.apache.commons.io.filefilter.WildcardFileFilter;
import org.apache.druid.data.input.AbstractInputSource;
import org.apache.druid.data.input.InputEntity;
import org.apache.druid.data.input.InputFileAttribute;
+import org.apache.druid.data.input.InputFilePointer;
import org.apache.druid.data.input.InputFormat;
import org.apache.druid.data.input.InputRowSchema;
import org.apache.druid.data.input.InputSourceReader;
@@ -290,6 +291,15 @@ public class LocalInputSource
);
}
+ @Nullable
+ @Override
+ public List<InputFilePointer> asFilePointers()
+ {
+ // Intentionally not implemented, because this method is used for fetching external files to a local
+ // cache, and there is no point in copying local files from one place to another.
+ return null;
+ }
+
@Override
public boolean equals(Object o)
{
diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java b/processing/src/main/java/org/apache/druid/io/FilePopulator.java
similarity index 57%
copy from server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java
copy to processing/src/main/java/org/apache/druid/io/FilePopulator.java
index bd9c8c65def..f292a441443 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java
+++ b/processing/src/main/java/org/apache/druid/io/FilePopulator.java
@@ -17,27 +17,23 @@
* under the License.
*/
-package org.apache.druid.segment.loading;
+package org.apache.druid.io;
-import org.junit.Assert;
-import org.junit.Test;
+import java.io.File;
+import java.io.IOException;
-public class SegmentLoaderConfigTest
+/**
+ * Lambda interface for populating a file with downloaded content.
+ */
+@FunctionalInterface
+public interface FilePopulator
{
- @Test
- public void testSetVirtualStorage()
- {
- final SegmentLoaderConfig config = new SegmentLoaderConfig();
-
- // Verify default values
- Assert.assertFalse(config.isVirtualStorage());
- Assert.assertFalse(config.isVirtualStorageEphemeral());
-
- // Set both to true
- config.setVirtualStorage(true, true);
-
- // Verify both fields are set
- Assert.assertTrue(config.isVirtualStorage());
- Assert.assertTrue(config.isVirtualStorageEphemeral());
- }
+ /**
+ * Populate the given file with content.
+ * The file's parent directory is guaranteed to exist.
+ *
+ * @param file The file to populate
+ * @throws IOException if population fails
+ */
+ void populate(File file) throws IOException;
}
diff --git a/processing/src/main/java/org/apache/druid/segment/AsyncCursorHolder.java b/processing/src/main/java/org/apache/druid/segment/AsyncCursorHolder.java
index e682f111619..d980022d325 100644
--- a/processing/src/main/java/org/apache/druid/segment/AsyncCursorHolder.java
+++ b/processing/src/main/java/org/apache/druid/segment/AsyncCursorHolder.java
@@ -19,37 +19,19 @@
package org.apache.druid.segment;
-import com.google.errorprone.annotations.concurrent.GuardedBy;
+import org.apache.druid.collections.ResourceHolder;
+import org.apache.druid.common.asyncresource.SettableAsyncResource;
import org.apache.druid.error.DruidException;
import org.apache.druid.frame.processor.ReturnOrAwait;
-import org.apache.druid.java.util.common.Either;
-import org.apache.druid.java.util.common.logger.Logger;
import javax.annotation.Nullable;
-import java.io.Closeable;
-import java.util.ArrayList;
-import java.util.List;
/**
* Closeable wrapper around an asynchronously-loaded {@link CursorHolder}, returned by
* {@link CursorFactory#makeCursorHolderAsync}. Designed to make lifecycle management safe even when the holder is
* still loading: callers receive a single Closeable handle and can register it once with their cleanup machinery,
* regardless of where the underlying load is in its lifecycle.
- * <p>
- * The hazard this exists to avoid: returning a {@code ListenableFuture<CursorHolder>} (or similar future-of-Closeable)
- * makes correct cleanup error-prone, where canceling the future or letting a caller fail before consuming the future
- * can orphan the produced holder, leaking the underlying resources. By exposing a Closeable that internally tracks the
- * load and disposes whatever has materialized, callers don't have to write that bookkeeping themselves.
- * <p>
- * <h3>Producer protocol</h3>
- * Producers feed results in via {@link #set(CursorHolder)} or {@link #setException(Throwable)}, both of which return
- * a boolean. If they return {@code false}, this wrapper has already been closed and the producer is responsible for
- * closing whatever holder it just produced.
- * Producers may pass a {@link Runnable} canceler at construction time which runs on {@link #close()} when the wrapper
- * is closed before the {@link #set} has been called, giving the producer an opportunity to abort its work. The canceler
- * is best-effort: a producer may have already produced the holder by the time it observes cancellation, in which case
- * its {@link #set} call will return false and it must close the holder it tried to set.
- * <p>
+ *
* <h3>Consumer protocol</h3>
* Consumers wait for {@link #isReady()} via {@link #addReadyCallback}, and {@link #release()} to transfer ownership of
* the {@link CursorHolder} (or throw the producer exception). Calling {@link #release()} before {@link #isReady()}
@@ -70,36 +52,8 @@ import java.util.List;
* // ... use holder; close it when done (or hand it to a component that owns its lifecycle) ...
* }</pre>
*/
-public class AsyncCursorHolder implements Closeable
+public class AsyncCursorHolder extends SettableAsyncResource<CursorHolder>
{
- private static final Logger LOG = new Logger(AsyncCursorHolder.class);
-
- /**
- * Completed {@link AsyncCursorHolder} backed by an already available {@link CursorHolder}
- */
- public static AsyncCursorHolder completed(CursorHolder holder)
- {
- final AsyncCursorHolder result = new AsyncCursorHolder(null);
- result.set(holder);
- return result;
- }
-
- @Nullable
- private final Runnable canceler;
-
- @GuardedBy("this")
- @Nullable
- private CursorHolder result = null;
- @GuardedBy("this")
- @Nullable
- private Throwable error = null;
- @GuardedBy("this")
- private boolean closed = false;
- @GuardedBy("this")
- private boolean disposed = false;
- @GuardedBy("this")
- private final List<Runnable> readyCallbacks = new ArrayList<>();
-
/**
* @param canceler optional callback invoked from {@link #close()} when the wrapper is closed before the load has
* completed ({@link #set} or {@link #setException}). Producers that support cancellation should
@@ -108,207 +62,31 @@ public class AsyncCursorHolder implements Closeable
*/
public AsyncCursorHolder(@Nullable Runnable canceler)
{
- this.canceler = canceler;
- }
-
- /**
- * Allows producer to mark the load successful with the given holder. Returns {@code true} if accepted, {@code false}
- * if this wrapper has already been closed, in which case the producer is responsible for closing {@link CursorHolder}
- * itself. Throws {@link DruidException} if the load was already completed (from prior calls to this method or
- * {@link #setException}).
- * <p>
- * Callbacks registered via {@link #addReadyCallback} fire outside the lock to avoid lock-ordering deadlocks and
- * unbounded lock holds.
- */
- public boolean set(CursorHolder holder)
- {
- if (holder == null) {
- throw DruidException.defensive("CursorHolder cannot be null");
- }
- return setInternal(Either.value(holder));
- }
-
- /**
- * Allows producer to mark the load as failed. Returns {@code true} if accepted, {@code false} if this wrapper has
- * already been closed (no holder was produced, so there's nothing for the producer to clean up). Throws
- * {@link DruidException} if the load was already completed (from prior calls to this method or {@link #set}).
- * <p>
- * Callbacks registered via {@link #addReadyCallback} fire outside the lock to avoid lock-ordering deadlocks and
- * unbounded lock holds.
- */
- public boolean setException(Throwable t)
- {
- return setInternal(Either.error(t));
- }
-
- private boolean setInternal(Either<Throwable, CursorHolder> value)
- {
- final List<Runnable> callbacksToFire;
- synchronized (this) {
- if (closed) {
- return false;
- }
- if (result != null || error != null) {
- throw DruidException.defensive("AsyncCursorHolder is already completed");
- }
- if (value.isError()) {
- error = value.error();
- } else {
- result = value.valueOrThrow();
- }
- callbacksToFire = drainCallbacks();
- }
- fireCallbacks(callbacksToFire);
- return true;
+ super(true);
+ setCanceler(canceler);
}
/**
- * Whether the load has completed (successfully or with failure). Once true, stays true. Callers that need to wait
- * for readiness without blocking the current thread should register a {@link #addReadyCallback} and yield.
+ * Convenience setter.
*/
- public synchronized boolean isReady()
+ public boolean set(CursorHolder cursorHolder)
{
- return result != null || error != null;
+ return super.set(cursorHolder, cursorHolder);
}
- /**
- * Take ownership of the underlying {@link CursorHolder}. After this returns, {@link #close()} on this
- * {@code AsyncCursorHolder} is a no-op; the caller is responsible for closing the returned holder. Useful when
- * passing the holder to another component (e.g. a cursor-lifecycle manager) that takes ownership of it.
- * <p>
- * Throws {@link DruidException} if the holder is not yet ready, has already been released, or this wrapper
- * has been closed. Wraps and rethrows the failure if the underlying load failed. Does not block; callers must
- * check {@link #isReady()} first (typically by yielding via a {@link #addReadyCallback}-driven wait pattern).
- */
+ @Override // Overridden to change access from protected to public
public synchronized CursorHolder release()
{
- if (closed) {
- throw DruidException.defensive("AsyncCursorHolder is already closed");
- }
- if (disposed) {
- throw DruidException.defensive("AsyncCursorHolder has already been released");
- }
- if (error != null) {
- // pass through as is
- if (error instanceof RuntimeException runtime) {
- throw runtime;
- } else if (error instanceof Error e) {
- throw e;
- }
- throw DruidException.forPersona(DruidException.Persona.DEVELOPER)
- .ofCategory(DruidException.Category.UNCATEGORIZED)
- .build(error, error.getMessage());
- }
- if (result == null) {
- throw DruidException.defensive("AsyncCursorHolder is not ready yet");
- }
- final CursorHolder resultToReturn = result;
- // clear result so it can be eligible for gc
- result = null;
- disposed = true;
- return resultToReturn;
+ return super.release();
}
/**
- * Register a callback to fire when {@link #isReady()} becomes true (whether the load succeeded or failed). If the
- * holder is already ready, the callback fires synchronously on the calling thread. Otherwise it fires on whatever
- * thread invokes {@link #set} or {@link #setException}, outside the wrapper's lock so the callback may safely
- * re-enter the wrapper. Multiple callbacks may be registered; each fires once.
- */
- public void addReadyCallback(Runnable callback)
- {
- final boolean fireImmediately;
- synchronized (this) {
- if (result != null || error != null) {
- fireImmediately = true;
- } else {
- readyCallbacks.add(callback);
- fireImmediately = false;
- }
- }
- if (fireImmediately) {
- callback.run();
- }
- }
-
- /**
- * Close the wrapper. Safe at any lifecycle point and idempotent:
- * <ul>
- * <li>Already-loaded: closes the underlying {@link CursorHolder} immediately.</li>
- * <li>Loading in progress: invokes the canceler (if one was supplied at construction). The producer may still
- * call {@link #set} / {@link #setException} after this; if the producer wins the race and calls {@code set}
- * with a holder, {@code set} returns false and the producer is responsible for closing it.</li>
- * <li>Load failed: no-op (no holder was produced).</li>
- * <li>Already released: no-op.</li>
- * <li>Already closed: throws {@link DruidException}.</li>
- * </ul>
+ * Completed {@link AsyncCursorHolder} backed by an already available {@link CursorHolder}
*/
- @Override
- public void close()
- {
- final CursorHolder holderToClose;
- final Runnable cancelerToRun;
- synchronized (this) {
- if (closed) {
- throw DruidException.defensive("Already closed");
- }
- closed = true;
- if (disposed) {
- // Ownership was already transferred via release(); the caller manages the holder lifecycle.
- return;
- }
- if (result != null) {
- // Result is here and no one has released it; we close it.
- disposed = true;
- holderToClose = result;
- cancelerToRun = null;
- } else if (error != null) {
- // Load already failed; nothing to dispose.
- holderToClose = null;
- cancelerToRun = null;
- } else {
- // Load not yet completed; signal cancellation to the producer (if any).
- holderToClose = null;
- cancelerToRun = canceler;
- }
- }
- if (holderToClose != null) {
- try {
- holderToClose.close();
- }
- catch (Throwable ignored) {
- // Best-effort cleanup
- }
- }
- if (cancelerToRun != null) {
- try {
- cancelerToRun.run();
- }
- catch (Throwable t) {
- // Best-effort cancel
- LOG.warn(t, "AsyncCursorHolder canceler exception");
- }
- }
- }
-
- @GuardedBy("this")
- private List<Runnable> drainCallbacks()
- {
- final List<Runnable> snapshot = List.copyOf(readyCallbacks);
- readyCallbacks.clear();
- return snapshot;
- }
-
- private static void fireCallbacks(List<Runnable> callbacks)
+ public static AsyncCursorHolder completed(CursorHolder holder)
{
- for (Runnable cb : callbacks) {
- try {
- cb.run();
- }
- catch (Throwable t) {
- // Best-effort; one bad callback shouldn't break others.
- LOG.warn(t, "AsyncCursorHolder callback exception");
- }
- }
+ final AsyncCursorHolder result = new AsyncCursorHolder(null);
+ result.set(ResourceHolder.fromCloseable(holder));
+ return result;
}
}
diff --git a/processing/src/main/java/org/apache/druid/segment/incremental/NoopRowIngestionMeters.java b/processing/src/main/java/org/apache/druid/segment/incremental/NoopRowIngestionMeters.java
index 67eb80a2dc8..6b7786d5605 100644
--- a/processing/src/main/java/org/apache/druid/segment/incremental/NoopRowIngestionMeters.java
+++ b/processing/src/main/java/org/apache/druid/segment/incremental/NoopRowIngestionMeters.java
@@ -23,9 +23,8 @@ import java.util.Collections;
import java.util.Map;
/**
- * This class is used only in {@code DartFrameContext}.
- *
- * Consider using {@link RowIngestionMetersFactory} instead.
+ * This class is used when stats are not needed. Consider using {@link RowIngestionMetersFactory} instead
+ * when stats are needed.
*/
public class NoopRowIngestionMeters implements RowIngestionMeters
{
diff --git a/processing/src/test/java/org/apache/druid/common/asyncresource/CollectAsyncResourceTest.java b/processing/src/test/java/org/apache/druid/common/asyncresource/CollectAsyncResourceTest.java
new file mode 100644
index 00000000000..40ce3e9b4bc
--- /dev/null
+++ b/processing/src/test/java/org/apache/druid/common/asyncresource/CollectAsyncResourceTest.java
@@ -0,0 +1,131 @@
+/*
+ * 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.common.asyncresource;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class CollectAsyncResourceTest
+{
+ @Test
+ public void testEmptyListCompletesImmediatelyWithEmptyResult()
+ {
+ final List<AsyncResource<String>> sources = List.of();
+ final AsyncResource<List<String>> collected = AsyncResources.collect(sources);
+
+ Assertions.assertTrue(collected.isReady(), "collect of an empty list must be immediately ready");
+ Assertions.assertEquals(List.of(), collected.get());
+ collected.close();
+ }
+
+ @Test
+ public void testReadyOnlyAfterAllSourcesReadyAndPreservesOrder()
+ {
+ final SettableAsyncResource<String> a = new SettableAsyncResource<>();
+ final SettableAsyncResource<String> b = new SettableAsyncResource<>();
+ final SettableAsyncResource<String> c = new SettableAsyncResource<>();
+ final List<AsyncResource<String>> sources = List.of(a, b, c);
+ final AsyncResource<List<String>> collected = AsyncResources.collect(sources);
+
+ Assertions.assertFalse(collected.isReady());
+
+ // Complete out of order; the result must still be in source order.
+ c.set("c", null);
+ Assertions.assertFalse(collected.isReady());
+ a.set("a", null);
+ Assertions.assertFalse(collected.isReady());
+ b.set("b", null);
+
+ Assertions.assertTrue(collected.isReady());
+ Assertions.assertEquals(List.of("a", "b", "c"), collected.get());
+ collected.close();
+ }
+
+ @Test
+ public void testAlreadyReadySourcesCompleteOnConstruction()
+ {
+ final SettableAsyncResource<String> a = new SettableAsyncResource<>();
+ final SettableAsyncResource<String> b = new SettableAsyncResource<>();
+ a.set("a", null);
+ b.set("b", null);
+
+ final List<AsyncResource<String>> sources = List.of(a, b);
+ final AsyncResource<List<String>> collected = AsyncResources.collect(sources);
+
+ Assertions.assertTrue(collected.isReady());
+ Assertions.assertEquals(List.of("a", "b"), collected.get());
+ collected.close();
+ }
+
+ @Test
+ public void testSourceFailurePropagates()
+ {
+ final SettableAsyncResource<String> a = new SettableAsyncResource<>();
+ final SettableAsyncResource<String> b = new SettableAsyncResource<>();
+ final List<AsyncResource<String>> sources = List.of(a, b);
+ final AsyncResource<List<String>> collected = AsyncResources.collect(sources);
+
+ final RuntimeException failure = new RuntimeException("boom");
+ a.set("a", null);
+ b.setException(failure);
+
+ Assertions.assertTrue(collected.isReady());
+ final RuntimeException thrown = Assertions.assertThrows(RuntimeException.class, collected::get);
+ Assertions.assertSame(failure, thrown);
+ collected.close();
+ }
+
+ @Test
+ public void testCloseClosesAllReadySources()
+ {
+ final AtomicInteger aClose = new AtomicInteger();
+ final AtomicInteger bClose = new AtomicInteger();
+ final SettableAsyncResource<String> a = new SettableAsyncResource<>();
+ final SettableAsyncResource<String> b = new SettableAsyncResource<>();
+ a.set("a", aClose::incrementAndGet);
+ b.set("b", bClose::incrementAndGet);
+
+ final List<AsyncResource<String>> sources = List.of(a, b);
+ final AsyncResource<List<String>> collected = AsyncResources.collect(sources);
+ Assertions.assertTrue(collected.isReady());
+
+ collected.close();
+ Assertions.assertEquals(1, aClose.get());
+ Assertions.assertEquals(1, bClose.get());
+ }
+
+ @Test
+ public void testCloseCancelsPendingSources()
+ {
+ final AtomicInteger aCancel = new AtomicInteger();
+ final SettableAsyncResource<String> a = new SettableAsyncResource<>();
+ a.setCanceler(aCancel::incrementAndGet);
+
+ final List<AsyncResource<String>> sources = List.of(a);
+ final AsyncResource<List<String>> collected = AsyncResources.collect(sources);
+ Assertions.assertFalse(collected.isReady());
+
+ collected.close();
+ Assertions.assertEquals(1, aCancel.get(), "closing the collect must cancel pending sources");
+ }
+}
diff --git a/processing/src/test/java/org/apache/druid/common/asyncresource/RecoverAsyncResourceTest.java b/processing/src/test/java/org/apache/druid/common/asyncresource/RecoverAsyncResourceTest.java
new file mode 100644
index 00000000000..9a3c695de61
--- /dev/null
+++ b/processing/src/test/java/org/apache/druid/common/asyncresource/RecoverAsyncResourceTest.java
@@ -0,0 +1,184 @@
+/*
+ * 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.common.asyncresource;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+public class RecoverAsyncResourceTest
+{
+ @Test
+ public void testSourceSuccessPassesThroughAndDoesNotRecover()
+ {
+ final TrackingAsyncResource<String> source = new TrackingAsyncResource<>();
+ final AtomicInteger recoveryCalls = new AtomicInteger();
+
+ final AsyncResource<String> recovering = AsyncResources.recover(
+ source,
+ e -> {
+ recoveryCalls.incrementAndGet();
+ return "fallback";
+ }
+ );
+
+ source.delegate.set("value", null);
+
+ Assertions.assertTrue(recovering.isReady());
+ Assertions.assertEquals("value", recovering.get());
+ Assertions.assertEquals(0, recoveryCalls.get(), "recovery must not be called on success");
+ Assertions.assertEquals(0, source.closeCount.get(), "source must not be closed before close()");
+
+ recovering.close();
+ Assertions.assertEquals(1, source.closeCount.get(), "close() must close the source");
+ }
+
+ @Test
+ public void testSourceFailureRecoversWithFallbackValue()
+ {
+ final TrackingAsyncResource<String> source = new TrackingAsyncResource<>();
+ final AtomicReference<Throwable> recoveryArg = new AtomicReference<>();
+
+ final AsyncResource<String> recovering = AsyncResources.recover(
+ source,
+ e -> {
+ recoveryArg.set(e);
+ return "fallback";
+ }
+ );
+
+ final RuntimeException failure = new RuntimeException("boom");
+ source.delegate.setException(failure);
+
+ Assertions.assertTrue(recovering.isReady());
+ Assertions.assertEquals("fallback", recovering.get());
+ Assertions.assertSame(failure, recoveryArg.get(), "recovery receives the source's error");
+ Assertions.assertEquals(
+ 1,
+ source.closeCount.get(),
+ "source must be closed eagerly during recovery, before the wrapper is closed"
+ );
+
+ // Closing the wrapper does not double-close the source.
+ recovering.close();
+ Assertions.assertEquals(1, source.closeCount.get());
+ }
+
+ @Test
+ public void testSourceFailureWithNullRecoveryPropagatesError()
+ {
+ final TrackingAsyncResource<String> source = new TrackingAsyncResource<>();
+
+ final AsyncResource<String> recovering = AsyncResources.recover(
+ source,
+ e -> null
+ );
+
+ final RuntimeException failure = new RuntimeException("boom");
+ source.delegate.setException(failure);
+
+ Assertions.assertTrue(recovering.isReady());
+ final RuntimeException thrown = Assertions.assertThrows(RuntimeException.class, recovering::get);
+ Assertions.assertSame(failure, thrown);
+ Assertions.assertEquals(0, source.closeCount.get(), "source must not be closed until close() when not recovering");
+
+ recovering.close();
+ Assertions.assertEquals(1, source.closeCount.get());
+ }
+
+ @Test
+ public void testRecoveryThrowsPropagatesWithOriginalSuppressed()
+ {
+ final TrackingAsyncResource<String> source = new TrackingAsyncResource<>();
+ final RuntimeException recoveryFailure = new RuntimeException("recovery failed");
+
+ final AsyncResource<String> recovering = AsyncResources.recover(
+ source,
+ e -> {
+ throw recoveryFailure;
+ }
+ );
+
+ final RuntimeException original = new RuntimeException("boom");
+ source.delegate.setException(original);
+
+ Assertions.assertTrue(recovering.isReady());
+ final RuntimeException thrown = Assertions.assertThrows(RuntimeException.class, recovering::get);
+ Assertions.assertSame(original, thrown);
+ Assertions.assertEquals(1, thrown.getSuppressed().length);
+ Assertions.assertSame(recoveryFailure, thrown.getSuppressed()[0]);
+ }
+
+ @Test
+ public void testCloseBeforeReadyClosesSource()
+ {
+ final TrackingAsyncResource<String> source = new TrackingAsyncResource<>();
+
+ final AsyncResource<String> recovering = AsyncResources.recover(
+ source,
+ e -> "fallback"
+ );
+
+ // Close the wrapper while the source is still pending.
+ recovering.close();
+ Assertions.assertEquals(1, source.closeCount.get());
+ Assertions.assertFalse(recovering.isReady());
+
+ // Completing the source after close is a no-op and must not throw.
+ Assertions.assertFalse(source.delegate.set("late", null));
+ }
+
+ /**
+ * An {@link AsyncResource} that delegates to a {@link SettableAsyncResource} and counts {@link #close()} calls,
+ * so tests can verify how the source resource's lifecycle is managed.
+ */
+ private static class TrackingAsyncResource<T> implements AsyncResource<T>
+ {
+ private final SettableAsyncResource<T> delegate = new SettableAsyncResource<>();
+ private final AtomicInteger closeCount = new AtomicInteger();
+
+ @Override
+ public boolean isReady()
+ {
+ return delegate.isReady();
+ }
+
+ @Override
+ public void addReadyCallback(Runnable callback)
+ {
+ delegate.addReadyCallback(callback);
+ }
+
+ @Override
+ public T get()
+ {
+ return delegate.get();
+ }
+
+ @Override
+ public void close()
+ {
+ closeCount.incrementAndGet();
+ delegate.close();
+ }
+ }
+}
diff --git a/processing/src/test/java/org/apache/druid/common/asyncresource/SettableAsyncResourceTest.java b/processing/src/test/java/org/apache/druid/common/asyncresource/SettableAsyncResourceTest.java
new file mode 100644
index 00000000000..43def411a36
--- /dev/null
+++ b/processing/src/test/java/org/apache/druid/common/asyncresource/SettableAsyncResourceTest.java
@@ -0,0 +1,276 @@
+/*
+ * 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.common.asyncresource;
+
+import org.apache.druid.collections.ResourceHolder;
+import org.apache.druid.error.DruidException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class SettableAsyncResourceTest
+{
+ @Test
+ public void testSetThenGetAndReady()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ Assertions.assertFalse(resource.isReady());
+
+ Assertions.assertTrue(resource.set("value", null));
+ Assertions.assertTrue(resource.isReady());
+ Assertions.assertEquals("value", resource.get());
+ }
+
+ @Test
+ public void testGetBeforeReadyThrows()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ Assertions.assertThrows(DruidException.class, resource::get);
+ }
+
+ @Test
+ public void testSetNullObjectThrows()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ Assertions.assertThrows(DruidException.class, () -> resource.set(null, null));
+ }
+
+ @Test
+ public void testCloseClosesTheResourceExactlyOnce()
+ {
+ final AtomicInteger closeCount = new AtomicInteger();
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.set("value", closeCount::incrementAndGet);
+
+ resource.close();
+ Assertions.assertEquals(1, closeCount.get());
+ }
+
+ @Test
+ public void testSetWithResourceHolder()
+ {
+ final AtomicInteger closeCount = new AtomicInteger();
+ final TrackedCloseable closeable = new TrackedCloseable("value", closeCount);
+ final SettableAsyncResource<TrackedCloseable> resource = new SettableAsyncResource<>();
+ Assertions.assertTrue(resource.set(ResourceHolder.fromCloseable(closeable)));
+
+ Assertions.assertSame(closeable, resource.get());
+ resource.close();
+ Assertions.assertEquals(1, closeCount.get());
+ }
+
+ @Test
+ public void testDoubleSetThrows()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.set("value", null);
+ Assertions.assertThrows(DruidException.class, () -> resource.set("again", null));
+ }
+
+ @Test
+ public void testSetAfterExceptionThrows()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.setException(new RuntimeException("boom"));
+ Assertions.assertThrows(DruidException.class, () -> resource.set("value", null));
+ }
+
+ @Test
+ public void testSetExceptionUncheckedRethrownAsIs()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ final RuntimeException failure = new RuntimeException("boom");
+ resource.setException(failure);
+
+ Assertions.assertTrue(resource.isReady());
+ final RuntimeException thrown = Assertions.assertThrows(RuntimeException.class, resource::get);
+ Assertions.assertSame(failure, thrown);
+ }
+
+ @Test
+ public void testSetExceptionCheckedWrappedInDruidException()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.setException(new IOException("boom"));
+
+ Assertions.assertTrue(resource.isReady());
+ Assertions.assertThrows(DruidException.class, resource::get);
+ }
+
+ @Test
+ public void testReadyCallbackFiresImmediatelyWhenAlreadyReady()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.set("value", null);
+
+ final AtomicInteger fired = new AtomicInteger();
+ resource.addReadyCallback(fired::incrementAndGet);
+ Assertions.assertEquals(1, fired.get());
+ }
+
+ @Test
+ public void testReadyCallbackFiresOnSet()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+
+ final AtomicInteger fired = new AtomicInteger();
+ resource.addReadyCallback(fired::incrementAndGet);
+ Assertions.assertEquals(0, fired.get());
+
+ resource.set("value", null);
+ Assertions.assertEquals(1, fired.get());
+ }
+
+ @Test
+ public void testReadyCallbackFiresOnSetException()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+
+ final AtomicInteger fired = new AtomicInteger();
+ resource.addReadyCallback(fired::incrementAndGet);
+ resource.setException(new RuntimeException("boom"));
+ Assertions.assertEquals(1, fired.get());
+ }
+
+ @Test
+ public void testAddReadyCallbackAfterCloseThrows()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.close();
+ Assertions.assertThrows(DruidException.class, () -> resource.addReadyCallback(() -> {}));
+ }
+
+ @Test
+ public void testCloseInNewStateRunsCanceler()
+ {
+ final AtomicInteger canceled = new AtomicInteger();
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.setCanceler(canceled::incrementAndGet);
+
+ resource.close();
+ Assertions.assertEquals(1, canceled.get());
+ }
+
+ @Test
+ public void testCancelerNotRunOnceReady()
+ {
+ final AtomicInteger canceled = new AtomicInteger();
+ final AtomicInteger closeCount = new AtomicInteger();
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.setCanceler(canceled::incrementAndGet);
+
+ // Becoming ready clears the canceler; close() must then close the resource, not cancel.
+ resource.set("value", closeCount::incrementAndGet);
+ resource.close();
+
+ Assertions.assertEquals(0, canceled.get(), "canceler must not run once the resource is ready");
+ Assertions.assertEquals(1, closeCount.get());
+ }
+
+ @Test
+ public void testSetCancelerTwiceThrows()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.setCanceler(() -> {});
+ Assertions.assertThrows(DruidException.class, () -> resource.setCanceler(() -> {}));
+ }
+
+ @Test
+ public void testSetAfterCloseReturnsFalseAndDoesNotCloseTheResource()
+ {
+ final AtomicInteger closeCount = new AtomicInteger();
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.close();
+
+ // Producer racing with close: set() returns false, and the resource is NOT closed by the wrapper, so the
+ // producer remains responsible for closing it.
+ Assertions.assertFalse(resource.set("value", closeCount::incrementAndGet));
+ Assertions.assertEquals(0, closeCount.get());
+ }
+
+ @Test
+ public void testDoubleCloseThrows()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.set("value", null);
+ resource.close();
+ Assertions.assertThrows(DruidException.class, resource::close);
+ }
+
+ @Test
+ public void testAwaitReturnsValueWhenReady() throws InterruptedException
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.set("value", null);
+ Assertions.assertEquals("value", resource.await());
+ }
+
+ @Test
+ public void testAwaitWithTimeoutReturnsValueWhenReady() throws Exception
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ resource.set("value", null);
+ Assertions.assertEquals("value", resource.await(1000));
+ }
+
+ @Test
+ public void testAwaitWithTimeoutThrowsWhenNotReady()
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ Assertions.assertThrows(TimeoutException.class, () -> resource.await(10));
+ }
+
+ @Test
+ public void testAwaitWakesUpWhenSetFromAnotherThread() throws Exception
+ {
+ final SettableAsyncResource<String> resource = new SettableAsyncResource<>();
+ final Thread setter = new Thread(() -> resource.set("value", null));
+ setter.start();
+ Assertions.assertEquals("value", resource.await());
+ setter.join();
+ }
+
+ private static class TrackedCloseable implements java.io.Closeable
+ {
+ private final String value;
+ private final AtomicInteger closeCount;
+
+ TrackedCloseable(final String value, final AtomicInteger closeCount)
+ {
+ this.value = value;
+ this.closeCount = closeCount;
+ }
+
+ @Override
+ public String toString()
+ {
+ return value;
+ }
+
+ @Override
+ public void close()
+ {
+ closeCount.incrementAndGet();
+ }
+ }
+}
diff --git a/processing/src/test/java/org/apache/druid/common/asyncresource/TransformAsyncResourceTest.java b/processing/src/test/java/org/apache/druid/common/asyncresource/TransformAsyncResourceTest.java
new file mode 100644
index 00000000000..2a066eef7fc
--- /dev/null
+++ b/processing/src/test/java/org/apache/druid/common/asyncresource/TransformAsyncResourceTest.java
@@ -0,0 +1,156 @@
+/*
+ * 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.common.asyncresource;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.io.Closeable;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class TransformAsyncResourceTest
+{
+ @Test
+ public void testTransformAppliesFunction()
+ {
+ final SettableAsyncResource<Integer> source = new SettableAsyncResource<>();
+ final AsyncResource<String> transformed = AsyncResources.transform(source, i -> "v" + i);
+
+ Assertions.assertFalse(transformed.isReady());
+ source.set(42, null);
+
+ Assertions.assertTrue(transformed.isReady());
+ Assertions.assertEquals("v42", transformed.get());
+ transformed.close();
+ }
+
+ @Test
+ public void testTransformOfAlreadyReadySourceFiresImmediately()
+ {
+ final SettableAsyncResource<Integer> source = new SettableAsyncResource<>();
+ source.set(7, null);
+
+ final AsyncResource<String> transformed = AsyncResources.transform(source, i -> "v" + i);
+ Assertions.assertTrue(transformed.isReady());
+ Assertions.assertEquals("v7", transformed.get());
+ transformed.close();
+ }
+
+ @Test
+ public void testFunctionCalledOnceAndLazily()
+ {
+ final AtomicInteger calls = new AtomicInteger();
+ final SettableAsyncResource<Integer> source = new SettableAsyncResource<>();
+ final AsyncResource<String> transformed = AsyncResources.transform(
+ source,
+ i -> {
+ calls.incrementAndGet();
+ return "v" + i;
+ }
+ );
+
+ Assertions.assertEquals(0, calls.get(), "function must not run before the source is ready");
+ source.set(1, null);
+ Assertions.assertEquals(1, calls.get());
+
+ // Repeated get() must not re-run the function.
+ transformed.get();
+ transformed.get();
+ Assertions.assertEquals(1, calls.get());
+ transformed.close();
+ }
+
+ @Test
+ public void testSourceFailurePropagatesAndFunctionNotCalled()
+ {
+ final AtomicInteger calls = new AtomicInteger();
+ final SettableAsyncResource<Integer> source = new SettableAsyncResource<>();
+ final AsyncResource<String> transformed = AsyncResources.transform(
+ source,
+ i -> {
+ calls.incrementAndGet();
+ return "v" + i;
+ }
+ );
+
+ final RuntimeException failure = new RuntimeException("boom");
+ source.setException(failure);
+
+ Assertions.assertTrue(transformed.isReady());
+ final RuntimeException thrown = Assertions.assertThrows(RuntimeException.class, transformed::get);
+ Assertions.assertSame(failure, thrown);
+ Assertions.assertEquals(0, calls.get(), "function must not run when the source failed");
+ transformed.close();
+ }
+
+ @Test
+ public void testFunctionThrowsPropagates()
+ {
+ final RuntimeException failure = new RuntimeException("function boom");
+ final SettableAsyncResource<Integer> source = new SettableAsyncResource<>();
+ final AsyncResource<String> transformed = AsyncResources.transform(
+ source,
+ i -> {
+ throw failure;
+ }
+ );
+
+ source.set(1, null);
+
+ Assertions.assertTrue(transformed.isReady());
+ final RuntimeException thrown = Assertions.assertThrows(RuntimeException.class, transformed::get);
+ Assertions.assertSame(failure, thrown);
+ transformed.close();
+ }
+
+ @Test
+ public void testCloseClosesSourceButNotTransformedValue()
+ {
+ final AtomicInteger sourceClose = new AtomicInteger();
+ final AtomicInteger targetClose = new AtomicInteger();
+ final SettableAsyncResource<Integer> source = new SettableAsyncResource<>();
+
+ // The function's result is Closeable, but transform must NOT manage its lifecycle (it only closes the source).
+ final Closeable target = targetClose::incrementAndGet;
+ final AsyncResource<Closeable> transformed = AsyncResources.transform(source, i -> target);
+
+ source.set(1, sourceClose::incrementAndGet);
+ Assertions.assertTrue(transformed.isReady());
+ Assertions.assertSame(target, transformed.get());
+
+ transformed.close();
+ Assertions.assertEquals(1, sourceClose.get(), "the source resource must be closed");
+ Assertions.assertEquals(0, targetClose.get(), "the transformed value must not be closed by the transform");
+ }
+
+ @Test
+ public void testCloseBeforeReadyCancelsSource()
+ {
+ final AtomicInteger sourceCancel = new AtomicInteger();
+ final SettableAsyncResource<Integer> source = new SettableAsyncResource<>();
+ source.setCanceler(sourceCancel::incrementAndGet);
+
+ final AsyncResource<String> transformed = AsyncResources.transform(source, i -> "v" + i);
+ Assertions.assertFalse(transformed.isReady());
+
+ transformed.close();
+ Assertions.assertEquals(1, sourceCancel.get());
+ }
+}
diff --git a/processing/src/test/java/org/apache/druid/data/input/ResourceInputSource.java b/processing/src/test/java/org/apache/druid/data/input/ResourceInputSource.java
index b052d505848..1ecdaf0500b 100644
--- a/processing/src/test/java/org/apache/druid/data/input/ResourceInputSource.java
+++ b/processing/src/test/java/org/apache/druid/data/input/ResourceInputSource.java
@@ -119,8 +119,13 @@ public class ResourceInputSource extends AbstractInputSource
@Override
public InputStream open() throws IOException
{
- final InputStream resourceStream = classLoader.getResourceAsStream(resourceFile);
- return CompressionUtils.decompress(resourceStream, resourceFile);
+ return CompressionUtils.decompress(openRaw(), resourceFile);
+ }
+
+ @Override
+ public InputStream openRaw()
+ {
+ return classLoader.getResourceAsStream(resourceFile);
}
}
}
diff --git a/processing/src/test/java/org/apache/druid/data/input/impl/CloudObjectInputSourceTest.java b/processing/src/test/java/org/apache/druid/data/input/impl/CloudObjectInputSourceTest.java
index e0daac8c7e9..a0ac99f1446 100644
--- a/processing/src/test/java/org/apache/druid/data/input/impl/CloudObjectInputSourceTest.java
+++ b/processing/src/test/java/org/apache/druid/data/input/impl/CloudObjectInputSourceTest.java
@@ -19,22 +19,33 @@
package org.apache.druid.data.input.impl;
+import com.google.common.base.Predicates;
import com.google.common.collect.Lists;
import org.apache.druid.data.input.InputEntity;
+import org.apache.druid.data.input.InputFilePointer;
import org.apache.druid.data.input.InputSplit;
import org.apache.druid.data.input.MaxSizeSplitHintSpec;
+import org.apache.druid.data.input.impl.systemfield.SystemField;
import org.apache.druid.data.input.impl.systemfield.SystemFields;
import org.apache.druid.java.util.common.parsers.JSONPathSpec;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import org.mockito.ArgumentMatchers;
import org.mockito.Mockito;
+import java.io.ByteArrayInputStream;
+import java.io.File;
+import java.io.IOException;
import java.net.URI;
+import java.nio.charset.StandardCharsets;
import java.nio.file.FileSystems;
+import java.nio.file.Files;
import java.nio.file.PathMatcher;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.Collections;
+import java.util.EnumSet;
import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;
@@ -226,6 +237,110 @@ public class CloudObjectInputSourceTest
Assertions.assertEquals(OBJECTS.size(), entities.size());
}
+ @Test
+ public void test_asFilePointers_withUris()
+ {
+ final CloudObjectInputSource inputSource = Mockito.mock(CloudObjectInputSource.class, Mockito.withSettings()
+ .useConstructor(SCHEME, URIS2, null, null, null, SystemFields.none())
+ .defaultAnswer(Mockito.CALLS_REAL_METHODS)
+ );
+ Mockito.when(inputSource.getSplitWidget()).thenReturn(new MockSplitWidget(123L));
+
+ final List<InputFilePointer> pointers = inputSource.asFilePointers();
+
+ Assertions.assertEquals(
+ URIS2,
+ pointers.stream().map(InputFilePointer::uri).collect(Collectors.toList())
+ );
+
+ // sizeSupplier defers to the split widget.
+ Assertions.assertEquals(123L, pointers.get(0).sizeSupplier().getAsLong());
+ }
+
+ @Test
+ public void test_asFilePointers_withObjects()
+ {
+ final CloudObjectInputSource inputSource = Mockito.mock(CloudObjectInputSource.class, Mockito.withSettings()
+ .useConstructor(SCHEME, null, null, OBJECTS_BEFORE_GLOB, null, SystemFields.none())
+ .defaultAnswer(Mockito.CALLS_REAL_METHODS)
+ );
+ Mockito.when(inputSource.getSplitWidget()).thenReturn(new MockSplitWidget());
+
+ final List<InputFilePointer> pointers = inputSource.asFilePointers();
+
+ Assertions.assertEquals(
+ OBJECTS_BEFORE_GLOB.stream().map(object -> object.toUri(SCHEME)).collect(Collectors.toList()),
+ pointers.stream().map(InputFilePointer::uri).collect(Collectors.toList())
+ );
+ }
+
+ @Test
+ public void test_asFilePointers_withObjectGlob()
+ {
+ final CloudObjectInputSource inputSource = Mockito.mock(CloudObjectInputSource.class, Mockito.withSettings()
+ .useConstructor(SCHEME, null, null, OBJECTS_BEFORE_GLOB, "**.csv", SystemFields.none())
+ .defaultAnswer(Mockito.CALLS_REAL_METHODS)
+ );
+ Mockito.when(inputSource.getSplitWidget()).thenReturn(new MockSplitWidget());
+
+ final List<InputFilePointer> pointers = inputSource.asFilePointers();
+
+ // Only the .csv object survives the glob; the .parquet object is filtered out.
+ Assertions.assertEquals(
+ URIS,
+ pointers.stream().map(InputFilePointer::uri).collect(Collectors.toList())
+ );
+ }
+
+ @Test
+ public void test_asFilePointers_withSystemFieldsReturnsNull()
+ {
+ final CloudObjectInputSource inputSource = Mockito.mock(CloudObjectInputSource.class, Mockito.withSettings()
+ .useConstructor(SCHEME, URIS, null, null, null, new SystemFields(EnumSet.of(SystemField.URI)))
+ .defaultAnswer(Mockito.CALLS_REAL_METHODS)
+ );
+
+ // System fields cannot be expressed as file pointers.
+ Assertions.assertNull(inputSource.asFilePointers());
+ }
+
+ @Test
+ public void test_asFilePointers_withPrefixesReturnsNull()
+ {
+ final CloudObjectInputSource inputSource = Mockito.mock(CloudObjectInputSource.class, Mockito.withSettings()
+ .useConstructor(SCHEME, null, PREFIXES, null, null, SystemFields.none())
+ .defaultAnswer(Mockito.CALLS_REAL_METHODS)
+ );
+
+ // Prefixes are not expressed as file pointers.
+ Assertions.assertNull(inputSource.asFilePointers());
+ }
+
+ @Test
+ public void test_asFilePointers_populatorFetchesContent(@TempDir File tempDir) throws IOException
+ {
+ final byte[] content = "hello,world\n1,2\n".getBytes(StandardCharsets.UTF_8);
+
+ final InputEntity entity = Mockito.mock(InputEntity.class);
+ Mockito.when(entity.openRaw()).thenReturn(new ByteArrayInputStream(content));
+ Mockito.when(entity.getRetryCondition()).thenReturn(Predicates.alwaysFalse());
+
+ final CloudObjectInputSource inputSource = Mockito.mock(CloudObjectInputSource.class, Mockito.withSettings()
+ .useConstructor(SCHEME, URIS, null, null, null, SystemFields.none())
+ .defaultAnswer(Mockito.CALLS_REAL_METHODS)
+ );
+ Mockito.when(inputSource.getSplitWidget()).thenReturn(new MockSplitWidget());
+ Mockito.doReturn(entity).when(inputSource).createEntity(ArgumentMatchers.any());
+
+ final List<InputFilePointer> pointers = inputSource.asFilePointers();
+ Assertions.assertEquals(1, pointers.size());
+
+ final File dstFile = new File(tempDir, "fetched.csv");
+ pointers.get(0).populator().populate(dstFile);
+
+ Assertions.assertArrayEquals(content, Files.readAllBytes(dstFile.toPath()));
+ }
+
@Test
public void testGlobSubdirectories()
{
@@ -240,6 +355,18 @@ public class CloudObjectInputSourceTest
private static class MockSplitWidget implements CloudObjectSplitWidget
{
+ private final long objectSize;
+
+ MockSplitWidget()
+ {
+ this(0);
+ }
+
+ MockSplitWidget(final long objectSize)
+ {
+ this.objectSize = objectSize;
+ }
+
@Override
public Iterator<LocationWithSize> getDescriptorIteratorForPrefixes(List<URI> prefixes)
{
@@ -249,7 +376,7 @@ public class CloudObjectInputSourceTest
@Override
public long getObjectSize(CloudObjectLocation descriptor)
{
- return 0;
+ return objectSize;
}
}
}
diff --git a/processing/src/test/java/org/apache/druid/data/input/impl/FileEntityTest.java b/processing/src/test/java/org/apache/druid/data/input/impl/FileEntityTest.java
new file mode 100644
index 00000000000..315c942bc45
--- /dev/null
+++ b/processing/src/test/java/org/apache/druid/data/input/impl/FileEntityTest.java
@@ -0,0 +1,125 @@
+/*
+ * 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.data.input.impl;
+
+import org.apache.commons.io.IOUtils;
+import org.apache.druid.data.input.InputEntity;
+import org.apache.druid.data.input.InputEntity.CleanableFile;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.utils.CompressionUtils;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+
+public class FileEntityTest
+{
+ private static final String CONTENT = "the quick brown fox\n";
+
+ @TempDir
+ File tempDir;
+
+ @Test
+ public void test_openRaw_returnsRawBytesWithoutDecompressing() throws IOException
+ {
+ final File gzFile = gzipFile("data.txt.gz", CONTENT);
+ try (InputStream in = new FileEntity(gzFile).openRaw()) {
+ Assertions.assertArrayEquals(Files.readAllBytes(gzFile.toPath()), IOUtils.toByteArray(in));
+ }
+ }
+
+ @Test
+ public void test_open_decompressesBasedOnFileName() throws IOException
+ {
+ final File gzFile = gzipFile("data.txt.gz", CONTENT);
+ try (InputStream in = new FileEntity(gzFile).open()) {
+ Assertions.assertEquals(CONTENT, IOUtils.toString(in, StandardCharsets.UTF_8));
+ }
+ }
+
+ @Test
+ public void test_open_returnsContentForUncompressedFile() throws IOException
+ {
+ final File file = writeFile("data.txt", CONTENT);
+ try (InputStream in = new FileEntity(file).open()) {
+ Assertions.assertEquals(CONTENT, IOUtils.toString(in, StandardCharsets.UTF_8));
+ }
+ }
+
+ @Test
+ public void test_fetch_uncompressedReturnsSameFileWithoutCopying() throws IOException
+ {
+ final File file = writeFile("data.txt", CONTENT);
+ final File fetchDir = makeDir("fetch-uncompressed");
+ try (CleanableFile fetched = new FileEntity(file).fetch(fetchDir, fetchBuffer())) {
+ Assertions.assertEquals(file, fetched.file());
+ }
+ // close() is a no-op for uncompressed files; the original file remains.
+ Assertions.assertTrue(file.exists());
+ }
+
+ @Test
+ public void test_fetch_compressedDecompressesIntoTempDirAndCleansUp() throws IOException
+ {
+ final File gzFile = gzipFile("data.txt.gz", CONTENT);
+ final File fetchDir = makeDir("fetch-compressed");
+ final File fetchedFile;
+ try (CleanableFile fetched = new FileEntity(gzFile).fetch(fetchDir, fetchBuffer())) {
+ fetchedFile = fetched.file();
+ Assertions.assertEquals(fetchDir, fetchedFile.getParentFile());
+ Assertions.assertEquals(CONTENT, StringUtils.fromUtf8(Files.readAllBytes(fetchedFile.toPath())));
+ }
+ // close() removes the decompressed temp file; the original compressed file remains.
+ Assertions.assertFalse(fetchedFile.exists());
+ Assertions.assertTrue(gzFile.exists());
+ }
+
+ private File writeFile(final String name, final String content) throws IOException
+ {
+ final File file = new File(tempDir, name);
+ Files.write(file.toPath(), StringUtils.toUtf8(content));
+ return file;
+ }
+
+ private File gzipFile(final String name, final String content) throws IOException
+ {
+ final File source = writeFile("source-for-" + name, content);
+ final File gzFile = new File(tempDir, name);
+ CompressionUtils.gzip(source, gzFile);
+ return gzFile;
+ }
+
+ private File makeDir(final String name) throws IOException
+ {
+ final File dir = new File(tempDir, name);
+ Files.createDirectories(dir.toPath());
+ return dir;
+ }
+
+ private static byte[] fetchBuffer()
+ {
+ return new byte[InputEntity.DEFAULT_FETCH_BUFFER_SIZE];
+ }
+}
diff --git a/processing/src/test/java/org/apache/druid/query/aggregation/AggregationTestHelper.java b/processing/src/test/java/org/apache/druid/query/aggregation/AggregationTestHelper.java
index 0c3c846f487..b8a31fe584b 100644
--- a/processing/src/test/java/org/apache/druid/query/aggregation/AggregationTestHelper.java
+++ b/processing/src/test/java/org/apache/druid/query/aggregation/AggregationTestHelper.java
@@ -430,7 +430,7 @@ public class AggregationTestHelper implements Closeable
}
@Override
- public InputStream open()
+ public InputStream openRaw()
{
return inputDataStream;
}
diff --git a/server/src/main/java/org/apache/druid/guice/LocalDataStorageDruidModule.java b/server/src/main/java/org/apache/druid/guice/LocalDataStorageDruidModule.java
index 4cdd28017ac..8f1d44c3fe3 100644
--- a/server/src/main/java/org/apache/druid/guice/LocalDataStorageDruidModule.java
+++ b/server/src/main/java/org/apache/druid/guice/LocalDataStorageDruidModule.java
@@ -34,8 +34,6 @@ import org.apache.druid.segment.loading.LocalDataSegmentPusher;
import org.apache.druid.segment.loading.LocalDataSegmentPusherConfig;
import org.apache.druid.segment.loading.LocalFileTimestampVersionFinder;
import org.apache.druid.segment.loading.LocalLoadSpec;
-import org.apache.druid.segment.loading.SegmentCacheManager;
-import org.apache.druid.segment.loading.SegmentLocalCacheManager;
import java.util.List;
@@ -48,8 +46,6 @@ public class LocalDataStorageDruidModule implements DruidModule
@Override
public void configure(Binder binder)
{
- binder.bind(SegmentCacheManager.class).to(SegmentLocalCacheManager.class).in(LazySingleton.class);
-
bindDeepStorageLocal(binder);
PolyBind.createChoice(
diff --git a/server/src/main/java/org/apache/druid/guice/StorageNodeModule.java b/server/src/main/java/org/apache/druid/guice/StorageNodeModule.java
index 979dce5c3a6..09b953204a4 100644
--- a/server/src/main/java/org/apache/druid/guice/StorageNodeModule.java
+++ b/server/src/main/java/org/apache/druid/guice/StorageNodeModule.java
@@ -33,9 +33,14 @@ import org.apache.druid.java.util.emitter.EmittingLogger;
import org.apache.druid.query.DruidProcessingConfig;
import org.apache.druid.segment.DefaultColumnFormatConfig;
import org.apache.druid.segment.column.ColumnConfig;
+import org.apache.druid.segment.loading.SegmentCacheManager;
import org.apache.druid.segment.loading.SegmentLoaderConfig;
+import org.apache.druid.segment.loading.SegmentLocalCacheManager;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
import org.apache.druid.segment.loading.StorageLocation;
import org.apache.druid.segment.loading.StorageLocationSelectorStrategy;
+import org.apache.druid.segment.loading.external.StorageLocationVirtualStorageManager;
+import org.apache.druid.segment.loading.external.VirtualStorageManager;
import org.apache.druid.server.DruidNode;
import org.apache.druid.server.coordination.DruidServerMetadata;
import org.apache.druid.server.coordination.ServerType;
@@ -62,6 +67,8 @@ public class StorageNodeModule implements Module
bindLocationSelectorStrategy(binder);
binder.bind(ServerTypeConfig.class).toProvider(Providers.of(null));
binder.bind(ColumnConfig.class).to(DruidProcessingConfig.class).in(LazySingleton.class);
+ binder.bind(SegmentCacheManager.class).to(SegmentLocalCacheManager.class).in(LazySingleton.class);
+ binder.bind(VirtualStorageManager.class).to(StorageLocationVirtualStorageManager.class).in(LazySingleton.class);
MetricsModule.register(binder, StorageMonitor.class);
}
@@ -120,6 +127,13 @@ public class StorageNodeModule implements Module
);
}
+ @Provides
+ @ManageLifecycle
+ public StorageLoadingThreadPool getStorageLoadingThreadPool(SegmentLoaderConfig config)
+ {
+ return StorageLoadingThreadPool.createFromConfig(config);
+ }
+
@Provides
@LazySingleton
@Named(IS_SEGMENT_CACHE_CONFIGURED)
diff --git a/server/src/main/java/org/apache/druid/metadata/input/SqlEntity.java b/server/src/main/java/org/apache/druid/metadata/input/SqlEntity.java
index dcf315b81f8..e77e5e2fbcf 100644
--- a/server/src/main/java/org/apache/druid/metadata/input/SqlEntity.java
+++ b/server/src/main/java/org/apache/druid/metadata/input/SqlEntity.java
@@ -84,7 +84,7 @@ public class SqlEntity implements InputEntity
}
@Override
- public InputStream open()
+ public InputStream openRaw()
{
throw new UnsupportedOperationException("Please use fetch() instead");
}
diff --git a/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java b/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
index 393b0ac8831..d3fa1b84f2d 100644
--- a/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
+++ b/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
@@ -144,4 +144,9 @@ public interface SegmentCacheManager
* Returns the storage locations backing this cache manager.
*/
List<StorageLocation> getLocations();
+
+ /**
+ * Returns the loading thread pool backing this cache manager.
+ */
+ StorageLoadingThreadPool getLoadingThreadPool();
}
diff --git a/server/src/main/java/org/apache/druid/segment/loading/SegmentLoaderConfig.java b/server/src/main/java/org/apache/druid/segment/loading/SegmentLoaderConfig.java
index a8d3981ad70..1fbe3679ad0 100644
--- a/server/src/main/java/org/apache/druid/segment/loading/SegmentLoaderConfig.java
+++ b/server/src/main/java/org/apache/druid/segment/loading/SegmentLoaderConfig.java
@@ -204,14 +204,19 @@ public class SegmentLoaderConfig
}
/**
- * Sets {@link #virtualStorage} and {@link #virtualStorageIsEphemeral}.
+ * Sets {@link #virtualStorage}.
*/
- public SegmentLoaderConfig setVirtualStorage(
- boolean virtualStorage,
- boolean virtualStorageFabricEphemeral
- )
+ public SegmentLoaderConfig setVirtualStorage(boolean virtualStorage)
{
this.virtualStorage = virtualStorage;
+ return this;
+ }
+
+ /**
+ * Sets {@link #virtualStorageIsEphemeral}.
+ */
+ public SegmentLoaderConfig setVirtualStorageIsEphemeral(boolean virtualStorageFabricEphemeral)
+ {
this.virtualStorageIsEphemeral = virtualStorageFabricEphemeral;
return this;
}
@@ -225,9 +230,19 @@ public class SegmentLoaderConfig
{
return this.getLocations()
.stream()
- .map(locationConfig -> new StorageLocation(locationConfig.getPath(),
- locationConfig.getMaxSize(),
- locationConfig.getFreeSpacePercent()))
+ .map(locationConfig -> {
+ final StorageLocation location = new StorageLocation(
+ locationConfig.getPath(),
+ locationConfig.getMaxSize(),
+ locationConfig.getFreeSpacePercent()
+ );
+
+ if (isVirtualStorageEphemeral()) {
+ location.setAreWeakEntriesEphemeral(true);
+ }
+
+ return location;
+ })
.collect(Collectors.toList());
}
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 5372baf4c24..a43400efcb1 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
@@ -24,7 +24,6 @@ import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Suppliers;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
-import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import com.google.inject.Inject;
@@ -115,13 +114,7 @@ public class SegmentLocalCacheManager implements SegmentCacheManager
private final IndexIO indexIO;
- private final ListeningExecutorService virtualStorageLoadOnDemandExec;
- /**
- * Bounds the number of in-flight on-demand loads when running with virtual threads. Null when virtual storage is
- * disabled or when the legacy fixed-pool path is selected (the pool size is the implicit limit there).
- */
- @Nullable
- private final Semaphore virtualStorageLoadOnDemandPermits;
+ private final StorageLoadingThreadPool virtualStorageLoadingThreadPool;
private ExecutorService loadOnBootstrapExec = null;
private ExecutorService loadOnDownloadExec = null;
@@ -129,64 +122,29 @@ public class SegmentLocalCacheManager implements SegmentCacheManager
public SegmentLocalCacheManager(
List<StorageLocation> locations,
SegmentLoaderConfig config,
+ StorageLoadingThreadPool virtualStorageLoadingThreadPool,
@Nonnull StorageLocationSelectorStrategy strategy,
IndexIO indexIO,
@Json ObjectMapper mapper
)
{
- this.config = config;
- this.jsonMapper = mapper;
this.locations = locations;
+ this.config = config;
+ this.virtualStorageLoadingThreadPool = virtualStorageLoadingThreadPool;
this.strategy = strategy;
this.indexIO = indexIO;
+ this.jsonMapper = mapper;
log.info("Using storage location strategy[%s].", this.strategy.getClass().getSimpleName());
if (config.isVirtualStorage()) {
if (config.getNumThreadsToLoadSegmentsIntoPageCacheOnDownload() > 0) {
- throw DruidException.defensive("Invalid configuration: virtualStorage is incompatible with numThreadsToLoadSegmentsIntoPageCacheOnDownload");
+ throw DruidException.defensive(
+ "Invalid configuration: virtualStorage is incompatible with numThreadsToLoadSegmentsIntoPageCacheOnDownload");
}
if (config.getNumThreadsToLoadSegmentsIntoPageCacheOnBootstrap() > 0) {
- throw DruidException.defensive("Invalid configuration: virtualStorage is incompatible with numThreadsToLoadSegmentsIntoPageCacheOnBootstrap");
- }
- 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.isVirtualStorageEphemeral()) {
- for (StorageLocation location : locations) {
- location.setAreWeakEntriesEphemeral(true);
- }
- }
- if (config.isVirtualStorageUseVirtualThreads()) {
- log.info(
- "Using virtual storage mode with virtual threads - max concurrent on demand loads: [%d].",
- config.getVirtualStorageLoadThreads()
- );
- virtualStorageLoadOnDemandPermits = new Semaphore(config.getVirtualStorageLoadThreads());
- virtualStorageLoadOnDemandExec = MoreExecutors.listeningDecorator(
- Executors.newThreadPerTaskExecutor(
- Thread.ofVirtual()
- .name("VirtualStorageOnDemandLoadingThread-", 0)
- .factory()
- )
- );
- } else {
- log.info(
- "Using virtual storage mode with fixed platform thread pool - on demand load threads: [%d].",
- config.getVirtualStorageLoadThreads()
- );
- virtualStorageLoadOnDemandPermits = null;
- virtualStorageLoadOnDemandExec = MoreExecutors.listeningDecorator(
- Executors.newFixedThreadPool(
- config.getVirtualStorageLoadThreads(),
- Execs.makeThreadFactory("VirtualStorageOnDemandLoadingThread-%s")
- )
- );
+ throw DruidException.defensive(
+ "Invalid configuration: virtualStorage is incompatible with numThreadsToLoadSegmentsIntoPageCacheOnBootstrap");
}
} else {
log.info(
@@ -208,8 +166,6 @@ public class SegmentLocalCacheManager implements SegmentCacheManager
Execs.makeThreadFactory("LoadSegmentsIntoPageCacheOnDownload-%s")
);
}
- virtualStorageLoadOnDemandExec = null;
- virtualStorageLoadOnDemandPermits = null;
}
}
@@ -697,9 +653,6 @@ public class SegmentLocalCacheManager implements SegmentCacheManager
if (loadOnDownloadExec != null) {
loadOnDownloadExec.shutdown();
}
- if (virtualStorageLoadOnDemandExec != null) {
- virtualStorageLoadOnDemandExec.shutdown();
- }
}
@VisibleForTesting
@@ -714,6 +667,12 @@ public class SegmentLocalCacheManager implements SegmentCacheManager
return locations;
}
+ @Override
+ public StorageLoadingThreadPool getLoadingThreadPool()
+ {
+ return virtualStorageLoadingThreadPool;
+ }
+
/**
* Checks whether a segment is already cached. This method does not confirm if the segment is actually mounted in
* the location, or even that the segment files in some location are valid, just that some files exist in the
@@ -788,12 +747,13 @@ public class SegmentLocalCacheManager implements SegmentCacheManager
return Suppliers.memoize(
() -> {
final long startTime = System.nanoTime();
- return virtualStorageLoadOnDemandExec.submit(
+ return virtualStorageLoadingThreadPool.getExecutorService().submit(
() -> {
// Acquire a permit inside the task so that waiting parks the (virtual) thread instead of blocking the
// submitter. In fixed-pool mode permits is null and the pool size itself bounds concurrency.
- if (virtualStorageLoadOnDemandPermits != null) {
- virtualStorageLoadOnDemandPermits.acquireUninterruptibly();
+ final Semaphore loadingPermits = virtualStorageLoadingThreadPool.getPermits();
+ if (loadingPermits != null) {
+ loadingPermits.acquireUninterruptibly();
}
try {
final long execStartTime = System.nanoTime();
@@ -807,8 +767,8 @@ public class SegmentLocalCacheManager implements SegmentCacheManager
);
}
finally {
- if (virtualStorageLoadOnDemandPermits != null) {
- virtualStorageLoadOnDemandPermits.release();
+ if (loadingPermits != null) {
+ loadingPermits.release();
}
}
}
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
new file mode 100644
index 00000000000..9c8de953b9d
--- /dev/null
+++ b/server/src/main/java/org/apache/druid/segment/loading/StorageLoadingThreadPool.java
@@ -0,0 +1,143 @@
+/*
+ * 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.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.guice.StorageNodeModule;
+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.loading.external.VirtualStorageManager;
+
+import javax.annotation.Nullable;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
+
+/**
+ * Holds the thread pool used for background loading by {@link SegmentLocalCacheManager} and
+ * {@link VirtualStorageManager}.
+ */
+public class StorageLoadingThreadPool
+{
+ private static final Logger log = new Logger(StorageNodeModule.class);
+
+ private final ListeningExecutorService exec;
+ private final Semaphore permits;
+
+ public StorageLoadingThreadPool(
+ @Nullable final ListeningExecutorService exec,
+ @Nullable final Semaphore permits
+ )
+ {
+ this.exec = exec;
+ this.permits = permits;
+ }
+
+ public static StorageLoadingThreadPool createFromConfig(final SegmentLoaderConfig config)
+ {
+ final ListeningExecutorService exec;
+ final Semaphore permits;
+
+ 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()
+ );
+ permits = new Semaphore(config.getVirtualStorageLoadThreads());
+ exec = MoreExecutors.listeningDecorator(
+ Executors.newThreadPerTaskExecutor(
+ Thread.ofVirtual()
+ .name("VirtualStorageOnDemandLoadingThread-", 0)
+ .factory()
+ )
+ );
+ } else {
+ log.info(
+ "Using virtual storage mode with fixed platform thread pool - on demand load threads: [%d].",
+ config.getVirtualStorageLoadThreads()
+ );
+ permits = null;
+ exec = MoreExecutors.listeningDecorator(
+ Executors.newFixedThreadPool(
+ config.getVirtualStorageLoadThreads(),
+ Execs.makeThreadFactory("VirtualStorageOnDemandLoadingThread-%s")
+ )
+ );
+ }
+ } else {
+ exec = null;
+ permits = null;
+ }
+
+ return new StorageLoadingThreadPool(exec, permits);
+ }
+
+ /**
+ * Returns an instance representing "no thread pool". Calling {@link #getExecutorService()} will return an error.
+ */
+ public static StorageLoadingThreadPool none()
+ {
+ return new StorageLoadingThreadPool(null, null);
+ }
+
+ public boolean isAvailable()
+ {
+ return exec != null;
+ }
+
+ public ListeningExecutorService getExecutorService()
+ {
+ if (exec == null) {
+ throw DruidException.defensive("No exec set, we thought we wouldn't need this");
+ }
+ return exec;
+ }
+
+ /**
+ * Bounds the number of in-flight on-demand loads when running with virtual threads. If null, no permits are used;
+ * the natural limit of the {@link #getExecutorService()} is sufficient. If nonnull, callers must acquire a permit
+ * in the tasks submitted to {@link #getExecutorService()}, before doing real work. Permits must be released before
+ * the task completes.
+ */
+ @Nullable
+ public Semaphore getPermits()
+ {
+ return permits;
+ }
+
+ @LifecycleStop
+ public void stop()
+ {
+ if (exec != null) {
+ exec.shutdownNow();
+ }
+ }
+}
diff --git a/server/src/main/java/org/apache/druid/segment/loading/external/CachedFile.java b/server/src/main/java/org/apache/druid/segment/loading/external/CachedFile.java
new file mode 100644
index 00000000000..875e801d214
--- /dev/null
+++ b/server/src/main/java/org/apache/druid/segment/loading/external/CachedFile.java
@@ -0,0 +1,87 @@
+/*
+ * 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.external;
+
+import org.apache.druid.segment.loading.CacheEntry;
+import org.apache.druid.segment.loading.StorageLocation;
+
+import java.io.Closeable;
+import java.io.File;
+import java.util.function.Function;
+
+/**
+ * Handle to a cached file that prevents eviction while held.
+ * Must be closed when no longer needed to make the file eligible for eviction.
+ */
+public class CachedFile implements Closeable
+{
+ private final StorageLocation.ReservationHold<CacheEntry> hold;
+ private final DownloadableCacheEntry entry;
+ private final String identifier;
+
+ CachedFile(StorageLocation.ReservationHold<CacheEntry> hold, String identifier)
+ {
+ this.hold = hold;
+ this.entry = (DownloadableCacheEntry) hold.getEntry();
+ this.identifier = identifier;
+ }
+
+ /**
+ * Get the identifier for this cached file.
+ */
+ public String getIdentifier()
+ {
+ return identifier;
+ }
+
+ /**
+ * Get the File object for this cached file.
+ * The file is guaranteed to exist and be populated while this handle is open.
+ */
+ public File getFile()
+ {
+ return entry.getFile();
+ }
+
+ /**
+ * Extend this cached file with additional functionality. The lifecycle of the extended functionality is
+ * tracked along with the lifecycle of the underlying cache entry. The provided supplier is called immediately unless
+ * there is already a mapping for the class. Either way, the class is returned.
+ *
+ * <p>Do not call "close" on the returned object. Its lifecycle is associated with the underlying cache entry,
+ * and it may be reused beyond the lifecycle of this particular CachedFile object. It will be closed automatically
+ * when the underlying cache entry is unmounted.
+ */
+ public <T extends Closeable> T extend(Class<T> clazz, Function<CachedFile, T> supplier)
+ {
+ return entry.extend(clazz, () -> supplier.apply(this));
+ }
+
+ /**
+ * Release the hold on this cached file, making it eligible for eviction.
+ *
+ * <p>Usage of the File returned from this object after the hold is closed is undefined and expected to end poorly.
+ */
+ @Override
+ public void close()
+ {
+ hold.close();
+ }
+}
diff --git a/server/src/main/java/org/apache/druid/segment/loading/external/DownloadableCacheEntry.java b/server/src/main/java/org/apache/druid/segment/loading/external/DownloadableCacheEntry.java
new file mode 100644
index 00000000000..9a736b193ad
--- /dev/null
+++ b/server/src/main/java/org/apache/druid/segment/loading/external/DownloadableCacheEntry.java
@@ -0,0 +1,314 @@
+/*
+ * 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.external;
+
+import com.google.common.hash.Hashing;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.io.FilePopulator;
+import org.apache.druid.java.util.common.FileUtils;
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.segment.loading.CacheEntry;
+import org.apache.druid.segment.loading.CacheEntryIdentifier;
+import org.apache.druid.segment.loading.StorageLocation;
+import org.apache.druid.utils.CloseableUtils;
+
+import javax.annotation.Nullable;
+import java.io.Closeable;
+import java.io.File;
+import java.io.IOException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
+
+/**
+ * CacheEntry implementation that calls a FilePopulator lambda when mounted.
+ */
+class DownloadableCacheEntry implements CacheEntry
+{
+ private static final Logger log = new Logger(DownloadableCacheEntry.class);
+
+ private static final int MAX_PREFIX_SUFFIX_LENGTH = 50;
+
+ private final StringCacheEntryIdentifier identifier;
+ private final long sizeBytes;
+ private final FilePopulator populator;
+ private final File locationPath;
+
+ private volatile boolean mounted = false;
+ private volatile File mountedFile = null;
+
+ /**
+ * Map for keeping track of functionality added by {@link #extend}.
+ */
+ private final ConcurrentHashMap<Class<? extends Closeable>, Closeable> extendMap = new ConcurrentHashMap<>();
+
+ public DownloadableCacheEntry(
+ StringCacheEntryIdentifier identifier,
+ long sizeBytes,
+ FilePopulator populator,
+ File locationPath
+ )
+ {
+ this.identifier = identifier;
+ this.sizeBytes = sizeBytes;
+ this.populator = populator;
+ this.locationPath = locationPath;
+ }
+
+ @Override
+ public CacheEntryIdentifier getId()
+ {
+ return identifier;
+ }
+
+ @Override
+ public long getSize()
+ {
+ return sizeBytes;
+ }
+
+ @Override
+ public boolean isMounted()
+ {
+ return mounted;
+ }
+
+ @Override
+ public void mount(StorageLocation location)
+ {
+ if (mounted) {
+ return; // Already mounted
+ }
+
+ // Determine file path - delegate to the same sanitization logic
+ File file = getFileForIdentifier(locationPath, identifier.value());
+
+ // Ensure parent directory exists
+ File parentDir = file.getParentFile();
+ if (!parentDir.exists()) {
+ try {
+ FileUtils.mkdirp(parentDir);
+ }
+ catch (IOException e) {
+ throw DruidException.forPersona(DruidException.Persona.OPERATOR)
+ .ofCategory(DruidException.Category.RUNTIME_FAILURE)
+ .build(e, "Failed to create parent directory [%s]", parentDir);
+ }
+ }
+
+ if (!(file.exists() && file.length() == sizeBytes)) {
+ // file doesn't exist (or is the wrong length), so we need populate it
+
+ if (file.exists()) {
+ // It wasn't the right length, let's delete it to avoid confusion for the populator.
+ if (!file.delete()) {
+ log.info("Problem deleting file [%s] before populating for VSM", file);
+ }
+ }
+
+ try {
+ populator.populate(file);
+ }
+ catch (Throwable e) {
+ // Clean up partial file on failure
+ if (file.exists()) {
+ file.delete();
+ }
+ throw DruidException.forPersona(DruidException.Persona.OPERATOR)
+ .ofCategory(DruidException.Category.RUNTIME_FAILURE)
+ .build(e, "Failed to populate file [%s]", file);
+ }
+
+ // Verify the file was created
+ if (!file.exists()) {
+ throw DruidException.forPersona(DruidException.Persona.OPERATOR)
+ .ofCategory(DruidException.Category.RUNTIME_FAILURE)
+ .build("Populator did not create file [%s]", file);
+ }
+ }
+
+ this.mountedFile = file;
+ this.mounted = true;
+ }
+
+ @Override
+ public void unmount()
+ {
+ final Closer extendCloser = CloseableUtils.forIterable(extendMap.values());
+ extendMap.clear();
+ CloseableUtils.closeAndSuppressExceptions(extendCloser, e -> log.warn(e, "Failed to close extended functionality"));
+
+ // StorageLocation will call this when it needs to reclaim space
+ if (mountedFile != null && mountedFile.exists()) {
+ if (!mountedFile.delete()) {
+ log.warn("Failed to delete file[%s] on unmount(). Leaving it.", mountedFile);
+ }
+ }
+ mountedFile = null;
+ mounted = false;
+ }
+
+ /**
+ * Get the mounted file.
+ *
+ * @return The file, or null if not mounted
+ */
+ @Nullable
+ File getFile()
+ {
+ return mountedFile;
+ }
+
+ /**
+ * Extend this cache entry with additional functionality. The lifecycle of the extended functionality is
+ * tracked along with the lifecycle of the cache entry. The provided supplier is called immediately unless
+ * there is already a mapping for the class. Either way, the class is returned.
+ */
+ @SuppressWarnings("unchecked")
+ public <T extends Closeable> T extend(Class<T> clazz, Supplier<T> supplier)
+ {
+ return (T) extendMap.computeIfAbsent(clazz, k -> supplier.get());
+ }
+
+ private File getFileForIdentifier(File locationPath, String identifier)
+ {
+ String sanitized = sanitizePath(identifier);
+ return new File(locationPath, sanitized);
+ }
+
+ /**
+ * Converts an arbitrary cache entry identifier (typically a URI or file path) into a filename. The result has
+ * the form {@code <prefix>-<hash>-<suffix>}. The purpose of this sanitization is to provide identifiers for
+ * external paths that filesystem-safe (hence the removal of special characters and shortening of the overall
+ * name), non-colliding (hence the inclusion of a hash), and human-recognizable (hence the inclusion of prefix
+ * and suffix).
+ *
+ * <p>Typically the prefix and suffix are of equal length. However, care is taken to preserve extensions, which
+ * means that the suffix may end up longer than the prefix if this is required to avoid mangling an extension.
+ * This allows downstream code to e.g. detect format or compression based on file extension.
+ *
+ * @throws DruidException if the identifier sanitizes to an empty string
+ */
+ static String sanitizePath(final String originalPath)
+ {
+ String path = originalPath;
+ path = path.replace('\\', '/');
+
+ int startIndex = 0;
+ if (path.startsWith("/")) {
+ startIndex = 1;
+ }
+ int endIndex = path.endsWith("/") ? path.length() - 1 : path.length();
+
+ // Split the kept portion into a prefix and a suffix so that file extensions (e.g. ".gz") are
+ // preserved at the end of the sanitized name. The hashcode disambiguates collisions when
+ // multiple paths share the same prefix+suffix.
+ final int prefixEnd;
+ final int suffixStart;
+ if (endIndex - startIndex > MAX_PREFIX_SUFFIX_LENGTH * 2) {
+ prefixEnd = startIndex + MAX_PREFIX_SUFFIX_LENGTH;
+ suffixStart = endIndex - MAX_PREFIX_SUFFIX_LENGTH;
+ } else {
+ final int mid = startIndex + (endIndex - startIndex) / 2;
+ final int extensionPos = indexOfExtension(path, startIndex, endIndex);
+ final int split;
+ if (extensionPos >= 0 && extensionPos < mid && endIndex - extensionPos < MAX_PREFIX_SUFFIX_LENGTH) {
+ // Adjust the split point so the suffix contains the entire extension.
+ split = extensionPos;
+ } else {
+ split = mid;
+ }
+ prefixEnd = split;
+ suffixStart = split;
+ }
+
+ final StringBuilder sanitized = new StringBuilder();
+ appendSanitized(sanitized, path, startIndex, prefixEnd);
+ sanitized.append('-')
+ .append(Hashing.sha512().hashUnencodedChars(originalPath).toString(), 0, 16)
+ .append('-');
+ appendSanitized(sanitized, path, suffixStart, endIndex);
+
+ final String result = sanitized.toString();
+ if (result.isEmpty()) {
+ throw DruidException.defensive("Identifier resulted in empty path after sanitization");
+ }
+
+ return result;
+ }
+
+ private static void appendSanitized(StringBuilder dest, String path, int start, int end)
+ {
+ for (int i = start; i < end; i++) {
+ char c = path.charAt(i);
+ if (Character.isLetterOrDigit(c) || c == '-' || c == '_' || c == '.') {
+ dest.append(c);
+ } else {
+ dest.append('_');
+ }
+ }
+ }
+
+ /**
+ * Returns the position in {@code path} where the extension begins, or -1 if there is no extension.
+ * Searches from startIndex (inclusive) to endIndex (exclusive). Only alphanumeric strings are
+ * considered possible extensions.
+ */
+ private static int indexOfExtension(String path, int startIndex, int endIndex)
+ {
+ // Find the filename.
+ int basenameStart = startIndex;
+ for (int i = endIndex - 1; i >= startIndex; i--) {
+ if (path.charAt(i) == '/') {
+ basenameStart = i + 1;
+ break;
+ }
+ }
+
+ // Find the extension of the filename.
+ int pos = endIndex;
+ int extensionPos = -1;
+ while (pos > basenameStart) {
+ final int initialPos = pos;
+ while (pos > basenameStart && isAlphaNumeric(path.charAt(pos - 1))) {
+ pos--;
+ }
+ if (pos == initialPos) {
+ // Saw some non-alphanumeric characters.
+ break;
+ }
+ // Check if there's a dot. Use basenameStart + 1 because we don't want to consider ".dotfile" as an extension.
+ if (pos > basenameStart + 1 && path.charAt(pos - 1) == '.') {
+ // Found an extension. Remember the position and keep searching (for "xyz.log.gz" we want to capture ".log.gz").
+ pos--;
+ extensionPos = pos;
+ } else {
+ break;
+ }
+ }
+
+ return extensionPos;
+ }
+
+ private static boolean isAlphaNumeric(char c)
+ {
+ return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z');
+ }
+}
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
new file mode 100644
index 00000000000..943d22506a4
--- /dev/null
+++ b/server/src/main/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManager.java
@@ -0,0 +1,299 @@
+/*
+ * 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.external;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.inject.Inject;
+import org.apache.druid.collections.ResourceHolder;
+import org.apache.druid.common.asyncresource.AsyncResource;
+import org.apache.druid.common.asyncresource.AsyncResources;
+import org.apache.druid.common.asyncresource.SettableAsyncResource;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.io.FilePopulator;
+import org.apache.druid.java.util.emitter.EmittingLogger;
+import org.apache.druid.segment.loading.CacheEntry;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
+import org.apache.druid.segment.loading.StorageLocation;
+import org.apache.druid.segment.loading.StorageLocationSelectorStrategy;
+
+import javax.annotation.Nullable;
+import java.io.File;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.LongSupplier;
+
+/**
+ * Default implementation of VirtualStorageManager that delegates to StorageLocation.
+ * Uses weak reservations for all cached files, making them eligible for eviction.
+ */
+public class StorageLocationVirtualStorageManager implements VirtualStorageManager
+{
+ private static final EmittingLogger log = new EmittingLogger(StorageLocationVirtualStorageManager.class);
+
+ private final StorageLocationSelectorStrategy strategy;
+ private final StorageLoadingThreadPool loadingThreadPool;
+
+ /**
+ * Per-identifier locks to ensure only one thread populates a given identifier.
+ * Once a location is resolved, it's stored in the lock so subsequent threads can check it.
+ */
+ private final ConcurrentHashMap<String, PopulationLock> populationLocks = new ConcurrentHashMap<>();
+
+ @Inject
+ public StorageLocationVirtualStorageManager(
+ List<StorageLocation> locations,
+ StorageLocationSelectorStrategy strategy,
+ StorageLoadingThreadPool loadingThreadPool
+ )
+ {
+ this.strategy = strategy;
+ this.loadingThreadPool = loadingThreadPool;
+ log.info("Initialized VirtualStorageManager with [%d] storage locations", locations.size());
+ }
+
+ @Nullable
+ @Override
+ public CachedFile get(String identifier)
+ {
+ // We only consult the lock to discover which location the entry was resolved to. No need for synchronization,
+ // because we don't intend to populate it.
+ final PopulationLock lock = populationLocks.get(identifier);
+ if (lock == null) {
+ return null;
+ }
+
+ final StorageLocation resolvedLocation = lock.getResolvedLocation();
+ if (resolvedLocation == null) {
+ // Population is still in flight: the lock exists but doesn't have a location yet.
+ // Return that the file doesn't exist instead of blocking.
+ return null;
+ }
+
+ final StorageLocation.ReservationHold<CacheEntry> hold =
+ resolvedLocation.addWeakReservationHoldIfExists(new StringCacheEntryIdentifier(identifier));
+ return hold == null ? null : new CachedFile(hold, identifier);
+ }
+
+ @Override
+ public CachedFile reserveAndPopulate(
+ String identifier,
+ LongSupplier sizeSupplier,
+ FilePopulator populator
+ )
+ {
+ // Get or create lock for this identifier
+ final PopulationLock lock = populationLocks.computeIfAbsent(identifier, ignored -> new PopulationLock());
+
+ synchronized (lock) {
+ final StringCacheEntryIdentifier cacheId = new StringCacheEntryIdentifier(identifier);
+
+ // If multiple threads are trying to reserve the same location, the first one will update the state on the lock
+ // so check that first.
+ final StorageLocation resolvedLocation = lock.getResolvedLocation();
+ if (resolvedLocation == null) {
+ // Determining the size often requires an external system call, so we use this supplier to defer it until
+ // we absolutely need it
+ final long sizeBytes = sizeSupplier.getAsLong();
+
+ // Try to reserve in each location according to strategy
+ final Iterator<StorageLocation> locationIter = strategy.getLocations();
+ Throwable lastException = null;
+
+ while (locationIter.hasNext()) {
+ final StorageLocation location = locationIter.next();
+ final File locationFile = location.getPath();
+ try {
+ // 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)
+ {
+ final AtomicBoolean mounted = new AtomicBoolean(false);
+
+ @Override
+ public void mount(StorageLocation location)
+ {
+ super.mount(location);
+ mounted.set(true);
+ }
+
+ @Override
+ public void unmount()
+ {
+ if (mounted.get()) {
+ populationLocks.remove(identifier, lock);
+ }
+ super.unmount();
+ }
+ }
+ );
+
+ if (hold == null) {
+ log.debug(
+ "Failed to reserve [%d] bytes for [%s] in location [%s], trying next",
+ sizeBytes,
+ identifier,
+ locationFile
+ );
+ continue;
+ }
+
+ // Mount the entry (calls populator)
+ try {
+ hold.getEntry().mount(location);
+ }
+ catch (Throwable e) {
+ try {
+ hold.close();
+ }
+ catch (Throwable e2) {
+ e.addSuppressed(e2);
+ }
+ throw e;
+ }
+
+ // Store the resolved location for anything that is waiting on this same lock
+ lock.setResolvedLocation(location);
+
+ return new CachedFile(hold, identifier);
+ }
+ catch (Throwable e) {
+ lastException = e;
+ log.debug(e, "Failed to reserve and populate in location [%s], trying next", locationFile);
+ }
+ }
+ // We have exited the loop, which should have returned, nothing must've worked...
+
+ if (lastException != null) {
+ throw DruidException.forPersona(DruidException.Persona.OPERATOR)
+ .ofCategory(DruidException.Category.RUNTIME_FAILURE)
+ .build(
+ lastException,
+ "Failed to reserve and populate entry [%s] in any location",
+ identifier
+ );
+ } else {
+ throw DruidException.forPersona(DruidException.Persona.OPERATOR)
+ .ofCategory(DruidException.Category.CAPACITY_EXCEEDED)
+ .build(
+ "No space available to reserve [%,d] bytes for entry [%s]",
+ sizeBytes,
+ identifier
+ );
+ }
+ } else {
+ // Try to get from the resolved location first
+ StorageLocation.ReservationHold<CacheEntry> hold = resolvedLocation.addWeakReservationHoldIfExists(cacheId);
+
+ if (hold != null) {
+ return new CachedFile(hold, identifier);
+ } else {
+ // hold == null means the entry was evicted before we had a chance to add our hold.
+ // Drop the stale lock and re-resolve.
+ populationLocks.remove(identifier, lock);
+ return reserveAndPopulate(identifier, sizeSupplier, populator);
+ }
+ }
+ }
+ }
+
+ @Override
+ public AsyncResource<CachedFile> reserveAndPopulateAsync(
+ String identifier,
+ LongSupplier sizeSupplier,
+ FilePopulator populator
+ )
+ {
+ final CachedFile cachedFile = get(identifier);
+ if (cachedFile != null) {
+ return AsyncResources.ofCloseable(cachedFile);
+ }
+
+ if (loadingThreadPool.isAvailable()) {
+ final SettableAsyncResource<CachedFile> resource = new SettableAsyncResource<>();
+ final ListenableFuture<?> future = loadingThreadPool.getExecutorService().submit(
+ () -> {
+ try {
+ final Semaphore loadingPermits = loadingThreadPool.getPermits();
+ if (loadingPermits != null) {
+ loadingPermits.acquire();
+ }
+ try {
+ final CachedFile theCachedFile = reserveAndPopulate(identifier, sizeSupplier, populator);
+ if (!resource.set(ResourceHolder.fromCloseable(theCachedFile))) {
+ theCachedFile.close();
+ }
+ }
+ finally {
+ if (loadingPermits != null) {
+ loadingPermits.release();
+ }
+ }
+ }
+ catch (Throwable e) {
+ resource.setException(e);
+ }
+ }
+ );
+
+ resource.setCanceler(() -> future.cancel(true));
+ return resource;
+ } else {
+ final SettableAsyncResource<CachedFile> resource = new SettableAsyncResource<>();
+ try {
+ resource.set(ResourceHolder.fromCloseable(reserveAndPopulate(identifier, sizeSupplier, populator)));
+ }
+ catch (Throwable e) {
+ resource.setException(e);
+ }
+ return resource;
+ }
+ }
+
+ /**
+ * Lock object that tracks which StorageLocation was used for a given identifier.
+ * Once resolved, subsequent threads can check the resolved location first.
+ */
+ private static class PopulationLock
+ {
+ private final AtomicReference<StorageLocation> resolvedLocation = new AtomicReference<>();
+
+ @Nullable
+ StorageLocation getResolvedLocation()
+ {
+ return resolvedLocation.get();
+ }
+
+ void setResolvedLocation(StorageLocation location)
+ {
+ if (!resolvedLocation.compareAndSet(null, location)) {
+ throw DruidException.defensive(
+ "Resolved location already set to [%s], cannot change to [%s]",
+ resolvedLocation.get().getPath(),
+ location.getPath()
+ );
+ }
+ }
+ }
+}
diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java b/server/src/main/java/org/apache/druid/segment/loading/external/StringCacheEntryIdentifier.java
similarity index 58%
copy from server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java
copy to server/src/main/java/org/apache/druid/segment/loading/external/StringCacheEntryIdentifier.java
index bd9c8c65def..906f6b96b9d 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java
+++ b/server/src/main/java/org/apache/druid/segment/loading/external/StringCacheEntryIdentifier.java
@@ -17,27 +17,18 @@
* under the License.
*/
-package org.apache.druid.segment.loading;
+package org.apache.druid.segment.loading.external;
-import org.junit.Assert;
-import org.junit.Test;
+import org.apache.druid.segment.loading.CacheEntryIdentifier;
-public class SegmentLoaderConfigTest
+/**
+ * Simple CacheEntryIdentifier implementation that wraps a string.
+ */
+public record StringCacheEntryIdentifier(String value) implements CacheEntryIdentifier
{
- @Test
- public void testSetVirtualStorage()
+ @Override
+ public String toString()
{
- final SegmentLoaderConfig config = new SegmentLoaderConfig();
-
- // Verify default values
- Assert.assertFalse(config.isVirtualStorage());
- Assert.assertFalse(config.isVirtualStorageEphemeral());
-
- // Set both to true
- config.setVirtualStorage(true, true);
-
- // Verify both fields are set
- Assert.assertTrue(config.isVirtualStorage());
- Assert.assertTrue(config.isVirtualStorageEphemeral());
+ return value;
}
}
diff --git a/server/src/main/java/org/apache/druid/segment/loading/external/VirtualStorageManager.java b/server/src/main/java/org/apache/druid/segment/loading/external/VirtualStorageManager.java
new file mode 100644
index 00000000000..074b132d127
--- /dev/null
+++ b/server/src/main/java/org/apache/druid/segment/loading/external/VirtualStorageManager.java
@@ -0,0 +1,97 @@
+/*
+ * 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.external;
+
+import org.apache.druid.common.asyncresource.AsyncResource;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.io.FilePopulator;
+import org.apache.druid.segment.loading.StorageLocation;
+import org.apache.druid.utils.Throwables;
+
+import javax.annotation.Nullable;
+import java.util.function.LongSupplier;
+
+/**
+ * Manages a cache of downloaded files across multiple {@link StorageLocation}.
+ * Provides a simple API for checking if files exist and reserving space for new downloads.
+ */
+public interface VirtualStorageManager
+{
+ /**
+ * Check if a cached file exists and acquire a hold on it to prevent eviction.
+ *
+ * @param identifier Unique identifier for the cached file
+ * @return CachedFile handle if the file exists, or null if not found
+ */
+ @Nullable
+ CachedFile get(String identifier);
+
+ /**
+ * Reserve space for a new file, populate it via the provided lambda, and return
+ * a handle to access it. Population is done in the calling thread.
+ *
+ * <p>If an entry with the same identifier already exists, returns the existing entry
+ * without calling the populator.
+ *
+ * @param identifier Unique identifier for the cached file
+ * @param sizeSupplier Supplier that provides the number of bytes to reserve
+ * @param populator Lambda that will be called with a File to populate the content
+ * @return CachedFile handle to access the populated file
+ */
+ CachedFile reserveAndPopulate(
+ String identifier,
+ LongSupplier sizeSupplier,
+ FilePopulator populator
+ );
+
+ /**
+ * Reserve space for a new file, populate it via the provided lambda, and return
+ * a handle to access it. Population is done asynchronously.
+ *
+ * <p>Close the {@link AsyncResource} that is returned from
+ * this method when done. Do not close the inner {@link CachedFile}, as this will lead
+ * to a double-close.
+ *
+ * <p>If an entry with the same identifier already exists, returns the existing entry
+ * without calling the populator.
+ *
+ * @param identifier Unique identifier for the cached file
+ * @param sizeSupplier Supplier that provides the number of bytes to reserve
+ * @param populator Lambda that will be called with a File to populate the content
+ *
+ * @return handle to access the populated file
+ */
+ AsyncResource<CachedFile> reserveAndPopulateAsync(
+ String identifier,
+ LongSupplier sizeSupplier,
+ FilePopulator populator
+ );
+
+ /**
+ * Whether the given throwable indicates that a reservation failed because there was not enough space in the
+ * storage locations to hold a file.
+ */
+ static boolean isInsufficientStorage(final Throwable e)
+ {
+ final DruidException druidException = Throwables.getCauseOfType(e, DruidException.class);
+ return druidException != null
+ && druidException.getCategory() == DruidException.Category.CAPACITY_EXCEEDED;
+ }
+}
diff --git a/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java b/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java
index ae118afa7bb..1fed47d38ea 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java
+++ b/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java
@@ -121,6 +121,12 @@ public class NoopSegmentCacheManager implements SegmentCacheManager
@Override
public List<StorageLocation> getLocations()
{
- return List.of();
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public StorageLoadingThreadPool getLoadingThreadPool()
+ {
+ throw new UnsupportedOperationException();
}
}
diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java
index bd9c8c65def..cda83b1ebb7 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java
+++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLoaderConfigTest.java
@@ -34,7 +34,7 @@ public class SegmentLoaderConfigTest
Assert.assertFalse(config.isVirtualStorageEphemeral());
// Set both to true
- config.setVirtualStorage(true, true);
+ config.setVirtualStorage(true).setVirtualStorageIsEphemeral(true);
// Verify both fields are set
Assert.assertTrue(config.isVirtualStorage());
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 dc4fe4b384f..ddc6094bf6c 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
@@ -175,6 +175,7 @@ class SegmentLocalCacheManagerConcurrencyTest
manager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestIndex.INDEX_IO,
jsonMapper
@@ -182,6 +183,7 @@ class SegmentLocalCacheManagerConcurrencyTest
virtualStorageManager = new SegmentLocalCacheManager(
storageLocations,
vsfLoaderConfig,
+ StorageLoadingThreadPool.createFromConfig(vsfLoaderConfig),
new RoundRobinStorageLocationSelectorStrategy(storageLocations),
TestIndex.INDEX_IO,
jsonMapper
diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerTest.java
index 79c84af538b..c6b7e61f12a 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerTest.java
+++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerTest.java
@@ -132,6 +132,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
manager = new SegmentLocalCacheManager(
ImmutableList.of(),
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(ImmutableList.of()),
TestIndex.INDEX_IO,
jsonMapper
@@ -145,9 +146,11 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
final ImmutableList<StorageLocation> locations = ImmutableList.of(
new StorageLocation(localSegmentCacheDir, 10000000000L, null)
);
+ final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig();
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
locations,
- new SegmentLoaderConfig(),
+ loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(locations),
TestIndex.INDEX_IO,
jsonMapper
@@ -158,9 +161,11 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
@Test
public void testCanHandleSegmentsWithEmptyLocationsAndConfigLocations()
{
+ final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig();
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
ImmutableList.of(),
- new SegmentLoaderConfig(),
+ loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(ImmutableList.of()),
TestIndex.INDEX_IO,
jsonMapper
@@ -171,9 +176,11 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
@Test
public void testGetCachedSegmentsWhenCanHandleSegmentsIsFalse()
{
+ final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig();
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
null,
- new SegmentLoaderConfig(),
+ loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(null),
TestIndex.INDEX_IO,
jsonMapper
@@ -361,6 +368,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
locations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(locations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -412,6 +420,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -449,7 +458,8 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
final List<StorageLocation> storageLocations = loaderConfig.toStorageLocations();
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
- new SegmentLoaderConfig().setLocations(locations),
+ loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -489,7 +499,8 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
final List<StorageLocation> storageLocations = loaderConfig.toStorageLocations();
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
- new SegmentLoaderConfig().setLocations(locations),
+ loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -528,7 +539,8 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
final List<StorageLocation> storageLocations = loaderConfig.toStorageLocations();
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
- new SegmentLoaderConfig().setLocations(locations),
+ loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -580,9 +592,11 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
);
}
+ final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig().setLocations(locationConfigs);
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
locations,
- new SegmentLoaderConfig().setLocations(locationConfigs),
+ loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new RoundRobinStorageLocationSelectorStrategy(locations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -661,6 +675,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
new SegmentLoaderConfig().setLocations(locations),
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -743,6 +758,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
locations,
segmentLoaderConfig,
+ StorageLoadingThreadPool.createFromConfig(segmentLoaderConfig),
new RandomStorageLocationSelectorStrategy(locations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -824,6 +840,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -861,6 +878,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -908,6 +926,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -977,6 +996,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
() -> new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -1023,6 +1043,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -1084,6 +1105,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -1172,6 +1194,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -1256,6 +1279,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
final SegmentLocalCacheManager manager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
@@ -1380,6 +1404,7 @@ public class SegmentLocalCacheManagerTest extends InitializedNullHandlingTest
return new SegmentLocalCacheManager(
locations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(locations),
TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
jsonMapper
diff --git a/server/src/test/java/org/apache/druid/segment/loading/external/DownloadableCacheEntryTest.java b/server/src/test/java/org/apache/druid/segment/loading/external/DownloadableCacheEntryTest.java
new file mode 100644
index 00000000000..7440cf91e4a
--- /dev/null
+++ b/server/src/test/java/org/apache/druid/segment/loading/external/DownloadableCacheEntryTest.java
@@ -0,0 +1,225 @@
+/*
+ * 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.external;
+
+import com.google.common.hash.Hashing;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class DownloadableCacheEntryTest
+{
+ @Test
+ public void test_simpleName_splitInHalf()
+ {
+ final String path = "foo";
+ final String expected = "f-" + hash(path) + "-oo";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_singleExtension_suffixIsExtension()
+ {
+ final String path = "foo.gz";
+ final String expected = "foo-" + hash(path) + "-.gz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_chainedExtensionsShortBasename_suffixIsFullExtensionChain()
+ {
+ final String path = "a.log.gz";
+ final String expected = "a-" + hash(path) + "-.log.gz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_chainedExtensionsLongBasename_suffixIncludesFullExtensionChain()
+ {
+ final String path = "container.log.gz";
+ final String expected = "containe-" + hash(path) + "-r.log.gz";
+ final String result = DownloadableCacheEntry.sanitizePath(path);
+ Assertions.assertEquals(expected, result);
+ Assertions.assertTrue(result.endsWith(".log.gz"));
+ }
+
+ @Test
+ public void test_uriWithExtensions_slashesReplacedAndSuffixEndsWithExtensionChain()
+ {
+ final String path = "file:/var/folders/raw/container.log.gz";
+ final String expected = "file__var_folders_r-" + hash(path) + "-aw_container.log.gz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_nonAlnumChar_breaksExtensionChain_onlyTrailingAlnumExtensionsCount()
+ {
+ final String path = "foo.my_thing.gz";
+ final String expected = "foo.my_-" + hash(path) + "-thing.gz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_underscoreInBasename_noTrailingExtension_splitInHalf()
+ {
+ final String path = "foo.bar_baz";
+ final String expected = "foo.b-" + hash(path) + "-ar_baz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_dotsInDirectoryComponents_ignored_basenameHasNoExtension()
+ {
+ final String path = "foo.bar/baz";
+ final String expected = "foo.b-" + hash(path) + "-ar_baz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_dotsInDirectoryAndBasename_onlyBasenameExtensionsConsidered()
+ {
+ final String path = "foo.bar/baz.qux";
+ final String expected = "foo.bar-" + hash(path) + "-_baz.qux";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_trailingDot_notTreatedAsExtension()
+ {
+ final String path = "foo.bar.";
+ final String expected = "foo.-" + hash(path) + "-bar.";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_dotfile_leadingDotNotTreatedAsExtension()
+ {
+ final String path = "foo/.gitignore";
+ final String expected = "foo_.gi-" + hash(path) + "-tignore";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_dotfileWithExtension_leadingDotIgnoredOnlyTrailingExtensionUsed()
+ {
+ final String path = "foo/.hidden.gz";
+ final String expected = "foo_.hi-" + hash(path) + "-dden.gz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_basenameIsBareDotfile_noExtension()
+ {
+ final String path = ".gitignore";
+ final String expected = ".giti-" + hash(path) + "-gnore";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_backslashesNormalizedToSlashes()
+ {
+ final String path = "foo\\bar\\baz.gz";
+ final String expected = "foo_bar-" + hash(path) + "-_baz.gz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_leadingSlash_stripped()
+ {
+ final String path = "/foo/bar.gz";
+ final String expected = "foo_b-" + hash(path) + "-ar.gz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_trailingSlash_dropped()
+ {
+ final String path = "foo/bar/";
+ final String expected = "foo-" + hash(path) + "-_bar";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_longPath_truncatedToFirst50AndLast50_suffixEndsWithExtension()
+ {
+ final StringBuilder pathBuilder = new StringBuilder("/aaaa");
+ while (pathBuilder.length() < 200) {
+ pathBuilder.append("/bbbb");
+ }
+ pathBuilder.append("/finalfile.gz");
+ final String path = pathBuilder.toString();
+
+ final String result = DownloadableCacheEntry.sanitizePath(path);
+
+ Assertions.assertTrue(
+ result.endsWith(".gz"),
+ "truncated output should still end with the original extension, got: " + result
+ );
+ Assertions.assertTrue(
+ result.contains("-" + hash(path) + "-"),
+ "hashCode should be embedded between '-' delimiters between the prefix and suffix, got: " + result
+ );
+ Assertions.assertEquals(
+ 100 + 2 + hash(path).length(),
+ result.length(),
+ "kept portion should be exactly 100 chars (50 each side) plus delimiters and hashcode"
+ );
+ }
+
+ @Test
+ public void test_disallowedChars_replacedWithUnderscores()
+ {
+ final String path = "foo bar*baz?.gz";
+ final String expected = "foo_bar-" + hash(path) + "-_baz_.gz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_hyphenPreserved()
+ {
+ final String path = "foo-bar.gz";
+ final String expected = "foo-b-" + hash(path) + "-ar.gz";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_emptyPath_resultIsJustHashWithDelimiters()
+ {
+ final String path = "";
+ final String expected = "-" + hash(path) + "-";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ @Test
+ public void test_slashOnly_resultIsJustHashWithDelimiters()
+ {
+ final String path = "/";
+ final String expected = "-" + hash(path) + "-";
+ Assertions.assertEquals(expected, DownloadableCacheEntry.sanitizePath(path));
+ }
+
+ /**
+ * Computes the disambiguating hash the way {@link DownloadableCacheEntry#sanitizePath(String)} does, so the tests
+ * below can assert on the prefix/suffix splitting (the actual behavior under test) without hardcoding hash values.
+ */
+ private static String hash(final String path)
+ {
+ return Hashing.sha512().hashUnencodedChars(path).toString().substring(0, 16);
+ }
+}
diff --git a/server/src/test/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManagerTest.java b/server/src/test/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManagerTest.java
new file mode 100644
index 00000000000..7c0d5cbdf32
--- /dev/null
+++ b/server/src/test/java/org/apache/druid/segment/loading/external/StorageLocationVirtualStorageManagerTest.java
@@ -0,0 +1,1013 @@
+/*
+ * 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.external;
+
+import org.apache.druid.common.asyncresource.AsyncResource;
+import org.apache.druid.common.asyncresource.AsyncResources;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.error.DruidExceptionMatcher;
+import org.apache.druid.java.util.common.concurrent.Execs;
+import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStrategy;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
+import org.apache.druid.segment.loading.StorageLocation;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class StorageLocationVirtualStorageManagerTest
+{
+ @TempDir
+ public File tempFolder;
+
+ private StorageLocation location;
+ private StorageLocationVirtualStorageManager manager;
+
+ @BeforeEach
+ public void setup() throws IOException
+ {
+ File locationPath = new File(tempFolder, "storage");
+ Files.createDirectories(locationPath.toPath());
+ location = new StorageLocation(locationPath, 10_000_000L, null);
+ manager = new StorageLocationVirtualStorageManager(
+ Collections.singletonList(location),
+ new LeastBytesUsedStorageLocationSelectorStrategy(Collections.singletonList(location)),
+ StorageLoadingThreadPool.none()
+ );
+ }
+
+ @AfterEach
+ public void teardown()
+ {
+ location.reset();
+ }
+
+ /**
+ * Helper method to run test logic with an ExecutorService that is guaranteed to be shut down.
+ */
+ private void withExecutor(ExecutorConsumer testLogic) throws Exception
+ {
+ ExecutorService executorService = Execs.multiThreaded(10, "test-thread-%d");
+ try {
+ testLogic.accept(executorService);
+ }
+ finally {
+ executorService.shutdownNow();
+ }
+ }
+
+ @FunctionalInterface
+ private interface ExecutorConsumer
+ {
+ void accept(ExecutorService executor) throws Exception;
+ }
+
+ @Test
+ public void testGetReturnsNullWhenNotExists()
+ {
+ CachedFile result = manager.get("nonexistent");
+ Assertions.assertNull(result);
+ }
+
+ @Test
+ public void testReserveAndPopulateCreatesFile() throws Exception
+ {
+ String identifier = "test-file";
+ String content = "Hello, World!";
+
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(0, location.getWeakEntryCount());
+
+ try (CachedFile cachedFile = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8))
+ )) {
+ Assertions.assertNotNull(cachedFile);
+ Assertions.assertEquals(identifier, cachedFile.getIdentifier());
+
+ // Validate there's an active hold while the file is open
+ Assertions.assertEquals(1, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+
+ File file = cachedFile.getFile();
+ Assertions.assertNotNull(file);
+ Assertions.assertTrue(file.exists());
+
+ String readContent = new String(Files.readAllBytes(file.toPath()), StandardCharsets.UTF_8);
+ Assertions.assertEquals(content, readContent);
+ }
+
+ // After close, hold should be released
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ // Entry remains in cache, just not held
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+ }
+
+ @Test
+ public void testGetReturnsFileAfterPopulate() throws Exception
+ {
+ String identifier = "test-file";
+ String content = "Test content";
+
+ // First populate
+ try (CachedFile cachedFile = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8))
+ )) {
+ Assertions.assertNotNull(cachedFile);
+ Assertions.assertEquals(1, location.getWeakStats().getHoldCount());
+ }
+
+ // After close, hold should be released but entry remains
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+
+ // Then get
+ try (CachedFile cachedFile = manager.get(identifier)) {
+ Assertions.assertNotNull(cachedFile);
+ Assertions.assertEquals(identifier, cachedFile.getIdentifier());
+
+ // Hold should be re-acquired
+ Assertions.assertEquals(1, location.getWeakStats().getHoldCount());
+
+ String readContent = new String(
+ Files.readAllBytes(cachedFile.getFile().toPath()),
+ StandardCharsets.UTF_8
+ );
+ Assertions.assertEquals(content, readContent);
+ }
+
+ // After second close, hold should be released again
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+ }
+
+ @Test
+ public void testGetReturnsNullAfterEphemeralEviction() throws Exception
+ {
+ location.setAreWeakEntriesEphemeral(true);
+
+ String identifier = "test-file";
+ String content = "Hello, World!";
+
+ try (CachedFile cachedFile = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8))
+ )) {
+ try (CachedFile fromGet = manager.get(identifier)) {
+ Assertions.assertNotNull(fromGet);
+ }
+ }
+ // Last hold released → ephemeral eviction → entry unmounted
+
+ Assertions.assertNull(manager.get(identifier));
+ Assertions.assertEquals(0, location.getWeakEntryCount());
+ }
+
+ @Test
+ public void testPopulationLockClearedOnEphemeralEviction() throws Exception
+ {
+ location.setAreWeakEntriesEphemeral(true);
+
+ String identifier = "test-file";
+ AtomicInteger firstSupplierCalls = new AtomicInteger();
+ AtomicInteger secondSupplierCalls = new AtomicInteger();
+ AtomicInteger firstPopulatorCalls = new AtomicInteger();
+ AtomicInteger secondPopulatorCalls = new AtomicInteger();
+
+ try (CachedFile firstFile = manager.reserveAndPopulate(
+ identifier,
+ () -> {
+ firstSupplierCalls.incrementAndGet();
+ return 4L;
+ },
+ file -> {
+ firstPopulatorCalls.incrementAndGet();
+ Files.write(file.toPath(), "v1".getBytes(StandardCharsets.UTF_8));
+ }
+ )) {
+ Assertions.assertNotNull(firstFile);
+ }
+ // Hold released → ephemeral eviction → unmount → populationLocks.remove(id, lock)
+
+ try (CachedFile secondFile = manager.reserveAndPopulate(
+ identifier,
+ () -> {
+ secondSupplierCalls.incrementAndGet();
+ return 4L;
+ },
+ file -> {
+ secondPopulatorCalls.incrementAndGet();
+ Files.write(file.toPath(), "v2".getBytes(StandardCharsets.UTF_8));
+ }
+ )) {
+ Assertions.assertNotNull(secondFile);
+ String readContent = new String(
+ Files.readAllBytes(secondFile.getFile().toPath()),
+ StandardCharsets.UTF_8
+ );
+ Assertions.assertEquals("v2", readContent);
+ }
+
+ Assertions.assertEquals(1, firstSupplierCalls.get(), "first reserve must be a miss");
+ Assertions.assertEquals(1, firstPopulatorCalls.get(), "first populator must run");
+ Assertions.assertEquals(
+ 1,
+ secondSupplierCalls.get(),
+ "second reserve must also be a miss — proving the lock was cleared on unmount"
+ );
+ Assertions.assertEquals(1, secondPopulatorCalls.get(), "second populator must run");
+ }
+
+ @Test
+ public void testPopulationLockClearedOnCapacityEviction() throws Exception
+ {
+ long capacity = location.availableSizeBytes();
+ long smallSize = capacity / 4;
+ long bigSize = capacity - 100; // forces reclaim of the small entry
+
+ String smallId = "small";
+ String bigId = "big";
+
+ AtomicInteger populatorCalls = new AtomicInteger();
+ try (CachedFile smallFile = manager.reserveAndPopulate(
+ smallId,
+ () -> smallSize,
+ file -> {
+ populatorCalls.incrementAndGet();
+ Files.write(file.toPath(), new byte[]{1, 2, 3});
+ }
+ )) {
+ Assertions.assertNotNull(smallFile);
+ }
+ // Small entry is now weakly cached, eligible for eviction.
+
+ try (CachedFile bigFile = manager.reserveAndPopulate(
+ bigId,
+ () -> bigSize,
+ file -> {
+ populatorCalls.incrementAndGet();
+ Files.write(file.toPath(), new byte[]{4, 5, 6});
+ }
+ )) {
+ Assertions.assertNotNull(bigFile);
+ }
+ // Reserving `big` should have forced eviction of `small`, invoking unmount → cleanup.
+
+ AtomicInteger supplierCalls = new AtomicInteger();
+ try (CachedFile reborn = manager.reserveAndPopulate(
+ smallId,
+ () -> {
+ supplierCalls.incrementAndGet();
+ return smallSize;
+ },
+ file -> {
+ populatorCalls.incrementAndGet();
+ Files.write(file.toPath(), new byte[]{7, 8, 9});
+ }
+ )) {
+ Assertions.assertNotNull(reborn);
+ }
+
+ Assertions.assertEquals(
+ 1,
+ supplierCalls.get(),
+ "re-reserving an evicted id must take the miss path"
+ );
+ Assertions.assertEquals(3, populatorCalls.get());
+ }
+
+ @Test
+ public void testPopulationLockClearedOnLocationReset() throws Exception
+ {
+ // location.reset() unmounts every weakly-held entry; the cleanup hook must fire for each.
+ String identifier = "test-file";
+ String content = "Hello, World!";
+
+ try (CachedFile cachedFile = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8))
+ )) {
+ Assertions.assertNotNull(cachedFile);
+ }
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+
+ location.reset();
+ Assertions.assertEquals(0, location.getWeakEntryCount());
+
+ AtomicInteger supplierCalls = new AtomicInteger();
+ AtomicInteger populatorCalls = new AtomicInteger();
+ try (CachedFile reborn = manager.reserveAndPopulate(
+ identifier,
+ () -> {
+ supplierCalls.incrementAndGet();
+ return (long) content.length();
+ },
+ file -> {
+ populatorCalls.incrementAndGet();
+ Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8));
+ }
+ )) {
+ Assertions.assertNotNull(reborn);
+ }
+
+ Assertions.assertEquals(
+ 1,
+ supplierCalls.get(),
+ "re-reserving after location.reset() must take the miss path"
+ );
+ Assertions.assertEquals(1, populatorCalls.get());
+ }
+
+ @Test
+ public void testMultipleCallsToReserveAndPopulateWithSameIdentifier() throws Exception
+ {
+ String identifier = "test-file";
+ String content1 = "First content";
+ String content2 = "Second content";
+
+ // First call
+ try (CachedFile cachedFile1 = manager.reserveAndPopulate(
+ identifier,
+ content1::length,
+ file -> Files.write(file.toPath(), content1.getBytes(StandardCharsets.UTF_8))
+ )) {
+ Assertions.assertNotNull(cachedFile1);
+ Assertions.assertEquals(1, location.getWeakStats().getHoldCount());
+
+ // Second call with same identifier should return existing file
+ try (CachedFile cachedFile2 = manager.reserveAndPopulate(
+ identifier,
+ content2::length,
+ file -> {
+ // This populator should NOT be called
+ throw new RuntimeException("Populator should not be called for existing file");
+ }
+ )) {
+ Assertions.assertNotNull(cachedFile2);
+
+ Assertions.assertEquals(2, location.getWeakStats().getHoldCount());
+
+ // Should have first content, not second
+ String readContent = new String(
+ Files.readAllBytes(cachedFile2.getFile().toPath()),
+ StandardCharsets.UTF_8
+ );
+ Assertions.assertEquals(content1, readContent);
+ }
+
+ // Inner hold released, outer hold still active
+ Assertions.assertEquals(1, location.getWeakStats().getHoldCount());
+ }
+
+ // All holds released
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+ }
+
+ @Test
+ public void testSizeSupplierIsNotCalledOnCacheHit()
+ {
+ String identifier = "test-file";
+ String content = "Hello, World!";
+ AtomicInteger missSupplierCalls = new AtomicInteger();
+ AtomicInteger hitSupplierCalls = new AtomicInteger();
+
+ try (CachedFile first = manager.reserveAndPopulate(
+ identifier,
+ () -> {
+ missSupplierCalls.incrementAndGet();
+ return content.length();
+ },
+ file -> Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8))
+ )) {
+ Assertions.assertNotNull(first);
+ Assertions.assertEquals(1, missSupplierCalls.get(), "supplier should be invoked on cache miss");
+
+ // Hit path: supplier must not be invoked.
+ try (CachedFile second = manager.reserveAndPopulate(
+ identifier,
+ () -> {
+ hitSupplierCalls.incrementAndGet();
+ return content.length();
+ },
+ ignored -> {
+ throw new RuntimeException("Populator should not be called for existing file");
+ }
+ )) {
+ Assertions.assertNotNull(second);
+ }
+ }
+
+ Assertions.assertEquals(
+ 0,
+ hitSupplierCalls.get(),
+ "size supplier must not be invoked when the cached entry already exists"
+ );
+ }
+
+ @Test
+ public void testSizeSupplierIsCalledExactlyOnceOnCacheMiss()
+ {
+ // Companion to the cache-hit case: a successful miss path should evaluate the
+ // supplier exactly once, even though the implementation iterates storage locations
+ // looking for one with capacity.
+ String identifier = "test-file";
+ String content = "Hello, World!";
+ AtomicInteger supplierCalls = new AtomicInteger();
+
+ try (CachedFile cachedFile = manager.reserveAndPopulate(
+ identifier,
+ () -> {
+ supplierCalls.incrementAndGet();
+ return content.length();
+ },
+ file -> Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8))
+ )) {
+ Assertions.assertNotNull(cachedFile);
+ }
+
+ Assertions.assertEquals(1, supplierCalls.get(), "size supplier must be invoked exactly once on cache miss");
+ }
+
+ @Test
+ public void testPopulatorThrowsException()
+ {
+ String identifier = "test-file";
+
+ new DruidExceptionMatcher(
+ DruidException.Persona.OPERATOR,
+ DruidException.Category.RUNTIME_FAILURE,
+ "general"
+ )
+ .expectMessageContains("populate")
+ .assertThrowsAndMatches(() -> manager.reserveAndPopulate(
+ identifier,
+ () -> 100,
+ file -> {
+ throw new IOException("Populator failed");
+ }
+ ));
+
+ // Note: The entry may still be in the cache but in an unmounted state.
+ // StorageLocation will eventually evict it. The important thing is that
+ // the exception was properly thrown and the file was not successfully created.
+ }
+
+ @Test
+ public void testInsufficientSpace()
+ {
+ String identifier = "huge-file";
+ long hugeSize = location.availableSizeBytes() + 1;
+
+ new DruidExceptionMatcher(
+ DruidException.Persona.OPERATOR,
+ DruidException.Category.CAPACITY_EXCEEDED,
+ "general"
+ )
+ .expectMessageContains("No space available")
+ .assertThrowsAndMatches(() -> manager.reserveAndPopulate(
+ identifier,
+ () -> hugeSize,
+ file -> Files.write(file.toPath(), new byte[1000])
+ ));
+ }
+
+ @Test
+ public void testReserveAndPopulateAsyncRecoversFromInsufficientStorage()
+ {
+ // Verify something similar to what ExternalInputSliceReader does: fall back to recovery logic if the
+ // storage location cannot hold a file.
+ final long tooLargeFileSize = location.availableSizeBytes() + 1;
+
+ final AsyncResource<String> recovered = AsyncResources.recover(
+ AsyncResources.transform(
+ manager.reserveAndPopulateAsync(
+ "huge-file",
+ () -> tooLargeFileSize,
+ file -> Files.write(file.toPath(), new byte[]{1, 2, 3})
+ ),
+ cachedFile -> "fetched"
+ ),
+ e -> VirtualStorageManager.isInsufficientStorage(e) ? "recovery" : null
+ );
+
+ try (recovered) {
+ Assertions.assertTrue(recovered.isReady());
+ Assertions.assertEquals("recovery", recovered.get());
+ }
+ }
+
+ @Test
+ public void testReserveAndPopulateAsyncDoesNotRecoverFromOtherErrors()
+ {
+ // A populate failure that is not an insufficient-storage error must be propagated, not recovered: the recovery
+ // function returns null, so the original error surfaces from get().
+ final AsyncResource<String> recovered = AsyncResources.recover(
+ AsyncResources.transform(
+ manager.reserveAndPopulateAsync(
+ "boom-file",
+ () -> 100L,
+ file -> {
+ throw new IOException("populate failed");
+ }
+ ),
+ cachedFile -> "fetched"
+ ),
+ e -> VirtualStorageManager.isInsufficientStorage(e) ? "streamed" : null
+ );
+
+ try (recovered) {
+ Assertions.assertTrue(recovered.isReady());
+ final DruidException e = Assertions.assertThrows(DruidException.class, recovered::get);
+ Assertions.assertEquals(DruidException.Category.RUNTIME_FAILURE, e.getCategory());
+ }
+ }
+
+ @Test
+ public void testFilePathSanitization()
+ {
+ String identifier = "path/with/slashes/and-special_chars.txt";
+ String content = "Test content";
+
+ try (CachedFile cachedFile = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8))
+ )) {
+ File file = cachedFile.getFile();
+ Assertions.assertNotNull(file);
+ Assertions.assertTrue(file.exists());
+
+ // Verify file is in a subdirectory structure
+ Assertions.assertTrue(file.getAbsolutePath().contains("path"));
+ Assertions.assertTrue(file.getAbsolutePath().contains("slashes"));
+ }
+ }
+
+ // ========== LIFECYCLE TESTS ==========
+
+ @Test
+ public void testLifecycleWithMultipleFiles()
+ {
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(0, location.getWeakEntryCount());
+
+ // Create multiple files
+ CachedFile file1 = manager.reserveAndPopulate(
+ "file1",
+ () -> 5,
+ file -> Files.write(file.toPath(), "File1".getBytes(StandardCharsets.UTF_8))
+ );
+ Assertions.assertEquals(1, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+
+ CachedFile file2 = manager.reserveAndPopulate(
+ "file2",
+ () -> 5,
+ file -> Files.write(file.toPath(), "File2".getBytes(StandardCharsets.UTF_8))
+ );
+ Assertions.assertEquals(2, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(2, location.getWeakEntryCount());
+
+ CachedFile file3 = manager.reserveAndPopulate(
+ "file3",
+ () -> 5,
+ file -> Files.write(file.toPath(), "File3".getBytes(StandardCharsets.UTF_8))
+ );
+ Assertions.assertEquals(3, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(3, location.getWeakEntryCount());
+
+ // Close one file
+ file1.close();
+ Assertions.assertEquals(2, location.getWeakStats().getHoldCount());
+ // Entry remains
+ Assertions.assertEquals(3, location.getWeakEntryCount());
+
+ // Close another
+ file2.close();
+ Assertions.assertEquals(1, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(3, location.getWeakEntryCount());
+
+ // Close last
+ file3.close();
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(3, location.getWeakEntryCount());
+ }
+
+ @Test
+ public void testLifecycleWithMultipleHoldsOnSameFile()
+ {
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(0, location.getWeakEntryCount());
+
+ String identifier = "shared-file";
+ String content = "Shared content";
+
+ // Create first hold
+ CachedFile hold1 = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8))
+ );
+ Assertions.assertEquals(1, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+
+ // Get second hold on same file
+ CachedFile hold2 = manager.get(identifier);
+ Assertions.assertNotNull(hold2);
+ Assertions.assertEquals(2, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+
+ // Get third hold
+ CachedFile hold3 = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> {
+ throw new RuntimeException("Should not populate again");
+ }
+ );
+ Assertions.assertEquals(3, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+
+ // Release holds one by one - entry stays held until all are released
+ hold1.close();
+ // Entry is still held by other handles
+ Assertions.assertEquals(2, location.getWeakStats().getHoldCount());
+
+ hold2.close();
+ Assertions.assertEquals(1, location.getWeakStats().getHoldCount());
+
+ // After last close, no holds remain
+ hold3.close();
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+ }
+
+ // ========== THREAD-SAFETY TESTS ==========
+
+ @Test
+ public void testConcurrentReserveAndPopulateSameIdentifier() throws Exception
+ {
+ withExecutor(executorService -> {
+ String identifier = "concurrent-file";
+ String content = "Test content";
+ int numThreads = 10;
+
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+
+ AtomicInteger populatorCallCount = new AtomicInteger(0);
+ CountDownLatch startLatch = new CountDownLatch(1);
+ CountDownLatch completionLatch = new CountDownLatch(numThreads);
+
+ List<Future<CachedFile>> futures = new ArrayList<>();
+
+ // Launch multiple threads that all try to populate the same identifier
+ for (int i = 0; i < numThreads; i++) {
+ Future<CachedFile> future = executorService.submit(() -> {
+ try {
+ // Wait for all threads to be ready
+ startLatch.await();
+
+ CachedFile result = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> {
+ populatorCallCount.incrementAndGet();
+ // Add a small sleep to increase chance of race conditions if locking is broken
+ try {
+ Thread.sleep(10);
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8));
+ }
+ );
+
+ return result;
+ }
+ finally {
+ completionLatch.countDown();
+ }
+ });
+ futures.add(future);
+ }
+
+ // Start all threads at once
+ startLatch.countDown();
+
+ // Wait for all to complete
+ Assertions.assertTrue(completionLatch.await(10, TimeUnit.SECONDS));
+
+ // Verify populator was called exactly once
+ Assertions.assertEquals(
+ 1,
+ populatorCallCount.get(),
+ "Populator should only be called once despite concurrent access"
+ );
+
+ Assertions.assertEquals(numThreads, location.getWeakStats().getHoldCount());
+
+ // Verify all threads got a valid CachedFile and close them
+ for (Future<CachedFile> future : futures) {
+ CachedFile cachedFile = future.get();
+ Assertions.assertNotNull(cachedFile);
+ cachedFile.close();
+ }
+
+ // After all closes, no active holds
+ Assertions.assertEquals(0, location.getWeakStats().getHoldCount());
+ Assertions.assertEquals(1, location.getWeakEntryCount());
+ });
+ }
+
+ @Test
+ public void testConcurrentGetWhilePopulating() throws Exception
+ {
+ withExecutor(executorService -> {
+ String identifier = "concurrent-get-file";
+ String content = "Test content";
+ int numGetters = 5;
+
+ CountDownLatch populateStarted = new CountDownLatch(1);
+ CountDownLatch populateCanComplete = new CountDownLatch(1);
+ CountDownLatch allComplete = new CountDownLatch(numGetters + 1);
+
+ AtomicInteger nullGetCount = new AtomicInteger(0);
+ AtomicInteger nonNullGetCount = new AtomicInteger(0);
+
+ // Start populator thread
+ Future<CachedFile> populateFuture = executorService.submit(() -> {
+ try {
+ CachedFile result = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> {
+ populateStarted.countDown();
+ // Wait before completing population
+ try {
+ populateCanComplete.await();
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8));
+ }
+ );
+ return result;
+ }
+ finally {
+ allComplete.countDown();
+ }
+ });
+
+ // Wait for population to start
+ Assertions.assertTrue(populateStarted.await(5, TimeUnit.SECONDS));
+
+ // Launch getter threads while population is in progress
+ List<Future<CachedFile>> getterFutures = new ArrayList<>();
+ for (int i = 0; i < numGetters; i++) {
+ Future<CachedFile> future = executorService.submit(() -> {
+ try {
+ CachedFile result = manager.get(identifier);
+ if (result == null) {
+ nullGetCount.incrementAndGet();
+ } else {
+ nonNullGetCount.incrementAndGet();
+ }
+ return result;
+ }
+ finally {
+ allComplete.countDown();
+ }
+ });
+ getterFutures.add(future);
+ }
+
+ // Let population complete
+ populateCanComplete.countDown();
+
+ // Wait for all threads
+ Assertions.assertTrue(allComplete.await(10, TimeUnit.SECONDS));
+
+ // Clean up
+ CachedFile populated = populateFuture.get();
+ Assertions.assertNotNull(populated);
+ populated.close();
+
+ for (Future<CachedFile> future : getterFutures) {
+ CachedFile result = future.get();
+ if (result != null) {
+ result.close();
+ }
+ }
+
+ // Some gets might return null (before population), some might succeed (after)
+ // Both are valid outcomes, we just verify no exceptions occurred
+ Assertions.assertEquals(nullGetCount.get() + nonNullGetCount.get(), numGetters);
+ });
+ }
+
+ @Test
+ public void testConcurrentPopulationOfDifferentIdentifiers() throws Exception
+ {
+ withExecutor(executorService -> {
+ int numIdentifiers = 10;
+ String contentPrefix = "Content for file ";
+
+ CyclicBarrier startBarrier = new CyclicBarrier(numIdentifiers);
+ CountDownLatch completionLatch = new CountDownLatch(numIdentifiers);
+ AtomicInteger successCount = new AtomicInteger(0);
+
+ List<Future<CachedFile>> futures = new ArrayList<>();
+
+ for (int i = 0; i < numIdentifiers; i++) {
+ final int fileNum = i;
+ Future<CachedFile> future = executorService.submit(() -> {
+ try {
+ // Wait for all threads to be ready
+ startBarrier.await();
+
+ String identifier = "file-" + fileNum;
+ String content = contentPrefix + fileNum;
+
+ CachedFile result = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> {
+ // Simulate some work
+ try {
+ Thread.sleep(5);
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8));
+ }
+ );
+
+ successCount.incrementAndGet();
+ return result;
+ }
+ finally {
+ completionLatch.countDown();
+ }
+ });
+ futures.add(future);
+ }
+
+ // Wait for all to complete
+ Assertions.assertTrue(completionLatch.await(10, TimeUnit.SECONDS));
+
+ // Verify all succeeded
+ Assertions.assertEquals(numIdentifiers, successCount.get());
+
+ // Verify all files are correct
+ for (int i = 0; i < numIdentifiers; i++) {
+ try (CachedFile cachedFile = futures.get(i).get()) {
+ Assertions.assertNotNull(cachedFile);
+
+ String expectedContent = contentPrefix + i;
+ String actualContent = new String(
+ Files.readAllBytes(cachedFile.getFile().toPath()),
+ StandardCharsets.UTF_8
+ );
+ Assertions.assertEquals(expectedContent, actualContent);
+ }
+ }
+ });
+ }
+
+ @Test
+ public void testResolvedLocationIsReused()
+ {
+ String identifier = "resolved-location-file";
+ String content = "Test content";
+
+ // First reservation
+ try (CachedFile cachedFile1 = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8))
+ )) {
+ Assertions.assertNotNull(cachedFile1);
+
+ // Second reservation should use the cached resolved location
+ // and find the existing file without trying other locations
+ try (CachedFile cachedFile2 = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> {
+ throw new RuntimeException("Should not populate again");
+ }
+ )) {
+ Assertions.assertNotNull(cachedFile2);
+
+ // Both should point to same underlying file
+ Assertions.assertEquals(
+ cachedFile1.getFile().getAbsolutePath(),
+ cachedFile2.getFile().getAbsolutePath()
+ );
+ }
+ }
+ }
+
+ @Test
+ public void testConcurrentResolvedLocationLookup() throws Exception
+ {
+ withExecutor(executorService -> {
+ String identifier = "concurrent-resolved-file";
+ String content = "Test content";
+ int numThreads = 10;
+
+ // First, populate the file
+ try (CachedFile initialFile = manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> Files.write(file.toPath(), content.getBytes(StandardCharsets.UTF_8))
+ )) {
+ Assertions.assertNotNull(initialFile);
+ }
+
+ // Now have multiple threads try to reserve the same identifier
+ CyclicBarrier startBarrier = new CyclicBarrier(numThreads);
+ CountDownLatch completionLatch = new CountDownLatch(numThreads);
+ AtomicInteger populatorCallCount = new AtomicInteger(0);
+
+ List<Future<CachedFile>> futures = new ArrayList<>();
+
+ for (int i = 0; i < numThreads; i++) {
+ Future<CachedFile> future = executorService.submit(() -> {
+ try {
+ startBarrier.await();
+
+ return manager.reserveAndPopulate(
+ identifier,
+ content::length,
+ file -> {
+ populatorCallCount.incrementAndGet();
+ throw new RuntimeException("Should not be called");
+ }
+ );
+ }
+ finally {
+ completionLatch.countDown();
+ }
+ });
+ futures.add(future);
+ }
+
+ Assertions.assertTrue(completionLatch.await(10, TimeUnit.SECONDS));
+
+ // Verify populator was never called (file already exists)
+ Assertions.assertEquals(0, populatorCallCount.get());
+
+ // Verify all got valid results
+ for (Future<CachedFile> future : futures) {
+ CachedFile cachedFile = future.get();
+ Assertions.assertNotNull(cachedFile);
+ cachedFile.close();
+ }
+ });
+ }
+}
diff --git a/server/src/test/java/org/apache/druid/server/SegmentManagerTest.java b/server/src/test/java/org/apache/druid/server/SegmentManagerTest.java
index 2246237f5a0..3b0d6c6e9db 100644
--- a/server/src/test/java/org/apache/druid/server/SegmentManagerTest.java
+++ b/server/src/test/java/org/apache/druid/server/SegmentManagerTest.java
@@ -49,6 +49,7 @@ import org.apache.druid.segment.loading.LocalLoadSpec;
import org.apache.druid.segment.loading.SegmentLoaderConfig;
import org.apache.druid.segment.loading.SegmentLoadingException;
import org.apache.druid.segment.loading.SegmentLocalCacheManager;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
import org.apache.druid.segment.loading.StorageLocation;
import org.apache.druid.segment.loading.StorageLocationConfig;
import org.apache.druid.server.SegmentManager.DataSourceState;
@@ -165,6 +166,7 @@ public class SegmentManagerTest extends InitializedNullHandlingTest
cacheManager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestIndex.INDEX_IO,
objectMapper
@@ -175,6 +177,7 @@ public class SegmentManagerTest extends InitializedNullHandlingTest
virtualCacheManager = new SegmentLocalCacheManager(
virtualStorageLocations,
virtualLoaderConfig,
+ StorageLoadingThreadPool.createFromConfig(virtualLoaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(virtualStorageLocations),
TestIndex.INDEX_IO,
objectMapper
diff --git a/server/src/test/java/org/apache/druid/server/coordination/SegmentCacheBootstrapperCacheTest.java b/server/src/test/java/org/apache/druid/server/coordination/SegmentCacheBootstrapperCacheTest.java
index bf8edc61956..7878630d2af 100644
--- a/server/src/test/java/org/apache/druid/server/coordination/SegmentCacheBootstrapperCacheTest.java
+++ b/server/src/test/java/org/apache/druid/server/coordination/SegmentCacheBootstrapperCacheTest.java
@@ -31,6 +31,7 @@ import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStr
import org.apache.druid.segment.loading.SegmentLoaderConfig;
import org.apache.druid.segment.loading.SegmentLoadingException;
import org.apache.druid.segment.loading.SegmentLocalCacheManager;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
import org.apache.druid.segment.loading.StorageLocation;
import org.apache.druid.segment.loading.StorageLocationConfig;
import org.apache.druid.server.SegmentManager;
@@ -99,6 +100,7 @@ public class SegmentCacheBootstrapperCacheTest
cacheManager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestIndex.INDEX_IO,
objectMapper
@@ -119,6 +121,7 @@ public class SegmentCacheBootstrapperCacheTest
new SegmentLocalCacheManager(
emptyLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(emptyLocations),
TestIndex.INDEX_IO,
objectMapper
diff --git a/server/src/test/java/org/apache/druid/server/coordination/SegmentManagerBroadcastJoinIndexedTableTest.java b/server/src/test/java/org/apache/druid/server/coordination/SegmentManagerBroadcastJoinIndexedTableTest.java
index fdcada7f941..5ee311c389b 100644
--- a/server/src/test/java/org/apache/druid/server/coordination/SegmentManagerBroadcastJoinIndexedTableTest.java
+++ b/server/src/test/java/org/apache/druid/server/coordination/SegmentManagerBroadcastJoinIndexedTableTest.java
@@ -55,6 +55,7 @@ import org.apache.druid.segment.loading.SegmentLoaderConfig;
import org.apache.druid.segment.loading.SegmentLoadingException;
import org.apache.druid.segment.loading.SegmentLocalCacheManager;
import org.apache.druid.segment.loading.SegmentizerFactory;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
import org.apache.druid.segment.loading.StorageLocation;
import org.apache.druid.segment.loading.StorageLocationConfig;
import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory;
@@ -149,6 +150,7 @@ public class SegmentManagerBroadcastJoinIndexedTableTest extends InitializedNull
segmentCacheManager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestIndex.INDEX_IO,
objectMapper
diff --git a/server/src/test/java/org/apache/druid/server/coordination/SegmentManagerThreadSafetyTest.java b/server/src/test/java/org/apache/druid/server/coordination/SegmentManagerThreadSafetyTest.java
index 43c79372c50..8695cef6fcf 100644
--- a/server/src/test/java/org/apache/druid/server/coordination/SegmentManagerThreadSafetyTest.java
+++ b/server/src/test/java/org/apache/druid/server/coordination/SegmentManagerThreadSafetyTest.java
@@ -43,6 +43,7 @@ import org.apache.druid.segment.loading.SegmentLoaderConfig;
import org.apache.druid.segment.loading.SegmentLoadingException;
import org.apache.druid.segment.loading.SegmentLocalCacheManager;
import org.apache.druid.segment.loading.SegmentizerFactory;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
import org.apache.druid.segment.loading.StorageLocation;
import org.apache.druid.segment.loading.StorageLocationConfig;
import org.apache.druid.server.SegmentManager;
@@ -119,6 +120,7 @@ public class SegmentManagerThreadSafetyTest
segmentCacheManager = new SegmentLocalCacheManager(
storageLocations,
loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
TestIndex.INDEX_IO,
objectMapper
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org