consul/agent/consul/multilimiter/multilimiter_test.go
hashicorp-copywrite[bot] 5fb9df1640
[COMPLIANCE] License changes (#18443)
* Adding explicit MPL license for sub-package

This directory and its subdirectories (packages) contain files licensed with the MPLv2 `LICENSE` file in this directory and are intentionally licensed separately from the BSL `LICENSE` file at the root of this repository.

* Adding explicit MPL license for sub-package

This directory and its subdirectories (packages) contain files licensed with the MPLv2 `LICENSE` file in this directory and are intentionally licensed separately from the BSL `LICENSE` file at the root of this repository.

* Updating the license from MPL to Business Source License

Going forward, this project will be licensed under the Business Source License v1.1. Please see our blog post for more details at <Blog URL>, FAQ at www.hashicorp.com/licensing-faq, and details of the license at www.hashicorp.com/bsl.

* add missing license headers

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

* Update copyright file headers to BUSL-1.1

---------

Co-authored-by: hashicorp-copywrite[bot] <110428419+hashicorp-copywrite[bot]@users.noreply.github.com>
2023-08-11 09:12:13 -04:00

689 lines
20 KiB
Go

// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: BUSL-1.1
package multilimiter
import (
"bytes"
"context"
"encoding/binary"
"fmt"
"math/rand"
"sync"
"testing"
"time"
"github.com/hashicorp/consul/sdk/testutil/retry"
radix "github.com/hashicorp/go-immutable-radix"
"github.com/stretchr/testify/require"
"golang.org/x/time/rate"
)
type Limited struct {
key KeyType
}
func (l Limited) Key() []byte {
return l.key
}
func TestNewMultiLimiter(t *testing.T) {
c := Config{}
m := NewMultiLimiter(c)
require.NotNil(t, m)
require.NotNil(t, m.limiters)
}
func TestRateLimiterUpdate(t *testing.T) {
c := Config{ReconcileCheckLimit: 1 * time.Hour, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
key := Key([]byte("test"))
c1 := LimiterConfig{Rate: 10}
m.UpdateConfig(c1, key)
//Allow a key
m.Allow(Limited{key: key})
storeLimiter(m)
limiters := m.limiters.Load()
l1, ok1 := limiters.Get(key)
retry.RunWith(&retry.Timer{Wait: 100 * time.Millisecond, Timeout: 2 * time.Second}, t, func(r *retry.R) {
limiters = m.limiters.Load()
l1, ok1 = limiters.Get(key)
// check key exist
require.True(r, ok1)
require.NotNil(r, l1)
})
la1 := l1.(*Limiter).lastAccess.Load()
// Sleep a bit just to make sure time change
time.Sleep(time.Millisecond)
// allow same key again
m.Allow(Limited{key: key})
limiters = m.limiters.Load()
l2, ok2 := limiters.Get(key)
// check it exist and it's same key
require.True(t, ok2)
require.NotNil(t, l2)
require.Equal(t, l1, l2)
// last access should be different
la2 := l1.(*Limiter).lastAccess.Load()
require.NotEqual(t, la1, la2)
}
func TestRateLimiterCleanup(t *testing.T) {
// Create a limiter and Allow a key, check the key exists
c := Config{ReconcileCheckLimit: 1 * time.Second, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
limiters := m.limiters.Load()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
m.Run(ctx)
key := Key([]byte("test"))
m.UpdateConfig(LimiterConfig{Rate: 0.1}, key)
m.Allow(Limited{key: key})
retry.RunWith(&retry.Timer{Wait: 100 * time.Millisecond, Timeout: 10 * time.Second}, t, func(r *retry.R) {
l := m.limiters.Load()
require.NotEqual(r, limiters, l)
limiters = l
})
retry.RunWith(&retry.Timer{Wait: 100 * time.Millisecond, Timeout: 10 * time.Second}, t, func(r *retry.R) {
v, ok := limiters.Get(key)
require.True(r, ok)
require.NotNil(r, v)
})
time.Sleep(c.ReconcileCheckInterval)
// Wait > ReconcileCheckInterval and check that the key was cleaned up
retry.RunWith(&retry.Timer{Wait: 100 * time.Millisecond, Timeout: 2 * time.Second}, t, func(r *retry.R) {
l := m.limiters.Load()
require.NotEqual(r, limiters, l)
limiters = l
v, ok := limiters.Get(key)
require.False(r, ok)
require.Nil(r, v)
})
}
func storeLimiter(m *MultiLimiter) {
txn := m.limiters.Load().Txn()
mockTicker := mockTicker{tickerCh: make(chan time.Time, 1)}
ctx := context.Background()
reconcileCheckLimit := m.defaultConfig.Load().ReconcileCheckLimit
m.reconcile(ctx, &mockTicker, txn, reconcileCheckLimit)
mockTicker.tickerCh <- time.Now()
m.reconcile(ctx, &mockTicker, txn, reconcileCheckLimit)
}
func reconcile(m *MultiLimiter) {
txn := m.limiters.Load().Txn()
mockTicker := mockTicker{tickerCh: make(chan time.Time, 1)}
ctx := context.Background()
reconcileCheckLimit := m.defaultConfig.Load().ReconcileCheckLimit
mockTicker.tickerCh <- time.Now()
txn = m.reconcile(ctx, &mockTicker, txn, reconcileCheckLimit)
m.limiters.Store(txn.Commit())
}
func TestRateLimiterStore(t *testing.T) {
// Create a MultiLimiter m with a defaultConfig c and check the defaultConfig is applied
t.Run("Store multiple transactions", func(t *testing.T) {
c := Config{ReconcileCheckLimit: 100 * time.Millisecond, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
ipNoPrefix1 := Key([]byte(""), []byte("127.0.0.1"))
c1 := LimiterConfig{Rate: 1}
ipNoPrefix2 := Key([]byte(""), []byte("127.0.0.2"))
c2 := LimiterConfig{Rate: 2}
{
// Update config for ipNoPrefix1 and check it's applied
m.UpdateConfig(c1, ipNoPrefix1)
m.Allow(ipLimited{key: ipNoPrefix1})
storeLimiter(m)
l, ok := m.limiters.Load().Get(ipNoPrefix1)
require.True(t, ok)
require.NotNil(t, l)
limiter := l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
}
{
// Update config for ipNoPrefix2 and check it's applied
m.UpdateConfig(c2, ipNoPrefix2)
m.Allow(ipLimited{key: ipNoPrefix2})
storeLimiter(m)
l, ok := m.limiters.Load().Get(ipNoPrefix2)
require.True(t, ok)
require.NotNil(t, l)
limiter := l.(*Limiter)
require.True(t, c2.isApplied(limiter.limiter))
//Check that ipNoPrefix1 is unchanged
l, ok = m.limiters.Load().Get(ipNoPrefix1)
require.True(t, ok)
require.NotNil(t, l)
limiter = l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
}
})
t.Run("runStore store multiple Limiters", func(t *testing.T) {
c := Config{ReconcileCheckLimit: 10 * time.Second, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
ctx, cancel := context.WithCancel(context.Background())
m.Run(ctx)
defer cancel()
// Create a limiter for ipNoPrefix1
ipNoPrefix1 := Key([]byte(""), []byte("127.0.0.1"))
c1 := LimiterConfig{Rate: 1}
limiters := m.limiters.Load()
m.UpdateConfig(c1, ipNoPrefix1)
m.Allow(ipLimited{key: ipNoPrefix1})
retry.RunWith(&retry.Timer{Wait: 1 * time.Second, Timeout: 5 * time.Second}, t, func(r *retry.R) {
l := m.limiters.Load()
require.NotEqual(r, limiters, l)
limiters = l
})
// Check that ipNoPrefix1 have the expected limiter
l, ok := m.limiters.Load().Get(ipNoPrefix1)
require.True(t, ok)
require.NotNil(t, l)
limiter := l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
// Create a limiter for ipNoPrefix2
ipNoPrefix2 := Key([]byte(""), []byte("127.0.0.2"))
c2 := LimiterConfig{Rate: 2}
m.UpdateConfig(c2, ipNoPrefix2)
m.Allow(ipLimited{key: ipNoPrefix2})
retry.RunWith(&retry.Timer{Wait: 1 * time.Second, Timeout: 5 * time.Second}, t, func(r *retry.R) {
l := m.limiters.Load()
require.NotEqual(r, limiters, l)
limiters = l
})
// Check that ipNoPrefix1 have the expected limiter
l, ok = m.limiters.Load().Get(ipNoPrefix1)
require.True(t, ok)
require.NotNil(t, l)
limiter = l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
// Check that ipNoPrefix2 have the expected limiter
l, ok = m.limiters.Load().Get(ipNoPrefix2)
require.True(t, ok)
require.NotNil(t, l)
limiter = l.(*Limiter)
require.True(t, c2.isApplied(limiter.limiter))
})
}
func TestRateLimiterUpdateConfig(t *testing.T) {
// Create a MultiLimiter m with a defaultConfig c and check the defaultConfig is applied
t.Run("Allow a key and check defaultConfig is applied to that key", func(t *testing.T) {
//Create a multilimiter
c := Config{ReconcileCheckLimit: 100 * time.Millisecond, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
//Create a limiter for ipNoPrefix
ipNoPrefix := Key([]byte(""), []byte("127.0.0.1"))
c1 := LimiterConfig{Rate: 1}
m.UpdateConfig(c1, ipNoPrefix)
m.Allow(ipLimited{key: ipNoPrefix})
storeLimiter(m)
// Verify the expected limiter is applied
l, ok := m.limiters.Load().Get(ipNoPrefix)
require.True(t, ok)
require.NotNil(t, l)
limiter := l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
})
t.Run("Update nil prefix and make sure it's written in the root", func(t *testing.T) {
//Create a multilimiter
c := Config{ReconcileCheckLimit: 100 * time.Millisecond, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
//Create a limiter for nil
prefix := []byte(nil)
c1 := LimiterConfig{Rate: 1}
m.UpdateConfig(c1, prefix)
// Verify the expected limiter is applied
v, ok := m.limitersConfigs.Load().Get([]byte(""))
require.True(t, ok)
require.NotNil(t, v)
c2 := v.(*LimiterConfig)
require.Equal(t, c1, *c2)
})
t.Run("Allow 2 keys with prefix and check defaultConfig is applied to those keys", func(t *testing.T) {
//Create a multilimiter
c := Config{ReconcileCheckLimit: 100 * time.Millisecond, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
//Create a limiter for ip
prefix := []byte("namespace.write")
ip := Key(prefix, []byte("127.0.0.1"))
c1 := LimiterConfig{Rate: 1}
m.UpdateConfig(c1, ip)
m.Allow(ipLimited{key: ip})
storeLimiter(m)
//Create a limiter for ip2
ip2 := Key(prefix, []byte("127.0.0.2"))
c2 := LimiterConfig{Rate: 2}
m.UpdateConfig(c2, ip2)
m.Allow(ipLimited{key: ip2})
//Verify the config is applied for ip
l, ok := m.limiters.Load().Get(ip)
require.True(t, ok)
require.NotNil(t, l)
limiter := l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
})
t.Run("Apply a defaultConfig to 'namespace.write' check the defaultConfig is applied to existing keys under that prefix", func(t *testing.T) {
//Create a multilimiter
c := Config{ReconcileCheckLimit: 100 * time.Millisecond, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
//Create a limiter for ip
prefix := []byte("namespace.write")
ip := Key(prefix, []byte("127.0.0.1"))
c3 := LimiterConfig{Rate: 2}
m.UpdateConfig(c3, prefix)
// call reconcileLimitedOnce to make sure the update is applied
m.reconcileConfig(m.limiters.Load().Txn())
m.Allow(ipLimited{key: ip})
storeLimiter(m)
//Verify the config is applied for ip
l3, ok3 := m.limiters.Load().Get(ip)
require.True(t, ok3)
require.NotNil(t, l3)
limiter3 := l3.(*Limiter)
require.True(t, c3.isApplied(limiter3.limiter))
})
t.Run("Allow an IP with prefix and check prefix defaultConfig is applied to new keys under that prefix", func(t *testing.T) {
//Create a multilimiter
c := Config{ReconcileCheckLimit: 100 * time.Millisecond, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
//Create a limiter for ip
c1 := LimiterConfig{Rate: 3}
prefix := []byte("namespace.read")
m.UpdateConfig(c1, prefix)
ip := Key(prefix, []byte("127.0.0.1"))
m.Allow(ipLimited{key: ip})
storeLimiter(m)
//Verify the config is applied for ip
l, ok := m.limiters.Load().Get(ip)
require.True(t, ok)
require.NotNil(t, l)
limiter := l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
})
t.Run("Allow an IP with prefix and check prefix defaultConfig is applied to new keys under that prefix, delete config and check default applied", func(t *testing.T) {
//Create a multilimiter
c := Config{ReconcileCheckLimit: 100 * time.Second, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
//Create a limiter for ip
c1 := LimiterConfig{Rate: 3}
prefix := []byte("namespace.read")
m.UpdateConfig(c1, prefix)
ip := Key(prefix, []byte("127.0.0.1"))
m.Allow(ipLimited{key: ip})
storeLimiter(m)
//Verify the config is applied for ip
l, ok := m.limiters.Load().Get(ip)
require.True(t, ok)
require.NotNil(t, l)
limiter := l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
// Delete the prefix
m.DeleteConfig(prefix)
reconcile(m)
// Verify the limiter is removed
_, ok = m.limiters.Load().Get(ip)
require.False(t, ok)
})
t.Run("Allow an IP with prefix and check prefix config is applied to new keys under that prefix", func(t *testing.T) {
//Create a multilimiter
c := Config{ReconcileCheckLimit: 100 * time.Millisecond, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
//Create a limiter for ip
c1 := LimiterConfig{Rate: 3}
prefix := Key([]byte("ip.ratelimit"), []byte("127.0"))
m.UpdateConfig(c1, prefix)
ip := Key([]byte("ip.ratelimit"), []byte("127.0.0.1"))
m.Allow(ipLimited{key: ip})
storeLimiter(m)
//Verify the config is applied for ip
load := m.limiters.Load()
l, ok := load.Get(ip)
require.True(t, ok)
require.NotNil(t, l)
limiter := l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
})
t.Run("Allow an IP with 2 prefixes and check prefix config is applied to new keys under that prefix", func(t *testing.T) {
//Create a multilimiter
c := Config{ReconcileCheckLimit: 100 * time.Millisecond, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
//Create a limiter for "127.0" ip with config c1
c1 := LimiterConfig{Rate: 3}
prefix := Key([]byte("ip.ratelimit"), []byte("127.0"))
m.UpdateConfig(c1, prefix)
//Create a limiter for "127.0.0" ip with config c2
prefix = Key([]byte("ip.ratelimit"), []byte("127.0.0"))
c2 := LimiterConfig{Rate: 6}
m.UpdateConfig(c2, prefix)
ip := Key([]byte("ip.ratelimit"), []byte("127.0.0.1"))
m.Allow(ipLimited{key: ip})
storeLimiter(m)
// Verify that "127.0.0.1" have the right limiter config
load := m.limiters.Load()
l, ok := load.Get(ip)
require.True(t, ok)
require.NotNil(t, l)
limiter := l.(*Limiter)
require.True(t, c2.isApplied(limiter.limiter))
//Create a limiter for "127.0.1.1" ip with config c2
ip = Key([]byte("ip.ratelimit"), []byte("127.0.1.1"))
m.Allow(ipLimited{key: ip})
storeLimiter(m)
// Verify that "127.0.1.1" have the right limiter config
load = m.limiters.Load()
l, ok = load.Get(ip)
require.True(t, ok)
require.NotNil(t, l)
limiter = l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
})
t.Run("Allow an IP with prefix and check after it's cleaned new Allow would give it the right defaultConfig", func(t *testing.T) {
//Create a multilimiter
c := Config{ReconcileCheckLimit: 100 * time.Millisecond, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(t, *m.defaultConfig.Load(), c)
//Create a limiter for ip
prefix := []byte("namespace.read")
ip := Key(prefix, []byte("127.0.0.1"))
c1 := LimiterConfig{Rate: 1}
m.UpdateConfig(c1, prefix)
// call reconcileLimitedOnce to make sure the update is applied
reconcile(m)
m.Allow(ipLimited{key: ip})
storeLimiter(m)
//Verify the config is applied for ip
l, ok := m.limiters.Load().Get(ip)
require.True(t, ok)
require.NotNil(t, l)
limiter := l.(*Limiter)
require.True(t, c1.isApplied(limiter.limiter))
})
}
func FuzzSingleConfig(f *testing.F) {
c := Config{ReconcileCheckLimit: 100 * time.Millisecond, ReconcileCheckInterval: 10 * time.Millisecond}
m := NewMultiLimiter(c)
require.Equal(f, *m.defaultConfig.Load(), c)
f.Add(Key(randIP()))
f.Add(Key(randIP(), randIP()))
f.Add(Key(randIP(), randIP(), randIP()))
f.Add(Key(randIP(), randIP(), randIP(), randIP()))
c1 := LimiterConfig{
Rate: 100,
Burst: 123,
}
f.Fuzz(func(t *testing.T, ff []byte) {
m.UpdateConfig(c1, ff)
m.Allow(Limited{key: ff})
storeLimiter(m)
checkLimiter(t, ff, m.limiters.Load().Txn())
checkTree(t, m.limiters.Load().Txn())
})
}
func FuzzSplitKey(f *testing.F) {
f.Add(Key(randIP(), randIP()))
f.Add(Key(randIP(), randIP(), randIP()))
f.Add(Key(randIP(), randIP(), randIP(), randIP()))
f.Add([]byte(""))
f.Fuzz(func(t *testing.T, ff []byte) {
prefix, suffix := splitKey(ff)
require.NotNil(t, prefix)
require.NotNil(t, suffix)
if len(prefix) == 0 && len(suffix) == 0 {
return
}
joined := bytes.Join([][]byte{prefix, suffix}, []byte(separator))
require.Equal(t, ff, joined)
require.False(t, bytes.Contains(prefix, []byte(separator)))
})
}
func checkLimiter(t require.TestingT, ff []byte, Tree *radix.Txn) {
v, ok := Tree.Get(ff)
require.True(t, ok)
require.NotNil(t, v)
}
func FuzzUpdateConfig(f *testing.F) {
f.Add(bytes.Join([][]byte{[]byte(""), Key(randIP()), Key(randIP(), randIP()), Key(randIP(), randIP(), randIP()), Key(randIP(), randIP(), randIP(), randIP())}, []byte(",")))
f.Fuzz(func(t *testing.T, ff []byte) {
cm := Config{ReconcileCheckLimit: 1 * time.Millisecond, ReconcileCheckInterval: 1 * time.Millisecond}
m := NewMultiLimiter(cm)
ctx, cancel := context.WithCancel(context.Background())
m.Run(ctx)
defer cancel()
keys := bytes.Split(ff, []byte(","))
for _, f := range keys {
prefix, _ := splitKey(f)
c := LimiterConfig{Rate: rate.Limit(rand.Float64()), Burst: rand.Int()}
m.UpdateConfig(c, prefix)
go m.Allow(Limited{key: f})
}
m.reconcileConfig(m.limiters.Load().Txn())
checkTree(t, m.limiters.Load().Txn())
})
}
func checkTree(t require.TestingT, tree *radix.Txn) {
iterator := tree.Root().Iterator()
kp, v, ok := iterator.Next()
for ok {
switch c := v.(type) {
case *Limiter:
if c.limiter != nil {
prefix, _ := splitKey(kp)
v, _ := tree.Get(prefix)
switch c2 := v.(type) {
case *LimiterConfig:
if c2 != nil {
applied := c2.isApplied(c.limiter)
require.True(t, applied, fmt.Sprintf("%v,%v", kp, prefix))
}
}
}
default:
require.Nil(t, v)
}
kp, v, ok = iterator.Next()
}
}
type ipLimited struct {
key []byte
}
func (i ipLimited) Key() []byte {
return i.key
}
func BenchmarkTestRateLimiterFixedIP(b *testing.B) {
var Config = Config{ReconcileCheckLimit: time.Microsecond, ReconcileCheckInterval: time.Millisecond}
m := NewMultiLimiter(Config)
ctx, cancel := context.WithCancel(context.Background())
m.Run(ctx)
defer cancel()
ip := []byte{244, 233, 0, 1}
for j := 0; j < b.N; j++ {
m.Allow(ipLimited{key: ip})
}
}
func BenchmarkTestRateLimiterAllowPrefill(b *testing.B) {
cases := []struct {
name string
prefill uint64
}{
{name: "no prefill", prefill: 0},
{name: "prefill with 1K keys", prefill: 1000},
{name: "prefill with 10K keys", prefill: 10_000},
{name: "prefill with 100K keys", prefill: 100_000},
}
for _, tc := range cases {
b.Run(tc.name, func(b *testing.B) {
var Config = Config{ReconcileCheckLimit: time.Second, ReconcileCheckInterval: time.Second}
m := NewMultiLimiter(Config)
var i uint64
for i = 0xdeaddead; i < 0xdeaddead+tc.prefill; i++ {
buf := make([]byte, 8)
binary.LittleEndian.PutUint64(buf, i)
m.Allow(ipLimited{key: buf})
}
b.ResetTimer()
for j := 0; j < b.N; j++ {
buf := make([]byte, 4)
binary.LittleEndian.PutUint32(buf, uint32(j))
m.Allow(ipLimited{key: buf})
}
})
}
}
func BenchmarkTestRateLimiterAllowConcurrencyPrefill(b *testing.B) {
cases := []struct {
name string
prefill uint64
}{
{name: "no prefill", prefill: 0},
{name: "prefill with 1K keys", prefill: 1000},
{name: "prefill with 10K keys", prefill: 10_000},
{name: "prefill with 100K keys", prefill: 100_000},
}
for _, tc := range cases {
b.Run(tc.name, func(b *testing.B) {
var Config = Config{ReconcileCheckLimit: time.Second, ReconcileCheckInterval: 100 * time.Second}
m := NewMultiLimiter(Config)
ctx, cancel := context.WithCancel(context.Background())
m.Run(ctx)
defer cancel()
var i uint64
for i = 0xdeaddead; i < 0xdeaddead+tc.prefill; i++ {
buf := make([]byte, 8)
binary.LittleEndian.PutUint64(buf, i)
m.Allow(ipLimited{key: buf})
}
wg := sync.WaitGroup{}
b.ResetTimer()
for j := 0; j < b.N; j++ {
wg.Add(1)
go func(n int) {
buf := make([]byte, 4)
binary.LittleEndian.PutUint32(buf, uint32(n))
m.Allow(ipLimited{key: buf})
wg.Done()
}(j)
}
wg.Wait()
})
}
}
func BenchmarkTestRateLimiterRandomIP(b *testing.B) {
var Config = Config{ReconcileCheckLimit: time.Microsecond, ReconcileCheckInterval: time.Millisecond}
m := NewMultiLimiter(Config)
ctx, cancel := context.WithCancel(context.Background())
m.Run(ctx)
defer cancel()
for j := 0; j < b.N; j++ {
m.Allow(ipLimited{key: randIP()})
}
}
func randIP() []byte {
buf := make([]byte, 4)
ip := rand.Uint32()
binary.LittleEndian.PutUint32(buf, ip)
return buf
}
type mockTicker struct {
tickerCh chan time.Time
}
func (m *mockTicker) Ticker() <-chan time.Time {
return m.tickerCh
}
func splitKey(key []byte) ([]byte, []byte) {
ret := bytes.SplitN(key, []byte(separator), 2)
if len(ret) != 2 {
return []byte(""), []byte("")
}
return ret[0], ret[1]
}