Replaced thread unsafe usage of hasNext()

1 files changed, 13 insertions, 4 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index 011a28c..eaf2cc9 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;
@@ -64,6 +65,9 @@ public class Scanner implements AutoCloseable {
         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);
 
@@ -81,15 +85,17 @@ public class Scanner implements AutoCloseable {
                 }
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
+            } finally {
+                producerWorking.set(false);
             }
         });
 
         // 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 
+        // 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 ((queue.hasNext() && !cancelled) || inPipeline.get() > 0) {
+        while ((producerWorking.get() && !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);
@@ -119,7 +125,7 @@ public class Scanner implements AutoCloseable {
         return scan.getResults();
     }
 
-    private ScanResult scanHostPorts(String host, String ports, Scan scan) {
+    private ScanResult scanHostPorts(final String host, final String ports, final Scan scan) {
         PortRange portRange = new PortRange(ports);
         BitSet openPorts = new BitSet(PortRange.MAX_PORT);
         BitSet filteredPorts = new BitSet(PortRange.MAX_PORT);
@@ -140,6 +146,7 @@ public class Scanner implements AutoCloseable {
         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(() -> {
@@ -157,11 +164,13 @@ public class Scanner implements AutoCloseable {
                 }
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
+            } finally {
+                producerWorking.set(false);
             }
         });
 
         // consumer is the main thread
-        while ((portRange.hasNext() && !cancelled) || inPipeline.get() > 0) {
+        while ((producerWorking.get() && !cancelled) || inPipeline.get() > 0) {
             try {
                 Future<PortResult> portResultFuture = completionService.poll(10, TimeUnit.MILLISECONDS);
                 if (portResultFuture == null) {