1 files changed, 54 insertions, 63 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index eac5dc7..192f81a 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -10,10 +10,8 @@ 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 {
 
@@ -67,7 +65,7 @@ public class Scanner implements AutoCloseable {
         scan.start();
 
         // start producer thread
-        final ProducerState state = startProducer(
+        final ProducerState<ScanResult> state = startProducer(
                 scan.getHosts().iterator(),
                 maxHostsLimit,
                 host -> () -> scanHostPorts(host, scan)
@@ -109,35 +107,6 @@ 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);
@@ -154,41 +123,22 @@ public class Scanner implements AutoCloseable {
             return new ScanResult(host, openPorts, filteredPorts, serviceTypes);
         }
 
-        CompletionService<PortResult> completionService = new ExecutorCompletionService<>(executor);
         List<Throwable> errors = new ArrayList<>();
         final AtomicLong portSlotFactory = new AtomicLong(System.nanoTime());
-        final Semaphore activePerHostWorkers = new Semaphore(maxWorkersPerHost);
-        final AtomicInteger inPipeline = new AtomicInteger(0);
-        final AtomicBoolean producerWorking = new AtomicBoolean(true);
-
         // 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();
-                    try {
-                        completionService.submit(() -> {
-                            waitForSlot(portSlotFactory);
-                            return checkPort(host, currentPort, scan);
-                        });
-                    } catch (RuntimeException e) {
-                        inPipeline.decrementAndGet();
-                        throw e;
-                    }
+        ProducerState<PortResult> state = startProducer(
+                portRange,
+                maxWorkersPerHost,
+                port -> () -> {
+                    waitForSlot(portSlotFactory);
+                    return checkPort(host, port, scan);
                 }
-            } catch (InterruptedException e) {
-                Thread.currentThread().interrupt();
-            } finally {
-                producerWorking.set(false);
-            }
-        });
+        );
+        CompletionService<PortResult> completionService = state.completionService();
+
 
         // consumer is the main thread
-        while ((producerWorking.get() && !cancelled) || inPipeline.get() > 0) {
+        while ((state.running().get() && !cancelled) || state.inPipeline().get() > 0) {
             try {
                 Future<PortResult> portResultFuture = completionService.poll(10, TimeUnit.MILLISECONDS);
                 if (portResultFuture == null) {
@@ -214,8 +164,8 @@ public class Scanner implements AutoCloseable {
                         }
                     }
                 } finally {
-                    activePerHostWorkers.release();
-                    inPipeline.decrementAndGet();
+                    state.activeWorkers().release();
+                    state.inPipeline().decrementAndGet();
                     progress.done().incrementAndGet();
                 }
             } catch (InterruptedException ignored) {
@@ -232,6 +182,47 @@ public class Scanner implements AutoCloseable {
         return new ScanResult(host, openPorts, filteredPorts, serviceTypes);
     }
 
+    /**
+     * This method helps to cleanup the code a bit and remove redundancy
+     * The producer for providing hosts and the one for providing ports
+     * are similar and the small differences can be handled using a function
+     *
+     * @param queue
+     * @param maxWorkers
+     * @param taskFactory
+     * @param <INPUT>
+     * @param <OUTPUT>
+     * @return
+     */
+    private <INPUT, OUTPUT> ProducerState startProducer(Iterator<INPUT> queue, int maxWorkers, Function<INPUT, Callable<OUTPUT>> taskFactory) {
+        final AtomicInteger inPipeline = new AtomicInteger(0);
+        final AtomicBoolean running = new AtomicBoolean(true);
+        final Semaphore activeWorkers = new Semaphore(maxWorkers);
+        CompletionService<OUTPUT> 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 INPUT 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 boolean checkHostOnline(String host) {
         try {
             return InetAddress.getByName(host).isReachable(timeoutInMillis);