Improve BlockingEnumeration

Author: passaro

hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/select/BlockingEnumeration.java

@@ -21,6 +21,7 @@
 import java.util.Enumeration;
 import java.util.NoSuchElementException;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.LinkedBlockingQueue;
 
 import org.reactivestreams.Subscriber;
@@ -32,103 +33,70 @@
 /**
  * Implements the {@link Enumeration} interface by subscribing to a
  * {@link SdkPublisher} instance. The enumeration will buffer a fixed
- * number of events received from the publisher and block until the first is
- * consumed. Calls to {@link #hasMoreElements()} and {@link #nextElement()}
- * can also block while waiting for new events.
+ * number of elements and only request new ones from the publisher
+ * when they are consumed. Calls to {@link #hasMoreElements()} and
+ * {@link #nextElement()} may block while waiting for new elements.
  * @param <T> the type of element.
  */
 public final class BlockingEnumeration<T> implements Enumeration<T> {
-  private static final class Token<T> {
-    public T element;
-    public Throwable error;
+  private static final class Signal<T> {
+    public final T element;
+    public final Throwable error;
 
-    public Token() {
-    }
-
-    public Token(T element) {
+    public Signal(T element) {
       this.element = element;
+      this.error = null;
     }
 
-    public Token(Throwable error) {
+    public Signal(Throwable error) {
+      this.element = null;
       this.error = error;
     }
   }
 
-  private final Token<T> END_TOKEN = new Token<>();
-  private final BlockingQueue<Token<T>> blockingQueue;
-  private Token<T> current = null;
+  private final Signal<T> END_SIGNAL = new Signal<>((Throwable)null);
+  private final CompletableFuture<Subscription> subscription = new CompletableFuture<>();
+  private final BlockingQueue<Signal<T>> signalQueue;
+  private final int bufferSize;
+  private Signal<T> current = null;
 
   /**
-   * Create an enumeration with a fixed buffer size.
+   * Create an enumeration with a fixed buffer size and an
+   * optional injected first element.
    * @param publisher the publisher feeding the enumeration.
-   * @param limit the buffer size.
+   * @param bufferSize the buffer size.
+   * @param firstElement (optional) first element the enumeration will return.
    */
-  public BlockingEnumeration(SdkPublisher<T> publisher, final int limit) {
-    this.blockingQueue = new LinkedBlockingQueue<>(limit);
-    publisher.subscribe(new Subscriber<T>() {
-      private Subscription subscription;
-
-      @Override
-      public void onSubscribe(Subscription s) {
-        this.subscription = s;
-        this.subscription.request(limit);
-      }
-
-      @Override
-      public void onNext(T t) {
-        enqueue(new Token<>(t));
-        subscription.request(1);
-      }
-
-      @Override
-      public void onError(Throwable t) {
-        enqueue(new Token<>(t));
-      }
-
-      @Override
-      public void onComplete() {
-        enqueue(END_TOKEN);
-      }
-
-      private void enqueue(Token<T> token) {
-        try {
-          blockingQueue.put(token);
-        } catch (InterruptedException e) {
-          // TODO: log?
-        }
-      }
-    });
+  public BlockingEnumeration(SdkPublisher<T> publisher,
+      final int bufferSize,
+      final T firstElement) {
+    this.signalQueue = new LinkedBlockingQueue<>();
+    this.bufferSize = bufferSize;
+    if (firstElement != null) {
+      this.current = new Signal<>(firstElement);
+    }
+    publisher.subscribe(new EnumerationSubscriber());
   }
 
   /**
-   * Create an enumeration with a fixed buffer size and an
-   * injected first element.
+   * Create an enumeration with a fixed buffer size.
    * @param publisher the publisher feeding the enumeration.
-   * @param limit the buffer size.
-   * @param firstElement first element the enumeration will return.
+   * @param bufferSize the buffer size.
    */
   public BlockingEnumeration(SdkPublisher<T> publisher,
-      final int limit,
-      T firstElement) {
-    this(publisher, limit);
-    this.current = new Token<>(firstElement);
+      final int bufferSize) {
+    this(publisher, bufferSize, null);
   }
 
   @Override
   public boolean hasMoreElements() {
     if (current == null) {
-      current = dequeue();
-    }
-    return current != END_TOKEN;
-  }
-
-  @Override
-  public T nextElement() {
-    if (current == null) {
-      current = dequeue();
-    }
-    if (current == END_TOKEN) {
-      throw new NoSuchElementException();
+      try {
+        current = signalQueue.take();
+      } catch (InterruptedException e) {
+        current = new Signal<>(e);
+        subscription.thenAccept(Subscription::cancel);
+      }
     }
     if (current.error != null) {
       if (current.error instanceof SdkException) {
@@ -137,17 +105,47 @@ public T nextElement() {
         throw SdkException.create("Unexpected error", current.error);
       }
     }
+    return current != END_SIGNAL;
+  }
+
+  @Override
+  public T nextElement() {
+    if (!hasMoreElements()) {
+      throw new NoSuchElementException();
+    }
     T element = current.element;
     current = null;
+    subscription.thenAccept(s -> s.request(1));
     return element;
   }
 
-  private Token<T> dequeue() {
-    try {
-      return blockingQueue.take();
-    } catch (InterruptedException e) {
-      // TODO: log?
-      return END_TOKEN;
+  private final class EnumerationSubscriber implements Subscriber<T> {
+
+    @Override
+    public void onSubscribe(Subscription s) {
+      long request = bufferSize;
+      if (current != null) {
+        request--;
+      }
+      if (request > 0) {
+        s.request(request);
+      }
+      subscription.complete(s);
+    }
+
+    @Override
+    public void onNext(T t) {
+      signalQueue.add(new Signal<>(t));
+    }
+
+    @Override
+    public void onError(Throwable t) {
+      signalQueue.add(new Signal<>(t));
+    }
+
+    @Override
+    public void onComplete() {
+      signalQueue.add(END_SIGNAL);
     }
   }
 }