fix: dedicated ephemeral storage loading pool for tasks

indexing-service/src/main/java/org/apache/druid/indexing/common/SegmentCacheManagerFactory.java

@@ -20,9 +20,12 @@
 package org.apache.druid.indexing.common;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.inject.Inject;
+import org.apache.druid.guice.annotations.EphemeralStorageLoading;
 import org.apache.druid.guice.annotations.Json;
 import org.apache.druid.segment.IndexIO;
+import org.apache.druid.segment.loading.AcquireMode;
 import org.apache.druid.segment.loading.LeastBytesUsedStorageLocationSelectorStrategy;
 import org.apache.druid.segment.loading.SegmentCacheManager;
 import org.apache.druid.segment.loading.SegmentLoaderConfig;
@@ -43,26 +46,45 @@ public class SegmentCacheManagerFactory
 {
   private final IndexIO indexIO;
   private final ObjectMapper jsonMapper;
+  private final StorageLoadingThreadPool loadingThreadPool;
 
   @Inject
   public SegmentCacheManagerFactory(
       IndexIO indexIO,
-      @Json ObjectMapper mapper
+      @Json ObjectMapper mapper,
+      @EphemeralStorageLoading StorageLoadingThreadPool loadingThreadPool
   )
   {
     this.indexIO = indexIO;
     this.jsonMapper = mapper;
+    this.loadingThreadPool = loadingThreadPool;
   }
 
   /**
-   * Creates a new {@link SegmentCacheManager} backed by a new storage location in {@code storageDir}, and a new
-   * loading thread pool of default size.
+   * Convenience constructor for tests and manual construction. Builds a private, always-virtual loading pool of default
+   * size rather than sharing the process-wide {@link EphemeralStorageLoading} pool. The pool is not lifecycle-managed;
+   * this is fine for short-lived test JVMs.
+   */
+  @VisibleForTesting
+  public SegmentCacheManagerFactory(
+      IndexIO indexIO,
+      ObjectMapper mapper
+  )
+  {
+    this(indexIO, mapper, StorageLoadingThreadPool.createForEphemeral(new SegmentLoaderConfig()));
+  }
+
+  /**
+   * Creates a new {@link SegmentCacheManager} backed by a new storage location in {@code storageDir}. When
+   * {@code virtualStorage} is true, the returned manager uses this factory's process-wide
+   * {@link EphemeralStorageLoading} loading pool (shared across all per-task caches and stopped by the lifecycle)
+   * rather than creating its own.
    *
    * @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.
+   *                       loading is triggered by {@link SegmentCacheManager#acquireSegment(DataSegment, AcquireMode)},
+   *                       and segment files are deleted as soon as all holds are closed.
    */
   public SegmentCacheManager manufacturate(File storageDir, Long maxSize, boolean virtualStorage)
   {
@@ -77,11 +99,10 @@ public SegmentCacheManager manufacturate(File storageDir, Long maxSize, boolean
             .setVirtualStorage(virtualStorage)
             .setVirtualStorageIsEphemeral(virtualStorage);
     final List<StorageLocation> storageLocations = loaderConfig.toStorageLocations();
-    final StorageLoadingThreadPool loadingThreadPool = StorageLoadingThreadPool.createFromConfig(loaderConfig);
     return new SegmentLocalCacheManager(
         storageLocations,
         loaderConfig,
-        loadingThreadPool,
+        virtualStorage ? loadingThreadPool : StorageLoadingThreadPool.none(),
         new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
         indexIO,
         jsonMapper

indexing-service/src/test/java/org/apache/druid/indexing/common/SegmentCacheManagerFactoryTest.java

@@ -0,0 +1,105 @@
+/*
+ * 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.indexing.common;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.java.util.emitter.EmittingLogger;
+import org.apache.druid.segment.TestHelper;
+import org.apache.druid.segment.TestIndex;
+import org.apache.druid.segment.loading.SegmentCacheManager;
+import org.apache.druid.segment.loading.SegmentLoaderConfig;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
+import org.apache.druid.server.metrics.NoopServiceEmitter;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.File;
+
+class SegmentCacheManagerFactoryTest
+{
+  @TempDir
+  File tempDir;
+
+  private ObjectMapper jsonMapper;
+
+  @BeforeAll
+  static void setUpClass()
+  {
+    EmittingLogger.registerEmitter(new NoopServiceEmitter());
+  }
+
+  @BeforeEach
+  void setUp()
+  {
+    jsonMapper = TestHelper.makeJsonMapper();
+  }
+
+  @Test
+  void testVirtualStorageManagersShareTheInjectedPool()
+  {
+    final StorageLoadingThreadPool shared = StorageLoadingThreadPool.createForEphemeral(new SegmentLoaderConfig());
+    try {
+      final SegmentCacheManagerFactory factory = new SegmentCacheManagerFactory(TestIndex.INDEX_IO, jsonMapper, shared);
+      final SegmentCacheManager m1 = factory.manufacturate(new File(tempDir, "a"), null, true);
+      final SegmentCacheManager m2 = factory.manufacturate(new File(tempDir, "b"), null, true);
+
+      Assertions.assertSame(shared, m1.getLoadingThreadPool());
+      Assertions.assertSame(shared, m2.getLoadingThreadPool());
+    }
+    finally {
+      shared.stop();
+    }
+  }
+
+  @Test
+  void testNonVirtualStorageManagerDoesNotUseTheSharedPool()
+  {
+    final StorageLoadingThreadPool shared = StorageLoadingThreadPool.createForEphemeral(new SegmentLoaderConfig());
+    try {
+      final SegmentCacheManagerFactory factory = new SegmentCacheManagerFactory(TestIndex.INDEX_IO, jsonMapper, shared);
+      final SegmentCacheManager m = factory.manufacturate(new File(tempDir, "c"), null, false);
+
+      Assertions.assertNotSame(shared, m.getLoadingThreadPool());
+      Assertions.assertFalse(m.getLoadingThreadPool().isAvailable());
+    }
+    finally {
+      shared.stop();
+    }
+  }
+
+  @Test
+  void testConvenienceConstructorBuildsOneAvailablePoolPerFactory()
+  {
+    final SegmentCacheManagerFactory factory = new SegmentCacheManagerFactory(TestIndex.INDEX_IO, jsonMapper);
+    final SegmentCacheManager m1 = factory.manufacturate(new File(tempDir, "d"), null, true);
+    final SegmentCacheManager m2 = factory.manufacturate(new File(tempDir, "e"), null, true);
+    try {
+      Assertions.assertTrue(m1.getLoadingThreadPool().isAvailable());
+      // The convenience constructor builds one pool and shares it across the factory's manufacturate calls.
+      Assertions.assertSame(m1.getLoadingThreadPool(), m2.getLoadingThreadPool());
+    }
+    finally {
+      m1.getLoadingThreadPool().stop();
+    }
+  }
+}

processing/src/main/java/org/apache/druid/guice/annotations/EphemeralStorageLoading.java

@@ -0,0 +1,37 @@
+/*
+ * 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.guice.annotations;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Binding annotation for the {@code StorageLoadingThreadPool} used by ephemeral, per-task segment caches
+ */
+@Target({ElementType.FIELD, ElementType.PARAMETER, ElementType.METHOD})
+@Retention(RetentionPolicy.RUNTIME)
+@BindingAnnotation
+public @interface EphemeralStorageLoading
+{
+}

server/src/main/java/org/apache/druid/guice/StorageNodeModule.java

@@ -28,6 +28,7 @@
 import com.google.inject.util.Providers;
 import org.apache.druid.client.DruidServerConfig;
 import org.apache.druid.discovery.DataNodeService;
+import org.apache.druid.guice.annotations.EphemeralStorageLoading;
 import org.apache.druid.guice.annotations.Self;
 import org.apache.druid.java.util.emitter.EmittingLogger;
 import org.apache.druid.query.DruidProcessingConfig;
@@ -134,6 +135,19 @@ public StorageLoadingThreadPool getStorageLoadingThreadPool(SegmentLoaderConfig
     return StorageLoadingThreadPool.createFromConfig(config);
   }
 
+  /**
+   * Process-wide, always-virtual on-demand loading pool shared by the ephemeral per-task segment caches built via
+   * {@code SegmentCacheManagerFactory} (tasks and MSQ workers). Lifecycle-managed so those caches need not each create
+   * (and leak) their own pool on long-lived workers.
+   */
+  @Provides
+  @ManageLifecycle
+  @EphemeralStorageLoading
+  public StorageLoadingThreadPool getEphemeralStorageLoadingThreadPool(SegmentLoaderConfig config)
+  {
+    return StorageLoadingThreadPool.createForEphemeral(config);
+  }
+
   @Provides
   @LazySingleton
   @Named(IS_SEGMENT_CACHE_CONFIGURED)

server/src/main/java/org/apache/druid/segment/loading/StorageLoadingThreadPool.java

@@ -63,49 +63,59 @@ public StorageLoadingThreadPool(
 
   public static StorageLoadingThreadPool createFromConfig(final SegmentLoaderConfig config)
   {
-    final ListeningExecutorService exec;
+    return new StorageLoadingThreadPool(config.isVirtualStorage() ? createOnDemandLoadingExecutor(config) : null);
+  }
+
+  /**
+   * Build a pool configured for virtual-storage on-demand loading <b>regardless of</b>
+   * {@link SegmentLoaderConfig#isVirtualStorage()}. Used for the process-wide loading pool shared by ephemeral,
+   * per-task segment caches, whose host process may not itself run in virtual-storage mode. The executor is created
+   * eagerly but spawns no threads until work is submitted, so an unused pool is cheap.
+   *
+   * @see org.apache.druid.guice.annotations.EphemeralStorageLoading
+   */
+  public static StorageLoadingThreadPool createForEphemeral(final SegmentLoaderConfig config)
+  {
+    return new StorageLoadingThreadPool(createOnDemandLoadingExecutor(config));
+  }
 
-    if (config.isVirtualStorage()) {
-      if (config.getVirtualStorageLoadThreads() <= 0) {
-        throw DruidException.forPersona(DruidException.Persona.OPERATOR)
-                            .ofCategory(DruidException.Category.INVALID_INPUT)
-                            .build(
-                                "virtualStorageLoadThreads must be greater than 0, got [%d]",
-                                config.getVirtualStorageLoadThreads()
-                            );
-      }
-      if (config.isVirtualStorageUseVirtualThreads()) {
-        log.info(
-            "Using virtual storage mode with virtual threads - max concurrent on demand loads: [%d].",
-            config.getVirtualStorageLoadThreads()
-        );
-        exec = new PermitBoundedListeningExecutorService(
-            MoreExecutors.listeningDecorator(
-                Executors.newThreadPerTaskExecutor(
-                    Thread.ofVirtual()
-                          .name("VirtualStorageOnDemandLoadingThread-", 0)
-                          .factory()
-                )
-            ),
-            new Semaphore(config.getVirtualStorageLoadThreads())
-        );
-      } else {
-        log.info(
-            "Using virtual storage mode with fixed platform thread pool - on demand load threads: [%d].",
-            config.getVirtualStorageLoadThreads()
-        );
-        exec = MoreExecutors.listeningDecorator(
-            Executors.newFixedThreadPool(
-                config.getVirtualStorageLoadThreads(),
-                Execs.makeThreadFactory("VirtualStorageOnDemandLoadingThread-%s")
-            )
-        );
-      }
+  private static ListeningExecutorService createOnDemandLoadingExecutor(final SegmentLoaderConfig config)
+  {
+    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()
+      );
+      return new PermitBoundedListeningExecutorService(
+          MoreExecutors.listeningDecorator(
+              Executors.newThreadPerTaskExecutor(
+                  Thread.ofVirtual()
+                        .name("VirtualStorageOnDemandLoadingThread-", 0)
+                        .factory()
+              )
+          ),
+          new Semaphore(config.getVirtualStorageLoadThreads())
+      );
     } else {
-      exec = null;
+      log.info(
+          "Using virtual storage mode with fixed platform thread pool - on demand load threads: [%d].",
+          config.getVirtualStorageLoadThreads()
+      );
+      return MoreExecutors.listeningDecorator(
+          Executors.newFixedThreadPool(
+              config.getVirtualStorageLoadThreads(),
+              Execs.makeThreadFactory("VirtualStorageOnDemandLoadingThread-%s")
+          )
+      );
     }
-
-    return new StorageLoadingThreadPool(exec);
   }
 
   /**

server/src/test/java/org/apache/druid/guice/StorageNodeModuleTest.java

@@ -27,10 +27,12 @@
 import com.google.inject.util.Modules;
 import org.apache.druid.discovery.DataNodeService;
 import org.apache.druid.error.ExceptionMatcher;
+import org.apache.druid.guice.annotations.EphemeralStorageLoading;
 import org.apache.druid.guice.annotations.Self;
 import org.apache.druid.initialization.Initialization;
 import org.apache.druid.query.DruidProcessingConfig;
 import org.apache.druid.segment.loading.SegmentLoaderConfig;
+import org.apache.druid.segment.loading.StorageLoadingThreadPool;
 import org.apache.druid.segment.loading.StorageLocationConfig;
 import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.coordination.DruidServerMetadata;
@@ -167,6 +169,25 @@ public void testDruidServerMetadataWithNoServerTypeConfigShouldThrowProvisionExc
         .assertThrowsAndMatches(() -> injector.getInstance(DruidServerMetadata.class));
   }
 
+  @Test
+  public void testEphemeralStorageLoadingThreadPoolIsInjectedAndAvailable()
+  {
+    // The qualified ephemeral loading pool must resolve from the core injector (StorageNodeModule is universal via
+    // CoreInjectorBuilder), so SegmentCacheManagerFactory's @Inject constructor can be satisfied on every process.
+    // It is always-virtual regardless of the node's virtualStorage flag.
+    Mockito.when(segmentLoaderConfig.getVirtualStorageLoadThreads()).thenReturn(4);
+
+    final StorageLoadingThreadPool pool = injector().getInstance(
+        Key.get(StorageLoadingThreadPool.class, EphemeralStorageLoading.class)
+    );
+    try {
+      Assert.assertTrue(pool.isAvailable());
+    }
+    finally {
+      pool.stop();
+    }
+  }
+
   private Injector injector()
   {
     return makeInjector(INJECT_SERVER_TYPE_CONFIG);