Master node changes for master task throttling

server/src/main/java/org/opensearch/cluster/ClusterStateTaskExecutor.java

@@ -78,6 +78,12 @@ default String describeTasks(List<T> tasks) {
         return String.join(", ", tasks.stream().map(t -> (CharSequence) t.toString()).filter(t -> t.length() > 0)::iterator);
     }
 
+    public static final String DEFAULT_MASTER_THROTTLING_KEY = "";
+
+    default String getMasterThrottlingKey() {
+        return DEFAULT_MASTER_THROTTLING_KEY;
+    }
+
     /**
      * Represents the result of a batched execution of cluster state update tasks
      * @param <T> the type of the cluster state update task

server/src/main/java/org/opensearch/cluster/service/MasterService.java

@@ -121,6 +121,18 @@ public class MasterService extends AbstractLifecycleComponent {
 
     private volatile PrioritizedOpenSearchThreadPoolExecutor threadPoolExecutor;
     private volatile Batcher taskBatcher;
+    /**
+     * Throttler which performs throttling based on pending task count per task type.
+     * It acquire permits before task goes into queue and release the permit when task is removed from the queue.
+     *
+     * It acquire permits at one place, while submitting the tasks.
+     * It releases from three places.
+     * 1. If submit task is failed due to some exception (Since we already acquire permit and submit is failing so we need to release it).
+     * 2. Release before execution ( Task is removed from queue and went into execution).
+     * 3. Task is timed out in queue(Task is timed out in queue, so master will throw exception to data nodes. Releasing as task is removed from queue)
+     */
+    protected final MasterTaskThrottler masterTaskThrottler;
+    private final MasterThrottlingStats throttlingStats;
 
     public MasterService(Settings settings, ClusterSettings clusterSettings, ThreadPool threadPool) {
         this.nodeName = Objects.requireNonNull(Node.NODE_NAME_SETTING.get(settings));
@@ -131,6 +143,8 @@ public MasterService(Settings settings, ClusterSettings clusterSettings, ThreadP
             this::setSlowTaskLoggingThreshold
         );
 
+        this.throttlingStats = new MasterThrottlingStats();
+        this.masterTaskThrottler = new MasterTaskThrottler(clusterSettings, this, throttlingStats);
         this.threadPool = threadPool;
     }
 
@@ -151,7 +165,7 @@ protected synchronized void doStart() {
         Objects.requireNonNull(clusterStatePublisher, "please set a cluster state publisher before starting");
         Objects.requireNonNull(clusterStateSupplier, "please set a cluster state supplier before starting");
         threadPoolExecutor = createThreadPoolExecutor();
-        taskBatcher = new Batcher(logger, threadPoolExecutor);
+        taskBatcher = new Batcher(logger, threadPoolExecutor, masterTaskThrottler);
     }
 
     protected PrioritizedOpenSearchThreadPoolExecutor createThreadPoolExecutor() {
@@ -166,8 +180,8 @@ protected PrioritizedOpenSearchThreadPoolExecutor createThreadPoolExecutor() {
     @SuppressWarnings("unchecked")
     class Batcher extends TaskBatcher {
 
-        Batcher(Logger logger, PrioritizedOpenSearchThreadPoolExecutor threadExecutor) {
-            super(logger, threadExecutor);
+        Batcher(Logger logger, PrioritizedOpenSearchThreadPoolExecutor threadExecutor, TaskBatcherListener taskBatcherListener) {
+            super(logger, threadExecutor, taskBatcherListener);
         }
 
         @Override
@@ -570,6 +584,13 @@ public List<PendingClusterTask> pendingTasks() {
         }).collect(Collectors.toList());
     }
 
+    /**
+     * Returns the number of throttled pending tasks.
+     */
+    public long numberOfThrottledPendingTasks() {
+        return throttlingStats.getTotalThrottledTaskCount();
+    }
+
     /**
      * Returns the number of currently pending tasks.
      */

server/src/main/java/org/opensearch/cluster/service/MasterTaskThrottler.java

@@ -0,0 +1,146 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.cluster.service;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.Version;
+import org.opensearch.cluster.ClusterStateTaskExecutor;
+import org.opensearch.common.settings.ClusterSettings;
+import org.opensearch.common.settings.Setting;
+import org.opensearch.common.settings.Settings;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+/**
+ * This class is extension of {@link Throttler} and does throttling of master tasks.
+ *
+ * This class does throttling on task submission to master, it uses class name of request of tasks as key for
+ * throttling. Throttling will be performed over task executor's class level, different task types have different executors class.
+ *
+ * Set specific setting to for setting the threshold of throttling of particular task type.
+ * e.g : Set "master.throttling.thresholds.put_mapping" to set throttling limit of "put mapping" tasks,
+ *       Set it to default value(-1) to disable the throttling for this task type.
+ */
+public class MasterTaskThrottler implements TaskBatcherListener {
+    private static final Logger logger = LogManager.getLogger(MasterTaskThrottler.class);
+
+    public static final Setting<Settings> THRESHOLD_SETTINGS = Setting.groupSetting(
+        "master.throttling.thresholds.",
+        Setting.Property.Dynamic,
+        Setting.Property.NodeScope
+    );
+
+    /**
+     * To configure more task for throttling, override getMasterThrottlingKey method with task name in task executor.
+     * Verify that throttled tasks would be retry.
+     *
+     * Added retry mechanism in TransportMasterNodeAction so it would be retried for customer generated tasks.
+     */
+    public static Set<String> CONFIGURED_TASK_FOR_THROTTLING = Collections.unmodifiableSet(
+        // TODO
+        // Add throttling key for other master task as well.
+        // Will be added as part of upcoming PRs,
+        // Right now just added "put-mapping" for ref only.
+        new HashSet<>(Arrays.asList("put-mapping"))
+    );
+
+    private final int DEFAULT_THRESHOLD_VALUE = -1; // Disabled throttling
+    private final MasterService masterService;
+    private final MasterTaskThrottlerListener masterTaskThrottlerListener;
+
+    private final ConcurrentMap<String, Long> tasksCount;
+    private final ConcurrentMap<String, Long> tasksThreshold;
+
+    public MasterTaskThrottler(
+        final ClusterSettings clusterSettings,
+        final MasterService masterService,
+        final MasterTaskThrottlerListener masterTaskThrottlerListener
+    ) {
+        clusterSettings.addSettingsUpdateConsumer(THRESHOLD_SETTINGS, this::updateSetting, this::validateSetting);
+        this.masterService = masterService;
+        this.masterTaskThrottlerListener = masterTaskThrottlerListener;
+        tasksCount = new ConcurrentHashMap<>(128); // setting initial capacity so each task will lend in different segment
+        tasksThreshold = new ConcurrentHashMap<>(128); // setting initial capacity so each task will lend in different segment
+    }
+
+    public void validateSetting(final Settings settings) {
+        /**
+         * TODO: Change the version number of check as per version in which this change will be merged.
+         */
+        if (masterService.state().nodes().getMinNodeVersion().compareTo(Version.V_3_0_0) < 0) {
+            throw new IllegalArgumentException("All the nodes in cluster should be on version later than or equal to 3.0.0");
+        }
+        Map<String, Settings> groups = settings.getAsGroups();
+        for (String key : groups.keySet()) {
+            if (!CONFIGURED_TASK_FOR_THROTTLING.contains(key)) {
+                throw new IllegalArgumentException("Master task throttling is not configured for given task type: " + key);
+            }
+            int threshold = groups.get(key).getAsInt("value", DEFAULT_THRESHOLD_VALUE);
+            if (threshold < DEFAULT_THRESHOLD_VALUE) {
+                throw new IllegalArgumentException("Provide positive integer for limit or -1 for disabling throttling");
+            }
+        }
+    }
+
+    public void updateSetting(final Settings settings) {
+        Map<String, Settings> groups = settings.getAsGroups();
+        for (String key : groups.keySet()) {
+            updateLimit(key, groups.get(key).getAsInt("value", DEFAULT_THRESHOLD_VALUE));
+        }
+    }
+
+    protected void updateLimit(final String taskKey, final int limit) {
+        assert limit >= DEFAULT_THRESHOLD_VALUE;
+        if (limit == DEFAULT_THRESHOLD_VALUE) {
+            tasksThreshold.remove(taskKey);
+        } else {
+            tasksThreshold.put(taskKey, (long)limit);
+        }
+    }
+
+    @Override
+    public void beforeSubmit(List<? extends TaskBatcher.BatchedTask> tasks) {
+        String masterTaskKey = ((ClusterStateTaskExecutor<Object>) tasks.get(0).batchingKey).getMasterThrottlingKey();
+        tasksCount.putIfAbsent(masterTaskKey, 0L);
+        tasksCount.computeIfPresent(masterTaskKey, (key, count) -> {
+            long size = tasks.size();
+            Long threshold = tasksThreshold.get(masterTaskKey);
+            if(threshold != null && (count + size > threshold)) {
+                logger.warn(
+                    "Throwing Throttling Exception for [{}]. Trying to add [{}] tasks to queue, limit is set to [{}]",
+                    masterTaskKey,
+                    tasks.size(),
+                    threshold
+                );
+                throw new MasterTaskThrottlingException("Throttling Exception : Limit exceeded for " + masterTaskKey);
+            }
+            return count + size;
+        });
+    }
+
+    @Override
+    public void beforeExecute(List<? extends TaskBatcher.BatchedTask> tasks) {
+        String masterTaskKey = ((ClusterStateTaskExecutor<Object>) tasks.get(0).batchingKey).getMasterThrottlingKey();
+        tasksCount.computeIfPresent(masterTaskKey, (key, count) -> count - tasks.size());
+    }
+
+    @Override
+    public void beforeTimeout(List<? extends TaskBatcher.BatchedTask> tasks) {
+        String masterTaskKey = ((ClusterStateTaskExecutor<Object>) tasks.get(0).batchingKey).getMasterThrottlingKey();
+        tasksCount.computeIfPresent(masterTaskKey, (key, count) -> count - tasks.size());
+    }
+}

server/src/main/java/org/opensearch/cluster/service/MasterTaskThrottlerListener.java

@@ -0,0 +1,16 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.cluster.service;
+
+/**
+ * Listener interface for master task throttling
+ */
+public interface MasterTaskThrottlerListener {
+    void onThrottle(String type, final int permits);
+}

server/src/main/java/org/opensearch/cluster/service/MasterThrottlingStats.java

@@ -0,0 +1,45 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.cluster.service;
+
+import org.opensearch.common.metrics.CounterMetric;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Contains stats of Master Task Throttling.
+ * It stores the total cumulative count of throttled tasks per task type.
+ */
+public class MasterThrottlingStats implements MasterTaskThrottlerListener {
+
+    private Map<String, CounterMetric> throttledTasksCount = new ConcurrentHashMap<>();
+
+    private void incrementThrottlingCount(String type, final int permits) {
+        if (!throttledTasksCount.containsKey(type)) {
+            throttledTasksCount.computeIfAbsent(type, k -> new CounterMetric());
+        }
+        throttledTasksCount.get(type).inc(permits);
+    }
+
+    public long getThrottlingCount(String type) {
+        return throttledTasksCount.get(type).count();
+    }
+
+    public long getTotalThrottledTaskCount() {
+        CounterMetric totalCount = new CounterMetric();
+        throttledTasksCount.forEach((aClass, counterMetric) -> { totalCount.inc(counterMetric.count()); });
+        return totalCount.count();
+    }
+
+    @Override
+    public void onThrottle(String type, int permits) {
+        incrementThrottlingCount(type, permits);
+    }
+}

server/src/main/java/org/opensearch/cluster/service/TaskBatcher.java

@@ -62,10 +62,12 @@ public abstract class TaskBatcher {
     private final PrioritizedOpenSearchThreadPoolExecutor threadExecutor;
     // package visible for tests
     final Map<Object, LinkedHashSet<BatchedTask>> tasksPerBatchingKey = new HashMap<>();
+    private final TaskBatcherListener taskBatcherListener;
 
-    public TaskBatcher(Logger logger, PrioritizedOpenSearchThreadPoolExecutor threadExecutor) {
+    public TaskBatcher(Logger logger, PrioritizedOpenSearchThreadPoolExecutor threadExecutor, TaskBatcherListener taskBatcherListener) {
         this.logger = logger;
         this.threadExecutor = threadExecutor;
+        this.taskBatcherListener = taskBatcherListener;
     }
 
     public void submitTasks(List<? extends BatchedTask> tasks, @Nullable TimeValue timeout) throws OpenSearchRejectedExecutionException {
@@ -75,6 +77,11 @@ public void submitTasks(List<? extends BatchedTask> tasks, @Nullable TimeValue t
         final BatchedTask firstTask = tasks.get(0);
         assert tasks.stream().allMatch(t -> t.batchingKey == firstTask.batchingKey)
             : "tasks submitted in a batch should share the same batching key: " + tasks;
+        assert tasks.stream().allMatch(t -> t.getTask().getClass() == firstTask.getTask().getClass())
+            : "tasks submitted in a batch should be of same class: " + tasks;
+
+        taskBatcherListener.beforeSubmit(tasks);
+
         // convert to an identity map to check for dups based on task identity
         final Map<Object, BatchedTask> tasksIdentity = tasks.stream()
             .collect(
@@ -136,6 +143,7 @@ private void onTimeoutInternal(List<? extends BatchedTask> tasks, TimeValue time
                     }
                 }
             }
+            taskBatcherListener.beforeTimeout(toRemove);
             onTimeout(toRemove, timeout);
         }
     }
@@ -173,6 +181,7 @@ void runIfNotProcessed(BatchedTask updateTask) {
                     return tasks.isEmpty() ? entry.getKey() : entry.getKey() + "[" + tasks + "]";
                 }).reduce((s1, s2) -> s1 + ", " + s2).orElse("");
 
+                taskBatcherListener.beforeExecute(toExecute);
                 run(updateTask.batchingKey, toExecute, tasksSummary);
             }
         }

server/src/main/java/org/opensearch/common/settings/ClusterSettings.java

@@ -82,6 +82,7 @@
 import org.opensearch.cluster.service.ClusterApplierService;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.cluster.service.MasterService;
+import org.opensearch.cluster.service.MasterTaskThrottler;
 import org.opensearch.common.logging.Loggers;
 import org.opensearch.common.network.NetworkModule;
 import org.opensearch.common.network.NetworkService;
@@ -569,7 +570,8 @@ public void apply(Settings value, Settings current, Settings previous) {
                 ShardIndexingPressureMemoryManager.THROUGHPUT_DEGRADATION_LIMITS,
                 ShardIndexingPressureMemoryManager.SUCCESSFUL_REQUEST_ELAPSED_TIMEOUT,
                 ShardIndexingPressureMemoryManager.MAX_OUTSTANDING_REQUESTS,
-                IndexingPressure.MAX_INDEXING_BYTES
+                IndexingPressure.MAX_INDEXING_BYTES,
+                MasterTaskThrottler.THRESHOLD_SETTINGS
             )
         )
     );