From 4075ca3f4b1159e162544fcb958b706659c7421d Mon Sep 17 00:00:00 2001 From: Vladimir Mihailenco Date: Sun, 29 Dec 2019 10:46:22 +0200 Subject: [PATCH] Tweak API --- .travis.yml | 3 ++- Makefile | 1 - README.md | 13 ++++--------- example_test.go | 9 ++------- go.mod | 2 +- go.sum | 4 ++-- rate.go | 45 ++++++++++++++++++++++++++++++++++----------- rate_test.go | 28 +++++++++++----------------- 8 files changed, 56 insertions(+), 49 deletions(-) diff --git a/.travis.yml b/.travis.yml index 8294e35..2b3216f 100644 --- a/.travis.yml +++ b/.travis.yml @@ -8,6 +8,7 @@ services: go: - 1.11.x - 1.12.x + - 1.13.x - tip matrix: @@ -20,4 +21,4 @@ env: go_import_path: github.com/go-redis/redis_rate before_install: - - curl -sfL https://install.goreleaser.com/github.com/golangci/golangci-lint.sh | sh -s -- -b $(go env GOPATH)/bin v1.17.1 + - curl -sfL https://install.goreleaser.com/github.com/golangci/golangci-lint.sh | sh -s -- -b $(go env GOPATH)/bin v1.21.0 diff --git a/Makefile b/Makefile index 9c3b8dc..57914e3 100644 --- a/Makefile +++ b/Makefile @@ -3,5 +3,4 @@ all: go test ./... -short -race go test ./... -run=NONE -bench=. -benchmem env GOOS=linux GOARCH=386 go test ./... - go vet golangci-lint run diff --git a/README.md b/README.md index b381b7c..bca118f 100644 --- a/README.md +++ b/README.md @@ -3,14 +3,13 @@ [![Build Status](https://travis-ci.org/go-redis/redis_rate.svg?branch=master)](https://travis-ci.org/go-redis/redis_rate) [![GoDoc](https://godoc.org/github.com/go-redis/redis_rate?status.svg)](https://godoc.org/github.com/go-redis/redis_rate) -This package is based on [rwz/redis-gcra](https://github.com/rwz/redis-gcra) and implements [GCRA](https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm) for rate limiting based on Redis. The code requires Redis version 3.2 or newer since it relies on [replicate_commands](https://redis.io/commands/eval#replicating-commands-instead-of-scripts) feature. +This package is based on [rwz/redis-gcra](https://github.com/rwz/redis-gcra) and implements [GCRA](https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm) (aka leaky bucket) for rate limiting based on Redis. The code requires Redis version 3.2 or newer since it relies on [replicate_commands](https://redis.io/commands/eval#replicating-commands-instead-of-scripts) feature. -```go +``` go package redis_rate_test import ( "fmt" - "time" "github.com/go-redis/redis/v7" "github.com/go-redis/redis_rate/v8" @@ -22,12 +21,8 @@ func ExampleNewLimiter() { }) _ = rdb.FlushDB().Err() - limiter := redis_rate.NewLimiter(rdb, &redis_rate.Limit{ - Burst: 10, - Rate: 10, - Period: time.Second, - }) - res, err := limiter.Allow("project:123") + limiter := redis_rate.NewLimiter(rdb) + res, err := limiter.Allow("project:123", redis_rate.PerSecond(10)) if err != nil { panic(err) } diff --git a/example_test.go b/example_test.go index 8232a62..8f8fb8a 100644 --- a/example_test.go +++ b/example_test.go @@ -2,7 +2,6 @@ package redis_rate_test import ( "fmt" - "time" "github.com/go-redis/redis/v7" "github.com/go-redis/redis_rate/v8" @@ -14,12 +13,8 @@ func ExampleNewLimiter() { }) _ = rdb.FlushDB().Err() - limiter := redis_rate.NewLimiter(rdb, &redis_rate.Limit{ - Burst: 10, - Rate: 10, - Period: time.Second, - }) - res, err := limiter.Allow("project:123") + limiter := redis_rate.NewLimiter(rdb) + res, err := limiter.Allow("project:123", redis_rate.PerSecond(10)) if err != nil { panic(err) } diff --git a/go.mod b/go.mod index 5d995df..47291e7 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,6 @@ module github.com/go-redis/redis_rate/v8 go 1.13 require ( - github.com/go-redis/redis/v7 v7.0.0-beta.3 + github.com/go-redis/redis/v7 v7.0.0-beta.4 github.com/stretchr/testify v1.4.0 ) diff --git a/go.sum b/go.sum index 58b555a..6518e6e 100644 --- a/go.sum +++ b/go.sum @@ -2,8 +2,8 @@ github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I= github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= -github.com/go-redis/redis/v7 v7.0.0-beta.3 h1:p5g3p597f83HyVxg6qvGLLDqCYWodYWC9XmSbKPwuto= -github.com/go-redis/redis/v7 v7.0.0-beta.3/go.mod h1:dohSoK1cSNPaisjbZhSk7RYyPhVx2k+4sAbJdPK5KPs= +github.com/go-redis/redis/v7 v7.0.0-beta.4 h1:p6z7Pde69EGRWvlC++y8aFcaWegyrKHzOBGo0zUACTQ= +github.com/go-redis/redis/v7 v7.0.0-beta.4/go.mod h1:xhhSbUMTsleRPur+Vgx9sUHtyN33bdjxY+9/0n9Ig8s= github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM= github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/hpcloud/tail v1.0.0 h1:nfCOvKYfkgYP8hkirhJocXT2+zOD8yUNjXaWfTlyFKI= diff --git a/rate.go b/rate.go index b4e93f0..ad59ff8 100644 --- a/rate.go +++ b/rate.go @@ -22,29 +22,52 @@ type Limit struct { Burst int } +func PerSecond(rate int) *Limit { + return &Limit{ + Rate: rate, + Period: time.Second, + Burst: rate, + } +} + +func PerMinute(rate int) *Limit { + return &Limit{ + Rate: rate, + Period: time.Minute, + Burst: rate, + } +} + +func PerHour(rate int) *Limit { + return &Limit{ + Rate: rate, + Period: time.Hour, + Burst: rate, + } +} + +//------------------------------------------------------------------------------ + // Limiter controls how frequently events are allowed to happen. type Limiter struct { - rdb rediser - limit *Limit + rdb rediser } -// NewLimiter returns a new Limiter that allows events up to rate r -// and permits bursts of at most b tokens. -func NewLimiter(rdb rediser, limit *Limit) *Limiter { +// NewLimiter returns a new Limiter. +func NewLimiter(rdb rediser) *Limiter { return &Limiter{ - rdb: rdb, - limit: limit, + rdb: rdb, } } // Allow is shorthand for AllowN(key, 1). -func (l *Limiter) Allow(key string) (*Result, error) { - return l.AllowN(key, 1) +func (l *Limiter) Allow(key string, limit *Limit) (*Result, error) { + return l.AllowN(key, limit, 1) } // AllowN reports whether n events may happen at time now. -func (l *Limiter) AllowN(key string, n int) (*Result, error) { - values := []interface{}{l.limit.Burst, l.limit.Rate, l.limit.Period.Seconds(), n} +func (l *Limiter) AllowN(key string, limit *Limit, n int) (*Result, error) { + values := []interface{}{limit.Burst, limit.Rate, limit.Period.Seconds(), n} v, err := gcra.Run(l.rdb, []string{redisPrefix + key}, values...).Result() if err != nil { return nil, err diff --git a/rate_test.go b/rate_test.go index 63afeb8..fc4d54e 100644 --- a/rate_test.go +++ b/rate_test.go @@ -10,38 +10,35 @@ import ( "github.com/go-redis/redis_rate/v8" ) -func rateLimiter(limit *redis_rate.Limit) *redis_rate.Limiter { +func rateLimiter() *redis_rate.Limiter { ring := redis.NewRing(&redis.RingOptions{ Addrs: map[string]string{"server0": ":6379"}, }) - if err := ring.FlushDb().Err(); err != nil { + if err := ring.FlushDB().Err(); err != nil { panic(err) } - return redis_rate.NewLimiter(ring, limit) + return redis_rate.NewLimiter(ring) } func TestAllow(t *testing.T) { - l := rateLimiter(&redis_rate.Limit{ - Burst: 10, - Rate: 10, - Period: time.Second, - }) + l := rateLimiter() + limit := redis_rate.PerSecond(10) - res, err := l.Allow("test_id") + res, err := l.Allow("test_id", limit) assert.Nil(t, err) assert.True(t, res.Allowed) assert.Equal(t, res.Remaining, 9) assert.Equal(t, res.RetryAfter, time.Duration(-1)) assert.InDelta(t, res.ResetAfter, 100*time.Millisecond, float64(10*time.Millisecond)) - res, err = l.AllowN("test_id", 2) + res, err = l.AllowN("test_id", limit, 2) assert.Nil(t, err) assert.True(t, res.Allowed) assert.Equal(t, res.Remaining, 7) assert.Equal(t, res.RetryAfter, time.Duration(-1)) assert.InDelta(t, res.ResetAfter, 300*time.Millisecond, float64(10*time.Millisecond)) - res, err = l.AllowN("test_id", 1000) + res, err = l.AllowN("test_id", limit, 1000) assert.Nil(t, err) assert.False(t, res.Allowed) assert.Equal(t, res.Remaining, 0) @@ -50,17 +47,14 @@ func TestAllow(t *testing.T) { } func BenchmarkAllow(b *testing.B) { - l := rateLimiter(&redis_rate.Limit{ - Burst: 1000, - Rate: 1000, - Period: time.Second, - }) + l := rateLimiter() + limit := redis_rate.PerSecond(10000) b.ResetTimer() b.RunParallel(func(pb *testing.PB) { for pb.Next() { - _, err := l.Allow("foo") + _, err := l.Allow("foo", limit) if err != nil { b.Fatal(err) }