[PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["yashmayya (via GitHub)" <gi...@apache.org>]

yashmayya opened a new pull request, #14574:
URL: https://github.com/apache/pinot/pull/14574

   - Currently, there is no limit on the number of multi-stage queries that are executed concurrently. Every request that comes in to a broker is immediately compiled into a query plan and dispatched to the servers for query execution.
   - On the servers, each stage in the query plan is submitted to a cached thread pool executor. This means that we don't have an effective way to limit the number of threads that are being utilized to execute multi-stage queries on servers. This can result in a very large number of multi-stage query executor threads in high QPS / complex query workloads leading to resource contention and performance degradation.
   - Simply using a fixed thread pool executor instead is not a good solution because this could lead to distributed deadlocks and query timeouts / failures. Similarly, it isn't trivial to implement the max concurrent query execution logic in the servers without considering global ordering.
   - Consider this example from @gortiz - 
     - Given a query `Q1` is formed by stages `S11` and `S12`, where `S11` depends on data from `S12`.
     - And a query `Q2` is formed by stages `S21` and `S22`, where `S21` depends on data from `S22`
     - If `Q1` and `Q2` arrive at the same time and:
       - Server `Srv1` receives `Q1` first, starts executing `S11` and ends up waiting for server `Srv2` to execute `S12`
       - Server `Srv2` receives `Q2` first, starts executing `S21` and ends up waiting for server `Srv1` to execute `S22`
       - Server `Srv1` receives `Q2` and would execute `S22`, but sees that it has reached max number of concurrent queries, so it waits.
       - Server Srv2 receives `Q1` and would execute `S12`, but sees that it has reached max number of concurrent queries, so it waits.
       - This leads to a situation where both the queries are blocked and will eventually timeout and fail.
   - This patch takes an alternative approach of limiting the number of concurrently executing multi-stage queries in the broker itself.
   - A new cluster level config `pinot.multistage.engine.max.concurrent.queries` is introduced. This is disabled by default and is currently an opt-in throttling mechanism.
   - If set to a value > 0, the value will be divided among all the brokers - i.e., if `pinot.multistage.engine.max.concurrent.queries` is set to `12` and there are 3 brokers in the cluster, each broker will be able to execute 4 multi-stage queries concurrently. If there are already 4 queries currently executing, any subsequent incoming queries will wait for a slot to open up.
   - This is implemented using a custom semaphore which reacts to changes in the cluster config value or the number of brokers in the cluster in order to recalculate the total available permits. A custom semaphore implementation is needed here because we want to use the [Semaphore::reducePermits](https://docs.oracle.com/javase/8/docs/api/java/util/concurrent/Semaphore.html#reducePermits-int-) method (when # of brokers increases or the cluster config value for max concurrent queries decreases) which has `protected` access. We need to use this method in order to be able to reduce the total number of permits for the semaphore without blocking.
   - The name of the cluster config doesn't indicate that it's a broker level mechanism and hence can be re-used in the future if we move this logic to the servers. 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888523744


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).
+ */
+public class MultiStageQueryThrottler implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQueryThrottler.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _numServers;
+  // If _maxConcurrentQueries is <= 0, it means that the cluster is not configured to limit the number of multi-stage
+  // queries that can be executed concurrently. In this case, we should not block the query.
+  private int _maxConcurrentQueries;
+  private AdjustableSemaphore _semaphore;
+
+  @Override
+  public void init(HelixManager helixManager) {
+    _helixManager = helixManager;
+    _helixAdmin = _helixManager.getClusterManagmentTool();
+    _helixConfigScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(
+        _helixManager.getClusterName()).build();
+
+    _maxConcurrentQueries = Integer.parseInt(
+        _helixAdmin.getConfig(_helixConfigScope,
+                Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+            .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+    List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+    _numBrokers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+        .count());
+    _numServers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+        .count());
+
+    if (_maxConcurrentQueries > 0) {
+      _semaphore = new AdjustableSemaphore(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers), true);
+    }
+  }
+
+  /**
+   * Returns true if the query can be executed (waiting until it can be executed if necessary), false otherwise.
+   * <p>
+   * {@link #release()} should be called after the query is done executing. It is the responsibility of the caller to
+   * ensure that {@link #release()} is called exactly once for each call to this method.
+   *
+   * @param timeout the maximum time to wait
+   * @param unit the time unit of the timeout argument
+   * @throws InterruptedException if the current thread is interrupted
+   */
+  public boolean tryAcquire(long timeout, TimeUnit unit)
+      throws InterruptedException {
+    if (_maxConcurrentQueries <= 0) {
+      return true;
+    }
+    return _semaphore.tryAcquire(timeout, unit);
+  }
+
+  /**
+   * Should be called after the query is done executing. It is the responsibility of the caller to ensure that this
+   * method is called exactly once for each call to {@link #tryAcquire(long, TimeUnit)}.
+   */
+  public void release() {
+    if (_maxConcurrentQueries > 0) {
+      _semaphore.release();
+    }
+  }
+
+  @Override
+  public void processClusterChange(HelixConstants.ChangeType changeType) {
+    Preconditions.checkArgument(
+        changeType == HelixConstants.ChangeType.EXTERNAL_VIEW || changeType == HelixConstants.ChangeType.CLUSTER_CONFIG,
+        "MultiStageQuerySemaphore can only handle EXTERNAL_VIEW and CLUSTER_CONFIG changes");
+
+    if (changeType == HelixConstants.ChangeType.EXTERNAL_VIEW) {
+      List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+      int numBrokers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+          .count());
+      int numServers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+          .count());
+
+      if (numBrokers != _numBrokers || numServers != _numServers) {
+        _numBrokers = numBrokers;
+        _numServers = numServers;
+        if (_maxConcurrentQueries > 0) {
+          _semaphore.setPermits(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers));
+        }
+      }
+    } else {
+      int maxConcurrentQueries = Integer.parseInt(
+          _helixAdmin.getConfig(_helixConfigScope,
+                  Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+              .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                  CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+      if (_maxConcurrentQueries == maxConcurrentQueries) {
+        return;
+      }
+
+      if (_maxConcurrentQueries <= 0 && maxConcurrentQueries > 0
+          || _maxConcurrentQueries > 0 && maxConcurrentQueries <= 0) {
+        // This operation isn't safe to do while queries are running so we require a restart of the broker for this
+        // change to take effect.
+        LOGGER.warn("Enabling or disabling limitation of the maximum number of multi-stage queries running "
+            + "concurrently requires a restart of the broker to take effect");
+        return;
+      }
+
+      if (maxConcurrentQueries > 0) {
+        _semaphore.setPermits(Math.max(1, maxConcurrentQueries * _numServers / _numBrokers));
+      }

Review Comment:
   I have two proposals to fix it:
   1. Modify `tryAcquire` to return an object that:
     * has a method to check whether we are allowed to go
     * is AutoCloseable. Its close method release the lock if and only if the semaphore didn't change (we can use a version stored in a AtomicInt to verify that without locking).
   2. Keep in this object a set of threads and modify the release method to only release the permit if the caller thread is listed in the set. We can clean the set when config changes. This requires less changes in the code, but the synchronization may be expensive.
   
   In general I think we need to synchronize in this class anyway



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888382486


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).
+ */
+public class MultiStageQueryThrottler implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQueryThrottler.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _numServers;
+  // If _maxConcurrentQueries is <= 0, it means that the cluster is not configured to limit the number of multi-stage
+  // queries that can be executed concurrently. In this case, we should not block the query.
+  private int _maxConcurrentQueries;
+  private AdjustableSemaphore _semaphore;
+
+  @Override
+  public void init(HelixManager helixManager) {
+    _helixManager = helixManager;
+    _helixAdmin = _helixManager.getClusterManagmentTool();
+    _helixConfigScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(
+        _helixManager.getClusterName()).build();
+
+    _maxConcurrentQueries = Integer.parseInt(
+        _helixAdmin.getConfig(_helixConfigScope,
+                Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+            .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+    List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+    _numBrokers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+        .count());
+    _numServers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+        .count());
+
+    if (_maxConcurrentQueries > 0) {
+      _semaphore = new AdjustableSemaphore(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers), true);
+    }
+  }
+
+  /**
+   * Returns true if the query can be executed (waiting until it can be executed if necessary), false otherwise.
+   * <p>
+   * {@link #release()} should be called after the query is done executing. It is the responsibility of the caller to
+   * ensure that {@link #release()} is called exactly once for each call to this method.
+   *
+   * @param timeout the maximum time to wait
+   * @param unit the time unit of the timeout argument
+   * @throws InterruptedException if the current thread is interrupted
+   */
+  public boolean tryAcquire(long timeout, TimeUnit unit)
+      throws InterruptedException {
+    if (_maxConcurrentQueries <= 0) {
+      return true;
+    }
+    return _semaphore.tryAcquire(timeout, unit);
+  }
+
+  /**
+   * Should be called after the query is done executing. It is the responsibility of the caller to ensure that this
+   * method is called exactly once for each call to {@link #tryAcquire(long, TimeUnit)}.
+   */
+  public void release() {
+    if (_maxConcurrentQueries > 0) {
+      _semaphore.release();
+    }
+  }
+
+  @Override
+  public void processClusterChange(HelixConstants.ChangeType changeType) {
+    Preconditions.checkArgument(
+        changeType == HelixConstants.ChangeType.EXTERNAL_VIEW || changeType == HelixConstants.ChangeType.CLUSTER_CONFIG,
+        "MultiStageQuerySemaphore can only handle EXTERNAL_VIEW and CLUSTER_CONFIG changes");
+
+    if (changeType == HelixConstants.ChangeType.EXTERNAL_VIEW) {
+      List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+      int numBrokers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+          .count());
+      int numServers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+          .count());
+
+      if (numBrokers != _numBrokers || numServers != _numServers) {
+        _numBrokers = numBrokers;
+        _numServers = numServers;
+        if (_maxConcurrentQueries > 0) {
+          _semaphore.setPermits(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers));
+        }
+      }
+    } else {
+      int maxConcurrentQueries = Integer.parseInt(
+          _helixAdmin.getConfig(_helixConfigScope,
+                  Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+              .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                  CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+      if (_maxConcurrentQueries == maxConcurrentQueries) {
+        return;
+      }
+
+      if (_maxConcurrentQueries <= 0 && maxConcurrentQueries > 0
+          || _maxConcurrentQueries > 0 && maxConcurrentQueries <= 0) {
+        // This operation isn't safe to do while queries are running so we require a restart of the broker for this
+        // change to take effect.
+        LOGGER.warn("Enabling or disabling limitation of the maximum number of multi-stage queries running "
+            + "concurrently requires a restart of the broker to take effect");
+        return;
+      }
+
+      if (maxConcurrentQueries > 0) {
+        _semaphore.setPermits(Math.max(1, maxConcurrentQueries * _numServers / _numBrokers));
+      }

Review Comment:
   Maybe we can add a test to verify my scenario.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1886491028


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java:
##########
@@ -224,67 +227,89 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
       return new BrokerResponseNative(QueryException.getException(QueryException.QUOTA_EXCEEDED_ERROR, errorMessage));
     }
 
-    Tracing.ThreadAccountantOps.setupRunner(String.valueOf(requestId), ThreadExecutionContext.TaskType.MSE);
-
-    long executionStartTimeNs = System.nanoTime();
-    QueryDispatcher.QueryResult queryResults;
+    Timer queryTimer = new Timer(queryTimeoutMs);
     try {
-      queryResults =
-          _queryDispatcher.submitAndReduce(requestContext, dispatchableSubPlan, queryTimeoutMs, queryOptions);
-    } catch (TimeoutException e) {
-      for (String table : tableNames) {
-        _brokerMetrics.addMeteredTableValue(table, BrokerMeter.BROKER_RESPONSES_WITH_TIMEOUTS, 1);
+      // It's fine to block in this thread because we use a separate thread pool from the main Jersey server to process
+      // these requests.
+      if (!_querySemaphore.tryAcquire(queryTimeoutMs, TimeUnit.MILLISECONDS)) {
+        LOGGER.warn("Timed out waiting to execute request {}: {}", requestId, query);
+        requestContext.setErrorCode(QueryException.EXECUTION_TIMEOUT_ERROR_CODE);
+        return new BrokerResponseNative(QueryException.EXECUTION_TIMEOUT_ERROR);
       }
-      LOGGER.warn("Timed out executing request {}: {}", requestId, query);
+    } catch (InterruptedException e) {
+      LOGGER.warn("Interrupt received while waiting to execute request {}: {}", requestId, query);
       requestContext.setErrorCode(QueryException.EXECUTION_TIMEOUT_ERROR_CODE);
       return new BrokerResponseNative(QueryException.EXECUTION_TIMEOUT_ERROR);

Review Comment:
   Maybe not in this PR, but we will probably need to have a way to tell the user where we were waiting most of the time. It could be useful for example to tell users that their queries are timing out because other queries are consuming all threads.



##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQuerySemaphore.java:
##########
@@ -0,0 +1,174 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future.
+ */
+public class MultiStageQuerySemaphore extends Semaphore implements ClusterChangeHandler {

Review Comment:
   IMHO this class does too many things. I would suggest to split it in two different classes:
   
   1. The semaphore itself. It doesn't need to know the number of brokers, max concurrent queries, etc. It has a `setTotalPermits` for a new number of permits, which calls `super.reducePermits` and `super.release` if needed.
   2. The ClusterChangeHandler, which has a reference to the semaphore. It knows the number of brokers and max concurrent queries and calls `setTotalPermits` each time there are cluster changes.
   
   By doing so the semaphore has no knowledge related to this feature and it is just a semaphore with the ability to change permits. It is easier to understand and we can use it in other developments. And different handlers can be used as needed.



##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQuerySemaphore.java:
##########
@@ -0,0 +1,174 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future.
+ */
+public class MultiStageQuerySemaphore extends Semaphore implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQuerySemaphore.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _maxConcurrentQueries;
+  private int _totalPermits;
+
+  public MultiStageQuerySemaphore(int numBrokers, int maxConcurrentQueries) {
+    super(Math.max(1, maxConcurrentQueries / Math.max(numBrokers, 1)), true);
+    _maxConcurrentQueries = maxConcurrentQueries;
+    _numBrokers = Math.max(1, numBrokers);
+    _totalPermits =
+        maxConcurrentQueries > 0 ? Math.max(1, maxConcurrentQueries / Math.max(numBrokers, 1)) : maxConcurrentQueries;
+  }

Review Comment:
   I think we are using too many _max_es here. IMO it would be better to create a static factory method that does all the max logic and then call a simple constructor.
   



##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQuerySemaphore.java:
##########
@@ -0,0 +1,174 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future.
+ */
+public class MultiStageQuerySemaphore extends Semaphore implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQuerySemaphore.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _maxConcurrentQueries;
+  private int _totalPermits;
+
+  public MultiStageQuerySemaphore(int numBrokers, int maxConcurrentQueries) {
+    super(Math.max(1, maxConcurrentQueries / Math.max(numBrokers, 1)), true);
+    _maxConcurrentQueries = maxConcurrentQueries;
+    _numBrokers = Math.max(1, numBrokers);
+    _totalPermits =
+        maxConcurrentQueries > 0 ? Math.max(1, maxConcurrentQueries / Math.max(numBrokers, 1)) : maxConcurrentQueries;
+  }
+
+  @Override
+  public void acquire()
+      throws InterruptedException {
+    // If _totalPermits is <= 0, it means that the cluster is not configured to limit the number of multi-stage queries

Review Comment:
   It is fine to keep it here, but given this is an important attribute invariant, I think it is worth to add it as javadoc in the attribute itself.



##########
pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java:
##########
@@ -335,13 +339,26 @@ public void start()
     MultiStageBrokerRequestHandler multiStageBrokerRequestHandler = null;
     QueryDispatcher queryDispatcher = null;
     if (_brokerConf.getProperty(Helix.CONFIG_OF_MULTI_STAGE_ENGINE_ENABLED, Helix.DEFAULT_MULTI_STAGE_ENGINE_ENABLED)) {
+      int numBrokers = Math.max(1, (int) _helixAdmin
+          .getInstancesInCluster(_spectatorHelixManager.getClusterName())
+          .stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+          .count());
+      int maxConcurrentQueries = Integer.parseInt(
+          _helixAdmin.getConfig(new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(
+                      _spectatorHelixManager.getClusterName()).build(),

Review Comment:
   I don't think this code belongs here. It should be done automatically by `MultiStageQuerySemaphore`



##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQuerySemaphore.java:
##########
@@ -0,0 +1,174 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future.
+ */
+public class MultiStageQuerySemaphore extends Semaphore implements ClusterChangeHandler {

Review Comment:
   Probably you can even generalize the handler to depend on some function callbacks or abstract methods so you don't need to mock all the helix attributes. I would prefer that (the less mocking the better) but I understand it can also be seen as overengineering.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["yashmayya (via GitHub)" <gi...@apache.org>]

yashmayya commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888541652


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).
+ */
+public class MultiStageQueryThrottler implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQueryThrottler.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _numServers;
+  // If _maxConcurrentQueries is <= 0, it means that the cluster is not configured to limit the number of multi-stage
+  // queries that can be executed concurrently. In this case, we should not block the query.
+  private int _maxConcurrentQueries;
+  private AdjustableSemaphore _semaphore;
+
+  @Override
+  public void init(HelixManager helixManager) {
+    _helixManager = helixManager;
+    _helixAdmin = _helixManager.getClusterManagmentTool();
+    _helixConfigScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(
+        _helixManager.getClusterName()).build();
+
+    _maxConcurrentQueries = Integer.parseInt(
+        _helixAdmin.getConfig(_helixConfigScope,
+                Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+            .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+    List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+    _numBrokers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+        .count());
+    _numServers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+        .count());
+
+    if (_maxConcurrentQueries > 0) {
+      _semaphore = new AdjustableSemaphore(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers), true);
+    }
+  }
+
+  /**
+   * Returns true if the query can be executed (waiting until it can be executed if necessary), false otherwise.
+   * <p>
+   * {@link #release()} should be called after the query is done executing. It is the responsibility of the caller to
+   * ensure that {@link #release()} is called exactly once for each call to this method.
+   *
+   * @param timeout the maximum time to wait
+   * @param unit the time unit of the timeout argument
+   * @throws InterruptedException if the current thread is interrupted
+   */
+  public boolean tryAcquire(long timeout, TimeUnit unit)
+      throws InterruptedException {
+    if (_maxConcurrentQueries <= 0) {
+      return true;
+    }
+    return _semaphore.tryAcquire(timeout, unit);
+  }
+
+  /**
+   * Should be called after the query is done executing. It is the responsibility of the caller to ensure that this
+   * method is called exactly once for each call to {@link #tryAcquire(long, TimeUnit)}.
+   */
+  public void release() {
+    if (_maxConcurrentQueries > 0) {
+      _semaphore.release();
+    }
+  }
+
+  @Override
+  public void processClusterChange(HelixConstants.ChangeType changeType) {
+    Preconditions.checkArgument(
+        changeType == HelixConstants.ChangeType.EXTERNAL_VIEW || changeType == HelixConstants.ChangeType.CLUSTER_CONFIG,
+        "MultiStageQuerySemaphore can only handle EXTERNAL_VIEW and CLUSTER_CONFIG changes");
+
+    if (changeType == HelixConstants.ChangeType.EXTERNAL_VIEW) {
+      List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+      int numBrokers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+          .count());
+      int numServers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+          .count());
+
+      if (numBrokers != _numBrokers || numServers != _numServers) {
+        _numBrokers = numBrokers;
+        _numServers = numServers;
+        if (_maxConcurrentQueries > 0) {
+          _semaphore.setPermits(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers));
+        }
+      }
+    } else {
+      int maxConcurrentQueries = Integer.parseInt(
+          _helixAdmin.getConfig(_helixConfigScope,
+                  Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+              .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                  CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+      if (_maxConcurrentQueries == maxConcurrentQueries) {
+        return;
+      }
+
+      if (_maxConcurrentQueries <= 0 && maxConcurrentQueries > 0
+          || _maxConcurrentQueries > 0 && maxConcurrentQueries <= 0) {
+        // This operation isn't safe to do while queries are running so we require a restart of the broker for this
+        // change to take effect.
+        LOGGER.warn("Enabling or disabling limitation of the maximum number of multi-stage queries running "
+            + "concurrently requires a restart of the broker to take effect");
+        return;
+      }
+
+      if (maxConcurrentQueries > 0) {
+        _semaphore.setPermits(Math.max(1, maxConcurrentQueries * _numServers / _numBrokers));
+      }

Review Comment:
   > But... isn't that actually going to be problematic when moving from N to M permits
   
   Hmm, I don't think so - we're increasing / decreasing the number of permits based on the previous total number of permits, not the current available number of permits. You can see some of the unit tests like `testIncreaseNumBrokers` to see how this would look like in action.
   
   > In general I think we need to synchronize in this class anyway
   
   The cluster change handling should be done in a single threaded fashion and the semaphore should support concurrent modification of the number of permits while other threads are trying to acquire / release so I don't think additional external synchronization is necessary here?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["yashmayya (via GitHub)" <gi...@apache.org>]

yashmayya commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888485790


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).
+ */
+public class MultiStageQueryThrottler implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQueryThrottler.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _numServers;
+  // If _maxConcurrentQueries is <= 0, it means that the cluster is not configured to limit the number of multi-stage
+  // queries that can be executed concurrently. In this case, we should not block the query.
+  private int _maxConcurrentQueries;
+  private AdjustableSemaphore _semaphore;
+
+  @Override
+  public void init(HelixManager helixManager) {
+    _helixManager = helixManager;
+    _helixAdmin = _helixManager.getClusterManagmentTool();
+    _helixConfigScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(
+        _helixManager.getClusterName()).build();
+
+    _maxConcurrentQueries = Integer.parseInt(
+        _helixAdmin.getConfig(_helixConfigScope,
+                Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+            .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+    List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+    _numBrokers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+        .count());
+    _numServers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+        .count());
+
+    if (_maxConcurrentQueries > 0) {
+      _semaphore = new AdjustableSemaphore(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers), true);
+    }
+  }
+
+  /**
+   * Returns true if the query can be executed (waiting until it can be executed if necessary), false otherwise.
+   * <p>
+   * {@link #release()} should be called after the query is done executing. It is the responsibility of the caller to
+   * ensure that {@link #release()} is called exactly once for each call to this method.
+   *
+   * @param timeout the maximum time to wait
+   * @param unit the time unit of the timeout argument
+   * @throws InterruptedException if the current thread is interrupted
+   */
+  public boolean tryAcquire(long timeout, TimeUnit unit)
+      throws InterruptedException {
+    if (_maxConcurrentQueries <= 0) {
+      return true;
+    }
+    return _semaphore.tryAcquire(timeout, unit);
+  }
+
+  /**
+   * Should be called after the query is done executing. It is the responsibility of the caller to ensure that this
+   * method is called exactly once for each call to {@link #tryAcquire(long, TimeUnit)}.
+   */
+  public void release() {
+    if (_maxConcurrentQueries > 0) {
+      _semaphore.release();
+    }
+  }
+
+  @Override
+  public void processClusterChange(HelixConstants.ChangeType changeType) {
+    Preconditions.checkArgument(
+        changeType == HelixConstants.ChangeType.EXTERNAL_VIEW || changeType == HelixConstants.ChangeType.CLUSTER_CONFIG,
+        "MultiStageQuerySemaphore can only handle EXTERNAL_VIEW and CLUSTER_CONFIG changes");
+
+    if (changeType == HelixConstants.ChangeType.EXTERNAL_VIEW) {
+      List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+      int numBrokers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+          .count());
+      int numServers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+          .count());
+
+      if (numBrokers != _numBrokers || numServers != _numServers) {
+        _numBrokers = numBrokers;
+        _numServers = numServers;
+        if (_maxConcurrentQueries > 0) {
+          _semaphore.setPermits(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers));
+        }
+      }
+    } else {
+      int maxConcurrentQueries = Integer.parseInt(
+          _helixAdmin.getConfig(_helixConfigScope,
+                  Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+              .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                  CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+      if (_maxConcurrentQueries == maxConcurrentQueries) {
+        return;
+      }
+
+      if (_maxConcurrentQueries <= 0 && maxConcurrentQueries > 0
+          || _maxConcurrentQueries > 0 && maxConcurrentQueries <= 0) {
+        // This operation isn't safe to do while queries are running so we require a restart of the broker for this
+        // change to take effect.
+        LOGGER.warn("Enabling or disabling limitation of the maximum number of multi-stage queries running "
+            + "concurrently requires a restart of the broker to take effect");
+        return;
+      }
+
+      if (maxConcurrentQueries > 0) {
+        _semaphore.setPermits(Math.max(1, maxConcurrentQueries * _numServers / _numBrokers));
+      }

Review Comment:
   > I think there is a possible NPE here. Imagine that the broker is initialized with -1 concurrent queries, _semaphore is kept as null. Then the config is changed to set an actual max concurrent queries, then we try to call _semaphore.setPermits here without actually creating the semaphore.
   
   Unless I'm misunderstanding you, that condition is already being checked just before this line and we're disallowing such a change right? This makes things a lot more simple to reason about - transitioning from disabled to enabled or vice versa while queries are running will create lots of edge cases to handle with the acquire / release mechanism.
   
   > Maybe we can add a test to verify my scenario.
   
   There's already tests verifying disabled -> enabled and enabled -> disabled transitions.



##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).

Review Comment:
   Makes sense, I was assuming that for the 2nd option the check would happen before dispatch instead of during planning. Either way, this is just a rough sketch for a future idea so I've just removed this part of the Javadoc to avoid confusion. This PR can always be re-visited for more context in the future.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888366484


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).

Review Comment:
   nit: I think we are merging two different things here. In my mind there are 3 options:
   
   1. Broader broker check, the one we are implementing, which simplifies the check by assuming evenly distributed queries across brokers and servers.
   2. Ask servers if they can execute queries, where the broker ask the servers whether they can execute more queries or not (probably this should be done while workers/plan fragments are assigned). A server needs to be able to execute all the associated workers of the query atomically. This means that a query that spawns a lot of workers on a single server may consume a lot of threads in the server.
   3. Ask servers if they can execute workers. This is what implies there could be blocks if we don't care about ordering.
   
   The difference between 2 and 3 is that queries that require tons of workers may end up consuming a lot of threads in 2, while in 3 we would execute only the some workers and once they finish the other will start to run. In order to do so we need to execute the dag in a consistent manner (and probably buffer some data).



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888374701


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).
+ */
+public class MultiStageQueryThrottler implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQueryThrottler.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _numServers;
+  // If _maxConcurrentQueries is <= 0, it means that the cluster is not configured to limit the number of multi-stage
+  // queries that can be executed concurrently. In this case, we should not block the query.
+  private int _maxConcurrentQueries;

Review Comment:
   nit: It would be better to set this as javadoc `/**` to be able to use IDE intregrations



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888518378


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).
+ */
+public class MultiStageQueryThrottler implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQueryThrottler.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _numServers;
+  // If _maxConcurrentQueries is <= 0, it means that the cluster is not configured to limit the number of multi-stage
+  // queries that can be executed concurrently. In this case, we should not block the query.
+  private int _maxConcurrentQueries;
+  private AdjustableSemaphore _semaphore;
+
+  @Override
+  public void init(HelixManager helixManager) {
+    _helixManager = helixManager;
+    _helixAdmin = _helixManager.getClusterManagmentTool();
+    _helixConfigScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(
+        _helixManager.getClusterName()).build();
+
+    _maxConcurrentQueries = Integer.parseInt(
+        _helixAdmin.getConfig(_helixConfigScope,
+                Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+            .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+    List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+    _numBrokers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+        .count());
+    _numServers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+        .count());
+
+    if (_maxConcurrentQueries > 0) {
+      _semaphore = new AdjustableSemaphore(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers), true);
+    }
+  }
+
+  /**
+   * Returns true if the query can be executed (waiting until it can be executed if necessary), false otherwise.
+   * <p>
+   * {@link #release()} should be called after the query is done executing. It is the responsibility of the caller to
+   * ensure that {@link #release()} is called exactly once for each call to this method.
+   *
+   * @param timeout the maximum time to wait
+   * @param unit the time unit of the timeout argument
+   * @throws InterruptedException if the current thread is interrupted
+   */
+  public boolean tryAcquire(long timeout, TimeUnit unit)
+      throws InterruptedException {
+    if (_maxConcurrentQueries <= 0) {
+      return true;
+    }
+    return _semaphore.tryAcquire(timeout, unit);
+  }
+
+  /**
+   * Should be called after the query is done executing. It is the responsibility of the caller to ensure that this
+   * method is called exactly once for each call to {@link #tryAcquire(long, TimeUnit)}.
+   */
+  public void release() {
+    if (_maxConcurrentQueries > 0) {
+      _semaphore.release();
+    }
+  }
+
+  @Override
+  public void processClusterChange(HelixConstants.ChangeType changeType) {
+    Preconditions.checkArgument(
+        changeType == HelixConstants.ChangeType.EXTERNAL_VIEW || changeType == HelixConstants.ChangeType.CLUSTER_CONFIG,
+        "MultiStageQuerySemaphore can only handle EXTERNAL_VIEW and CLUSTER_CONFIG changes");
+
+    if (changeType == HelixConstants.ChangeType.EXTERNAL_VIEW) {
+      List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+      int numBrokers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+          .count());
+      int numServers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+          .count());
+
+      if (numBrokers != _numBrokers || numServers != _numServers) {
+        _numBrokers = numBrokers;
+        _numServers = numServers;
+        if (_maxConcurrentQueries > 0) {
+          _semaphore.setPermits(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers));
+        }
+      }
+    } else {
+      int maxConcurrentQueries = Integer.parseInt(
+          _helixAdmin.getConfig(_helixConfigScope,
+                  Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+              .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                  CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+      if (_maxConcurrentQueries == maxConcurrentQueries) {
+        return;
+      }
+
+      if (_maxConcurrentQueries <= 0 && maxConcurrentQueries > 0
+          || _maxConcurrentQueries > 0 && maxConcurrentQueries <= 0) {
+        // This operation isn't safe to do while queries are running so we require a restart of the broker for this
+        // change to take effect.
+        LOGGER.warn("Enabling or disabling limitation of the maximum number of multi-stage queries running "
+            + "concurrently requires a restart of the broker to take effect");
+        return;
+      }
+
+      if (maxConcurrentQueries > 0) {
+        _semaphore.setPermits(Math.max(1, maxConcurrentQueries * _numServers / _numBrokers));
+      }

Review Comment:
   Oh! I didn't _parse_ the previous check. It makes sense. In that case queries that were running when the feature was disabled will suddenly try to release the permit once they finish.
   
   I guess we can resolve that by keeping a reference to a virtual _semaphore_ (not an actual Semaphore, but something that implements the API) and then call release on that object, which will be a noop on queries started when the feature was disabled.
   
   But it is true that it is not a blocking issue and we can assume it requires a restart in this case.
   
   But... isn't that actually going to be problematic when moving from N to M permits? I mean:
   
   1. The semaphore is configured with N (where N > 0) permits.
   1. Query starts, semaphore has N (where N > 0) permits, so it consumes one and continues. The semaphore has now N-1 permits.
   2. The config changes and now we need to use use M permits. We restart the semaphore to use that number.
   3. The query ends and release its permit. Now the semaphore has M+1 permits.
   
   In the worst case scenario we could end up having N+M permits. Luckily, this is not accumulative. But unless I'm wrong I think we need to fix the issue.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["yashmayya (via GitHub)" <gi...@apache.org>]

yashmayya commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888541652


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).
+ */
+public class MultiStageQueryThrottler implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQueryThrottler.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _numServers;
+  // If _maxConcurrentQueries is <= 0, it means that the cluster is not configured to limit the number of multi-stage
+  // queries that can be executed concurrently. In this case, we should not block the query.
+  private int _maxConcurrentQueries;
+  private AdjustableSemaphore _semaphore;
+
+  @Override
+  public void init(HelixManager helixManager) {
+    _helixManager = helixManager;
+    _helixAdmin = _helixManager.getClusterManagmentTool();
+    _helixConfigScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(
+        _helixManager.getClusterName()).build();
+
+    _maxConcurrentQueries = Integer.parseInt(
+        _helixAdmin.getConfig(_helixConfigScope,
+                Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+            .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+    List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+    _numBrokers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+        .count());
+    _numServers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+        .count());
+
+    if (_maxConcurrentQueries > 0) {
+      _semaphore = new AdjustableSemaphore(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers), true);
+    }
+  }
+
+  /**
+   * Returns true if the query can be executed (waiting until it can be executed if necessary), false otherwise.
+   * <p>
+   * {@link #release()} should be called after the query is done executing. It is the responsibility of the caller to
+   * ensure that {@link #release()} is called exactly once for each call to this method.
+   *
+   * @param timeout the maximum time to wait
+   * @param unit the time unit of the timeout argument
+   * @throws InterruptedException if the current thread is interrupted
+   */
+  public boolean tryAcquire(long timeout, TimeUnit unit)
+      throws InterruptedException {
+    if (_maxConcurrentQueries <= 0) {
+      return true;
+    }
+    return _semaphore.tryAcquire(timeout, unit);
+  }
+
+  /**
+   * Should be called after the query is done executing. It is the responsibility of the caller to ensure that this
+   * method is called exactly once for each call to {@link #tryAcquire(long, TimeUnit)}.
+   */
+  public void release() {
+    if (_maxConcurrentQueries > 0) {
+      _semaphore.release();
+    }
+  }
+
+  @Override
+  public void processClusterChange(HelixConstants.ChangeType changeType) {
+    Preconditions.checkArgument(
+        changeType == HelixConstants.ChangeType.EXTERNAL_VIEW || changeType == HelixConstants.ChangeType.CLUSTER_CONFIG,
+        "MultiStageQuerySemaphore can only handle EXTERNAL_VIEW and CLUSTER_CONFIG changes");
+
+    if (changeType == HelixConstants.ChangeType.EXTERNAL_VIEW) {
+      List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+      int numBrokers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+          .count());
+      int numServers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+          .count());
+
+      if (numBrokers != _numBrokers || numServers != _numServers) {
+        _numBrokers = numBrokers;
+        _numServers = numServers;
+        if (_maxConcurrentQueries > 0) {
+          _semaphore.setPermits(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers));
+        }
+      }
+    } else {
+      int maxConcurrentQueries = Integer.parseInt(
+          _helixAdmin.getConfig(_helixConfigScope,
+                  Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+              .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                  CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+      if (_maxConcurrentQueries == maxConcurrentQueries) {
+        return;
+      }
+
+      if (_maxConcurrentQueries <= 0 && maxConcurrentQueries > 0
+          || _maxConcurrentQueries > 0 && maxConcurrentQueries <= 0) {
+        // This operation isn't safe to do while queries are running so we require a restart of the broker for this
+        // change to take effect.
+        LOGGER.warn("Enabling or disabling limitation of the maximum number of multi-stage queries running "
+            + "concurrently requires a restart of the broker to take effect");
+        return;
+      }
+
+      if (maxConcurrentQueries > 0) {
+        _semaphore.setPermits(Math.max(1, maxConcurrentQueries * _numServers / _numBrokers));
+      }

Review Comment:
   > But... isn't that actually going to be problematic when moving from N to M permits
   
   Hmm, I don't think so - we're increasing / decreasing the number of permits based on the previous total number of permits, not the current available number of permits. You can see some of the unit tests like `testIncreaseNumBrokers` / `testIncreaseMaxConcurrentQueries` / `testDecreaseMaxConcurrentQueries` to see how this would look like in action.
   
   So for your example, when the config change from N to M occurs, we don't directly set the number of permits to M. Instead, we either increase the number of permits by `M - N` (if M > N) or decrease the number of permits by `N - M` (if M < N).
   
   > In general I think we need to synchronize in this class anyway
   
   The cluster change handling should be done in a single threaded fashion and the semaphore should support concurrent modification of the number of permits while other threads are trying to acquire / release so I don't think additional external synchronization is necessary here?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on PR #14574:
URL: https://github.com/apache/pinot/pull/14574#issuecomment-2548534525

   +1


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888566742


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).
+ */
+public class MultiStageQueryThrottler implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQueryThrottler.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _numServers;
+  // If _maxConcurrentQueries is <= 0, it means that the cluster is not configured to limit the number of multi-stage
+  // queries that can be executed concurrently. In this case, we should not block the query.
+  private int _maxConcurrentQueries;
+  private AdjustableSemaphore _semaphore;
+
+  @Override
+  public void init(HelixManager helixManager) {
+    _helixManager = helixManager;
+    _helixAdmin = _helixManager.getClusterManagmentTool();
+    _helixConfigScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(
+        _helixManager.getClusterName()).build();
+
+    _maxConcurrentQueries = Integer.parseInt(
+        _helixAdmin.getConfig(_helixConfigScope,
+                Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+            .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+    List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+    _numBrokers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+        .count());
+    _numServers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+        .count());
+
+    if (_maxConcurrentQueries > 0) {
+      _semaphore = new AdjustableSemaphore(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers), true);
+    }
+  }
+
+  /**
+   * Returns true if the query can be executed (waiting until it can be executed if necessary), false otherwise.
+   * <p>
+   * {@link #release()} should be called after the query is done executing. It is the responsibility of the caller to
+   * ensure that {@link #release()} is called exactly once for each call to this method.
+   *
+   * @param timeout the maximum time to wait
+   * @param unit the time unit of the timeout argument
+   * @throws InterruptedException if the current thread is interrupted
+   */
+  public boolean tryAcquire(long timeout, TimeUnit unit)
+      throws InterruptedException {
+    if (_maxConcurrentQueries <= 0) {
+      return true;
+    }
+    return _semaphore.tryAcquire(timeout, unit);
+  }
+
+  /**
+   * Should be called after the query is done executing. It is the responsibility of the caller to ensure that this
+   * method is called exactly once for each call to {@link #tryAcquire(long, TimeUnit)}.
+   */
+  public void release() {
+    if (_maxConcurrentQueries > 0) {
+      _semaphore.release();
+    }
+  }
+
+  @Override
+  public void processClusterChange(HelixConstants.ChangeType changeType) {
+    Preconditions.checkArgument(
+        changeType == HelixConstants.ChangeType.EXTERNAL_VIEW || changeType == HelixConstants.ChangeType.CLUSTER_CONFIG,
+        "MultiStageQuerySemaphore can only handle EXTERNAL_VIEW and CLUSTER_CONFIG changes");
+
+    if (changeType == HelixConstants.ChangeType.EXTERNAL_VIEW) {
+      List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+      int numBrokers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+          .count());
+      int numServers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+          .count());
+
+      if (numBrokers != _numBrokers || numServers != _numServers) {
+        _numBrokers = numBrokers;
+        _numServers = numServers;
+        if (_maxConcurrentQueries > 0) {
+          _semaphore.setPermits(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers));
+        }
+      }
+    } else {
+      int maxConcurrentQueries = Integer.parseInt(
+          _helixAdmin.getConfig(_helixConfigScope,
+                  Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+              .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                  CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+      if (_maxConcurrentQueries == maxConcurrentQueries) {
+        return;
+      }
+
+      if (_maxConcurrentQueries <= 0 && maxConcurrentQueries > 0
+          || _maxConcurrentQueries > 0 && maxConcurrentQueries <= 0) {
+        // This operation isn't safe to do while queries are running so we require a restart of the broker for this
+        // change to take effect.
+        LOGGER.warn("Enabling or disabling limitation of the maximum number of multi-stage queries running "
+            + "concurrently requires a restart of the broker to take effect");
+        return;
+      }
+
+      if (maxConcurrentQueries > 0) {
+        _semaphore.setPermits(Math.max(1, maxConcurrentQueries * _numServers / _numBrokers));
+      }

Review Comment:
   You are right. I though we were calling a _set_ method to change the number of permits instead of sending the increase/decrease.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["yashmayya (via GitHub)" <gi...@apache.org>]

yashmayya merged PR #14574:
URL: https://github.com/apache/pinot/pull/14574


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888384065


##########
pinot-common/src/main/java/org/apache/pinot/common/concurrency/AdjustableSemaphore.java:
##########
@@ -0,0 +1,49 @@
+/**
+ * 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.pinot.common.concurrency;
+
+import java.util.concurrent.Semaphore;
+
+
+/**
+ * A semaphore that allows adjusting the number of permits in a non-blocking way.
+ */
+public class AdjustableSemaphore extends Semaphore {
+
+  private int _totalPermits;
+
+  public AdjustableSemaphore(int permits) {
+    super(permits);
+    _totalPermits = permits;
+  }
+
+  public AdjustableSemaphore(int permits, boolean fair) {
+    super(permits, fair);
+    _totalPermits = permits;
+  }
+
+  public void setPermits(int permits) {

Review Comment:
   Just to fail fast (and in controlled way) I would check that permits > 0



##########
pinot-common/src/main/java/org/apache/pinot/common/concurrency/AdjustableSemaphore.java:
##########
@@ -0,0 +1,49 @@
+/**
+ * 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.pinot.common.concurrency;
+
+import java.util.concurrent.Semaphore;
+
+
+/**
+ * A semaphore that allows adjusting the number of permits in a non-blocking way.
+ */
+public class AdjustableSemaphore extends Semaphore {
+
+  private int _totalPermits;
+
+  public AdjustableSemaphore(int permits) {
+    super(permits);
+    _totalPermits = permits;
+  }
+
+  public AdjustableSemaphore(int permits, boolean fair) {
+    super(permits, fair);
+    _totalPermits = permits;
+  }
+
+  public void setPermits(int permits) {

Review Comment:
   nit: Just to fail fast (and in controlled way) I would check that permits > 0



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["codecov-commenter (via GitHub)" <gi...@apache.org>]

codecov-commenter commented on PR #14574:
URL: https://github.com/apache/pinot/pull/14574#issuecomment-2511160903

   ## [Codecov](https://app.codecov.io/gh/apache/pinot/pull/14574?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
   All modified and coverable lines are covered by tests :white_check_mark:
   > Project coverage is 0.00%. Comparing base [(`59551e4`)](https://app.codecov.io/gh/apache/pinot/commit/59551e45224f1535c4863fd577622b37366ccc97?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) to head [(`38ba313`)](https://app.codecov.io/gh/apache/pinot/commit/38ba313ef4f41e301db640252cf9e3de4c8b5da7?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
   > Report is 1408 commits behind head on master.
   
   > :exclamation:  There is a different number of reports uploaded between BASE (59551e4) and HEAD (38ba313). Click for more details.
   > 
   > <details><summary>HEAD has 53 uploads less than BASE</summary>
   >
   >| Flag | BASE (59551e4) | HEAD (38ba313) |
   >|------|------|------|
   >|integration|7|1|
   >|integration2|3|1|
   >|temurin|12|1|
   >|java-21|7|1|
   >|skip-bytebuffers-true|3|1|
   >|skip-bytebuffers-false|7|0|
   >|unittests|5|0|
   >|unittests1|2|0|
   >|java-11|5|0|
   >|unittests2|3|0|
   >|integration1|2|0|
   >|custom-integration1|2|0|
   ></details>
   
   <details><summary>Additional details and impacted files</summary>
   
   
   ```diff
   @@              Coverage Diff              @@
   ##             master   #14574       +/-   ##
   =============================================
   - Coverage     61.75%    0.00%   -61.76%     
   =============================================
     Files          2436        3     -2433     
     Lines        133233        6   -133227     
     Branches      20636        0    -20636     
   =============================================
   - Hits          82274        0    -82274     
   + Misses        44911        6    -44905     
   + Partials       6048        0     -6048     
   ```
   
   | [Flag](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
   |---|---|---|
   | [custom-integration1](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [integration](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `0.00% <ø> (-0.01%)` | :arrow_down: |
   | [integration1](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [integration2](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `0.00% <ø> (ø)` | |
   | [java-11](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [java-21](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `0.00% <ø> (-61.63%)` | :arrow_down: |
   | [skip-bytebuffers-false](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [skip-bytebuffers-true](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `0.00% <ø> (-27.73%)` | :arrow_down: |
   | [temurin](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `0.00% <ø> (-61.76%)` | :arrow_down: |
   | [unittests](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [unittests1](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   | [unittests2](https://app.codecov.io/gh/apache/pinot/pull/14574/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
   
   Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
   
   </details>
   
   [:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/apache/pinot/pull/14574?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).   
   :loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888355290


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster

Review Comment:
   nit: This javadoc says twice that we could move the check to server level.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["yashmayya (via GitHub)" <gi...@apache.org>]

yashmayya commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1886880302


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java:
##########
@@ -224,67 +227,89 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
       return new BrokerResponseNative(QueryException.getException(QueryException.QUOTA_EXCEEDED_ERROR, errorMessage));
     }
 
-    Tracing.ThreadAccountantOps.setupRunner(String.valueOf(requestId), ThreadExecutionContext.TaskType.MSE);
-
-    long executionStartTimeNs = System.nanoTime();
-    QueryDispatcher.QueryResult queryResults;
+    Timer queryTimer = new Timer(queryTimeoutMs);
     try {
-      queryResults =
-          _queryDispatcher.submitAndReduce(requestContext, dispatchableSubPlan, queryTimeoutMs, queryOptions);
-    } catch (TimeoutException e) {
-      for (String table : tableNames) {
-        _brokerMetrics.addMeteredTableValue(table, BrokerMeter.BROKER_RESPONSES_WITH_TIMEOUTS, 1);
+      // It's fine to block in this thread because we use a separate thread pool from the main Jersey server to process
+      // these requests.
+      if (!_querySemaphore.tryAcquire(queryTimeoutMs, TimeUnit.MILLISECONDS)) {
+        LOGGER.warn("Timed out waiting to execute request {}: {}", requestId, query);
+        requestContext.setErrorCode(QueryException.EXECUTION_TIMEOUT_ERROR_CODE);
+        return new BrokerResponseNative(QueryException.EXECUTION_TIMEOUT_ERROR);
       }
-      LOGGER.warn("Timed out executing request {}: {}", requestId, query);
+    } catch (InterruptedException e) {
+      LOGGER.warn("Interrupt received while waiting to execute request {}: {}", requestId, query);
       requestContext.setErrorCode(QueryException.EXECUTION_TIMEOUT_ERROR_CODE);
       return new BrokerResponseNative(QueryException.EXECUTION_TIMEOUT_ERROR);

Review Comment:
   We could either add a debug / info log showing this queueing / throttling mechanism based wait time per query or add a new metric to show the number of queries that timed out while waiting to be dispatched (or both). WDYT?



##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQuerySemaphore.java:
##########
@@ -0,0 +1,174 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future.
+ */
+public class MultiStageQuerySemaphore extends Semaphore implements ClusterChangeHandler {

Review Comment:
   Good point, I've split up this class into two new ones - `MultiStageQueryThrottler` which implements the `ClusterChangeHandler` interface and `AdjustableSemaphore` which is now a generic util class. I've preferred to keep the minimal interfaces because I don't feel like the amount of mocking required for testing here is too high, but I'd be happy to reconsider if you feel strongly about it.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888381962


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageQueryThrottler.java:
##########
@@ -0,0 +1,168 @@
+/**
+ * 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.pinot.broker.requesthandler;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants;
+import org.apache.helix.HelixManager;
+import org.apache.helix.model.HelixConfigScope;
+import org.apache.helix.model.builder.HelixConfigScopeBuilder;
+import org.apache.pinot.broker.broker.helix.ClusterChangeHandler;
+import org.apache.pinot.common.concurrency.AdjustableSemaphore;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * This class helps limit the number of multi-stage queries being executed concurrently. Currently, this limit is
+ * applied at the broker level, but could be moved to the server level in the future. Note that the cluster
+ * configuration is a "per server" value and the broker currently simply assumes that a query will be across all
+ * servers. Another assumption here is that queries are evenly distributed across brokers. Ideally, we want to move to a
+ * model where the broker asks each server whether it can execute a query stage before dispatching the query stage to
+ * the server. This would allow for more fine-grained control over the number of queries being executed concurrently
+ * (but there are some limitations around ordering and blocking that need to be solved first).
+ */
+public class MultiStageQueryThrottler implements ClusterChangeHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiStageQueryThrottler.class);
+
+  private HelixManager _helixManager;
+  private HelixAdmin _helixAdmin;
+  private HelixConfigScope _helixConfigScope;
+  private int _numBrokers;
+  private int _numServers;
+  // If _maxConcurrentQueries is <= 0, it means that the cluster is not configured to limit the number of multi-stage
+  // queries that can be executed concurrently. In this case, we should not block the query.
+  private int _maxConcurrentQueries;
+  private AdjustableSemaphore _semaphore;
+
+  @Override
+  public void init(HelixManager helixManager) {
+    _helixManager = helixManager;
+    _helixAdmin = _helixManager.getClusterManagmentTool();
+    _helixConfigScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).forCluster(
+        _helixManager.getClusterName()).build();
+
+    _maxConcurrentQueries = Integer.parseInt(
+        _helixAdmin.getConfig(_helixConfigScope,
+                Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+            .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+    List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+    _numBrokers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+        .count());
+    _numServers = Math.max(1, (int) clusterInstances.stream()
+        .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+        .count());
+
+    if (_maxConcurrentQueries > 0) {
+      _semaphore = new AdjustableSemaphore(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers), true);
+    }
+  }
+
+  /**
+   * Returns true if the query can be executed (waiting until it can be executed if necessary), false otherwise.
+   * <p>
+   * {@link #release()} should be called after the query is done executing. It is the responsibility of the caller to
+   * ensure that {@link #release()} is called exactly once for each call to this method.
+   *
+   * @param timeout the maximum time to wait
+   * @param unit the time unit of the timeout argument
+   * @throws InterruptedException if the current thread is interrupted
+   */
+  public boolean tryAcquire(long timeout, TimeUnit unit)
+      throws InterruptedException {
+    if (_maxConcurrentQueries <= 0) {
+      return true;
+    }
+    return _semaphore.tryAcquire(timeout, unit);
+  }
+
+  /**
+   * Should be called after the query is done executing. It is the responsibility of the caller to ensure that this
+   * method is called exactly once for each call to {@link #tryAcquire(long, TimeUnit)}.
+   */
+  public void release() {
+    if (_maxConcurrentQueries > 0) {
+      _semaphore.release();
+    }
+  }
+
+  @Override
+  public void processClusterChange(HelixConstants.ChangeType changeType) {
+    Preconditions.checkArgument(
+        changeType == HelixConstants.ChangeType.EXTERNAL_VIEW || changeType == HelixConstants.ChangeType.CLUSTER_CONFIG,
+        "MultiStageQuerySemaphore can only handle EXTERNAL_VIEW and CLUSTER_CONFIG changes");
+
+    if (changeType == HelixConstants.ChangeType.EXTERNAL_VIEW) {
+      List<String> clusterInstances = _helixAdmin.getInstancesInCluster(_helixManager.getClusterName());
+      int numBrokers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE))
+          .count());
+      int numServers = Math.max(1, (int) clusterInstances.stream()
+          .filter(instance -> instance.startsWith(CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE))
+          .count());
+
+      if (numBrokers != _numBrokers || numServers != _numServers) {
+        _numBrokers = numBrokers;
+        _numServers = numServers;
+        if (_maxConcurrentQueries > 0) {
+          _semaphore.setPermits(Math.max(1, _maxConcurrentQueries * _numServers / _numBrokers));
+        }
+      }
+    } else {
+      int maxConcurrentQueries = Integer.parseInt(
+          _helixAdmin.getConfig(_helixConfigScope,
+                  Collections.singletonList(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES))
+              .getOrDefault(CommonConstants.Helix.CONFIG_OF_MAX_CONCURRENT_MULTI_STAGE_QUERIES,
+                  CommonConstants.Helix.DEFAULT_MAX_CONCURRENT_MULTI_STAGE_QUERIES));
+
+      if (_maxConcurrentQueries == maxConcurrentQueries) {
+        return;
+      }
+
+      if (_maxConcurrentQueries <= 0 && maxConcurrentQueries > 0
+          || _maxConcurrentQueries > 0 && maxConcurrentQueries <= 0) {
+        // This operation isn't safe to do while queries are running so we require a restart of the broker for this
+        // change to take effect.
+        LOGGER.warn("Enabling or disabling limitation of the maximum number of multi-stage queries running "
+            + "concurrently requires a restart of the broker to take effect");
+        return;
+      }
+
+      if (maxConcurrentQueries > 0) {
+        _semaphore.setPermits(Math.max(1, maxConcurrentQueries * _numServers / _numBrokers));
+      }

Review Comment:
   I think there is a possible NPE here. Imagine that the broker is initialized with -1 concurrent queries, `_semaphore` is kept as null. Then the config is changed to set an actual max concurrent queries, then we try to call `_semaphore.setPermits` here without actually creating the semaphore.
   
   For simplicity, I would suggest to always allocate the semaphore (with whatever permits). Given don't delegate on the semaphore when `_maxConcurrentQueries <= 0`, that should be a clean and correct solution



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org

Re: [PR] Add cluster configuration to allow limiting the number of multi-stage queries running concurrently [pinot]

["gortiz (via GitHub)" <gi...@apache.org>]

gortiz commented on code in PR #14574:
URL: https://github.com/apache/pinot/pull/14574#discussion_r1888387562


##########
pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java:
##########
@@ -224,67 +227,89 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
       return new BrokerResponseNative(QueryException.getException(QueryException.QUOTA_EXCEEDED_ERROR, errorMessage));
     }
 
-    Tracing.ThreadAccountantOps.setupRunner(String.valueOf(requestId), ThreadExecutionContext.TaskType.MSE);
-
-    long executionStartTimeNs = System.nanoTime();
-    QueryDispatcher.QueryResult queryResults;
+    Timer queryTimer = new Timer(queryTimeoutMs);
     try {
-      queryResults =
-          _queryDispatcher.submitAndReduce(requestContext, dispatchableSubPlan, queryTimeoutMs, queryOptions);
-    } catch (TimeoutException e) {
-      for (String table : tableNames) {
-        _brokerMetrics.addMeteredTableValue(table, BrokerMeter.BROKER_RESPONSES_WITH_TIMEOUTS, 1);
+      // It's fine to block in this thread because we use a separate thread pool from the main Jersey server to process
+      // these requests.
+      if (!_querySemaphore.tryAcquire(queryTimeoutMs, TimeUnit.MILLISECONDS)) {
+        LOGGER.warn("Timed out waiting to execute request {}: {}", requestId, query);
+        requestContext.setErrorCode(QueryException.EXECUTION_TIMEOUT_ERROR_CODE);
+        return new BrokerResponseNative(QueryException.EXECUTION_TIMEOUT_ERROR);
       }
-      LOGGER.warn("Timed out executing request {}: {}", requestId, query);
+    } catch (InterruptedException e) {
+      LOGGER.warn("Interrupt received while waiting to execute request {}: {}", requestId, query);
       requestContext.setErrorCode(QueryException.EXECUTION_TIMEOUT_ERROR_CODE);
       return new BrokerResponseNative(QueryException.EXECUTION_TIMEOUT_ERROR);

Review Comment:
   a metric would be great, a log could be useful. But I was more worried about immediate the feedback to the user. I mean, the error information we return. But it was a wishful thinking more than an actual requirement.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org