├── policymanager ├── config.yml └── main.go ├── register ├── plugin │ └── plugin.go └── compile │ └── plugin.go ├── Makefile ├── generator ├── Dockerfile └── main.go ├── limit_test.go ├── go.mod ├── Dockerfile ├── LICENSE ├── rule_test.go ├── plugin_test.go ├── .travis.yml ├── rule.go ├── condition_limiter.go ├── bucket_limiter.go ├── condition_limiter_test.go ├── limit.go ├── go.sum ├── plugin.go └── README.md /policymanager/config.yml: -------------------------------------------------------------------------------- 1 | --- 2 | key: id 3 | default_limit: 1 4 | rules: 5 | - limit: 100 6 | selectors: 7 | id: foo 8 | -------------------------------------------------------------------------------- /register/plugin/plugin.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "github.com/elastic/beats/libbeat/processors" 5 | "github.com/ozonru/filebeat-throttle-plugin" 6 | ) 7 | 8 | var Bundle = processors.Plugin("throttle", throttleplugin.NewProcessor) 9 | -------------------------------------------------------------------------------- /policymanager/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "net/http" 5 | ) 6 | 7 | func GetHandler(w http.ResponseWriter, r *http.Request) { 8 | http.ServeFile(w, r, "config.yml") 9 | } 10 | 11 | func main() { 12 | http.HandleFunc("/policy", GetHandler) 13 | http.ListenAndServe(":8080", nil) 14 | } 15 | -------------------------------------------------------------------------------- /register/compile/plugin.go: -------------------------------------------------------------------------------- 1 | // package main contains only import to register throttle plugin in filebeat. 2 | package main 3 | 4 | import ( 5 | "github.com/elastic/beats/libbeat/processors" 6 | "github.com/ozonru/filebeat-throttle-plugin" 7 | ) 8 | 9 | func init() { 10 | processors.RegisterPlugin("throttle", throttleplugin.NewProcessor) 11 | } 12 | 13 | 14 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | all: test 2 | 3 | 4 | test: 5 | go test -race -v ./... 6 | 7 | linux_plugin: 8 | GOOS=linux go build -v -o output/filebeat_throttle_linux.so -buildmode=plugin github.com/ozonru/filebeat-throttle-plugin/register/plugin 9 | 10 | darwin_plugin: 11 | GOOS=darwin go build -v -o output/filebeat_throttle_darwin.so -buildmode=plugin github.com/ozonru/filebeat-throttle-plugin/register/plugin 12 | -------------------------------------------------------------------------------- /generator/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM golang:1.11-stretch 2 | 3 | COPY main.go /go/src/gitlab.ozon.ru/sre/logs-generator/main.go 4 | 5 | RUN go build -o /usr/local/bin/logs-generator /go/src/gitlab.ozon.ru/sre/logs-generator/main.go 6 | 7 | FROM debian:stretch-slim 8 | 9 | COPY --from=0 /usr/local/bin/logs-generator /usr/local/bin/logs-generator 10 | 11 | ENTRYPOINT ["/usr/local/bin/logs-generator"] 12 | CMD ["-id", "some_id"] 13 | -------------------------------------------------------------------------------- /generator/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "flag" 5 | "fmt" 6 | "time" 7 | ) 8 | 9 | var ( 10 | rate = flag.Float64("rate", 10, "records per second") 11 | id = flag.String("id", "generator", "`id` field value") 12 | ) 13 | 14 | func main() { 15 | flag.Parse() 16 | 17 | var j int64 18 | interval := time.Duration(1 / *rate * float64(time.Second)) 19 | 20 | for range time.Tick(interval) { 21 | j++ 22 | fmt.Printf(`{"ts":"%v", "message": "%v", "id":"%v"}`+"\n", time.Now().Format(time.RFC3339Nano), j, *id) 23 | } 24 | } 25 | 26 | -------------------------------------------------------------------------------- /limit_test.go: -------------------------------------------------------------------------------- 1 | package throttleplugin 2 | 3 | import ( 4 | "context" 5 | "net/http" 6 | "net/http/httptest" 7 | "testing" 8 | 9 | "github.com/stretchr/testify/assert" 10 | ) 11 | 12 | func testServer(t *testing.T, body []byte) (url string, closeFn func()) { 13 | h := func(w http.ResponseWriter, r *http.Request) { 14 | w.Write(body) 15 | } 16 | s := httptest.NewServer(http.HandlerFunc(h)) 17 | 18 | return s.URL, func() { s.Close() } 19 | } 20 | 21 | func TestRemoteLimiter_Update(t *testing.T) { 22 | response := `key: id 23 | default_limit: 1 24 | rules: 25 | - limit: 100 26 | selectors: 27 | id: foo` 28 | url, closeFn := testServer(t, []byte(response)) 29 | defer closeFn() 30 | 31 | l, _ := NewRemoteLimiter(url, 1, 10) 32 | assert.NoError(t, l.Update(context.Background())) 33 | } 34 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/ozonru/filebeat-throttle-plugin 2 | 3 | require ( 4 | github.com/elastic/beats v6.6.2+incompatible 5 | github.com/elastic/go-ucfg v0.7.0 // indirect 6 | github.com/gofrs/uuid v3.2.0+incompatible // indirect 7 | github.com/inconshreveable/mousetrap v1.0.0 // indirect 8 | github.com/joeshaw/multierror v0.0.0-20140124173710-69b34d4ec901 // indirect 9 | github.com/pkg/errors v0.8.1 10 | github.com/prometheus/client_golang v0.9.2 11 | github.com/spf13/cobra v0.0.3 // indirect 12 | github.com/spf13/pflag v1.0.3 // indirect 13 | github.com/stretchr/testify v1.3.0 14 | go.uber.org/atomic v1.3.2 // indirect 15 | go.uber.org/multierr v1.1.0 // indirect 16 | go.uber.org/zap v1.9.1 // indirect 17 | golang.org/x/sys v0.0.0-20190321052220-f7bb7a8bee54 // indirect 18 | gopkg.in/yaml.v2 v2.2.2 19 | ) 20 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | ARG filebeatVersion=6.5.4 2 | ARG goVersion=1.10.6 3 | FROM golang:$goVersion 4 | ARG filebeatVersion 5 | RUN curl -L --output /tmp/filebeat.tar.gz https://github.com/elastic/beats/archive/v$filebeatVersion.tar.gz 6 | 7 | RUN mkdir -p /go/src/github.com/elastic/beats && tar -xvzf /tmp/filebeat.tar.gz --strip-components=1 -C /go/src/github.com/elastic/beats 8 | RUN go get -v golang.org/x/vgo 9 | 10 | COPY . /go/src/github.com/ozonru/filebeat-throttle-plugin 11 | COPY register/plugin/plugin.go /go/src/github.com/elastic/beats/libbeat/processors/throttle/plugin.go 12 | 13 | RUN (cd /go/src/github.com/ozonru/filebeat-throttle-plugin && vgo mod vendor -v) 14 | RUN rm -rf /go/src/github.com/ozonru/filebeat-throttle-plugin/vendor/github.com/elastic 15 | RUN go build -v -o /output/filebeat_throttle_linux.so -buildmode=plugin github.com/elastic/beats/libbeat/processors/throttle 16 | 17 | 18 | FROM docker.elastic.co/beats/filebeat:$filebeatVersion 19 | COPY --from=0 /output/filebeat_throttle_linux.so /filebeat_throttle_linux.so 20 | 21 | CMD ["-e", "--plugin", "/filebeat_throttle_linux.so"] 22 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 OZON.ru 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /rule_test.go: -------------------------------------------------------------------------------- 1 | package throttleplugin 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/elastic/beats/libbeat/beat" 7 | "github.com/elastic/beats/libbeat/common" 8 | "github.com/stretchr/testify/assert" 9 | ) 10 | 11 | func TestNewRule(t *testing.T) { 12 | fields := map[string]string{ 13 | "a": "1", 14 | "b": "2", 15 | "c": "3", 16 | "d": "4", 17 | "e": "5", 18 | } 19 | 20 | r := NewRule(fields, 100) 21 | 22 | assert.Equal(t, []string{"a", "b", "c", "d", "e"}, r.keys) 23 | assert.Equal(t, []string{"1", "2", "3", "4", "5"}, r.values) 24 | } 25 | 26 | func TestMatch(t *testing.T) { 27 | event := &beat.Event{Fields: common.MapStr{}} 28 | event.PutValue("a", "1") 29 | event.PutValue("b", "2") 30 | 31 | t.Run("match", func(t *testing.T) { 32 | r := NewRule(map[string]string{"a": "1"}, 10) 33 | ok, key := r.Match(event) 34 | 35 | assert.True(t, ok) 36 | assert.Equal(t, "10:1:", key) 37 | }) 38 | } 39 | func BenchmarkMatch(b *testing.B) { 40 | fields := map[string]string{ 41 | "a": "1", 42 | "b": "2", 43 | "c": "3", 44 | "d": "4", 45 | "e": "5", 46 | } 47 | 48 | r := NewRule(fields, 100) 49 | event := &beat.Event{Fields: common.MapStr{}} 50 | event.PutValue("a", "1") 51 | for k, v := range fields { 52 | event.PutValue(k, v) 53 | } 54 | 55 | for i := 0; i < b.N; i++ { 56 | r.Match(event) 57 | } 58 | } 59 | -------------------------------------------------------------------------------- /plugin_test.go: -------------------------------------------------------------------------------- 1 | package throttleplugin 2 | 3 | import ( 4 | "fmt" 5 | "math/rand" 6 | "strings" 7 | "testing" 8 | 9 | "github.com/elastic/beats/libbeat/beat" 10 | "github.com/elastic/beats/libbeat/common" 11 | "github.com/stretchr/testify/assert" 12 | ) 13 | 14 | var config = `--- 15 | metric_name: %metric% 16 | policy_update_interval: 60m 17 | labels: 18 | - 19 | from: input.rrrr 20 | to: input 21 | - 22 | from: host.name 23 | to: host` 24 | 25 | func getConfig() []byte { 26 | name := fmt.Sprintf("metric_%v", rand.Int63()) 27 | cfg := strings.Replace(config, "%metric%", name, -1) 28 | 29 | return []byte(cfg) 30 | } 31 | 32 | func TestConfig_GetMetricLabels(t *testing.T) { 33 | c := Config{ 34 | MetricLabels: []LabelMapping{ 35 | {"from_1", "to_1"}, 36 | {"from_2", "to_2"}, 37 | {"from_3", "to_3"}, 38 | }, 39 | } 40 | 41 | assert.Equal(t, []string{"to_1", "to_2", "to_3"}, c.GetMetricLabels()) 42 | } 43 | 44 | func TestConfig_GetFields(t *testing.T) { 45 | c := Config{ 46 | MetricLabels: []LabelMapping{ 47 | {"from_1", "to_1"}, 48 | {"from_2", "to_2"}, 49 | {"from_3", "to_3"}, 50 | }, 51 | } 52 | 53 | assert.Equal(t, []string{"from_1", "from_2", "from_3"}, c.GetFields()) 54 | } 55 | 56 | func BenchmarkRun(b *testing.B) { 57 | cfg, err := common.NewConfigWithYAML(getConfig(), "test") 58 | if err != nil { 59 | b.Fatal(err) 60 | } 61 | mp, err := newProcessor(cfg) 62 | if err != nil { 63 | b.Fatal(err) 64 | } 65 | defer mp.Close() 66 | 67 | event := &beat.Event{ 68 | Fields: common.MapStr{}, 69 | } 70 | event.PutValue("input.type", "foo") 71 | event.PutValue("host.name", "bar") 72 | for i := 0; i < b.N; i++ { 73 | mp.Run(event) 74 | } 75 | } 76 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | services: 3 | - docker 4 | go: 5 | - 1.11.x 6 | cache: 7 | directories: 8 | - $HOME/.gocache 9 | env: 10 | matrix: 11 | - GO111MODULE=on 12 | global: 13 | - GOCACHE=$HOME/.gocache 14 | - TAG=`if [ "$TRAVIS_BRANCH" == "master" ]; then echo "latest"; else echo $TRAVIS_BRANCH ; fi` 15 | - secure: TtgIxjXZ6modBXExBWECdxz/cjOjf17vg5c8sL6t9u74lzDTswhmkuTKQ5HbcIfPcvKlBZ5J1l96qDSgW5S58zu2X22VaBXRMgUoolXe8u1RfzFeU7E7vCzpeCIqLCDLn+1NAvwHIKNx2+ZcjFBiGRcG0ol8WfG6OgvQjr9RbmLLsJFLMVxDEpWRM9kFZMTv98ijPucpbCQOOC3B6aS2XvV863vKhrUTavm9AaWzUVdmlXqikWEjaRZVbm1IdI4qs6sqLVJd3kV1a55dFwJl0SI52/kfIIDOg8c40wPx0uFu4wSY6g5enS5CwDQZ7RSRNtL3khfOaD6zWOOnztUK2YOV85oaKK3rO95o4FOxjRNDRhF4Tb2mDmcoXBt17OvOvQ7hrOm0GNH7SA6nFpl/emjBqg3r3FsM/eZTQd6p2A3HseGHQ8wpOpQngyF9kUrgL03ChCpoc11EziFQIdkjNkLfG3OsUxF/Z8WD74IkU2bSjZ+por6U0If+xK0I3xEpozHuQCns8v6dRp31bRwl9Om8VqCHa7T38ypzFeuSzcFTk0Y+f3b0l5PsWwytA/gMfUGeGwNPU703NMI4HtbyiL82FGaMS0+gXpHsRXHzAk/e+mgM0Mx2kH94MWyrefyxW5bbZrbEKzVhExHfOAg8OM1KhnUIIU92t1nIR3Zy0EU= 16 | - secure: iN6+yZpQH+klM8wJ16Y9Jle6an5saQ2bzOdTREiZFIohaa1t4gHgr4fx2spJsOv/1B9iYSyWZ/UkSDZb8ZfwzI9SdvNbJ0Oj8C1o64/R5uK/mRYOhNK9jzEK3Yehj717Rtaeh1zZr9f6oDF99Sqp/so3PBHh94W4VENgRvHLHuasrkAFF0wZ87ri0dBpQrv82ngxA6nddIGVvM6WifrY+YbPr2Jbb5rL/SzACQgRp3Vlhnt4OXf3KoICm4hakOfLY2p82mncWImShxEOYCGeApm9WgB0/OM/ywBHU2yOHyfolP7M1WlxrlFqTJt9+fu1AqKp1rzQAdXPrkvs4MHImYWmMT6jGioVY5BY+1uZ6IW6YbTRUTrf29bT12Q5d+C6wIAx4I2ZPuK7tAORXZjpBprGEZ+tJy16Om6EKZydj4Vu+xtWAdS2FwVBUYK1CZGgVFq9eN87tFFohw7F5ACGmYHyLXItuqd3QE5K1Ky9tfJI7VVXBDfONO99kKgFtHi6crt+Vuf553gWNja9wlPKuclYiFNurJ7Gqj9Ml0Z5xKWD2We9i2UqNTV3u7NQ6wQ8oLkDXbz6fEvT3EmO3+Xh5qIG9xcUeArKV1IPDOB+KPkKEPHbfdcGHj5hmM8ojd03h6xsSTOUuCjtyh9Go2WolPaM9jTvkIW3GixXwd9ZBmI= 17 | script: 18 | - make test 19 | - docker build -t ozonru/filebeat-throttle-plugin:${TAG} . 20 | after_success: 21 | - echo "$DOCKER_PASSWORD" | docker login -u "$DOCKER_USERNAME" --password-stdin 22 | - docker push ozonru/filebeat-throttle-plugin:${TAG} 23 | -------------------------------------------------------------------------------- /rule.go: -------------------------------------------------------------------------------- 1 | package throttleplugin 2 | 3 | import ( 4 | "bytes" 5 | "sort" 6 | "strconv" 7 | "sync" 8 | "unsafe" 9 | 10 | "github.com/elastic/beats/libbeat/beat" 11 | ) 12 | 13 | var sbPool sync.Pool 14 | 15 | func init() { 16 | sbPool = sync.Pool{ 17 | New: func() interface{} { 18 | return &bytes.Buffer{} 19 | }, 20 | } 21 | } 22 | 23 | type Rule struct { 24 | keys []string // sorted list of used keys is used for combining limiter key. 25 | values []string // values to check against. order is the same as for keys. 26 | limit int64 27 | 28 | // baseKey contains strings representation of limit to increase Match performance. 29 | // strconv.Itoa makes 2 allocations with 32 bytes for each call. 30 | baseKey string 31 | } 32 | 33 | // NewRule returns new Rule instance. 34 | func NewRule(fields map[string]string, limit int64) Rule { 35 | var ( 36 | keys = make([]string, 0, len(fields)) 37 | values = make([]string, len(fields)) 38 | ) 39 | 40 | for k := range fields { 41 | keys = append(keys, k) 42 | } 43 | 44 | sort.Strings(keys) 45 | 46 | for i, k := range keys { 47 | values[i] = fields[k] 48 | } 49 | 50 | return Rule{ 51 | keys: keys, 52 | values: values, 53 | limit: limit, 54 | baseKey: strconv.FormatInt(limit, 10), 55 | } 56 | } 57 | 58 | // Limit returns current limit. 59 | func (r Rule) Limit() int64 { 60 | return r.limit 61 | } 62 | 63 | // Match checks if event has the same field values as expected. 64 | func (r Rule) Match(e *beat.Event) (ok bool, key string) { 65 | b := sbPool.Get() 66 | sb := b.(*bytes.Buffer) 67 | sb.Reset() 68 | defer func() { 69 | sbPool.Put(b) 70 | }() 71 | sb.WriteString(r.baseKey) 72 | sb.WriteByte(':') 73 | for i, k := range r.keys { 74 | v, err := e.GetValue(k) 75 | if err != nil { 76 | return false, "" 77 | } 78 | 79 | sv, ok := v.(string) 80 | if !ok { 81 | // only strings values are supported 82 | return false, "" 83 | } 84 | 85 | if sv != r.values[i] { 86 | return false, "" 87 | } 88 | 89 | sb.WriteString(sv) 90 | sb.WriteByte(':') 91 | } 92 | 93 | buf := sb.Bytes() 94 | // zero-allocation convertion from []bytes to string. 95 | s := *(*string)(unsafe.Pointer(&buf)) 96 | 97 | return true, s 98 | } 99 | -------------------------------------------------------------------------------- /condition_limiter.go: -------------------------------------------------------------------------------- 1 | package throttleplugin 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | "time" 7 | 8 | "github.com/elastic/beats/libbeat/beat" 9 | "github.com/elastic/beats/libbeat/conditions" 10 | "github.com/pkg/errors" 11 | ) 12 | 13 | // ConditionLimiter first checks if event is valid for specified conditions and then applies rate limiting. 14 | type ConditionLimiter struct { 15 | condition conditions.Condition 16 | keys []string // sorted list of used keys is used for combining limiter key. 17 | fields map[string]string // used only for WriteStatus functionality, because it's hard to pretty print conditions. 18 | 19 | bl *BucketLimiter 20 | } 21 | 22 | // NewConditionLimiter returns new ConditionLimiter instance. 23 | func NewConditionLimiter(fields map[string]string, bucketInterval, limit, buckets int64, now time.Time) (*ConditionLimiter, error) { 24 | f := conditions.Fields{} 25 | if err := f.Unpack(prepareFields(fields)); err != nil { 26 | return nil, errors.Wrap(err, "failed to unpack fields") 27 | } 28 | 29 | cs, err := conditions.NewCondition(&conditions.Config{Equals: &f}) 30 | if err != nil { 31 | return nil, errors.Wrap(err, "failed to create conditions") 32 | } 33 | 34 | return &ConditionLimiter{ 35 | fields: fields, 36 | condition: cs, 37 | bl: NewBucketLimiter(bucketInterval, limit, buckets, now), 38 | }, nil 39 | } 40 | 41 | // WriteStatus writes text based status into Writer. 42 | func (cl *ConditionLimiter) WriteStatus(w io.Writer) error { 43 | fmt.Fprintf(w, "%v\n", cl.fields) 44 | 45 | return cl.bl.WriteStatus(w) 46 | } 47 | 48 | // SetLimit updates limit value. 49 | // Note: it's allowed only to change limit, not bucketInterval. 50 | func (cl *ConditionLimiter) SetLimit(limit int64) { 51 | cl.bl.SetLimit(limit) 52 | } 53 | 54 | // Check checks if event satisfies condition. 55 | func (cl *ConditionLimiter) Check(e *beat.Event) bool { 56 | return cl.condition.Check(e) 57 | } 58 | 59 | // Allow returns TRUE if event is allowed to be processed. 60 | func (cl *ConditionLimiter) Allow(t time.Time) bool { 61 | return cl.bl.Allow(t) 62 | } 63 | 64 | func prepareFields(m map[string]string) map[string]interface{} { 65 | r := make(map[string]interface{}, len(m)) 66 | for k, v := range m { 67 | r[k] = v 68 | } 69 | 70 | return r 71 | } 72 | -------------------------------------------------------------------------------- /bucket_limiter.go: -------------------------------------------------------------------------------- 1 | package throttleplugin 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | "sync" 7 | "time" 8 | ) 9 | 10 | type BucketLimiter struct { 11 | mu sync.Mutex 12 | bucketInterval int64 // bucket interval in seconds (60 = 1 min) 13 | limit int64 // maximum number of events per bucket 14 | minBucketID int64 // minimum bucket id 15 | buckets []int64 16 | lastUpdate time.Time 17 | } 18 | 19 | func NewBucketLimiter(bucketInterval, limit, buckets int64, now time.Time) *BucketLimiter { 20 | return &BucketLimiter{ 21 | bucketInterval: bucketInterval, 22 | limit: limit, 23 | minBucketID: timeToBucketID(now, bucketInterval) - buckets + 1, 24 | buckets: make([]int64, buckets), 25 | } 26 | } 27 | 28 | // Allow returns TRUE if event is allowed to be processed. 29 | func (bl *BucketLimiter) Allow(t time.Time) bool { 30 | index := timeToBucketID(t, bl.bucketInterval) 31 | 32 | bl.mu.Lock() 33 | defer bl.mu.Unlock() 34 | bl.lastUpdate = time.Now() 35 | 36 | max := bl.minBucketID + int64(len(bl.buckets)) - 1 37 | 38 | if index < bl.minBucketID { 39 | // limiter doesn't track that bucket anymore. 40 | return false 41 | } 42 | 43 | if index > max { 44 | // event from new bucket. We need to add N new buckets 45 | n := index - max 46 | for i := 0; int64(i) < n; i++ { 47 | bl.buckets = append(bl.buckets, 0) 48 | } 49 | 50 | // remove old ones 51 | bl.buckets = bl.buckets[n:] 52 | 53 | // and set new min index 54 | bl.minBucketID += n 55 | } 56 | 57 | return bl.increment(index) 58 | } 59 | 60 | // LastUpdate returns last Allow method call time. 61 | func (bl *BucketLimiter) LastUpdate() time.Time { 62 | bl.mu.Lock() 63 | defer bl.mu.Unlock() 64 | 65 | return bl.lastUpdate 66 | } 67 | 68 | // WriteStatus writes text based status into Writer. 69 | func (bl *BucketLimiter) WriteStatus(w io.Writer) error { 70 | bl.mu.Lock() 71 | defer bl.mu.Unlock() 72 | 73 | for i, value := range bl.buckets { 74 | fmt.Fprintf(w, "#%s: ", bucketIDToTime(int64(i)+bl.minBucketID, bl.bucketInterval)) 75 | progress(w, value, bl.limit, 20) 76 | fmt.Fprintf(w, " %d/%d\n", value, bl.limit) 77 | } 78 | return nil 79 | } 80 | 81 | // SetLimit updates limit value. 82 | // Note: it's allowed only to change limit, not bucketInterval. 83 | func (bl *BucketLimiter) SetLimit(limit int64) { 84 | bl.mu.Lock() 85 | bl.limit = limit 86 | bl.mu.Unlock() 87 | } 88 | 89 | // increment adds 1 to specified bucket. 90 | // Note: this func is not thread safe, so it must be guarded with lock. 91 | func (сl *BucketLimiter) increment(index int64) bool { 92 | i := index - сl.minBucketID 93 | if сl.buckets[i] >= int64(сl.limit) { 94 | return false 95 | } 96 | сl.buckets[i]++ 97 | 98 | return true 99 | } 100 | 101 | func progress(w io.Writer, current, limit, max int64) { 102 | p := float64(current) / float64(limit) * float64(max) 103 | 104 | fmt.Fprint(w, "[") 105 | for i := int64(0); i < max; i++ { 106 | if i < int64(p) { 107 | fmt.Fprint(w, "#") 108 | } else { 109 | fmt.Fprint(w, "_") 110 | } 111 | } 112 | fmt.Fprint(w, "]") 113 | } 114 | 115 | // bucketbucketIDToTime converts bucketID to time. This time is start of the bucket. 116 | func bucketIDToTime(id int64, interval int64) time.Time { 117 | return time.Unix(id*interval, 0) 118 | } 119 | 120 | // timeToBucketID converts time to bucketID. 121 | func timeToBucketID(t time.Time, interval int64) int64 { 122 | return t.Unix() / interval 123 | } 124 | -------------------------------------------------------------------------------- /condition_limiter_test.go: -------------------------------------------------------------------------------- 1 | package throttleplugin 2 | 3 | import ( 4 | "bytes" 5 | "testing" 6 | "time" 7 | 8 | "github.com/elastic/beats/libbeat/beat" 9 | "github.com/elastic/beats/libbeat/common" 10 | "github.com/stretchr/testify/assert" 11 | ) 12 | 13 | func TestAllow(t *testing.T) { 14 | now := time.Date(2018, 12, 19, 19, 30, 25, 0, time.UTC) 15 | 16 | t.Run("current bucket overflow", func(t *testing.T) { 17 | cl, err := NewConditionLimiter(map[string]string{}, 60, 2, 5, now) 18 | if err != nil { 19 | t.Fatal(err) 20 | } 21 | 22 | assert.True(t, cl.Allow(now)) 23 | assert.True(t, cl.Allow(now)) 24 | assert.False(t, cl.Allow(now), "bucket must be exceeded") 25 | }) 26 | 27 | t.Run("update limit", func(t *testing.T) { 28 | cl, err := NewConditionLimiter(map[string]string{}, 60, 2, 5, now) 29 | if err != nil { 30 | t.Fatal(err) 31 | } 32 | 33 | assert.True(t, cl.Allow(now)) 34 | assert.True(t, cl.Allow(now)) 35 | assert.False(t, cl.Allow(now)) 36 | 37 | cl.SetLimit(3) 38 | assert.True(t, cl.Allow(now)) 39 | assert.False(t, cl.Allow(now), "bucket must be exceeded") 40 | }) 41 | 42 | t.Run("add to oldest bucket", func(t *testing.T) { 43 | cl, err := NewConditionLimiter(map[string]string{}, 60, 2, 5, now) 44 | if err != nil { 45 | t.Fatal(err) 46 | } 47 | 48 | ts := now.Add(-4 * time.Minute) 49 | 50 | assert.True(t, cl.Allow(ts)) 51 | assert.True(t, cl.Allow(ts)) 52 | }) 53 | 54 | t.Run("old bucket shifts", func(t *testing.T) { 55 | cl, err := NewConditionLimiter(map[string]string{}, 60, 2, 5, now) 56 | if err != nil { 57 | t.Fatal(err) 58 | } 59 | 60 | oldest := now.Add(-4 * time.Minute) 61 | 62 | assert.True(t, cl.Allow(now.Add(time.Minute))) 63 | assert.False(t, cl.Allow(oldest), "oldest bucket must be shifted") 64 | }) 65 | } 66 | 67 | func TestCheck(t *testing.T) { 68 | cl, _ := NewConditionLimiter(map[string]string{"a": "1"}, 60, 2, 5, time.Now()) 69 | 70 | t.Run("valid", func(t *testing.T) { 71 | event := &beat.Event{Fields: common.MapStr{}} 72 | event.PutValue("a", "1") 73 | 74 | assert.True(t, cl.Check(event)) 75 | }) 76 | 77 | t.Run("invalid", func(t *testing.T) { 78 | event := &beat.Event{Fields: common.MapStr{}} 79 | event.PutValue("a", "2") 80 | 81 | assert.False(t, cl.Check(event)) 82 | }) 83 | } 84 | 85 | func BenchmarkAllow(b *testing.B) { 86 | b.Run("never overflows", func(b *testing.B) { 87 | t := time.Now() 88 | cl, _ := NewConditionLimiter( 89 | map[string]string{}, 90 | 1, 91 | int64(b.N), 92 | 10, 93 | t, 94 | ) 95 | 96 | for i := 0; i < b.N; i++ { 97 | cl.Allow(t) 98 | } 99 | }) 100 | 101 | b.Run("always overflows", func(b *testing.B) { 102 | t := time.Now() 103 | cl, _ := NewConditionLimiter( 104 | map[string]string{}, 105 | 1, 106 | 0, 107 | 10, 108 | t, 109 | ) 110 | 111 | for i := 0; i < b.N; i++ { 112 | cl.Allow(t) 113 | } 114 | }) 115 | 116 | b.Run("always shifts", func(b *testing.B) { 117 | t := time.Now() 118 | cl, _ := NewConditionLimiter( 119 | map[string]string{}, 120 | 1, 121 | 0, 122 | 10, 123 | t, 124 | ) 125 | 126 | for i := 0; i < b.N; i++ { 127 | cl.Allow(t.Add(time.Duration(i) * time.Second)) 128 | } 129 | }) 130 | } 131 | 132 | func TestConditionLimiter_WriteStatus(t *testing.T) { 133 | current := time.Local 134 | time.Local = time.UTC 135 | defer func() { 136 | time.Local = current 137 | }() 138 | tt := time.Date(2018, 12, 19, 19, 30, 25, 0, time.UTC) 139 | cl, _ := NewConditionLimiter(map[string]string{"a": "1"}, 60, 2, 2, tt) 140 | cl.Allow(tt) 141 | 142 | var b bytes.Buffer 143 | cl.WriteStatus(&b) 144 | 145 | assert.Equal(t, `map[a:1] 146 | #2018-12-19 19:29:00 +0000 UTC: [____________________] 0/2 147 | #2018-12-19 19:30:00 +0000 UTC: [##########__________] 1/2 148 | `, b.String()) 149 | } 150 | 151 | func TestProgress(t *testing.T) { 152 | var b bytes.Buffer 153 | 154 | progress(&b, 3, 5, 10) 155 | assert.Equal(t, "[######____]", b.String()) 156 | } 157 | 158 | func BenchmarkTimeToMinute(b *testing.B) { 159 | f := time.Now() 160 | for i := 0; i < b.N; i++ { 161 | timeToBucketID(f, 1) 162 | } 163 | } 164 | 165 | func BenchmarkIDToTime(b *testing.B) { 166 | for i := 0; i < b.N; i++ { 167 | bucketIDToTime(1, 1) 168 | } 169 | } 170 | -------------------------------------------------------------------------------- /limit.go: -------------------------------------------------------------------------------- 1 | package throttleplugin 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "io" 7 | "io/ioutil" 8 | "net/http" 9 | "sync" 10 | "time" 11 | 12 | "github.com/elastic/beats/libbeat/beat" 13 | "github.com/elastic/beats/libbeat/logp" 14 | "github.com/pkg/errors" 15 | "gopkg.in/yaml.v2" 16 | ) 17 | 18 | type RemoteConfig struct { 19 | Key string `yaml:"key"` 20 | DefaultLimit int64 `yaml:"default_limit"` 21 | Rules []RuleConfig `yaml:"rules"` 22 | } 23 | 24 | type RuleConfig struct { 25 | Limit int64 `yaml:"limit"` 26 | Selectors map[string]string `yaml:"selectors"` 27 | } 28 | 29 | type RemoteLimiter struct { 30 | url string 31 | client *http.Client 32 | bucketInterval int64 33 | buckets int64 34 | 35 | mu sync.RWMutex 36 | key string 37 | rules []Rule 38 | limiters map[string]*BucketLimiter 39 | } 40 | 41 | // NewRemoteLimiter creates new remote limiter instance. 42 | func NewRemoteLimiter(url string, bucketInterval, buckets int64) (*RemoteLimiter, error) { 43 | rl := &RemoteLimiter{ 44 | url: url, 45 | client: http.DefaultClient, 46 | bucketInterval: bucketInterval, 47 | buckets: buckets, 48 | limiters: make(map[string]*BucketLimiter), 49 | } 50 | 51 | return rl, nil 52 | } 53 | 54 | // Allow returns TRUE if event is allowed to be processed. 55 | func (rl *RemoteLimiter) Allow(e *beat.Event) bool { 56 | var ts time.Time 57 | 58 | if tsString, err := e.GetValue("ts"); err == nil { 59 | ts, _ = time.Parse(time.RFC3339, tsString.(string)) 60 | } 61 | 62 | if ts.IsZero() { 63 | ts = time.Now() 64 | } 65 | 66 | keyValue, _ := e.GetValue(rl.key) 67 | kv := fmt.Sprintf("%v", keyValue) 68 | 69 | rl.mu.Lock() 70 | defer rl.mu.Unlock() 71 | 72 | for _, r := range rl.rules { 73 | if matched, key := r.Match(e); matched { 74 | key = kv + key 75 | // check if we already have limiter 76 | limiter, ok := rl.limiters[key] 77 | if !ok { 78 | limiter = NewBucketLimiter(rl.bucketInterval, r.Limit(), rl.buckets, ts) 79 | rl.limiters[key] = limiter 80 | } 81 | 82 | return limiter.Allow(ts) 83 | } 84 | } 85 | 86 | return true 87 | } 88 | 89 | // Update retrieves policies from Policy Manager. 90 | func (rl *RemoteLimiter) Update(ctx context.Context) error { 91 | r, err := http.NewRequest("GET", rl.url, nil) 92 | if err != nil { 93 | return errors.Wrap(err, "failed to create request") 94 | } 95 | 96 | res, err := rl.client.Do(r) 97 | if err != nil { 98 | return err 99 | } 100 | defer res.Body.Close() 101 | 102 | body, err := ioutil.ReadAll(res.Body) 103 | if err != nil { 104 | return err 105 | } 106 | 107 | var c RemoteConfig 108 | 109 | if err := yaml.Unmarshal(body, &c); err != nil { 110 | return errors.Wrap(err, "failed to unpack config") 111 | } 112 | 113 | rules := make([]Rule, 0, len(c.Rules)+1) 114 | 115 | for _, l := range c.Rules { 116 | rules = append(rules, NewRule(l.Selectors, l.Limit)) 117 | } 118 | 119 | defaultRule := NewRule(map[string]string{}, c.DefaultLimit) 120 | rules = append(rules, defaultRule) 121 | 122 | limiterTTL := time.Duration(rl.bucketInterval*rl.buckets) * time.Second 123 | limiterThreshold := time.Now().Add(-limiterTTL) 124 | 125 | rl.mu.Lock() 126 | defer rl.mu.Unlock() 127 | rl.key = c.Key 128 | rl.rules = rules 129 | for id, l := range rl.limiters { 130 | if l.LastUpdate().Before(limiterThreshold) { 131 | delete(rl.limiters, id) 132 | } 133 | } 134 | 135 | return nil 136 | } 137 | 138 | // UpdateWithInterval runs update with some interval. 139 | func (rl *RemoteLimiter) UpdateWithInterval(ctx context.Context, interval time.Duration) error { 140 | t := time.NewTicker(interval) 141 | defer t.Stop() 142 | 143 | for { 144 | select { 145 | case <-ctx.Done(): 146 | return ctx.Err() 147 | case <-t.C: 148 | if err := rl.Update(ctx); err != nil { 149 | logp.Err("failed to update limit policies: %v", err) 150 | } 151 | } 152 | } 153 | } 154 | 155 | func (rl *RemoteLimiter) WriteStatus(w io.Writer) error { 156 | rl.mu.Lock() 157 | defer rl.mu.Unlock() 158 | 159 | for key, cl := range rl.limiters { 160 | fmt.Fprintf(w, "#%v\n\n", key) 161 | if err := cl.WriteStatus(w); err != nil { 162 | return err 163 | } 164 | fmt.Fprintln(w, "---------") 165 | } 166 | 167 | fmt.Fprintf(w, "rules: \n\n%#v", rl.rules) 168 | 169 | return nil 170 | } 171 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973 h1:xJ4a3vCFaGF/jqvzLMYoU8P317H5OQ+Via4RmuPwCS0= 2 | github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q= 3 | github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= 4 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 5 | github.com/elastic/beats v6.6.2+incompatible h1:bxz0RLQYMqhsVN7vk2VYpR8EloyUNCyvEzBd9lYx9eU= 6 | github.com/elastic/beats v6.6.2+incompatible/go.mod h1:7cX7zGsOwJ01FLkZs9Tg5nBdnQi6XB3hYAyWekpKgeY= 7 | github.com/elastic/go-ucfg v0.7.0 h1:1+C/sZdJKww8hKl7XtLPTjs4cFslhQF2fazKTF+ZE+4= 8 | github.com/elastic/go-ucfg v0.7.0/go.mod h1:iaiY0NBIYeasNgycLyTvhJftQlQEUO2hpF+FX0JKxzo= 9 | github.com/gofrs/uuid v3.2.0+incompatible h1:y12jRkkFxsd7GpqdSZ+/KCs/fJbqpEXSGd4+jfEaewE= 10 | github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM= 11 | github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM= 12 | github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= 13 | github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM= 14 | github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8= 15 | github.com/joeshaw/multierror v0.0.0-20140124173710-69b34d4ec901 h1:rp+c0RAYOWj8l6qbCUTSiRLG/iKnW3K3/QfPPuSsBt4= 16 | github.com/joeshaw/multierror v0.0.0-20140124173710-69b34d4ec901/go.mod h1:Z86h9688Y0wesXCyonoVr47MasHilkuLMqGhRZ4Hpak= 17 | github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU= 18 | github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= 19 | github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I= 20 | github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 21 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 22 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 23 | github.com/prometheus/client_golang v0.9.2 h1:awm861/B8OKDd2I/6o1dy3ra4BamzKhYOiGItCeZ740= 24 | github.com/prometheus/client_golang v0.9.2/go.mod h1:OsXs2jCmiKlQ1lTBmv21f2mNfw4xf/QclQDMrYNZzcM= 25 | github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910 h1:idejC8f05m9MGOsuEi1ATq9shN03HrxNkD/luQvxCv8= 26 | github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo= 27 | github.com/prometheus/common v0.0.0-20181126121408-4724e9255275 h1:PnBWHBf+6L0jOqq0gIVUe6Yk0/QMZ640k6NvkxcBf+8= 28 | github.com/prometheus/common v0.0.0-20181126121408-4724e9255275/go.mod h1:daVV7qP5qjZbuso7PdcryaAu0sAZbrN9i7WWcTMWvro= 29 | github.com/prometheus/procfs v0.0.0-20181204211112-1dc9a6cbc91a h1:9a8MnZMP0X2nLJdBg+pBmGgkJlSaKC2KaQmTCk1XDtE= 30 | github.com/prometheus/procfs v0.0.0-20181204211112-1dc9a6cbc91a/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk= 31 | github.com/spf13/cobra v0.0.3 h1:ZlrZ4XsMRm04Fr5pSFxBgfND2EBVa1nLpiy1stUsX/8= 32 | github.com/spf13/cobra v0.0.3/go.mod h1:1l0Ry5zgKvJasoi3XT1TypsSe7PqH0Sj9dhYf7v3XqQ= 33 | github.com/spf13/pflag v1.0.3 h1:zPAT6CGy6wXeQ7NtTnaTerfKOsV6V6F8agHXFiazDkg= 34 | github.com/spf13/pflag v1.0.3/go.mod h1:DYY7MBk1bdzusC3SYhjObp+wFpr4gzcvqqNjLnInEg4= 35 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 36 | github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q= 37 | github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= 38 | go.uber.org/atomic v1.3.2 h1:2Oa65PReHzfn29GpvgsYwloV9AVFHPDk8tYxt2c2tr4= 39 | go.uber.org/atomic v1.3.2/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE= 40 | go.uber.org/multierr v1.1.0 h1:HoEmRHQPVSqub6w2z2d2EOVs2fjyFRGyofhKuyDq0QI= 41 | go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0= 42 | go.uber.org/zap v1.9.1 h1:XCJQEf3W6eZaVwhRBof6ImoYGJSITeKWsyeh3HFu/5o= 43 | go.uber.org/zap v1.9.1/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q= 44 | golang.org/x/net v0.0.0-20181201002055-351d144fa1fc/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= 45 | golang.org/x/sync v0.0.0-20181108010431-42b317875d0f h1:Bl/8QSvNqXvPGPGXa2z5xUTmV7VDcZyvRZ+QQXkXTZQ= 46 | golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 47 | golang.org/x/sys v0.0.0-20190321052220-f7bb7a8bee54 h1:xe1/2UUJRmA9iDglQSlkx8c5n3twv58+K0mPpC2zmhA= 48 | golang.org/x/sys v0.0.0-20190321052220-f7bb7a8bee54/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 49 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 50 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 51 | gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw= 52 | gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= 53 | -------------------------------------------------------------------------------- /plugin.go: -------------------------------------------------------------------------------- 1 | package throttleplugin 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "net/http" 7 | _ "net/http/pprof" 8 | "sync" 9 | "time" 10 | 11 | "github.com/elastic/beats/libbeat/beat" 12 | "github.com/elastic/beats/libbeat/common" 13 | "github.com/elastic/beats/libbeat/logp" 14 | "github.com/elastic/beats/libbeat/processors" 15 | "github.com/pkg/errors" 16 | "github.com/prometheus/client_golang/prometheus" 17 | "github.com/prometheus/client_golang/prometheus/promhttp" 18 | ) 19 | 20 | var unknownValue = "UNKNOWN" 21 | 22 | // Config defines processor configuration. 23 | type Config struct { 24 | PolicyHost string `config:"policy_host"` 25 | PolicyUpdateInterval time.Duration `config:"policy_update_interval"` 26 | PrometheusPort int `config:"prometheus_port"` 27 | 28 | BucketSize int64 `config:"bucket_size"` 29 | Buckets int64 `config:"buckets"` 30 | 31 | MetricName string `config:"metric_name"` 32 | MetricLabels []LabelMapping `config:"metric_labels"` 33 | } 34 | 35 | type LabelMapping struct { 36 | From string `config:"from"` 37 | To string `config:"to"` 38 | } 39 | 40 | func (c Config) GetMetricLabels() []string { 41 | labels := make([]string, len(c.MetricLabels)) 42 | for i, lm := range c.MetricLabels { 43 | labels[i] = lm.To 44 | } 45 | 46 | return labels 47 | } 48 | 49 | func (c Config) GetFields() []string { 50 | fields := make([]string, len(c.MetricLabels)) 51 | for i, lm := range c.MetricLabels { 52 | fields[i] = lm.From 53 | } 54 | 55 | return fields 56 | } 57 | 58 | // compile-time check that Processor implements processors.Processor interface. 59 | var _ processors.Processor = &Processor{} 60 | 61 | // Processor make rate-limiting for messages. 62 | type Processor struct { 63 | metric *prometheus.CounterVec 64 | fields []string // list of fields that will be used as labels in metrics 65 | valuesPool sync.Pool // used for zero allocation metric observations 66 | 67 | limiter *RemoteLimiter 68 | throttled int64 69 | 70 | httpServer *http.Server 71 | } 72 | 73 | // NewProcessor returns new processor instance. 74 | func NewProcessor(cfg *common.Config) (processors.Processor, error) { 75 | return newProcessor(cfg) 76 | } 77 | 78 | func newProcessor(cfg *common.Config) (*Processor, error) { 79 | var c Config 80 | if err := cfg.Unpack(&c); err != nil { 81 | return nil, fmt.Errorf("failed to unpack config: %v", err) 82 | } 83 | 84 | vec := prometheus.NewCounterVec( 85 | prometheus.CounterOpts{ 86 | Namespace: "filebeat", 87 | Name: c.MetricName, 88 | Help: c.MetricName, 89 | }, 90 | append(c.GetMetricLabels(), "throttled"), 91 | ) 92 | 93 | go func() { 94 | t := time.NewTicker(time.Minute) 95 | defer t.Stop() 96 | 97 | for range t.C { 98 | vec.Reset() 99 | } 100 | }() 101 | prometheus.MustRegister(vec) 102 | 103 | limiter, err := NewRemoteLimiter(c.PolicyHost, c.BucketSize, c.Buckets) 104 | if err != nil { 105 | return nil, errors.Wrap(err, "failed to create RemoteLimiter") 106 | } 107 | 108 | processor := &Processor{ 109 | metric: vec, 110 | fields: c.GetFields(), 111 | valuesPool: sync.Pool{ 112 | New: func() interface{} { 113 | // "+1" is used because last label is always "throttled". 114 | size := len(c.GetFields()) + 1 115 | return make([]string, size) 116 | }, 117 | }, 118 | limiter: limiter, 119 | } 120 | 121 | logp.Info("listening prometheus handler on port: %v", c.PrometheusPort) 122 | processor.RunHTTPHandlers(c.PrometheusPort) 123 | 124 | logp.Info("initial update for policies...") 125 | ctx, cancel := context.WithTimeout(context.Background(), time.Second) 126 | defer cancel() 127 | 128 | if err := limiter.Update(ctx); err != nil { 129 | logp.Err("failed to make initial policy update: %v. Using default", err) 130 | } 131 | 132 | logp.Info("limit policy url: %v, updateInterval: %v", c.PolicyHost, c.PolicyUpdateInterval) 133 | go limiter.UpdateWithInterval(context.Background(), c.PolicyUpdateInterval) 134 | 135 | return processor, nil 136 | } 137 | 138 | // RunHTTPHandlers runs prometheus handler on specified port. 139 | func (mp *Processor) RunHTTPHandlers(port int) { 140 | h := promhttp.Handler() 141 | mux := http.NewServeMux() 142 | mux.Handle("/metrics", h) 143 | mux.HandleFunc("/status", func(w http.ResponseWriter, r *http.Request) { 144 | w.Header().Add("Content-Type", "text/plain") 145 | mp.limiter.WriteStatus(w) 146 | }) 147 | logp.Info("starting prometheus handler on :%v", port) 148 | srv := &http.Server{Addr: fmt.Sprintf(":%v", port), Handler: mux} 149 | mp.httpServer = srv 150 | 151 | go func() { 152 | if err := srv.ListenAndServe(); err != nil { 153 | logp.Err("failed to run prometheus handler: %v", err) 154 | } 155 | }() 156 | } 157 | 158 | func (mp *Processor) Close() error { 159 | return mp.httpServer.Close() 160 | } 161 | 162 | func (mp *Processor) Run(event *beat.Event) (*beat.Event, error) { 163 | v := mp.valuesPool.Get() 164 | values := v.([]string) 165 | defer func() { 166 | // no reslicing is required before returning to pool because slice size is always the same. 167 | mp.valuesPool.Put(v) 168 | }() 169 | 170 | for i, f := range mp.fields { 171 | v, err := event.GetValue(f) 172 | if err == common.ErrKeyNotFound { 173 | values[i] = unknownValue 174 | continue 175 | } 176 | 177 | switch t := v.(type) { 178 | case string: 179 | values[i] = t 180 | default: 181 | values[i] = fmt.Sprintf("%v", v) 182 | } 183 | } 184 | 185 | if !mp.limiter.Allow(event) { 186 | mp.throttled++ 187 | values[len(values)-1] = "y" 188 | mp.metric.WithLabelValues(values...).Inc() 189 | return nil, nil 190 | } 191 | 192 | mp.throttled = 0 193 | values[len(values)-1] = "n" 194 | mp.metric.WithLabelValues(values...).Inc() 195 | 196 | return event, nil 197 | } 198 | 199 | func (mp *Processor) update(ctx context.Context) error { 200 | return nil 201 | } 202 | 203 | func (mp *Processor) String() string { 204 | return "metrics_processor" 205 | } 206 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![Build Status](https://travis-ci.org/ozonru/filebeat-throttle-plugin.svg?branch=master)](https://travis-ci.org/ozonru/filebeat-throttle-plugin) 2 | [![Go Report Card](https://goreportcard.com/badge/github.com/ozonru/filebeat-throttle-plugin)](https://goreportcard.com/report/github.com/ozonru/filebeat-throttle-plugin) 3 | 4 | # filebeat-throttle-plugin 5 | 6 | This plugins allows to throttle beat events. 7 | 8 | Throttle processor retrieves configuration from separate component called `policy manager`. 9 | 10 | 11 | ``` 12 | ┌──────┐ ┌──────┐ ┌──────┐ 13 | │Node 1│ │Node 2│ │Node 3│ 14 | ├──────┴─────────────┐ ├──────┴─────────────┐ ├──────┴─────────────┐ 15 | │ ┌──────┐ ┌──────┐ │ │ ┌──────┐ ┌──────┐ │ │ ┌──────┐ ┌──────┐ │ 16 | │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ 17 | │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ 18 | │ └──────┘ └──────┘ │ │ └──────┘ └──────┘ │ │ └──────┘ └──────┘ │ 19 | │ ┌──────┐ ┌──────┐ │ │ ┌──────┐ ┌──────┐ │ │ ┌──────┐ ┌──────┐ │ 20 | │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ 21 | │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ 22 | │ └──────┘ └──────┘ │ │ └──────┘ └──────┘ │ │ └──────┘ └──────┘ │ 23 | │ ┌──────┐ ┌──────┐ │ │ ┌──────┐ ┌──────┐ │ │ ┌──────┐ ┌──────┐ │ 24 | │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ 25 | │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ 26 | │ └──────┘ └──────┘ │ │ └──────┘ └──────┘ │ │ └──────┘ └──────┘ │ 27 | │ │ │ │ │ │ 28 | │ │ │ │ │ │ 29 | │ ┌────────────────┐ │ │ ┌────────────────┐ │ │ ┌────────────────┐ │ 30 | │ │ filebeat │ │ │ │ filebeat │ │ │ │ filebeat │ │ 31 | │ └────────────────┘ │ │ └────────────────┘ │ │ └────────────────┘ │ 32 | │ │ │ │ │ │ │ │ │ 33 | └──────────┼─────────┘ └──────────┼─────────┘ └──────────┼─────────┘ 34 | │ │ │ 35 | └───────────────┐ │ ┌───────────────┘ 36 | │ │ │ 37 | │ │ │ 38 | ▼ ▼ ▼ 39 | ┌─────────────────────────────┐ 40 | │ │ 41 | │ Policy Manager │ 42 | │ │ 43 | └─────────────────────────────┘ 44 | HTTP /policy endpoint 45 | ``` 46 | 47 | ## Configuration 48 | 49 | To enable throttling you have to add `throttle` processor to configuration: 50 | 51 | ``` 52 | - throttle: 53 | prometheus_port: 9090 54 | metric_name: proccessed_records 55 | metric_labels: 56 | - from: kubernetes_container_name 57 | to: container_name 58 | - from: "labels.app" 59 | to: app 60 | policy_host: "http://policymanager.local:8080/policy" 61 | policy_update_interval: 1s 62 | bucket_size: 1 63 | buckets: 1000 64 | ``` 65 | 66 | - `prometheus_port` - prometheus metrics handler to listen on 67 | - `metric_name` - name of counter metric with number of processed/throttled events 68 | - `metric_labels` - additional fields that will be converted to metric labels 69 | - `policy_host` - policy manager host 70 | - `policy_update_interval` - how often processor refresh policies 71 | - `buckets` - number of buckets 72 | - `bucket_size` - bucket duration (in seconds) 73 | 74 | ## Policy Manager 75 | 76 | Policy manager exposes configuration by `/policy` endpoint in following format: 77 | ```yaml 78 | --- 79 | limits: 80 | - value: 500 81 | conditions: 82 | kubernetes_container_name: "simple-generator" 83 | - value: 5000 84 | conditions: 85 | kubernetes_namespace: "bx" 86 | ``` 87 | 88 | `value` specifies maximum number of events that will be passed in interval `bucket_size`. 89 | In `conditions` section you use any fields from your events. All conditions works as `equal`. 90 | 91 | 92 | ## Throttling algorithm 93 | 94 | We use `token bucket` algorithm for throttling. In the simplest way we can use only single bucket for events limit. But in real life beats can be down (maintance, some failures, etc): in this case all events (new and old ones) use tokens from same bucket and some events can be skipped because of overflow. To avoid such situations we need to keep N last bucket and use event timestamp to choose bucket. 95 | 96 | Imagine that processor has following configuration: 97 | ``` 98 | bucket_size: 1 // 1 second 99 | buckets: 5 100 | ``` 101 | 102 | And rule with `limit: 10` 103 | 104 | Then in the time moment T we'll have buckets: 105 | ``` 106 | T - 5 T - 4 T - 3 T - 2 T - 1 T 107 | ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ 108 | │ 0/10 │ │ 0/10 │ │ 0/10 │ │ 0/10 │ │ 0/10 │ 109 | └──────────┘ └──────────┘ └──────────┘ └──────────┘ └──────────┘ 110 | ``` 111 | 112 | After adding event in interval between `T-3` and `T-2`: 113 | ``` 114 | T - 5 T - 4 T - 3 T - 2 T - 1 T 115 | ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ 116 | │ 0/10 │ │ 0/10 │ │ 1/10 │ │ 0/10 │ │ 0/10 │ 117 | └──────────┘ └──────────┘ └──────────┘ └──────────┘ └──────────┘ 118 | ``` 119 | 120 | If one one of buckets overflows all new events will be ignored. 121 | 122 | At `T + 1` all buckets shift left: 123 | ``` 124 | T - 4 T - 3 T - 2 T - 1 T T + 1 125 | ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ 126 | │ 0/10 │ │ 1/10 │ │ 0/10 │ │ 0/10 │ │ 0/10 │ 127 | └──────────┘ └──────────┘ └──────────┘ └──────────┘ └──────────┘ 128 | ``` 129 | 130 | All events with timestamp earlier than `T-4` will be ignored. 131 | 132 | 133 | ## Building 134 | 135 | There are two options how you can build and use throttle processor: 136 | - build beat binary with throttle processor inside 137 | - use it as separate plugin 138 | 139 | ### Compile-in 140 | 141 | You have to copy `register/compile/plugin.go` to the `main` package of required beat. 142 | ``` 143 | cp $GOPATH/github.com/ozonru/filebeat-throttle-plugin/register/compile/plugin.go $GOPATH/github.com/elastic/beats/filebeat 144 | go build github.com/elastic/beats/filebeat 145 | ``` 146 | 147 | ### Plugin 148 | 149 | You can build plugin both for linux and MacOS: 150 | 151 | ``` 152 | make linux_plugin 153 | make darwin_plugin 154 | ``` 155 | 156 | To use plugin make sure that you have the same Go version as beat binary was built. 157 | 158 | ## Log generator (for development) 159 | 160 | In `generator` you can find simple log generator that can be used for local development & testing). 161 | ``` 162 | ➜ go get -u gitlab.ozon.ru/sre/filebeat-ratelimit-plugin/generator 163 | ➜ generator -h 164 | Usage of generator: 165 | -id id 166 | id field value (default "generator") 167 | -rate float 168 | records per second (default 10) 169 | Usage of generator: 170 | -id id 171 | id field value (default "generator") 172 | -rate float 173 | records per second (default 10) 174 | ➜ generator -id foo 175 | {"ts":"2019-01-09T18:59:56.109522+03:00", "message": "1", "id":"foo"} 176 | {"ts":"2019-01-09T18:59:56.207788+03:00", "message": "2", "id":"foo"} 177 | {"ts":"2019-01-09T18:59:56.310223+03:00", "message": "3", "id":"foo"} 178 | {"ts":"2019-01-09T18:59:56.409879+03:00", "message": "4", "id":"foo"} 179 | {"ts":"2019-01-09T18:59:56.509572+03:00", "message": "5", "id":"foo"} 180 | {"ts":"2019-01-09T18:59:56.608653+03:00", "message": "6", "id":"foo"} 181 | {"ts":"2019-01-09T18:59:56.708547+03:00", "message": "7", "id":"foo"} 182 | {"ts":"2019-01-09T18:59:56.809872+03:00", "message": "8", "id":"foo"} 183 | ^C 184 | ``` 185 | --------------------------------------------------------------------------------