-
Notifications
You must be signed in to change notification settings - Fork 2.3k
/
workpool.go
91 lines (81 loc) · 2.89 KB
/
workpool.go
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
84
85
86
87
88
89
90
91
package core
import (
"context"
"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
syncutil "github.com/projectdiscovery/utils/sync"
)
// WorkPool implements an execution pool for executing different
// types of task with different concurrency requirements.
//
// It also allows Configuration of such requirements. This is used
// for per-module like separate headless concurrency etc.
type WorkPool struct {
Headless *syncutil.AdaptiveWaitGroup
Default *syncutil.AdaptiveWaitGroup
config WorkPoolConfig
}
// WorkPoolConfig is the configuration for work pool
type WorkPoolConfig struct {
// InputConcurrency is the concurrency for inputs values.
InputConcurrency int
// TypeConcurrency is the concurrency for the request type templates.
TypeConcurrency int
// HeadlessInputConcurrency is the concurrency for headless inputs values.
HeadlessInputConcurrency int
// TypeConcurrency is the concurrency for the headless request type templates.
HeadlessTypeConcurrency int
}
// NewWorkPool returns a new WorkPool instance
func NewWorkPool(config WorkPoolConfig) *WorkPool {
headlessWg, _ := syncutil.New(syncutil.WithSize(config.HeadlessTypeConcurrency))
defaultWg, _ := syncutil.New(syncutil.WithSize(config.TypeConcurrency))
return &WorkPool{
config: config,
Headless: headlessWg,
Default: defaultWg,
}
}
// Wait waits for all the work pool wait groups to finish
func (w *WorkPool) Wait() {
w.Default.Wait()
w.Headless.Wait()
}
// InputPool returns a work pool for an input type
func (w *WorkPool) InputPool(templateType types.ProtocolType) *syncutil.AdaptiveWaitGroup {
var count int
if templateType == types.HeadlessProtocol {
count = w.config.HeadlessInputConcurrency
} else {
count = w.config.InputConcurrency
}
swg, _ := syncutil.New(syncutil.WithSize(count))
return swg
}
func (w *WorkPool) RefreshWithConfig(config WorkPoolConfig) {
if w.config.TypeConcurrency != config.TypeConcurrency {
w.config.TypeConcurrency = config.TypeConcurrency
}
if w.config.HeadlessTypeConcurrency != config.HeadlessTypeConcurrency {
w.config.HeadlessTypeConcurrency = config.HeadlessTypeConcurrency
}
if w.config.InputConcurrency != config.InputConcurrency {
w.config.InputConcurrency = config.InputConcurrency
}
if w.config.HeadlessInputConcurrency != config.HeadlessInputConcurrency {
w.config.HeadlessInputConcurrency = config.HeadlessInputConcurrency
}
w.Refresh(context.Background())
}
func (w *WorkPool) Refresh(ctx context.Context) {
if w.Default.Size != w.config.TypeConcurrency {
if err := w.Default.Resize(ctx, w.config.TypeConcurrency); err != nil {
gologger.Warning().Msgf("Could not resize workpool: %s\n", err)
}
}
if w.Headless.Size != w.config.HeadlessTypeConcurrency {
if err := w.Headless.Resize(ctx, w.config.HeadlessTypeConcurrency); err != nil {
gologger.Warning().Msgf("Could not resize workpool: %s\n", err)
}
}
}