Load ShardHeapUsageSupplier via SPI

Author: nicktindall

server/src/internalClusterTest/java/org/elasticsearch/index/shard/IndexShardIT.java

@@ -10,6 +10,7 @@
 
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.store.LockObtainFailedException;
+import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.index.IndexRequest;
@@ -816,36 +817,43 @@ private static void assertAllIndicesRemovedAndDeletionCompleted(Iterable<Indices
         }
     }
 
-    private static class BogusShardShardHeapUsageSupplier implements ShardHeapUsageSupplier {
+    public static class BogusShardShardHeapUsageSupplier implements ShardHeapUsageSupplier {
 
-        private final ClusterService clusterService;
+        private final BogusShardHeapUsagePlugin plugin;
 
-        private BogusShardShardHeapUsageSupplier(ClusterService clusterService) {
-            this.clusterService = clusterService;
+        public BogusShardShardHeapUsageSupplier(BogusShardHeapUsagePlugin plugin) {
+            this.plugin = plugin;
         }
 
         @Override
         public void getClusterHeapUsage(ActionListener<Map<String, ShardHeapUsage>> listener) {
             ActionListener.completeWith(
                 listener,
-                () -> clusterService.state().nodes().stream().collect(Collectors.toUnmodifiableMap(DiscoveryNode::getId, node -> {
-                    final long maxHeap = randomNonNegativeLong();
-                    final long usedHeap = (long) (randomFloat() * maxHeap);
-                    return new ShardHeapUsage(node.getId(), node.getName(), maxHeap, usedHeap);
-                }))
+                () -> plugin.getClusterService()
+                    .state()
+                    .nodes()
+                    .stream()
+                    .collect(Collectors.toUnmodifiableMap(DiscoveryNode::getId, node -> {
+                        final long maxHeap = randomNonNegativeLong();
+                        final long usedHeap = (long) (randomFloat() * maxHeap);
+                        return new ShardHeapUsage(node.getId(), node.getName(), maxHeap, usedHeap);
+                    }))
             );
         }
     }
 
     public static class BogusShardHeapUsagePlugin extends Plugin implements ClusterPlugin {
 
-        public BogusShardHeapUsagePlugin() {}
+        private final SetOnce<ClusterService> clusterService = new SetOnce<>();
 
         @Override
         public Collection<?> createComponents(PluginServices services) {
-            BogusShardShardHeapUsageSupplier bogusHeapUsageSupplier = new BogusShardShardHeapUsageSupplier(services.clusterService());
-            services.allocationService().setHeapUsageSupplier(bogusHeapUsageSupplier);
+            clusterService.set(services.clusterService());
             return List.of();
         }
+
+        public ClusterService getClusterService() {
+            return clusterService.get();
+        }
     }
 }

server/src/main/java/org/elasticsearch/cluster/InternalClusterInfoService.java

@@ -97,19 +97,26 @@ public class InternalClusterInfoService implements ClusterInfoService, ClusterSt
 
     private final Object mutex = new Object();
     private final List<ActionListener<ClusterInfo>> nextRefreshListeners = new ArrayList<>();
+    private final ShardHeapUsageSupplier shardHeapUsageSupplier;
 
-    private ShardHeapUsageSupplier shardHeapUsageSupplier = ShardHeapUsageSupplier.EMPTY;
     private AsyncRefresh currentRefresh;
     private RefreshScheduler refreshScheduler;
 
     @SuppressWarnings("this-escape")
-    public InternalClusterInfoService(Settings settings, ClusterService clusterService, ThreadPool threadPool, Client client) {
+    public InternalClusterInfoService(
+        Settings settings,
+        ClusterService clusterService,
+        ThreadPool threadPool,
+        Client client,
+        ShardHeapUsageSupplier shardHeapUsageSupplier
+    ) {
         this.leastAvailableSpaceUsages = Map.of();
         this.mostAvailableSpaceUsages = Map.of();
         this.shardHeapUsages = Map.of();
         this.indicesStatsSummary = IndicesStatsSummary.EMPTY;
         this.threadPool = threadPool;
         this.client = client;
+        this.shardHeapUsageSupplier = shardHeapUsageSupplier;
         this.updateFrequency = INTERNAL_CLUSTER_INFO_UPDATE_INTERVAL_SETTING.get(settings);
         this.fetchTimeout = INTERNAL_CLUSTER_INFO_TIMEOUT_SETTING.get(settings);
         this.diskThresholdEnabled = DiskThresholdSettings.CLUSTER_ROUTING_ALLOCATION_DISK_THRESHOLD_ENABLED_SETTING.get(settings);
@@ -134,16 +141,6 @@ void setUpdateFrequency(TimeValue updateFrequency) {
         this.updateFrequency = updateFrequency;
     }
 
-    /**
-     * This can be provided by plugins, which are initialised long after the ClusterInfoService is created
-     *
-     * @param shardHeapUsageSupplier The HeapUsageSupplier to use
-     */
-    public void setShardHeapUsageSupplier(ShardHeapUsageSupplier shardHeapUsageSupplier) {
-        assert this.shardHeapUsageSupplier == ShardHeapUsageSupplier.EMPTY;
-        this.shardHeapUsageSupplier = shardHeapUsageSupplier;
-    }
-
     @Override
     public void clusterChanged(ClusterChangedEvent event) {
         final Runnable newRefresh;

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

@@ -15,9 +15,7 @@
 import org.elasticsearch.cluster.ClusterChangedEvent;
 import org.elasticsearch.cluster.ClusterInfoService;
 import org.elasticsearch.cluster.ClusterState;
-import org.elasticsearch.cluster.InternalClusterInfoService;
 import org.elasticsearch.cluster.RestoreInProgress;
-import org.elasticsearch.cluster.ShardHeapUsageSupplier;
 import org.elasticsearch.cluster.health.ClusterHealthStatus;
 import org.elasticsearch.cluster.metadata.AutoExpandReplicas;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
@@ -639,12 +637,6 @@ public void addAllocFailuresResetListenerTo(ClusterService clusterService) {
         });
     }
 
-    public void setHeapUsageSupplier(ShardHeapUsageSupplier shardHeapUsageSupplier) {
-        if (clusterInfoService instanceof InternalClusterInfoService internalClusterInfoService) {
-            internalClusterInfoService.setShardHeapUsageSupplier(shardHeapUsageSupplier);
-        }
-    }
-
     /**
      *  We should reset allocation/relocation failure count to allow further retries when:
      *

server/src/main/java/org/elasticsearch/node/NodeServiceProvider.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.ClusterInfoService;
 import org.elasticsearch.cluster.InternalClusterInfoService;
+import org.elasticsearch.cluster.ShardHeapUsageSupplier;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.breaker.CircuitBreaker;
@@ -74,7 +75,17 @@ ClusterInfoService newClusterInfoService(
         ThreadPool threadPool,
         NodeClient client
     ) {
-        final InternalClusterInfoService service = new InternalClusterInfoService(settings, clusterService, threadPool, client);
+        final ShardHeapUsageSupplier shardHeapUsageSupplier = pluginsService.loadSingletonServiceProvider(
+            ShardHeapUsageSupplier.class,
+            () -> ShardHeapUsageSupplier.EMPTY
+        );
+        final InternalClusterInfoService service = new InternalClusterInfoService(
+            settings,
+            clusterService,
+            threadPool,
+            client,
+            shardHeapUsageSupplier
+        );
         if (DiscoveryNode.isMasterNode(settings)) {
             // listen for state changes (this node starts/stops being the elected master, or new nodes are added)
             clusterService.addListener(service);

server/src/test/java/org/elasticsearch/cluster/InternalClusterInfoServiceSchedulingTests.java

@@ -76,9 +76,14 @@ protected PrioritizedEsThreadPoolExecutor createThreadPoolExecutor() {
         final ClusterService clusterService = new ClusterService(settings, clusterSettings, masterService, clusterApplierService);
 
         final FakeClusterInfoServiceClient client = new FakeClusterInfoServiceClient(threadPool);
-        final InternalClusterInfoService clusterInfoService = new InternalClusterInfoService(settings, clusterService, threadPool, client);
         final ShardHeapUsageSupplier mockShardHeapUsageSupplier = spy(new StubShardShardHeapUsageSupplier());
-        clusterInfoService.setShardHeapUsageSupplier(mockShardHeapUsageSupplier);
+        final InternalClusterInfoService clusterInfoService = new InternalClusterInfoService(
+            settings,
+            clusterService,
+            threadPool,
+            client,
+            mockShardHeapUsageSupplier
+        );
         clusterService.addListener(clusterInfoService);
         clusterInfoService.addListener(ignored -> {});
 

server/src/test/resources/META-INF/services/org.elasticsearch.cluster.ShardHeapUsageSupplier

@@ -0,0 +1,10 @@
+#
+# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+# or more contributor license agreements. Licensed under the "Elastic License
+# 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+# Public License v 1"; you may not use this file except in compliance with, at
+# your election, the "Elastic License 2.0", the "GNU Affero General Public
+# License v3.0 only", or the "Server Side Public License, v 1".
+#
+
+org.elasticsearch.index.shard.IndexShardIT$BogusShardShardHeapUsageSupplier

test/framework/src/main/java/org/elasticsearch/cluster/MockInternalClusterInfoService.java

@@ -43,7 +43,7 @@ public static class TestPlugin extends Plugin {}
     private volatile BiFunction<DiscoveryNode, FsInfo.Path, FsInfo.Path> diskUsageFunction;
 
     public MockInternalClusterInfoService(Settings settings, ClusterService clusterService, ThreadPool threadPool, NodeClient client) {
-        super(settings, clusterService, threadPool, client);
+        super(settings, clusterService, threadPool, client, ShardHeapUsageSupplier.EMPTY);
     }
 
     public void setDiskUsageFunctionAndRefresh(BiFunction<DiscoveryNode, FsInfo.Path, FsInfo.Path> diskUsageFn) {