fix Flow.take waiting for n+1 elements instead of terminating at n

flows/src/main/java/com/softwaremill/jox/flows/Flow.java

@@ -1,10 +1,7 @@
 package com.softwaremill.jox.flows;
 
-import static com.softwaremill.jox.Select.defaultClause;
-import static com.softwaremill.jox.Select.selectOrClosed;
-import static com.softwaremill.jox.flows.Flows.usingEmit;
-import static com.softwaremill.jox.structured.Scopes.supervised;
-import static java.lang.Thread.sleep;
+import com.softwaremill.jox.*;
+import com.softwaremill.jox.structured.*;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -28,8 +25,11 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.*;
 
-import com.softwaremill.jox.*;
-import com.softwaremill.jox.structured.*;
+import static com.softwaremill.jox.Select.defaultClause;
+import static com.softwaremill.jox.Select.selectOrClosed;
+import static com.softwaremill.jox.flows.Flows.usingEmit;
+import static com.softwaremill.jox.structured.Scopes.supervised;
+import static java.lang.Thread.sleep;
 
 /**
  * Describes an asynchronous transformation pipeline. When run, emits elements of type `T`.
@@ -834,27 +834,24 @@ public <U> Flow<U> flatMap(ThrowingFunction<T, Flow<U>> mappingFunction) {
      * emitting `n` elements, the returned flow completes as well.
      */
     public Flow<T> take(int n) {
-        return usingEmit(
-                emit -> {
-                    AtomicInteger taken = new AtomicInteger(0);
-                    try {
-                        last.run(
-                                t -> {
-                                    if (taken.getAndIncrement() < n) {
-                                        emit.apply(t);
-                                    } else {
-                                        throw new BreakException();
-                                    }
-                                });
-                    } catch (JoxScopeExecutionException e) {
-                        if (!(e.getCause() instanceof BreakException)) {
-                            throw e;
-                        }
-                        // ignore
-                    } catch (BreakException e) {
-                        // ignore
+        return usingEmit(emit -> {
+            AtomicInteger taken = new AtomicInteger(0);
+            try {
+                last.run(t -> {
+                    final var next = taken.incrementAndGet();
+
+                    emit.apply(t);
+
+                    if (next >= n) {
+                        throw new BreakException();
                     }
                 });
+            } catch (JoxScopeExecutionException e) {
+                if (!(e.getCause() instanceof BreakException)) {
+                    throw e;
+                }
+            } catch (BreakException _) {}
+        });
     }
 
     /**

flows/src/test/java/com/softwaremill/jox/flows/FlowTest.java

@@ -1,20 +1,21 @@
 package com.softwaremill.jox.flows;
 
-import static com.softwaremill.jox.structured.Scopes.supervised;
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.contains;
-import static org.junit.jupiter.api.Assertions.*;
+import com.softwaremill.jox.Channel;
+import com.softwaremill.jox.ChannelError;
+import com.softwaremill.jox.Source;
+import com.softwaremill.jox.structured.JoxScopeExecutionException;
+import org.junit.jupiter.api.Test;
 
 import java.time.Duration;
 import java.util.*;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.IntStream;
 
-import org.junit.jupiter.api.Test;
-
-import com.softwaremill.jox.ChannelError;
-import com.softwaremill.jox.Source;
-import com.softwaremill.jox.structured.JoxScopeExecutionException;
+import static com.softwaremill.jox.structured.Scopes.supervised;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.contains;
+import static org.junit.jupiter.api.Assertions.*;
 
 class FlowTest {
 
@@ -45,12 +46,10 @@ void shouldRunToList() throws Throwable {
 
     @Test
     void shouldThrowExceptionForFailedFlow() {
-        assertThrows(
-                IllegalStateException.class,
-                () -> {
-                    Flows.failed(new IllegalStateException())
-                            .runFold(0, (acc, n) -> Integer.valueOf(acc.toString() + n));
-                });
+        assertThrows(IllegalStateException.class, () -> {
+            Flows.failed(new IllegalStateException())
+                .runFold(0, (acc, n) -> Integer.valueOf(acc.toString() + n));
+        });
     }
 
     @Test
@@ -101,6 +100,52 @@ void shouldTakeAllIfFlowEndsSooner() throws Exception {
         assertEquals(List.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10), result);
     }
 
+    @Test
+    void shouldTakeExactlyAllElementsFromFlowWithoutWaitingForMore() throws Exception {
+        supervised(scope -> {
+            var ch = Channel.<Integer>newUnlimitedChannel();
+
+            scope.fork(() -> {
+                ch.send(1);
+                ch.send(2);
+                ch.send(3);
+                return null;
+            });
+
+            var result = Flows.fromSource(ch).take(3).runToList();
+            assertEquals(List.of(1, 2, 3), result);
+            return null;
+        });
+    }
+
+    @Test
+    void shouldBreakImmediatelyAfterTakingNElements() throws Exception {
+        supervised(scope -> {
+            var ch = Channel.<Integer>newUnlimitedChannel();
+            var processedCount = new AtomicInteger(0);
+
+            scope.fork(() -> {
+                for (int i = 1; i <= 10; i++) {
+                    ch.send(i);
+                }
+                return null;
+            });
+
+            var result = Flows.fromSource(ch)
+                .map(
+                    x -> {
+                        processedCount.incrementAndGet();
+                        return x;
+                    })
+                .take(3)
+                .runToList();
+
+            assertEquals(List.of(1, 2, 3), result);
+            assertEquals(3, processedCount.get());
+            return null;
+        });
+    }
+
     @Test
     void shouldWorkWithASingleAsyncBoundary() throws Throwable {
         // given