[Backport 2.x] Added resource usage trackers for in-flight cancellati…

…on of SearchShardTask (opensearch-project#4805)

1. CpuUsageTracker: cancels tasks if they consume too much CPU
2. ElapsedTimeTracker: cancels tasks if they consume too much time
3. HeapUsageTracker: cancels tasks if they consume too much heap

Signed-off-by: Ketan Verma <ketan9495@gmail.com>

CHANGELOG.md

@@ -26,6 +26,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Update GeoGrid base class access modifier to support extensibility ([#4921](https://github.com/opensearch-project/OpenSearch/pull/4921))
 - Build no-jdk distributions as part of release build ([#4902](https://github.com/opensearch-project/OpenSearch/pull/4902))
 - Added in-flight cancellation of SearchShardTask based on resource consumption ([#4565](https://github.com/opensearch-project/OpenSearch/pull/4565))
+- Added resource usage trackers for in-flight cancellation of SearchShardTask ([#4805](https://github.com/opensearch-project/OpenSearch/pull/4805))
 
 ### Dependencies
 - Bumps `com.diffplug.spotless` from 6.9.1 to 6.10.0

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

@@ -45,6 +45,9 @@
 import org.opensearch.search.backpressure.settings.NodeDuressSettings;
 import org.opensearch.search.backpressure.settings.SearchBackpressureSettings;
 import org.opensearch.search.backpressure.settings.SearchShardTaskSettings;
+import org.opensearch.search.backpressure.trackers.CpuUsageTracker;
+import org.opensearch.search.backpressure.trackers.ElapsedTimeTracker;
+import org.opensearch.search.backpressure.trackers.HeapUsageTracker;
 import org.opensearch.tasks.TaskManager;
 import org.opensearch.tasks.TaskResourceTrackingService;
 import org.opensearch.watcher.ResourceWatcherService;
@@ -596,11 +599,12 @@ public void apply(Settings value, Settings current, Settings previous) {
                 NodeDuressSettings.SETTING_NUM_SUCCESSIVE_BREACHES,
                 NodeDuressSettings.SETTING_CPU_THRESHOLD,
                 NodeDuressSettings.SETTING_HEAP_THRESHOLD,
-                SearchShardTaskSettings.SETTING_TOTAL_HEAP_THRESHOLD,
-                SearchShardTaskSettings.SETTING_HEAP_THRESHOLD,
-                SearchShardTaskSettings.SETTING_HEAP_VARIANCE_THRESHOLD,
-                SearchShardTaskSettings.SETTING_CPU_TIME_THRESHOLD,
-                SearchShardTaskSettings.SETTING_ELAPSED_TIME_THRESHOLD
+                SearchShardTaskSettings.SETTING_TOTAL_HEAP_PERCENT_THRESHOLD,
+                HeapUsageTracker.SETTING_HEAP_PERCENT_THRESHOLD,
+                HeapUsageTracker.SETTING_HEAP_VARIANCE_THRESHOLD,
+                HeapUsageTracker.SETTING_HEAP_MOVING_AVERAGE_WINDOW_SIZE,
+                CpuUsageTracker.SETTING_CPU_TIME_MILLIS_THRESHOLD,
+                ElapsedTimeTracker.SETTING_ELAPSED_TIME_MILLIS_THRESHOLD
             )
         )
     );

server/src/main/java/org/opensearch/common/util/TokenBucket.java

@@ -101,9 +101,9 @@ public boolean request() {
      */
     private static class State {
         final double tokens;
-        final double lastRefilledAt;
+        final long lastRefilledAt;
 
-        public State(double tokens, double lastRefilledAt) {
+        public State(double tokens, long lastRefilledAt) {
             this.tokens = tokens;
             this.lastRefilledAt = lastRefilledAt;
         }
@@ -113,7 +113,7 @@ public boolean equals(Object o) {
             if (this == o) return true;
             if (o == null || getClass() != o.getClass()) return false;
             State state = (State) o;
-            return Double.compare(state.tokens, tokens) == 0 && Double.compare(state.lastRefilledAt, lastRefilledAt) == 0;
+            return Double.compare(state.tokens, tokens) == 0 && lastRefilledAt == state.lastRefilledAt;
         }
 
         @Override

server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java

@@ -17,6 +17,9 @@
 import org.opensearch.monitor.jvm.JvmStats;
 import org.opensearch.monitor.process.ProcessProbe;
 import org.opensearch.search.backpressure.settings.SearchBackpressureSettings;
+import org.opensearch.search.backpressure.trackers.CpuUsageTracker;
+import org.opensearch.search.backpressure.trackers.ElapsedTimeTracker;
+import org.opensearch.search.backpressure.trackers.HeapUsageTracker;
 import org.opensearch.search.backpressure.trackers.NodeDuressTracker;
 import org.opensearch.search.backpressure.trackers.TaskResourceUsageTracker;
 import org.opensearch.tasks.CancellableTask;
@@ -29,7 +32,6 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Optional;
@@ -84,7 +86,7 @@ public SearchBackpressureService(
                     () -> JvmStats.jvmStats().getMem().getHeapUsedPercent() / 100.0 >= settings.getNodeDuressSettings().getHeapThreshold()
                 )
             ),
-            Collections.emptyList()
+            List.of(new CpuUsageTracker(settings), new HeapUsageTracker(settings), new ElapsedTimeTracker(settings, System::nanoTime))
         );
     }
 
@@ -97,7 +99,7 @@ public SearchBackpressureService(
         List<TaskResourceUsageTracker> taskResourceUsageTrackers
     ) {
         this.settings = settings;
-        this.settings.setListener(this);
+        this.settings.addListener(this);
         this.taskResourceTrackingService = taskResourceTrackingService;
         this.taskResourceTrackingService.addTaskCompletionListener(this);
         this.threadPool = threadPool;
@@ -181,10 +183,10 @@ boolean isNodeInDuress() {
      * Returns true if the increase in heap usage is due to search requests.
      */
     boolean isHeapUsageDominatedBySearch(List<CancellableTask> searchShardTasks) {
-        long runningTasksHeapUsage = searchShardTasks.stream().mapToLong(task -> task.getTotalResourceStats().getMemoryInBytes()).sum();
-        long searchTasksHeapThreshold = getSettings().getSearchShardTaskSettings().getTotalHeapThresholdBytes();
-        if (runningTasksHeapUsage < searchTasksHeapThreshold) {
-            logger.debug("heap usage not dominated by search requests [{}/{}]", runningTasksHeapUsage, searchTasksHeapThreshold);
+        long usage = searchShardTasks.stream().mapToLong(task -> task.getTotalResourceStats().getMemoryInBytes()).sum();
+        long threshold = getSettings().getSearchShardTaskSettings().getTotalHeapBytesThreshold();
+        if (usage < threshold) {
+            logger.debug("heap usage not dominated by search requests [{}/{}]", usage, threshold);
             return false;
         }
 
@@ -213,7 +215,7 @@ TaskCancellation getTaskCancellation(CancellableTask task) {
         List<Runnable> callbacks = new ArrayList<>();
 
         for (TaskResourceUsageTracker tracker : taskResourceUsageTrackers) {
-            Optional<TaskCancellation.Reason> reason = tracker.cancellationReason(task);
+            Optional<TaskCancellation.Reason> reason = tracker.checkAndMaybeGetCancellationReason(task);
             if (reason.isPresent()) {
                 reasons.add(reason.get());
                 callbacks.add(tracker::incrementCancellations);

server/src/main/java/org/opensearch/search/backpressure/settings/SearchBackpressureSettings.java

@@ -8,13 +8,16 @@
 
 package org.opensearch.search.backpressure.settings;
 
-import org.apache.lucene.util.SetOnce;
+import org.opensearch.ExceptionsHelper;
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 /**
  * Settings related to search backpressure and cancellation of in-flight requests.
@@ -23,7 +26,7 @@
  */
 public class SearchBackpressureSettings {
     private static class Defaults {
-        private static final long INTERVAL = 1000;
+        private static final long INTERVAL_MILLIS = 1000;
 
         private static final boolean ENABLED = true;
         private static final boolean ENFORCED = false;
@@ -37,9 +40,9 @@ private static class Defaults {
      * Defines the interval (in millis) at which the SearchBackpressureService monitors and cancels tasks.
      */
     private final TimeValue interval;
-    public static final Setting<Long> SETTING_INTERVAL = Setting.longSetting(
-        "search_backpressure.interval",
-        Defaults.INTERVAL,
+    public static final Setting<Long> SETTING_INTERVAL_MILLIS = Setting.longSetting(
+        "search_backpressure.interval_millis",
+        Defaults.INTERVAL_MILLIS,
         1,
         Setting.Property.NodeScope
     );
@@ -116,15 +119,19 @@ public interface Listener {
         void onCancellationBurstChanged();
     }
 
-    private final SetOnce<Listener> listener = new SetOnce<>();
+    private final List<Listener> listeners = new ArrayList<>();
+    private final Settings settings;
+    private final ClusterSettings clusterSettings;
     private final NodeDuressSettings nodeDuressSettings;
     private final SearchShardTaskSettings searchShardTaskSettings;
 
     public SearchBackpressureSettings(Settings settings, ClusterSettings clusterSettings) {
+        this.settings = settings;
+        this.clusterSettings = clusterSettings;
         this.nodeDuressSettings = new NodeDuressSettings(settings, clusterSettings);
         this.searchShardTaskSettings = new SearchShardTaskSettings(settings, clusterSettings);
 
-        interval = new TimeValue(SETTING_INTERVAL.get(settings));
+        interval = new TimeValue(SETTING_INTERVAL_MILLIS.get(settings));
 
         enabled = SETTING_ENABLED.get(settings);
         clusterSettings.addSettingsUpdateConsumer(SETTING_ENABLED, this::setEnabled);
@@ -142,8 +149,16 @@ public SearchBackpressureSettings(Settings settings, ClusterSettings clusterSett
         clusterSettings.addSettingsUpdateConsumer(SETTING_CANCELLATION_BURST, this::setCancellationBurst);
     }
 
-    public void setListener(Listener listener) {
-        this.listener.set(listener);
+    public void addListener(Listener listener) {
+        listeners.add(listener);
+    }
+
+    public Settings getSettings() {
+        return settings;
+    }
+
+    public ClusterSettings getClusterSettings() {
+        return clusterSettings;
     }
 
     public NodeDuressSettings getNodeDuressSettings() {
@@ -180,9 +195,7 @@ public double getCancellationRatio() {
 
     private void setCancellationRatio(double cancellationRatio) {
         this.cancellationRatio = cancellationRatio;
-        if (listener.get() != null) {
-            listener.get().onCancellationRatioChanged();
-        }
+        notifyListeners(Listener::onCancellationRatioChanged);
     }
 
     public double getCancellationRate() {
@@ -195,9 +208,7 @@ public double getCancellationRateNanos() {
 
     private void setCancellationRate(double cancellationRate) {
         this.cancellationRate = cancellationRate;
-        if (listener.get() != null) {
-            listener.get().onCancellationRateChanged();
-        }
+        notifyListeners(Listener::onCancellationRateChanged);
     }
 
     public double getCancellationBurst() {
@@ -206,8 +217,20 @@ public double getCancellationBurst() {
 
     private void setCancellationBurst(double cancellationBurst) {
         this.cancellationBurst = cancellationBurst;
-        if (listener.get() != null) {
-            listener.get().onCancellationBurstChanged();
+        notifyListeners(Listener::onCancellationBurstChanged);
+    }
+
+    private void notifyListeners(Consumer<Listener> consumer) {
+        List<Exception> exceptions = new ArrayList<>();
+
+        for (Listener listener : listeners) {
+            try {
+                consumer.accept(listener);
+            } catch (Exception e) {
+                exceptions.add(e);
+            }
         }
+
+        ExceptionsHelper.maybeThrowRuntimeAndSuppress(exceptions);
     }
 }

server/src/main/java/org/opensearch/search/backpressure/settings/SearchShardTaskSettings.java

@@ -22,139 +22,37 @@ public class SearchShardTaskSettings {
     private static final long HEAP_SIZE_BYTES = JvmStats.jvmStats().getMem().getHeapMax().getBytes();
 
     private static class Defaults {
-        private static final double TOTAL_HEAP_THRESHOLD = 0.05;
-        private static final double HEAP_THRESHOLD = 0.005;
-        private static final double HEAP_VARIANCE_THRESHOLD = 2.0;
-        private static final long CPU_TIME_THRESHOLD = 15;
-        private static final long ELAPSED_TIME_THRESHOLD = 30000;
+        private static final double TOTAL_HEAP_PERCENT_THRESHOLD = 0.05;
     }
 
     /**
      * Defines the heap usage threshold (in percentage) for the sum of heap usages across all search shard tasks
      * before in-flight cancellation is applied.
      */
-    private volatile double totalHeapThreshold;
-    public static final Setting<Double> SETTING_TOTAL_HEAP_THRESHOLD = Setting.doubleSetting(
-        "search_backpressure.search_shard_task.total_heap_threshold",
-        Defaults.TOTAL_HEAP_THRESHOLD,
+    private volatile double totalHeapPercentThreshold;
+    public static final Setting<Double> SETTING_TOTAL_HEAP_PERCENT_THRESHOLD = Setting.doubleSetting(
+        "search_backpressure.search_shard_task.total_heap_percent_threshold",
+        Defaults.TOTAL_HEAP_PERCENT_THRESHOLD,
         0.0,
         1.0,
         Setting.Property.Dynamic,
         Setting.Property.NodeScope
     );
 
-    /**
-     * Defines the heap usage threshold (in percentage) for an individual task before it is considered for cancellation.
-     */
-    private volatile double heapThreshold;
-    public static final Setting<Double> SETTING_HEAP_THRESHOLD = Setting.doubleSetting(
-        "search_backpressure.search_shard_task.heap_threshold",
-        Defaults.HEAP_THRESHOLD,
-        0.0,
-        1.0,
-        Setting.Property.Dynamic,
-        Setting.Property.NodeScope
-    );
-
-    /**
-     * Defines the heap usage variance for an individual task before it is considered for cancellation.
-     * A task is considered for cancellation when taskHeapUsage is greater than or equal to heapUsageMovingAverage * variance.
-     */
-    private volatile double heapVarianceThreshold;
-    public static final Setting<Double> SETTING_HEAP_VARIANCE_THRESHOLD = Setting.doubleSetting(
-        "search_backpressure.search_shard_task.heap_variance",
-        Defaults.HEAP_VARIANCE_THRESHOLD,
-        0.0,
-        Setting.Property.Dynamic,
-        Setting.Property.NodeScope
-    );
-
-    /**
-     * Defines the CPU usage threshold (in millis) for an individual task before it is considered for cancellation.
-     */
-    private volatile long cpuTimeThreshold;
-    public static final Setting<Long> SETTING_CPU_TIME_THRESHOLD = Setting.longSetting(
-        "search_backpressure.search_shard_task.cpu_time_threshold",
-        Defaults.CPU_TIME_THRESHOLD,
-        0,
-        Setting.Property.Dynamic,
-        Setting.Property.NodeScope
-    );
-
-    /**
-     * Defines the elapsed time threshold (in millis) for an individual task before it is considered for cancellation.
-     */
-    private volatile long elapsedTimeThreshold;
-    public static final Setting<Long> SETTING_ELAPSED_TIME_THRESHOLD = Setting.longSetting(
-        "search_backpressure.search_shard_task.elapsed_time_threshold",
-        Defaults.ELAPSED_TIME_THRESHOLD,
-        0,
-        Setting.Property.Dynamic,
-        Setting.Property.NodeScope
-    );
-
     public SearchShardTaskSettings(Settings settings, ClusterSettings clusterSettings) {
-        totalHeapThreshold = SETTING_TOTAL_HEAP_THRESHOLD.get(settings);
-        clusterSettings.addSettingsUpdateConsumer(SETTING_TOTAL_HEAP_THRESHOLD, this::setTotalHeapThreshold);
-
-        heapThreshold = SETTING_HEAP_THRESHOLD.get(settings);
-        clusterSettings.addSettingsUpdateConsumer(SETTING_HEAP_THRESHOLD, this::setHeapThreshold);
-
-        heapVarianceThreshold = SETTING_HEAP_VARIANCE_THRESHOLD.get(settings);
-        clusterSettings.addSettingsUpdateConsumer(SETTING_HEAP_VARIANCE_THRESHOLD, this::setHeapVarianceThreshold);
-
-        cpuTimeThreshold = SETTING_CPU_TIME_THRESHOLD.get(settings);
-        clusterSettings.addSettingsUpdateConsumer(SETTING_CPU_TIME_THRESHOLD, this::setCpuTimeThreshold);
-
-        elapsedTimeThreshold = SETTING_ELAPSED_TIME_THRESHOLD.get(settings);
-        clusterSettings.addSettingsUpdateConsumer(SETTING_ELAPSED_TIME_THRESHOLD, this::setElapsedTimeThreshold);
-    }
-
-    public double getTotalHeapThreshold() {
-        return totalHeapThreshold;
-    }
-
-    public long getTotalHeapThresholdBytes() {
-        return (long) (HEAP_SIZE_BYTES * getTotalHeapThreshold());
-    }
-
-    private void setTotalHeapThreshold(double totalHeapThreshold) {
-        this.totalHeapThreshold = totalHeapThreshold;
-    }
-
-    public double getHeapThreshold() {
-        return heapThreshold;
-    }
-
-    public long getHeapThresholdBytes() {
-        return (long) (HEAP_SIZE_BYTES * getHeapThreshold());
-    }
-
-    private void setHeapThreshold(double heapThreshold) {
-        this.heapThreshold = heapThreshold;
-    }
-
-    public double getHeapVarianceThreshold() {
-        return heapVarianceThreshold;
-    }
-
-    private void setHeapVarianceThreshold(double heapVarianceThreshold) {
-        this.heapVarianceThreshold = heapVarianceThreshold;
-    }
-
-    public long getCpuTimeThreshold() {
-        return cpuTimeThreshold;
+        totalHeapPercentThreshold = SETTING_TOTAL_HEAP_PERCENT_THRESHOLD.get(settings);
+        clusterSettings.addSettingsUpdateConsumer(SETTING_TOTAL_HEAP_PERCENT_THRESHOLD, this::setTotalHeapPercentThreshold);
     }
 
-    private void setCpuTimeThreshold(long cpuTimeThreshold) {
-        this.cpuTimeThreshold = cpuTimeThreshold;
+    public double getTotalHeapPercentThreshold() {
+        return totalHeapPercentThreshold;
     }
 
-    public long getElapsedTimeThreshold() {
-        return elapsedTimeThreshold;
+    public long getTotalHeapBytesThreshold() {
+        return (long) (HEAP_SIZE_BYTES * getTotalHeapPercentThreshold());
     }
 
-    private void setElapsedTimeThreshold(long elapsedTimeThreshold) {
-        this.elapsedTimeThreshold = elapsedTimeThreshold;
+    private void setTotalHeapPercentThreshold(double totalHeapPercentThreshold) {
+        this.totalHeapPercentThreshold = totalHeapPercentThreshold;
     }
 }