1 files changed, 35 insertions, 17 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index 286afde..12cfe15 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -4,6 +4,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.net.*;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.*;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -23,6 +24,7 @@ public class Scanner implements AutoCloseable {
     private final int maxWorkersPerHost;
     private final int timeoutInMillis;
     private final long delayInNanos;
+    private final Phaser runningScans = new Phaser(0);
 
     public Scanner(
             int socketLimit,
@@ -53,15 +55,20 @@ public class Scanner implements AutoCloseable {
         }
 
         scan.start();
+        runningScans.register();
         try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor()) {
             CompletionService<ScanResult> completionService = new ExecutorCompletionService<>(executor);
             Queue<String> hostsInQueue = new ArrayDeque<>(scan.getHosts());
 
             int hostsToProcess = hostsInQueue.size();
             int activeHostWorkers = 0;
-            while (hostsToProcess > 0 && !cancelled) {
+            while ((hostsToProcess > 0 && !cancelled) || activeHostWorkers > 0) {
                 // process queue
-                while (!hostsInQueue.isEmpty() && activeHostWorkers < maxHostsLimit && !cancelled) {
+                while (
+                        !hostsInQueue.isEmpty() // continue as long as we have hosts to process
+                                && activeHostWorkers < maxHostsLimit // limit is not reached
+                                && !cancelled // and scanner not closed
+                ) {
                     activeHostWorkers++;
                     final String host = hostsInQueue.poll();
                     completionService.submit(() -> scanHostPorts(host, scan.getPorts(), scan));
@@ -81,10 +88,8 @@ public class Scanner implements AutoCloseable {
                     System.out.printf("%s -> %s%n", e.getClass().getSimpleName(), e.getMessage());
                 }
             }
-
-            if (cancelled) {
-                executor.shutdown();
-            }
+        } finally {
+            runningScans.arriveAndDeregister();
         }
 
         scan.stop();
@@ -113,8 +118,12 @@ public class Scanner implements AutoCloseable {
             AtomicLong portSlotFactory = new AtomicLong(System.nanoTime());
             int activePerHostWorkers = 0;
             int portsToProcess = portRange.getTotal();
-            while (portsToProcess > 0 && !cancelled) {
-                while (!cancelled && portRange.hasNext() && activePerHostWorkers < maxWorkersPerHost) {
+            while ((portsToProcess > 0 && !cancelled) || activePerHostWorkers > 0) {
+                while (
+                        portRange.hasNext() // continue as long as we have ports
+                                && activePerHostWorkers < maxWorkersPerHost // and we have not reached our limit 
+                                && !cancelled // and the scanner is not about to close
+                ) {
                     activePerHostWorkers++;
                     final int currentPort = portRange.next();
                     completionService.submit(() -> {
@@ -128,19 +137,24 @@ public class Scanner implements AutoCloseable {
                     if (portResultFuture == null) {
                         continue;
                     }
+
+                    PortResult portResult = portResultFuture.get();
+
                     activePerHostWorkers--;
                     portsToProcess--;
                     progress.done().incrementAndGet();
-
-                    // bitset is not thread-safe, so it is set
-                    // outside the other virtual threads that update progress
-                    PortResult portResult = portResultFuture.get();
                     switch (portResult.getState()) {
                         case OPEN -> {
+                            // bitset is not thread-safe, so it is set
+                            // outside the other virtual threads that update progress
                             openPorts.set(portResult.getPort());
+                            progress.open().incrementAndGet();
                             serviceTypes.put(portResult.getPort(), ServiceDetector.detect(portResult.getBanner()));
                         }
-                        case FILTERED -> filteredPorts.set(portResult.getPort());
+                        case FILTERED -> {
+                            filteredPorts.set(portResult.getPort());
+                            progress.filtered().incrementAndGet();
+                        }
                     }
                 } catch (InterruptedException ignored) {
                     Thread.currentThread().interrupt();
@@ -149,10 +163,6 @@ public class Scanner implements AutoCloseable {
                 }
             }
 
-            if (!cancelled) {
-                executor.shutdown();
-            }
-
             if (!errors.isEmpty()) {
                 System.out.printf("Errors happened during scan of host %s%nErrors:%s -> %s", host, errors.size(), errors);
             }
@@ -233,5 +243,13 @@ public class Scanner implements AutoCloseable {
     @Override
     public void close() throws Exception {
         cancel();
+        awaitTermination(Duration.ofMillis(10000));
+    }
+
+    public void awaitTermination(Duration duration) throws InterruptedException, TimeoutException {
+        runningScans.awaitAdvanceInterruptibly(
+                runningScans.getPhase(),
+                duration.toMillis(),
+                TimeUnit.MILLISECONDS);
     }
 }