forked from zeromicro/go-zero
-
Notifications
You must be signed in to change notification settings - Fork 0
/
chunkexecutor.go
113 lines (94 loc) · 2.38 KB
/
chunkexecutor.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
102
103
104
105
106
107
108
109
110
111
112
113
package executors
import "time"
const defaultChunkSize = 1024 * 1024 // 1M
type (
// ChunkOption defines the method to customize a ChunkExecutor.
ChunkOption func(options *chunkOptions)
// A ChunkExecutor is an executor to execute tasks when either requirement meets:
// 1. up to given chunk size
// 2. flush interval elapsed
ChunkExecutor struct {
executor *PeriodicalExecutor
container *chunkContainer
}
chunkOptions struct {
chunkSize int
flushInterval time.Duration
}
)
// NewChunkExecutor returns a ChunkExecutor.
func NewChunkExecutor(execute Execute, opts ...ChunkOption) *ChunkExecutor {
options := newChunkOptions()
for _, opt := range opts {
opt(&options)
}
container := &chunkContainer{
execute: execute,
maxChunkSize: options.chunkSize,
}
executor := &ChunkExecutor{
executor: NewPeriodicalExecutor(options.flushInterval, container),
container: container,
}
return executor
}
// Add adds task with given chunk size into ce.
func (ce *ChunkExecutor) Add(task interface{}, size int) error {
ce.executor.Add(chunk{
val: task,
size: size,
})
return nil
}
// Flush forces ce to flush and execute tasks.
func (ce *ChunkExecutor) Flush() {
ce.executor.Flush()
}
// Wait waits the execution to be done.
func (ce *ChunkExecutor) Wait() {
ce.executor.Wait()
}
// WithChunkBytes customizes a ChunkExecutor with the given chunk size.
func WithChunkBytes(size int) ChunkOption {
return func(options *chunkOptions) {
options.chunkSize = size
}
}
// WithFlushInterval customizes a ChunkExecutor with the given flush interval.
func WithFlushInterval(duration time.Duration) ChunkOption {
return func(options *chunkOptions) {
options.flushInterval = duration
}
}
func newChunkOptions() chunkOptions {
return chunkOptions{
chunkSize: defaultChunkSize,
flushInterval: defaultFlushInterval,
}
}
type chunkContainer struct {
tasks []interface{}
execute Execute
size int
maxChunkSize int
}
func (bc *chunkContainer) AddTask(task interface{}) bool {
ck := task.(chunk)
bc.tasks = append(bc.tasks, ck.val)
bc.size += ck.size
return bc.size >= bc.maxChunkSize
}
func (bc *chunkContainer) Execute(tasks interface{}) {
vals := tasks.([]interface{})
bc.execute(vals)
}
func (bc *chunkContainer) RemoveAll() interface{} {
tasks := bc.tasks
bc.tasks = nil
bc.size = 0
return tasks
}
type chunk struct {
val interface{}
size int
}