Moved producer thread into custom method which encapsulates logic for

that producer -> reduced complexity of runScan method

2 files changed, 65 insertions, 38 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/ProducerState.java b/src/main/java/com/it_jaros/jscanner/ProducerState.java
new file mode 100644
index 0000000..36a5841
--- /dev/null
+++ b/src/main/java/com/it_jaros/jscanner/ProducerState.java
@@ -0,0 +1,13 @@
+package com.it_jaros.jscanner;
+
+import java.util.concurrent.CompletionService;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public record ProducerState<OUTPUT>(
+        AtomicBoolean running,
+        AtomicInteger inPipeline,
+        Semaphore activeWorkers,
+        CompletionService<OUTPUT> completionService
+) {}
diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index 7e52775..eac5dc7 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -10,7 +10,10 @@ 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.atomic.AtomicReference;
 import java.util.concurrent.locks.LockSupport;
+import java.util.function.Function;
+import java.util.stream.Stream;
 
 public class Scanner implements AutoCloseable {
 
@@ -56,50 +59,27 @@ public class Scanner implements AutoCloseable {
      * @param scan
      * @return
      */
-    public List<ScanResult> runScan(Scan scan) {
+    public List<ScanResult> runScan(final Scan scan) {
         if (scan == null) {
             throw new IllegalArgumentException("Scan argument cannot be null");
         }
 
         scan.start();
-        CompletionService<ScanResult> completionService = new ExecutorCompletionService<>(executor);
-
-        Iterator<String> queue = scan.getHosts().iterator();
-        // since queue.hasNext() is not thread safe, we need atomic boolean to check
-        // if the producer is still there
-        final AtomicBoolean producerWorking = new AtomicBoolean(true);
-        final Semaphore activeHostWorkers = new Semaphore(maxHostsLimit);
-        final AtomicInteger inPipeline = new AtomicInteger(0);
 
         // start producer thread
-        executor.submit(() -> {
-            try {
-                while (queue.hasNext() && !cancelled) {
-                    activeHostWorkers.acquire();
-                    // just in case while waiting something changed
-                    if (cancelled) {
-                        break;
-                    }
-
-                    // here we are filling the completion service
-                    // host <-> virtual thread
-                    final String host = queue.next();
-                    inPipeline.incrementAndGet();
-                    completionService.submit(() -> scanHostPorts(host, scan));
-                }
-            } catch (InterruptedException e) {
-                Thread.currentThread().interrupt();
-            } finally {
-                producerWorking.set(false);
-            }
-        });
+        final ProducerState state = startProducer(
+                scan.getHosts().iterator(),
+                maxHostsLimit,
+                host -> () -> scanHostPorts(host, scan)
+        );
+        CompletionService<ScanResult> completionService = state.completionService();
 
         // 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 ((producerWorking.get() && !cancelled) || inPipeline.get() > 0) {
+        while ((state.running().get() && !cancelled) || state.inPipeline().get() > 0) {
             try {
                 // let's check for results and give add them to our scan data holder object
                 Future<ScanResult> finishedHost = completionService.poll(10, TimeUnit.MILLISECONDS);
@@ -112,11 +92,11 @@ public class Scanner implements AutoCloseable {
                     ScanResult result = finishedHost.get();
                     // we only want results that give value and not cost RAM for nothing
                     if (!result.openPorts().isEmpty() || !result.filteredPorts().isEmpty()) {
-                        scan.addScanResult(finishedHost.get());
+                        scan.addScanResult(result);
                     }
                 } finally {
-                    activeHostWorkers.release();
-                    inPipeline.decrementAndGet();
+                    state.activeWorkers().release();
+                    state.inPipeline().decrementAndGet();
                 }
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
@@ -129,6 +109,35 @@ public class Scanner implements AutoCloseable {
         return scan.getResults();
     }
 
+    private ProducerState startProducer(Iterator<String> queue, int maxWorkers, Function<String, Callable<ScanResult>> taskFactory) {
+        final AtomicInteger inPipeline = new AtomicInteger(0);
+        final AtomicBoolean running = new AtomicBoolean(true);
+        final Semaphore activeWorkers = new Semaphore(maxWorkers);
+        CompletionService<ScanResult> completionService = new ExecutorCompletionService<>(executor);
+        executor.submit(() -> {
+            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 String item = queue.next();
+                    inPipeline.incrementAndGet();
+                    completionService.submit(taskFactory.apply(item));
+                }
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+            } finally {
+                running.set(false);
+            }
+        });
+        return new ProducerState(running, inPipeline, activeWorkers, completionService);
+    }
+
     private ScanResult scanHostPorts(final String host, final Scan scan) {
         PortRange portRange = new PortRange(scan.getPorts());
         BitSet openPorts = new BitSet(PortRange.MAX_PORT);
@@ -161,10 +170,15 @@ public class Scanner implements AutoCloseable {
                     // port <-> virtual thread
                     final int currentPort = portRange.next();
                     inPipeline.incrementAndGet();
-                    completionService.submit(() -> {
-                        waitForSlot(portSlotFactory);
-                        return checkPort(host, currentPort, scan);
-                    });
+                    try {
+                        completionService.submit(() -> {
+                            waitForSlot(portSlotFactory);
+                            return checkPort(host, currentPort, scan);
+                        });
+                    } catch (RuntimeException e) {
+                        inPipeline.decrementAndGet();
+                        throw e;
+                    }
                 }
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();