2 files changed, 43 insertions, 24 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/App.java b/src/main/java/com/it_jaros/jscanner/App.java
index 8fa3cca..700b3f0 100644
--- a/src/main/java/com/it_jaros/jscanner/App.java
+++ b/src/main/java/com/it_jaros/jscanner/App.java
@@ -1,17 +1,14 @@
 package com.it_jaros.jscanner;
 
-import java.time.Duration;
 import java.util.BitSet;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 public class App {
 
     private static Thread shutdownHook;
     private static final CountDownLatch shutdownLatch = new CountDownLatch(1);
-    private static final int shutdownTimeoutInMillis = 5000;
 
     public static void main(String[] args) {
         try {
@@ -113,10 +110,7 @@ public class App {
     private static void shutdown(Scanner scanner) {
         try {
             scanner.cancel();
-            if (!scanner.awaitTermination(Duration.ofMillis(shutdownTimeoutInMillis))) {
-                scanner.cancelNow();
-            }
-            shutdownLatch.await(shutdownTimeoutInMillis, TimeUnit.MILLISECONDS);
+            shutdownLatch.await();
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
         } catch (Exception e) {
diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index ae02d3d..40c2e23 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -10,7 +10,6 @@ import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.locks.LockSupport;
 import java.util.function.Function;
 
 public class Scanner implements AutoCloseable {
@@ -73,10 +72,12 @@ public class Scanner implements AutoCloseable {
 
         // the main thread is the consumer
         // Let the consumer run as long as
-        // if there is something in queue and the scanner is not canceled
-        //   -> the producer did not have time yet to add new tasks
-        // or if in queue is still some workers left that need to finish
-        while ((state.running().get() && !cancelled) || state.inPipeline().get() > 0) {
+        // as the producer runs
+        // or if still tasks are pending in pipeline
+        // we do not listen to canceled here because we want
+        // because we want all results (also partial) ones collected here to print
+        // whatever is already there
+        while (state.running().get() || state.inPipeline().get() > 0) {
             try {
                 // let's check for results and give add them to our scan data holder object
                 Future<ScanResult> finishedHost = state.completionService().poll(10, TimeUnit.MILLISECONDS);
@@ -108,6 +109,7 @@ public class Scanner implements AutoCloseable {
 
     /**
      * Scans the ports of a given host
+     *
      * @param host
      * @param scan
      * @return
@@ -136,12 +138,15 @@ public class Scanner implements AutoCloseable {
                 maxWorkersPerHost,
                 port -> () -> {
                     waitForSlot(portSlotFactory);
+                    if (cancelled) return null;
                     return checkPort(host, port, scan);
                 }
         );
 
         // consumer is the main thread
-        while ((state.running().get() && !cancelled) || state.inPipeline().get() > 0) {
+        // we run as long as the producer is running or as long as things are in pipeline to be processed
+        // only exception is when cancelled is set
+        while ((state.running().get() || state.inPipeline().get() > 0) && !cancelled) {
             try {
                 Future<PortResult> portResultFuture = state.completionService().poll(10, TimeUnit.MILLISECONDS);
                 if (portResultFuture == null) {
@@ -206,16 +211,36 @@ public class Scanner implements AutoCloseable {
             try {
                 while (queue.hasNext() && !cancelled) {
                     activeWorkers.acquire();
-                    // just in case while waiting something changed
-                    if (cancelled) {
-                        break;
-                    }
 
-                    // here we are filling the completion service
-                    // host <-> virtual thread
-                    final INPUT item = queue.next();
-                    inPipeline.incrementAndGet();
-                    completionService.submit(taskFactory.apply(item));
+                    // remember if we submitted anything
+                    // so we can release the semaphore
+                    boolean submitted = false;
+                    try {
+                        // just in case something changed while waiting
+                        // for the semaphore
+                        if (cancelled) {
+                            break;
+                        }
+
+                        // get next item and create callable
+                        // using lambda expression
+                        final INPUT item = queue.next();
+                        Callable<OUTPUT> task = taskFactory.apply(item);
+                        inPipeline.incrementAndGet();
+                        try {
+                            // here we are filling the completion service
+                            // host <-> virtual thread
+                            completionService.submit(task);
+                            submitted = true;
+                        } catch (Throwable e) {
+                            inPipeline.decrementAndGet();
+                            throw e;
+                        }
+                    } finally {
+                        if (!submitted) {
+                            activeWorkers.release();
+                        }
+                    }
                 }
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
@@ -236,12 +261,12 @@ public class Scanner implements AutoCloseable {
         return false;
     }
 
-    private void waitForSlot(AtomicLong scanSlotFactory) {
+    private void waitForSlot(AtomicLong scanSlotFactory) throws InterruptedException {
         if (delayInNanos > 0) {
             long slot = scanSlotFactory.getAndAdd(delayInNanos);
             long wait = slot - System.nanoTime();
             if (wait > 0)
-                LockSupport.parkNanos(wait);
+                Thread.sleep(Duration.ofNanos(wait));
         }
     }