[7.x] Add support for getting a key set for ImmutableOpenMap (#77897)

client/rest-high-level/src/test/java/org/elasticsearch/client/AbstractResponseTestCase.java

@@ -19,8 +19,6 @@
 import org.elasticsearch.test.ESTestCase;
 
 import java.io.IOException;
-import java.util.HashSet;
-import java.util.Iterator;
 import java.util.Map;
 import java.util.Set;
 
@@ -80,12 +78,7 @@ protected ToXContent.Params getParams() {
     }
 
     protected static <T> void assertMapEquals(ImmutableOpenMap<String, T> expected, Map<String, T> actual) {
-        Set<String> expectedKeys = new HashSet<>();
-        Iterator<String> keysIt = expected.keysIt();
-        while (keysIt.hasNext()) {
-            expectedKeys.add(keysIt.next());
-        }
-
+        Set<String> expectedKeys = expected.keySet();
         assertEquals(expectedKeys, actual.keySet());
         for (String key : expectedKeys) {
             assertEquals(expected.get(key), actual.get(key));

server/src/main/java/org/elasticsearch/action/admin/indices/alias/get/TransportGetAliasesAction.java

@@ -34,7 +34,6 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Predicate;
@@ -135,8 +134,7 @@ private static void checkSystemIndexAccess(GetAliasesRequest request, SystemIndi
 
         List<String> netNewSystemIndices = new ArrayList<>();
         List<String> systemIndicesNames = new ArrayList<>();
-        for (Iterator<String> it = aliasesMap.keysIt(); it.hasNext(); ) {
-            String indexName = it.next();
+        aliasesMap.keySet().forEach(indexName -> {
             IndexMetadata index = state.metadata().index(indexName);
             if (index != null && index.isSystem()) {
                 if (systemIndexAccessAllowPredicate.test(index) == false) {
@@ -147,7 +145,7 @@ private static void checkSystemIndexAccess(GetAliasesRequest request, SystemIndi
                     }
                 }
             }
-        }
+        });
         if (systemIndicesNames.isEmpty() == false) {
             deprecationLogger.critical(DeprecationCategory.API, "open_system_index_access",
                 "this request accesses system indices: {}, but in a future major version, direct access to system " +

server/src/main/java/org/elasticsearch/action/admin/indices/resolve/ResolveIndexAction.java

@@ -26,16 +26,16 @@
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.service.ClusterService;
-import org.elasticsearch.core.Nullable;
-import org.elasticsearch.common.xcontent.ParseField;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.util.concurrent.CountDown;
+import org.elasticsearch.common.xcontent.ParseField;
 import org.elasticsearch.common.xcontent.ToXContentObject;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.RemoteClusterAware;
@@ -51,9 +51,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.SortedMap;
-import java.util.Spliterators;
 import java.util.TreeMap;
-import java.util.stream.StreamSupport;
 
 public class ResolveIndexAction extends ActionType<ResolveIndexAction.Response> {
 
@@ -542,10 +540,7 @@ private static void enrichIndexAbstraction(String indexAbstraction, SortedMap<St
                     case CONCRETE_INDEX:
                         IndexAbstraction.Index index = (IndexAbstraction.Index) ia;
 
-                        String[] aliasNames = StreamSupport.stream(
-                            Spliterators.spliteratorUnknownSize(index.getWriteIndex().getAliases().keysIt(), 0), false)
-                            .toArray(String[]::new);
-                        Arrays.sort(aliasNames);
+                        String[] aliasNames = index.getWriteIndex().getAliases().keySet().stream().sorted().toArray(String[]::new);
 
                         List<String> attributes = new ArrayList<>();
                         attributes.add(index.getWriteIndex().getState() == IndexMetadata.State.OPEN ? "open" : "closed");

server/src/main/java/org/elasticsearch/common/collect/ImmutableOpenMap.java

@@ -21,8 +21,10 @@
 import com.carrotsearch.hppc.procedures.ObjectObjectProcedure;
 
 import java.util.AbstractMap;
+import java.util.AbstractSet;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Set;
 import java.util.Spliterator;
 import java.util.Spliterators;
 import java.util.function.Consumer;
@@ -136,6 +138,29 @@ public void remove() {
         };
     }
 
+    /**
+     * Returns a {@link Set} view of the keys contained in this map.
+     */
+    public Set<KType> keySet() {
+        return new AbstractSet<KType>() {
+            @Override
+            public Iterator<KType> iterator() {
+                return keysIt();
+            }
+
+            @Override
+            public int size() {
+                return map.size();
+            }
+
+            @Override
+            @SuppressWarnings("unchecked")
+            public boolean contains(Object o) {
+                return map.containsKey((KType) o);
+            }
+        };
+    }
+
     /**
      * @return Returns a container with all values stored in this map.
      */

server/src/main/java/org/elasticsearch/indices/AssociatedIndexDescriptor.java

@@ -13,10 +13,10 @@
 import org.apache.lucene.util.automaton.RegExp;
 import org.elasticsearch.cluster.metadata.Metadata;
 
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
+import java.util.stream.Collectors;
 
 /**
  * An "associated index" is an index that is related to or derived from a system
@@ -107,14 +107,10 @@ static Automaton buildAutomaton(String pattern) {
      */
     @Override
     public List<String> getMatchingIndices(Metadata metadata) {
-        ArrayList<String> matchingIndices = new ArrayList<>();
-        metadata.indices().keysIt().forEachRemaining(indexName -> {
-            if (matchesIndexPattern(indexName)) {
-                matchingIndices.add(indexName);
-            }
-        });
-
-        return Collections.unmodifiableList(matchingIndices);
+        return metadata.indices().keySet()
+            .stream()
+            .filter(this::matchesIndexPattern)
+            .collect(Collectors.collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
     }
 
     /**

server/src/main/java/org/elasticsearch/indices/SystemDataStreamDescriptor.java

@@ -14,12 +14,12 @@
 import org.elasticsearch.cluster.metadata.DataStream;
 import org.elasticsearch.cluster.metadata.Metadata;
 
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.stream.Collectors;
 
 import static org.elasticsearch.indices.AssociatedIndexDescriptor.buildAutomaton;
 
@@ -94,14 +94,8 @@ public String getDataStreamName() {
      * @return List of names of backing indices
      */
     public List<String> getBackingIndexNames(Metadata metadata) {
-        ArrayList<String> matchingIndices = new ArrayList<>();
-        metadata.indices().keysIt().forEachRemaining(indexName -> {
-            if (this.characterRunAutomaton.run(indexName)) {
-                matchingIndices.add(indexName);
-            }
-        });
-
-        return Collections.unmodifiableList(matchingIndices);
+        return metadata.indices().keySet().stream().filter(this.characterRunAutomaton::run)
+            .collect(Collectors.collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
     }
 
     public String getDescription() {

server/src/main/java/org/elasticsearch/indices/SystemIndexDescriptor.java

@@ -32,6 +32,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
@@ -363,14 +364,10 @@ public boolean matchesIndexPattern(String index) {
      */
     @Override
     public List<String> getMatchingIndices(Metadata metadata) {
-        ArrayList<String> matchingIndices = new ArrayList<>();
-        metadata.indices().keysIt().forEachRemaining(indexName -> {
-            if (matchesIndexPattern(indexName)) {
-                matchingIndices.add(indexName);
-            }
-        });
-
-        return Collections.unmodifiableList(matchingIndices);
+        return metadata.indices().keySet()
+            .stream()
+            .filter(this::matchesIndexPattern)
+            .collect(Collectors.collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
     }
 
     /**

server/src/test/java/org/elasticsearch/action/admin/indices/alias/get/GetAliasesResponseTests.java

@@ -8,17 +8,15 @@
 
 package org.elasticsearch.action.admin.indices.alias.get;
 
-import org.elasticsearch.cluster.metadata.DataStreamTestHelper;
 import org.elasticsearch.cluster.metadata.AliasMetadata;
 import org.elasticsearch.cluster.metadata.AliasMetadata.Builder;
+import org.elasticsearch.cluster.metadata.DataStreamTestHelper;
 import org.elasticsearch.common.collect.ImmutableOpenMap;
-import org.elasticsearch.core.Tuple;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.core.Tuple;
 import org.elasticsearch.test.AbstractWireSerializingTestCase;
 
 import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Set;
 
@@ -52,12 +50,7 @@ private static ImmutableOpenMap<String, List<AliasMetadata>> mutateAliases(Immut
             return builder.build();
         }
 
-        Set<String> indices = new HashSet<>();
-        Iterator<String> keys = aliases.keysIt();
-        while (keys.hasNext()) {
-            indices.add(keys.next());
-        }
-
+        Set<String> indices = aliases.keySet();
         List<String> indicesToBeModified = randomSubsetOf(randomIntBetween(1, indices.size()), indices);
         ImmutableOpenMap.Builder<String, List<AliasMetadata>> builder = ImmutableOpenMap.builder();
 

server/src/test/java/org/elasticsearch/cluster/node/DiscoveryNodesTests.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.cluster.node;
 
 import com.carrotsearch.randomizedtesting.generators.RandomPicks;
+
 import org.elasticsearch.Version;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.transport.TransportAddress;
@@ -317,30 +318,22 @@ Set<String> matchingNodeIds(DiscoveryNodes nodes) {
         }, MASTER_ELIGIBLE(DiscoveryNodeRole.MASTER_ROLE.roleName() + ":true") {
             @Override
             Set<String> matchingNodeIds(DiscoveryNodes nodes) {
-                Set<String> ids = new HashSet<>();
-                nodes.getMasterNodes().keysIt().forEachRemaining(ids::add);
-                return ids;
+                return nodes.getMasterNodes().keySet();
             }
         }, DATA(DiscoveryNodeRole.DATA_ROLE.roleName() + ":true") {
             @Override
             Set<String> matchingNodeIds(DiscoveryNodes nodes) {
-                Set<String> ids = new HashSet<>();
-                nodes.getDataNodes().keysIt().forEachRemaining(ids::add);
-                return ids;
+                return nodes.getDataNodes().keySet();
             }
         }, INGEST(DiscoveryNodeRole.INGEST_ROLE.roleName() + ":true") {
             @Override
             Set<String> matchingNodeIds(DiscoveryNodes nodes) {
-                Set<String> ids = new HashSet<>();
-                nodes.getIngestNodes().keysIt().forEachRemaining(ids::add);
-                return ids;
+                return nodes.getIngestNodes().keySet();
             }
         }, COORDINATING_ONLY(DiscoveryNode.COORDINATING_ONLY + ":true") {
             @Override
             Set<String> matchingNodeIds(DiscoveryNodes nodes) {
-                Set<String> ids = new HashSet<>();
-                nodes.getCoordinatingOnlyNodes().keysIt().forEachRemaining(ids::add);
-                return ids;
+                return nodes.getCoordinatingOnlyNodes().keySet();
             }
         }, CUSTOM_ATTRIBUTE("attr:value") {
             @Override

server/src/test/java/org/elasticsearch/common/collect/ImmutableOpenMapTests.java

@@ -18,7 +18,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.SortedSet;
 import java.util.TreeMap;
+import java.util.TreeSet;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -45,11 +47,7 @@ public void testSortedStream() {
     }
 
     public void testStreamOperationsOnRandomMap() {
-        ImmutableOpenMap<Long, String> map = Randomness.get().longs(randomIntBetween(1, 1000))
-            .mapToObj(e -> Tuple.tuple(e, randomAlphaOfLength(8)))
-            .collect(() -> ImmutableOpenMap.<Long, String>builder(), (builder, t) -> builder.fPut(t.v1(), t.v2()),
-                ImmutableOpenMap.Builder::putAll)
-            .build();
+        ImmutableOpenMap<Long, String> map = randomImmutableOpenMap();
 
         int limit = randomIntBetween(0, map.size());
         Map<Long, List<String>> collectedViaStreams = map.stream()
@@ -79,4 +77,54 @@ public void testStreamOperationsOnRandomMap() {
     public void testEmptyStreamWorks() {
         assertThat(ImmutableOpenMap.of().stream().count(), equalTo(0L));
     }
+
+    public void testKeySetStreamOperationsAreSupported() {
+        assertThat(regionCurrencySymbols.keySet().stream().filter(e -> e.startsWith("U") == false).collect(Collectors.toSet()),
+            equalTo(Stream.of("Japan", "EU", "Korea").collect(Collectors.toSet())));
+    }
+
+    public void testSortedKeysSet() {
+        assertThat(regionCurrencySymbols.keySet(),
+            equalTo(Stream.of("EU", "Japan", "Korea", "UK", "USA").collect(Collectors.toSet())));
+    }
+
+    public void testStreamOperationsOnRandomMapKeys() {
+        ImmutableOpenMap<Long, String> map = randomImmutableOpenMap();
+
+        int limit = randomIntBetween(0, map.size());
+        List<Long> collectedViaStream = map.keySet()
+            .stream()
+            .filter(e -> e > 0)
+            .sorted()
+            .limit(limit)
+            .collect(Collectors.toList());
+
+        SortedSet<Long> positiveNumbers = new TreeSet<>();
+        for (ObjectObjectCursor<Long, String> cursor : map) {
+            if (cursor.key > 0) {
+                positiveNumbers.add(cursor.key);
+            }
+        }
+        int i = 0;
+        List<Long> collectedIteratively = new ArrayList<>();
+        for (Long l : positiveNumbers) {
+            if (i++ >= limit) {
+                break;
+            }
+            collectedIteratively.add(l);
+        }
+        assertThat(collectedViaStream, equalTo(collectedIteratively));
+    }
+
+    public void testEmptyKeySetWorks() {
+        assertThat(ImmutableOpenMap.of().keySet().size(), equalTo(0));
+    }
+
+    private static ImmutableOpenMap<Long, String> randomImmutableOpenMap() {
+        return Randomness.get().longs(randomIntBetween(1, 1000))
+            .mapToObj(e -> Tuple.tuple(e, randomAlphaOfLength(8)))
+            .collect(ImmutableOpenMap::<Long, String>builder, (builder, t) -> builder.fPut(t.v1(), t.v2()),
+                ImmutableOpenMap.Builder::putAll)
+            .build();
+    }
 }

x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/repository/CcrRepository.java

@@ -187,8 +187,7 @@ public void getSnapshotInfo(GetSnapshotInfoContext context) {
             .get(ccrSettings.getRecoveryActionTimeout());
         Metadata metadata = response.getState().metadata();
         ImmutableOpenMap<String, IndexMetadata> indicesMap = metadata.indices();
-        ArrayList<String> indices = new ArrayList<>(indicesMap.size());
-        indicesMap.keysIt().forEachRemaining(indices::add);
+        List<String> indices = new ArrayList<>(indicesMap.keySet());
 
         // fork to the snapshot meta pool because the context expects to run on it and asserts that it does
         threadPool.executor(ThreadPool.Names.SNAPSHOT_META).execute(() -> context.onResponse(