├── .github └── workflows │ └── test.yml ├── .gitignore ├── LICENSE ├── README.md ├── attribute_limiter.go ├── attribute_limiter_test.go ├── coverage.html ├── examples ├── http-server │ └── server.go └── simple │ └── examples.go ├── go.mod ├── go.sum ├── limiter.go ├── limiter_test.go └── window.go /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | on: [push, pull_request] 2 | name: Test 3 | jobs: 4 | test: 5 | strategy: 6 | matrix: 7 | go-version: [1.15.x, 1.16.x, 1.17.x] 8 | os: [ubuntu-latest] 9 | runs-on: ${{ matrix.os }} 10 | steps: 11 | - name: Install Go 12 | uses: actions/setup-go@v2 13 | with: 14 | go-version: ${{ matrix.go-version }} 15 | - name: Checkout code 16 | uses: actions/checkout@v2 17 | - name: Test 18 | run: go test ./ -v 19 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | 8 | # Test binary, built with `go test -c` 9 | *.test 10 | 11 | # Output of the go coverage tool, specifically when used with LiteIDE 12 | *.out 13 | 14 | # Dependency directories (remove the comment below to include it) 15 | # vendor/ 16 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Narasimha Prasanna HN 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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ratelimiter 2 | 3 | ![Tests](https://github.com/Narasimha1997/ratelimiter/actions/workflows/test.yml/badge.svg) 4 | [![Go Reference](https://pkg.go.dev/badge/github.com/Narasimha1997/ratelimiter.svg)](https://pkg.go.dev/github.com/Narasimha1997/ratelimiter) 5 | 6 | A generic concurrent rate limiter library for Golang based on Sliding-window rate limitng algorithm. 7 | 8 | The implementation of rate-limiter algorithm is based on Scalable Distributed Rate Limiter algorithm used in Kong API gateway. Read [this blog](https://konghq.com/blog/how-to-design-a-scalable-rate-limiting-algorithm/) for more details. 9 | 10 | This library can be used in your codebase to rate-limit literally anything. For example, you can integrate this library to provide rate-limiting for your REST/gRPC APIs or you can use this library to 11 | rate-limit the number of go-routines spawned or number of tasks submitted to a function/module per given time interval. This library provides generic rate check APIs that can be used anywhere. The library is built with concurrency in mind from the groud up, the rate-limiter can be used across go-routines without having to worry about synchronization issues. This library also provides capability to create and manage multiple rate-limiters with different configurations assiociated with unique keys. 12 | 13 | ### How is this different from Go's official rate package? 14 | The official Go [package](https://pkg.go.dev/golang.org/x/time/rate ) provides a rate limiter implementation which uses [Token Bucket Algorithm](https://en.wikipedia.org/wiki/Token_bucket). This repository (i.e the current repository) implements rate limiting functionalities using [Sliding Window Rate Limiting Algorithm](https://konghq.com/blog/how-to-design-a-scalable-rate-limiting-algorithm/) as used in Kong API gateway. Both of these libraries provide the same functionality, but there are trade-offs between these two algorithms, [this blog](https://konghq.com/blog/how-to-design-a-scalable-rate-limiting-algorithm/) properly explains these trade-offs. Understand the trade-offs and match them with your requirements to decide which algorithm to use. 15 | 16 | ### Installation: 17 | The package can be installed as a Go module. 18 | 19 | ``` 20 | go get github.com/Narasimha1997/ratelimiter 21 | ``` 22 | 23 | ### Using the library: 24 | There are two types of rate-limiters used. 25 | 26 | #### All APIs: 27 | 1. **Generic rate-limiter**: 28 | ```go 29 | /* creates an instance of DefaultLimiter and returns it's pointer. 30 | Parameters: 31 | limit: The number of tasks to be allowd 32 | size: duration 33 | */ 34 | func NewDefaultLimiter(limit uint64, size time.Duration) *DefaultLimiter 35 | 36 | /* 37 | Kill the limiter, returns error if the limiter has been killed already. 38 | */ 39 | func (s *DefaultLimiter) Kill() error 40 | 41 | /* 42 | Makes decison whether n tasks can be allowed or not. 43 | Parameters: 44 | n: number of tasks to be processed, set this as 1 for a single task. 45 | (Example: An HTTP request) 46 | Returns (bool, error), 47 | if limiter is inactive (or it is killed), returns an error 48 | the boolean flag is either true - i.e n tasks can be allowed or false otherwise. 49 | */ 50 | func (s *DefaultLimiter) ShouldAllow(n uint64) (bool, error) 51 | 52 | /* 53 | Kill the limiter, returns error if the limiter has been killed already. 54 | */ 55 | func (s *DefaultLimiter) Kill() error 56 | ``` 57 | 58 | 2. **On-demand rate-limiter** 59 | ```go 60 | /* creates an instance of SyncLimiter and returns it's pointer. 61 | Parameters: 62 | limit: The number of tasks to be allowd 63 | size: duration 64 | */ 65 | func NewSyncLimiter(limit uint64, size time.Duration) *SyncLimiter 66 | 67 | /* 68 | Kill the limiter, returns error if the limiter has been killed already. 69 | */ 70 | func (s *SyncLimiter) Kill() error 71 | 72 | /* 73 | Makes decison whether n tasks can be allowed or not. 74 | Parameters: 75 | n: number of tasks to be processed, set this as 1 for a single task. 76 | (Example: An HTTP request) 77 | Returns (bool, error), 78 | if limiter is inactive (or it is killed), returns an error 79 | the boolean flag is either true - i.e n tasks can be allowed or false otherwise. 80 | */ 81 | func (s *SyncLimiter) ShouldAllow(n uint64) (bool, error) 82 | 83 | /* 84 | Kill the limiter, returns error if the limiter has been killed already. 85 | */ 86 | func (s *SyncLimiter) Kill() error 87 | ``` 88 | 89 | 3. **Attribute based Rate Limiter** 90 | ```go 91 | /* 92 | Creates an instance of AttributeBasedLimiter and returns it's pointer. 93 | Parameters: 94 | backgroundSliding: if set to true, DefaultLimiter will be used as an underlying limiter. 95 | else, SyncLimiter will be used. 96 | */ 97 | func NewAttributeBasedLimiter(backgroundSliding bool) *AttributeBasedLimiter 98 | 99 | /* 100 | Check if AttributeBasedLimiter has a limiter for the key. 101 | Parameters: 102 | key: a unique key string, example: IP address, token, uuid etc 103 | Returns a boolean flag, if true, the key is already present, false otherwise. 104 | */ 105 | func (a *AttributeBasedLimiter) HasKey(key string) bool 106 | 107 | /* 108 | Create a new key-limiter assiociation. 109 | Parameters: 110 | key: a unique key string, example: IP address, token, uuid etc 111 | limit: The number of tasks to be allowd 112 | size: duration 113 | Returns error if the key already exist. 114 | */ 115 | 116 | func (a *AttributeBasedLimiter) CreateNewKey( 117 | key string, limit uint64, 118 | size time.Duration, 119 | ) error 120 | 121 | /* 122 | check if AttributeBasedLimiter has a limiter for the key. 123 | Create a new key-limiter assiociation if the key not exists. 124 | Parameters: 125 | key: a unique key string, example: IP address, token, uuid etc. 126 | limit: The number of tasks to be allowd 127 | size: duration 128 | Return true if the key exists or is created successfully. 129 | */ 130 | func (a *AttributeBasedLimiter) HasOrCreateKey(key string, limit uint64, size time.Duration); 131 | 132 | /* 133 | Makes decison whether n tasks can be allowed or not. 134 | Parameters: 135 | key: a unique key string, example: IP address, token, uuid etc 136 | n: number of tasks to be processed, set this as 1 for a single task. 137 | (Example: An HTTP request) 138 | Returns (bool, error), 139 | if limiter is inactive (or it is killed) or key is not present, returns an error 140 | the boolean flag is either true - i.e n tasks can be allowed or false otherwise. 141 | */ 142 | func (a *AttributeBasedLimiter) ShouldAllow(key string, n uint64) (bool, error) 143 | 144 | /* 145 | MustShouldAllow makes decison whether n tasks can be allowed or not. 146 | Creates a new key if it does not exist. 147 | Parameters: 148 | key: a unique key string, example: IP address, token, uuid etc 149 | n: number of tasks to be processed, set this as 1 for a single task. 150 | (Example: An HTTP request) 151 | limit: The number of tasks to be allowd 152 | size: duration 153 | 154 | Returns bool. 155 | (false) when limiter is inactive (or it is killed) or n tasks can be not allowed. 156 | (true) when n tasks can be allowed or new key-limiter. 157 | */ 158 | func (a *AttributeBasedLimiter) MustShouldAllow(key string, n uint64, limit uint64, size time.Duration) bool 159 | 160 | /* 161 | Remove the key and kill its underlying limiter. 162 | Parameters: 163 | key: a unique key string, example: IP address, token, uuid etc 164 | Returns an error if the key is not present. 165 | */ 166 | func (a *AttributeBasedLimiter) DeleteKey(key string) error 167 | ``` 168 | 169 | ### Examples and Explanation of each type of rate-limiter: 170 | #### Generic rate-limiter 171 | The generic rate-limiter instance can be created if you want to have a single rate-limiter with single configuration for everything. The generic rate-limiter can be created by calling `NewDefaultLimiter()` function and by passing the `limit` and `size` as parameters. Example: 172 | 173 | ```go 174 | func GenericRateLimiter() { 175 | /* create an instance of Limiter. 176 | format: NewLimiter(limit uint64, size time.Duration), 177 | where: 178 | limit: The number of tasks/items that should be allowed. 179 | size: The window size, i.e the time interval during which the limit 180 | should be imposed. 181 | To summarize, if limit = 100 and duration = 5s, then allow 100 items per 5 seconds 182 | */ 183 | 184 | limiter := ratelimiter.NewDefaultLimiter( 185 | 100, time.Second*5, 186 | ) 187 | 188 | /* 189 | Cleaning up the limiter: Once the limiter is no longer required, 190 | the underlying goroutines and resources used by the limiter can be cleaned up. 191 | This can be done using: 192 | limiter.Kill(), 193 | Returns an error if the limiter is already being killed. 194 | */ 195 | 196 | defer limiter.Kill() 197 | 198 | /* 199 | the limiter provides ShouldAllow(N uint64) function which 200 | returns true/false if N items/tasks can be allowed during current 201 | time interval. 202 | 203 | An error is returned if the limiter is already killed. 204 | */ 205 | 206 | // ShouldAllow(N uint64) -> returns bool, error 207 | 208 | // should return true 209 | fmt.Println(limiter.ShouldAllow(60)) 210 | // should return false, because (60 + 50 = 110) > 100 during this window 211 | fmt.Println(limiter.ShouldAllow(50)) 212 | // sleep for some time 213 | time.Sleep(5 * time.Second) 214 | // should return true, because the previous window has been slided over 215 | fmt.Println(limiter.ShouldAllow(20)) 216 | } 217 | ``` 218 | 219 | #### On demand window sliding: 220 | The previous method i.e the Generic Rate limiter spins up a background goroutine that takes care of sliding the rate-limiting window whenever it's size expires, because of this, rate-limiting check function `ShouldAllow` has fewer steps and takes very less time to make decision. But if your application manages a large number of Limiters, for example a web-server that performs rate-limiting across hundreds of different IPs, then your `AttributeBasedRateLimiter` spins up a goroutine for each unique IP and thus lot of such routines needs to be manitanied, this might induce scheduling pressure. 221 | 222 | An alternative solution is to use a rate-limiter does not require a background routine, instead the window is sliding is taken care by `ShouldAllow` function itself, this method can be used to maintain large number of rate limiters without any scheduling pressure. This limiter is called `SyncLimiter` and can be used just like `DefaultLimiter`, because `SyncLimiter` and `DefaultLimiter` are built on top of the same `Limiter` interface. To use this, just replace `NewDefaultLimiter` with `NewSyncLimiter` 223 | ```go 224 | ...... 225 | 226 | limiter := ratelimiter.NewSyncLimiter( 227 | 100, time.Second*5, 228 | ) 229 | ...... 230 | ``` 231 | 232 | #### Attribute based rate-limiter: 233 | Attribute based rate-limiter can hold multiple rate-limiters with different configurations in a map 234 | of type. Each limiter is uniquely identified by a key. Calling `NewAttributeBasedLimiter()` will create an empty rate limiter with no entries. 235 | 236 | ```go 237 | func AttributeRateLimiter() { 238 | /* 239 | Attribute based rate-limiter can hold multiple 240 | rate-limiters with different configurations in a map 241 | of type. Each limiter is uniquely identified 242 | by a key. Calling NewAttributeBasedLimiter() will create an empty 243 | rate limiter with no entries. 244 | */ 245 | /* 246 | Attribute based rate-limiter has a boolean parameter called: 247 | `backgroundSliding` - if set to true, the attribute based rate-limiter 248 | uses Limiter instance and each Limiter instance have it's own background goroutine 249 | to manage sliding window. This might be resource expensive for large number of attributes, 250 | but is faster than SyncLimiter. 251 | 252 | Disable this, i.e pass `false` if you want to manage large number of attributes 253 | in less memory and compute, sacrifcing a minimal amount of performance. 254 | */ 255 | limiter := ratelimiter.NewAttributeBasedLimiter(true) 256 | 257 | /* 258 | Now we are adding a new entry to the limiter, we pass: 259 | key: A string that is used to uniquely identify the rate-limiter. 260 | limit: The number of tasks/items that should be allowed. 261 | size: The window size, i.e the time interval during which the limit 262 | should be imposed. 263 | 264 | returns error if the key already exists in the map. 265 | */ 266 | // we have two articles here (for example) 267 | article_ids := []string{"article_id=10", "article_id=11"} 268 | 269 | // for article_id=10, allow 10 tasks/items per every second 270 | err := limiter.CreateNewKey(&article_ids[0], 10, 5*time.Second) 271 | if err != nil { 272 | log.Fatalln(err) 273 | } 274 | // for article_id=11, allow 100 tasks/items per every 6 minutes 275 | err = limiter.CreateNewKey(&article_ids[1], 100, 6*time.Minute) 276 | if err != nil { 277 | log.Fatalln(err) 278 | } 279 | // rates can be checked by passing key and N as parameters 280 | // Can I make 8 requests to article_id=10 during this time window? 281 | 282 | // ShouldAllow(key *string, N uint64) returns (bool, error) 283 | // the bool is true/false, true if it can be allowed 284 | // false if it cant be allowed. 285 | // error if key is not found. 286 | 287 | fmt.Println(limiter.ShouldAllow(&article_ids[0], 8)) 288 | // Can I make 104 requests to article_id=11 during this time window? 289 | fmt.Println(limiter.ShouldAllow(&article_ids[0], 104)) 290 | 291 | /* 292 | Other functions: 293 | 1. HasKey: to check if the attribute already has given key 294 | call: HasKey(key string) function. 295 | Example: limiter.HasKey(&article_id[0]) 296 | Returns a bool, true if exists, false otherwise 297 | 298 | 2. DeleteKey: to remove the key from attribute map 299 | call: DeleteKey(key string) function. 300 | Example: limiter.DeleteKey(&article_id[1]) 301 | Returns an error, if key was not in the map. 302 | */ 303 | } 304 | ``` 305 | 306 | ### Using ratelimiter as a middleware with HTTP web server: 307 | ratelimiter is pluggable and can be used anywhere. This code snippet shows how it can be used with 308 | Go's standard HTTP library when building a web server: 309 | 310 | ```go 311 | ..... 312 | // allow 100 requests every 5 seconds 313 | limiter := ratelimiter.NewSyncLimiter(100, time.Second * 5) 314 | 315 | // register the handler 316 | rateLimiterHandler := func(next http.Handler) http.Handler { 317 | return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 318 | allowed, err := limiter.ShouldAllow(1) 319 | if err != nil { 320 | log.Fatalln(err) 321 | } 322 | if allowed { 323 | next.ServeHTTP(w, r) 324 | } 325 | }) 326 | } 327 | 328 | // create a test route handler 329 | ponger := func(w http.ResponseWriter, r *http.Request) { 330 | w.Write([]byte("Pong!!")) 331 | } 332 | 333 | // attach the ratelimiter middleware: 334 | muxServer := http.NewServeMux() 335 | muxServer.Handle("/", rateLimiterHandler( 336 | http.HandlerFunc(ponger), 337 | )) 338 | 339 | // start the server 340 | err := http.ListenAndServe(":6000", muxServer) 341 | if err != nil { 342 | log.Fatalln(err) 343 | } 344 | ``` 345 | The complete example can be found at `examples/http-server/server.go`. 346 | `curl` was used to simulate `X` requests per second and following was the output, as logged. 347 | ``` 348 | ................................ 349 | 2021/10/05 14:12:38 Iteration: 7, Requests received: 522, Allowed: 99 350 | 2021/10/05 14:12:43 Iteration: 8, Requests received: 533, Allowed: 101 351 | 2021/10/05 14:12:48 Iteration: 9, Requests received: 515, Allowed: 100 352 | 2021/10/05 14:12:53 Iteration: 10, Requests received: 505, Allowed: 100 353 | 2021/10/05 14:12:58 Iteration: 11, Requests received: 508, Allowed: 100 354 | 2021/10/05 14:13:03 Iteration: 12, Requests received: 474, Allowed: 100 355 | 2021/10/05 14:13:08 Iteration: 13, Requests received: 495, Allowed: 100 356 | 2021/10/05 14:13:13 Iteration: 14, Requests received: 478, Allowed: 100 357 | .................................. 358 | ``` 359 | The ratelimiter was able to balance the requested limit as specified. 360 | If you have installed the package, you can simply run the webserver as follows: 361 | ``` 362 | go run examples/http-server/server.go 363 | ``` 364 | 365 | ### Testing 366 | Tests are written in `attribute_limiter_test.go` and `limiter_test.go` files. To execute the tests, 367 | simply run: 368 | ``` 369 | go test ./ -v 370 | ``` 371 | 372 | These are some of the results from tests: 373 | 1. **Single goroutine, Generic limiter**: This test configures the rate-limiter to allow 100 requests/sec and fires 500 requests/sec with a time gap of 2ms each, allowed requests are counted and is tested with difference +/- 3. The same test is run for 10 samples. Here are the results: 374 | 375 | ``` 376 | === RUN TestLimiterAccuracy 377 | Iteration 1, Allowed tasks: 100, passed rate limiting accuracy test. 378 | Iteration 2, Allowed tasks: 101, passed rate limiting accuracy test. 379 | Iteration 3, Allowed tasks: 100, passed rate limiting accuracy test. 380 | Iteration 4, Allowed tasks: 100, passed rate limiting accuracy test. 381 | Iteration 5, Allowed tasks: 100, passed rate limiting accuracy test. 382 | Iteration 6, Allowed tasks: 100, passed rate limiting accuracy test. 383 | Iteration 7, Allowed tasks: 101, passed rate limiting accuracy test. 384 | Iteration 8, Allowed tasks: 100, passed rate limiting accuracy test. 385 | Iteration 9, Allowed tasks: 100, passed rate limiting accuracy test. 386 | Iteration 10, Allowed tasks: 100, passed rate limiting accuracy test. 387 | --- PASS: TestLimiterAccuracy (10.01s) 388 | ``` 389 | 390 | 2. **4 goroutines, Generic Limiter**: This test configures the limiter to allow 100 requests/sec and spins up 4 goroutines, the same limiter is shared across all the routines. Each goroutine generates 500 requests/sec with 2ms time gap between 2 requests. Allowed requests are counted per each goroutine, the result sum of all counts should be almost equal to 100. The accuracy is measured considering +/- 3 as error offset. The same test is conducted 10 times. Here are the results: 391 | 392 | ``` 393 | === RUN TestConcurrentLimiterAccuracy 394 | Iteration 1, Allowed tasks: 101, passed rate limiting accuracy test. 395 | Iteration 2, Allowed tasks: 100, passed rate limiting accuracy test. 396 | Iteration 3, Allowed tasks: 100, passed rate limiting accuracy test. 397 | Iteration 4, Allowed tasks: 100, passed rate limiting accuracy test. 398 | Iteration 5, Allowed tasks: 100, passed rate limiting accuracy test. 399 | Iteration 6, Allowed tasks: 100, passed rate limiting accuracy test. 400 | Iteration 7, Allowed tasks: 100, passed rate limiting accuracy test. 401 | Iteration 8, Allowed tasks: 100, passed rate limiting accuracy test. 402 | Iteration 9, Allowed tasks: 100, passed rate limiting accuracy test. 403 | Iteration 10, Allowed tasks: 100, passed rate limiting accuracy test. 404 | --- PASS: TestConcurrentLimiterAccuracy (10.01s) 405 | ``` 406 | 407 | 3. **2 goroutines, 2 attribute keys, Attribute based limiter**: An attribute based limiter is created with 2 keys, these keys are configured to allow 100 requests/sec and 123 requests/sec respectively. Two goroutines are created and same attribute based limiter is shared across. Each goroutine produces 500 requests/sec per key. The overall count is then verified for each goroutine with error offset of +/- 3. Here are the results: 408 | 409 | ``` 410 | === RUN TestAttributeBasedLimiterAccuracy 411 | Iteration 1, Allowed tasks: 100, passed rate limiting accuracy test. 412 | Iteration 1, Allowed tasks: 123, passed rate limiting accuracy test. 413 | Iteration 2, Allowed tasks: 101, passed rate limiting accuracy test. 414 | Iteration 2, Allowed tasks: 124, passed rate limiting accuracy test. 415 | Iteration 3, Allowed tasks: 100, passed rate limiting accuracy test. 416 | Iteration 3, Allowed tasks: 123, passed rate limiting accuracy test. 417 | Iteration 4, Allowed tasks: 100, passed rate limiting accuracy test. 418 | Iteration 4, Allowed tasks: 123, passed rate limiting accuracy test. 419 | Iteration 5, Allowed tasks: 100, passed rate limiting accuracy test. 420 | Iteration 5, Allowed tasks: 123, passed rate limiting accuracy test. 421 | --- PASS: TestAttributeBasedLimiterAccuracy (5.00s) 422 | ``` 423 | 424 | **Code coverage**: 425 | To generate code coverage report, execute: 426 | ``` 427 | go test -coverprofile=c.out 428 | ``` 429 | 430 | This should print the following after running all the tests. 431 | ``` 432 | coverage: 99.0% of statements 433 | ok github.com/Narasimha1997/ratelimiter 25.099s 434 | ``` 435 | 436 | You can also save the results as HTML for more detailed code view of the coverage. 437 | ``` 438 | go tool cover -html=c.out -o coverage.html 439 | ``` 440 | 441 | This will generate a file called `coverage.html`. The `coverage.html` is provided in the repo which is pre-generated. 442 | 443 | **Benchmarks**: 444 | Benchmarks can be executed by running: 445 | ``` 446 | go test -bench=. 447 | ``` 448 | 449 | Current benchmarks are as follows: 450 | ``` 451 | goos: linux 452 | goarch: amd64 453 | pkg: github.com/Narasimha1997/ratelimiter 454 | cpu: Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz 455 | BenchmarkDefaultLimiter-12 11732958 85.61 ns/op 456 | BenchmarkSyncLimiter-12 7047988 175.9 ns/op 457 | BenchmarkConcurrentDefaultLimiter-12 7017625 163.9 ns/op 458 | BenchmarkConcurrentSyncLimiter-12 4132976 256.3 ns/op 459 | PASS 460 | ok github.com/Narasimha1997/ratelimiter 46.408s 461 | ``` 462 | 463 | #### Notes on test: 464 | The testing code produces 500 requests/sec with `2ms` precision time gap between each request. The accuracy of this `2ms` time tick generation can differ from platform to platform, even a small difference of 500 micorseconds can add up together and give more time for test to run in the end because of clock drift, as a result the error offset +/- 3 might not always work. 465 | ### Contributing 466 | Feel free to raise issues, make pull requests or suggest new features. 467 | -------------------------------------------------------------------------------- /attribute_limiter.go: -------------------------------------------------------------------------------- 1 | package ratelimiter 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | "time" 7 | ) 8 | 9 | // AttributeMap is a custom map type of string key and Limiter instance as value 10 | type AttributeMap map[string]Limiter 11 | 12 | // AttributeBasedLimiter is an instance that can manage multiple rate limiter instances 13 | // with different configutations. 14 | type AttributeBasedLimiter struct { 15 | attributeMap AttributeMap 16 | m sync.Mutex 17 | syncMode bool 18 | } 19 | 20 | // HasKey check if AttributeBasedLimiter has a limiter for the key. 21 | // 22 | // Parameters: 23 | // 24 | // 1. key: a unique key string, example: IP address, token, uuid etc 25 | // 26 | // Returns a boolean flag, if true, the key is already present, false otherwise. 27 | func (a *AttributeBasedLimiter) HasKey(key string) bool { 28 | a.m.Lock() 29 | _, ok := a.attributeMap[key] 30 | a.m.Unlock() 31 | return ok 32 | } 33 | 34 | // CreateNewKey create a new key-limiter assiociation. 35 | // 36 | // Parameters: 37 | // 38 | // 1. key: a unique key string, example: IP address, token, uuid etc 39 | // 40 | // 2. limit: The number of tasks to be allowd 41 | // 42 | // 3. size: duration 43 | // 44 | // Returns error if the key already exists. 45 | func (a *AttributeBasedLimiter) CreateNewKey(key string, limit uint64, size time.Duration) error { 46 | a.m.Lock() 47 | defer a.m.Unlock() 48 | 49 | return a.createNewKey(key, limit, size) 50 | } 51 | 52 | func (a *AttributeBasedLimiter) createNewKey(key string, limit uint64, size time.Duration) error { 53 | if _, ok := a.attributeMap[key]; ok { 54 | return fmt.Errorf( 55 | "key %s is already defined", key, 56 | ) 57 | } 58 | 59 | // create a new entry: 60 | if !a.syncMode { 61 | a.attributeMap[key] = NewDefaultLimiter(limit, size) 62 | } else { 63 | a.attributeMap[key] = NewSyncLimiter(limit, size) 64 | } 65 | return nil 66 | } 67 | 68 | // HasOrCreateKey check if AttributeBasedLimiter has a limiter for the key. 69 | // Create a new key-limiter assiociation if the key not exists. 70 | // 71 | // Parameters: 72 | // 73 | // 1. key: a unique key string, example: IP address, token, uuid etc 74 | // 75 | // 2. limit: The number of tasks to be allowd 76 | // 77 | // 3. size: duration 78 | // 79 | // Return true if the key exists or is created successfully. 80 | func (a *AttributeBasedLimiter) HasOrCreateKey(key string, limit uint64, size time.Duration) bool { 81 | a.m.Lock() 82 | defer a.m.Unlock() 83 | 84 | if _, ok := a.attributeMap[key]; ok { 85 | return true 86 | } 87 | 88 | if err := a.createNewKey(key, limit, size); err == nil { 89 | return true 90 | } 91 | 92 | return false 93 | } 94 | 95 | // ShouldAllow makes decison whether n tasks can be allowed or not. 96 | // 97 | // Parameters: 98 | // 99 | // key: a unique key string, example: IP address, token, uuid etc 100 | // 101 | // n: number of tasks to be processed, set this as 1 for a single task. 102 | // (Example: An HTTP request) 103 | // 104 | // Returns (bool, error). 105 | // (false, error) when limiter is inactive (or it is killed) or key is not present. 106 | // (true/false, nil) if key exists and n tasks can be allowed or not. 107 | func (a *AttributeBasedLimiter) ShouldAllow(key string, n uint64) (bool, error) { 108 | a.m.Lock() 109 | defer a.m.Unlock() 110 | 111 | limiter, ok := a.attributeMap[key] 112 | if ok { 113 | return limiter.ShouldAllow(n) 114 | } 115 | 116 | return false, fmt.Errorf("key %s not found", key) 117 | } 118 | 119 | // MustShouldAllow makes decison whether n tasks can be allowed or not. 120 | // 121 | // Parameters: 122 | // 123 | // key: a unique key string, example: IP address, token, uuid etc 124 | // 125 | // n: number of tasks to be processed, set this as 1 for a single task. 126 | // (Example: An HTTP request) 127 | // 128 | // limit: The number of tasks to be allowd 129 | // 130 | // size: duration 131 | // 132 | // Returns bool. 133 | // (false) when limiter is inactive (or it is killed) or n tasks can be not allowed. 134 | // (true) when n tasks can be allowed or new key-limiter. 135 | func (a *AttributeBasedLimiter) MustShouldAllow(key string, n uint64, limit uint64, size time.Duration) bool { 136 | a.m.Lock() 137 | defer a.m.Unlock() 138 | 139 | if limiter, ok := a.attributeMap[key]; ok { 140 | allowed, err := limiter.ShouldAllow(n) 141 | return allowed && err == nil 142 | } 143 | 144 | err := a.createNewKey(key, limit, size) 145 | if err != nil { 146 | return err == nil 147 | } 148 | 149 | // check ratelimiter on newly created key: 150 | limiter := a.attributeMap[key] 151 | allowed, err := limiter.ShouldAllow(n) 152 | return allowed && err == nil 153 | } 154 | 155 | // DeleteKey remove the key and kill its underlying limiter. 156 | // 157 | // Parameters: 158 | // 159 | // 1.key: a unique key string, example: IP address, token, uuid etc 160 | // 161 | // Returns an error if the key is not present. 162 | func (a *AttributeBasedLimiter) DeleteKey(key string) error { 163 | 164 | a.m.Lock() 165 | defer a.m.Unlock() 166 | 167 | if limiter, ok := a.attributeMap[key]; ok { 168 | err := limiter.Kill() 169 | if err != nil { 170 | return err 171 | } 172 | delete(a.attributeMap, key) 173 | return nil 174 | } 175 | 176 | return fmt.Errorf("key %s not found", key) 177 | } 178 | 179 | // NewAttributeBasedLimiter creates an instance of AttributeBasedLimiter and returns it's pointer. 180 | // 181 | // Parameters: 182 | // 183 | // 1. backgroundSliding: if set to true, DefaultLimiter will be used as an underlying limiter, 184 | // else, SyncLimiter will be used. 185 | func NewAttributeBasedLimiter(backgroundSliding bool) *AttributeBasedLimiter { 186 | return &AttributeBasedLimiter{ 187 | attributeMap: make(AttributeMap), 188 | syncMode: !backgroundSliding, 189 | } 190 | } 191 | -------------------------------------------------------------------------------- /attribute_limiter_test.go: -------------------------------------------------------------------------------- 1 | package ratelimiter 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | "testing" 7 | "time" 8 | ) 9 | 10 | func TestAttributeMapGetSetDelete(t *testing.T) { 11 | 12 | duration := 1 * time.Second 13 | limit := 100 14 | 15 | attributeLimiter := NewAttributeBasedLimiter(true) 16 | 17 | // create a new key attribute: 18 | // Example scenario, the rate-limiter 19 | testKey1 := "/api/getArticle?id=10" 20 | testKey2 := "/api/getArticle?id=20" 21 | testKey3 := "/api/getArticle?id=30" 22 | 23 | // check keys: 24 | if attributeLimiter.HasKey(testKey1) { 25 | t.Fatalf( 26 | "AttributeBasedLimiter.HasKey() failed, returned true for non-existing key %s", 27 | testKey1, 28 | ) 29 | } 30 | 31 | if attributeLimiter.HasKey(testKey2) { 32 | t.Fatalf( 33 | "AttributeBasedLimiter.HasKey() failed, returned true for non-existing key %s", 34 | testKey2, 35 | ) 36 | } 37 | 38 | // create key: 39 | if err := attributeLimiter.CreateNewKey(testKey1, uint64(limit), duration); err != nil { 40 | t.Fatalf( 41 | "AttributeBasedLimiter.CreateNewKey() failed, returned error on creating key %s, Error: %v\n", 42 | testKey1, err, 43 | ) 44 | } 45 | 46 | if err := attributeLimiter.CreateNewKey(testKey2, uint64(limit), duration); err != nil { 47 | t.Fatalf( 48 | "AttributeBasedLimiter.CreateNewKey() failed, returned error on creating key %s, Error: %v\n", 49 | testKey2, err, 50 | ) 51 | } 52 | 53 | // create an already existing key: 54 | if err := attributeLimiter.CreateNewKey(testKey2, uint64(limit), duration); err == nil { 55 | t.Fatalf( 56 | "AttributeBasedLimiter.CreateNewKey() failed, did not return error when creating existing key %s\n", 57 | testKey2, 58 | ) 59 | } 60 | 61 | // create key: 62 | if ok := attributeLimiter.HasOrCreateKey(testKey3, uint64(limit), duration); !ok { 63 | t.Fatalf( 64 | "AttributeBasedLimiter.HasOrCreateKey() failed, returned false on creating key %s", 65 | testKey3, 66 | ) 67 | } 68 | 69 | // create an already existing key: 70 | if ok := attributeLimiter.HasOrCreateKey(testKey2, uint64(limit), duration); !ok { 71 | t.Fatalf( 72 | "AttributeBasedLimiter.HasOrCreateKey() failed, returned false for existing key %s", 73 | testKey2, 74 | ) 75 | } 76 | 77 | // check existing keys: 78 | if !attributeLimiter.HasKey(testKey1) { 79 | t.Fatalf( 80 | "AttributeBasedLimiter.HasKey() failed, returned false for existing key %s", 81 | testKey1, 82 | ) 83 | } 84 | 85 | if !attributeLimiter.HasKey(testKey1) { 86 | t.Fatalf( 87 | "AttributeBasedLimiter.HasKey() failed, returned false for existing key %s", 88 | testKey2, 89 | ) 90 | } 91 | 92 | // remove key 93 | if err := attributeLimiter.DeleteKey(testKey1); err != nil { 94 | t.Fatalf( 95 | "AttributeBasedLimiter.DeleteKey() failed, returned error when removing existing key %s, Error: %v", 96 | testKey1, err, 97 | ) 98 | } 99 | 100 | if err := attributeLimiter.DeleteKey(testKey2); err != nil { 101 | t.Fatalf( 102 | "AttributeBasedLimiter.DeleteKey() failed, returned error when removing existing key %s, Error: %v", 103 | testKey2, err, 104 | ) 105 | } 106 | 107 | // check keys again: 108 | if attributeLimiter.HasKey(testKey1) { 109 | t.Fatalf( 110 | "AttributeBasedLimiter.HasKey() failed, returned true for non-existing key %s", 111 | testKey1, 112 | ) 113 | } 114 | 115 | if attributeLimiter.HasKey(testKey2) { 116 | t.Fatalf( 117 | "AttributeBasedLimiter.HasKey() failed, returned true for non-existing key %s", 118 | testKey2, 119 | ) 120 | } 121 | 122 | // check ShouldAllow on non-existing key: 123 | if _, err := attributeLimiter.ShouldAllow("noKey", 5); err == nil { 124 | t.Fatalf( 125 | "AttributeBasedLimiter.ShouldAllow() failed, did not return error when checking non-existing key.", 126 | ) 127 | } 128 | 129 | // check ShouldAllow on non-existing key: 130 | if ok := attributeLimiter.MustShouldAllow("newKey", 5, uint64(limit), duration); !ok { 131 | t.Fatalf( 132 | "AttributeBasedLimiter.MustShouldAllow() failed, did not return false when checking non-existing key.", 133 | ) 134 | } 135 | 136 | // check ShouldAllow on existing key: 137 | if _, err := attributeLimiter.ShouldAllow("newKey", 5); err != nil { 138 | t.Fatalf( 139 | "AttributeBasedLimiter.ShouldAllow() failed, did not return error when checking existing key.", 140 | ) 141 | } 142 | 143 | // check ShouldAllow on existing key: 144 | if ok := attributeLimiter.MustShouldAllow("newKey", 5, uint64(limit), duration); !ok { 145 | t.Fatalf( 146 | "AttributeBasedLimiter.MustShouldAllow() failed, did not return false when checking non-existing key.", 147 | ) 148 | } 149 | 150 | // Remove the non-existing key: 151 | if err := attributeLimiter.DeleteKey("noKey"); err == nil { 152 | t.Fatalf( 153 | "AttributeBasedLimiter.DeleteKey failed, did not return error when deleting non-existing key.", 154 | ) 155 | } 156 | } 157 | 158 | func TestAttributeBasedLimiterAccuracy(t *testing.T) { 159 | 160 | // number of unique keys to be tested 161 | keys := []string{"/api/getArticle?id=10", "/api/getArticle?id=20"} 162 | 163 | // key1 has limit of 100 hits/sec and key2 has 123 hits/sec allowed. 164 | limits := []uint64{100, 123} 165 | counters := make([]uint64, len(keys)) 166 | 167 | // per second window 168 | duration := 1 * time.Second 169 | 170 | // 10 samples will be executed. 171 | nRuns := 5 172 | 173 | // test with accuracy +/- 3, modify this variable to 174 | // test accuracy for various error offsets, 0 is the most 175 | // ideal case. 176 | var allowanceRange uint64 = 15 177 | 178 | sharedLimiter := NewAttributeBasedLimiter(true) 179 | 180 | for idx, key := range keys { 181 | err := sharedLimiter.CreateNewKey(key, limits[idx], duration) 182 | if err != nil { 183 | t.Fatalf("%v", err) 184 | } 185 | } 186 | 187 | routine := func(key string, idx int, wg *sync.WaitGroup) { 188 | defer wg.Done() 189 | j := 0 190 | counters[idx] = 0 191 | for range time.Tick(2 * time.Millisecond) { 192 | allowed, err := sharedLimiter.ShouldAllow(key, 1) 193 | if err != nil { 194 | break 195 | } 196 | 197 | if allowed { 198 | counters[idx]++ 199 | } 200 | 201 | j++ 202 | 203 | if j%500 == 0 { 204 | break 205 | } 206 | } 207 | } 208 | 209 | // run for nRuns: 210 | for i := 0; i < nRuns; i++ { 211 | wg := sync.WaitGroup{} 212 | for idx, key := range keys { 213 | wg.Add(1) 214 | go routine(key, idx, &wg) 215 | } 216 | 217 | wg.Wait() 218 | 219 | // loop over the keys and check rate-limit: 220 | for idx, count := range counters { 221 | limit := limits[idx] 222 | 223 | // check accuracy of counter 224 | if (limit-allowanceRange) <= count && count <= (limit+allowanceRange) { 225 | fmt.Printf( 226 | "Iteration %d, Allowed tasks: %d, passed rate limiting accuracy test.\n", 227 | i+1, count, 228 | ) 229 | } else { 230 | t.Fatalf( 231 | "Accuracy test failed, expected results to be in +/- %d error range, but got %d", 232 | allowanceRange, count, 233 | ) 234 | } 235 | } 236 | } 237 | } 238 | 239 | func TestAttributeBasedLimiterAccuracySync(t *testing.T) { 240 | 241 | // number of unique keys to be tested 242 | keys := []string{"/api/getArticle?id=10", "/api/getArticle?id=20"} 243 | 244 | // key1 has limit of 100 hits/sec and key2 has 123 hits/sec allowed. 245 | limits := []uint64{100, 123} 246 | counters := make([]uint64, len(keys)) 247 | 248 | // per second window 249 | duration := 1 * time.Second 250 | 251 | // 10 samples will be executed. 252 | nRuns := 6 253 | 254 | // test with accuracy +/- 3, modify this variable to 255 | // test accuracy for various error offsets, 0 is the most 256 | // ideal case. 257 | var allowanceRange uint64 = 15 258 | 259 | sharedLimiter := NewAttributeBasedLimiter(false) 260 | 261 | isDry := true 262 | 263 | for idx, key := range keys { 264 | err := sharedLimiter.CreateNewKey(key, limits[idx], duration) 265 | if err != nil { 266 | t.Fatalf("%v", err) 267 | } 268 | } 269 | 270 | routine := func(key string, idx int, wg *sync.WaitGroup) { 271 | defer wg.Done() 272 | j := 0 273 | counters[idx] = 0 274 | for range time.Tick(2 * time.Millisecond) { 275 | allowed, err := sharedLimiter.ShouldAllow(key, 1) 276 | if err != nil { 277 | break 278 | } 279 | 280 | if allowed { 281 | counters[idx]++ 282 | } 283 | 284 | j++ 285 | 286 | if j%500 == 0 { 287 | break 288 | } 289 | } 290 | } 291 | 292 | // run for nRuns: 293 | for i := 0; i < nRuns; i++ { 294 | wg := sync.WaitGroup{} 295 | for idx, key := range keys { 296 | wg.Add(1) 297 | go routine(key, idx, &wg) 298 | } 299 | 300 | wg.Wait() 301 | 302 | // loop over the keys and check rate-limit: 303 | if !isDry { 304 | for idx, count := range counters { 305 | limit := limits[idx] 306 | 307 | // check accuracy of counter 308 | if (limit-allowanceRange) <= count && count <= (limit+allowanceRange) { 309 | fmt.Printf( 310 | "Iteration %d, Allowed tasks: %d, passed rate limiting accuracy test.\n", 311 | i, count, 312 | ) 313 | } else { 314 | t.Fatalf( 315 | "Accuracy test failed, expected results to be in +/- %d error range, but got %d", 316 | allowanceRange, count, 317 | ) 318 | } 319 | } 320 | } 321 | 322 | isDry = false 323 | } 324 | } 325 | -------------------------------------------------------------------------------- /coverage.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | ratelimiter: Go Coverage Report 7 | 52 | 53 | 54 |
55 | 66 |
67 | not tracked 68 | 69 | not covered 70 | covered 71 | 72 |
73 |
74 |
75 | 76 | 263 | 264 | 499 | 500 | 550 | 551 |
552 | 553 | 580 | 581 | -------------------------------------------------------------------------------- /examples/http-server/server.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "log" 5 | "net/http" 6 | "time" 7 | 8 | "github.com/Narasimha1997/ratelimiter" 9 | ) 10 | 11 | func main() { 12 | perIntervalRecv := 0 13 | perIntervalAllowed := 0 14 | nIterations := 0 15 | 16 | duration := time.Second * 5 17 | requestsAllowed := uint64(100) 18 | 19 | reporter := func() { 20 | for { 21 | time.Sleep(duration) 22 | log.Printf( 23 | "Iteration: %d, Requests received: %d, Allowed: %d", 24 | nIterations+1, perIntervalRecv, perIntervalAllowed, 25 | ) 26 | 27 | perIntervalRecv = 0 28 | perIntervalAllowed = 0 29 | nIterations++ 30 | } 31 | } 32 | 33 | // add a middleware: 34 | limiter := ratelimiter.NewSyncLimiter(requestsAllowed, duration) 35 | 36 | rateLimiterHandler := func(next http.Handler) http.Handler { 37 | return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 38 | allowed, err := limiter.ShouldAllow(1) 39 | if err != nil { 40 | log.Fatalln(err) 41 | } 42 | 43 | perIntervalRecv++ 44 | 45 | if allowed { 46 | perIntervalAllowed++ 47 | next.ServeHTTP(w, r) 48 | } 49 | }) 50 | } 51 | 52 | ponger := func(w http.ResponseWriter, r *http.Request) { 53 | w.Write([]byte("Pong!!")) 54 | } 55 | 56 | // attach the ratelimiter middleware: 57 | muxServer := http.NewServeMux() 58 | muxServer.Handle("/", rateLimiterHandler( 59 | http.HandlerFunc(ponger), 60 | )) 61 | 62 | // start reporter routine: 63 | go reporter() 64 | err := http.ListenAndServe(":6000", muxServer) 65 | if err != nil { 66 | log.Fatalln(err) 67 | } 68 | } 69 | -------------------------------------------------------------------------------- /examples/simple/examples.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "log" 6 | "sync" 7 | "time" 8 | 9 | "github.com/Narasimha1997/ratelimiter" 10 | ) 11 | 12 | func GenericRateLimiter() { 13 | /* create an instance of Limiter. 14 | format: NewLimiter(limit uint64, size time.Duration), 15 | where: 16 | limit: The number of tasks/items that should be allowed. 17 | size: The window size, i.e the time interval during which the limit 18 | should be imposed. 19 | To summarize, if limit = 100 and duration = 5s, then allow 100 items per 5 seconds 20 | */ 21 | 22 | limiter := ratelimiter.NewSyncLimiter( 23 | 100, time.Second*5, 24 | ) 25 | 26 | /* 27 | Cleaning up the limiter: Once the limiter is no longer required, 28 | the underlying goroutines and resources used by the limiter can be cleaned up. 29 | This can be done using: 30 | limiter.Kill(), 31 | Returns an error if the limiter is already being killed. 32 | */ 33 | 34 | defer limiter.Kill() 35 | 36 | /* 37 | the limiter provides ShouldAllow(N uint64) function which 38 | returns true/false if N items/tasks can be allowed during current 39 | time interval. 40 | 41 | An error is returned if the limiter is already killed. 42 | */ 43 | 44 | // ShouldAllow(N uint64) -> returns bool, error 45 | 46 | // should return true 47 | fmt.Println(limiter.ShouldAllow(60)) 48 | // should return false, because (60 + 50 = 110) > 100 during this window 49 | fmt.Println(limiter.ShouldAllow(50)) 50 | // sleep for some time 51 | time.Sleep(5 * time.Second) 52 | // should return true, because the previous window has been slided over 53 | fmt.Println(limiter.ShouldAllow(80)) 54 | } 55 | 56 | func AttributeRateLimiter() { 57 | /* 58 | Attribute based rate-limiter can hold multiple 59 | rate-limiters with different configurations in a map 60 | of type. Each limiter is uniquely identified 61 | by a key. Calling NewAttributeBasedLimiter() will create an empty 62 | rate limiter with no entries. 63 | */ 64 | limiter := ratelimiter.NewAttributeBasedLimiter(true) 65 | 66 | /* 67 | Now we are adding a new entry to the limiter, we pass: 68 | key: A string that is used to uniquely identify the rate-limiter. 69 | limit: The number of tasks/items that should be allowed. 70 | size: The window size, i.e the time interval during which the limit 71 | should be imposed. 72 | 73 | returns error if the key already exists in the map. 74 | */ 75 | // we have two articles here (for example) 76 | article_ids := []string{"article_id=10", "article_id=11"} 77 | 78 | // for article_id=10, allow 10 tasks/items per every second 79 | err := limiter.CreateNewKey(article_ids[0], 10, 5*time.Second) 80 | if err != nil { 81 | log.Fatalln(err) 82 | } 83 | // for article_id=11, allow 100 tasks/items per every 6 minutes 84 | err = limiter.CreateNewKey(article_ids[1], 100, 6*time.Minute) 85 | if err != nil { 86 | log.Fatalln(err) 87 | } 88 | // rates can be checked by passing key and N as parameters 89 | // Can I make 8 requests to article_id=10 during this time window? 90 | 91 | // ShouldAllow(key string, N uint64) returns (bool, error) 92 | // the bool is true/false, true if it can be allowed 93 | // false if it cant be allowed. 94 | // error if key is not found. 95 | 96 | fmt.Println(limiter.ShouldAllow(article_ids[0], 8)) 97 | // Can I make 104 requests to article_id=11 during this time window? 98 | fmt.Println(limiter.ShouldAllow(article_ids[0], 104)) 99 | 100 | /* 101 | Other functions: 102 | 1. HasKey: to check if the attribute already has given key 103 | call: HasKey(key *string) function. 104 | Example: limiter.HasKey(&article_id[0]) 105 | Returns a bool, true if exists, false otherwise 106 | 107 | 2. DeleteKey: to remove the key from attribute map 108 | call: DeleteKey(key *string) function. 109 | Example: limiter.DeleteKey(&article_id[1]) 110 | Returns an error, if key was not in the map. 111 | */ 112 | } 113 | 114 | func TestLoop() { 115 | limiter := ratelimiter.NewAttributeBasedLimiter(false) 116 | limiter.CreateNewKey("test", 10, time.Second*5) 117 | 118 | counter := 0 119 | lockVal := sync.Mutex{} 120 | 121 | go (func() { 122 | 123 | for { 124 | allowed, _ := limiter.ShouldAllow("test", 1) 125 | if allowed { 126 | lockVal.Lock() 127 | counter++ 128 | lockVal.Unlock() 129 | } 130 | 131 | time.Sleep(300 * time.Millisecond) 132 | 133 | } 134 | })() 135 | 136 | for { 137 | time.Sleep(time.Second * 5) 138 | fmt.Println(counter) 139 | lockVal.Lock() 140 | counter = 0 141 | lockVal.Unlock() 142 | } 143 | } 144 | 145 | /*func main() { 146 | 147 | fmt.Println("Generic rate limiter:") 148 | GenericRateLimiter() 149 | fmt.Println("Attribute based rate limiter:") 150 | AttributeRateLimiter() 151 | 152 | fmt.Println("Test loop with rate limited values:") 153 | TestLoop() 154 | }*/ 155 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/Narasimha1997/ratelimiter 2 | 3 | go 1.15 4 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | go.opentelemetry.io v0.1.0 h1:EANZoRCOP+A3faIlw/iN6YEWoYb1vleZRKm1EvH8T48= 2 | go.opentelemetry.io/otel v1.0.0 h1:qTTn6x71GVBvoafHK/yaRUmFzI4LcONZD0/kXxl5PHI= 3 | -------------------------------------------------------------------------------- /limiter.go: -------------------------------------------------------------------------------- 1 | package ratelimiter 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "sync" 7 | "time" 8 | ) 9 | 10 | // Limiter is an interface that is implemented by DefaultLimiter and SyncLimiter 11 | type Limiter interface { 12 | Kill() error 13 | ShouldAllow(n uint64) (bool, error) 14 | } 15 | 16 | // DefaultLimiter maintains all the structures used for rate limting using a background goroutine. 17 | type DefaultLimiter struct { 18 | previous *Window 19 | current *Window 20 | lock sync.Mutex 21 | size time.Duration 22 | limit uint64 23 | killed bool 24 | windowContext context.Context 25 | cancelFn func() 26 | } 27 | 28 | // ShouldAllow makes decison whether n tasks can be allowed or not. 29 | // 30 | // Parameters: 31 | // 32 | // 1. n: number of tasks to be processed, set this as 1 for a single task. (Example: An HTTP request) 33 | // 34 | // Returns (bool, error). (false, error) if limiter is inactive (or it is killed). Otherwise, 35 | // (true/false, nil) depending on whether n tasks can be allowed or not. 36 | func (l *DefaultLimiter) ShouldAllow(n uint64) (bool, error) { 37 | l.lock.Lock() 38 | defer l.lock.Unlock() 39 | 40 | if l.killed { 41 | return false, fmt.Errorf("function ShouldAllow called on an inactive instance") 42 | } 43 | 44 | if l.limit == 0 || l.size < time.Millisecond { 45 | return false, fmt.Errorf("invalid limiter configuration") 46 | } 47 | 48 | currentTime := time.Now() 49 | currentWindowBoundary := currentTime.Sub(l.current.getStartTime()) 50 | 51 | w := float64(l.size-currentWindowBoundary) / float64(l.size) 52 | 53 | currentSlidingRequests := uint64(w*float64(l.previous.count)) + l.current.count 54 | 55 | if currentSlidingRequests+n > l.limit { 56 | return false, nil 57 | } 58 | 59 | // add current request count to window of current count 60 | l.current.updateCount(n) 61 | return true, nil 62 | } 63 | 64 | func (l *DefaultLimiter) progressiveWindowSlider() { 65 | for { 66 | select { 67 | case <-l.windowContext.Done(): 68 | return 69 | default: 70 | toSleepDuration := l.size - time.Since(l.current.getStartTime()) 71 | time.Sleep(toSleepDuration) 72 | l.lock.Lock() 73 | // make current as previous and create a new current window 74 | l.previous.setStateFrom(l.current) 75 | l.current.resetToTime(time.Now()) 76 | l.lock.Unlock() 77 | } 78 | } 79 | } 80 | 81 | // Kill the limiter, returns error if the limiter has been killed already. 82 | func (l *DefaultLimiter) Kill() error { 83 | l.lock.Lock() 84 | defer l.lock.Unlock() 85 | 86 | if l.killed { 87 | return fmt.Errorf("called Kill on already killed limiter") 88 | } 89 | 90 | defer l.cancelFn() 91 | l.killed = true 92 | return nil 93 | } 94 | 95 | // NewDefaultLimiter creates an instance of DefaultLimiter and returns it's pointer. 96 | // 97 | // Parameters: 98 | // 99 | // 1. limit: The number of tasks to be allowd 100 | // 101 | // 2. size: duration 102 | func NewDefaultLimiter(limit uint64, size time.Duration) *DefaultLimiter { 103 | previous := NewWindow(0, time.Unix(0, 0)) 104 | current := NewWindow(0, time.Unix(0, 0)) 105 | 106 | childCtx, cancelFn := context.WithCancel(context.Background()) 107 | 108 | limiter := &DefaultLimiter{ 109 | previous: previous, 110 | current: current, 111 | lock: sync.Mutex{}, 112 | size: size, 113 | limit: limit, 114 | killed: false, 115 | windowContext: childCtx, 116 | cancelFn: cancelFn, 117 | } 118 | 119 | go limiter.progressiveWindowSlider() 120 | return limiter 121 | } 122 | 123 | // SyncLimiter maintains all the structures used for rate limting on demand. 124 | type SyncLimiter struct { 125 | previous *Window 126 | current *Window 127 | lock sync.Mutex 128 | size time.Duration 129 | limit uint64 130 | killed bool 131 | } 132 | 133 | func (s *SyncLimiter) getNSlidesSince(now time.Time) (time.Duration, time.Time) { 134 | sizeAlignedTime := now.Truncate(s.size) 135 | timeSinceStart := sizeAlignedTime.Sub(s.current.getStartTime()) 136 | 137 | return timeSinceStart / s.size, sizeAlignedTime 138 | } 139 | 140 | // ShouldAllow makes decison whether n tasks can be allowed or not. 141 | // 142 | // Parameters: 143 | // 144 | // 1. n: number of tasks to be processed, set this as 1 for a single task. (Example: An HTTP request) 145 | // 146 | // Returns (bool, error). (false, error) if limiter is inactive (or it is killed). Otherwise, 147 | // (true/false, error) depending on whether n tasks can be allowed or not. 148 | func (s *SyncLimiter) ShouldAllow(n uint64) (bool, error) { 149 | s.lock.Lock() 150 | defer s.lock.Unlock() 151 | 152 | if s.killed { 153 | return false, fmt.Errorf("function ShouldAllow called on an inactive instance") 154 | } 155 | 156 | if s.limit == 0 || s.size < time.Millisecond { 157 | return false, fmt.Errorf("invalid limiter configuration") 158 | } 159 | 160 | currentTime := time.Now() 161 | 162 | // advance the window on demand, as this doesn't make use of goroutine. 163 | nSlides, alignedCurrentTime := s.getNSlidesSince(currentTime) 164 | 165 | // window slide shares both current and previous windows. 166 | if nSlides == 1 { 167 | s.previous.setToState( 168 | alignedCurrentTime.Add(-s.size), 169 | s.current.count, 170 | ) 171 | 172 | s.current.resetToTime( 173 | alignedCurrentTime, 174 | ) 175 | 176 | } else if nSlides > 1 { 177 | s.previous.resetToTime( 178 | alignedCurrentTime.Add(-s.size), 179 | ) 180 | s.current.resetToTime( 181 | alignedCurrentTime, 182 | ) 183 | } 184 | 185 | currentWindowBoundary := currentTime.Sub(s.current.getStartTime()) 186 | 187 | w := float64(s.size-currentWindowBoundary) / float64(s.size) 188 | 189 | currentSlidingRequests := uint64(w*float64(s.previous.count)) + s.current.count 190 | 191 | if currentSlidingRequests+n > s.limit { 192 | return false, nil 193 | } 194 | 195 | // add current request count to window of current count 196 | s.current.updateCount(n) 197 | return true, nil 198 | } 199 | 200 | // Kill the limiter, returns error if the limiter has been killed already. 201 | func (s *SyncLimiter) Kill() error { 202 | s.lock.Lock() 203 | defer s.lock.Unlock() 204 | 205 | if s.killed { 206 | return fmt.Errorf("called Kill on already killed limiter") 207 | } 208 | 209 | // kill is a dummy implementation for SyncLimiter, 210 | // because there is no need of stopping a go-routine. 211 | s.killed = true 212 | return nil 213 | } 214 | 215 | // NewSyncLimiter creates an instance of SyncLimiter and returns it's pointer. 216 | // 217 | // Parameters: 218 | // 219 | // 1. limit: The number of tasks to be allowd 220 | // 221 | // 2. size: duration 222 | func NewSyncLimiter(limit uint64, size time.Duration) *SyncLimiter { 223 | current := NewWindow(0, time.Unix(0, 0)) 224 | previous := NewWindow(0, time.Unix(0, 0)) 225 | 226 | return &SyncLimiter{ 227 | previous: previous, 228 | current: current, 229 | lock: sync.Mutex{}, 230 | killed: false, 231 | size: size, 232 | limit: limit, 233 | } 234 | } 235 | -------------------------------------------------------------------------------- /limiter_test.go: -------------------------------------------------------------------------------- 1 | package ratelimiter 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | "testing" 7 | "time" 8 | ) 9 | 10 | func TestInvalidLimiterConfiguration(t *testing.T) { 11 | limiter := NewDefaultLimiter(10, time.Microsecond*800) 12 | if _, err := limiter.ShouldAllow(3); err == nil { 13 | t.Fatalf("ShouldAllow() failed, did not throw error when window size <= 1 millisecond") 14 | } 15 | 16 | limiter1 := NewSyncLimiter(0, 10*time.Second) 17 | if _, err := limiter1.ShouldAllow(10); err == nil { 18 | t.Fatalf("ShouldAllow() failed, did not throw error when limit == 0") 19 | } 20 | } 21 | 22 | func TestLimiterAccuracy(t *testing.T) { 23 | 24 | nRuns := 10 25 | var count uint64 = 0 26 | 27 | // Time duration of the window. 28 | duration := time.Second * 1 29 | 30 | // 100 tasks must be allowed to execute 31 | // for every `duration` interval. 32 | var limit uint64 = 100 33 | 34 | // test with accuracy +/- 3, modify this variable to 35 | // test accuracy for various error offsets, 0 is the most 36 | // ideal case. 37 | var allowanceRange uint64 = 20 38 | 39 | // will be set to true once the go routine completes all `nRuns` 40 | 41 | limiter := NewDefaultLimiter(limit, duration) 42 | defer limiter.Kill() 43 | 44 | for i := 0; i < nRuns; i++ { 45 | count = 0 46 | nTicks := 0 47 | for range time.Tick(time.Millisecond * 2) { 48 | 49 | canAllow, err := limiter.ShouldAllow(1) 50 | if err != nil { 51 | t.Fatalf("%v", err) 52 | } 53 | 54 | if canAllow { 55 | count++ 56 | } 57 | 58 | nTicks++ 59 | 60 | if nTicks%500 == 0 { 61 | break 62 | } 63 | } 64 | 65 | if (limit-allowanceRange) <= count && count <= (limit+allowanceRange) { 66 | fmt.Printf( 67 | "Iteration %d, Allowed tasks: %d, passed rate limiting accuracy test.\n", 68 | i+1, count, 69 | ) 70 | count = 0 71 | } else { 72 | t.Fatalf( 73 | "Accuracy test failed, expected results to be in +/- %d error range, but got %d", 74 | allowanceRange, count, 75 | ) 76 | } 77 | } 78 | } 79 | 80 | func TestConcurrentLimiterAccuracy(t *testing.T) { 81 | nRuns := 10 82 | duration := time.Second * 1 83 | 84 | // 100 tasks must be allowed to execute 85 | // for every `duration` interval. 86 | var limit uint64 = 100 87 | 88 | // create a limiter, that is shared across go routines: 89 | sharedLimiter := NewDefaultLimiter(limit, duration) 90 | defer sharedLimiter.Kill() 91 | 92 | // launch N go-routines: 93 | nRoutines := 4 94 | 95 | // test with accuracy +/- 3, modify this variable to 96 | // test accuracy for various error offsets, 0 is the most 97 | // ideal case. 98 | var allowanceRange uint64 = 20 99 | 100 | counterSlice := make([]uint64, nRoutines) 101 | 102 | routine := func(idx int, wg *sync.WaitGroup) { 103 | 104 | defer wg.Done() 105 | 106 | // no need of mutex locking the counterSlice 107 | // because each goroutine has access to only a 108 | // unique index `idx` of the slice. 109 | counterSlice[idx] = 0 110 | j := 0 111 | 112 | // Use of time.Tick in production is discouraged. 113 | // time.Tick cannot be stopped, we are using it because 114 | // this is a test code. 115 | for range time.Tick(2 * time.Millisecond) { 116 | canAllow, err := sharedLimiter.ShouldAllow(1) 117 | if err != nil { 118 | break 119 | } 120 | 121 | if canAllow { 122 | counterSlice[idx]++ 123 | } 124 | 125 | j++ 126 | if j%500 == 0 { 127 | break 128 | } 129 | } 130 | } 131 | 132 | for i := 0; i < nRuns; i++ { 133 | // create a wait group and 134 | wg := sync.WaitGroup{} 135 | for j := 0; j < nRoutines; j++ { 136 | wg.Add(1) 137 | go routine(j, &wg) 138 | } 139 | 140 | wg.Wait() 141 | 142 | // sum over the counterSlice and check accuracy: 143 | var count uint64 = 0 144 | for _, partialCount := range counterSlice { 145 | count += partialCount 146 | } 147 | 148 | // check accuracy of counter 149 | if (limit-allowanceRange) <= count && count <= (limit+allowanceRange) { 150 | fmt.Printf( 151 | "Iteration %d, Allowed tasks: %d, passed rate limiting accuracy test.\n", 152 | i+1, count, 153 | ) 154 | } else { 155 | t.Fatalf( 156 | "Accuracy test failed, expected results to be in +/- %d error range, but got %d", 157 | allowanceRange, count, 158 | ) 159 | } 160 | } 161 | } 162 | 163 | func TestConcurrentSyncLimiter(t *testing.T) { 164 | nRuns := 10 165 | duration := time.Second * 1 166 | 167 | // 100 tasks must be allowed to execute 168 | // for every `duration` interval. 169 | var limit uint64 = 100 170 | 171 | // create a limiter, that is shared across go routines: 172 | sharedLimiter := NewSyncLimiter(limit, duration) 173 | defer sharedLimiter.Kill() 174 | 175 | // launch N go-routines: 176 | nRoutines := 4 177 | 178 | // dry run, this will allow rate-limiter to stabilize: 179 | isDry := true 180 | 181 | // test with accuracy +/- 3, modify this variable to 182 | // test accuracy for various error offsets, 0 is the most 183 | // ideal case. 184 | var allowanceRange uint64 = 20 185 | 186 | counterSlice := make([]uint64, nRoutines) 187 | 188 | routine := func(idx int, wg *sync.WaitGroup) { 189 | 190 | defer wg.Done() 191 | 192 | // no need of mutex locking the counterSlice 193 | // because each goroutine has access to only a 194 | // unique index `idx` of the slice. 195 | counterSlice[idx] = 0 196 | j := 0 197 | 198 | // Use of time.Tick in production is discouraged. 199 | // time.Tick cannot be stopped, we are using it because 200 | // this is a test code. 201 | for range time.Tick(2 * time.Millisecond) { 202 | canAllow, err := sharedLimiter.ShouldAllow(1) 203 | if err != nil { 204 | break 205 | } 206 | 207 | if canAllow { 208 | counterSlice[idx]++ 209 | } 210 | 211 | j++ 212 | if j%500 == 0 { 213 | break 214 | } 215 | } 216 | } 217 | 218 | for i := 0; i < nRuns; i++ { 219 | // create a wait group and 220 | wg := sync.WaitGroup{} 221 | for j := 0; j < nRoutines; j++ { 222 | wg.Add(1) 223 | go routine(j, &wg) 224 | } 225 | 226 | wg.Wait() 227 | 228 | // sum over the counterSlice and check accuracy: 229 | var count uint64 = 0 230 | for _, partialCount := range counterSlice { 231 | count += partialCount 232 | } 233 | 234 | // check accuracy of counter 235 | if !isDry { 236 | if (limit-allowanceRange) <= count && count <= (limit+allowanceRange) { 237 | fmt.Printf( 238 | "Iteration %d, Allowed tasks: %d, passed rate limiting accuracy test.\n", 239 | i, count, 240 | ) 241 | } else { 242 | t.Fatalf( 243 | "Accuracy test failed, expected results to be in +/- %d error range, but got %d", 244 | allowanceRange, count, 245 | ) 246 | } 247 | } 248 | 249 | isDry = false 250 | } 251 | } 252 | 253 | func TestLimiterCleanup(t *testing.T) { 254 | var limit uint64 = 10 255 | var size time.Duration = 5 * time.Second 256 | 257 | limiter := NewDefaultLimiter(limit, size) 258 | 259 | // call allow check on limiter: 260 | _, err := limiter.ShouldAllow(1) 261 | if err != nil { 262 | t.Fatalf("Error when calling ShouldAllow() on active limiter, Error: %v", err) 263 | } 264 | 265 | // kill the limiter: 266 | if err = limiter.Kill(); err != nil { 267 | t.Fatalf("Failed to kill an active limiter, Error: %v", err) 268 | } 269 | 270 | // try to call kill again on already killed limiter: 271 | if err = limiter.Kill(); err == nil { 272 | t.Fatalf("Failed to throw error when Kill() was called on the same limiter twice.") 273 | } 274 | 275 | // call ShouldAllow() on inactive limiter, this should throw an error 276 | _, err = limiter.ShouldAllow(4) 277 | if err == nil { 278 | t.Fatalf("Calling ShouldAllow() on inactive limiter did not throw any errors.") 279 | } 280 | } 281 | 282 | func TestSyncLimiterCleanup(t *testing.T) { 283 | var limit uint64 = 10 284 | var size time.Duration = 5 * time.Second 285 | 286 | limiter := NewSyncLimiter(limit, size) 287 | 288 | // call allow check on limiter: 289 | _, err := limiter.ShouldAllow(1) 290 | if err != nil { 291 | t.Fatalf("Error when calling ShouldAllow() on active limiter, Error: %v", err) 292 | } 293 | 294 | // kill the limiter: 295 | if err = limiter.Kill(); err != nil { 296 | t.Fatalf("Failed to kill an active limiter, Error: %v", err) 297 | } 298 | 299 | // try to call kill again on already killed limiter: 300 | if err = limiter.Kill(); err == nil { 301 | t.Fatalf("Failed to throw error when Kill() was called on the same limiter twice.") 302 | } 303 | 304 | // call ShouldAllow() on inactive limiter, this should throw an error 305 | _, err = limiter.ShouldAllow(4) 306 | if err == nil { 307 | t.Fatalf("Calling ShouldAllow() on inactive limiter did not throw any errors.") 308 | } 309 | } 310 | 311 | func BenchmarkDefaultLimiter(b *testing.B) { 312 | limiter := NewDefaultLimiter(100, 1*time.Second) 313 | 314 | for i := 0; i < b.N; i++ { 315 | _, err := limiter.ShouldAllow(1) 316 | if err != nil { 317 | b.Fatalf("Error when calling ShouldAllow() on active limiter, Error: %v", err) 318 | } 319 | } 320 | } 321 | 322 | func BenchmarkSyncLimiter(b *testing.B) { 323 | limiter := NewSyncLimiter(100, 1*time.Second) 324 | 325 | for i := 0; i < b.N; i++ { 326 | _, err := limiter.ShouldAllow(1) 327 | if err != nil { 328 | b.Fatalf("Error when calling ShouldAllow() on active limiter, Error: %v", err) 329 | } 330 | } 331 | } 332 | 333 | func BenchmarkConcurrentDefaultLimiter(b *testing.B) { 334 | limiter := NewDefaultLimiter(100, 1*time.Second) 335 | 336 | b.RunParallel(func(p *testing.PB) { 337 | for p.Next() { 338 | _, err := limiter.ShouldAllow(1) 339 | if err != nil { 340 | b.Fatalf("Error when calling ShouldAllow() on active limiter, Error: %v", err) 341 | } 342 | } 343 | }) 344 | } 345 | 346 | func BenchmarkConcurrentSyncLimiter(b *testing.B) { 347 | limiter := NewSyncLimiter(100, 1*time.Second) 348 | 349 | b.RunParallel(func(p *testing.PB) { 350 | for p.Next() { 351 | _, err := limiter.ShouldAllow(1) 352 | if err != nil { 353 | b.Fatalf("Error when calling ShouldAllow() on active limiter, Error: %v", err) 354 | } 355 | } 356 | }) 357 | } 358 | -------------------------------------------------------------------------------- /window.go: -------------------------------------------------------------------------------- 1 | package ratelimiter 2 | 3 | import ( 4 | "time" 5 | ) 6 | 7 | // Window represents the structure of timing-window at given point of time. 8 | type Window struct { 9 | count uint64 10 | startTime time.Time 11 | } 12 | 13 | func (w *Window) updateCount(n uint64) { 14 | w.count += n 15 | } 16 | 17 | func (w *Window) getStartTime() time.Time { 18 | return w.startTime 19 | } 20 | 21 | func (w *Window) setStateFrom(other *Window) { 22 | w.count = other.count 23 | w.startTime = other.startTime 24 | } 25 | 26 | func (w *Window) resetToTime(startTime time.Time) { 27 | w.count = 0 28 | w.startTime = startTime 29 | } 30 | 31 | func (w *Window) setToState(startTime time.Time, count uint64) { 32 | w.startTime = startTime 33 | w.count = count 34 | } 35 | 36 | // Creates and returns a pointer to the new Window instance. 37 | // 38 | // Parameters: 39 | // 40 | // 1. count: The initial count of the window. 41 | // 42 | // 2. startTime: The initial starting time of the window. 43 | func NewWindow(count uint64, startTime time.Time) *Window { 44 | 45 | return &Window{ 46 | count: count, 47 | startTime: startTime, 48 | } 49 | } 50 | --------------------------------------------------------------------------------