CancellableCollectSpliterator: race fixed

amaembo · Sep 11, 2015 · 0b2786e · 0b2786e
1 parent 276f100
commit 0b2786e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/src/main/java/javax/util/streamex/CancellableCollectSpliterator.java b/src/main/java/javax/util/streamex/CancellableCollectSpliterator.java
@@ -52,21 +52,23 @@ public boolean tryAdvance(Consumer<? super A> action) {
 	    Spliterator<T> source = this.source;
 	    if(source == null)
 	        return false;
-	    this.source = null;
 	    acc = supplier.get();
 	    if(cancelled == null) {
+	        this.source = null;
 	        // sequential mode
             while(!cancelPredicate.test(acc) && source.tryAdvance(this)) {
                 // empty
             }
 	    } else {
 	        do {
                 if(cancelPredicate.test(acc)) {
+                    this.source = null;
                     if(isFinished())
                         cancelled.set(true);
                     break;
                 }
 	        } while(!cancelled.get() && source.tryAdvance(this));
+	        this.source = null;
 	    }
 	    action.accept(acc);
 		return true;

diff --git a/src/test/java/javax/util/streamex/MoreCollectorsTest.java b/src/test/java/javax/util/streamex/MoreCollectorsTest.java
@@ -153,6 +153,13 @@ public void testFirstLast() {
             assertFalse(supplier.toString(), supplier.get().collect(MoreCollectors.last()).isPresent());
         }
     }
+
+    @Test
+    public void testHeadTailParallel() {
+        for(int i=0; i<100; i++) {
+            assertEquals("#"+i, Arrays.asList(0, 1), IntStreamEx.range(1000).boxed().parallel().collect(MoreCollectors.head(2)));
+        }
+    }
 
     @Test
     public void testHeadTail() {