[ML] Release cluster state (#136769) (#136835)

Refactoring TransportMlMemoryAction to not retain cluster state through
the lifecycle of the request so that it can be garbage collected as soon
as possible.

Fix #123243

Author: prwhelan

docs/changelog/136769.yaml

@@ -0,0 +1,6 @@
+pr: 136769
+summary: Release cluster state
+area: Machine Learning
+type: bug
+issues:
+ - 123243

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportMlMemoryAction.java

@@ -16,6 +16,7 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.client.internal.OriginSettingClient;
 import org.elasticsearch.client.internal.ParentTaskAssigningClient;
+import org.elasticsearch.cluster.ClusterName;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.block.ClusterBlockException;
 import org.elasticsearch.cluster.block.ClusterBlockLevel;
@@ -35,6 +36,7 @@
 import org.elasticsearch.xpack.core.ml.action.MlMemoryAction.Response.MlMemoryStats;
 import org.elasticsearch.xpack.core.ml.action.TrainedModelCacheInfoAction;
 import org.elasticsearch.xpack.core.ml.action.TrainedModelCacheInfoAction.Response.CacheInfo;
+import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignmentMetadata;
 import org.elasticsearch.xpack.ml.job.NodeLoad;
 import org.elasticsearch.xpack.ml.job.NodeLoadDetector;
 import org.elasticsearch.xpack.ml.process.MlMemoryTracker;
@@ -88,6 +90,9 @@ protected void masterOperation(
 
         ClusterSettings clusterSettings = clusterService.getClusterSettings();
 
+        var clusterName = state.getClusterName();
+        var trainedModelAssignmentMetadata = TrainedModelAssignmentMetadata.fromState(state);
+        PersistentTasksCustomMetadata persistentTasksCustomMetadata = state.getMetadata().custom(PersistentTasksCustomMetadata.TYPE);
         // Resolve the node specification to some concrete nodes
         String[] nodeIds = state.nodes().resolveNodes(request.getNodeId());
 
@@ -116,7 +121,9 @@ protected void masterOperation(
                         trainedModelCacheInfoRequest,
                         delegate2.delegateFailureAndWrap(
                             (l, trainedModelCacheInfoResponse) -> handleResponses(
-                                state,
+                                clusterName,
+                                persistentTasksCustomMetadata,
+                                trainedModelAssignmentMetadata,
                                 clusterSettings,
                                 nodesStatsResponse,
                                 trainedModelCacheInfoResponse,
@@ -131,12 +138,14 @@ protected void masterOperation(
         if (memoryTracker.isEverRefreshed()) {
             memoryTrackerRefreshListener.onResponse(null);
         } else {
-            memoryTracker.refresh(state.getMetadata().custom(PersistentTasksCustomMetadata.TYPE), memoryTrackerRefreshListener);
+            memoryTracker.refresh(persistentTasksCustomMetadata, memoryTrackerRefreshListener);
         }
     }
 
     void handleResponses(
-        ClusterState state,
+        ClusterName clusterName,
+        PersistentTasksCustomMetadata persistentTasks,
+        TrainedModelAssignmentMetadata assignmentMetadata,
         ClusterSettings clusterSettings,
         NodesStatsResponse nodesStatsResponse,
         TrainedModelCacheInfoAction.Response trainedModelCacheInfoResponse,
@@ -174,7 +183,8 @@ void handleResponses(
             ByteSizeValue mlNativeInference;
             if (node.getRoles().contains(DiscoveryNodeRole.ML_ROLE)) {
                 NodeLoad nodeLoad = nodeLoadDetector.detectNodeLoad(
-                    state,
+                    persistentTasks,
+                    assignmentMetadata,
                     node,
                     maxOpenJobsPerNode,
                     maxMachineMemoryPercent,
@@ -220,7 +230,7 @@ void handleResponses(
             );
         }
 
-        listener.onResponse(new MlMemoryAction.Response(state.getClusterName(), nodeResponses, failures));
+        listener.onResponse(new MlMemoryAction.Response(clusterName, nodeResponses, failures));
     }
 
     @Override

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/NodeLoadDetector.java

@@ -87,7 +87,24 @@ public NodeLoad detectNodeLoad(
         int maxMachineMemoryPercent,
         boolean useAutoMachineMemoryCalculation
     ) {
-        PersistentTasksCustomMetadata persistentTasks = clusterState.getMetadata().custom(PersistentTasksCustomMetadata.TYPE);
+        return detectNodeLoad(
+            clusterState.getMetadata().custom(PersistentTasksCustomMetadata.TYPE),
+            assignmentMetadata,
+            node,
+            maxNumberOfOpenJobs,
+            maxMachineMemoryPercent,
+            useAutoMachineMemoryCalculation
+        );
+    }
+
+    public NodeLoad detectNodeLoad(
+        PersistentTasksCustomMetadata persistentTasks,
+        TrainedModelAssignmentMetadata assignmentMetadata,
+        DiscoveryNode node,
+        int maxNumberOfOpenJobs,
+        int maxMachineMemoryPercent,
+        boolean useAutoMachineMemoryCalculation
+    ) {
         Map<String, String> nodeAttributes = node.getAttributes();
         List<String> errors = new ArrayList<>();
         OptionalLong maxMlMemory = NativeMemoryCalculator.allowedBytesForMl(node, maxMachineMemoryPercent, useAutoMachineMemoryCalculation);