/
pipeline.go
133 lines (118 loc) · 2.7 KB
/
pipeline.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
package pipeline
import (
"context"
"encoding/json"
"go.uber.org/zap"
"golang.org/x/sync/errgroup"
)
type pipelineState uint8
const (
chanSize = 0
)
type processMetrics struct {
SrcWait *metricsResult `json:"source_wait"`
Proc *metricsResult `json:"processing_time"`
EmitWait *metricsResult `json:"emit_wait"`
}
type process interface {
getID() string
setup(srcChan, dstChan chan interface{}, logger *zap.Logger)
getDstChan() chan interface{}
run(context.Context) error
metrics() *processMetrics
}
type sinkMetrics struct {
SrcWait *metricsResult `json:"source_wait"`
Sink *metricsResult `json:"sink_time"`
}
type sink interface {
getID() string
setup(srcChan chan interface{}, logger *zap.Logger)
run(context.Context) error
metrics() *sinkMetrics
}
type Pipeline struct {
srcChan chan interface{}
process []process
sink sink
logger *zap.Logger
}
func NewPipeline(logger *zap.Logger) *Pipeline {
return &Pipeline{
srcChan: make(chan interface{}, chanSize),
process: make([]process, 0),
logger: logger,
}
}
func (p *Pipeline) Add(proc process) error {
dstChan := make(chan interface{}, chanSize)
srcChan := p.srcChan
if l := len(p.process); l > 0 {
srcChan = p.process[l-1].getDstChan()
}
proc.setup(srcChan, dstChan, p.logger)
p.process = append(p.process, proc)
return nil
}
func (p *Pipeline) Run(ctx context.Context, s sink) error {
dstChan := p.srcChan
if len(p.process) > 0 {
dstChan = p.process[len(p.process)-1].getDstChan()
}
s.setup(dstChan, p.logger)
p.sink = s
cctx, cancel := context.WithCancel(ctx)
defer cancel()
g, gctx := errgroup.WithContext(cctx)
for _, pw := range p.process {
func(p process) {
g.Go(func() error {
if err := p.run(gctx); err != nil {
return err
}
return nil
})
}(pw)
}
g.Go(func() error {
if err := p.sink.run(gctx); err != nil {
return err
}
return nil
})
if err := g.Wait(); err != nil {
return err
}
return nil
}
func (p *Pipeline) Submit(ctx context.Context, item interface{}) error {
select {
case <-ctx.Done():
return nil
case p.srcChan <- item:
}
return nil
}
func (p *Pipeline) Metrics() string {
type pipelineMetric struct {
Id string `json:"id"`
Metrics interface{} `json:"metrics"`
}
metrics := make([]*pipelineMetric, 0)
for _, proc := range p.process {
metrics = append(metrics, &pipelineMetric{proc.getID(), proc.metrics()})
}
if p.sink != nil {
metrics = append(metrics, &pipelineMetric{p.sink.getID(), p.sink.metrics()})
}
json, err := json.MarshalIndent(metrics, "", " ")
if err != nil {
panic(err)
}
return string(json)
}
func (p *Pipeline) Shutdown() error {
close(p.srcChan)
p.logger.Debug("pipeline draining for shutdown")
return nil
}