forked from zeromicro/go-zero
-
Notifications
You must be signed in to change notification settings - Fork 0
/
bulkexecutor.go
101 lines (83 loc) · 2.16 KB
/
bulkexecutor.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
92
93
94
95
96
97
98
99
100
101
package executors
import "time"
const defaultBulkTasks = 1000
type (
// BulkOption defines the method to customize a BulkExecutor.
BulkOption func(options *bulkOptions)
// A BulkExecutor is an executor that can execute tasks on either requirement meets:
// 1. up to given size of tasks
// 2. flush interval time elapsed
BulkExecutor struct {
executor *PeriodicalExecutor
container *bulkContainer
}
bulkOptions struct {
cachedTasks int
flushInterval time.Duration
}
)
// NewBulkExecutor returns a BulkExecutor.
func NewBulkExecutor(execute Execute, opts ...BulkOption) *BulkExecutor {
options := newBulkOptions()
for _, opt := range opts {
opt(&options)
}
container := &bulkContainer{
execute: execute,
maxTasks: options.cachedTasks,
}
executor := &BulkExecutor{
executor: NewPeriodicalExecutor(options.flushInterval, container),
container: container,
}
return executor
}
// Add adds task into be.
func (be *BulkExecutor) Add(task interface{}) error {
be.executor.Add(task)
return nil
}
// Flush forces be to flush and execute tasks.
func (be *BulkExecutor) Flush() {
be.executor.Flush()
}
// Wait waits be to done with the task execution.
func (be *BulkExecutor) Wait() {
be.executor.Wait()
}
// WithBulkTasks customizes a BulkExecutor with given tasks limit.
func WithBulkTasks(tasks int) BulkOption {
return func(options *bulkOptions) {
options.cachedTasks = tasks
}
}
// WithBulkInterval customizes a BulkExecutor with given flush interval.
func WithBulkInterval(duration time.Duration) BulkOption {
return func(options *bulkOptions) {
options.flushInterval = duration
}
}
func newBulkOptions() bulkOptions {
return bulkOptions{
cachedTasks: defaultBulkTasks,
flushInterval: defaultFlushInterval,
}
}
type bulkContainer struct {
tasks []interface{}
execute Execute
maxTasks int
}
func (bc *bulkContainer) AddTask(task interface{}) bool {
bc.tasks = append(bc.tasks, task)
return len(bc.tasks) >= bc.maxTasks
}
func (bc *bulkContainer) Execute(tasks interface{}) {
vals := tasks.([]interface{})
bc.execute(vals)
}
func (bc *bulkContainer) RemoveAll() interface{} {
tasks := bc.tasks
bc.tasks = nil
return tasks
}