Use min of scheduler threads and server threads for subquery guardrails.

extensions-contrib/moving-average-query/src/test/java/org/apache/druid/query/movingaverage/MovingAverageQueryTest.java

@@ -72,11 +72,13 @@
 import org.apache.druid.segment.join.MapJoinableFactory;
 import org.apache.druid.server.ClientQuerySegmentWalker;
 import org.apache.druid.server.QueryStackTests;
+import org.apache.druid.server.SubqueryGuardrailHelper;
 import org.apache.druid.server.initialization.ServerConfig;
 import org.apache.druid.server.metrics.NoopServiceEmitter;
 import org.apache.druid.server.metrics.SubqueryCountStatsProvider;
 import org.apache.druid.testing.InitializedNullHandlingTest;
 import org.apache.druid.timeline.TimelineLookup;
+import org.apache.druid.utils.JvmUtils;
 import org.hamcrest.core.IsInstanceOf;
 import org.joda.time.Interval;
 import org.junit.Assert;
@@ -388,7 +390,7 @@ public void emit(Event event)
         serverConfig,
         null,
         new CacheConfig(),
-        null,
+        new SubqueryGuardrailHelper(null, JvmUtils.getRuntimeInfo().getMaxHeapSizeBytes(), 1),
         new SubqueryCountStatsProvider()
     );
 

server/pom.xml

@@ -443,6 +443,8 @@
                       <!-- Tested in the SQL layer, but, oddly, not in this module. -->
                         <exclude>org/apache/druid/server/QueryResponse.class</exclude>
                         <exclude>org/apache/druid/curator/CuratorModule.class</exclude>
+                        <!-- Guice providers where unit tests would be tightly-coupled and not useful -->
+                        <exclude>org/apache/druid/server/SubqueryGuardrailHelperProvider.class</exclude>
                     </excludes>
                 </configuration>
             </plugin>

server/src/main/java/org/apache/druid/server/ClientQuerySegmentWalker.java

@@ -58,7 +58,6 @@
 import org.apache.druid.query.SegmentDescriptor;
 import org.apache.druid.query.TableDataSource;
 import org.apache.druid.query.context.ResponseContext;
-import org.apache.druid.query.lookup.LookupExtractorFactoryContainerProvider;
 import org.apache.druid.query.planning.DataSourceAnalysis;
 import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.segment.join.JoinableFactory;
@@ -120,7 +119,7 @@ public ClientQuerySegmentWalker(
       ServerConfig serverConfig,
       Cache cache,
       CacheConfig cacheConfig,
-      LookupExtractorFactoryContainerProvider lookupManager,
+      SubqueryGuardrailHelper subqueryGuardrailHelper,
       SubqueryCountStatsProvider subqueryStatsProvider
   )
   {
@@ -134,11 +133,7 @@ public ClientQuerySegmentWalker(
     this.serverConfig = serverConfig;
     this.cache = cache;
     this.cacheConfig = cacheConfig;
-    this.subqueryGuardrailHelper = new SubqueryGuardrailHelper(
-        lookupManager,
-        Runtime.getRuntime().maxMemory(),
-        serverConfig.getNumThreads()
-    );
+    this.subqueryGuardrailHelper = subqueryGuardrailHelper;
     this.subqueryStatsProvider = subqueryStatsProvider;
   }
 
@@ -154,7 +149,7 @@ public ClientQuerySegmentWalker(
       ServerConfig serverConfig,
       Cache cache,
       CacheConfig cacheConfig,
-      LookupExtractorFactoryContainerProvider lookupManager,
+      SubqueryGuardrailHelper subqueryGuardrailHelper,
       SubqueryCountStatsProvider subqueryStatsProvider
   )
   {
@@ -169,7 +164,7 @@ public ClientQuerySegmentWalker(
         serverConfig,
         cache,
         cacheConfig,
-        lookupManager,
+        subqueryGuardrailHelper,
         subqueryStatsProvider
     );
   }

server/src/main/java/org/apache/druid/server/SubqueryGuardrailHelper.java

@@ -48,11 +48,11 @@ public class SubqueryGuardrailHelper
   public SubqueryGuardrailHelper(
       final LookupExtractorFactoryContainerProvider lookupManager,
       final long maxMemoryInJvm,
-      final int brokerNumHttpConnections
+      final int maxConcurrentQueries
   )
   {
     final DateTime start = DateTimes.nowUtc();
-    autoLimitBytes = computeLimitBytesForAuto(lookupManager, maxMemoryInJvm, brokerNumHttpConnections);
+    autoLimitBytes = computeLimitBytesForAuto(lookupManager, maxMemoryInJvm, maxConcurrentQueries);
     final long startupTimeMs = new Interval(start, DateTimes.nowUtc()).toDurationMillis();
 
     log.info("Took [%d] ms to initialize the SubqueryGuardrailHelper.", startupTimeMs);
@@ -114,13 +114,13 @@ public long convertSubqueryLimitStringToLong(final String maxSubqueryLimit)
   private static long computeLimitBytesForAuto(
       final LookupExtractorFactoryContainerProvider lookupManager,
       final long maxMemoryInJvm,
-      final int brokerNumHttpConnections
+      final int maxConcurrentQueries
   )
   {
     long memoryInJvmWithoutLookups = maxMemoryInJvm - computeLookupFootprint(lookupManager);
     long memoryInJvmForSubqueryResultsInlining = (long) (memoryInJvmWithoutLookups * SUBQUERY_MEMORY_BYTES_FRACTION);
-    long memoryInJvmForSubqueryResultsInliningPerQuery = memoryInJvmForSubqueryResultsInlining
-                                                         / brokerNumHttpConnections;
+    long memoryInJvmForSubqueryResultsInliningPerQuery =
+        memoryInJvmForSubqueryResultsInlining / maxConcurrentQueries;
     return Math.max(memoryInJvmForSubqueryResultsInliningPerQuery, 1L);
   }
 

server/src/main/java/org/apache/druid/server/SubqueryGuardrailHelperProvider.java

@@ -0,0 +1,68 @@
+/*
+ * 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.server;
+
+import com.google.inject.Inject;
+import com.google.inject.Provider;
+import org.apache.druid.guice.LazySingleton;
+import org.apache.druid.query.lookup.LookupExtractorFactoryContainerProvider;
+import org.apache.druid.server.initialization.ServerConfig;
+import org.apache.druid.utils.JvmUtils;
+
+public class SubqueryGuardrailHelperProvider implements Provider<SubqueryGuardrailHelper>
+{
+  private final LookupExtractorFactoryContainerProvider lookupExtractorFactoryContainerProvider;
+  private final ServerConfig serverConfig;
+  private final QuerySchedulerConfig querySchedulerConfig;
+
+  @Inject
+  public SubqueryGuardrailHelperProvider(
+      LookupExtractorFactoryContainerProvider lookupExtractorFactoryContainerProvider,
+      ServerConfig serverConfig,
+      QuerySchedulerConfig querySchedulerConfig
+  )
+  {
+    this.lookupExtractorFactoryContainerProvider = lookupExtractorFactoryContainerProvider;
+    this.serverConfig = serverConfig;
+    this.querySchedulerConfig = querySchedulerConfig;
+  }
+
+  @Override
+  @LazySingleton
+  public SubqueryGuardrailHelper get()
+  {
+    final int maxConcurrentQueries;
+
+    if (querySchedulerConfig.getNumThreads() > 0) {
+      maxConcurrentQueries = Math.min(
+          querySchedulerConfig.getNumThreads(),
+          serverConfig.getNumThreads()
+      );
+    } else {
+      maxConcurrentQueries = serverConfig.getNumThreads();
+    }
+
+    return new SubqueryGuardrailHelper(
+        lookupExtractorFactoryContainerProvider,
+        JvmUtils.getRuntimeInfo().getMaxHeapSizeBytes(),
+        maxConcurrentQueries
+    );
+  }
+}

server/src/test/java/org/apache/druid/server/ClientQuerySegmentWalkerTest.java

@@ -1439,8 +1439,7 @@ public <T> QueryRunner<T> getQueryRunnerForSegments(Query<T> query, Iterable<Seg
         ),
         conglomerate,
         joinableFactory,
-        serverConfig,
-        null
+        serverConfig
     );
   }
 

server/src/test/java/org/apache/druid/server/QueryStackTests.java

@@ -83,6 +83,7 @@
 import org.apache.druid.server.scheduling.ManualQueryPrioritizationStrategy;
 import org.apache.druid.server.scheduling.NoQueryLaningStrategy;
 import org.apache.druid.timeline.VersionedIntervalTimeline;
+import org.apache.druid.utils.JvmUtils;
 import org.junit.Assert;
 
 import javax.annotation.Nullable;
@@ -117,8 +118,7 @@ public static ClientQuerySegmentWalker createClientQuerySegmentWalker(
       final QuerySegmentWalker localWalker,
       final QueryRunnerFactoryConglomerate conglomerate,
       final JoinableFactory joinableFactory,
-      final ServerConfig serverConfig,
-      final LookupExtractorFactoryContainerProvider lookupManager
+      final ServerConfig serverConfig
   )
   {
     return new ClientQuerySegmentWalker(
@@ -164,7 +164,7 @@ public boolean isUseResultLevelCache()
             return false;
           }
         },
-        lookupManager,
+        new SubqueryGuardrailHelper(null, JvmUtils.getRuntimeInfo().getMaxHeapSizeBytes(), 1),
         new SubqueryCountStatsProvider()
     );
   }

services/src/main/java/org/apache/druid/cli/CliBroker.java

@@ -61,6 +61,8 @@
 import org.apache.druid.server.ClientQuerySegmentWalker;
 import org.apache.druid.server.ResponseContextConfig;
 import org.apache.druid.server.SegmentManager;
+import org.apache.druid.server.SubqueryGuardrailHelper;
+import org.apache.druid.server.SubqueryGuardrailHelperProvider;
 import org.apache.druid.server.coordination.ServerType;
 import org.apache.druid.server.coordination.ZkCoordinator;
 import org.apache.druid.server.http.BrokerResource;
@@ -146,6 +148,7 @@ protected List<? extends Module> getModules()
 
           binder.bind(BrokerQueryResource.class).in(LazySingleton.class);
           Jerseys.addResource(binder, BrokerQueryResource.class);
+          binder.bind(SubqueryGuardrailHelper.class).toProvider(SubqueryGuardrailHelperProvider.class);
           binder.bind(QueryCountStatsProvider.class).to(BrokerQueryResource.class).in(LazySingleton.class);
           binder.bind(SubqueryCountStatsProvider.class).toInstance(new SubqueryCountStatsProvider());
           Jerseys.addResource(binder, BrokerResource.class);

sql/src/test/java/org/apache/druid/sql/calcite/util/SpecificSegmentsQuerySegmentWalker.java

@@ -116,8 +116,7 @@ public SpecificSegmentsQuerySegmentWalker(
         ),
         conglomerate,
         joinableFactoryWrapper.getJoinableFactory(),
-        new ServerConfig(),
-        LOOKUP_EXTRACTOR_FACTORY_CONTAINER_PROVIDER
+        new ServerConfig()
     );
   }