[ML] Rebalance should not notify listener before update is applied (#88012)

When we are rebalancing trained model assignments, we eventually
update cluster state with the new metadata. We only want to
notify the listener after we have applied the cluster state update.
This commit fixes a bug where we could notify before the update
was actually applied resulting in NPE when writing the response
of the `CreateTrainedModelAssignmentAction`.

This is a fix following the work done in #87366.

Author: dimitris-athanasiou

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

@@ -450,6 +450,7 @@ public boolean test(ClusterState clusterState) {
                 .orElse(null);
             if (trainedModelAssignment == null) {
                 // Something weird happened, it should NEVER be null...
+                logger.trace(() -> format("[%s] assignment was null while waiting for state [%s]", modelId, waitForState));
                 return true;
             }
 

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentClusterService.java

@@ -358,6 +358,7 @@ private static ClusterState update(ClusterState currentState, TrainedModelAssign
     }
 
     private static ClusterState forceUpdate(ClusterState currentState, TrainedModelAssignmentMetadata.Builder modelAssignments) {
+        logger.debug(() -> format("updated assignments: %s", modelAssignments.build()));
         Metadata.Builder metadata = Metadata.builder(currentState.metadata());
         if (currentState.getNodes().getMinNodeVersion().onOrAfter(RENAME_ALLOCATION_TO_ASSIGNMENT_VERSION)) {
             metadata.putCustom(TrainedModelAssignmentMetadata.NAME, modelAssignments.build())
@@ -389,10 +390,14 @@ private void rebalanceAssignments(
             }
 
             submitUnbatchedTask(reason, new ClusterStateUpdateTask() {
+
+                private volatile boolean isUpdated;
+
                 @Override
                 public ClusterState execute(ClusterState currentState) {
 
                     if (areClusterStatesCompatibleForRebalance(clusterState, currentState)) {
+                        isUpdated = true;
                         return update(currentState, rebalancedMetadata);
                     }
                     rebalanceAssignments(currentState, modelToAdd, reason, listener);
@@ -406,7 +411,9 @@ public void onFailure(Exception e) {
 
                 @Override
                 public void clusterStateProcessed(ClusterState oldState, ClusterState newState) {
-                    listener.onResponse(TrainedModelAssignmentMetadata.fromState(newState));
+                    if (isUpdated) {
+                        listener.onResponse(TrainedModelAssignmentMetadata.fromState(newState));
+                    }
                 }
             });
         });

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentMetadata.java

@@ -16,6 +16,7 @@
 import org.elasticsearch.cluster.NamedDiff;
 import org.elasticsearch.cluster.SimpleDiffable;
 import org.elasticsearch.cluster.metadata.Metadata;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -149,6 +150,11 @@ public int hashCode() {
         return Objects.hash(modelRoutingEntries);
     }
 
+    @Override
+    public String toString() {
+        return Strings.toString(this);
+    }
+
     public static class Builder {
 
         public static Builder empty() {