-
Notifications
You must be signed in to change notification settings - Fork 106
/
filterratelimit.go
88 lines (73 loc) · 1.77 KB
/
filterratelimit.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
package filterratelimit
import (
"context"
"time"
"github.com/tsaikd/gogstash/config"
"github.com/tsaikd/gogstash/config/goglog"
"github.com/tsaikd/gogstash/config/logevent"
)
// ModuleName is the name used in config file
const ModuleName = "rate_limit"
// FilterConfig holds the configuration json fields and internal objects
type FilterConfig struct {
config.FilterConfig
Rate int64 `json:"rate"` // event number per second
Burst int64 `json:"burst"` // burst limit
throttle chan time.Time
}
// DefaultFilterConfig returns an FilterConfig struct with default values
func DefaultFilterConfig() FilterConfig {
return FilterConfig{
FilterConfig: config.FilterConfig{
CommonConfig: config.CommonConfig{
Type: ModuleName,
},
},
Burst: 100,
}
}
// InitHandler initialize the filter plugin
func InitHandler(
ctx context.Context,
raw config.ConfigRaw,
control config.Control,
) (config.TypeFilterConfig, error) {
conf := DefaultFilterConfig()
if err := config.ReflectConfig(raw, &conf); err != nil {
return nil, err
}
if conf.Rate <= 0 {
goglog.Logger.Warn("filter ratelimit config rate should > 0, ignored")
return &conf, nil
}
conf.throttle = make(chan time.Time, conf.Burst)
tick := time.NewTicker(time.Second / time.Duration(conf.Rate))
go func() {
defer tick.Stop()
for {
select {
case <-ctx.Done():
return
case t := <-tick.C:
select {
case <-ctx.Done():
return
case conf.throttle <- t:
default:
}
}
}
}()
return &conf, nil
}
// Event the main filter event
func (f *FilterConfig) Event(ctx context.Context, event logevent.LogEvent) (logevent.LogEvent, bool) {
if event.Extra == nil {
event.Extra = map[string]any{}
}
if f.throttle == nil {
return event, false
}
<-f.throttle
return event, true
}