[ML] checking if p-tasks metadata is null before updating state (#41091) (#41124)

* [ML] checking if p-tasks metadata is null before updating state

* Adding test that validates fix

* removing debug println

Author: benwtrent

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

@@ -246,10 +246,14 @@ public ClusterState execute(ClusterState currentState) {
                         currentState.metaData().custom(PersistentTasksCustomMetaData.TYPE), currentState.nodes());
 
                 ClusterState.Builder newState = ClusterState.builder(currentState);
-                newState.metaData(MetaData.builder(currentState.getMetaData())
-                        .putCustom(MlMetadata.TYPE, removed.mlMetadata)
-                        .putCustom(PersistentTasksCustomMetaData.TYPE, updatedTasks)
-                        .build());
+                MetaData.Builder metaDataBuilder = MetaData.builder(currentState.getMetaData())
+                    .putCustom(MlMetadata.TYPE, removed.mlMetadata);
+
+                // If there are no tasks in the cluster state metadata to begin with, this could be null.
+                if (updatedTasks != null) {
+                    metaDataBuilder = metaDataBuilder.putCustom(PersistentTasksCustomMetaData.TYPE, updatedTasks);
+                }
+                newState.metaData(metaDataBuilder.build());
                 return newState.build();
             }
 

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/integration/MlConfigMigratorIT.java

@@ -5,6 +5,7 @@
  */
 package org.elasticsearch.xpack.ml.integration;
 
+import com.carrotsearch.hppc.cursors.ObjectCursor;
 import org.elasticsearch.Version;
 import org.elasticsearch.action.DocWriteRequest;
 import org.elasticsearch.action.index.IndexRequest;
@@ -47,10 +48,12 @@
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.SortedMap;
 import java.util.TreeMap;
@@ -148,9 +151,13 @@ public void testMigrateConfigs() throws InterruptedException, IOException {
                 .routingTable(routingTable.build())
                 .build();
         when(clusterService.state()).thenReturn(clusterState);
-
+        List<MetaData.Custom> customs = new ArrayList<>();
         doAnswer(invocation -> {
                 ClusterStateUpdateTask listener = (ClusterStateUpdateTask) invocation.getArguments()[1];
+                ClusterState result = listener.execute(clusterState);
+                for (ObjectCursor<MetaData.Custom> value : result.metaData().customs().values()){
+                    customs.add(value.value);
+                }
                 listener.clusterStateProcessed("source", mock(ClusterState.class), mock(ClusterState.class));
                 return null;
         }).when(clusterService).submitStateUpdateTask(eq("remove-migrated-ml-configs"), any());
@@ -164,6 +171,9 @@ public void testMigrateConfigs() throws InterruptedException, IOException {
         blockingCall(actionListener -> mlConfigMigrator.migrateConfigs(clusterState, actionListener),
                 responseHolder, exceptionHolder);
 
+        // Verify that we have custom values in the new cluster state and that none of them is null
+        assertThat(customs.size(), greaterThan(0));
+        assertThat(customs.stream().anyMatch(Objects::isNull), is(false));
         assertNull(exceptionHolder.get());
         assertTrue(responseHolder.get());
         assertSnapshot(mlMetadata.build());