1 files changed, 28 insertions, 16 deletions

diff --git a/src/main/java/com/it_jaros/jns/scan/engine/ScanHostTask.java b/src/main/java/com/it_jaros/jns/scan/engine/ScanHostTask.java
index 8441172..afe0e8e 100644
--- a/src/main/java/com/it_jaros/jns/scan/engine/ScanHostTask.java
+++ b/src/main/java/com/it_jaros/jns/scan/engine/ScanHostTask.java
@@ -75,24 +75,25 @@ public class ScanHostTask implements Callable<ScanResult> {
                 port -> new ScanPortTask(hostAddress, port, context, hostContext)
         );
 
-        // consumer is the main thread
-        // we run as long as the producer is running OR
-        // as long as things are in pipeline waiting to be processed
-        // ONLY exception is when cancelled is set
+        runConsumerLoop(state);
+        return accumulator.build();
+    }
+
+    /**
+     * consumer is the main thread
+     * we run as long as the producer is running OR
+     * as long as things are in pipeline waiting to be processed
+     * ONLY exception is when cancelled is set
+     */
+    private void runConsumerLoop(ProducerState<PortResult> state) {
         while (!cancelledToken.isCancelled() && (state.running().get() || state.inPipeline().get() > 0)) {
-            try {
-                PollState<PortResult> poll = getPortResult(state);
-                if (poll instanceof PollState.Success<PortResult>(PortResult value)) {
-                    accumulator.add(value);
-                } else if (poll instanceof PollState.Failure(Throwable error)) {
-                    accumulator.recordFailure(error, -1);
-                }
-            } catch (InterruptedException ignored) {
-                Thread.currentThread().interrupt();
+            PollState<PortResult> poll = getPortResult(state);
+            if (poll instanceof PollState.Success<PortResult>(PortResult value)) {
+                accumulator.add(value);
+            } else if (poll instanceof PollState.Failure(Throwable error)) {
+                accumulator.recordFailure(error, -1);
             }
         }
-
-        return accumulator.build();
     }
 
     private boolean checkHostOnline() {
@@ -109,8 +110,19 @@ public class ScanHostTask implements Callable<ScanResult> {
         return false;
     }
 
-    private PollState<PortResult> getPortResult(ProducerState<PortResult> state) throws InterruptedException {
+    private PollState<PortResult> getPortResult(ProducerState<PortResult> state) {
+        try {
+            return pollPortResult(state);
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+        }
+
+        return new PollState.Unavailable<>();
+    }
+
+    private static PollState<PortResult> pollPortResult(ProducerState<PortResult> state) throws InterruptedException {
         Future<PortResult> portResultFuture = state.completionService().poll(ProducerThread.pollInterval.toMillis(), TimeUnit.MILLISECONDS);
+
         if (portResultFuture == null) {
             return new PollState.Unavailable<>();
         }