forked from benthosdev/benthos
/
sample.go
113 lines (94 loc) · 3.01 KB
/
sample.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
package processor
import (
"math/rand"
"sync"
"time"
"github.com/dafanshu/benthos/v3/internal/docs"
"github.com/dafanshu/benthos/v3/lib/log"
"github.com/dafanshu/benthos/v3/lib/metrics"
"github.com/dafanshu/benthos/v3/lib/response"
"github.com/dafanshu/benthos/v3/lib/types"
)
//------------------------------------------------------------------------------
func init() {
Constructors[TypeSample] = TypeSpec{
constructor: NewSample,
Status: docs.StatusDeprecated,
Footnotes: `
## Alternatives
All functionality of this processor has been superseded by the
[bloblang](/docs/components/processors/bloblang) processor.`,
FieldSpecs: docs.FieldSpecs{
docs.FieldCommon("retain", "The percentage of messages to keep."),
docs.FieldCommon("seed", "A seed for pseudo-random sampling."),
},
}
}
//------------------------------------------------------------------------------
// SampleConfig contains configuration fields for the Sample processor.
type SampleConfig struct {
Retain float64 `json:"retain" yaml:"retain"`
RandomSeed int64 `json:"seed" yaml:"seed"`
}
// NewSampleConfig returns a SampleConfig with default values.
func NewSampleConfig() SampleConfig {
return SampleConfig{
Retain: 10.0, // 10%
RandomSeed: 0,
}
}
//------------------------------------------------------------------------------
// Sample is a processor that drops messages based on a random sample.
type Sample struct {
conf Config
log log.Modular
stats metrics.Type
retain float64
gen *rand.Rand
mut sync.Mutex
mCount metrics.StatCounter
mDropped metrics.StatCounter
mSent metrics.StatCounter
mBatchSent metrics.StatCounter
}
// NewSample returns a Sample processor.
func NewSample(
conf Config, mgr types.Manager, log log.Modular, stats metrics.Type,
) (Type, error) {
gen := rand.New(rand.NewSource(conf.Sample.RandomSeed))
return &Sample{
conf: conf,
log: log,
stats: stats,
retain: conf.Sample.Retain / 100.0,
gen: gen,
mCount: stats.GetCounter("count"),
mDropped: stats.GetCounter("dropped"),
mSent: stats.GetCounter("sent"),
mBatchSent: stats.GetCounter("batch.sent"),
}, nil
}
//------------------------------------------------------------------------------
// ProcessMessage applies the processor to a message, either creating >0
// resulting messages or a response to be sent back to the message source.
func (s *Sample) ProcessMessage(msg types.Message) ([]types.Message, types.Response) {
s.mCount.Incr(1)
s.mut.Lock()
defer s.mut.Unlock()
if s.gen.Float64() > s.retain {
s.mDropped.Incr(1)
return nil, response.NewAck()
}
s.mBatchSent.Incr(1)
s.mSent.Incr(int64(msg.Len()))
msgs := [1]types.Message{msg}
return msgs[:], nil
}
// CloseAsync shuts down the processor and stops processing requests.
func (s *Sample) CloseAsync() {
}
// WaitForClose blocks until the processor has closed down.
func (s *Sample) WaitForClose(timeout time.Duration) error {
return nil
}
//------------------------------------------------------------------------------