├── README.md ├── gcra_benchmark_test.go ├── leakybucket_benchmark_test.go ├── tokenbucket_benchmark_test.go ├── script.go ├── basebucket.go ├── tokenbucket.go ├── helpers_test.go ├── gcra.go ├── leakybucket.go ├── ratelimiter_examples_test.go ├── tokenbucket_test.go ├── gcra_test.go └── leakybucket_test.go /README.md: -------------------------------------------------------------------------------- 1 | # ratelimiter 2 | 3 | A distributed token-bucket rate limiter backed by Redis. 4 | 5 | 6 | ## Algorithm 7 | 8 | See https://en.wikipedia.org/wiki/Token_bucket. 9 | 10 | 11 | ## Usage 12 | 13 | For usage and examples, see the [Godoc][1]. 14 | 15 | 16 | [1]: https://godoc.org/github.com/RussellLuo/ratelimiter -------------------------------------------------------------------------------- /gcra_benchmark_test.go: -------------------------------------------------------------------------------- 1 | package ratelimiter_test 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/RussellLuo/ratelimiter" 8 | "github.com/go-redis/redis" 9 | ) 10 | 11 | func BenchmarkGCRA_Transmit(b *testing.B) { 12 | gcra := ratelimiter.NewGCRA( 13 | &Redis{redis.NewClient(&redis.Options{ 14 | Addr: "localhost:6379", 15 | })}, 16 | "ratelimiter:gcra:benchmark", 17 | &ratelimiter.Config{ 18 | Interval: 1 * time.Second / 2, 19 | Capacity: 5, 20 | }, 21 | ) 22 | for i := 0; i < b.N; i++ { 23 | gcra.Transmit(1) 24 | } 25 | } 26 | -------------------------------------------------------------------------------- /leakybucket_benchmark_test.go: -------------------------------------------------------------------------------- 1 | package ratelimiter_test 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/RussellLuo/ratelimiter" 8 | "github.com/go-redis/redis" 9 | ) 10 | 11 | func BenchmarkLeakyBucket_Give(b *testing.B) { 12 | lb := ratelimiter.NewLeakyBucket( 13 | &Redis{redis.NewClient(&redis.Options{ 14 | Addr: "localhost:6379", 15 | })}, 16 | "ratelimiter:tokenbucket:benchmark", 17 | &ratelimiter.Config{ 18 | Interval: 1 * time.Second / 2, 19 | Capacity: 5, 20 | }, 21 | ) 22 | for i := 0; i < b.N; i++ { 23 | lb.Give(1) 24 | } 25 | } 26 | -------------------------------------------------------------------------------- /tokenbucket_benchmark_test.go: -------------------------------------------------------------------------------- 1 | package ratelimiter_test 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/RussellLuo/ratelimiter" 8 | "github.com/go-redis/redis" 9 | ) 10 | 11 | func BenchmarkTokenBucket_Take(b *testing.B) { 12 | tb := ratelimiter.NewTokenBucket( 13 | &Redis{redis.NewClient(&redis.Options{ 14 | Addr: "localhost:6379", 15 | })}, 16 | "ratelimiter:tokenbucket:benchmark", 17 | &ratelimiter.Config{ 18 | Interval: 1 * time.Second / 2, 19 | Capacity: 5, 20 | }, 21 | ) 22 | for i := 0; i < b.N; i++ { 23 | tb.Take(1) 24 | } 25 | } 26 | -------------------------------------------------------------------------------- /script.go: -------------------------------------------------------------------------------- 1 | package ratelimiter 2 | 3 | import ( 4 | "crypto/sha1" 5 | "encoding/hex" 6 | "io" 7 | ) 8 | 9 | type Redis interface { 10 | Eval(script string, keys []string, args ...interface{}) (interface{}, error) 11 | EvalSha(sha1 string, keys []string, args ...interface{}) (interface{}, error, bool) 12 | } 13 | 14 | type Script struct { 15 | redis Redis 16 | src string 17 | hash string 18 | } 19 | 20 | func NewScript(redis Redis, src string) *Script { 21 | h := sha1.New() 22 | io.WriteString(h, src) 23 | return &Script{ 24 | redis: redis, 25 | src: src, 26 | hash: hex.EncodeToString(h.Sum(nil)), 27 | } 28 | } 29 | 30 | func (s *Script) Run(keys []string, args ...interface{}) (interface{}, error) { 31 | result, err, noScript := s.redis.EvalSha(s.hash, keys, args...) 32 | if noScript { 33 | result, err = s.redis.Eval(s.src, keys, args...) 34 | } 35 | return result, err 36 | } 37 | -------------------------------------------------------------------------------- /basebucket.go: -------------------------------------------------------------------------------- 1 | package ratelimiter 2 | 3 | import ( 4 | "sync" 5 | "time" 6 | ) 7 | 8 | // Config is the bucket configuration. 9 | // Both the leaky and token bucket algorithms share the same bucket configuration. 10 | type Config struct { 11 | // Token bucket: 12 | // the interval between each addition of one token 13 | // Leaky bucket: 14 | // the interval between each leak of one unit of water 15 | Interval time.Duration 16 | 17 | // the capacity of the bucket 18 | Capacity int64 19 | } 20 | 21 | // baseBucket is a basic structure both for TokenBucket and LeakyBucket. 22 | type baseBucket struct { 23 | mu sync.RWMutex 24 | config *Config 25 | } 26 | 27 | // Config returns the bucket configuration in a concurrency-safe way. 28 | func (b *baseBucket) Config() Config { 29 | b.mu.RLock() 30 | config := *b.config 31 | b.mu.RUnlock() 32 | return config 33 | } 34 | 35 | // SetConfig updates the bucket configuration in a concurrency-safe way. 36 | func (b *baseBucket) SetConfig(config *Config) { 37 | b.mu.Lock() 38 | b.config = config 39 | b.mu.Unlock() 40 | } 41 | -------------------------------------------------------------------------------- /tokenbucket.go: -------------------------------------------------------------------------------- 1 | package ratelimiter 2 | 3 | import ( 4 | "time" 5 | ) 6 | 7 | // the Lua script that implements the Token Bucket Algorithm. 8 | // bucket.tc represents the token count. 9 | // bucket.ts represents the timestamp of the last time the bucket was refilled. 10 | const luaTokenBucket = ` 11 | local key = KEYS[1] 12 | local interval = tonumber(ARGV[1]) 13 | local capacity = tonumber(ARGV[2]) 14 | local now = tonumber(ARGV[3]) 15 | local amount = tonumber(ARGV[4]) 16 | 17 | local bucket = {tc=capacity, ts=now} 18 | local value = redis.call("get", key) 19 | if value then 20 | bucket = cjson.decode(value) 21 | end 22 | 23 | local added = math.floor((now - bucket.ts) / interval) 24 | if added > 0 then 25 | bucket.tc = math.min(bucket.tc + added, capacity) 26 | bucket.ts = bucket.ts + added * interval 27 | end 28 | 29 | if bucket.tc >= amount then 30 | bucket.tc = bucket.tc - amount 31 | bucket.ts = string.format("%.f", bucket.ts) 32 | if redis.call("set", key, cjson.encode(bucket)) then 33 | return 1 34 | end 35 | end 36 | 37 | return 0 38 | ` 39 | 40 | // TokenBucket implements the Token Bucket Algorithm. 41 | // See https://en.wikipedia.org/wiki/Token_bucket. 42 | type TokenBucket struct { 43 | baseBucket 44 | 45 | script *Script 46 | key string 47 | } 48 | 49 | // NewTokenBucket returns a new token-bucket rate limiter special for key in redis 50 | // with the specified bucket configuration. 51 | func NewTokenBucket(redis Redis, key string, config *Config) *TokenBucket { 52 | return &TokenBucket{ 53 | baseBucket: baseBucket{config: config}, 54 | script: NewScript(redis, luaTokenBucket), 55 | key: key, 56 | } 57 | } 58 | 59 | // Take takes amount tokens from the bucket. 60 | func (b *TokenBucket) Take(amount int64) (bool, error) { 61 | config := b.Config() 62 | if amount > config.Capacity { 63 | return false, nil 64 | } 65 | 66 | now := time.Now().UnixNano() 67 | result, err := b.script.Run( 68 | []string{b.key}, 69 | int64(config.Interval/time.Microsecond), 70 | config.Capacity, 71 | int64(time.Duration(now)/time.Microsecond), 72 | amount, 73 | ) 74 | if err != nil { 75 | return false, err 76 | } else { 77 | return result == int64(1), nil 78 | } 79 | } 80 | -------------------------------------------------------------------------------- /helpers_test.go: -------------------------------------------------------------------------------- 1 | package ratelimiter_test 2 | 3 | import ( 4 | "sort" 5 | "sync" 6 | "time" 7 | ) 8 | 9 | const delayedError = 6 * time.Millisecond 10 | 11 | type Func func(int64) (bool, time.Duration, error) 12 | 13 | type arg struct { 14 | WaitDuration time.Duration 15 | Concurrency int 16 | Amount int64 17 | } 18 | 19 | type rv struct { 20 | ok bool 21 | delayed time.Duration 22 | err error 23 | } 24 | 25 | type result struct { 26 | Passed int 27 | Dropped int 28 | Delayed int 29 | DelayDurations []time.Duration 30 | } 31 | 32 | func concurrentlyDo(f Func, args []arg) []result { 33 | times := len(args) 34 | rvChans := make([]chan rv, times) 35 | 36 | var wg sync.WaitGroup 37 | for i, a := range args { 38 | rvChans[i] = make(chan rv, a.Concurrency) 39 | for j := 0; j < a.Concurrency; j++ { 40 | wg.Add(1) 41 | go func(i int, a arg) { 42 | time.Sleep(a.WaitDuration) 43 | ok, delayed, err := f(a.Amount) 44 | rvChans[i] <- rv{ok: ok, delayed: delayed, err: err} 45 | wg.Done() 46 | }(i, a) 47 | } 48 | } 49 | wg.Wait() 50 | 51 | for _, c := range rvChans { 52 | close(c) 53 | } 54 | 55 | result := make([]result, times) 56 | for i, c := range rvChans { 57 | for rv := range c { 58 | if rv.ok { 59 | if rv.delayed < delayedError { 60 | result[i].Passed++ 61 | } else { 62 | result[i].Delayed++ 63 | result[i].DelayDurations = append(result[i].DelayDurations, rv.delayed) 64 | } 65 | } else { 66 | result[i].Dropped++ 67 | } 68 | } 69 | 70 | // sorted in ascending order 71 | d := result[i].DelayDurations 72 | sort.SliceStable(d, func(i, j int) bool { 73 | return d[i] < d[j] 74 | }) 75 | } 76 | 77 | return result 78 | } 79 | 80 | func deepEqual(got []result, want []result, careDelayed bool) bool { 81 | for i, r := range got { 82 | if !(r.Passed == want[i].Passed && r.Dropped == want[i].Dropped && 83 | r.Delayed == want[i].Delayed) { 84 | return false 85 | } 86 | if careDelayed { 87 | for j, d := range r.DelayDurations { 88 | wantD := want[i].DelayDurations[j] 89 | if !(d > wantD-delayedError && d < wantD+delayedError) { 90 | return false 91 | } 92 | } 93 | } 94 | } 95 | return true 96 | } 97 | -------------------------------------------------------------------------------- /gcra.go: -------------------------------------------------------------------------------- 1 | package ratelimiter 2 | 3 | import ( 4 | "time" 5 | ) 6 | 7 | // the Lua script that implements the generic cell rate algorithm. 8 | const luaGCRA = ` 9 | local key = KEYS[1] 10 | local interval = tonumber(ARGV[1]) 11 | local tolerance = tonumber(ARGV[2]) 12 | local now = tonumber(ARGV[3]) 13 | local amount = tonumber(ARGV[4]) 14 | 15 | local tat = redis.call("get", key) 16 | if tat then 17 | tat = tonumber(tat) 18 | else 19 | tat = now 20 | end 21 | 22 | local new_tat = math.max(now, tat) + amount 23 | 24 | if now >= new_tat - tolerance - interval then 25 | local ttl = math.ceil((new_tat - now)/ 1000000) 26 | if redis.call("setex", key, ttl, new_tat) then 27 | return math.max(tat - now, 0) 28 | end 29 | end 30 | 31 | return -1 32 | ` 33 | 34 | // GCRA implements the generic cell rate algorithm. 35 | // See https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm. 36 | type GCRA struct { 37 | baseBucket 38 | 39 | script *Script 40 | key string 41 | } 42 | 43 | // NewGCRA returns a new GCRA rate limiter special for key in redis 44 | // with the specified parameters. 45 | func NewGCRA(redis Redis, key string, config *Config) *GCRA { 46 | return &GCRA{ 47 | baseBucket: baseBucket{config: config}, 48 | script: NewScript(redis, luaGCRA), 49 | key: key, 50 | } 51 | } 52 | 53 | // Transmit transmits a message to the bucket. 54 | // Think of count 1 represents a message containing only one cell, and 55 | // count greater than 1 represents a message containing multiple cells. 56 | func (g *GCRA) Transmit(amount int64) (bool, time.Duration, error) { 57 | config := g.Config() 58 | if amount > config.Capacity { 59 | return false, -1, nil 60 | } 61 | 62 | // the interval between each arrival of one cell 63 | emissionInterval := config.Interval 64 | // how much earlier a cell can arrive than it would 65 | delayVariationTolerance := time.Duration(config.Capacity-1) * config.Interval 66 | 67 | now := time.Now().UnixNano() 68 | result, err := g.script.Run( 69 | []string{g.key}, 70 | int64(emissionInterval/time.Microsecond), 71 | int64(delayVariationTolerance/time.Microsecond), 72 | int64(time.Duration(now)/time.Microsecond), 73 | amount*int64(emissionInterval/time.Microsecond), 74 | ) 75 | if err != nil { 76 | return false, -1, err 77 | } else { 78 | switch delayed := result.(int64); delayed { 79 | case -1: 80 | return false, -1, nil 81 | default: 82 | return true, time.Duration(delayed) * time.Microsecond, nil 83 | } 84 | } 85 | } 86 | -------------------------------------------------------------------------------- /leakybucket.go: -------------------------------------------------------------------------------- 1 | package ratelimiter 2 | 3 | import ( 4 | "time" 5 | ) 6 | 7 | // the Lua script that implements the Leaky Bucket Algorithm as a meter. 8 | // bucket.wl represents the water level, 9 | // bucket.ts represents the timestamp of the last time the bucket was refilled. 10 | const luaLeakyBucket = ` 11 | local key = KEYS[1] 12 | local interval = tonumber(ARGV[1]) 13 | local capacity = tonumber(ARGV[2]) 14 | local now = tonumber(ARGV[3]) 15 | local amount = tonumber(ARGV[4]) 16 | 17 | local bucket = {wl=0, ts=now} 18 | local value = redis.call("get", key) 19 | if value then 20 | bucket = cjson.decode(value) 21 | end 22 | 23 | local leaks = math.floor((now - bucket.ts) / interval) 24 | if leaks > 0 then 25 | bucket.wl = math.max(bucket.wl - leaks, 0) 26 | bucket.ts = bucket.ts + leaks * interval 27 | end 28 | 29 | if bucket.wl + amount <= capacity then 30 | local delayed = math.max(bucket.wl * interval - (now - bucket.ts), 0) 31 | bucket.wl = bucket.wl + amount 32 | bucket.ts = string.format("%.f", bucket.ts) 33 | if redis.call("set", key, cjson.encode(bucket)) then 34 | return delayed 35 | end 36 | end 37 | 38 | return -1 39 | ` 40 | 41 | // LeakyBucket implements the Leaky Bucket Algorithm as a meter. 42 | // See https://en.wikipedia.org/wiki/Leaky_bucket#The_Leaky_Bucket_Algorithm_as_a_Meter. 43 | type LeakyBucket struct { 44 | baseBucket 45 | 46 | script *Script 47 | key string 48 | } 49 | 50 | // NewLeakyBucket returns a new leaky-bucket rate limiter special for key in redis 51 | // with the specified bucket configuration. 52 | func NewLeakyBucket(redis Redis, key string, config *Config) *LeakyBucket { 53 | return &LeakyBucket{ 54 | baseBucket: baseBucket{config: config}, 55 | script: NewScript(redis, luaLeakyBucket), 56 | key: key, 57 | } 58 | } 59 | 60 | // Give gives amount units of water into the bucket. 61 | func (b *LeakyBucket) Give(amount int64) (bool, time.Duration, error) { 62 | config := b.Config() 63 | if amount > config.Capacity { 64 | return false, -1, nil 65 | } 66 | 67 | now := time.Now().UnixNano() 68 | result, err := b.script.Run( 69 | []string{b.key}, 70 | int64(config.Interval/time.Microsecond), 71 | config.Capacity, 72 | int64(time.Duration(now)/time.Microsecond), 73 | amount, 74 | ) 75 | if err != nil { 76 | return false, -1, err 77 | } else { 78 | switch delayed := result.(int64); delayed { 79 | case -1: 80 | return false, -1, nil 81 | default: 82 | return true, time.Duration(delayed) * time.Microsecond, nil 83 | } 84 | } 85 | } 86 | -------------------------------------------------------------------------------- /ratelimiter_examples_test.go: -------------------------------------------------------------------------------- 1 | package ratelimiter_test 2 | 3 | import ( 4 | "fmt" 5 | "strings" 6 | "time" 7 | 8 | "github.com/RussellLuo/ratelimiter" 9 | "github.com/go-redis/redis" 10 | ) 11 | 12 | type Redis struct { 13 | client *redis.Client 14 | } 15 | 16 | func (r *Redis) Eval(script string, keys []string, args ...interface{}) (interface{}, error) { 17 | return r.client.Eval(script, keys, args...).Result() 18 | } 19 | 20 | func (r *Redis) EvalSha(sha1 string, keys []string, args ...interface{}) (interface{}, error, bool) { 21 | result, err := r.client.EvalSha(sha1, keys, args...).Result() 22 | noScript := err != nil && strings.HasPrefix(err.Error(), "NOSCRIPT ") 23 | return result, err, noScript 24 | } 25 | 26 | func ExampleTokenBucket_Take() { 27 | tb := ratelimiter.NewTokenBucket( 28 | &Redis{redis.NewClient(&redis.Options{ 29 | Addr: "localhost:6379", 30 | })}, 31 | "ratelimiter:tokenbucket:example", 32 | &ratelimiter.Config{ 33 | Interval: 1 * time.Second / 2, 34 | Capacity: 5, 35 | }, 36 | ) 37 | if ok, err := tb.Take(1); ok { 38 | fmt.Println("PASS") 39 | } else { 40 | if err != nil { 41 | fmt.Println(err.Error()) 42 | } 43 | fmt.Println("DROP") 44 | } 45 | // Output: 46 | // PASS 47 | } 48 | 49 | func ExampleLeakyBucket_Give() { 50 | lb := ratelimiter.NewLeakyBucket( 51 | &Redis{redis.NewClient(&redis.Options{ 52 | Addr: "localhost:6379", 53 | })}, 54 | "ratelimiter:leakybucket:example", 55 | &ratelimiter.Config{ 56 | Interval: 1 * time.Second / 2, 57 | Capacity: 5, 58 | }, 59 | ) 60 | if ok, delayed, err := lb.Give(1); ok { 61 | if delayed == 0 { 62 | fmt.Println("PASS") 63 | } else { 64 | fmt.Println("DELAY") 65 | } 66 | } else { 67 | if err != nil { 68 | fmt.Println(err.Error()) 69 | } 70 | fmt.Println("DROP") 71 | } 72 | // Output: 73 | // PASS 74 | } 75 | 76 | func ExampleGCRA_Transmit() { 77 | gcra := ratelimiter.NewGCRA( 78 | &Redis{redis.NewClient(&redis.Options{ 79 | Addr: "localhost:6379", 80 | })}, 81 | "ratelimiter:gcra:example", 82 | &ratelimiter.Config{ 83 | Interval: 1 * time.Second / 2, 84 | Capacity: 5, 85 | }, 86 | ) 87 | if ok, delayed, err := gcra.Transmit(1); ok { 88 | if delayed == 0 { 89 | fmt.Println("PASS") 90 | } else { 91 | fmt.Println("DELAY") 92 | } 93 | } else { 94 | if err != nil { 95 | fmt.Println(err.Error()) 96 | } 97 | fmt.Println("DROP") 98 | } 99 | // Output: 100 | // PASS 101 | } -------------------------------------------------------------------------------- /tokenbucket_test.go: -------------------------------------------------------------------------------- 1 | package ratelimiter_test 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/RussellLuo/ratelimiter" 8 | "github.com/go-redis/redis" 9 | ) 10 | 11 | func TestTokenBucket_Take(t *testing.T) { 12 | client := redis.NewClient(&redis.Options{ 13 | Addr: "localhost:6379", 14 | }) 15 | key := "ratelimiter:tokenbucket:test" 16 | 17 | bucket := ratelimiter.NewTokenBucket( 18 | &Redis{client}, 19 | key, 20 | &ratelimiter.Config{ 21 | Interval: 1 * time.Second / 2, 22 | Capacity: 5, 23 | }, 24 | ) 25 | 26 | f := func(amount int64) (bool, time.Duration, error) { 27 | ok, err := bucket.Take(amount) 28 | return ok, 0, err 29 | } 30 | 31 | cases := []struct { 32 | in []arg 33 | want []result 34 | }{ 35 | { 36 | in: []arg{ 37 | { 38 | WaitDuration: 0 * time.Second, 39 | Concurrency: 2, 40 | Amount: 1, 41 | }, 42 | { 43 | WaitDuration: 1 * time.Second, 44 | Concurrency: 2, 45 | Amount: 1, 46 | }, 47 | { 48 | WaitDuration: 2 * time.Second, 49 | Concurrency: 2, 50 | Amount: 1, 51 | }, 52 | }, 53 | want: []result{ 54 | { 55 | Passed: 2, 56 | Dropped: 0, 57 | }, 58 | { 59 | Passed: 2, 60 | Dropped: 0, 61 | }, 62 | { 63 | Passed: 2, 64 | Dropped: 0, 65 | }, 66 | }, 67 | }, 68 | { 69 | in: []arg{ 70 | { 71 | WaitDuration: 0 * time.Second, 72 | Concurrency: 4, 73 | Amount: 1, 74 | }, 75 | { 76 | WaitDuration: 1 * time.Second, 77 | Concurrency: 4, 78 | Amount: 1, 79 | }, 80 | { 81 | WaitDuration: 2 * time.Second, 82 | Concurrency: 4, 83 | Amount: 1, 84 | }, 85 | }, 86 | want: []result{ 87 | { 88 | Passed: 4, 89 | Dropped: 0, 90 | }, 91 | { 92 | Passed: 3, 93 | Dropped: 1, 94 | }, 95 | { 96 | Passed: 2, 97 | Dropped: 2, 98 | }, 99 | }, 100 | }, 101 | { 102 | in: []arg{ 103 | { 104 | WaitDuration: 0 * time.Second, 105 | Concurrency: 1, 106 | Amount: 5, 107 | }, 108 | { 109 | WaitDuration: 500 * time.Millisecond, 110 | Concurrency: 2, 111 | Amount: 1, 112 | }, 113 | { 114 | WaitDuration: 2 * time.Second, 115 | Concurrency: 5, 116 | Amount: 1, 117 | }, 118 | }, 119 | want: []result{ 120 | { 121 | Passed: 1, 122 | Dropped: 0, 123 | }, 124 | { 125 | Passed: 1, 126 | Dropped: 1, 127 | }, 128 | { 129 | Passed: 3, 130 | Dropped: 2, 131 | }, 132 | }, 133 | }, 134 | } 135 | for _, c := range cases { 136 | client.Del(key) 137 | got := concurrentlyDo(f, c.in) 138 | if !deepEqual(got, c.want, false) { 139 | t.Errorf("Got (%#v) != Want (%#v)", got, c.want) 140 | } 141 | } 142 | } 143 | -------------------------------------------------------------------------------- /gcra_test.go: -------------------------------------------------------------------------------- 1 | package ratelimiter_test 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/RussellLuo/ratelimiter" 8 | "github.com/go-redis/redis" 9 | ) 10 | 11 | func TestGCRA_Transmit(t *testing.T) { 12 | client := redis.NewClient(&redis.Options{ 13 | Addr: "localhost:6379", 14 | }) 15 | key := "ratelimiter:gcra:test" 16 | 17 | gcra := ratelimiter.NewGCRA( 18 | &Redis{client}, 19 | key, 20 | &ratelimiter.Config{ 21 | Interval: 1 * time.Second / 2, 22 | Capacity: 5, 23 | }, 24 | ) 25 | 26 | cases := []struct { 27 | in []arg 28 | want []result 29 | }{ 30 | { 31 | in: []arg{ 32 | { 33 | WaitDuration: 0 * time.Second, 34 | Concurrency: 2, 35 | Amount: 1, 36 | }, 37 | { 38 | WaitDuration: 1 * time.Second, 39 | Concurrency: 2, 40 | Amount: 1, 41 | }, 42 | { 43 | WaitDuration: 2 * time.Second, 44 | Concurrency: 2, 45 | Amount: 1, 46 | }, 47 | }, 48 | want: []result{ 49 | { 50 | Passed: 1, 51 | Dropped: 0, 52 | Delayed: 1, 53 | DelayDurations: []time.Duration{ 54 | time.Duration(500 * time.Millisecond), 55 | }, 56 | }, 57 | { 58 | Passed: 1, 59 | Dropped: 0, 60 | Delayed: 1, 61 | DelayDurations: []time.Duration{ 62 | time.Duration(500 * time.Millisecond), 63 | }, 64 | }, 65 | { 66 | Passed: 1, 67 | Dropped: 0, 68 | Delayed: 1, 69 | DelayDurations: []time.Duration{ 70 | time.Duration(500 * time.Millisecond), 71 | }, 72 | }, 73 | }, 74 | }, 75 | { 76 | in: []arg{ 77 | { 78 | WaitDuration: 0 * time.Second, 79 | Concurrency: 4, 80 | Amount: 1, 81 | }, 82 | { 83 | WaitDuration: 1 * time.Second, 84 | Concurrency: 4, 85 | Amount: 1, 86 | }, 87 | { 88 | WaitDuration: 2 * time.Second, 89 | Concurrency: 4, 90 | Amount: 1, 91 | }, 92 | }, 93 | want: []result{ 94 | { 95 | Passed: 1, 96 | Dropped: 0, 97 | Delayed: 3, 98 | DelayDurations: []time.Duration{ 99 | time.Duration(500 * time.Millisecond), 100 | time.Duration(1 * time.Second), 101 | time.Duration(1500 * time.Millisecond), 102 | }, 103 | }, 104 | { 105 | Passed: 0, 106 | Dropped: 1, 107 | Delayed: 3, 108 | DelayDurations: []time.Duration{ 109 | time.Duration(1 * time.Second), 110 | time.Duration(1500 * time.Millisecond), 111 | time.Duration(2 * time.Second), 112 | }, 113 | }, 114 | { 115 | Passed: 0, 116 | Dropped: 2, 117 | Delayed: 2, 118 | DelayDurations: []time.Duration{ 119 | time.Duration(1500 * time.Millisecond), 120 | time.Duration(2 * time.Second), 121 | }, 122 | }, 123 | }, 124 | }, 125 | { 126 | in: []arg{ 127 | { 128 | WaitDuration: 0 * time.Second, 129 | Concurrency: 1, 130 | Amount: 5, 131 | }, 132 | { 133 | WaitDuration: 500 * time.Millisecond, 134 | Concurrency: 2, 135 | Amount: 1, 136 | }, 137 | { 138 | WaitDuration: 2 * time.Second, 139 | Concurrency: 5, 140 | Amount: 1, 141 | }, 142 | }, 143 | want: []result{ 144 | { 145 | Passed: 1, 146 | Dropped: 0, 147 | Delayed: 0, 148 | }, 149 | { 150 | Passed: 0, 151 | Dropped: 1, 152 | Delayed: 1, 153 | DelayDurations: []time.Duration{ 154 | time.Duration(2 * time.Second), 155 | }, 156 | }, 157 | { 158 | Passed: 0, 159 | Dropped: 2, 160 | Delayed: 3, 161 | DelayDurations: []time.Duration{ 162 | time.Duration(1 * time.Second), 163 | time.Duration(1500 * time.Millisecond), 164 | time.Duration(2 * time.Second), 165 | }, 166 | }, 167 | }, 168 | }, 169 | } 170 | for _, c := range cases { 171 | client.Del(key) 172 | got := concurrentlyDo(gcra.Transmit, c.in) 173 | if !deepEqual(got, c.want, true) { 174 | t.Errorf("Got (%#v) != Want (%#v)", got, c.want) 175 | } 176 | } 177 | } 178 | -------------------------------------------------------------------------------- /leakybucket_test.go: -------------------------------------------------------------------------------- 1 | package ratelimiter_test 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/RussellLuo/ratelimiter" 8 | "github.com/go-redis/redis" 9 | ) 10 | 11 | func TestLeakyBucket_Give(t *testing.T) { 12 | client := redis.NewClient(&redis.Options{ 13 | Addr: "localhost:6379", 14 | }) 15 | key := "ratelimiter:leakybucket:test" 16 | 17 | bucket := ratelimiter.NewLeakyBucket( 18 | &Redis{client}, 19 | key, 20 | &ratelimiter.Config{ 21 | Interval: 1 * time.Second / 2, 22 | Capacity: 5, 23 | }, 24 | ) 25 | 26 | cases := []struct { 27 | in []arg 28 | want []result 29 | }{ 30 | { 31 | in: []arg{ 32 | { 33 | WaitDuration: 0 * time.Second, 34 | Concurrency: 2, 35 | Amount: 1, 36 | }, 37 | { 38 | WaitDuration: 1 * time.Second, 39 | Concurrency: 2, 40 | Amount: 1, 41 | }, 42 | { 43 | WaitDuration: 2 * time.Second, 44 | Concurrency: 2, 45 | Amount: 1, 46 | }, 47 | }, 48 | want: []result{ 49 | { 50 | Passed: 1, 51 | Dropped: 0, 52 | Delayed: 1, 53 | DelayDurations: []time.Duration{ 54 | time.Duration(500 * time.Millisecond), 55 | }, 56 | }, 57 | { 58 | Passed: 1, 59 | Dropped: 0, 60 | Delayed: 1, 61 | DelayDurations: []time.Duration{ 62 | time.Duration(500 * time.Millisecond), 63 | }, 64 | }, 65 | { 66 | Passed: 1, 67 | Dropped: 0, 68 | Delayed: 1, 69 | DelayDurations: []time.Duration{ 70 | time.Duration(500 * time.Millisecond), 71 | }, 72 | }, 73 | }, 74 | }, 75 | { 76 | in: []arg{ 77 | { 78 | WaitDuration: 0 * time.Second, 79 | Concurrency: 4, 80 | Amount: 1, 81 | }, 82 | { 83 | WaitDuration: 1 * time.Second, 84 | Concurrency: 4, 85 | Amount: 1, 86 | }, 87 | { 88 | WaitDuration: 2 * time.Second, 89 | Concurrency: 4, 90 | Amount: 1, 91 | }, 92 | }, 93 | want: []result{ 94 | { 95 | Passed: 1, 96 | Dropped: 0, 97 | Delayed: 3, 98 | DelayDurations: []time.Duration{ 99 | time.Duration(500 * time.Millisecond), 100 | time.Duration(1 * time.Second), 101 | time.Duration(1500 * time.Millisecond), 102 | }, 103 | }, 104 | { 105 | Passed: 0, 106 | Dropped: 1, 107 | Delayed: 3, 108 | DelayDurations: []time.Duration{ 109 | time.Duration(1 * time.Second), 110 | time.Duration(1500 * time.Millisecond), 111 | time.Duration(2 * time.Second), 112 | }, 113 | }, 114 | { 115 | Passed: 0, 116 | Dropped: 2, 117 | Delayed: 2, 118 | DelayDurations: []time.Duration{ 119 | time.Duration(1500 * time.Millisecond), 120 | time.Duration(2 * time.Second), 121 | }, 122 | }, 123 | }, 124 | }, 125 | { 126 | in: []arg{ 127 | { 128 | WaitDuration: 0 * time.Second, 129 | Concurrency: 1, 130 | Amount: 5, 131 | }, 132 | { 133 | WaitDuration: 500 * time.Millisecond, 134 | Concurrency: 2, 135 | Amount: 1, 136 | }, 137 | { 138 | WaitDuration: 2 * time.Second, 139 | Concurrency: 5, 140 | Amount: 1, 141 | }, 142 | }, 143 | want: []result{ 144 | { 145 | Passed: 1, 146 | Dropped: 0, 147 | Delayed: 0, 148 | }, 149 | { 150 | Passed: 0, 151 | Dropped: 1, 152 | Delayed: 1, 153 | DelayDurations: []time.Duration{ 154 | time.Duration(2 * time.Second), 155 | }, 156 | }, 157 | { 158 | Passed: 0, 159 | Dropped: 2, 160 | Delayed: 3, 161 | DelayDurations: []time.Duration{ 162 | time.Duration(1 * time.Second), 163 | time.Duration(1500 * time.Millisecond), 164 | time.Duration(2 * time.Second), 165 | }, 166 | }, 167 | }, 168 | }, 169 | } 170 | for _, c := range cases { 171 | client.Del(key) 172 | got := concurrentlyDo(bucket.Give, c.in) 173 | if !deepEqual(got, c.want, true) { 174 | t.Errorf("Got (%#v) != Want (%#v)", got, c.want) 175 | } 176 | } 177 | } 178 | --------------------------------------------------------------------------------