Add cache support in TransportGetAllocationStatsAction

Adds a new setting
TransportGetAllocationStatsAction.CACHE_MAX_AGE_SETTING to configure
the max age for cached AllocationStats on the master.  The default
value is currently 1 minute per the suggestion in issue 110716.

Closes #110716

Author: JeremyDahlgren

server/src/main/java/org/elasticsearch/action/admin/cluster/allocation/TransportGetAllocationStatsAction.java

@@ -30,6 +30,7 @@
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
@@ -42,13 +43,24 @@
 import java.io.IOException;
 import java.util.EnumSet;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 
 public class TransportGetAllocationStatsAction extends TransportMasterNodeReadAction<
     TransportGetAllocationStatsAction.Request,
     TransportGetAllocationStatsAction.Response> {
 
     public static final ActionType<TransportGetAllocationStatsAction.Response> TYPE = new ActionType<>("cluster:monitor/allocation/stats");
 
+    public static final long CACHE_DISABLED = 0L;
+    public static final Setting<TimeValue> CACHE_MAX_AGE_SETTING = Setting.timeSetting(
+        "cluster.transport.get.allocation.stats.action.cache.max_age",
+        TimeValue.timeValueMinutes(1),
+        TimeValue.timeValueMillis(CACHE_DISABLED),
+        TimeValue.timeValueMinutes(10),
+        Setting.Property.NodeScope
+    );
+
+    private final AllocationStatsCache allocationStatsCache;
     private final SingleResultDeduplicator<Map<String, NodeAllocationStats>> allocationStatsSupplier;
     private final DiskThresholdSettings diskThresholdSettings;
 
@@ -73,13 +85,29 @@ public TransportGetAllocationStatsAction(
             EsExecutors.DIRECT_EXECUTOR_SERVICE
         );
         final var managementExecutor = threadPool.executor(ThreadPool.Names.MANAGEMENT);
+        this.allocationStatsCache = new AllocationStatsCache(clusterService.getClusterSettings().get(CACHE_MAX_AGE_SETTING).millis());
         this.allocationStatsSupplier = new SingleResultDeduplicator<>(
             threadPool.getThreadContext(),
-            l -> managementExecutor.execute(ActionRunnable.supply(l, allocationStatsService::stats))
+            l -> managementExecutor.execute(ActionRunnable.supply(l, () -> {
+                final var cachedStats = allocationStatsCache.get();
+
+                if (cachedStats != null) {
+                    return cachedStats;
+                }
+
+                final var stats = allocationStatsService.stats();
+                allocationStatsCache.put(stats);
+                return stats;
+            }))
         );
         this.diskThresholdSettings = new DiskThresholdSettings(clusterService.getSettings(), clusterService.getClusterSettings());
     }
 
+    // Package access, intended for unit testing only.
+    void setCacheMaxAge(TimeValue maxAge) {
+        this.allocationStatsCache.setMaxAgeMsecs(maxAge.millis());
+    }
+
     @Override
     protected void doExecute(Task task, Request request, ActionListener<Response> listener) {
         if (clusterService.state().getMinTransportVersion().before(TransportVersions.V_8_14_0)) {
@@ -184,4 +212,42 @@ public DiskThresholdSettings getDiskThresholdSettings() {
             return diskThresholdSettings;
         }
     }
+
+    private record CachedAllocationStats(Map<String, NodeAllocationStats> stats, long timestampMsecs) {}
+
+    private static class AllocationStatsCache {
+        private volatile long maxAgeMsecs;
+        private final AtomicReference<CachedAllocationStats> cachedStats;
+
+        AllocationStatsCache(long maxAgeMsecs) {
+            this.maxAgeMsecs = maxAgeMsecs;
+            this.cachedStats = new AtomicReference<>();
+        }
+
+        void setMaxAgeMsecs(long maxAgeMsecs) {
+            this.maxAgeMsecs = maxAgeMsecs;
+        }
+
+        Map<String, NodeAllocationStats> get() {
+
+            if (maxAgeMsecs == CACHE_DISABLED) {
+                return null;
+            }
+
+            final var stats = cachedStats.get();
+
+            if (stats == null || System.currentTimeMillis() - stats.timestampMsecs > maxAgeMsecs) {
+                return null;
+            }
+
+            return stats.stats;
+        }
+
+        void put(Map<String, NodeAllocationStats> stats) {
+
+            if (maxAgeMsecs > CACHE_DISABLED) {
+                cachedStats.set(new CachedAllocationStats(stats, System.currentTimeMillis()));
+            }
+        }
+    }
 }

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

@@ -9,6 +9,7 @@
 package org.elasticsearch.common.settings;
 
 import org.apache.logging.log4j.LogManager;
+import org.elasticsearch.action.admin.cluster.allocation.TransportGetAllocationStatsAction;
 import org.elasticsearch.action.admin.cluster.configuration.TransportAddVotingConfigExclusionsAction;
 import org.elasticsearch.action.admin.indices.close.TransportCloseIndexAction;
 import org.elasticsearch.action.bulk.IncrementalBulkService;
@@ -631,6 +632,7 @@ public void apply(Settings value, Settings current, Settings previous) {
         DataStreamGlobalRetentionSettings.DATA_STREAMS_DEFAULT_RETENTION_SETTING,
         DataStreamGlobalRetentionSettings.DATA_STREAMS_MAX_RETENTION_SETTING,
         ShardsAvailabilityHealthIndicatorService.REPLICA_UNASSIGNED_BUFFER_TIME,
-        DataStream.isFailureStoreFeatureFlagEnabled() ? DataStreamFailureStoreSettings.DATA_STREAM_FAILURE_STORED_ENABLED_SETTING : null
+        DataStream.isFailureStoreFeatureFlagEnabled() ? DataStreamFailureStoreSettings.DATA_STREAM_FAILURE_STORED_ENABLED_SETTING : null,
+        TransportGetAllocationStatsAction.CACHE_MAX_AGE_SETTING
     ).filter(Objects::nonNull).collect(toSet());
 }

server/src/test/java/org/elasticsearch/action/admin/cluster/allocation/TransportGetAllocationStatsActionTests.java

@@ -9,13 +9,16 @@
 
 package org.elasticsearch.action.admin.cluster.allocation;
 
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.admin.cluster.node.stats.NodesStatsRequestParameters.Metric;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.routing.allocation.AllocationStatsService;
+import org.elasticsearch.cluster.routing.allocation.NodeAllocationStats;
 import org.elasticsearch.cluster.routing.allocation.NodeAllocationStatsTests;
 import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.core.CheckedConsumer;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskId;
@@ -35,12 +38,14 @@
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static org.hamcrest.Matchers.anEmptyMap;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.not;
 import static org.mockito.Mockito.mock;
+
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -173,4 +178,46 @@ public void testDeduplicatesStatsComputations() throws InterruptedException {
             thread.join();
         }
     }
+
+    public void testGetStatsWithCachingEnabled() throws Exception {
+
+        final AtomicReference<Map<String, NodeAllocationStats>> allocationStats = new AtomicReference<>();
+        allocationStats.set(Map.of(randomIdentifier(), NodeAllocationStatsTests.randomNodeAllocationStats()));
+        when(allocationStatsService.stats()).thenReturn(allocationStats.get());
+
+        final CheckedConsumer<ActionListener<Void>, Exception> threadTask = l -> {
+            final var request = new TransportGetAllocationStatsAction.Request(
+                TimeValue.ONE_MINUTE,
+                new TaskId(randomIdentifier(), randomNonNegativeLong()),
+                EnumSet.of(Metric.ALLOCATIONS)
+            );
+
+            final var future = new PlainActionFuture<TransportGetAllocationStatsAction.Response>();
+            action.masterOperation(mock(Task.class), request, ClusterState.EMPTY_STATE, future);
+            final var response = future.get();
+            assertSame("Expected the cached allocation stats to be returned", response.getNodeAllocationStats(), allocationStats.get());
+
+            l.onResponse(null);
+        };
+
+        // Start with a high cache max age, all threads should get the same value.
+        action.setCacheMaxAge(TimeValue.timeValueMinutes(5));
+        ESTestCase.startInParallel(between(1, 5), threadNumber -> safeAwait(threadTask));
+        verify(allocationStatsService, times(1)).stats();
+
+        // Force the cached stats to expire.
+        action.setCacheMaxAge(TimeValue.timeValueMillis(1));
+        Thread.sleep(2L);
+
+        // Expect a call to the stats service.
+        allocationStats.set(Map.of(randomIdentifier(), NodeAllocationStatsTests.randomNodeAllocationStats()));
+        when(allocationStatsService.stats()).thenReturn(allocationStats.get());
+        threadTask.accept(ActionListener.noop());
+        verify(allocationStatsService, times(2)).stats();
+
+        // Set a large max age for the cache again so all the threads in the next run can get the cached value.
+        action.setCacheMaxAge(TimeValue.timeValueMinutes(5));
+        ESTestCase.startInParallel(between(1, 5), threadNumber -> safeAwait(threadTask));
+        verify(allocationStatsService, times(2)).stats();
+    }
 }