Merge pull request #20 from projectdiscovery/bugfix-go-routine-leak

Add Stop method

go.mod

@@ -2,7 +2,10 @@ module github.com/projectdiscovery/ratelimit
 
 go 1.18
 
-require github.com/stretchr/testify v1.8.1
+require (
+	github.com/stretchr/testify v1.8.1
+	golang.org/x/exp v0.0.0-20221217163422-3c43f8badb15
+)
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect

go.sum

@@ -10,6 +10,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=
+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=

keyratelimit.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 	"time"
+
+	"golang.org/x/exp/maps"
 )
 
 // Options of MultiLimiter
@@ -74,6 +76,24 @@ func (m *MultiLimiter) Take(key string) error {
 	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)...)
+	}
+}
+
+// stopWithKeys stops the internal limiters matching keys
+func (m *MultiLimiter) stopWithKeys(keys ...string) {
+	for _, key := range keys {
+		if limiter, ok := m.limiters[key]; ok {
+			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 {

ratelimit.go

@@ -68,6 +68,13 @@ func (ratelimiter *Limiter) SleepandReset(sleepTime time.Duration, newLimit uint
 	go ratelimiter.run(ctx)
 }
 
+// Stop the rate limiter canceling the internal context
+func (ratelimiter *Limiter) Stop() {
+	if ratelimiter.cancelFunc != nil {
+		ratelimiter.cancelFunc()
+	}
+}
+
 // New creates a new limiter instance with the tokens amount and the interval
 func New(ctx context.Context, max uint, duration time.Duration) *Limiter {
 	internalctx, cancel := context.WithCancel(context.TODO())