Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(libbeat): mitigate race condition in ratelimit processor #42966

Open
wants to merge 12 commits into
base: main
Choose a base branch
from
1 change: 1 addition & 0 deletions CHANGELOG.next.asciidoc
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ https://github.com/elastic/beats/compare/v8.8.1\...main[Check the HEAD diff]
- Fix templates and docs to use correct `--` version of command line arguments. {issue}42038[42038] {pull}42060[42060]
- removed support for a single `-` to precede multi-letter command line arguments. Use `--` instead. {issue}42117[42117] {pull}42209[42209]
- Removed encryption from diskqueue V2 for fips compliance {issue}4534[4534]{pull}42848[42848]
- Fix race conditions in the global ratelimit processor that could drop events or apply rate limiting incorrectly.

*Auditbeat*

Expand Down
13 changes: 13 additions & 0 deletions libbeat/processors/ratelimit/rate_limit_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -176,3 +176,16 @@ func TestRateLimit(t *testing.T) {
})
}
}

func BenchmarkRateLimit(b *testing.B) {
p, err := new(conf.MustNewConfigFrom(mapstr.M{
"limit": "100/s",
}))
require.NoError(b, err)
event := beat.Event{Fields: mapstr.M{"field": 1}}

b.ResetTimer()
for i := 0; i < b.N; i++ {
p.Run(&event) //nolint:errcheck // ignore
}
}
66 changes: 38 additions & 28 deletions libbeat/processors/ratelimit/token_bucket.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,10 +35,36 @@ func init() {
}

type bucket struct {
mu sync.Mutex

tokens float64
lastReplenish time.Time
}

func (b *bucket) withdraw() bool {
b.mu.Lock()
defer b.mu.Unlock()

if b.tokens < 1 {
return false
}

b.tokens--
return true
}

func (b *bucket) replenish(rate rate, clock clockwork.Clock) float64 {
b.mu.Lock()
defer b.mu.Unlock()

secsSinceLastReplenish := clock.Now().Sub(b.lastReplenish).Seconds()
tokensToReplenish := secsSinceLastReplenish * rate.valuePerSecond()

b.tokens += tokensToReplenish
b.lastReplenish = clock.Now()
return b.tokens
}

type tokenBucket struct {
mu unison.Mutex

Expand Down Expand Up @@ -122,35 +148,20 @@ func (t *tokenBucket) setClock(c clockwork.Clock) {
}

func (t *tokenBucket) getBucket(key uint64) *bucket {
v, exists := t.buckets.LoadOrStore(key, &bucket{
tokens: t.depth,
lastReplenish: t.clock.Now(),
})
//nolint:errcheck // ignore
b := v.(*bucket)

v, exists := t.buckets.Load(key)
if exists {
//nolint:errcheck // ignore
b := v.(*bucket)
b.replenish(t.limit, t.clock)
return b
}

return b
}

func (b *bucket) withdraw() bool {
if b.tokens < 1 {
return false
}
b.tokens--
return true
}

func (b *bucket) replenish(rate rate, clock clockwork.Clock) {
secsSinceLastReplenish := clock.Now().Sub(b.lastReplenish).Seconds()
tokensToReplenish := secsSinceLastReplenish * rate.valuePerSecond()

b.tokens += tokensToReplenish
b.lastReplenish = clock.Now()
v, _ = t.buckets.LoadOrStore(key, &bucket{
tokens: t.depth,
lastReplenish: t.clock.Now(),
})
//nolint:errcheck // ignore
return v.(*bucket)
}

func (t *tokenBucket) runGC() {
Expand All @@ -177,9 +188,8 @@ func (t *tokenBucket) runGC() {
//nolint:errcheck // ignore
b := v.(*bucket)

b.replenish(t.limit, t.clock)

if b.tokens >= t.depth {
tokens := b.replenish(t.limit, t.clock)
if tokens >= t.depth {
toDelete = append(toDelete, key)
}

Expand All @@ -193,7 +203,7 @@ func (t *tokenBucket) runGC() {
}

// Reset GC metrics
t.gc.metrics.numCalls = atomic.Uint64{}
t.gc.metrics.numCalls.Store(0)

gcDuration := time.Since(gcStartTime)
numBucketsDeleted := len(toDelete)
Expand Down
Loading