Streams v2: non-rewindable streaming

Author: Burak Serdar

crud/src/main/java/com/redhat/lightblue/assoc/CompositeFindImpl.java

@@ -21,6 +21,7 @@
 import java.util.Set;
 import java.util.HashSet;
 import java.util.List;
+import java.util.ArrayList;
 
 import java.util.concurrent.Executors;
 
@@ -143,7 +144,9 @@ public void explain(OperationContext ctx,
         initialize(ctx,req);
         ExecutionContext executionContext = new ExecutionContext(ctx,null);
         JsonDoc doc=new JsonDoc(executionPlan.explain(executionContext));
-        ctx.setDocumentStream(new ListDocumentStream(new DocCtx(doc)));
+        ArrayList<DocCtx> l=new ArrayList<>(1);
+        l.add(new DocCtx(doc));
+        ctx.setDocumentStream(new ListDocumentStream<DocCtx>(l));
     }
 
     @Override

crud/src/main/java/com/redhat/lightblue/assoc/ep/AbstractSearchStep.java

@@ -20,6 +20,11 @@
 
 import java.util.List;
 import java.util.ArrayList;
+import java.util.Spliterator;
+import java.util.Spliterators;
+
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -91,7 +96,12 @@ public void setQueries(List<Conjunct> conjuncts) {
 
     @Override
     public StepResult<ResultDocument> getResults(ExecutionContext ctx) {
-        return new DocumentStreamStepResult<ResultDocument>(getSearchResults(ctx));
+        return new StepResult<ResultDocument>() {
+            @Override
+            public Stream<ResultDocument> stream() {
+                return StreamSupport.stream(Spliterators.spliteratorUnknownSize(getSearchResults(ctx),Spliterator.IMMUTABLE),false);
+            }
+        };
     }
 
     protected abstract DocumentStream<ResultDocument> getSearchResults(ExecutionContext ctx);

crud/src/main/java/com/redhat/lightblue/assoc/ep/DocumentStreamStepResult.java

@@ -57,12 +57,7 @@ public JoinSearch(ExecutionBlock block, Source<JoinTuple> source) {
 
     @Override
     protected DocumentStream<ResultDocument> getSearchResults(final ExecutionContext ctx) {
-        return new DocumentStream<ResultDocument>() {
-            @Override
-            public Iterator<ResultDocument> getDocuments() {
-                return new BatchQueryIterator(256,ctx);
-            }
-        };
+        return new BatchQueryIterator(256,ctx);
     }
 
     /**
@@ -72,12 +67,12 @@ public Iterator<ResultDocument> getDocuments() {
      * batch of queries, computes results, and streams them to the
      * caller
      */
-    private class BatchQueryIterator implements Iterator<ResultDocument> {
+    private class BatchQueryIterator implements DocumentStream<ResultDocument> {
         private final int batchSize;
         private final ExecutionContext ctx;
         private final Iterator<JoinTuple> sourceStream;
 
-        private Iterator<DocCtx> currentIterator;
+        private DocumentStream<DocCtx> currentIterator;
         private boolean done=false; // Are we still iterating, or are we done?
 
         public BatchQueryIterator(int batchSize,ExecutionContext ctx) {
@@ -110,12 +105,21 @@ public ResultDocument next() {
             }
             throw new NoSuchElementException();
         }
+
+        @Override
+        public void close() {
+            if(currentIterator!=null)
+                currentIterator.close();
+        }            
 
         private void retrieveNextBatch() {
             do {
                 int n=0;
                 ArrayList<QueryExpression> qBatch=new ArrayList<>(batchSize);
-                currentIterator=null;
+                if(currentIterator!=null) {
+                    currentIterator.close();
+                    currentIterator=null;
+                }
                 while(sourceStream.hasNext()&&n<batchSize) {
                     JoinTuple t=sourceStream.next();
                     qBatch.addAll(Searches.writeQueriesForJoinTuple(t, block));
@@ -130,9 +134,11 @@ private void retrieveNextBatch() {
                     findRequest.setFrom(from);
                     findRequest.setTo(to);
                     OperationContext opctx = search(ctx, findRequest);
-                    currentIterator=opctx.getDocumentStream().getDocuments();
-                    if(!currentIterator.hasNext())
+                    currentIterator=opctx.getDocumentStream();
+                    if(!currentIterator.hasNext()) {
+                        currentIterator.close();
                         currentIterator=null;
+                    }
                 } else {
                     done=true;
                 }

crud/src/main/java/com/redhat/lightblue/assoc/ep/JoinSearch.java

@@ -38,7 +38,6 @@
 import com.redhat.lightblue.crud.CRUDFindRequest;
 import com.redhat.lightblue.crud.DocCtx;
 import com.redhat.lightblue.crud.DocumentStream;
-import com.redhat.lightblue.crud.MapDocumentStream;
 import com.redhat.lightblue.crud.ListDocumentStream;
 
 import com.redhat.lightblue.util.JsonDoc;
@@ -57,7 +56,7 @@ public Search(ExecutionBlock block) {
     }
 
     protected DocumentStream<ResultDocument> postProcess(OperationContext result, ExecutionContext ctx) {
-        return new MapDocumentStream<DocCtx,ResultDocument>(result.getDocumentStream(),x->new ResultDocument(block,x.getOutputDocument()));
+        return DocumentStream.map(result.getDocumentStream(),x->new ResultDocument(block,x.getOutputDocument()));
     }
 
     public OperationContext search(ExecutionContext ctx) {
@@ -96,13 +95,14 @@ public JsonNode explain(ExecutionContext ctx) {
         OperationContext searchCtx = ctx.getOperationContext().
             getDerivedOperationContext(block.getMetadata().getName(), req);
         new SimpleFindImpl(block.getMetadata(), searchCtx.getFactory()).explain(searchCtx,req);
-        Iterator<DocCtx> itr=searchCtx.getDocumentStream().getDocuments();
+        DocumentStream<DocCtx> itr=searchCtx.getDocumentStream();
         List<JsonDoc> docs=new ArrayList<>();
         while(itr.hasNext()) {
             DocCtx doc=itr.next();
             if(!doc.hasErrors())
                 docs.add(doc.getOutputDocument());
         }
+        itr.close();
         if(!docs.isEmpty()) {
             if(docs.size()==1) {
                 node.set("implementation",docs.get(0).getRoot());

crud/src/main/java/com/redhat/lightblue/assoc/ep/Search.java

@@ -11,14 +11,22 @@
  */
 public class StepResultDocumentStream implements DocumentStream<DocCtx> {
 
-    private final StepResult<DocCtx> result;
+    private final Iterator<DocCtx> itr;
 
     public StepResultDocumentStream(StepResult<DocCtx> result) {
-        this.result=result;
+        this.itr=result.stream().iterator();
     }
 
     @Override
-    public Iterator<DocCtx> getDocuments() {
-        return result.stream().iterator();
+    public boolean hasNext() {
+        return itr.hasNext();
     }
+
+    @Override
+    public DocCtx next() {
+        return itr.next();
+    }
+
+    @Override
+    public void close() {}
 }

crud/src/main/java/com/redhat/lightblue/assoc/ep/StepResultDocumentStream.java

@@ -71,8 +71,9 @@ public CRUDOperationContext(CRUDOperation op,
         this.CRUDOperation = op;
         this.entityName = entityName;
         this.factory = f;
-        if(docs!=null)
-        	documents=docs.stream().map(x->new DocCtx(x)).collect(Collectors.toList());
+        if(docs!=null) {
+        	setInputDocuments(docs.stream().map(x->new DocCtx(x)).collect(Collectors.toList()));
+        }
         this.hookManager = new HookManager(factory.getHookResolver(), factory.getNodeFactory());
         this.callerRoles = new HashSet<>();
         this.executionOptions = eo;
@@ -88,8 +89,9 @@ public CRUDOperationContext(CRUDOperation op,
         this.CRUDOperation = op;
         this.entityName = entityName;
         this.factory = f;
-        if(docs!=null)
-        	documents=docs.stream().map(x->new DocCtx(x)).collect(Collectors.toList());
+        if(docs!=null) {
+        	setInputDocuments(docs.stream().map(x->new DocCtx(x)).collect(Collectors.toList()));
+        }
         this.callerRoles = callerRoles;
         this.hookManager = hookManager;
         this.executionOptions = eo;
@@ -109,7 +111,7 @@ public CRUDOperationContext(CRUDOperation op,
         this.CRUDOperation = op;
         this.entityName = entityName;
         this.factory = f;
-        this.documents = docs;
+        setInputDocuments(docs);
         this.callerRoles = callerRoles;
         this.hookManager = hookManager;
         this.executionOptions = eo;
@@ -183,26 +185,22 @@ public void setDocumentStream(DocumentStream<DocCtx> stream) {
      * to iterate through those documents.
      */
     public DocumentStream<DocCtx> getDocumentStream() {
-        return documentStream==null?new ListDocumentStream<DocCtx>(documents==null?new ArrayList<DocCtx>():documents):documentStream;
+        return documentStream;
     }
 
     /**
-     * Returns the list of documents in the context.
-     *
-     * If the list of documents is not set, but the document stream is set, this call constructs a list from the
-     * document stream and returns that
+     * Returns the list of input documents in the context.
      */
-    public List<DocCtx> getDocuments() {
-        if(documents==null&&documentStream!=null) {
-            documents=new ArrayList<DocCtx>();
-            for(Iterator<DocCtx> itr=documentStream.getDocuments();itr.hasNext();)
-                documents.add(itr.next());
-        }
+    public List<DocCtx> getInputDocuments() {
         return documents;
     }
 
-    public void setDocuments(List<DocCtx> docs) {
+    public void setInputDocuments(List<DocCtx> docs) {
         documents = docs;
+        if(documents!=null) {
+            // Set the documentStream to input documents. A later call to setDocumentStream may set the output stream
+            documentStream=new ListDocumentStream(documents);
+        }
     }
 
     /**

crud/src/main/java/com/redhat/lightblue/crud/CRUDOperationContext.java

@@ -19,11 +19,32 @@
 package com.redhat.lightblue.crud;
 
 import java.util.Iterator;
+import java.util.function.Function;
 
 /**
  * This interface is used to stream documents. A call to getDocuments
  * should return an iterator of the document resultset.
  */
-public interface DocumentStream<T> {
-    Iterator<T> getDocuments();
+public interface DocumentStream<T> extends Iterator<T>{
+    /**
+     * Close the document stream
+     */
+    void close();
+
+    public static <S,D> DocumentStream<D> map(final DocumentStream<S> source,final Function<S,D> map) {
+        return new DocumentStream<D>() {
+            @Override
+            public boolean hasNext() {
+                return source.hasNext();
+            }
+            @Override
+            public void close() {
+                source.close();
+            }
+            @Override
+            public D next() {
+                return map.apply(source.next());
+            }
+        };
+    }
 }

crud/src/main/java/com/redhat/lightblue/crud/DocumentStream.java

@@ -18,28 +18,39 @@
  */
 package com.redhat.lightblue.crud;
 
-import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 
 /**
- * Default trivial implementation of document stream using a list
+ * Default trivial implementation of document stream that uses a list to hold its elements
  */
 public class ListDocumentStream<T> implements DocumentStream<T> {
 
     private final List<T> documents;
-    
-    public ListDocumentStream(T item) {
-    	this.documents=new ArrayList<T>(1);
-    	this.documents.add(item);
-    }
+    private final Iterator<T> itr;
 
     public ListDocumentStream(List<T> list) {
         this.documents=list;
+        this.itr=documents.iterator();
+    }
+
+    /**
+     * Return the underlying list
+     */
+    public List<T> getDocuments() {
+        return documents;
     }
 
     @Override
-    public Iterator<T> getDocuments() {
-        return documents.iterator();
+    public boolean hasNext() {
+        return itr.hasNext();
     }
+
+    @Override
+    public T next() {
+        return itr.next();
+    }
+
+    @Override
+    public void close() {}
 }