├── LICENSE ├── README.md ├── examples └── main.go └── limiter.go /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2015 Josh Rendek 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Why? 2 | 3 | 4 | If you: 5 | 6 | * Have several distributed workers that are dependent on a rate limited service 7 | * Need to control how fast a group of workers and how quickly they hit an API/Service 8 | 9 | Then this is the package for you. 10 | 11 | 12 | ## Installation 13 | 14 | `go get github.com/joshrendek/redis-rate-limit` 15 | 16 | ## Guarantees 17 | 18 | * Global lock deadlock prevention with a TTL 19 | * Individual worker lock deadlock prevention with TTLs 20 | 21 | ## Usage 22 | 23 | See `examples/main.go` 24 | 25 | Importing: 26 | 27 | ``` 28 | import( 29 | "github.com/joshrendek/redis-rate-limit" 30 | ) 31 | ``` 32 | 33 | Using with a regular WG - we can have other parts of a job run at one 34 | concurrency (lets say this is CPU intensive) and then sending to an API. 35 | 36 | We want to run N workers but also limit the amount of workers that can 37 | concurrently access the limited resource: 38 | 39 | ``` go 40 | package main 41 | 42 | import ( 43 | "fmt" 44 | "github.com/joshrendek/redis-rate-limit" 45 | "math/rand" 46 | "sync" 47 | "time" 48 | ) 49 | 50 | var ( 51 | wg sync.WaitGroup 52 | tasks = make(chan bool, 40) 53 | ) 54 | 55 | func startWorkers() { 56 | opts := limit.Options{ 57 | Address: "localhost:6379", 58 | LockName: "wg", 59 | MaxRate: 15, 60 | LockWaitDuration: 100 * time.Millisecond, 61 | WorkerTimeout: 5 * time.Second, 62 | } 63 | limiter := limit.NewRateLimit(opts) 64 | for i := 0; i < 40; i++ { 65 | wg.Add(1) 66 | go func() { 67 | for data := range tasks { 68 | uid := limiter.Add(1) 69 | fmt.Printf("Working .... %+v\n", data) 70 | time.Sleep(time.Duration(rand.Intn(2000)) * time.Millisecond) 71 | // do some work on data 72 | limiter.Done(uid) 73 | } 74 | wg.Done() 75 | }() 76 | } 77 | } 78 | 79 | func main() { 80 | go startWorkers() 81 | 82 | for i := 0; i < 5000; i++ { 83 | tasks <- true 84 | time.Sleep(100 * time.Millisecond) 85 | } 86 | 87 | close(tasks) 88 | wg.Wait() 89 | } 90 | 91 | ``` 92 | 93 | 94 | ## License 95 | 96 | ``` 97 | 98 | The MIT License (MIT) 99 | 100 | Copyright (c) 2015 Josh Rendek 101 | 102 | Permission is hereby granted, free of charge, to any person obtaining a copy 103 | of this software and associated documentation files (the "Software"), to deal 104 | in the Software without restriction, including without limitation the rights 105 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 106 | copies of the Software, and to permit persons to whom the Software is 107 | furnished to do so, subject to the following conditions: 108 | 109 | The above copyright notice and this permission notice shall be included in all 110 | copies or substantial portions of the Software. 111 | 112 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 113 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 114 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 115 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 116 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 117 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 118 | SOFTWARE. 119 | ``` 120 | -------------------------------------------------------------------------------- /examples/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "github.com/joshrendek/redis-rate-limit" 6 | "math/rand" 7 | "sync" 8 | "time" 9 | ) 10 | 11 | var ( 12 | wg sync.WaitGroup 13 | tasks = make(chan bool, 40) 14 | ) 15 | 16 | func startWorkers() { 17 | opts := limit.Options{ 18 | Address: "localhost:6379", 19 | LockName: "wg", 20 | MaxRate: 15, 21 | LockWaitDuration: 100 * time.Millisecond, 22 | WorkerTimeout: 5 * time.Second, 23 | } 24 | limiter := limit.NewRateLimit(opts) 25 | for i := 0; i < 40; i++ { 26 | wg.Add(1) 27 | go func() { 28 | for data := range tasks { 29 | uid := limiter.Add(1) 30 | fmt.Printf("Working .... %+v\n", data) 31 | time.Sleep(time.Duration(rand.Intn(2000)) * time.Millisecond) 32 | // do some work on data 33 | limiter.Done(uid) 34 | } 35 | wg.Done() 36 | }() 37 | } 38 | } 39 | 40 | func main() { 41 | go startWorkers() 42 | 43 | for i := 0; i < 5000; i++ { 44 | tasks <- true 45 | time.Sleep(100 * time.Millisecond) 46 | } 47 | 48 | close(tasks) 49 | wg.Wait() 50 | } 51 | -------------------------------------------------------------------------------- /limiter.go: -------------------------------------------------------------------------------- 1 | package limit 2 | 3 | import ( 4 | "fmt" 5 | "github.com/satori/go.uuid" 6 | "gopkg.in/redis.v3" 7 | "sync" 8 | "time" 9 | ) 10 | 11 | // Options is used to contain config options for the limiter 12 | // Address is the address of the redis server 13 | // MaxRate is the concurrency limit 14 | // LockWaitDuration is used as a duration for busy waiting on the lock key 15 | // WorkerTimeout is used as the TTL for the worker key, this should be set 16 | // longer than you expect your workers to take. 17 | type Options struct { 18 | Address string 19 | MaxRate int64 20 | LockWaitDuration time.Duration 21 | WorkerTimeout time.Duration 22 | LockName string 23 | } 24 | 25 | // RateLimit is the limiter object exposed to control concurrency 26 | type RateLimit struct { 27 | Redis *redis.Client 28 | Opts Options 29 | } 30 | 31 | // Add is for telling wait group something is starting 32 | // A lock using SetNX is created in redis based on the l.Opts.LockName 33 | // The lock busy waits until its free and then acquires it 34 | // LockWaitDuration should be set to something reasonable so the CPU 35 | // isn't wasting cycles. 36 | // Add handles generating the UUID for the lock set and also 37 | // the lock keys that are given a TTL 38 | // The return value should be stored so you can pass the 39 | // uid to Done() 40 | func (l *RateLimit) Add(i int) string { 41 | lockCheck := l.Redis.SetNX(l.Opts.LockName+":lock", "t", 0).Val() 42 | l.Redis.Expire(l.Opts.LockName+":lock", 10*time.Second) 43 | for { 44 | if lockCheck { 45 | break 46 | } 47 | time.Sleep(l.Opts.LockWaitDuration) 48 | lockCheck = l.Redis.SetNX(l.Opts.LockName+":lock", "t", 0).Val() 49 | } 50 | l.checkRate() 51 | uid := uuid.NewV4() 52 | l.Redis.Set(l.Opts.LockName+":"+uid.String(), uid.String(), l.Opts.WorkerTimeout) 53 | l.Redis.SAdd(l.Opts.LockName, uid.String()) 54 | l.Redis.Del(l.Opts.LockName + ":lock") 55 | l.cleanLocks() 56 | return uid.String() 57 | } 58 | 59 | func (l *RateLimit) cleanLocks() { 60 | workers := l.Redis.SMembers(l.Opts.LockName).Val() 61 | for _, w := range workers { 62 | if l.Redis.Exists(l.Opts.LockName + ":" + w).Val() { 63 | continue 64 | } else { 65 | // It expired! Lets make sure it gets removed from the lock set 66 | fmt.Println("Removing: ", w) 67 | l.Redis.SRem(l.Opts.LockName, w) 68 | } 69 | } 70 | } 71 | 72 | // Done Instructs the limiter to remove the lock key 73 | // and uuid from the lock set 74 | func (l *RateLimit) Done(uid string) { 75 | l.Redis.Del(l.Opts.LockName + ":" + uid) 76 | l.Redis.SRem(l.Opts.LockName, uid) 77 | } 78 | 79 | func (l *RateLimit) checkRate() { 80 | for { 81 | wgVal := l.Redis.SCard(l.Opts.LockName).Val() 82 | if wgVal < l.Opts.MaxRate { 83 | break 84 | } 85 | } 86 | } 87 | 88 | var ( 89 | wg sync.WaitGroup 90 | tasks = make(chan bool, 40) 91 | ) 92 | 93 | // NewRateLimit is for setting up a new rate limiter with options 94 | func NewRateLimit(opts Options) RateLimit { 95 | client := redis.NewClient(&redis.Options{ 96 | Addr: opts.Address, 97 | PoolSize: 40, 98 | }) 99 | _, err := client.Ping().Result() 100 | if err != nil { 101 | panic(err) 102 | } 103 | return RateLimit{Redis: client, Opts: opts} 104 | } 105 | --------------------------------------------------------------------------------