diff --git a/v1/client.go b/v1/client.go index 273dc44..f8b9c25 100644 --- a/v1/client.go +++ b/v1/client.go @@ -36,7 +36,7 @@ func (c *MgClient) WithLogger(logger BasicLogger) *MgClient { } // WithLimiter sets the provided limiter instance into the Client. -func (c *MgClient) WithLimiter(limiter *TokensBucket) *MgClient { +func (c *MgClient) WithLimiter(limiter Limiter) *MgClient { c.limiter = limiter return c } diff --git a/v1/rate_limit.go b/v1/rate_limit.go index 48e99f0..455443c 100644 --- a/v1/rate_limit.go +++ b/v1/rate_limit.go @@ -7,11 +7,21 @@ import ( "time" ) +// NoopLimiter implements Limiter but doesn't limit anything. +var NoopLimiter Limiter = &noopLimiter{} + type token struct { rps atomic.Uint32 lastUse atomic.Value } +// Limiter implements some form of rate limiting. +type Limiter interface { + // Obtain the right to send a request. Should lock the execution if current goroutine needs to wait. + Obtain(string) +} + +// TokensBucket implements basic Limiter with fixed window and fixed amount of tokens per window. type TokensBucket struct { maxRPS uint32 tokens sync.Map @@ -21,7 +31,8 @@ type TokensBucket struct { sleep sleeper } -func NewTokensBucket(maxRPS uint32, unusedTokenTime, checkTokenTime time.Duration) *TokensBucket { +// NewTokensBucket constructs TokensBucket with provided parameters. +func NewTokensBucket(maxRPS uint32, unusedTokenTime, checkTokenTime time.Duration) Limiter { bucket := &TokensBucket{ maxRPS: maxRPS, unusedTokenTime: unusedTokenTime, @@ -79,6 +90,11 @@ func (m *TokensBucket) deleteUnusedToken() { } } +type noopLimiter struct{} + +func (l *noopLimiter) Obtain(string) {} + +// sleeper sleeps. This thing is necessary for tests. type sleeper interface { Sleep(time.Duration) } diff --git a/v1/types.go b/v1/types.go index f4965a3..7535531 100644 --- a/v1/types.go +++ b/v1/types.go @@ -78,12 +78,12 @@ const ( // MgClient type. type MgClient struct { - URL string `json:"url"` - Token string `json:"token"` - Debug bool `json:"debug"` - httpClient *http.Client `json:"-"` - logger BasicLogger `json:"-"` - limiter *TokensBucket `json:"-"` + URL string `json:"url"` + Token string `json:"token"` + Debug bool `json:"debug"` + httpClient *http.Client `json:"-"` + logger BasicLogger `json:"-"` + limiter Limiter `json:"-"` } // Channel type.