Deduplicate allocation stats calls

docs/changelog/123246.yaml

@@ -0,0 +1,5 @@
+pr: 123246
+summary: Deduplicate allocation stats calls
+area: Allocation
+type: bug
+issues: []

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

@@ -13,9 +13,12 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.action.ActionRunnable;
 import org.elasticsearch.action.ActionType;
+import org.elasticsearch.action.SingleResultDeduplicator;
 import org.elasticsearch.action.admin.cluster.node.stats.NodesStatsRequestParameters.Metric;
 import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.action.support.master.MasterNodeReadRequest;
 import org.elasticsearch.action.support.master.TransportMasterNodeReadAction;
 import org.elasticsearch.cluster.ClusterState;
@@ -28,6 +31,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.util.concurrent.EsExecutors;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.injection.guice.Inject;
@@ -46,7 +50,7 @@ public class TransportGetAllocationStatsAction extends TransportMasterNodeReadAc
 
     public static final ActionType<TransportGetAllocationStatsAction.Response> TYPE = new ActionType<>("cluster:monitor/allocation/stats");
 
-    private final AllocationStatsService allocationStatsService;
+    private final SingleResultDeduplicator<Map<String, NodeAllocationStats>> allocationStatsSupplier;
     private final DiskThresholdSettings diskThresholdSettings;
 
     @Inject
@@ -66,9 +70,15 @@ public TransportGetAllocationStatsAction(
             actionFilters,
             TransportGetAllocationStatsAction.Request::new,
             TransportGetAllocationStatsAction.Response::new,
-            threadPool.executor(ThreadPool.Names.MANAGEMENT)
+            // DIRECT is ok here because we fork the allocation stats computation onto a MANAGEMENT thread if needed, or else we return
+            // very cheaply.
+            EsExecutors.DIRECT_EXECUTOR_SERVICE
+        );
+        final var managementExecutor = threadPool.executor(ThreadPool.Names.MANAGEMENT);
+        this.allocationStatsSupplier = new SingleResultDeduplicator<>(
+            threadPool.getThreadContext(),
+            l -> managementExecutor.execute(ActionRunnable.supply(l, allocationStatsService::stats))
         );
-        this.allocationStatsService = allocationStatsService;
         this.diskThresholdSettings = new DiskThresholdSettings(clusterService.getSettings(), clusterService.getClusterSettings());
     }
 
@@ -84,12 +94,15 @@ protected void doExecute(Task task, Request request, ActionListener<Response> li
 
     @Override
     protected void masterOperation(Task task, Request request, ClusterState state, ActionListener<Response> listener) throws Exception {
-        listener.onResponse(
-            new Response(
-                request.metrics().contains(Metric.ALLOCATIONS) ? allocationStatsService.stats() : Map.of(),
-                request.metrics().contains(Metric.FS) ? diskThresholdSettings : null
-            )
-        );
+        // NB we are still on a transport thread here - if adding more functionality here make sure to fork to a different pool
+
+        final SubscribableListener<Map<String, NodeAllocationStats>> allocationStatsStep = request.metrics().contains(Metric.ALLOCATIONS)
+            ? SubscribableListener.newForked(allocationStatsSupplier::execute)
+            : SubscribableListener.newSucceeded(Map.of());
+
+        allocationStatsStep.andThenApply(
+            allocationStats -> new Response(allocationStats, request.metrics().contains(Metric.FS) ? diskThresholdSettings : null)
+        ).addListener(listener);
     }
 
     @Override

server/src/main/java/org/elasticsearch/cluster/routing/allocation/AllocationStatsService.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.allocator.ShardsAllocator;
 import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.transport.Transports;
 
 import java.util.Map;
 import java.util.function.Supplier;
@@ -46,6 +47,8 @@ public AllocationStatsService(
      * Returns a map of node IDs to node allocation stats.
      */
     public Map<String, NodeAllocationStats> stats() {
+        assert Transports.assertNotTransportThread("too expensive for a transport worker");
+
         var clusterState = clusterService.state();
         var nodesStatsAndWeights = nodeAllocationStatsAndWeightsCalculator.nodesAllocationStatsAndWeights(
             clusterState.metadata(),

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

@@ -31,8 +31,13 @@
 import java.util.EnumSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 
 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.never;
@@ -112,4 +117,50 @@ public void testReturnsOnlyRequestedStats() throws Exception {
             assertNull(response.getDiskThresholdSettings());
         }
     }
+
+    public void testDeduplicatesStatsComputations() throws InterruptedException {
+        final var requestCounter = new AtomicInteger();
+        final var isExecuting = new AtomicBoolean();
+        when(allocationStatsService.stats()).thenAnswer(invocation -> {
+            try {
+                assertTrue(isExecuting.compareAndSet(false, true));
+                assertThat(Thread.currentThread().getName(), containsString("[management]"));
+                return Map.of(Integer.toString(requestCounter.incrementAndGet()), NodeAllocationStatsTests.randomNodeAllocationStats());
+            } finally {
+                Thread.yield();
+                assertTrue(isExecuting.compareAndSet(true, false));
+            }
+        });
+
+        final var threads = new Thread[between(1, 5)];
+        final var startBarrier = new CyclicBarrier(threads.length);
+        for (int i = 0; i < threads.length; i++) {
+            threads[i] = new Thread(() -> {
+                safeAwait(startBarrier);
+
+                final var minRequestIndex = requestCounter.get();
+
+                final TransportGetAllocationStatsAction.Response response = safeAwait(
+                    l -> action.masterOperation(
+                        mock(Task.class),
+                        new TransportGetAllocationStatsAction.Request(
+                            TEST_REQUEST_TIMEOUT,
+                            TaskId.EMPTY_TASK_ID,
+                            EnumSet.of(Metric.ALLOCATIONS)
+                        ),
+                        ClusterState.EMPTY_STATE,
+                        l
+                    )
+                );
+
+                final var requestIndex = Integer.valueOf(response.getNodeAllocationStats().keySet().iterator().next());
+                assertThat(requestIndex, greaterThanOrEqualTo(minRequestIndex)); // did not get a stale result
+            }, "thread-" + i);
+            threads[i].start();
+        }
+
+        for (final var thread : threads) {
+            thread.join();
+        }
+    }
 }