From fc4cf39f2a2484dc97e9e48e058b1a94da46f06f Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Sun, 7 May 2023 23:58:44 +0200 Subject: [PATCH 1/2] Using atomic for internal counters --- .gitignore | 2 ++ example/main.go | 2 +- keyratelimit.go | 4 ++-- ratelimit.go | 38 +++++++++++++++++++++----------------- 4 files changed, 26 insertions(+), 20 deletions(-) diff --git a/.gitignore b/.gitignore index 66fd13c..83540aa 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,5 @@ # Dependency directories (remove the comment below to include it) # vendor/ +.idea +.vscode \ No newline at end of file diff --git a/example/main.go b/example/main.go index 92a3520..e7d54fe 100644 --- a/example/main.go +++ b/example/main.go @@ -10,7 +10,7 @@ import ( func main() { // create a rate limiter by passing context, max tasks/tokens , time interval - limiter := ratelimit.New(context.Background(), 5, time.Duration(10*time.Second)) + limiter := ratelimit.New(context.Background(), 5, 10*time.Second) save := time.Now() diff --git a/keyratelimit.go b/keyratelimit.go index 719b62d..5bb269e 100644 --- a/keyratelimit.go +++ b/keyratelimit.go @@ -43,7 +43,7 @@ type MultiLimiter struct { ctx context.Context } -// Adds new bucket with key +// Add new bucket with key func (m *MultiLimiter) Add(opts *Options) error { if err := opts.Validate(); err != nil { return err @@ -54,7 +54,7 @@ func (m *MultiLimiter) Add(opts *Options) error { } else { rlimiter = New(m.ctx, opts.MaxCount, opts.Duration) } - // ok if true if key already exists + // ok is true if key already exists _, ok := m.limiters.LoadOrStore(opts.Key, rlimiter) if ok { return ErrKeyAlreadyExists.Msgf("key: %v", opts.Key) diff --git a/ratelimit.go b/ratelimit.go index bc8b8a4..a480f65 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -3,13 +3,17 @@ package ratelimit import ( "context" "math" + "sync/atomic" "time" ) +// equals to -1 +var minusOne = ^uint32(0) + // Limiter allows a burst of request during the defined duration type Limiter struct { - maxCount uint - count uint + maxCount uint32 + count atomic.Uint32 ticker *time.Ticker tokens chan struct{} ctx context.Context @@ -20,9 +24,9 @@ type Limiter struct { func (limiter *Limiter) run(ctx context.Context) { defer close(limiter.tokens) for { - if limiter.count == 0 { + if limiter.count.Load() == 0 { <-limiter.ticker.C - limiter.count = limiter.maxCount + limiter.count.Store(limiter.maxCount) } select { case <-ctx.Done(): @@ -33,21 +37,21 @@ func (limiter *Limiter) run(ctx context.Context) { limiter.ticker.Stop() return case limiter.tokens <- struct{}{}: - limiter.count-- + limiter.count.Add(minusOne) case <-limiter.ticker.C: - limiter.count = limiter.maxCount + limiter.count.Store(limiter.maxCount) } } } // Take one token from the bucket -func (rateLimiter *Limiter) Take() { - <-rateLimiter.tokens +func (limiter *Limiter) Take() { + <-limiter.tokens } // GetLimit returns current rate limit per given duration -func (ratelimiter *Limiter) GetLimit() uint { - return ratelimiter.maxCount +func (limiter *Limiter) GetLimit() uint { + return uint(limiter.maxCount) } // TODO: SleepandReset should be able to handle multiple calls without resetting multiple times @@ -72,9 +76,9 @@ func (ratelimiter *Limiter) GetLimit() uint { // } // Stop the rate limiter canceling the internal context -func (ratelimiter *Limiter) Stop() { - if ratelimiter.cancelFunc != nil { - ratelimiter.cancelFunc() +func (limiter *Limiter) Stop() { + if limiter.cancelFunc != nil { + limiter.cancelFunc() } } @@ -83,13 +87,13 @@ func New(ctx context.Context, max uint, duration time.Duration) *Limiter { internalctx, cancel := context.WithCancel(context.TODO()) limiter := &Limiter{ - maxCount: uint(max), - count: uint(max), + maxCount: uint32(max), ticker: time.NewTicker(duration), tokens: make(chan struct{}), ctx: ctx, cancelFunc: cancel, } + limiter.count.Store(uint32(max)) go limiter.run(internalctx) return limiter @@ -100,13 +104,13 @@ func NewUnlimited(ctx context.Context) *Limiter { internalctx, cancel := context.WithCancel(context.TODO()) limiter := &Limiter{ - maxCount: math.MaxUint, - count: math.MaxUint, + maxCount: math.MaxUint32, ticker: time.NewTicker(time.Millisecond), tokens: make(chan struct{}), ctx: ctx, cancelFunc: cancel, } + limiter.count.Store(math.MaxUint32) go limiter.run(internalctx) return limiter From 9871967beade04e133ee9daf5fa04867c2963e37 Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Mon, 8 May 2023 00:00:40 +0200 Subject: [PATCH 2/2] go version bump --- .github/workflows/build-test.yml | 2 +- .github/workflows/lint-test.yml | 2 +- go.mod | 9 ++++++--- go.sum | 10 +++++++++- 4 files changed, 17 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build-test.yml b/.github/workflows/build-test.yml index db61658..6d308a7 100644 --- a/.github/workflows/build-test.yml +++ b/.github/workflows/build-test.yml @@ -15,7 +15,7 @@ jobs: - name: Set up Go uses: actions/setup-go@v3 with: - go-version: 1.18 + go-version: 1.19 - name: Check out code uses: actions/checkout@v3 diff --git a/.github/workflows/lint-test.yml b/.github/workflows/lint-test.yml index a250797..4fb5ec8 100644 --- a/.github/workflows/lint-test.yml +++ b/.github/workflows/lint-test.yml @@ -13,7 +13,7 @@ jobs: - name: Set up Go uses: actions/setup-go@v3 with: - go-version: 1.18 + go-version: 1.19 - name: Run golangci-lint uses: golangci/golangci-lint-action@v3.4.0 with: diff --git a/go.mod b/go.mod index 50d0e8f..01c8513 100644 --- a/go.mod +++ b/go.mod @@ -1,12 +1,15 @@ module github.com/projectdiscovery/ratelimit -go 1.18 +go 1.19 -require github.com/stretchr/testify v1.8.1 +require ( + github.com/projectdiscovery/utils v0.0.6 + github.com/stretchr/testify v1.8.1 +) require ( github.com/davecgh/go-spew v1.1.1 // indirect + github.com/kr/pretty v0.3.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 b472b9c..a583648 100644 --- a/go.sum +++ b/go.sum @@ -1,10 +1,18 @@ +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 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/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= 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/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= +github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= 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= @@ -12,8 +20,8 @@ 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= -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/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=