blob: e7bcbcc1ddc86f09f7fb71e39ea94d8e99ae65d6 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
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);
}
}
|
