Improved shutdown logic that finally works properly

Now if the user quits, then the main thread continues to collect
all available results
the producer and the host scan port threads close as soon as
they check the flag and return whatever they have till now

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));
         }
     }