forked from v2ray/v2ray-core
-
Notifications
You must be signed in to change notification settings - Fork 0
/
task.go
186 lines (159 loc) · 3.47 KB
/
task.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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
package task
import (
"context"
"strings"
"v2ray.com/core/common"
"v2ray.com/core/common/signal/semaphore"
)
type Task func() error
type MultiError []error
func (e MultiError) Error() string {
var r strings.Builder
common.Must2(r.WriteString("multierr: "))
for _, err := range e {
common.Must2(r.WriteString(err.Error()))
common.Must2(r.WriteString(" | "))
}
return r.String()
}
type executionContext struct {
ctx context.Context
tasks []Task
onSuccess Task
onFailure Task
}
func (c *executionContext) executeTask() error {
if len(c.tasks) == 0 {
return nil
}
// Reuse current goroutine if we only have one task to run.
if len(c.tasks) == 1 && c.ctx == nil {
return c.tasks[0]()
}
ctx := context.Background()
if c.ctx != nil {
ctx = c.ctx
}
return executeParallel(ctx, c.tasks)
}
func (c *executionContext) run() error {
err := c.executeTask()
if err == nil && c.onSuccess != nil {
return c.onSuccess()
}
if err != nil && c.onFailure != nil {
return c.onFailure()
}
return err
}
type ExecutionOption func(*executionContext)
func WithContext(ctx context.Context) ExecutionOption {
return func(c *executionContext) {
c.ctx = ctx
}
}
func Parallel(tasks ...Task) ExecutionOption {
return func(c *executionContext) {
c.tasks = append(c.tasks, tasks...)
}
}
// Sequential runs all tasks sequentially, and returns the first error encountered.Sequential
// Once a task returns an error, the following tasks will not run.
func Sequential(tasks ...Task) ExecutionOption {
return func(c *executionContext) {
switch len(tasks) {
case 0:
return
case 1:
c.tasks = append(c.tasks, tasks[0])
default:
c.tasks = append(c.tasks, func() error {
return execute(tasks...)
})
}
}
}
func SequentialAll(tasks ...Task) ExecutionOption {
return func(c *executionContext) {
switch len(tasks) {
case 0:
return
case 1:
c.tasks = append(c.tasks, tasks[0])
default:
c.tasks = append(c.tasks, func() error {
var merr MultiError
for _, task := range tasks {
if err := task(); err != nil {
merr = append(merr, err)
}
}
if len(merr) == 0 {
return nil
}
return merr
})
}
}
}
func OnSuccess(task Task) ExecutionOption {
return func(c *executionContext) {
c.onSuccess = task
}
}
func OnFailure(task Task) ExecutionOption {
return func(c *executionContext) {
c.onFailure = task
}
}
func Single(task Task, opts ...ExecutionOption) Task {
return Run(append([]ExecutionOption{Sequential(task)}, opts...)...)
}
func Run(opts ...ExecutionOption) Task {
var c executionContext
for _, opt := range opts {
opt(&c)
}
return func() error {
return c.run()
}
}
// execute runs a list of tasks sequentially, returns the first error encountered or nil if all tasks pass.
func execute(tasks ...Task) error {
for _, task := range tasks {
if err := task(); err != nil {
return err
}
}
return nil
}
// executeParallel executes a list of tasks asynchronously, returns the first error encountered or nil if all tasks pass.
func executeParallel(ctx context.Context, tasks []Task) error {
n := len(tasks)
s := semaphore.New(n)
done := make(chan error, 1)
for _, task := range tasks {
<-s.Wait()
go func(f Task) {
err := f()
if err == nil {
s.Signal()
return
}
select {
case done <- err:
default:
}
}(task)
}
for i := 0; i < n; i++ {
select {
case err := <-done:
return err
case <-ctx.Done():
return ctx.Err()
case <-s.Wait():
}
}
return nil
}