Update ratelimit. Simplify solution

This commit is contained in:
Fedor Korotkiy 2021-03-05 14:06:01 +03:00
parent de5fdd1753
commit 8ead5eb586
3 changed files with 36 additions and 2 deletions

View file

@ -7,10 +7,15 @@
func NewLimiter(maxCount int, interval time.Duration) *Limiter func NewLimiter(maxCount int, interval time.Duration) *Limiter
func (l *Limiter) Acquire(ctx context.Context) error func (l *Limiter) Acquire(ctx context.Context) error
func (l *Limiter) Stop()
``` ```
`Limiter` должен гарантировать, что на любом интервале времени `interval`, не больше `maxCount` вызовов `Limiter` должен гарантировать, что на любом интервале времени `interval`, не больше `maxCount` вызовов
`Acquire` могут завершиться без ошибки. `Acquire` могут завершиться без ошибки.
Каждый вызов `Acquire` должен либо завершаться успешно, либо завершаться с ошибкой в случае если `ctx` отменили Каждый вызов `Acquire` должен либо завершаться успешно, либо завершаться с ошибкой в случае если `ctx` отменили
во время ожидания. во время ожидания. Об отмене `ctx` нужно нужно узнавать по закрытию канала `ctx.Done()`. Если `ctx` отменён,
нужно возвращать ошибку `ctx.Err()`.
Вызовы `Acquire()` после `Stop()`, должны сразу завершаться с ошибкой ErrStopped.

View file

@ -4,6 +4,7 @@ package ratelimit
import ( import (
"context" "context"
"errors"
"time" "time"
) )
@ -11,6 +12,8 @@ import (
type Limiter struct { type Limiter struct {
} }
var ErrStopped = errors.New("limiter stopped")
// NewLimiter returns limiter that throttles rate of successful Acquire() calls // NewLimiter returns limiter that throttles rate of successful Acquire() calls
// to maxSize events at any given interval. // to maxSize events at any given interval.
func NewLimiter(maxCount int, interval time.Duration) *Limiter { func NewLimiter(maxCount int, interval time.Duration) *Limiter {
@ -20,3 +23,7 @@ func NewLimiter(maxCount int, interval time.Duration) *Limiter {
func (l *Limiter) Acquire(ctx context.Context) error { func (l *Limiter) Acquire(ctx context.Context) error {
panic("not implemented") panic("not implemented")
} }
func (l *Limiter) Stop() {
panic("not implemented")
}

View file

@ -9,11 +9,15 @@ import (
"time" "time"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
"go.uber.org/goleak"
"golang.org/x/sync/errgroup" "golang.org/x/sync/errgroup"
) )
func TestNoRateLimit(t *testing.T) { func TestNoRateLimit(t *testing.T) {
defer goleak.VerifyNone(t)
limit := NewLimiter(1, 0) limit := NewLimiter(1, 0)
defer limit.Stop()
ctx := context.Background() ctx := context.Background()
@ -22,7 +26,12 @@ func TestNoRateLimit(t *testing.T) {
} }
func TestBlockedRateLimit(t *testing.T) { func TestBlockedRateLimit(t *testing.T) {
limit := NewLimiter(0, time.Minute) defer goleak.VerifyNone(t)
limit := NewLimiter(1, time.Minute)
defer limit.Stop()
require.NoError(t, limit.Acquire(context.Background()))
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond) ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel() defer cancel()
@ -32,7 +41,10 @@ func TestBlockedRateLimit(t *testing.T) {
} }
func TestSimpleLimitCancel(t *testing.T) { func TestSimpleLimitCancel(t *testing.T) {
defer goleak.VerifyNone(t)
limit := NewLimiter(1, time.Minute) limit := NewLimiter(1, time.Minute)
defer limit.Stop()
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond) ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel() defer cancel()
@ -44,7 +56,10 @@ func TestSimpleLimitCancel(t *testing.T) {
} }
func TestTimeDistribution(t *testing.T) { func TestTimeDistribution(t *testing.T) {
defer goleak.VerifyNone(t)
limit := NewLimiter(100, time.Second) limit := NewLimiter(100, time.Second)
defer limit.Stop()
var lock sync.Mutex var lock sync.Mutex
okTimes := []time.Duration{} okTimes := []time.Duration{}
@ -94,12 +109,15 @@ func TestTimeDistribution(t *testing.T) {
} }
func TestStressBlocking(t *testing.T) { func TestStressBlocking(t *testing.T) {
defer goleak.VerifyNone(t)
const ( const (
N = 100 N = 100
G = 100 G = 100
) )
limit := NewLimiter(N, time.Millisecond*10) limit := NewLimiter(N, time.Millisecond*10)
defer limit.Stop()
var eg errgroup.Group var eg errgroup.Group
for i := 0; i < G; i++ { for i := 0; i < G; i++ {
@ -118,12 +136,15 @@ func TestStressBlocking(t *testing.T) {
} }
func TestStressNoBlocking(t *testing.T) { func TestStressNoBlocking(t *testing.T) {
defer goleak.VerifyNone(t)
const ( const (
N = 100 N = 100
G = 100 G = 100
) )
limit := NewLimiter(N, time.Millisecond*10) limit := NewLimiter(N, time.Millisecond*10)
defer limit.Stop()
var eg errgroup.Group var eg errgroup.Group
for i := 0; i < G; i++ { for i := 0; i < G; i++ {
@ -148,6 +169,7 @@ func BenchmarkNoBlocking(b *testing.B) {
b.SetBytes(1) b.SetBytes(1)
limit := NewLimiter(1, 0) limit := NewLimiter(1, 0)
defer limit.Stop()
ctx := context.Background() ctx := context.Background()