Replaced nested while loops with producer thread and consumer stays in

main thread

1 files changed, 37 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 7ff838f..9433841 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -7,6 +7,7 @@ import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.*;
 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;
@@ -51,6 +52,7 @@ public class Scanner implements AutoCloseable {
 
     /**
      * Starts a given scan.
+     *
      * @param scan
      * @return
      */
@@ -63,31 +65,49 @@ public class Scanner implements AutoCloseable {
         CompletionService<ScanResult> completionService = new ExecutorCompletionService<>(executor);
 
         Iterator<String> queue = scan.getHosts().iterator();
-        int activeHostWorkers = 0;
-        while ((queue.hasNext() && !cancelled) || activeHostWorkers > 0) {
-            while (
-                    queue.hasNext() // continue as long as we have hosts to process
-                            && activeHostWorkers < maxHostsLimit // limit is not reached
-                            && !cancelled // and scanner not closed
-            ) {
-                // here we are filling the completion service
-                // host <-> virtual thread
-                final String host = queue.next();
-                activeHostWorkers++;
-                completionService.submit(() -> scanHostPorts(host, scan.getPorts(), scan));
+        final Semaphore activeHostWorkers = new Semaphore(maxHostsLimit);
+        final AtomicInteger inPipeline = new AtomicInteger(0);
+
+        // start producer thread
+        executor.submit(() -> {
+            try {
+                while (queue.hasNext() && !cancelled) {
+                    activeHostWorkers.acquire();
+
+                    // here we are filling the completion service
+                    // host <-> virtual thread
+                    final String host = queue.next();
+                    inPipeline.incrementAndGet();
+                    completionService.submit(() -> scanHostPorts(host, scan.getPorts(), scan));
+                }
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
             }
+        });
 
+        // the main thread is the consumer
+        // Let the consumer run as long as
+        // if there is something in queue and the scanner is not cancelled 
+        //   -> 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 ((queue.hasNext() && !cancelled) || 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);
                 if (finishedHost == null) {
                     continue;
                 }
-                activeHostWorkers--;
-                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());
+
+                // No matter what happens we have to free the resources after getting ScanResult
+                try {
+                    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());
+                    }
+                } finally {
+                    activeHostWorkers.release();
+                    inPipeline.decrementAndGet();
                 }
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();