2 files changed, 51 insertions, 17 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/Scan.java b/src/main/java/com/it_jaros/jscanner/Scan.java
index 3bd62bd..8321783 100644
--- a/src/main/java/com/it_jaros/jscanner/Scan.java
+++ b/src/main/java/com/it_jaros/jscanner/Scan.java
@@ -17,6 +17,7 @@ public class Scan {
     private final Runnable onDone;
     private final Counter hostCounter;
     private final Counter hostTotalCounter;
+    private final Counter portTotalCounter;
     private final Counter socketCounter;
     private final Counter threadCounter;
     private long start = 0;
@@ -32,6 +33,7 @@ public class Scan {
         this.onDone = onDone;
         this.hostCounter = new Counter();
         this.hostTotalCounter = new Counter();
+        this.portTotalCounter = new Counter();
         this.socketCounter = new Counter();
         this.threadCounter = new Counter();
     }
@@ -123,4 +125,40 @@ public class Scan {
         }
         return stop - start;
     }
+
+    public void portStart() {
+        this.portTotalCounter.inc();
+        this.threadCounter.inc();
+    }
+
+    public void portFinish() {
+        this.threadCounter.dec();
+    }
+
+    public void hostStart() {
+        this.hostCounter.inc();
+        this.hostTotalCounter.inc();
+        this.threadCounter.inc();
+    }
+
+    public void hostFinish() {
+        this.threadCounter.dec();
+        this.hostCounter.dec();
+    }
+
+    public void socketStart() {
+        this.socketCounter.inc();
+    }
+
+    public void socketFinish() {
+        this.socketCounter.dec();
+    }
+
+    public void producerStart() {
+        this.threadCounter.inc();
+    }
+
+    public void producerStop() {
+        this.threadCounter.dec();
+    }
 }
diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index f5f12a7..664723f 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -62,11 +62,11 @@ public class Scanner implements AutoCloseable {
         scan.start();
 
         // start producer thread
+        scan.producerStart();
         final ProducerState<ScanResult> state = startProducer(
                 scan.getHosts().iterator(),
                 maxHostsLimit,
-                host -> new ScanHostTask(scan, host),
-                scan.getThreadCounter()
+                host -> new ScanHostTask(scan, host)
         );
 
         // the main thread is the consumer
@@ -87,6 +87,7 @@ public class Scanner implements AutoCloseable {
                 Thread.currentThread().interrupt();
             }
         }
+        scan.producerStop();
 
         scan.stop();
     }
@@ -123,8 +124,7 @@ public class Scanner implements AutoCloseable {
     private <INPUT, OUTPUT> ProducerState<OUTPUT> startProducer(
             Iterator<INPUT> queue,
             int maxWorkers,
-            Function<INPUT, Callable<OUTPUT>> taskFactory,
-            Counter threadCounter
+            Function<INPUT, Callable<OUTPUT>> taskFactory
     ) {
         final AtomicInteger inPipeline = new AtomicInteger(0);
         final AtomicBoolean running = new AtomicBoolean(true);
@@ -132,7 +132,6 @@ public class Scanner implements AutoCloseable {
         CompletionService<OUTPUT> completionService = new ExecutorCompletionService<>(executor);
         executor.submit(() -> {
             try {
-                threadCounter.inc();
                 while (!cancelled && queue.hasNext()) {
                     // get semaphore and remember if task got submitted
                     // so in case we fail to submit we release the semaphore
@@ -166,7 +165,6 @@ public class Scanner implements AutoCloseable {
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
             } finally {
-                threadCounter.dec();
                 running.set(false);
             }
         });
@@ -211,16 +209,13 @@ public class Scanner implements AutoCloseable {
         @Override
         public ScanResult call() {
             try {
-                scan.getThreadCounter().inc();
-                scan.getHostCounter().inc();
-                scan.getHostTotalCounter().inc();
+                scan.hostStart();
                 if (cancelled) {
                     return ScanResult.empty(host);
                 }
                 return scanHostPorts();
             } finally {
-                scan.getThreadCounter().dec();
-                scan.getHostCounter().dec();
+                scan.hostFinish();
             }
         }
 
@@ -232,11 +227,11 @@ public class Scanner implements AutoCloseable {
 
             final PortRange portRange = new PortRange(scan.getPorts());
             // producer thread
+            scan.producerStart();
             ProducerState<PortResult> state = startProducer(
                     portRange.iterator(),
                     maxWorkersPerHost,
-                    port -> new ScanPortTask(host, port, scan, rateLimiter),
-                    scan.getThreadCounter()
+                    port -> new ScanPortTask(host, port, scan, rateLimiter)
             );
 
             // consumer is the main thread
@@ -255,6 +250,7 @@ public class Scanner implements AutoCloseable {
                     Thread.currentThread().interrupt();
                 }
             }
+            scan.producerStop();
 
             return accumulator.build();
         }
@@ -333,14 +329,14 @@ public class Scanner implements AutoCloseable {
         @Override
         public PortResult call() throws Exception {
             try {
-                scan.getThreadCounter().inc();
+                scan.portStart();
                 portScanRateLimiter.apply();
                 if (cancelled) {
                     return PortResult.empty();
                 }
                 return checkPort();
             } finally {
-                scan.getThreadCounter().dec();
+                scan.portFinish();
             }
         }
 
@@ -350,7 +346,7 @@ public class Scanner implements AutoCloseable {
             result.setState(PortState.UNKNOWN);
 
             socketLimit.acquire();
-            scan.getSocketCounter().inc();
+            scan.socketStart();
             try (Socket socket = new Socket()) {
                 socket.connect(new InetSocketAddress(host, port), timeoutInMillis);
                 result.setState(PortState.OPEN);
@@ -367,7 +363,7 @@ public class Scanner implements AutoCloseable {
             } catch (IOException e) {
                 result.setException(e);
             } finally {
-                scan.getSocketCounter().dec();
+                scan.socketFinish();
                 socketLimit.release();
             }