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
|
package worker
import (
"context"
"errors"
"sync"
"github.com/sirupsen/logrus"
)
type (
// A simple worker to deal with list.
ChanProcessor[T any] interface {
Query(context.Context) (<-chan T, error)
Process(context.Context, T) error
}
BatchProcessor[T any] interface {
Query(context.Context) ([]T, error)
Process(context.Context, T) error
}
chanProcessorWorker[T any] struct {
chanProcessor ChanProcessor[T]
logrus *logrus.Entry
scheduler *Scheduler
}
batchProcessorWorker[T any] struct {
batchProcessor BatchProcessor[T]
logrus *logrus.Entry
scheduler *Scheduler
}
)
func NewWorkerFromBatchProcessor[T any](
batchProcessor BatchProcessor[T],
scheduler *Scheduler,
logrus *logrus.Entry,
) Worker {
return &batchProcessorWorker[T]{
batchProcessor: batchProcessor,
scheduler: scheduler,
logrus: logrus,
}
}
func NewWorkerFromChanProcessor[T any](
chanProcessor ChanProcessor[T],
scheduler *Scheduler,
logrus *logrus.Entry,
) Worker {
return &chanProcessorWorker[T]{
chanProcessor: chanProcessor,
scheduler: scheduler,
logrus: logrus,
}
}
func (l *batchProcessorWorker[T]) Start(ctx context.Context) error {
for {
values, err := l.batchProcessor.Query(ctx)
if err != nil {
return err
}
select {
case <-ctx.Done():
return ctx.Err()
default:
}
if len(values) == 0 {
return nil
}
var wg sync.WaitGroup
for _, v := range values {
wg.Add(1)
l.scheduler.Take()
go func(v T) {
defer l.scheduler.Return()
defer wg.Done()
if err := l.batchProcessor.Process(ctx, v); err != nil && !errors.Is(err, context.Canceled) {
l.logrus.WithError(err).Error("Error processing batch")
}
}(v)
}
wg.Wait()
}
}
func (l *chanProcessorWorker[T]) Start(ctx context.Context) error {
c, err := l.chanProcessor.Query(ctx)
if err != nil {
return err
}
for {
select {
case <-ctx.Done():
return ctx.Err()
case v, ok := <-c:
if !ok {
return nil
}
l.scheduler.Take()
go func(v T) {
defer l.scheduler.Return()
if err := l.chanProcessor.Process(ctx, v); err != nil {
l.logrus.WithError(err).Error("Error processing chan")
}
}(v)
}
}
}
|
