-
Notifications
You must be signed in to change notification settings - Fork 17
/
worker.go
84 lines (68 loc) · 1.5 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
package consumer
import (
"context"
"sync"
"github.com/makasim/amqpextra/logger"
amqp "github.com/rabbitmq/amqp091-go"
)
type Worker interface {
Serve(ctx context.Context, h Handler, msgCh <-chan amqp.Delivery)
}
type DefaultWorker struct {
Logger logger.Logger
}
func (dw *DefaultWorker) Serve(ctx context.Context, h Handler, msgCh <-chan amqp.Delivery) {
defer dw.Logger.Printf("[DEBUG] worker stopped")
dw.Logger.Printf("[DEBUG] worker starting")
for {
select {
case msg, ok := <-msgCh:
if !ok {
return
}
if res := h.Handle(ctx, msg); res != nil {
dw.Logger.Printf("[ERROR] handler return non nil result: %#v", res)
}
case <-ctx.Done():
return
}
}
}
type ParallelWorker struct {
Num int
Logger logger.Logger
}
func NewParallelWorker(num int) *ParallelWorker {
if num < 1 {
panic("num workers must be greater than zero")
}
return &ParallelWorker{
Num: num,
Logger: logger.Discard,
}
}
func (pw *ParallelWorker) Serve(ctx context.Context, h Handler, msgCh <-chan amqp.Delivery) {
defer pw.Logger.Printf("[DEBUG] worker stopped")
pw.Logger.Printf("[DEBUG] worker starting")
wg := &sync.WaitGroup{}
for i := 0; i < pw.Num; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for {
select {
case msg, ok := <-msgCh:
if !ok {
return
}
if res := h.Handle(ctx, msg); res != nil {
pw.Logger.Printf("[ERROR] handler return non nil result: %#v", res)
}
case <-ctx.Done():
return
}
}
}()
}
wg.Wait()
}