Same producer pattern for ports as for hosts

1 files changed, 44 insertions, 30 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index 9433841..bfb7ac6 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -138,45 +138,59 @@ public class Scanner implements AutoCloseable {
 
         CompletionService<PortResult> completionService = new ExecutorCompletionService<>(executor);
         List<Throwable> errors = new ArrayList<>();
-        AtomicLong portSlotFactory = new AtomicLong(System.nanoTime());
-        int activePerHostWorkers = 0;
-        while ((portRange.hasNext() && !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
-            ) {
-                // Here we are filling the completion service
-                // port <-> virtual thread
-                final int currentPort = portRange.next();
-                activePerHostWorkers++;
-                completionService.submit(() -> {
-                    waitForSlot(portSlotFactory);
-                    return checkPort(host, currentPort, scan);
-                });
+        final AtomicLong portSlotFactory = new AtomicLong(System.nanoTime());
+        final Semaphore activePerHostWorkers = new Semaphore(maxWorkersPerHost);
+        final AtomicInteger inPipeline = new AtomicInteger(0);
+
+        // producer thread
+        executor.submit(() -> {
+            try {
+                while (portRange.hasNext() && !cancelled) {
+                    activePerHostWorkers.acquire();
+                    // Here we are filling the completion service
+                    // port <-> virtual thread
+                    final int currentPort = portRange.next();
+                    inPipeline.incrementAndGet();
+                    completionService.submit(() -> {
+                        waitForSlot(portSlotFactory);
+                        return checkPort(host, currentPort, scan);
+                    });
+                }
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
             }
+        });
 
+        // consumer is the main thread
+        while ((portRange.hasNext() && !cancelled) || inPipeline.get() > 0) {
             try {
                 Future<PortResult> portResultFuture = completionService.poll(10, TimeUnit.MILLISECONDS);
                 if (portResultFuture == null) {
                     continue;
                 }
 
-                PortResult portResult = portResultFuture.get();
-                activePerHostWorkers--;
-                progress.done().incrementAndGet();
-                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());
-                        progress.filtered().incrementAndGet();
+                try {
+                    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());
+                            progress.filtered().incrementAndGet();
+                        }
+                        default -> {
+                            // sonarcube glücklich machen
+                        }
                     }
+                } finally {
+                    activePerHostWorkers.release();
+                    inPipeline.decrementAndGet();
+                    progress.done().incrementAndGet();
                 }
             } catch (InterruptedException ignored) {
                 Thread.currentThread().interrupt();