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

Summary:
Pull Request resolved: prestodb#25689

Admission control scheduling policy

**Logic**
Gather worker overload data from the added end point in PR - prestodb#25687
Based on configured policies (cnt of overloaded workers or pct of overloaded workers) and cluster overload, queue the queries

**Background**
RFC PR: prestodb/rfcs#42

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

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

Differential Revision: D79470181

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,21 +18,25 @@
 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;
 import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.resourceGroups.AdmissionControlBypassConfig;
 import com.facebook.presto.spi.resourceGroups.ResourceGroup;
 import com.facebook.presto.spi.resourceGroups.ResourceGroupId;
 import com.facebook.presto.spi.resourceGroups.ResourceGroupQueryLimits;
 import com.facebook.presto.spi.resourceGroups.ResourceGroupState;
 import com.facebook.presto.spi.resourceGroups.SchedulingPolicy;
+import com.facebook.presto.spi.resourceGroups.SelectionCriteria;
 import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.ThreadSafe;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import org.weakref.jmx.Managed;
 import org.weakref.jmx.Nested;
 
+import java.security.Principal;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -96,6 +100,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
     // =============
@@ -123,6 +128,8 @@ public class InternalResourceGroup
     private boolean jmxExport;
     @GuardedBy("root")
     private ResourceGroupQueryLimits perQueryLimits = NO_LIMITS;
+    @GuardedBy("root")
+    private AdmissionControlBypassConfig admissionControlPolicyByPassConfig = AdmissionControlBypassConfig.NO_BYPASS;
 
     // Live data structures
     // ====================
@@ -166,12 +173,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);
@@ -653,6 +662,22 @@ public ResourceGroupQueryLimits getPerQueryLimits()
         }
     }
 
+    @Override
+    public void setAdmissionControlPolicyByPassConfig(AdmissionControlBypassConfig admissionControlPolicyByPassConfig)
+    {
+        synchronized (root) {
+            this.admissionControlPolicyByPassConfig = admissionControlPolicyByPassConfig;
+        }
+    }
+
+    @Override
+    public AdmissionControlBypassConfig getAdmissionControlPolicyByPassConfig()
+    {
+        synchronized (root) {
+            return admissionControlPolicyByPassConfig;
+        }
+    }
+
     public InternalResourceGroup getOrCreateSubGroup(String name, boolean staticSegment)
     {
         requireNonNull(name, "name is null");
@@ -671,7 +696,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);
@@ -713,25 +739,34 @@ public void run(ManagedQueryExecution query)
             if (!subGroups.isEmpty()) {
                 throw new PrestoException(INVALID_RESOURCE_GROUP, format("Cannot add queries to %s. It is not a leaf group.", id));
             }
-            // Check all ancestors for capacity
-            InternalResourceGroup group = this;
+
             boolean canQueue = true;
             boolean canRun = true;
-            while (true) {
-                canQueue &= group.canQueueMore();
-                canRun &= group.canRunMore();
-                if (!group.parent.isPresent()) {
-                    break;
+
+            // Check if admission control could be skipped for this query based on configured RG policies
+            SelectionCriteria selectionCriteria = new SelectionCriteria(query.getSession().getIdentity().getPrincipal().isPresent(), query.getSession().getUser(), query.getSession().getSource(), query.getSession().getClientTags(), query.getSession().getResourceEstimates(), query.getSession().getQueryType().map(Enum::name), query.getSession().getClientInfo(), query.getSession().getSchema(), query.getSession().getIdentity().getPrincipal().map(Principal::getName));
+            boolean shouldBypassAdmissionControl = admissionControlPolicyByPassConfig.shouldBypassAdmissionControl(selectionCriteria);
+
+            if (!shouldBypassAdmissionControl) {
+                // Check all ancestors for capacity
+                InternalResourceGroup group = this;
+                while (true) {
+                    canQueue &= group.canQueueMore();
+                    canRun &= group.canRunMore();
+                    if (!group.parent.isPresent()) {
+                        break;
+                    }
+                    group = group.parent.get();
                 }
-                group = group.parent.get();
             }
+
             if (!canQueue && !canRun) {
                 isQueryQueueFull = true;
             }
             else {
                 query.setResourceGroupQueryLimits(perQueryLimits);
-                if (canRun && queuedQueries.isEmpty()) {
-                    startInBackground(query);
+                if (shouldBypassAdmissionControl || (canRun && queuedQueries.isEmpty())) {
+                    startInBackground(query, shouldBypassAdmissionControl);
                 }
                 else {
                     enqueueQuery(query);
@@ -770,7 +805,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) {
@@ -790,18 +825,20 @@ private void updateEligibility()
         }
     }
 
-    private void startInBackground(ManagedQueryExecution query)
+    private void startInBackground(ManagedQueryExecution query, boolean shouldBypassAdmissionControl)
     {
         checkState(Thread.holdsLock(root), "Must hold lock to start a query");
         synchronized (root) {
-            runningQueries.add(query);
             InternalResourceGroup group = this;
-            while (group.parent.isPresent()) {
-                group.parent.get().descendantRunningQueries++;
-                group.parent.get().dirtySubGroups.add(group);
-                group = group.parent.get();
+            if (!shouldBypassAdmissionControl) {
+                runningQueries.add(query);
+                while (group.parent.isPresent()) {
+                    group.parent.get().descendantRunningQueries++;
+                    group.parent.get().dirtySubGroups.add(group);
+                    group = group.parent.get();
+                }
+                updateEligibility();
             }
-            updateEligibility();
             executor.execute(query::startWaitingForResources);
             group = this;
             long lastRunningQueryStartTimeMillis = currentTimeMillis();
@@ -905,7 +942,7 @@ protected boolean internalStartNext()
 
             ManagedQueryExecution query = queuedQueries.poll();
             if (query != null) {
-                startInBackground(query);
+                startInBackground(query, false);
                 return true;
             }
 
@@ -1019,6 +1056,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 +1177,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 +1187,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()