Refactor all iterator management service operations to optimize for s…

…mall iterator batches.
atomix · Aug 8, 2018 · b6e2d3e · b6e2d3e
1 parent dfecbd5
commit b6e2d3e
Show file tree

Hide file tree

Showing 17 changed files with 282 additions and 107 deletions.
diff --git a/core/src/main/java/io/atomix/core/collection/impl/DefaultDistributedCollectionService.java b/core/src/main/java/io/atomix/core/collection/impl/DefaultDistributedCollectionService.java
@@ -34,6 +34,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 
 import static io.atomix.core.collection.impl.CollectionUpdateResult.noop;
 import static io.atomix.core.collection.impl.CollectionUpdateResult.ok;
@@ -183,9 +184,23 @@ public void unlisten() {
   }
 
   @Override
-  public long iterate() {
+  public IteratorBatch<E> iterate() {
-    iterators.put(getCurrentIndex(), new IteratorContext(getCurrentSession().sessionId().id()));
+    return iterate(IteratorContext::new);
-    return getCurrentIndex();
+  }
+
+  protected IteratorBatch<E> iterate(Function<Long, AbstractIteratorContext> contextFactory) {
+    AbstractIteratorContext iterator = contextFactory.apply(getCurrentSession().sessionId().id());
+    if (!iterator.iterator().hasNext()) {
+      return null;
+    }
+
+    long iteratorId = getCurrentIndex();
+    iterators.put(iteratorId, iterator);
+    IteratorBatch<E> batch = next(iteratorId, 0);
+    if (batch.complete()) {
+      iterators.remove(iteratorId);
+    }
+    return batch;
   }
 
   @Override
@@ -211,7 +226,7 @@ public IteratorBatch<E> next(long iteratorId, int position) {
     if (elements.isEmpty()) {
       return null;
     }
-    return new IteratorBatch<>(context.position(), elements);
+    return new IteratorBatch<>(iteratorId, context.position(), elements, !context.iterator().hasNext());
   }
 
   @Override

diff --git a/core/src/main/java/io/atomix/core/iterator/impl/IterableService.java b/core/src/main/java/io/atomix/core/iterator/impl/IterableService.java
@@ -29,7 +29,7 @@ public interface IterableService<E> {
    * @return the iterator ID
    */
   @Command
-  long iterate();
+  IteratorBatch<E> iterate();
 
   /**
    * Returns the next batch of elements for the given iterator.

diff --git a/core/src/main/java/io/atomix/core/iterator/impl/IteratorBatch.java b/core/src/main/java/io/atomix/core/iterator/impl/IteratorBatch.java
@@ -22,13 +22,26 @@
  * Iterator batch.
  */
 public final class IteratorBatch<T> implements Iterator<T> {
+  private final long id;
   private final int position;
   private final Collection<T> entries;
+  private final boolean complete;
   private transient volatile Iterator<T> iterator;
 
-  public IteratorBatch(int position, Collection<T> entries) {
+  public IteratorBatch(long id, int position, Collection<T> entries, boolean complete) {
+    this.id = id;
     this.position = position;
     this.entries = entries;
+    this.complete = complete;
+  }
+
+  /**
+   * Returns the iterator identifier.
+   *
+   * @return the iterator identifier
+   */
+  public long id() {
+    return id;
   }
 
   /**
@@ -49,6 +62,15 @@ public Collection<T> entries() {
     return entries;
   }
 
+  /**
+   * Returns a boolean indicating whether the batch is complete.
+   *
+   * @return indicates whether this batch completes iteration
+   */
+  public boolean complete() {
+    return complete;
+  }
+
   private Iterator<T> iterator() {
     if (iterator == null) {
       synchronized (this) {

diff --git a/core/src/main/java/io/atomix/core/iterator/impl/OpenFunction.java b/core/src/main/java/io/atomix/core/iterator/impl/OpenFunction.java
@@ -19,13 +19,14 @@
  * Iterator open function.
  */
 @FunctionalInterface
-public interface OpenFunction<S> {
+public interface OpenFunction<S, T> {
 
   /**
    * Opens the iterator using the given service proxy.
    *
    * @param service the service proxy
-   * @return
+   * @return the initial iterator batch
    */
-  long open(S service);
+  IteratorBatch<T> open(S service);
+
 }
diff --git a/core/src/main/java/io/atomix/core/iterator/impl/PartitionedProxyIterator.java b/core/src/main/java/io/atomix/core/iterator/impl/PartitionedProxyIterator.java
@@ -29,15 +29,15 @@
 public class PartitionedProxyIterator<S, T> implements AsyncIterator<T> {
   private final ProxyClient<S> client;
   private final Iterator<AsyncIterator<T>> partitions;
-  private final OpenFunction<S> openFunction;
+  private final OpenFunction<S, T> openFunction;
   private final NextFunction<S, T> nextFunction;
   private final CloseFunction<S> closeFunction;
   private volatile AsyncIterator<T> iterator;
   private AtomicBoolean closed = new AtomicBoolean();
 
   public PartitionedProxyIterator(
       ProxyClient<S> client,
-      OpenFunction<S> openFunction,
+      OpenFunction<S, T> openFunction,
       NextFunction<S, T> nextFunction,
       CloseFunction<S> closeFunction) {
     this.client = client;

diff --git a/core/src/main/java/io/atomix/core/iterator/impl/ProxyIterator.java b/core/src/main/java/io/atomix/core/iterator/impl/ProxyIterator.java
@@ -21,7 +21,6 @@
 import io.atomix.utils.concurrent.Futures;
 import io.atomix.utils.concurrent.OrderedFuture;
 
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
 import java.util.concurrent.CompletableFuture;
@@ -35,23 +34,22 @@ public class ProxyIterator<S, T> implements AsyncIterator<T> {
   private final PartitionId partitionId;
   private final NextFunction<S, T> nextFunction;
   private final CloseFunction<S> closeFunction;
-  private final CompletableFuture<Long> openFuture;
+  private final CompletableFuture<IteratorBatch<T>> openFuture;
   private volatile CompletableFuture<IteratorBatch<T>> batch;
   private volatile CompletableFuture<Void> closeFuture;
 
   public ProxyIterator(
       ProxyClient<S> client,
       PartitionId partitionId,
-      OpenFunction<S> openFunction,
+      OpenFunction<S, T> openFunction,
       NextFunction<S, T> nextFunction,
       CloseFunction<S> closeFunction) {
     this.client = client;
     this.partitionId = partitionId;
     this.nextFunction = nextFunction;
     this.closeFunction = closeFunction;
     this.openFuture = OrderedFuture.wrap(client.applyOn(partitionId, openFunction::open));
-    this.batch = CompletableFuture.completedFuture(
+    this.batch = openFuture;
-        new IteratorBatch<T>(0, Collections.emptyList()));
   }
 
   /**
@@ -76,21 +74,31 @@ private CompletableFuture<Iterator<T>> batch() {
    * @return the next batch of entries from the cluster
    */
   private CompletableFuture<IteratorBatch<T>> fetch(int position) {
-    return openFuture.thenCompose(id -> client.applyOn(partitionId, service -> nextFunction.next(service, id, position))
+    return openFuture.thenCompose(initialBatch -> {
-        .thenCompose(batch -> {
+      if (!initialBatch.complete()) {
-          if (batch == null) {
+        return client.applyOn(partitionId, service -> nextFunction.next(service, initialBatch.id(), position))
-            return close().thenApply(v -> null);
+            .thenCompose(nextBatch -> {
-          }
+              if (nextBatch == null) {
-          return CompletableFuture.completedFuture(batch);
+                return close().thenApply(v -> null);
-        }));
+              }
+              return CompletableFuture.completedFuture(nextBatch);
+            });
+      }
+      return CompletableFuture.completedFuture(null);
+    });
   }
 
   @Override
   public CompletableFuture<Void> close() {
     if (closeFuture == null) {
       synchronized (this) {
         if (closeFuture == null) {
-          closeFuture = openFuture.thenCompose(id -> client.acceptOn(partitionId, service -> closeFunction.close(service, id)));
+          closeFuture = openFuture.thenCompose(initialBatch -> {
+            if (initialBatch != null && !initialBatch.complete()) {
+              return client.acceptOn(partitionId, service -> closeFunction.close(service, initialBatch.id()));
+            }
+            return CompletableFuture.completedFuture(null);
+          });
         }
       }
     }

diff --git a/core/src/main/java/io/atomix/core/map/impl/AbstractAtomicMapService.java b/core/src/main/java/io/atomix/core/map/impl/AbstractAtomicMapService.java
@@ -46,6 +46,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
@@ -534,61 +536,80 @@ public void clear() {
   }
 
   @Override
-  public long iterateKeys() {
+  public IteratorBatch<K> iterateKeys() {
-    return iterateEntries();
+    return iterate(DefaultIterator::new, (k, v) -> k);
   }
 
   @Override
   public IteratorBatch<K> nextKeys(long iteratorId, int position) {
-    IteratorBatch<Map.Entry<K, Versioned<byte[]>>> batch = nextEntries(iteratorId, position);
+    return next(iteratorId, position, (k, v) -> k);
-    return batch == null ? null : new IteratorBatch<>(batch.position(), batch.entries().stream()
-        .map(Map.Entry::getKey)
-        .collect(Collectors.toList()));
   }
 
   @Override
   public void closeKeys(long iteratorId) {
-    closeEntries(iteratorId);
+    close(iteratorId);
   }
 
   @Override
-  public long iterateValues() {
+  public IteratorBatch<Versioned<byte[]>> iterateValues() {
-    return iterateEntries();
+    return iterate(DefaultIterator::new, (k, v) -> v);
   }
 
   @Override
   public IteratorBatch<Versioned<byte[]>> nextValues(long iteratorId, int position) {
-    IteratorBatch<Map.Entry<K, Versioned<byte[]>>> batch = nextEntries(iteratorId, position);
+    return next(iteratorId, position, (k, v) -> v);
-    return batch == null ? null : new IteratorBatch<>(batch.position(), batch.entries().stream()
-        .map(Map.Entry::getValue)
-        .collect(Collectors.toList()));
   }
 
   @Override
   public void closeValues(long iteratorId) {
-    closeEntries(iteratorId);
+    close(iteratorId);
   }
 
   @Override
-  public long iterateEntries() {
+  public IteratorBatch<Map.Entry<K, Versioned<byte[]>>> iterateEntries() {
-    entryIterators.put(getCurrentIndex(), new DefaultIterator(getCurrentSession().sessionId().id()));
+    return iterate(DefaultIterator::new, Maps::immutableEntry);
-    return getCurrentIndex();
   }
 
   @Override
   public IteratorBatch<Map.Entry<K, Versioned<byte[]>>> nextEntries(long iteratorId, int position) {
+    return next(iteratorId, position, Maps::immutableEntry);
+  }
+
+  @Override
+  public void closeEntries(long iteratorId) {
+    close(iteratorId);
+  }
+
+  protected <T> IteratorBatch<T> iterate(
+      Function<Long, IteratorContext> contextFactory,
+      BiFunction<K, Versioned<byte[]>, T> function) {
+    IteratorContext iterator = contextFactory.apply(getCurrentSession().sessionId().id());
+    if (!iterator.iterator().hasNext()) {
+      return null;
+    }
+
+    long iteratorId = getCurrentIndex();
+    entryIterators.put(iteratorId, iterator);
+    IteratorBatch<T> batch = next(iteratorId, 0, function);
+    if (batch.complete()) {
+      entryIterators.remove(iteratorId);
+    }
+    return batch;
+  }
+
+  protected <T> IteratorBatch<T> next(long iteratorId, int position, BiFunction<K, Versioned<byte[]>, T> function) {
     IteratorContext context = entryIterators.get(iteratorId);
     if (context == null) {
       return null;
     }
 
-    List<Map.Entry<K, Versioned<byte[]>>> entries = new ArrayList<>();
+    List<T> entries = new ArrayList<>();
     int size = 0;
     while (context.iterator().hasNext()) {
       context.incrementPosition();
       if (context.position() > position) {
         Map.Entry<K, MapEntryValue> entry = context.iterator().next();
-        entries.add(Maps.immutableEntry(entry.getKey(), toVersioned(entry.getValue())));
+        entries.add(function.apply(entry.getKey(), toVersioned(entry.getValue())));
         size += entry.getValue().value().length;
 
         if (size >= MAX_ITERATOR_BATCH_SIZE) {
@@ -600,11 +621,10 @@ public IteratorBatch<Map.Entry<K, Versioned<byte[]>>> nextEntries(long iteratorI
     if (entries.isEmpty()) {
       return null;
     }
-    return new IteratorBatch<>(context.position, entries);
+    return new IteratorBatch<>(iteratorId, context.position, entries, !context.iterator().hasNext());
   }
 
-  @Override
+  protected void close(long iteratorId) {
-  public void closeEntries(long iteratorId) {
     entryIterators.remove(iteratorId);
   }
 

diff --git a/core/src/main/java/io/atomix/core/map/impl/AbstractAtomicNavigableMapService.java b/core/src/main/java/io/atomix/core/map/impl/AbstractAtomicNavigableMapService.java
@@ -17,6 +17,7 @@
 package io.atomix.core.map.impl;
 
 import com.google.common.collect.Maps;
+import io.atomix.core.iterator.impl.IteratorBatch;
 import io.atomix.core.map.AtomicNavigableMapType;
 import io.atomix.core.transaction.TransactionId;
 import io.atomix.primitive.PrimitiveType;
@@ -253,15 +254,28 @@ public int subMapSize(K fromKey, boolean fromInclusive, K toKey, boolean toInclu
   }
 
   @Override
-  public long subMapIterate(K fromKey, boolean fromInclusive, K toKey, boolean toInclusive) {
+  public IteratorBatch<K> subMapIterateKeys(K fromKey, boolean fromInclusive, K toKey, boolean toInclusive) {
-    entryIterators.put(getCurrentIndex(), new AscendingIterator(getCurrentSession().sessionId().id(), fromKey, fromInclusive, toKey, toInclusive));
+    return iterate(sessionId -> new AscendingIterator(sessionId, fromKey, fromInclusive, toKey, toInclusive), (k, v) -> k);
-    return getCurrentIndex();
   }
 
   @Override
-  public long subMapIterateDescending(K fromKey, boolean fromInclusive, K toKey, boolean toInclusive) {
+  public IteratorBatch<Map.Entry<K, Versioned<byte[]>>> subMapIterateEntries(K fromKey, boolean fromInclusive, K toKey, boolean toInclusive) {
-    entryIterators.put(getCurrentIndex(), new DescendingIterator(getCurrentSession().sessionId().id(), fromKey, fromInclusive, toKey, toInclusive));
+    return iterate(sessionId -> new AscendingIterator(sessionId, fromKey, fromInclusive, toKey, toInclusive), Maps::immutableEntry);
-    return getCurrentIndex();
+  }
+
+  @Override
+  public IteratorBatch<Versioned<byte[]>> subMapIterateValues(K fromKey, boolean fromInclusive, K toKey, boolean toInclusive) {
+    return iterate(sessionId -> new AscendingIterator(sessionId, fromKey, fromInclusive, toKey, toInclusive), (k, v) -> v);
+  }
+
+  @Override
+  public IteratorBatch<K> subMapIterateDescendingKeys(K fromKey, boolean fromInclusive, K toKey, boolean toInclusive) {
+    return iterate(sessionId -> new DescendingIterator(sessionId, fromKey, fromInclusive, toKey, toInclusive), (k, v) -> k);
+  }
+
+  @Override
+  public IteratorBatch<Map.Entry<K, Versioned<byte[]>>> subMapIterateDescendingEntries(K fromKey, boolean fromInclusive, K toKey, boolean toInclusive) {
+    return iterate(sessionId -> new DescendingIterator(sessionId, fromKey, fromInclusive, toKey, toInclusive), Maps::immutableEntry);
   }
 
   @Override