Cleanup

1 files changed, 9 insertions, 6 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index 192f81a..ae02d3d 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -70,7 +70,6 @@ public class Scanner implements AutoCloseable {
                 maxHostsLimit,
                 host -> () -> scanHostPorts(host, scan)
         );
-        CompletionService<ScanResult> completionService = state.completionService();
 
         // the main thread is the consumer
         // Let the consumer run as long as
@@ -80,7 +79,7 @@ public class Scanner implements AutoCloseable {
         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);
+                Future<ScanResult> finishedHost = state.completionService().poll(10, TimeUnit.MILLISECONDS);
                 if (finishedHost == null) {
                     continue;
                 }
@@ -107,6 +106,12 @@ public class Scanner implements AutoCloseable {
         return scan.getResults();
     }
 
+    /**
+     * Scans the ports of a given host
+     * @param host
+     * @param scan
+     * @return
+     */
     private ScanResult scanHostPorts(final String host, final Scan scan) {
         PortRange portRange = new PortRange(scan.getPorts());
         BitSet openPorts = new BitSet(PortRange.MAX_PORT);
@@ -134,13 +139,11 @@ public class Scanner implements AutoCloseable {
                     return checkPort(host, port, scan);
                 }
         );
-        CompletionService<PortResult> completionService = state.completionService();
-
 
         // consumer is the main thread
         while ((state.running().get() && !cancelled) || state.inPipeline().get() > 0) {
             try {
-                Future<PortResult> portResultFuture = completionService.poll(10, TimeUnit.MILLISECONDS);
+                Future<PortResult> portResultFuture = state.completionService().poll(10, TimeUnit.MILLISECONDS);
                 if (portResultFuture == null) {
                     continue;
                 }
@@ -194,7 +197,7 @@ public class Scanner implements AutoCloseable {
      * @param <OUTPUT>
      * @return
      */
-    private <INPUT, OUTPUT> ProducerState startProducer(Iterator<INPUT> queue, int maxWorkers, Function<INPUT, Callable<OUTPUT>> taskFactory) {
+    private <INPUT, OUTPUT> ProducerState<OUTPUT> 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);