-
Notifications
You must be signed in to change notification settings - Fork 5
/
util.go
165 lines (145 loc) · 3.37 KB
/
util.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
package tos
import (
"fmt"
"os"
"time"
)
func min(a int, b int) int {
if a < b {
return a
}
return b
}
const (
EventPartSucceed = 3
EventPartFailed = 4
EventPartAborted = 5 // The task needs to be interrupted in case of 403, 404, 405 errors
)
type task interface {
do() (interface{}, error)
getBaseInput() interface{}
}
type checkPoint interface {
WriteToFile() error
UpdatePartsInfo(result interface{})
GetCheckPointFilePath() string
}
type taskGroup interface {
// Wait 等待执行结果, success 是此次成功的 task 数量
Wait() (success int, err error)
// RunWorker 启动worker
RunWorker()
// Scheduler 分发任务
Scheduler()
}
type postEvent interface {
PostEvent(eventType int, result interface{}, taskErr error)
}
type taskGroupImpl struct {
cancelHandle chan struct{}
abortHandle chan struct{}
errCh chan error
resultsCh chan interface{}
tasksCh chan task
routinesNum int
tasks []task
checkPoint checkPoint
enableCheckPoint bool
postEvent postEvent
}
func (t *taskGroupImpl) Wait() (int, error) {
successNum := 0
failNum := 0
Loop:
for successNum+failNum < len(t.tasks) {
select {
case <-t.abortHandle:
break Loop
case <-t.cancelHandle:
break Loop
case part := <-t.resultsCh:
successNum++
t.checkPoint.UpdatePartsInfo(part)
if t.enableCheckPoint {
t.checkPoint.WriteToFile()
}
t.postEvent.PostEvent(EventPartSucceed, part, nil)
case taskErr := <-t.errCh:
if StatusCode(taskErr) == 403 || StatusCode(taskErr) == 404 || StatusCode(taskErr) == 405 {
close(t.abortHandle)
_ = os.Remove(t.checkPoint.GetCheckPointFilePath())
t.postEvent.PostEvent(EventPartAborted, nil, taskErr)
return successNum, fmt.Errorf("status code not service error, err:%s. ", taskErr.Error())
}
t.postEvent.PostEvent(EventPartFailed, nil, taskErr)
failNum++
}
}
return successNum, nil
}
func newTaskGroup(cancelHandle chan struct{}, routinesNum int, checkPoint checkPoint, postEvent postEvent, enableCheckPoint bool, tasks []task) taskGroup {
taskBufferSize := min(routinesNum, DefaultTaskBufferSize)
tasksCh := make(chan task, taskBufferSize)
return &taskGroupImpl{
cancelHandle: cancelHandle,
abortHandle: make(chan struct{}),
errCh: make(chan error),
resultsCh: make(chan interface{}),
tasksCh: tasksCh,
routinesNum: routinesNum,
tasks: tasks,
checkPoint: checkPoint,
enableCheckPoint: enableCheckPoint,
postEvent: postEvent,
}
}
func (t *taskGroupImpl) RunWorker() {
for i := 0; i < t.routinesNum; i++ {
go t.worker()
}
}
func (t *taskGroupImpl) Scheduler() {
go func() {
for _, task := range t.tasks {
select {
case <-t.cancelHandle:
return
case <-t.abortHandle:
return
default:
t.tasksCh <- task
}
}
close(t.tasksCh)
}()
}
func (t *taskGroupImpl) worker() {
for {
select {
case <-t.cancelHandle:
return
case <-t.abortHandle:
return
case task, ok := <-t.tasksCh:
if !ok {
return
}
result, err := task.do()
if err != nil {
t.errCh <- err
}
if result != nil {
t.resultsCh <- result
}
}
}
}
func GetUnixTimeMs() int64 {
return ToMillis(time.Now())
}
func ToMillis(t time.Time) int64 {
return t.UnixNano() / int64(time.Millisecond)
}
func StringPtr(input string) *string {
return &input
}