Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .travis.yml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ services:
go:
- 1.11.x
- 1.12.x
- 1.13.x
- tip

matrix:
Expand All @@ -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
1 change: 0 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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
13 changes: 4 additions & 9 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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)
}
Expand Down
9 changes: 2 additions & 7 deletions example_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package redis_rate_test

import (
"fmt"
"time"

"github.com/go-redis/redis/v7"
"github.com/go-redis/redis_rate/v8"
Expand All @@ -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)
}
Expand Down
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
4 changes: 2 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
45 changes: 34 additions & 11 deletions rate.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
28 changes: 11 additions & 17 deletions rate_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
}
Expand Down