-
Notifications
You must be signed in to change notification settings - Fork 45
/
batch.go
109 lines (92 loc) · 2.38 KB
/
batch.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
/*
Copyright (c) 2021 - Present. Blend Labs, Inc. All rights reserved
Use of this source code is governed by a MIT license that can be found in the LICENSE file.
*/
package async
import (
"context"
"runtime"
)
// NewBatch creates a new batch processor.
// Batch processes are a known quantity of work that needs to be processed in parallel.
func NewBatch(work chan interface{}, action WorkAction, options ...BatchOption) *Batch {
b := Batch{
Action: action,
Work: work,
Parallelism: runtime.NumCPU(),
}
for _, option := range options {
option(&b)
}
return &b
}
// BatchOption is an option for the batch worker.
type BatchOption func(*Batch)
// OptBatchErrors sets the batch worker error return channel.
func OptBatchErrors(errors chan error) BatchOption {
return func(i *Batch) {
i.Errors = errors
}
}
// OptBatchParallelism sets the batch worker parallelism, or the number of workers to create.
func OptBatchParallelism(parallelism int) BatchOption {
return func(i *Batch) {
i.Parallelism = parallelism
}
}
// Batch is a batch of work executed by a fixed count of workers.
type Batch struct {
Action WorkAction
Parallelism int
Work chan interface{}
Errors chan error
}
// Process executes the action for all the work items.
func (b *Batch) Process(ctx context.Context) {
allWorkers := make([]*Worker, b.Parallelism)
availableWorkers := make(chan *Worker, b.Parallelism)
// return worker is a local finalizer
// that grabs a reference to the workers set.
returnWorker := func(ctx context.Context, worker *Worker) error {
availableWorkers <- worker
return nil
}
// create and start workers.
for x := 0; x < b.Parallelism; x++ {
worker := NewWorker(b.Action)
worker.Context = ctx
worker.Errors = b.Errors
worker.Finalizer = returnWorker
workerStarted := worker.NotifyStarted()
go func() { _ = worker.Start() }()
<-workerStarted
allWorkers[x] = worker
availableWorkers <- worker
}
defer func() {
for x := 0; x < len(allWorkers); x++ {
_ = allWorkers[x].Stop()
}
}()
numWorkItems := len(b.Work)
var worker *Worker
var workItem interface{}
for x := 0; x < numWorkItems; x++ {
select {
case <-ctx.Done():
return
default:
}
select {
case workItem = <-b.Work:
select {
case worker = <-availableWorkers:
worker.Enqueue(workItem)
case <-ctx.Done():
return
}
case <-ctx.Done():
return
}
}
}