-
Notifications
You must be signed in to change notification settings - Fork 0
/
worker.go
203 lines (162 loc) · 3.29 KB
/
worker.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
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
package do
import (
"context"
"errors"
"log"
"os"
"sync"
"time"
"github.com/smallnest/chanx"
)
const (
errCount = 100
defaultCount = 50
)
var logger = log.New(os.Stdout, "[Worker]", log.LstdFlags|log.Lshortfile)
// 错误
var (
ErrWorkerIsStop = errors.New("Worker is stop")
ErrNilJobDo = errors.New("Job do field is nil")
)
// DefaultWorker 默认Wroker
var DefaultWorker = NewWorker(defaultCount)
func init() {
DefaultWorker.Start()
}
// Job 工作
type Job struct {
doctx DoWithCtx
timeout time.Duration // 超时时间
errorHandler ErrorHandler // 错误处理方法
}
type DoWithCtx func(ctx context.Context) error
// ErrorHandler 错误处理方法
type ErrorHandler func(error)
func NewJob(do DoWithCtx, timeout time.Duration, eh ErrorHandler) *Job {
return &Job{
doctx: do,
timeout: timeout,
errorHandler: eh,
}
}
func (job *Job) run(ctx context.Context) error {
if job.doctx == nil {
return ErrNilJobDo
}
if err := job.doctx(ctx); err != nil {
if job.errorHandler != nil {
job.errorHandler(err)
} else {
return err
}
}
return nil
}
// Worker 工人
type Worker struct {
// 所有管道都要有make, read, write, close操作
limitChan chan struct{} // 并发控制管道
stopChan chan struct{} // 停止管道
jobChan *chanx.UnboundedChan[Job] // 工作管道
errChan chan error // 错误管道
wg *sync.WaitGroup
stop bool // 是否调用了Stop方法
}
// NewWorker new a worker with limit number
func NewWorker(n int) *Worker {
if n <= 0 {
n = defaultCount
}
return &Worker{
limitChan: make(chan struct{}, n),
stopChan: make(chan struct{}),
jobChan: chanx.NewUnboundedChan[Job](n),
errChan: make(chan error, errCount),
wg: new(sync.WaitGroup),
}
}
// Start 开始
func (w *Worker) Start() {
go w.handleError()
go w.start()
logger.Printf("Start.\n")
}
func (w *Worker) start() {
for {
select {
case job, ok := <-w.jobChan.Out: // 有工作
if !ok {
continue
}
w.do(job)
case <-w.stopChan:
w.close()
return
}
}
}
func (w *Worker) do(job Job) {
// 占据一个坑
w.limitChan <- struct{}{}
// 开始工作
go func(job Job) {
defer func() {
if r := recover(); r != nil {
logger.Printf("job: %+v\n", r)
}
// 释放一个坑
<-w.limitChan
w.wg.Done()
}()
// 执行
ctx := context.Background()
if job.timeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, job.timeout)
defer cancel()
}
if err := job.run(ctx); err != nil {
w.errChan <- err
}
}(job)
}
func (w *Worker) handleError() {
for err := range w.errChan {
logger.Printf("err is %v\n", err)
}
}
// Stop 停止
func (w *Worker) Stop() {
w.stop = true
w.wait()
w.stopChan <- struct{}{}
logger.Printf("Stop.\n")
}
func (w *Worker) close() {
// close管道时,有可能panic
defer func() {
if r := recover(); r != nil {
logger.Printf("close: %+v\n", r)
}
}()
close(w.stopChan)
close(w.errChan)
close(w.jobChan.In)
close(w.limitChan)
}
func (w *Worker) wait() {
// 等待所有工作完成
w.wg.Wait()
}
// Push 添加
func (w *Worker) Push(job Job) error {
if w.stop {
return ErrWorkerIsStop
}
if job.doctx == nil {
return ErrNilJobDo
}
w.jobChan.In <- job
w.wg.Add(1)
return nil
}