Add ratelimit code

This commit is contained in:
Fedor Korotkiy 2021-03-02 19:46:51 +03:00
parent 6f093dc55a
commit a011860727
4 changed files with 210 additions and 0 deletions

36
ratelimit/rate.go Normal file
View file

@ -0,0 +1,36 @@
package ratelimit
import (
"fmt"
"strconv"
"strings"
"time"
)
// ParseRate parses rate in form of "N/D", e.g "10/1s" or "100/1ms"
func ParseRate(rate string) (count int, interval time.Duration, err error) {
parts := strings.SplitN(rate, "/", 2)
if len(parts) != 2 {
err = fmt.Errorf("invalid rate format in %q: missing slash", rate)
return
}
count, err = strconv.Atoi(parts[0])
if err != nil {
err = fmt.Errorf("invalid rate format in %q: %v", rate, err)
return
}
interval, err = time.ParseDuration(parts[1])
if err != nil {
err = fmt.Errorf("invalid rate format in %q: %v", rate, err)
return
}
if interval < 0 {
err = fmt.Errorf("invalid rate format in %q: negative interval", rate)
return
}
return
}

32
ratelimit/rate_test.go Normal file
View file

@ -0,0 +1,32 @@
package ratelimit
import (
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestParseRate(t *testing.T) {
count, dt, err := ParseRate("10/1s")
require.NoError(t, err)
require.Equal(t, 10, count)
require.Equal(t, time.Second, dt)
count, dt, err = ParseRate("1/1ms")
require.NoError(t, err)
require.Equal(t, 1, count)
require.Equal(t, time.Millisecond, dt)
}
func TestInvalidRate(t *testing.T) {
for _, invalid := range []string{
"",
"1/2",
"/10m",
"abc",
} {
_, _, err := ParseRate(invalid)
require.Errorf(t, err, "rate %q is invalid", invalid)
}
}

22
ratelimit/ratelimit.go Normal file
View file

@ -0,0 +1,22 @@
// +build !solution
package ratelimit
import (
"context"
"time"
)
// Limiter is precise rate limiter with context support.
type Limiter struct {
}
// NewLimiter returns limiter that throttles rate of successful Acquire() calls
// to maxSize events at any given interval.
func NewLimiter(maxCount int, interval time.Duration) *Limiter {
panic("not implemented")
}
func (l *Limiter) Acquire(ctx context.Context) error {
panic("not implemented")
}

120
ratelimit/ratelimit_test.go Normal file
View file

@ -0,0 +1,120 @@
package ratelimit
import (
"context"
"sync"
"testing"
"time"
"github.com/stretchr/testify/require"
"golang.org/x/sync/errgroup"
)
func TestNoRateLimit(t *testing.T) {
limit := NewLimiter(1, 0)
ctx := context.Background()
require.NoError(t, limit.Acquire(ctx))
require.NoError(t, limit.Acquire(ctx))
}
func TestBlockedRateLimit(t *testing.T) {
limit := NewLimiter(0, time.Minute)
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel()
err := limit.Acquire(ctx)
require.Equal(t, context.DeadlineExceeded, err)
}
func TestSimpleLimitCancel(t *testing.T) {
limit := NewLimiter(1, time.Minute)
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel()
require.NoError(t, limit.Acquire(ctx))
err := limit.Acquire(ctx)
require.Equal(t, context.DeadlineExceeded, err)
}
func TestStressBlocking(t *testing.T) {
const (
N = 100
G = 100
)
limit := NewLimiter(N, time.Millisecond*10)
var eg errgroup.Group
for i := 0; i < G; i++ {
eg.Go(func() error {
for j := 0; j < N; j++ {
if err := limit.Acquire(context.Background()); err != nil {
return err
}
}
return nil
})
}
require.NoError(t, eg.Wait())
}
func TestStressNoBlocking(t *testing.T) {
const (
N = 100
G = 100
)
limit := NewLimiter(N, time.Millisecond*10)
var eg errgroup.Group
for i := 0; i < G; i++ {
eg.Go(func() error {
for j := 0; j < N; j++ {
if err := limit.Acquire(context.Background()); err != nil {
return err
}
time.Sleep(time.Millisecond * 11)
}
return nil
})
}
require.NoError(t, eg.Wait())
}
func BenchmarkNoBlocking(b *testing.B) {
b.ReportAllocs()
b.SetBytes(1)
limit := NewLimiter(1, 0)
ctx := context.Background()
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
if err := limit.Acquire(ctx); err != nil {
b.Errorf("acquire failed: %v", err)
}
}
})
}
func BenchmarkReferenceMutex(b *testing.B) {
var mu sync.Mutex
var j int
for i := 0; i < b.N; i++ {
mu.Lock()
j++
mu.Unlock()
}
}