Fix diff building for combined cluster and project persistent tasks (MP-1965)

When the combined tasks diff is used, we need to remove the separate
project and cluster tasks from the original diff.

Relates: MP-1945, MP-1938

Author: ywangd

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

@@ -25,6 +25,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -120,15 +121,31 @@ public static <K, T, T1 extends T, T2 extends T, M extends Map<K, T>> MapDiff<K,
     }
 
     /**
-     * Create a new MapDiff by removing the key from any of its deletes, diffs and upserts
+     * Create a new MapDiff by removing the keys from any of its deletes, diffs and upserts
      */
-    public static <K, T, M extends Map<K, T>> MapDiff<K, T, M> removeKey(MapDiff<K, T, M> diff, K key) {
-        final List<K> deletes = diff.getDeletes().stream().filter(k -> k.equals(key) == false).toList();
-        final List<Map.Entry<K, Diff<T>>> diffs = diff.getDiffs().stream().filter(entry -> entry.getKey().equals(key) == false).toList();
-        final List<Map.Entry<K, T>> upserts = diff.getUpserts().stream().filter(entry -> entry.getKey().equals(key) == false).toList();
+    public static <K, T, M extends Map<K, T>> MapDiff<K, T, M> removeKeys(MapDiff<K, T, M> diff, Set<K> keys) {
+        final List<K> deletes = diff.getDeletes().stream().filter(k -> keys.contains(k) == false).toList();
+        final List<Map.Entry<K, Diff<T>>> diffs = diff.getDiffs().stream().filter(entry -> keys.contains(entry.getKey()) == false).toList();
+        final List<Map.Entry<K, T>> upserts = diff.getUpserts().stream().filter(entry -> keys.contains(entry.getKey()) == false).toList();
         return new MapDiff<>(diff.keySerializer, diff.valueSerializer, deletes, diffs, upserts, diff.builderCtor);
     }
 
+    /**
+     * Check whether the specified MapDiff has any changes associated with the specified key
+     */
+    public static <K, T, M extends Map<K, T>> boolean hasKey(MapDiff<K, T, M> diff, K key) {
+        if (diff.getDeletes().contains(key)) {
+            return true;
+        }
+        if (diff.getDiffs().stream().map(Map.Entry::getKey).anyMatch(k -> Objects.equals(k, key))) {
+            return true;
+        }
+        if (diff.getUpserts().stream().map(Map.Entry::getKey).anyMatch(k -> Objects.equals(k, key))) {
+            return true;
+        }
+        return false;
+    }
+
     /**
      * Creates a MapDiff that applies a single entry diff to a map
      */

server/src/main/java/org/elasticsearch/cluster/metadata/Metadata.java

@@ -998,12 +998,20 @@ public void writeTo(StreamOutput out) throws IOException {
             if (combinedTasksDiff == null) {
                 // No combined diff means either (1) no tasks are involved or (2) the diff is from an old node
                 // In both cases, we can proceed without further changes.
+                // For both cases, no cluster persistent tasks should exist in the merged diff
+                assert DiffableUtils.hasKey(mergedClusterAndProjectCustomDiff, ClusterPersistentTasksCustomMetadata.TYPE) == false;
+                // Unless it is (2), the merge diff should not contain project persistent tasks
+                assert fromNodeBeforeMultiProjectsSupport
+                    || DiffableUtils.hasKey(mergedClusterAndProjectCustomDiff, PersistentTasksCustomMetadata.TYPE) == false;
                 return mergedClusterAndProjectCustomDiff;
             } else {
                 // We need first delete the persistent tasks entries from the diffs by cluster and project customs themselves.
                 // Then add the combined tasks diff to the result.
                 return DiffableUtils.merge(
-                    DiffableUtils.removeKey(mergedClusterAndProjectCustomDiff, PersistentTasksCustomMetadata.TYPE),
+                    DiffableUtils.removeKeys(
+                        mergedClusterAndProjectCustomDiff,
+                        Set.of(PersistentTasksCustomMetadata.TYPE, ClusterPersistentTasksCustomMetadata.TYPE)
+                    ),
                     combinedTasksDiff,
                     DiffableUtils.getStringKeySerializer(),
                     BWC_CUSTOM_VALUE_SERIALIZER

server/src/test/java/org/elasticsearch/cluster/metadata/MetadataPersistentTasksTests.java

@@ -48,6 +48,7 @@
 public class MetadataPersistentTasksTests extends ESTestCase {
 
     private NamedWriteableRegistry namedWriteableRegistry;
+    private NamedWriteableRegistry namedWriteableRegistryBwc;
 
     @Before
     public void initializeRegistries() {
@@ -74,6 +75,25 @@ public void initializeRegistries() {
                 )
             ).toList()
         );
+        namedWriteableRegistryBwc = new NamedWriteableRegistry(
+            Stream.concat(
+                ClusterModule.getNamedWriteables()
+                    .stream()
+                    .filter(entry -> entry.name.equals(ClusterPersistentTasksCustomMetadata.TYPE) == false),
+                Stream.of(
+                    new NamedWriteableRegistry.Entry(
+                        PersistentTaskParams.class,
+                        TestClusterPersistentTasksParams.NAME,
+                        TestClusterPersistentTasksParams::new
+                    ),
+                    new NamedWriteableRegistry.Entry(
+                        PersistentTaskParams.class,
+                        TestProjectPersistentTasksParams.NAME,
+                        TestProjectPersistentTasksParams::new
+                    )
+                )
+            ).toList()
+        );
     }
 
     public void testPersistentTasksSerialization() throws IOException {
@@ -114,7 +134,7 @@ public void testPersistentTasksSerializationBwc() throws IOException {
         out.setTransportVersion(previousVersion);
         orig.writeTo(out);
 
-        final var in = new NamedWriteableAwareStreamInput(out.bytes().streamInput(), namedWriteableRegistry);
+        final var in = new NamedWriteableAwareStreamInput(out.bytes().streamInput(), namedWriteableRegistryBwc);
         in.setTransportVersion(previousVersion);
         final Metadata fromStream = Metadata.readFrom(in);
 
@@ -148,7 +168,7 @@ public void testPersistentTasksDiffSerializationBwc() throws IOException {
         out.setTransportVersion(previousVersion);
         diff.writeTo(out);
 
-        final var in = new NamedWriteableAwareStreamInput(out.bytes().streamInput(), namedWriteableRegistry);
+        final var in = new NamedWriteableAwareStreamInput(out.bytes().streamInput(), namedWriteableRegistryBwc);
         in.setTransportVersion(previousVersion);
         final Diff<Metadata> diffFromStream = Metadata.readDiffFrom(in);
         final Metadata metadataFromDiff = diffFromStream.apply(before);