From afe5cac2b4ef300ec778af64a0f1604e58e6bf6c Mon Sep 17 00:00:00 2001 From: Tarun Koyalwar Date: Tue, 31 Jan 2023 01:38:18 +0530 Subject: [PATCH 1/6] fix race conditions --- adaptive_ratelimit_test.go | 37 ++++++++----------- go.mod | 5 +-- go.sum | 2 -- keyratelimit.go | 73 ++++++++++++++++++++------------------ keyratelimit_test.go | 8 ++--- ratelimit.go | 40 +++++++++++---------- ratelimit_test.go | 1 + 7 files changed, 79 insertions(+), 87 deletions(-) diff --git a/adaptive_ratelimit_test.go b/adaptive_ratelimit_test.go index 0c8bb6d..b507e78 100644 --- a/adaptive_ratelimit_test.go +++ b/adaptive_ratelimit_test.go @@ -1,26 +1,17 @@ package ratelimit_test -import ( - "context" - "testing" - "time" +// func TestAdaptiveRateLimit(t *testing.T) { +// limiter := ratelimit.NewUnlimited(context.Background()) +// start := time.Now() - "github.com/projectdiscovery/ratelimit" - "github.com/stretchr/testify/require" -) - -func TestAdaptiveRateLimit(t *testing.T) { - limiter := ratelimit.NewUnlimited(context.Background()) - start := time.Now() - - for i := 0; i < 132; i++ { - limiter.Take() - // got 429 / hit ratelimit after 100 - if i == 100 { - // Retry-After and new limiter (calibrate using different statergies) - // new expected ratelimit 30req every 5 sec - limiter.SleepandReset(time.Duration(5)*time.Second, 30, time.Duration(5)*time.Second) - } - } - require.Equal(t, time.Since(start).Round(time.Second), time.Duration(10)*time.Second) -} +// for i := 0; i < 132; i++ { +// limiter.Take() +// // got 429 / hit ratelimit after 100 +// if i == 100 { +// // Retry-After and new limiter (calibrate using different statergies) +// // new expected ratelimit 30req every 5 sec +// limiter.SleepandReset(time.Duration(5)*time.Second, 30, time.Duration(5)*time.Second) +// } +// } +// require.Equal(t, time.Since(start).Round(time.Second), time.Duration(10)*time.Second) +// } diff --git a/go.mod b/go.mod index b53ef57..54ec400 100644 --- a/go.mod +++ b/go.mod @@ -2,10 +2,7 @@ module github.com/projectdiscovery/ratelimit go 1.18 -require ( - github.com/stretchr/testify v1.8.1 - golang.org/x/exp v0.0.0-20221217163422-3c43f8badb15 -) +require github.com/stretchr/testify v1.8.1 require ( github.com/davecgh/go-spew v1.1.1 // indirect diff --git a/go.sum b/go.sum index faea689..2ec90f7 100644 --- a/go.sum +++ b/go.sum @@ -10,8 +10,6 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= -golang.org/x/exp v0.0.0-20221217163422-3c43f8badb15 h1:5oN1Pz/eDhCpbMbLstvIPa0b/BEQo6g6nwV3pLjfM6w= -golang.org/x/exp v0.0.0-20221217163422-3c43f8badb15/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/keyratelimit.go b/keyratelimit.go index c8cafb3..5679540 100644 --- a/keyratelimit.go +++ b/keyratelimit.go @@ -3,9 +3,8 @@ package ratelimit import ( "context" "fmt" + "sync" "time" - - "golang.org/x/exp/maps" ) // Options of MultiLimiter @@ -34,7 +33,7 @@ func (o *Options) Validate() error { // MultiLimiter is wrapper around Limiter than can limit based on a key type MultiLimiter struct { - limiters map[string]*Limiter + limiters sync.Map // map of limiters ctx context.Context } @@ -43,68 +42,72 @@ func (m *MultiLimiter) Add(opts *Options) error { if err := opts.Validate(); err != nil { return err } - _, ok := m.limiters[opts.Key] - if ok { - return fmt.Errorf("key already exists") - } var rlimiter *Limiter if opts.IsUnlimited { rlimiter = NewUnlimited(m.ctx) } else { rlimiter = New(m.ctx, opts.MaxCount, opts.Duration) } - m.limiters[opts.Key] = rlimiter + // ok if true if key already exists + _, ok := m.limiters.LoadOrStore(opts.Key, rlimiter) + if ok { + return fmt.Errorf("key already exists") + } return nil } // GetLimit returns current ratelimit of given key func (m *MultiLimiter) GetLimit(key string) (uint, error) { - limiter, ok := m.limiters[key] - if !ok || limiter == nil { - return 0, fmt.Errorf("key doesnot exist") + limiter, err := m.get(key) + if err != nil { + return 0, err } return limiter.GetLimit(), nil } // Take one token from bucket returns error if key not present func (m *MultiLimiter) Take(key string) error { - limiter, ok := m.limiters[key] - if !ok || limiter == nil { - return fmt.Errorf("key doesnot exist") + limiter, err := m.get(key) + if err != nil { + return err } limiter.Take() return nil } -// Stop internal limiters with defined keys or all if no key is provided -func (m *MultiLimiter) Stop(keys ...string) { - if len(keys) > 0 { - m.stopWithKeys(keys...) - } else { - m.stopWithKeys(maps.Keys(m.limiters)...) +// AddAndTake adds key if not present and then takes token from bucket +func (m *MultiLimiter) AddAndTake(opts *Options) { + if limiter, err := m.get(opts.Key); err == nil { + limiter.Take() + return } + m.Add(opts) + m.Take(opts.Key) } -// stopWithKeys stops the internal limiters matching keys -func (m *MultiLimiter) stopWithKeys(keys ...string) { - for _, key := range keys { - if limiter, ok := m.limiters[key]; ok { +// Stop internal limiters with defined keys or all if no key is provided +func (m *MultiLimiter) Stop(keys ...string) { + if len(keys) == 0 { + m.limiters.Range(func(key, value any) bool { + value.(*Limiter).Stop() + return true + }) + return + } + for _, v := range keys { + if limiter, err := m.get(v); err == nil { limiter.Stop() } } } -// SleepandReset stops timer removes all tokens and resets with new limit (used for Adaptive Ratelimiting) -func (m *MultiLimiter) SleepandReset(SleepTime time.Duration, opts *Options) error { - if err := opts.Validate(); err != nil { - return err +// get returns *Limiter instance +func (m *MultiLimiter) get(key string) (*Limiter, error) { + val, _ := m.limiters.Load(key) + if val == nil { + return nil, fmt.Errorf("multilimiter: key does not exist") } - limiter, ok := m.limiters[opts.Key] - if !ok || limiter == nil { - return fmt.Errorf("key doesnot exist") - } - limiter.SleepandReset(SleepTime, opts.MaxCount, opts.Duration) - return nil + return val.(*Limiter), nil } // NewMultiLimiter : Limits @@ -114,7 +117,7 @@ func NewMultiLimiter(ctx context.Context, opts *Options) (*MultiLimiter, error) } multilimiter := &MultiLimiter{ ctx: ctx, - limiters: map[string]*Limiter{}, + limiters: sync.Map{}, } return multilimiter, multilimiter.Add(opts) } diff --git a/keyratelimit_test.go b/keyratelimit_test.go index 47d4a44..ed0e678 100644 --- a/keyratelimit_test.go +++ b/keyratelimit_test.go @@ -24,11 +24,11 @@ func TestMultiLimiter(t *testing.T) { go func() { defer wg.Done() defaultStart := time.Now() - for i := 0; i < 201; i++ { + for i := 0; i < 202; i++ { errx := limiter.Take("default") require.Nil(t, errx, "failed to take") } - require.Greater(t, time.Since(defaultStart), time.Duration(6)*time.Second) + require.Greater(t, time.Since(defaultStart).Nanoseconds(), (time.Duration(6) * time.Second).Nanoseconds()) }() err = limiter.Add(&ratelimit.Options{ @@ -43,11 +43,11 @@ func TestMultiLimiter(t *testing.T) { go func() { defer wg.Done() oneStart := time.Now() - for i := 0; i < 201; i++ { + for i := 0; i < 202; i++ { errx := limiter.Take("one") require.Nil(t, errx) } - require.Greater(t, time.Since(oneStart), time.Duration(6)*time.Second) + require.Greater(t, time.Since(oneStart).Nanoseconds(), (time.Duration(6) * time.Second).Nanoseconds()) }() wg.Wait() } diff --git a/ratelimit.go b/ratelimit.go index 4bd48a5..bc8b8a4 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -18,6 +18,7 @@ type Limiter struct { } func (limiter *Limiter) run(ctx context.Context) { + defer close(limiter.tokens) for { if limiter.count == 0 { <-limiter.ticker.C @@ -49,28 +50,29 @@ func (ratelimiter *Limiter) GetLimit() uint { return ratelimiter.maxCount } -// SleepandReset stops timer removes all tokens and resets with new limit (used for Adaptive Ratelimiting) -func (ratelimiter *Limiter) SleepandReset(sleepTime time.Duration, newLimit uint, duration time.Duration) { - // stop existing Limiter using internalContext - ratelimiter.cancelFunc() - // drain any token - close(ratelimiter.tokens) - <-ratelimiter.tokens - // sleep - time.Sleep(sleepTime) - //reset and start - ratelimiter.maxCount = newLimit - ratelimiter.count = newLimit - ratelimiter.ticker = time.NewTicker(duration) - ratelimiter.tokens = make(chan struct{}) - ctx, cancel := context.WithCancel(context.TODO()) - ratelimiter.cancelFunc = cancel - go ratelimiter.run(ctx) -} +// TODO: SleepandReset should be able to handle multiple calls without resetting multiple times +// Which is not possible in this implementation +// // SleepandReset stops timer removes all tokens and resets with new limit (used for Adaptive Ratelimiting) +// func (ratelimiter *Limiter) SleepandReset(sleepTime time.Duration, newLimit uint, duration time.Duration) { +// // stop existing Limiter using internalContext +// ratelimiter.cancelFunc() +// // drain any token +// close(ratelimiter.tokens) +// <-ratelimiter.tokens +// // sleep +// time.Sleep(sleepTime) +// //reset and start +// ratelimiter.maxCount = newLimit +// ratelimiter.count = newLimit +// ratelimiter.ticker = time.NewTicker(duration) +// ratelimiter.tokens = make(chan struct{}) +// ctx, cancel := context.WithCancel(context.TODO()) +// ratelimiter.cancelFunc = cancel +// go ratelimiter.run(ctx) +// } // Stop the rate limiter canceling the internal context func (ratelimiter *Limiter) Stop() { - defer close(ratelimiter.tokens) if ratelimiter.cancelFunc != nil { ratelimiter.cancelFunc() } diff --git a/ratelimit_test.go b/ratelimit_test.go index 5c1e916..3c74aa6 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -67,6 +67,7 @@ func TestRateLimit(t *testing.T) { } wg.Wait() + limiter.Stop() took := time.Since(start) require.Equal(t, expected, int(count)) require.True(t, took >= time.Duration(10*time.Second)) From 7630feb9df39821511deacf9c917123d4da45821 Mon Sep 17 00:00:00 2001 From: Tarun Koyalwar Date: Tue, 31 Jan 2023 16:31:55 +0530 Subject: [PATCH 2/6] add -race to gh workflow and improvements --- .github/workflows/build-test.yml | 2 +- keyratelimit_test.go | 12 ++++++++---- ratelimit_test.go | 12 ++++++++++++ 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/.github/workflows/build-test.yml b/.github/workflows/build-test.yml index c86bd40..db61658 100644 --- a/.github/workflows/build-test.yml +++ b/.github/workflows/build-test.yml @@ -21,7 +21,7 @@ jobs: uses: actions/checkout@v3 - name: Test - run: go test ./... + run: go test -race ./... - name: Build Example run: go build example/main.go \ No newline at end of file diff --git a/keyratelimit_test.go b/keyratelimit_test.go index ed0e678..857bf04 100644 --- a/keyratelimit_test.go +++ b/keyratelimit_test.go @@ -24,11 +24,13 @@ func TestMultiLimiter(t *testing.T) { go func() { defer wg.Done() defaultStart := time.Now() - for i := 0; i < 202; i++ { + for i := 0; i < 201; i++ { errx := limiter.Take("default") require.Nil(t, errx, "failed to take") } - require.Greater(t, time.Since(defaultStart).Nanoseconds(), (time.Duration(6) * time.Second).Nanoseconds()) + timeTaken := time.Since(defaultStart) + expectedTime := time.Duration(6) * time.Second + require.GreaterOrEqualf(t, timeTaken.Nanoseconds(), expectedTime.Nanoseconds(), "more token sent than requested in given timeframe") }() err = limiter.Add(&ratelimit.Options{ @@ -43,11 +45,13 @@ func TestMultiLimiter(t *testing.T) { go func() { defer wg.Done() oneStart := time.Now() - for i := 0; i < 202; i++ { + for i := 0; i < 201; i++ { errx := limiter.Take("one") require.Nil(t, errx) } - require.Greater(t, time.Since(oneStart).Nanoseconds(), (time.Duration(6) * time.Second).Nanoseconds()) + timeTaken := time.Since(oneStart) + expectedTime := time.Duration(6) * time.Second + require.GreaterOrEqualf(t, timeTaken.Nanoseconds(), expectedTime.Nanoseconds(), "more token sent than requested in given timeframe") }() wg.Wait() } diff --git a/ratelimit_test.go b/ratelimit_test.go index 3c74aa6..98e3e3a 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -73,4 +73,16 @@ func TestRateLimit(t *testing.T) { require.True(t, took >= time.Duration(10*time.Second)) require.True(t, took <= time.Duration(12*time.Second)) }) + + t.Run("Time comparsion", func(t *testing.T) { + limiter := New(context.TODO(), 100, time.Duration(3)*time.Second) + // if ratelimit works correctly it should take at least 6 sec to take/consume 201 tokens + startTime := time.Now() + for i := 0; i < 201; i++ { + limiter.Take() + } + timetaken := time.Since(startTime) + expected := time.Duration(6) * time.Second + require.GreaterOrEqualf(t, timetaken.Nanoseconds(), expected.Nanoseconds(), "more tokens sent than expected with ratelimit") + }) } From 3b5e930c1b2af61d661483e54b997700438afab8 Mon Sep 17 00:00:00 2001 From: Tarun Koyalwar Date: Tue, 31 Jan 2023 16:54:23 +0530 Subject: [PATCH 3/6] fix lint error --- keyratelimit.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/keyratelimit.go b/keyratelimit.go index 5679540..cc47c2e 100644 --- a/keyratelimit.go +++ b/keyratelimit.go @@ -81,8 +81,8 @@ func (m *MultiLimiter) AddAndTake(opts *Options) { limiter.Take() return } - m.Add(opts) - m.Take(opts.Key) + _ = m.Add(opts) + _ = m.Take(opts.Key) } // Stop internal limiters with defined keys or all if no key is provided From 909d2e7edde3e556a42c5f6eb444870738cc0fbe Mon Sep 17 00:00:00 2001 From: Tarun Koyalwar Date: Tue, 31 Jan 2023 16:59:29 +0530 Subject: [PATCH 4/6] roundoff time to millisec --- keyratelimit_test.go | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/keyratelimit_test.go b/keyratelimit_test.go index 857bf04..38c8257 100644 --- a/keyratelimit_test.go +++ b/keyratelimit_test.go @@ -19,6 +19,7 @@ func TestMultiLimiter(t *testing.T) { }) require.Nil(t, err) wg := &sync.WaitGroup{} + expectedTime := (time.Duration(6) * time.Second).Round(time.Millisecond) wg.Add(1) go func() { @@ -28,8 +29,7 @@ func TestMultiLimiter(t *testing.T) { errx := limiter.Take("default") require.Nil(t, errx, "failed to take") } - timeTaken := time.Since(defaultStart) - expectedTime := time.Duration(6) * time.Second + timeTaken := time.Since(defaultStart).Round(time.Millisecond) require.GreaterOrEqualf(t, timeTaken.Nanoseconds(), expectedTime.Nanoseconds(), "more token sent than requested in given timeframe") }() @@ -49,8 +49,7 @@ func TestMultiLimiter(t *testing.T) { errx := limiter.Take("one") require.Nil(t, errx) } - timeTaken := time.Since(oneStart) - expectedTime := time.Duration(6) * time.Second + timeTaken := time.Since(oneStart).Round(time.Millisecond) require.GreaterOrEqualf(t, timeTaken.Nanoseconds(), expectedTime.Nanoseconds(), "more token sent than requested in given timeframe") }() wg.Wait() From a8ff0ab3e0d8b5af85aca61f28c509e5fd4e3b53 Mon Sep 17 00:00:00 2001 From: Tarun Koyalwar Date: Tue, 31 Jan 2023 17:14:14 +0530 Subject: [PATCH 5/6] fix time comparison --- ratelimit_test.go | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/ratelimit_test.go b/ratelimit_test.go index 98e3e3a..6c58074 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -21,13 +21,13 @@ func TestRateLimit(t *testing.T) { limiter.Take() count++ } - took := time.Since(start) + took := time.Since(start).Nanoseconds() require.Equal(t, count, 10) - require.True(t, took < expected) + require.Less(t, took, expected.Nanoseconds()) // take another one above max limiter.Take() - took = time.Since(start) - require.True(t, took >= expected) + took = time.Since(start).Nanoseconds() + require.GreaterOrEqual(t, took, expected.Nanoseconds()) }) t.Run("Unlimited Rate Limit", func(t *testing.T) { @@ -41,7 +41,7 @@ func TestRateLimit(t *testing.T) { } took := time.Since(start) require.Equal(t, count, 1000) - require.True(t, took < time.Duration(1*time.Second)) + require.Lessf(t, took.Nanoseconds(), time.Duration(1*time.Second).Nanoseconds(), "burst rate of unlimited ratelimit is too low") }) t.Run("Concurrent Rate Limit Use", func(t *testing.T) { @@ -70,8 +70,8 @@ func TestRateLimit(t *testing.T) { limiter.Stop() took := time.Since(start) require.Equal(t, expected, int(count)) - require.True(t, took >= time.Duration(10*time.Second)) - require.True(t, took <= time.Duration(12*time.Second)) + require.GreaterOrEqualf(t, took, time.Duration(10*time.Second).Nanoseconds(), "ratelimit timeframe mismatch should be > 10s") + require.LessOrEqualf(t, took, time.Duration(12*time.Second).Nanoseconds(), "ratelimit timeframe mismatch should be < 10s") }) t.Run("Time comparsion", func(t *testing.T) { From 69db5c58d624f5b46b5daba896edbf94bc748c2e Mon Sep 17 00:00:00 2001 From: Tarun Koyalwar Date: Tue, 31 Jan 2023 19:26:03 +0530 Subject: [PATCH 6/6] improved error handling --- go.mod | 1 + go.sum | 2 ++ keyratelimit.go | 27 +++++++++++++++++++-------- 3 files changed, 22 insertions(+), 8 deletions(-) diff --git a/go.mod b/go.mod index 54ec400..50d0e8f 100644 --- a/go.mod +++ b/go.mod @@ -7,5 +7,6 @@ require github.com/stretchr/testify v1.8.1 require ( github.com/davecgh/go-spew v1.1.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/projectdiscovery/utils v0.0.6 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 2ec90f7..b472b9c 100644 --- a/go.sum +++ b/go.sum @@ -3,6 +3,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/projectdiscovery/utils v0.0.6 h1:6SDn/5E5NxrAfcYrZ7omXPmiU9n8p0rKXZ4BAOQyzbw= +github.com/projectdiscovery/utils v0.0.6/go.mod h1:PCwA5YuCYWPgHaGiZmr53/SA9iGQmAnw7DSHuhr8VPQ= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= diff --git a/keyratelimit.go b/keyratelimit.go index cc47c2e..719b62d 100644 --- a/keyratelimit.go +++ b/keyratelimit.go @@ -2,9 +2,15 @@ package ratelimit import ( "context" - "fmt" "sync" "time" + + errorutil "github.com/projectdiscovery/utils/errors" +) + +var ( + ErrKeyAlreadyExists = errorutil.NewWithTag("MultiLimiter", "key already exists") + ErrKeyMissing = errorutil.NewWithTag("MultiLimiter", "key does not exist") ) // Options of MultiLimiter @@ -19,13 +25,13 @@ type Options struct { func (o *Options) Validate() error { if !o.IsUnlimited { if o.Key == "" { - return fmt.Errorf("empty keys not allowed") + return errorutil.NewWithTag("MultiLimiter", "empty keys not allowed") } if o.MaxCount == 0 { - return fmt.Errorf("maxcount cannot be zero") + return errorutil.NewWithTag("MultiLimiter", "maxcount cannot be zero") } if o.Duration == 0 { - return fmt.Errorf("time duration not set") + return errorutil.NewWithTag("MultiLimiter", "time duration not set") } } return nil @@ -51,7 +57,7 @@ func (m *MultiLimiter) Add(opts *Options) error { // ok if true if key already exists _, ok := m.limiters.LoadOrStore(opts.Key, rlimiter) if ok { - return fmt.Errorf("key already exists") + return ErrKeyAlreadyExists.Msgf("key: %v", opts.Key) } return nil } @@ -89,7 +95,9 @@ func (m *MultiLimiter) AddAndTake(opts *Options) { func (m *MultiLimiter) Stop(keys ...string) { if len(keys) == 0 { m.limiters.Range(func(key, value any) bool { - value.(*Limiter).Stop() + if limiter, ok := value.(*Limiter); ok { + limiter.Stop() + } return true }) return @@ -105,9 +113,12 @@ func (m *MultiLimiter) Stop(keys ...string) { func (m *MultiLimiter) get(key string) (*Limiter, error) { val, _ := m.limiters.Load(key) if val == nil { - return nil, fmt.Errorf("multilimiter: key does not exist") + return nil, ErrKeyMissing.Msgf("key: %v", key) + } + if limiter, ok := val.(*Limiter); ok { + return limiter, nil } - return val.(*Limiter), nil + return nil, errorutil.NewWithTag("MultiLimiter", "type assertion of rateLimiter failed in multiLimiter") } // NewMultiLimiter : Limits