1 files changed, 83 insertions, 0 deletions
diff --git a/src/main/java/com/it_jaros/jscanner/scan/engine/ProducerThread.java b/src/main/java/com/it_jaros/jscanner/scan/engine/ProducerThread.java
new file mode 100644
index 0000000..e7bcbcc
--- /dev/null
+++ b/src/main/java/com/it_jaros/jscanner/scan/engine/ProducerThread.java
@@ -0,0 +1,83 @@
+package com.it_jaros.jscanner.scan.engine;
+
+import java.time.Duration;
+import java.util.Iterator;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+
+public class ProducerThread {
+
+    public static final Duration pollInterval = Duration.ofSeconds(1);
+
+    private final ExecutorService executor;
+    private final CancelledToken cancelledToken;
+
+    public ProducerThread(ScanExecutionContext context) {
+        this.executor = context.executorService();
+        this.cancelledToken = context.cancelledToken();
+    }
+
+    /**
+     * This method helps to cleanup the code a bit and remove redundancy
+     * The producer for providing hosts and the one for providing ports
+     * are similar and the small differences can be handled using a function
+     *
+     * @param queue
+     * @param maxWorkers
+     * @param taskFactory
+     * @param <INPUT>
+     * @param <OUTPUT>
+     * @return
+     */
+    public <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);
+        CompletionService<OUTPUT> completionService = new ExecutorCompletionService<>(executor);
+        executor.submit(() -> {
+            try {
+                while (!cancelledToken.isCancelled() && queue.hasNext()) {
+                    // get semaphore and remember if task got submitted
+                    // so in case we fail to submit we release the semaphore
+                    activeWorkers.acquire();
+                    boolean isTaskSubmitted = false;
+                    try {
+                        // just in case something
+                        // changed while waiting
+                        if (cancelledToken.isCancelled()) {
+                            break;
+                        }
+
+                        // get next item and create callable
+                        // using lambda expression
+                        final INPUT item = queue.next();
+                        Callable<OUTPUT> task = taskFactory.apply(item);
+                        inPipeline.incrementAndGet();
+                        try {
+                            completionService.submit(task);
+                            isTaskSubmitted = true;
+                        } catch (Throwable e) {
+                            inPipeline.decrementAndGet();
+                            throw e;
+                        }
+                    } finally {
+                        if (!isTaskSubmitted) {
+                            activeWorkers.release();
+                        }
+                    }
+                }
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+            } finally {
+                running.set(false);
+            }
+        });
+        return new ProducerState<>(running, inPipeline, activeWorkers, completionService);
+    }
+}