Cleanup / move some serverless code to ES (#112360)

Author: mosche

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

@@ -10,9 +10,11 @@
 
 import org.elasticsearch.core.Nullable;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Iterator;
+import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.function.BiFunction;
@@ -233,6 +235,57 @@ public T next() {
         }
     }
 
+    /**
+     * Returns an iterator that yields at most the first {@code n} elements of the provided {@code input} iterator.
+     */
+    public static <T> Iterator<T> limit(Iterator<? extends T> input, int n) {
+        assert n >= 0 : "negative limit";
+        if (n > 0 && input.hasNext()) {
+            return new LimitIterator<>(input, n);
+        } else {
+            return Collections.emptyIterator();
+        }
+    }
+
+    private static final class LimitIterator<T> implements Iterator<T> {
+        private final Iterator<? extends T> input;
+        private final int limit;
+        private int current;
+
+        LimitIterator(Iterator<? extends T> input, int limit) {
+            this.input = input;
+            this.limit = limit;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return current < limit && input.hasNext();
+        }
+
+        @Override
+        public T next() {
+            if (current >= limit) {
+                throw new NoSuchElementException();
+            }
+            ++current;
+            return input.next();
+        }
+    }
+
+    /**
+     * Returns a list containing the elements of the provided {@code iterator}.
+     */
+    public static <T> List<T> toList(Iterator<T> iterator) {
+        if (iterator.hasNext()) {
+            var list = new ArrayList<T>();
+            while (iterator.hasNext()) {
+                list.add(iterator.next());
+            }
+            return Collections.unmodifiableList(list);
+        }
+        return Collections.emptyList();
+    }
+
     public static <T, U> Iterator<U> flatMap(Iterator<? extends T> input, Function<T, Iterator<? extends U>> fn) {
         while (input.hasNext()) {
             final var value = fn.apply(input.next());

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

@@ -28,6 +28,10 @@
 import java.util.function.ToIntFunction;
 import java.util.stream.IntStream;
 
+import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.is;
+
 public class IteratorsTests extends ESTestCase {
     public void testConcatentation() {
         List<Integer> threeTwoOne = Arrays.asList(3, 2, 1);
@@ -242,6 +246,25 @@ public void testFilter() {
         }
     }
 
+    public void testLimit() {
+        var result = Iterators.limit(Collections.emptyIterator(), 10);
+        assertThat(result.hasNext(), is(false));
+        assertThat(Iterators.toList(result), is(empty()));
+
+        var values = List.of(1, 2, 3);
+        result = Iterators.limit(values.iterator(), 10);
+        assertThat(result.hasNext(), is(true));
+        assertThat(Iterators.toList(result), contains(1, 2, 3));
+
+        result = Iterators.limit(values.iterator(), 2);
+        assertThat(result.hasNext(), is(true));
+        assertThat(Iterators.toList(result), contains(1, 2));
+
+        result = Iterators.limit(values.iterator(), 0);
+        assertThat(result.hasNext(), is(false));
+        assertThat(Iterators.toList(result), is(empty()));
+    }
+
     public void testFailFast() {
         final var array = randomIntegerArray();
         assertEmptyIterator(Iterators.failFast(Iterators.forArray(array), () -> true));