[Coordinator throttling] Scheduling Policies for Admission Control based on worker load (#25689)

**Admission control scheduling policy**

**Logic**
Gather worker overload data from the added end point in PR -
[https://github.com/prestodb/presto/pull/25687](https://l.facebook.com/l.php?u=https%3A%2F%2Fgithub.com%2Fprestodb%2Fpresto%2Fpull%2F25687&h=AT20QdO6Ld9h4KN4cMS_xFUtKTtNuKLc-glQIqSWbPgMh0KxW_3-UYhiBxO2gl-I8P_i8wG5gixSBWcGmglCWFShRKtOUlWnYZ5RqgBNHhmkGYrK2XJfKIjqXNj4Ltd4oWghm0eSCgnKL0UGvLr9mXLyJyA)
Based on configured policies (cnt of overloaded workers or pct of
overloaded workers) and cluster overload, queue the queries

**Background**
RFC PR:
[https://github.com/prestodb/rfcs/pull/42](https://l.facebook.com/l.php?u=https%3A%2F%2Fgithub.com%2Fprestodb%2Frfcs%2Fpull%2F42&h=AT20QdO6Ld9h4KN4cMS_xFUtKTtNuKLc-glQIqSWbPgMh0KxW_3-UYhiBxO2gl-I8P_i8wG5gixSBWcGmglCWFShRKtOUlWnYZ5RqgBNHhmkGYrK2XJfKIjqXNj4Ltd4oWghm0eSCgnKL0UGvLr9mXLyJyA)

**Metrics on queuing due to this feature:**
Added following metrics
- ClusterOverloadDuration
- ClusterOverloadCount

**Feature flag:** 
Right now feature is disabled. We can use coordinator configs to enable
/ add thresholds

## Summary by Sourcery

Add cluster overload-based admission control scheduling policy and
supporting infrastructure to throttle queries based on worker load.

New Features:
- Introduce ClusterResourceChecker for periodic cluster overload
detection and query throttling based on worker load metrics
- Implement CpuMemoryOverloadPolicy with count- and percentage-based
overload thresholds and expose overload detection count and duration
metrics
- Add AdmissionControlBypassConfig to allow certain queries (e.g., DDL)
to bypass admission control

Enhancements:
- Refactor DiscoveryNodeManager and SPI to use RemoteNodeStats for
asynchronous fetching of node load metrics
- Integrate cluster overload checks into InternalResourceGroup
scheduling logic and update eligibility propagation on overload state
changes
- Introduce ClusterOverloadPolicyModule and DI bindings for policy
selection and checker

Tests:
- Add unit tests for ClusterResourceChecker, CpuMemoryOverloadPolicy,
and AdmissionControlBypassConfig
- Update existing resource group and node manager tests to support
clustering resource checker mocks

Chores:
- Add new configuration properties for cluster-overload throttling and
internal-communication stats polling intervals

Author: prashantgolash

presto-main-base/src/main/java/com/facebook/presto/execution/ClusterOverloadConfig.java

@@ -0,0 +1,111 @@
+/*
+ * Licensed 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 com.facebook.presto.execution;
+
+import com.facebook.airlift.configuration.Config;
+
+public class ClusterOverloadConfig
+{
+    public static final String OVERLOAD_POLICY_CNT_BASED = "overload_worker_cnt_based_throttling";
+    public static final String OVERLOAD_POLICY_PCT_BASED = "overload_worker_pct_based_throttling";
+    private boolean clusterOverloadThrottlingEnabled;
+    private double allowedOverloadWorkersPct = 0.01;
+    private int allowedOverloadWorkersCnt;
+    private String overloadPolicyType = OVERLOAD_POLICY_CNT_BASED;
+    private int overloadCheckCacheTtlInSecs = 5;
+
+    /**
+     * Gets the time-to-live for the cached cluster overload state.
+     * This determines how frequently the system will re-evaluate whether the cluster is overloaded.
+     *
+     * @return the cache TTL duration
+     */
+    public int getOverloadCheckCacheTtlInSecs()
+    {
+        return overloadCheckCacheTtlInSecs;
+    }
+
+    /**
+     * Gets the time-to-live for the cached cluster overload state.
+     * This determines how frequently the system will re-evaluate whether the cluster is overloaded.
+     *
+     * @return the cache TTL duration
+     */
+    public int getOverloadCheckCacheTtlMillis()
+    {
+        return overloadCheckCacheTtlInSecs * 1000;
+    }
+
+    /**
+     * Sets the time-to-live for the cached cluster overload state.
+     *
+     * @param overloadCheckCacheTtlInSecs the cache TTL duration
+     * @return this for chaining
+     */
+    @Config("cluster.overload-check-cache-ttl-secs")
+    public ClusterOverloadConfig setOverloadCheckCacheTtlInSecs(int overloadCheckCacheTtlInSecs)
+    {
+        this.overloadCheckCacheTtlInSecs = overloadCheckCacheTtlInSecs;
+        return this;
+    }
+
+    @Config("cluster-overload.enable-throttling")
+    public ClusterOverloadConfig setClusterOverloadThrottlingEnabled(boolean clusterOverloadThrottlingEnabled)
+    {
+        this.clusterOverloadThrottlingEnabled = clusterOverloadThrottlingEnabled;
+        return this;
+    }
+
+    public boolean isClusterOverloadThrottlingEnabled()
+    {
+        return this.clusterOverloadThrottlingEnabled;
+    }
+
+    @Config("cluster-overload.allowed-overload-workers-pct")
+    public ClusterOverloadConfig setAllowedOverloadWorkersPct(Double allowedOverloadWorkersPct)
+    {
+        this.allowedOverloadWorkersPct = allowedOverloadWorkersPct;
+        return this;
+    }
+
+    public double getAllowedOverloadWorkersPct()
+    {
+        return this.allowedOverloadWorkersPct;
+    }
+
+    @Config("cluster-overload.allowed-overload-workers-cnt")
+    public ClusterOverloadConfig setAllowedOverloadWorkersCnt(int allowedOverloadWorkersCnt)
+    {
+        this.allowedOverloadWorkersCnt = allowedOverloadWorkersCnt;
+        return this;
+    }
+
+    public double getAllowedOverloadWorkersCnt()
+    {
+        return this.allowedOverloadWorkersCnt;
+    }
+
+    @Config("cluster-overload.overload-policy-type")
+    public ClusterOverloadConfig setOverloadPolicyType(String overloadPolicyType)
+    {
+        // validate
+        this.overloadPolicyType = overloadPolicyType;
+        return this;
+    }
+
+    public String getOverloadPolicyType()
+    {
+        return this.overloadPolicyType;
+    }
+}

presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java

@@ -18,6 +18,7 @@
 import com.facebook.airlift.units.Duration;
 import com.facebook.presto.execution.ManagedQueryExecution;
 import com.facebook.presto.execution.resourceGroups.WeightedFairQueue.Usage;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker;
 import com.facebook.presto.metadata.InternalNodeManager;
 import com.facebook.presto.server.QueryStateInfo;
 import com.facebook.presto.server.ResourceGroupInfo;
@@ -96,6 +97,7 @@ public class InternalResourceGroup
     private final Function<ResourceGroupId, Optional<ResourceGroupRuntimeInfo>> additionalRuntimeInfo;
     private final Predicate<InternalResourceGroup> shouldWaitForResourceManagerUpdate;
     private final InternalNodeManager nodeManager;
+    private final ClusterResourceChecker clusterResourceChecker;
 
     // Configuration
     // =============
@@ -166,12 +168,14 @@ protected InternalResourceGroup(
             boolean staticResourceGroup,
             Function<ResourceGroupId, Optional<ResourceGroupRuntimeInfo>> additionalRuntimeInfo,
             Predicate<InternalResourceGroup> shouldWaitForResourceManagerUpdate,
-            InternalNodeManager nodeManager)
+            InternalNodeManager nodeManager,
+            ClusterResourceChecker clusterResourceChecker)
     {
         this.parent = requireNonNull(parent, "parent is null");
         this.jmxExportListener = requireNonNull(jmxExportListener, "jmxExportListener is null");
         this.executor = requireNonNull(executor, "executor is null");
         this.nodeManager = requireNonNull(nodeManager, "node manager is null");
+        this.clusterResourceChecker = requireNonNull(clusterResourceChecker, "clusterResourceChecker is null");
         requireNonNull(name, "name is null");
         if (parent.isPresent()) {
             id = new ResourceGroupId(parent.get().id, name);
@@ -671,7 +675,8 @@ public InternalResourceGroup getOrCreateSubGroup(String name, boolean staticSegm
                     staticResourceGroup && staticSegment,
                     additionalRuntimeInfo,
                     shouldWaitForResourceManagerUpdate,
-                    nodeManager);
+                    nodeManager,
+                    clusterResourceChecker);
             // Sub group must use query priority to ensure ordering
             if (schedulingPolicy == QUERY_PRIORITY) {
                 subGroup.setSchedulingPolicy(QUERY_PRIORITY);
@@ -770,7 +775,7 @@ private void enqueueQuery(ManagedQueryExecution query)
     }
 
     // This method must be called whenever the group's eligibility to run more queries may have changed.
-    private void updateEligibility()
+    protected void updateEligibility()
     {
         checkState(Thread.holdsLock(root), "Must hold lock to update eligibility");
         synchronized (root) {
@@ -1019,6 +1024,11 @@ private boolean canRunMore()
     {
         checkState(Thread.holdsLock(root), "Must hold lock");
         synchronized (root) {
+            // Check if more queries can be run on the cluster based on cluster overload
+            if (clusterResourceChecker.isClusterCurrentlyOverloaded()) {
+                return false;
+            }
+
             if (cpuUsageMillis >= hardCpuLimitMillis) {
                 return false;
             }
@@ -1135,7 +1145,8 @@ public RootInternalResourceGroup(
                 Executor executor,
                 Function<ResourceGroupId, Optional<ResourceGroupRuntimeInfo>> additionalRuntimeInfo,
                 Predicate<InternalResourceGroup> shouldWaitForResourceManagerUpdate,
-                InternalNodeManager nodeManager)
+                InternalNodeManager nodeManager,
+                ClusterResourceChecker clusterResourceChecker)
         {
             super(Optional.empty(),
                     name,
@@ -1144,7 +1155,16 @@ public RootInternalResourceGroup(
                     true,
                     additionalRuntimeInfo,
                     shouldWaitForResourceManagerUpdate,
-                    nodeManager);
+                    nodeManager,
+                    clusterResourceChecker);
+        }
+
+        public synchronized void updateEligibilityRecursively(InternalResourceGroup group)
+        {
+            group.updateEligibility();
+            for (InternalResourceGroup subGroup : group.subGroups()) {
+                updateEligibilityRecursively(subGroup);
+            }
         }
 
         public synchronized void processQueuedQueries()

presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroupManager.java

@@ -19,6 +19,8 @@
 import com.facebook.presto.execution.ManagedQueryExecution;
 import com.facebook.presto.execution.QueryManagerConfig;
 import com.facebook.presto.execution.resourceGroups.InternalResourceGroup.RootInternalResourceGroup;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadStateListener;
+import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker;
 import com.facebook.presto.metadata.InternalNodeManager;
 import com.facebook.presto.resourcemanager.ResourceGroupService;
 import com.facebook.presto.server.ResourceGroupInfo;
@@ -81,7 +83,7 @@
 
 @ThreadSafe
 public final class InternalResourceGroupManager<C>
-        implements ResourceGroupManager<C>
+        implements ResourceGroupManager<C>, ClusterOverloadStateListener
 {
     private static final Logger log = Logger.get(InternalResourceGroupManager.class);
     private static final File RESOURCE_GROUPS_CONFIGURATION = new File("etc/resource-groups.properties");
@@ -112,6 +114,7 @@ public final class InternalResourceGroupManager<C>
     private final QueryManagerConfig queryManagerConfig;
     private final InternalNodeManager nodeManager;
     private AtomicBoolean isConfigurationManagerLoaded;
+    private final ClusterResourceChecker clusterResourceChecker;
 
     @Inject
     public InternalResourceGroupManager(
@@ -121,7 +124,8 @@ public InternalResourceGroupManager(
             MBeanExporter exporter,
             ResourceGroupService resourceGroupService,
             ServerConfig serverConfig,
-            InternalNodeManager nodeManager)
+            InternalNodeManager nodeManager,
+            ClusterResourceChecker clusterResourceChecker)
     {
         this.queryManagerConfig = requireNonNull(queryManagerConfig, "queryManagerConfig is null");
         this.exporter = requireNonNull(exporter, "exporter is null");
@@ -137,6 +141,7 @@ public InternalResourceGroupManager(
         this.resourceGroupRuntimeExecutor = new PeriodicTaskExecutor(resourceGroupRuntimeInfoRefreshInterval.toMillis(), refreshExecutor, this::refreshResourceGroupRuntimeInfo);
         configurationManagerFactories.putIfAbsent(LegacyResourceGroupConfigurationManager.NAME, new LegacyResourceGroupConfigurationManager.Factory());
         this.isConfigurationManagerLoaded = new AtomicBoolean(false);
+        this.clusterResourceChecker = clusterResourceChecker;
     }
 
     @Override
@@ -254,6 +259,8 @@ public ResourceGroupConfigurationManager<C> getConfigurationManager()
     @PreDestroy
     public void destroy()
     {
+        // Unregister from cluster overload state changes
+        clusterResourceChecker.removeListener(this);
         refreshExecutor.shutdownNow();
         resourceGroupRuntimeExecutor.stop();
     }
@@ -275,6 +282,9 @@ public void start()
             if (isResourceManagerEnabled) {
                 resourceGroupRuntimeExecutor.start();
             }
+
+            // Register as listener for cluster overload state changes
+            clusterResourceChecker.addListener(this);
         }
     }
 
@@ -396,7 +406,7 @@ private synchronized void createGroupIfNecessary(SelectionContext<C> context, Ex
             else {
                 RootInternalResourceGroup root;
                 if (!isResourceManagerEnabled) {
-                    root = new RootInternalResourceGroup(id.getSegments().get(0), this::exportGroup, executor, ignored -> Optional.empty(), rg -> false, nodeManager);
+                    root = new RootInternalResourceGroup(id.getSegments().get(0), this::exportGroup, executor, ignored -> Optional.empty(), rg -> false, nodeManager, clusterResourceChecker);
                 }
                 else {
                     root = new RootInternalResourceGroup(
@@ -409,7 +419,8 @@ private synchronized void createGroupIfNecessary(SelectionContext<C> context, Ex
                                     resourceGroupRuntimeInfosSnapshot::get,
                                     lastUpdatedResourceGroupRuntimeInfo::get,
                                     concurrencyThreshold),
-                            nodeManager);
+                            nodeManager,
+                            clusterResourceChecker);
                 }
                 group = root;
                 rootGroups.add(root);
@@ -463,6 +474,24 @@ public int getQueriesQueuedOnInternal()
         return queriesQueuedInternal;
     }
 
+    @Override
+    public void onClusterEnteredOverloadedState()
+    {
+        // Resource groups will handle overload state through their existing admission control logic
+        // No additional action needed here as queries will be queued automatically
+    }
+
+    @Override
+    public void onClusterExitedOverloadedState()
+    {
+        log.info("Cluster exited overloaded state, updating eligibility for all resource groups");
+        for (RootInternalResourceGroup rootGroup : rootGroups) {
+            synchronized (rootGroup) {
+                rootGroup.updateEligibilityRecursively(rootGroup);
+            }
+        }
+    }
+
     @Managed
     public long getLastSchedulingCycleRuntimeDelayMs()
     {

presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicy.java

@@ -0,0 +1,40 @@
+/*
+ * Licensed 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 com.facebook.presto.execution.scheduler.clusterOverload;
+
+import com.facebook.presto.metadata.InternalNodeManager;
+
+/**
+ * Interface for policies that determine if cluster is overloaded.
+ * Implementations can check various metrics from NodeStats to determine
+ * if a worker is overloaded and queries should be throttled.
+ */
+public interface ClusterOverloadPolicy
+{
+    /**
+     * Checks if cluster is overloaded.
+     *
+     * @param nodeManager The node manager to get node information
+     * @return true if cluster is overloaded, false otherwise
+     */
+    boolean isClusterOverloaded(InternalNodeManager nodeManager);
+
+    /**
+     * Gets the name of the policy.
+     *
+     * @return The name of the policy
+     */
+    String getName();
+}