├── .document ├── .github ├── CODEOWNERS └── workflows │ ├── ci.yml │ ├── publish.yml │ └── rails_main_testing.yml ├── .gitignore ├── Changelog.md ├── Gemfile ├── Gemfile.lock ├── LICENSE ├── README.md ├── Rakefile ├── bin └── bundle_all ├── gemfiles ├── common.rb ├── rails7.0.gemfile ├── rails7.0.gemfile.lock ├── rails7.1.gemfile ├── rails7.1.gemfile.lock ├── rails7.2.gemfile ├── rails7.2.gemfile.lock └── rails_main.gemfile ├── lib ├── prop.rb └── prop │ ├── interval_strategy.rb │ ├── key.rb │ ├── leaky_bucket_strategy.rb │ ├── limiter.rb │ ├── middleware.rb │ ├── options.rb │ ├── rate_limited.rb │ └── version.rb ├── prop.gemspec └── test ├── helper.rb ├── test_changelog.rb ├── test_interval_strategy.rb ├── test_key.rb ├── test_leaky_bucket_strategy.rb ├── test_limiter.rb ├── test_middleware.rb ├── test_options.rb ├── test_prop.rb └── test_rate_limited.rb /.document: -------------------------------------------------------------------------------- 1 | README.rdoc 2 | lib/**/*.rb 3 | bin/* 4 | features/**/*.feature 5 | LICENSE 6 | -------------------------------------------------------------------------------- /.github/CODEOWNERS: -------------------------------------------------------------------------------- 1 | ** @zendesk/bolt @zendesk/classic-operations 2 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | --- 2 | name: CI 3 | on: [push] 4 | jobs: 5 | main: 6 | name: Tests 7 | runs-on: ubuntu-latest 8 | strategy: 9 | fail-fast: false 10 | matrix: 11 | ruby: 12 | - '3.1' 13 | - '3.2' 14 | - '3.3' 15 | - '3.4' 16 | - 'jruby-9.4' 17 | gemfile: 18 | - rails7.0 19 | - rails7.1 20 | - rails7.2 21 | include: 22 | - { ruby: '3.4', gemfile: 'rails_main' } 23 | env: 24 | BUNDLE_GEMFILE: gemfiles/${{ matrix.gemfile }}.gemfile 25 | steps: 26 | - uses: actions/checkout@v4 27 | - uses: ruby/setup-ruby@v1 28 | with: 29 | ruby-version: ${{ matrix.ruby }} 30 | bundler-cache: true 31 | - run: bundle exec rake 32 | -------------------------------------------------------------------------------- /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Publish to RubyGems.org 2 | 3 | on: 4 | push: 5 | branches: main 6 | paths: lib/prop/version.rb 7 | workflow_dispatch: 8 | 9 | jobs: 10 | publish: 11 | runs-on: ubuntu-latest 12 | environment: rubygems-publish 13 | if: github.repository_owner == 'zendesk' 14 | permissions: 15 | id-token: write 16 | contents: write 17 | steps: 18 | - uses: actions/checkout@v4 19 | - name: Set up Ruby 20 | uses: ruby/setup-ruby@v1 21 | with: 22 | bundler-cache: false 23 | ruby-version: "3.4" 24 | - name: Install dependencies 25 | run: bundle install 26 | - uses: rubygems/release-gem@v1 27 | -------------------------------------------------------------------------------- /.github/workflows/rails_main_testing.yml: -------------------------------------------------------------------------------- 1 | name: Test against Rails main 2 | 3 | on: 4 | schedule: 5 | - cron: "0 0 * * *" # Run every day at 00:00 UTC 6 | workflow_dispatch: 7 | 8 | jobs: 9 | main: 10 | name: Tests 11 | runs-on: ubuntu-latest 12 | strategy: 13 | fail-fast: false 14 | matrix: 15 | ruby: 16 | - '3.4' 17 | env: 18 | BUNDLE_GEMFILE: gemfiles/rails_main.gemfile 19 | steps: 20 | - uses: actions/checkout@v4 21 | - uses: ruby/setup-ruby@v1 22 | with: 23 | ruby-version: ${{ matrix.ruby }} 24 | bundler-cache: true 25 | - run: bundle exec rake 26 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | ## MAC OS 2 | .DS_Store 3 | 4 | ## TEXTMATE 5 | *.tmproj 6 | tmtags 7 | 8 | ## EMACS 9 | *~ 10 | \#* 11 | .\#* 12 | 13 | ## VIM 14 | *.swp 15 | 16 | ## PROJECT::GENERAL 17 | coverage 18 | rdoc 19 | pkg 20 | vendor 21 | 22 | gemfiles/rails_main.gemfile.lock 23 | -------------------------------------------------------------------------------- /Changelog.md: -------------------------------------------------------------------------------- 1 | # Changes 2 | 3 | ## Unreleased 4 | 5 | * Drop support for Ruby below 3.1 6 | * Add tests with Active Support 7.2 and Rails main 7 | 8 | ## 2.9.0 9 | 10 | * Drop Ruby < 2.7 11 | * Test with Ruby 3.2 and 3.3 12 | * Run tests with both Active Support 7.0 and 7.1 13 | * Add after_evaluated callback support 14 | 15 | ## 2.8.0 16 | 17 | * Specify raw when reading raw cache entries [PR](https://github.com/zendesk/prop/pull/45) 18 | 19 | ## 2.7.0 20 | 21 | * Feature: Add threshold to Prop::RateLimited exception 22 | 23 | ## 2.6.1 24 | 25 | * Bugfix: Set expires_in on increment and decrement 26 | 27 | ## 2.6.0 28 | 29 | * Use interval value as the ttl when writing to cache 30 | 31 | ## 2.5.0 32 | 33 | * Bugfix: Fix leaky bucket implementation 34 | 35 | ## 2.4.0 36 | 37 | * Allow zero case for threshold when configure the prop 38 | * See [PR description](https://github.com/zendesk/prop/pull/37) 39 | 40 | ## 2.3.0 41 | 42 | * Bugfix: Fix concurrency bug 43 | * See [PR description](https://github.com/zendesk/prop/pull/33) 44 | 45 | ## 2.2.5 46 | 47 | * Add a reader method for `cache` to top level `Prop` module 48 | * Added compatibility with Rails 5.2 49 | 50 | ## 2.2.4 51 | 52 | * Added compatibility with Rails 5.1 53 | 54 | ## 2.2.3 55 | 56 | * Remove Fixnum and replace with Integer per 2.4.1 Deprecations 57 | * Supported Rubies: 2.4.1, 2.3.4, 2.2.7 58 | 59 | ## 2.2.2 60 | 61 | * Bugfix: Fix underflow error in decrement method 62 | * See: [PR Description](https://github.com/zendesk/prop/pull/26) 63 | 64 | ## 2.2.1 65 | 66 | * Support decrement method for LeakyBucketStrategy 67 | * Support multiple rails versions (3.2, 4.1, 4.2, 5.0) 68 | 69 | ## 2.2.0 (also 2.1.3) 70 | 71 | * Support decrement method for IntervalStrategy 72 | 73 | Decrement can be used to for example throttle before an expensive action and then give quota back when some condition is met. 74 | `:decrement` is only supported for `IntervalStrategy` for now 75 | 76 | In case of api failures we want to decrement the rate limit: 77 | 78 | `Prop.throttle!(:api_counts, request.remote_ip, decrement: 1)` 79 | 80 | ## 2.1.2 81 | 82 | * Freeze string literals 83 | -------------------------------------------------------------------------------- /Gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile "gemfiles/rails7.0.gemfile" 2 | -------------------------------------------------------------------------------- /Gemfile.lock: -------------------------------------------------------------------------------- 1 | PATH 2 | remote: . 3 | specs: 4 | prop (2.9.0) 5 | 6 | GEM 7 | remote: https://rubygems.org/ 8 | specs: 9 | activesupport (7.0.8) 10 | concurrent-ruby (~> 1.0, >= 1.0.2) 11 | i18n (>= 1.6, < 2) 12 | minitest (>= 5.1) 13 | tzinfo (~> 2.0) 14 | base64 (0.2.0) 15 | bigdecimal (3.1.9) 16 | bigdecimal (3.1.9-java) 17 | concurrent-ruby (1.2.3) 18 | i18n (1.14.1) 19 | concurrent-ruby (~> 1.0) 20 | maxitest (5.8.0) 21 | minitest (>= 5.14.0, < 5.26.0) 22 | minitest (5.25.4) 23 | mocha (2.1.0) 24 | ruby2_keywords (>= 0.0.5) 25 | mutex_m (0.3.0) 26 | rake (13.1.0) 27 | ruby2_keywords (0.0.5) 28 | tzinfo (2.0.6) 29 | concurrent-ruby (~> 1.0) 30 | 31 | PLATFORMS 32 | java 33 | ruby 34 | 35 | DEPENDENCIES 36 | activesupport (~> 7.0.0) 37 | base64 38 | bigdecimal 39 | maxitest 40 | mocha 41 | mutex_m 42 | prop! 43 | rake 44 | 45 | BUNDLED WITH 46 | 2.6.2 47 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Prop ![Build status](https://github.com/zendesk/prop/workflows/ci/badge.svg) 2 | 3 | A gem to rate limit requests/actions of any kind.
4 | Define thresholds, register usage and finally act on exceptions once thresholds get exceeded. 5 | 6 | Prop supports two limiting strategies: 7 | 8 | * Basic strategy (default): Prop will use an interval to define a window of time using simple div arithmetic. 9 | This means that it's a worst-case throttle that will allow up to two times the specified requests within the specified interval. 10 | * Leaky bucket strategy: Prop also supports the [Leaky Bucket](https://en.wikipedia.org/wiki/Leaky_bucket) algorithm, 11 | which is similar to the basic strategy but also supports bursts up to a specified threshold. 12 | 13 | To store values, prop needs a cache: 14 | 15 | ```ruby 16 | # config/initializers/prop.rb 17 | Prop.cache = Rails.cache # needs read/write/increment methods 18 | ``` 19 | 20 | When using the interval strategy, prop sets a key expiry to its interval. Because the leaky bucket strategy does not set a ttl, it is best to use memcached or similar for all prop caching, not redis. 21 | 22 | ## Setting a Callback 23 | 24 | You can define an optional callback that is invoked when a rate limit is reached. In a Rails application you 25 | could use such a handler to add notification support: 26 | 27 | ```ruby 28 | Prop.before_throttle do |handle, key, threshold, interval| 29 | ActiveSupport::Notifications.instrument('throttle.prop', handle: handle, key: key, threshold: threshold, interval: interval) 30 | end 31 | ``` 32 | 33 | ## Setting an After Evaluated Callback 34 | 35 | You can define an optional callback that is invoked when a rate limit is checked. The callback will be invoked regardless 36 | of the result of the evaluation. 37 | 38 | ```ruby 39 | Prop.after_evaluated do |handle, counter, options| 40 | Rails.logger.info "Prop #{handle} has just been check. current value: #{counter}" 41 | end 42 | ``` 43 | 44 | ## Defining thresholds 45 | 46 | Example: Limit on accepted emails per hour from a given user, by defining a threshold and interval (in seconds): 47 | 48 | ```ruby 49 | Prop.configure(:mails_per_hour, threshold: 100, interval: 1.hour, description: "Mail rate limit exceeded") 50 | 51 | # Block requests by setting threshold to 0 52 | Prop.configure(:mails_per_hour, threshold: 0, interval: 1.hour, description: "All mail is blocked") 53 | ``` 54 | 55 | ```ruby 56 | # Throws Prop::RateLimited if the threshold/interval has been reached 57 | Prop.throttle!(:mails_per_hour) 58 | 59 | # Prop can be used to guard a block of code 60 | Prop.throttle!(:expensive_request) { calculator.something_very_hard } 61 | 62 | # Returns true if the threshold/interval has been reached 63 | Prop.throttled?(:mails_per_hour) 64 | 65 | # Sets the throttle count to 0 66 | Prop.reset(:mails_per_hour) 67 | 68 | # Returns the value of this throttle, usually a count, but see below for more 69 | Prop.count(:mails_per_hour) 70 | ``` 71 | 72 | Prop will raise a `KeyError` if you attempt to operate on an undefined handle. 73 | 74 | ## Scoping a throttle 75 | 76 | Example: scope the throttling to a specific sender rather than running a global "mails per hour" throttle: 77 | 78 | ```ruby 79 | Prop.throttle!(:mails_per_hour, mail.from) 80 | Prop.throttled?(:mails_per_hour, mail.from) 81 | Prop.reset(:mails_per_hour, mail.from) 82 | Prop.query(:mails_per_hour, mail.from) 83 | ``` 84 | 85 | The throttle scope can also be an array of values: 86 | 87 | ```ruby 88 | Prop.throttle!(:mails_per_hour, [ account.id, mail.from ]) 89 | ``` 90 | 91 | ## Error handling 92 | 93 | If the threshold for a given handle and key combination is exceeded, Prop throws a `Prop::RateLimited`. 94 | This exception contains a "handle" reference and a "description" if specified during the configuration. 95 | The handle allows you to rescue `Prop::RateLimited` and differentiate action depending on the handle. 96 | For example, in Rails you can use this in e.g. `ApplicationController`: 97 | 98 | ```ruby 99 | rescue_from Prop::RateLimited do |e| 100 | if e.handle == :authorization_attempt 101 | render status: :forbidden, message: I18n.t(e.description) 102 | elsif ... 103 | 104 | end 105 | end 106 | ``` 107 | 108 | ### Using the Middleware 109 | 110 | Prop ships with a built-in Rack middleware that you can use to do all the exception handling. 111 | When a `Prop::RateLimited` error is caught, it will build an HTTP 112 | [429 Too Many Requests](http://tools.ietf.org/html/draft-nottingham-http-new-status-02#section-4) 113 | response and set the following headers: 114 | 115 | Retry-After: 32 116 | Content-Type: text/plain 117 | Content-Length: 72 118 | 119 | Where `Retry-After` is the number of seconds the client has to wait before retrying this end point. 120 | The body of this response is whatever description Prop has configured for the throttle that got violated, 121 | or a default string if there's none configured. 122 | 123 | If you wish to do manual error messaging in these cases, you can define an error handler in your Prop configuration. 124 | Here's how the default error handler looks - you use anything that responds to `.call` and 125 | takes the environment and a `RateLimited` instance as argument: 126 | 127 | ```ruby 128 | error_handler = Proc.new do |env, error| 129 | body = error.description || "This action has been rate limited" 130 | headers = { "Content-Type" => "text/plain", "Content-Length" => body.size, "Retry-After" => error.retry_after } 131 | 132 | [ 429, headers, [ body ]] 133 | end 134 | 135 | ActionController::Dispatcher.middleware.insert_before(ActionController::ParamsParser, error_handler: error_handler) 136 | ``` 137 | 138 | An alternative to this, is to extend `Prop::Middleware` and override the `render_response(env, error)` method. 139 | 140 | ## Disabling Prop 141 | 142 | In case you need to perform e.g. a manual bulk operation: 143 | 144 | ```ruby 145 | Prop.disabled do 146 | # No throttles will be tested here 147 | end 148 | ``` 149 | 150 | ## Overriding threshold 151 | 152 | You can chose to override the threshold for a given key: 153 | 154 | ```ruby 155 | Prop.throttle!(:mails_per_hour, mail.from, threshold: current_account.mail_throttle_threshold) 156 | ``` 157 | 158 | When `throttle` is invoked without argument, the key is nil and as such a scope of its own, i.e. these are equivalent: 159 | 160 | ```ruby 161 | Prop.throttle!(:mails_per_hour) 162 | Prop.throttle!(:mails_per_hour, nil) 163 | ``` 164 | 165 | The default (and smallest possible) increment is 1, you can set that to any integer value using 166 | `:increment` which is handy for building time based throttles: 167 | 168 | ```ruby 169 | Prop.configure(:execute_time, threshold: 10, interval: 1.minute) 170 | Prop.throttle!(:execute_time, account.id, increment: (Benchmark.realtime { execute }).to_i) 171 | ``` 172 | 173 | Decrement can be used to for example throttle before an expensive action and then give quota back when some condition is met. 174 | 175 | ```ruby 176 | Prop.throttle!(:api_counts, request.remote_ip, decrement: 1) 177 | ``` 178 | 179 | ## Optional configuration 180 | 181 | You can add optional configuration to a prop and retrieve it using `Prop.configurations[:foo]`: 182 | 183 | ```ruby 184 | Prop.configure(:api_query, threshold: 10, interval: 1.minute, category: :api) 185 | Prop.configure(:api_insert, threshold: 50, interval: 1.minute, category: :api) 186 | Prop.configure(:password_failure, threshold: 5, interval: 1.minute, category: :auth) 187 | ``` 188 | 189 | ``` 190 | Prop.configurations[:api_query][:category] 191 | ``` 192 | 193 | You can use `Prop::RateLimited#config` to distinguish between errors: 194 | 195 | ```ruby 196 | rescue Prop::RateLimited => e 197 | case e.config[:category] 198 | when :api 199 | raise APIRateLimit 200 | when :auth 201 | raise AuthFailure 202 | ... 203 | end 204 | ``` 205 | 206 | ## First throttled 207 | 208 | You can opt to be notified when the throttle is breached for the first time.
209 | This can be used to send notifications on breaches but prevent spam on multiple throttle breaches. 210 | 211 | ```Ruby 212 | Prop.configure(:mails_per_hour, threshold: 100, interval: 1.hour, first_throttled: true) 213 | 214 | throttled = Prop.throttle(:mails_per_hour, user.id, increment: 60) 215 | if throttled 216 | if throttled == :first_throttled 217 | ApplicationMailer.spammer_warning(user).deliver_now 218 | end 219 | Rails.logger.warn("Not sending emails") 220 | else 221 | send_emails 222 | end 223 | 224 | # return values of throttle are: false, :first_throttled, true 225 | 226 | Prop.first_throttled(:mails_per_hour, 1, increment: 60) # -> false 227 | Prop.first_throttled(:mails_per_hour, 1, increment: 60) # -> :first_throttled 228 | Prop.first_throttled(:mails_per_hour, 1, increment: 60) # -> true 229 | 230 | # can also be accesses on `Prop::RateLimited` exceptions as `.first_throttled` 231 | ``` 232 | 233 | ## Using Leaky Bucket Algorithm 234 | 235 | You can add two additional configurations: `:strategy` and `:burst_rate` to use the 236 | [leaky bucket algorithm](https://en.wikipedia.org/wiki/Leaky_bucket). 237 | Prop will handle the details after configured, and you don't have to specify `:strategy` 238 | again when using `throttle`, `throttle!` or any other methods. 239 | 240 | The leaky bucket algorithm used is "leaky bucket as a meter". 241 | 242 | ```ruby 243 | Prop.configure(:api_request, strategy: :leaky_bucket, burst_rate: 20, threshold: 5, interval: 1.minute) 244 | ``` 245 | 246 | * `:threshold` value here would be the "leak rate" of leaky bucket algorithm. 247 | 248 | ### Releasing a new version 249 | A new version is published to RubyGems.org every time a change to `version.rb` is pushed to the `main` branch. 250 | In short, follow these steps: 251 | 1. Update `version.rb`, 252 | 2. update version in all `Gemfile.lock` files, 253 | 3. merge this change into `main`, and 254 | 4. look at [the action](https://github.com/zendesk/prop/actions/workflows/publish.yml) for output. 255 | 256 | To create a pre-release from a non-main branch: 257 | 1. change the version in `version.rb` to something like `1.2.0.pre.1` or `2.0.0.beta.2`, 258 | 2. push this change to your branch, 259 | 3. go to [Actions → “Publish to RubyGems.org” on GitHub](https://github.com/zendesk/prop/actions/workflows/publish.yml), 260 | 4. click the “Run workflow” button, 261 | 5. pick your branch from a dropdown. 262 | 263 | ## License 264 | 265 | Copyright 2015 Zendesk 266 | 267 | Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. 268 | You may obtain a copy of the License at 269 | 270 | http://www.apache.org/licenses/LICENSE-2.0 271 | 272 | Unless required by applicable law or agreed to in writing, 273 | software distributed under the License is distributed on an "AS IS" BASIS, 274 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 275 | See the License for the specific language governing permissions and limitations under the License. 276 | -------------------------------------------------------------------------------- /Rakefile: -------------------------------------------------------------------------------- 1 | require 'bundler/setup' 2 | require 'bundler/gem_tasks' 3 | require 'rake/testtask' 4 | 5 | Rake::TestTask.new :default do |test| 6 | test.pattern = 'test/**/test_*.rb' 7 | test.verbose = true 8 | test.warning = false 9 | end 10 | -------------------------------------------------------------------------------- /bin/bundle_all: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | # Usage: ./bundle_all 3 | # This will update the lock files 4 | # in the gemfiles/ folder as well 5 | echo "bundle_all" 6 | for i in gemfiles/*.gemfile Gemfile 7 | do 8 | echo "Bundling for $i" 9 | echo "BUNDLE_GEMFILE=$i bundle $@" 10 | BUNDLE_GEMFILE=$i bundle $@ 11 | done 12 | -------------------------------------------------------------------------------- /gemfiles/common.rb: -------------------------------------------------------------------------------- 1 | source 'https://rubygems.org' 2 | 3 | gemspec path: '../' 4 | 5 | gem 'maxitest' 6 | gem 'mocha' 7 | gem 'rake' 8 | -------------------------------------------------------------------------------- /gemfiles/rails7.0.gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'common.rb' 2 | 3 | gem 'activesupport', '~> 7.0.0' 4 | gem 'base64' 5 | gem 'bigdecimal' 6 | gem 'mutex_m' 7 | -------------------------------------------------------------------------------- /gemfiles/rails7.0.gemfile.lock: -------------------------------------------------------------------------------- 1 | PATH 2 | remote: .. 3 | specs: 4 | prop (2.9.0) 5 | 6 | GEM 7 | remote: https://rubygems.org/ 8 | specs: 9 | activesupport (7.0.8) 10 | concurrent-ruby (~> 1.0, >= 1.0.2) 11 | i18n (>= 1.6, < 2) 12 | minitest (>= 5.1) 13 | tzinfo (~> 2.0) 14 | base64 (0.2.0) 15 | bigdecimal (3.1.9) 16 | bigdecimal (3.1.9-java) 17 | concurrent-ruby (1.2.2) 18 | i18n (1.14.1) 19 | concurrent-ruby (~> 1.0) 20 | maxitest (5.8.0) 21 | minitest (>= 5.14.0, < 5.26.0) 22 | minitest (5.25.4) 23 | mocha (2.1.0) 24 | ruby2_keywords (>= 0.0.5) 25 | mutex_m (0.3.0) 26 | rake (13.1.0) 27 | ruby2_keywords (0.0.5) 28 | tzinfo (2.0.6) 29 | concurrent-ruby (~> 1.0) 30 | 31 | PLATFORMS 32 | java 33 | ruby 34 | 35 | DEPENDENCIES 36 | activesupport (~> 7.0.0) 37 | base64 38 | bigdecimal 39 | maxitest 40 | mocha 41 | mutex_m 42 | prop! 43 | rake 44 | 45 | BUNDLED WITH 46 | 2.6.2 47 | -------------------------------------------------------------------------------- /gemfiles/rails7.1.gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'common.rb' 2 | 3 | gem 'activesupport', '~> 7.1.0' 4 | -------------------------------------------------------------------------------- /gemfiles/rails7.1.gemfile.lock: -------------------------------------------------------------------------------- 1 | PATH 2 | remote: .. 3 | specs: 4 | prop (2.9.0) 5 | 6 | GEM 7 | remote: https://rubygems.org/ 8 | specs: 9 | activesupport (7.1.2) 10 | base64 11 | bigdecimal 12 | concurrent-ruby (~> 1.0, >= 1.0.2) 13 | connection_pool (>= 2.2.5) 14 | drb 15 | i18n (>= 1.6, < 2) 16 | minitest (>= 5.1) 17 | mutex_m 18 | tzinfo (~> 2.0) 19 | base64 (0.2.0) 20 | bigdecimal (3.1.5) 21 | bigdecimal (3.1.5-java) 22 | concurrent-ruby (1.2.2) 23 | connection_pool (2.4.1) 24 | drb (2.2.0) 25 | ruby2_keywords 26 | i18n (1.14.1) 27 | concurrent-ruby (~> 1.0) 28 | maxitest (5.8.0) 29 | minitest (>= 5.14.0, < 5.26.0) 30 | minitest (5.25.4) 31 | mocha (2.1.0) 32 | ruby2_keywords (>= 0.0.5) 33 | mutex_m (0.2.0) 34 | rake (13.1.0) 35 | ruby2_keywords (0.0.5) 36 | tzinfo (2.0.6) 37 | concurrent-ruby (~> 1.0) 38 | 39 | PLATFORMS 40 | java 41 | ruby 42 | 43 | DEPENDENCIES 44 | activesupport (~> 7.1.0) 45 | maxitest 46 | mocha 47 | prop! 48 | rake 49 | 50 | BUNDLED WITH 51 | 2.6.2 52 | -------------------------------------------------------------------------------- /gemfiles/rails7.2.gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'common.rb' 2 | 3 | gem 'activesupport', '~> 7.2.0' 4 | -------------------------------------------------------------------------------- /gemfiles/rails7.2.gemfile.lock: -------------------------------------------------------------------------------- 1 | PATH 2 | remote: .. 3 | specs: 4 | prop (2.9.0) 5 | 6 | GEM 7 | remote: https://rubygems.org/ 8 | specs: 9 | activesupport (7.2.1) 10 | base64 11 | bigdecimal 12 | concurrent-ruby (~> 1.0, >= 1.3.1) 13 | connection_pool (>= 2.2.5) 14 | drb 15 | i18n (>= 1.6, < 2) 16 | logger (>= 1.4.2) 17 | minitest (>= 5.1) 18 | securerandom (>= 0.3) 19 | tzinfo (~> 2.0, >= 2.0.5) 20 | base64 (0.2.0) 21 | bigdecimal (3.1.8) 22 | bigdecimal (3.1.8-java) 23 | concurrent-ruby (1.3.4) 24 | connection_pool (2.4.1) 25 | drb (2.2.1) 26 | i18n (1.14.6) 27 | concurrent-ruby (~> 1.0) 28 | logger (1.6.1) 29 | maxitest (5.8.0) 30 | minitest (>= 5.14.0, < 5.26.0) 31 | minitest (5.25.4) 32 | mocha (2.4.5) 33 | ruby2_keywords (>= 0.0.5) 34 | rake (13.2.1) 35 | ruby2_keywords (0.0.5) 36 | securerandom (0.3.1) 37 | tzinfo (2.0.6) 38 | concurrent-ruby (~> 1.0) 39 | 40 | PLATFORMS 41 | java 42 | ruby 43 | 44 | DEPENDENCIES 45 | activesupport (~> 7.2.0) 46 | maxitest 47 | mocha 48 | prop! 49 | rake 50 | 51 | BUNDLED WITH 52 | 2.6.2 53 | -------------------------------------------------------------------------------- /gemfiles/rails_main.gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'common.rb' 2 | 3 | gem "activesupport", github: "rails/rails" 4 | 5 | -------------------------------------------------------------------------------- /lib/prop.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require "prop/limiter" 3 | require "forwardable" 4 | 5 | module Prop 6 | # Shorthand for accessing Prop::Limiter methods 7 | class << self 8 | extend Forwardable 9 | def_delegators :"Prop::Limiter", :read, :write, :cache, :cache=, :configure, :configurations, :disabled, :before_throttle 10 | def_delegators :"Prop::Limiter", :throttle, :throttle!, :throttled?, :count, :query, :reset, :after_evaluated 11 | end 12 | end 13 | -------------------------------------------------------------------------------- /lib/prop/interval_strategy.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require 'prop/options' 3 | require 'prop/key' 4 | 5 | module Prop 6 | class IntervalStrategy 7 | class << self 8 | def zero_counter 9 | 0 10 | end 11 | 12 | def counter(cache_key, options) 13 | cache.read(cache_key, raw: true).to_i 14 | end 15 | 16 | # options argument is kept for api consistency for all strategies 17 | def increment(cache_key, amount, options = {}) 18 | raise ArgumentError, "Change amount must be a Integer, was #{amount.class}" unless amount.is_a?(Integer) 19 | cache.increment(cache_key, amount, expires_in: options.fetch(:interval, nil)) || (cache.write(cache_key, amount, raw: true, expires_in: options.fetch(:interval, nil)) && amount) # WARNING: potential race condition 20 | end 21 | 22 | def decrement(cache_key, amount, options = {}) 23 | raise ArgumentError, "Change amount must be a Integer, was #{amount.class}" unless amount.is_a?(Integer) 24 | cache.decrement(cache_key, amount, expires_in: options.fetch(:interval, nil)) || (cache.write(cache_key, 0, raw: true, expires_in: options.fetch(:interval, nil)) && 0) # WARNING: potential race condition 25 | end 26 | 27 | def reset(cache_key, options = {}) 28 | cache.write(cache_key, zero_counter, raw: true, expires_in: options.fetch(:interval, nil)) 29 | end 30 | 31 | def compare_threshold?(counter, operator, options) 32 | return false unless counter 33 | counter.send operator, options.fetch(:threshold) 34 | end 35 | 36 | def first_throttled?(counter, options) 37 | (counter - options.fetch(:increment, 1)) <= options.fetch(:threshold) 38 | end 39 | 40 | # Builds the expiring cache key 41 | def build(options) 42 | key = options.fetch(:key) 43 | handle = options.fetch(:handle) 44 | interval = options.fetch(:interval) 45 | 46 | window = (Time.now.to_i / interval) 47 | cache_key = Prop::Key.normalize([ handle, key, window ]) 48 | 49 | "prop/v2/#{Digest::MD5.hexdigest(cache_key)}" 50 | end 51 | 52 | def threshold_reached(options) 53 | threshold = options.fetch(:threshold) 54 | 55 | "#{options[:handle]} threshold of #{threshold} tries per #{options[:interval]}s exceeded for key #{options[:key].inspect}, hash #{options[:cache_key]}" 56 | end 57 | 58 | def validate_options!(options) 59 | validate_threshold(options[:threshold], :threshold) 60 | validate_interval(options[:interval], :interval) 61 | 62 | amount = options[:increment] || options[:decrement] 63 | if amount 64 | raise ArgumentError.new(":increment or :decrement must be zero or a positive Integer") if !amount.is_a?(Integer) || amount < 0 65 | end 66 | end 67 | 68 | private 69 | 70 | def validate_threshold(option, key) 71 | raise ArgumentError.new("#{key.inspect} must be a non-negative Integer") if !option.is_a?(Integer) || option < 0 72 | end 73 | 74 | def validate_interval(option, key) 75 | raise ArgumentError.new("#{key.inspect} must be a positive Integer") if !option.is_a?(Integer) || option <= 0 76 | end 77 | 78 | def cache 79 | Prop::Limiter.cache 80 | end 81 | end 82 | end 83 | end 84 | -------------------------------------------------------------------------------- /lib/prop/key.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require "digest/md5" 3 | 4 | module Prop 5 | class Key 6 | 7 | # Simple key expansion only supports arrays and primitives 8 | def self.normalize(key) 9 | if key.is_a?(Array) 10 | key.flatten.join("/") 11 | else 12 | key.to_s 13 | end 14 | end 15 | 16 | end 17 | end 18 | -------------------------------------------------------------------------------- /lib/prop/leaky_bucket_strategy.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require 'prop/options' 3 | require 'prop/key' 4 | 5 | module Prop 6 | class LeakyBucketStrategy 7 | class << self 8 | def _throttle_leaky_bucket(handle, key, cache_key, options) 9 | (over_limit, bucket) = options.key?(:decrement) ? 10 | decrement(cache_key, options.fetch(:decrement), options) : 11 | increment(cache_key, options.fetch(:increment, 1), options) 12 | 13 | [over_limit, bucket] 14 | end 15 | 16 | def counter(cache_key, options) 17 | cache.read(cache_key) || zero_counter 18 | end 19 | 20 | def leak_amount(bucket, amount, options, now) 21 | leak_rate = (now - bucket.fetch(:last_leak_time, 0)) / options.fetch(:interval).to_f 22 | leak_amount = (leak_rate * options.fetch(:threshold).to_f) 23 | leak_amount.to_i 24 | end 25 | 26 | def update_bucket(current_bucket_size, max_bucket_size, amount) 27 | over_limit = (max_bucket_size-current_bucket_size) < amount 28 | updated_bucket_size = over_limit ? current_bucket_size : current_bucket_size + amount 29 | [over_limit, updated_bucket_size] 30 | end 31 | 32 | # WARNING: race condition 33 | # this increment is not atomic, so it might miss counts when used frequently 34 | def increment(cache_key, amount, options) 35 | bucket = counter(cache_key, options) 36 | now = Time.now.to_i 37 | max_bucket_size = options.fetch(:burst_rate) 38 | current_bucket_size = bucket.fetch(:bucket, 0) 39 | leak_amount = leak_amount(bucket, amount, options, now) 40 | if leak_amount > 0 41 | # maybe TODO, update last_leak_time to reflect the exact time for the current leak amount 42 | # the current strategy will always reflect a little less leakage, probably not an issue though 43 | bucket[:last_leak_time] = now 44 | current_bucket_size = [(current_bucket_size - leak_amount), 0].max 45 | end 46 | 47 | over_limit, updated_bucket_size = update_bucket(current_bucket_size, max_bucket_size, amount) 48 | bucket[:bucket] = updated_bucket_size 49 | bucket[:over_limit] = over_limit 50 | cache.write(cache_key, bucket) 51 | [over_limit, bucket] 52 | end 53 | 54 | def decrement(cache_key, amount, options) 55 | now = Time.now.to_i 56 | bucket = counter(cache_key, options) 57 | leak_amount = leak_amount(bucket, amount, options, now) 58 | bucket[:bucket] = [bucket[:bucket] - amount - leak_amount, 0].max 59 | bucket[:last_leak_time] = now if leak_amount > 0 60 | bucket[:over_limit] = false 61 | cache.write(cache_key, bucket) 62 | [false, bucket] 63 | end 64 | 65 | def reset(cache_key, options = {}) 66 | cache.write(cache_key, zero_counter, raw: true) 67 | end 68 | 69 | def compare_threshold?(bucket, operator, options) 70 | bucket.fetch(:over_limit, false) 71 | end 72 | 73 | def build(options) 74 | key = options.fetch(:key) 75 | handle = options.fetch(:handle) 76 | 77 | cache_key = Prop::Key.normalize([ handle, key ]) 78 | 79 | "prop/leaky_bucket/#{Digest::MD5.hexdigest(cache_key)}" 80 | end 81 | 82 | def threshold_reached(options) 83 | burst_rate = options.fetch(:burst_rate) 84 | threshold = options.fetch(:threshold) 85 | 86 | "#{options[:handle]} threshold of #{threshold} tries per #{options[:interval]}s and burst rate #{burst_rate} tries exceeded for key #{options[:key].inspect}, hash #{options[:cache_key]}" 87 | end 88 | 89 | def validate_options!(options) 90 | Prop::IntervalStrategy.validate_options!(options) 91 | 92 | if !options[:burst_rate].is_a?(Integer) || options[:burst_rate] < options[:threshold] 93 | raise ArgumentError.new(":burst_rate must be an Integer and not less than :threshold") 94 | end 95 | 96 | if options[:first_throttled] 97 | raise ArgumentError.new(":first_throttled is not supported") 98 | end 99 | end 100 | 101 | def zero_counter 102 | { bucket: 0, last_leak_time: 0, over_limit: false } 103 | end 104 | 105 | def cache 106 | Prop::Limiter.cache 107 | end 108 | end 109 | end 110 | end 111 | -------------------------------------------------------------------------------- /lib/prop/limiter.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require 'prop/rate_limited' 3 | require 'prop/key' 4 | require 'prop/options' 5 | require 'prop/interval_strategy' 6 | require 'prop/leaky_bucket_strategy' 7 | 8 | module Prop 9 | class Limiter 10 | 11 | class << self 12 | attr_accessor :handles, :before_throttle_callback, :cache, :after_evaluated_callback 13 | 14 | def read(&blk) 15 | raise "Use .cache = " 16 | end 17 | 18 | def write(&blk) 19 | raise "Use .cache = " 20 | end 21 | 22 | def cache=(cache) 23 | [:read, :write, :increment, :decrement].each do |method| 24 | next if cache.respond_to?(method) 25 | raise ArgumentError, "Cache needs to respond to #{method}" 26 | end 27 | 28 | # https://github.com/petergoldstein/dalli/pull/481 29 | if defined?(ActiveSupport::Cache::DalliStore) && 30 | cache.is_a?(ActiveSupport::Cache::DalliStore) && 31 | Gem::Version.new(Dalli::VERSION) <= Gem::Version.new("2.7.4") 32 | raise "Upgrade to dalli 2.7.5+ to use prop v2, it fixes a local_cache vs increment bug" 33 | end 34 | 35 | @cache = cache 36 | end 37 | 38 | def before_throttle(&blk) 39 | self.before_throttle_callback = blk 40 | end 41 | 42 | def after_evaluated(&blk) 43 | self.after_evaluated_callback = blk 44 | end 45 | 46 | # Public: Registers a handle for rate limiting 47 | # 48 | # handle - the name of the handle you wish to use in your code, e.g. :login_attempt 49 | # defaults - the settings for this handle, e.g. { threshold: 5, interval: 5.minutes } 50 | # 51 | # Raises Prop::RateLimited if the number if the threshold for this handle has been reached 52 | def configure(handle, defaults) 53 | Prop::Options.validate_options!(defaults) 54 | 55 | self.handles ||= {} 56 | self.handles[handle] = defaults 57 | end 58 | 59 | # Public: Disables Prop for a block of code 60 | # 61 | # block - a block of code within which Prop will not raise 62 | def disabled(&_block) 63 | @disabled = true 64 | yield 65 | ensure 66 | @disabled = false 67 | end 68 | 69 | # Public: Records a single action for the given handle/key combination. 70 | # 71 | # handle - the registered handle associated with the action 72 | # key - a custom request specific key, e.g. [ account.id, "download", request.remote_ip ] 73 | # options - request specific overrides to the defaults configured for this handle 74 | # (optional) a block of code that this throttle is guarding 75 | # 76 | # Returns true if the threshold for this handle has been reached, else returns false 77 | def throttle(handle, key = nil, options = {}) 78 | options, cache_key, strategy = prepare(handle, key, options) 79 | throttled = _throttle(strategy, handle, key, cache_key, options).first 80 | block_given? && !throttled ? yield : throttled 81 | end 82 | 83 | # Public: Records a single action for the given handle/key combination. 84 | # 85 | # handle - the registered handle associated with the action 86 | # key - a custom request specific key, e.g. [ account.id, "download", request.remote_ip ] 87 | # options - request specific overrides to the defaults configured for this handle 88 | # (optional) a block of code that this throttle is guarding 89 | # 90 | # Raises Prop::RateLimited if the threshold for this handle has been reached 91 | # Returns the value of the block if given a such, otherwise the current count of the throttle 92 | def throttle!(handle, key = nil, options = {}, &block) 93 | options, cache_key, strategy = prepare(handle, key, options) 94 | throttled, counter = _throttle(strategy, handle, key, cache_key, options) 95 | 96 | if throttled 97 | raise Prop::RateLimited.new(options.merge( 98 | cache_key: cache_key, 99 | handle: handle, 100 | first_throttled: (throttled == :first_throttled) 101 | )) 102 | end 103 | 104 | block_given? ? yield : counter 105 | end 106 | 107 | # Public: Is the given handle/key combination currently throttled ? 108 | # 109 | # handle - the throttle identifier 110 | # key - the associated key 111 | # 112 | # Returns true if a call to `throttle!` with same parameters would raise, otherwise false 113 | def throttled?(handle, key = nil, options = {}) 114 | options, cache_key, strategy = prepare(handle, key, options) 115 | counter = strategy.counter(cache_key, options) 116 | strategy.compare_threshold?(counter, :>=, options) 117 | end 118 | 119 | # Public: Resets a specific throttle 120 | # 121 | # handle - the throttle identifier 122 | # key - the associated key 123 | # 124 | # Returns nothing 125 | def reset(handle, key = nil, options = {}) 126 | _options, cache_key, strategy = prepare(handle, key, options) 127 | strategy.reset(cache_key, options) 128 | end 129 | 130 | # Public: Counts the number of times the given handle/key combination has been hit in the current window 131 | # 132 | # handle - the throttle identifier 133 | # key - the associated key 134 | # 135 | # Returns a count of hits in the current window 136 | def count(handle, key = nil, options = {}) 137 | options, cache_key, strategy = prepare(handle, key, options) 138 | strategy.counter(cache_key, options) 139 | end 140 | alias :query :count 141 | 142 | def handles 143 | @handles ||= {} 144 | end 145 | alias :configurations :handles 146 | 147 | private 148 | 149 | def leaky_bucket_strategy?(strategy) 150 | strategy == Prop::LeakyBucketStrategy 151 | end 152 | 153 | def _throttle(strategy, handle, key, cache_key, options) 154 | return [false, strategy.zero_counter] if disabled? 155 | 156 | if leaky_bucket_strategy?(strategy) 157 | is_over_limit, bucket_info_hash = Prop::LeakyBucketStrategy._throttle_leaky_bucket(handle, key, cache_key, options) 158 | bucket_counter = bucket_info_hash.fetch(:bucket) 159 | after_evaluated_callback.call(handle, bucket_counter, options.merge(bucket_info_hash)) if after_evaluated_callback 160 | return [is_over_limit, bucket_info_hash] 161 | end 162 | 163 | counter = options.key?(:decrement) ? 164 | strategy.decrement(cache_key, options.fetch(:decrement), options) : 165 | strategy.increment(cache_key, options.fetch(:increment, 1), options) 166 | 167 | after_evaluated_callback.call(handle, counter, options) if after_evaluated_callback 168 | 169 | if strategy.compare_threshold?(counter, :>, options) 170 | before_throttle_callback && 171 | before_throttle_callback.call(handle, key, options[:threshold], options[:interval]) 172 | 173 | result = if options[:first_throttled] && strategy.first_throttled?(counter, options) 174 | :first_throttled 175 | else 176 | true 177 | end 178 | 179 | [result, counter] 180 | else 181 | [false, counter] 182 | end 183 | end 184 | 185 | def disabled? 186 | defined?(@disabled) && !!@disabled 187 | end 188 | 189 | def prepare(handle, key, params) 190 | unless defaults = handles[handle] 191 | raise KeyError.new("No such handle configured: #{handle.inspect}") 192 | end 193 | 194 | options = Prop::Options.build(key: key, params: params, defaults: defaults) 195 | 196 | strategy = options.fetch(:strategy) 197 | 198 | cache_key = strategy.build(key: key, handle: handle, interval: options[:interval]) 199 | 200 | [ options, cache_key, strategy ] 201 | end 202 | end 203 | end 204 | end 205 | -------------------------------------------------------------------------------- /lib/prop/middleware.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | module Prop 3 | 4 | # Convenience middleware that conveys the message configured on a Prop handle as well 5 | # as time left before the current window has passed in a Retry-After header. 6 | class Middleware 7 | 8 | # Default error handler 9 | class DefaultErrorHandler 10 | def self.call(env, error) 11 | body = error.description || "This action has been rate limited" 12 | headers = { "Content-Type" => "text/plain", "Content-Length" => "#{body.size}", "Retry-After" => "#{error.retry_after}" } 13 | 14 | [ 429, headers, [ body ]] 15 | end 16 | end 17 | 18 | def initialize(app, options = {}) 19 | @app = app 20 | @options = options 21 | @handler = options[:error_handler] || DefaultErrorHandler 22 | end 23 | 24 | def call(env) 25 | begin 26 | @app.call(env) 27 | rescue Prop::RateLimited => e 28 | render_response(env, e) 29 | end 30 | end 31 | 32 | protected 33 | 34 | def render_response(env, error) 35 | @handler.call(env, error) 36 | end 37 | end 38 | 39 | end 40 | -------------------------------------------------------------------------------- /lib/prop/options.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require 'prop/key' 3 | 4 | module Prop 5 | class Options 6 | 7 | # Sanitizes the option set and sets defaults 8 | def self.build(options) 9 | key = options.fetch(:key) 10 | params = options.fetch(:params) 11 | defaults = options.fetch(:defaults) 12 | result = defaults.merge(params) 13 | 14 | result[:key] = Prop::Key.normalize(key) 15 | result[:strategy] = get_strategy(result) 16 | 17 | result[:strategy].validate_options!(result) 18 | result 19 | end 20 | 21 | def self.validate_options!(options) 22 | get_strategy(options).validate_options!(options) 23 | end 24 | 25 | def self.get_strategy(options) 26 | if leaky_bucket.include?(options[:strategy]) 27 | Prop::LeakyBucketStrategy 28 | elsif options[:strategy] == nil 29 | Prop::IntervalStrategy 30 | else 31 | options[:strategy] # allowing any new/unknown strategy to be used 32 | end 33 | end 34 | 35 | def self.leaky_bucket 36 | [:leaky_bucket, "leaky_bucket"] 37 | end 38 | end 39 | end 40 | -------------------------------------------------------------------------------- /lib/prop/rate_limited.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | module Prop 3 | class RateLimited < StandardError 4 | attr_accessor :handle, :cache_key, :retry_after, :description, :first_throttled, :threshold 5 | 6 | def initialize(options) 7 | self.handle = options.fetch(:handle) 8 | self.cache_key = options.fetch(:cache_key) 9 | self.first_throttled = options.fetch(:first_throttled) 10 | self.description = options[:description] 11 | 12 | interval = options.fetch(:interval).to_i 13 | self.retry_after = interval - Time.now.to_i % interval 14 | 15 | self.threshold = options.fetch(:threshold) 16 | 17 | super(options.fetch(:strategy).threshold_reached(options)) 18 | end 19 | 20 | def config 21 | Prop.configurations.fetch(@handle) 22 | end 23 | end 24 | end 25 | -------------------------------------------------------------------------------- /lib/prop/version.rb: -------------------------------------------------------------------------------- 1 | module Prop 2 | VERSION = "2.9.0" 3 | end 4 | -------------------------------------------------------------------------------- /prop.gemspec: -------------------------------------------------------------------------------- 1 | require_relative "lib/prop/version" 2 | 3 | Gem::Specification.new "prop", Prop::VERSION do |s| 4 | s.license = "Apache License Version 2.0" 5 | 6 | s.summary = "Gem for implementing rate limits." 7 | 8 | s.authors = ["Morten Primdahl"] 9 | s.email = 'primdahl@me.com' 10 | s.homepage = 'https://github.com/zendesk/prop' 11 | 12 | s.required_ruby_version = '>= 3.1' 13 | s.files = `git ls-files lib LICENSE README.md`.split("\n") 14 | end 15 | -------------------------------------------------------------------------------- /test/helper.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require 'bundler/setup' 3 | 4 | require "maxitest/global_must" 5 | require "maxitest/autorun" 6 | require 'mocha/minitest' 7 | 8 | require 'time' 9 | require 'prop' 10 | require 'active_support' 11 | require 'active_support/core_ext/numeric/time' 12 | require 'active_support/cache' 13 | require 'active_support/cache/memory_store' 14 | require 'active_support/notifications' 15 | 16 | Minitest::Test.class_eval do 17 | def setup_fake_store 18 | Prop.cache = ActiveSupport::Cache::MemoryStore.new 19 | end 20 | 21 | def freeze_time(time = Time.now) 22 | @time = time 23 | Time.stubs(:now).returns(@time) 24 | end 25 | end 26 | -------------------------------------------------------------------------------- /test/test_changelog.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | describe Prop::VERSION do 5 | it "should always add a changelog while bumping versions" do 6 | changes = File.read("#{File.dirname(__FILE__)}/../Changelog.md") 7 | expected = "## #{Prop::VERSION}" 8 | assert changes.include?(expected), "#{expected} not found in Changelog.md" 9 | "not found in Changelog.md, Please update the Changelog file" 10 | end 11 | end 12 | -------------------------------------------------------------------------------- /test/test_interval_strategy.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | describe Prop::IntervalStrategy do 5 | before do 6 | @key = "cache_key" 7 | setup_fake_store 8 | freeze_time 9 | end 10 | 11 | describe "#counter" do 12 | describe "when @store[@key] is nil" do 13 | it "returns the current count" do 14 | Prop::IntervalStrategy.counter(@key, nil).must_equal 0 15 | end 16 | end 17 | 18 | describe "when @store[@key] has an existing value" do 19 | before { Prop::Limiter.cache.write(@key, 1) } 20 | 21 | it "returns the current count" do 22 | Prop::IntervalStrategy.counter(@key, nil).must_equal 1 23 | end 24 | end 25 | end 26 | 27 | describe "#increment" do 28 | it "increments an empty bucket" do 29 | Prop::IntervalStrategy.increment(@key, 5) 30 | assert_equal 5, Prop::IntervalStrategy.counter(@key, nil) 31 | end 32 | 33 | it "increments a filled bucket" do 34 | Prop::IntervalStrategy.increment(@key, 5) 35 | Prop::IntervalStrategy.increment(@key, 5) 36 | assert_equal 10, Prop::IntervalStrategy.counter(@key, nil) 37 | end 38 | 39 | it "does not write non-integers" do 40 | assert_raises ArgumentError do 41 | Prop::IntervalStrategy.increment(@key, "WHOOPS") 42 | end 43 | end 44 | end 45 | 46 | describe "#decrement" do 47 | xit "returns 0 when decrements an empty bucket" do 48 | Prop::IntervalStrategy.decrement(@key, -5) 49 | assert_equal 0, Prop::IntervalStrategy.counter(@key, nil) 50 | end 51 | 52 | it "decrements a filled bucket" do 53 | Prop::IntervalStrategy.increment(@key, 5) 54 | Prop::IntervalStrategy.decrement(@key, 2) 55 | assert_equal 3, Prop::IntervalStrategy.counter(@key, nil) 56 | end 57 | 58 | it "does not write non-integers" do 59 | assert_raises ArgumentError do 60 | Prop::IntervalStrategy.decrement(@key, "WHOOPS") 61 | end 62 | end 63 | end 64 | 65 | describe "#reset" do 66 | before { Prop::Limiter.cache.write(@key, 100) } 67 | 68 | it "resets the bucket" do 69 | Prop::IntervalStrategy.reset(@key) 70 | Prop::IntervalStrategy.counter(@key, nil).must_equal 0 71 | end 72 | end 73 | 74 | describe "#compare_threshold?" do 75 | it "returns true when the limit has been reached" do 76 | assert Prop::IntervalStrategy.compare_threshold?(100, :>=, { threshold: 100 }) 77 | assert Prop::IntervalStrategy.compare_threshold?(101, :>, { threshold: 100 }) 78 | end 79 | 80 | it "returns false when the limit has not been reached" do 81 | refute Prop::IntervalStrategy.compare_threshold?(99, :>=, { threshold: 100 }) 82 | refute Prop::IntervalStrategy.compare_threshold?(100, :>, { threshold: 100 }) 83 | end 84 | 85 | it "returns false when the counter fails to increment" do 86 | refute Prop::IntervalStrategy.compare_threshold?(false, :>, { threshold: 100 }) 87 | refute Prop::IntervalStrategy.compare_threshold?(nil, :>, { threshold: 100 }) 88 | end 89 | end 90 | 91 | describe "#build" do 92 | it "returns a hexdigested key" do 93 | Prop::IntervalStrategy.build(handle: :hello, key: [ "foo", 2, :bar ], interval: 60).must_match(/prop\/v2\/[a-f0-9]+/) 94 | end 95 | end 96 | 97 | describe "#validate_options!" do 98 | describe "when :increment is zero" do 99 | it "does not raise exception" do 100 | arg = { threshold: 1, interval: 1, increment: 0} 101 | refute Prop::IntervalStrategy.validate_options!(arg) 102 | end 103 | end 104 | 105 | describe "when :threshold is set to zero to disable the prop" do 106 | it "does not raise exception" do 107 | arg = { threshold: 0, interval: 1, increment: 1} 108 | refute Prop::IntervalStrategy.validate_options!(arg) 109 | end 110 | end 111 | end 112 | end 113 | -------------------------------------------------------------------------------- /test/test_key.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | describe Prop::Key do 5 | describe "#normalize" do 6 | it "turn a Integer into a String" do 7 | Prop::Key.normalize(3).must_equal "3" 8 | end 9 | 10 | it "return a String" do 11 | Prop::Key.normalize("S").must_equal "S" 12 | end 13 | 14 | it "flatten and join an Array" do 15 | Prop::Key.normalize([ 1, "B", "3" ]).must_equal "1/B/3" 16 | end 17 | end 18 | end 19 | 20 | -------------------------------------------------------------------------------- /test/test_leaky_bucket_strategy.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | describe Prop::LeakyBucketStrategy do 5 | before do 6 | @key = "leaky_bucket_cache_key" 7 | setup_fake_store 8 | freeze_time 9 | end 10 | 11 | def increment_test_helper(expected_over_limit, expected_bucket, amount, options) 12 | assert_equal([expected_over_limit, expected_bucket], Prop::LeakyBucketStrategy.increment(@key, amount, options)) 13 | assert_equal(expected_bucket, Prop::Limiter.cache.read(@key)) 14 | end 15 | 16 | def decrement_test_helper(expected_over_limit, expected_bucket, amount, options) 17 | assert_equal([expected_over_limit, expected_bucket], Prop::LeakyBucketStrategy.decrement(@key, amount, options)) 18 | assert_equal(expected_bucket, Prop::Limiter.cache.read(@key)) 19 | end 20 | 21 | describe "#_throttle_leaky_bucket" do 22 | it "increments bucket by default" do 23 | Prop::LeakyBucketStrategy._throttle_leaky_bucket('handle', "foo", @key, interval: 1, threshold: 10, burst_rate: 15) 24 | assert_equal({bucket: 1, last_leak_time: @time.to_i, over_limit: false}, Prop::Limiter.cache.read(@key)) 25 | end 26 | 27 | it "increments bucket when specified" do 28 | Prop::LeakyBucketStrategy._throttle_leaky_bucket('handle', "foo", @key, interval: 1, threshold: 10, burst_rate: 15, increment: 2) 29 | assert_equal({bucket: 2, last_leak_time: @time.to_i, over_limit: false}, Prop::Limiter.cache.read(@key)) 30 | end 31 | 32 | it "decrements bucket" do 33 | Prop::Limiter.cache.write(@key, { bucket: 10, last_leak_time: @time.to_i}) 34 | Prop::LeakyBucketStrategy._throttle_leaky_bucket('handle', "foo", @key, interval: 1, threshold: 10, burst_rate: 15, decrement: 5) 35 | assert_equal({bucket: 5, last_leak_time: @time.to_i, over_limit: false}, Prop::Limiter.cache.read(@key)) 36 | end 37 | end 38 | 39 | describe "#counter" do 40 | describe "when cache[@key] is nil" do 41 | it "returns an empty bucket" do 42 | bucket_expected = { bucket: 0, last_leak_time: 0, over_limit: false} 43 | assert_equal(Prop::LeakyBucketStrategy.counter(@key, interval: 1, threshold: 10), bucket_expected) 44 | end 45 | end 46 | 47 | describe "when @store[@key] has an existing value" do 48 | before do 49 | Prop::Limiter.cache.write(@key, bucket: 100, last_leak_time: @time.to_i) 50 | end 51 | 52 | it "returns the current bucket" do 53 | bucket_expected = { bucket: 100, last_leak_time: @time.to_i } 54 | assert_equal(Prop::LeakyBucketStrategy.counter(@key, interval: 1, threshold: 10), bucket_expected) 55 | end 56 | end 57 | end 58 | 59 | describe "#increment" do 60 | it "increments an empty bucket" do 61 | expected_bucket = {bucket: 1, last_leak_time: @time.to_i, over_limit: false} 62 | increment_test_helper(false, expected_bucket,1, interval: 1, threshold: 10, burst_rate: 15) 63 | end 64 | 65 | describe "when increment amount is 1" do 66 | it "increments an existing bucket above burst rate and sets over_limit to true" do 67 | expected_bucket = {bucket: 15, last_leak_time: @time.to_i, over_limit: true} 68 | Prop::Limiter.cache.write(@key, expected_bucket) 69 | increment_test_helper(true, expected_bucket, 1, interval: 1, threshold: 10, burst_rate: 15) 70 | end 71 | 72 | it "increments an existing bucket to value below burst and over_limit to false" do 73 | expected_bucket = {bucket: 11, last_leak_time: @time.to_i, over_limit: false} 74 | Prop::Limiter.cache.write(@key, { bucket: 10, last_leak_time: @time.to_i}) 75 | increment_test_helper(false, expected_bucket, 1, interval: 1, threshold: 10, burst_rate: 15) 76 | end 77 | 78 | it "leaks bucket at proper rate" do 79 | expected_bucket = {bucket: 6, last_leak_time: @time.to_i, over_limit: false} 80 | Prop::Limiter.cache.write(@key, { bucket: 10, last_leak_time: @time.to_i - 4 }) 81 | increment_test_helper(false, expected_bucket, 0, interval: 60, threshold: 60, burst_rate: 100) 82 | end 83 | end 84 | 85 | describe "when increment amount > 1" do 86 | it "increments an existing bucket to value below burst and over_limit to false" do 87 | expected_bucket = {bucket: 15, last_leak_time: @time.to_i, over_limit: false} 88 | Prop::Limiter.cache.write(@key, { bucket: 10, last_leak_time: @time.to_i}) 89 | increment_test_helper(false, expected_bucket, 5, interval: 1, threshold: 10, burst_rate: 20) 90 | end 91 | 92 | it "sets value to exactly burst_rate and sets over_limit to false" do 93 | expected_bucket = {bucket: 20, last_leak_time: @time.to_i, over_limit: false} 94 | Prop::Limiter.cache.write(@key, { bucket: 9, last_leak_time: @time.to_i }) 95 | increment_test_helper(false, expected_bucket, 11, interval: 60, threshold: 10, burst_rate: 20) 96 | end 97 | 98 | it "leaks bucket at proper rate and update bucket" do 99 | # bucket updated 5 seconds ago with leak rate of 1/second (threshold / interval) 100 | expected_bucket = {bucket: 90, last_leak_time: @time.to_i, over_limit: false} 101 | Prop::Limiter.cache.write(@key, { bucket: 85, last_leak_time: @time.to_i - 5 }) 102 | increment_test_helper(false, expected_bucket, 10, interval: 60, threshold: 60, burst_rate: 100) 103 | end 104 | 105 | it "increment would exceed burst rate and doesn't change bucket value and over_limit to true" do 106 | # bucket updated 5 seconds ago with leak rate of 1/second (threshold / interval) 107 | expected_bucket = {bucket: 85, last_leak_time: @time.to_i, over_limit: true} 108 | Prop::Limiter.cache.write(@key, { bucket: 85, last_leak_time: @time.to_i }) 109 | increment_test_helper(true, expected_bucket, 20, interval: 60, threshold: 60, burst_rate: 100) 110 | end 111 | 112 | it "leaks bucket at proper rate and update bucket" do 113 | expected_bucket = {bucket: 90, last_leak_time: @time.to_i, over_limit: false} 114 | Prop::Limiter.cache.write(@key, { bucket: 85, last_leak_time: @time.to_i - 5 }) 115 | increment_test_helper(false, expected_bucket, 10, interval: 60, threshold: 60, burst_rate: 100) 116 | end 117 | end 118 | 119 | it "does not leak bucket when leak amount is less than one" do 120 | # leak_rate = (now - last_leak_time) / interval 121 | # leak_amount = leak_rate * threshold 122 | # in this case it should be leak_rate = 0.1, 1 second between updates / interval 123 | # leak_amount = 0.9, 0.1 * 9 which truncs to 0 so no leakage 124 | expected_bucket = {bucket: 5, last_leak_time: @time.to_i-2, over_limit: false} 125 | Prop::Limiter.cache.write(@key, { bucket: 5, last_leak_time: @time.to_i - 2 }) 126 | increment_test_helper(false, expected_bucket, 0, interval: 10, threshold: 4, burst_rate: 15) 127 | end 128 | 129 | end 130 | 131 | describe "#decrement" do 132 | it "returns 0 when decrement an empty bucket" do 133 | Prop::LeakyBucketStrategy.decrement(@key, 5, interval: 1, threshold: 10, burst_rate: 15) 134 | assert_equal(Prop::Limiter.cache.read(@key), {bucket: 0, last_leak_time: @time.to_i, over_limit: false}) 135 | end 136 | 137 | it "decrements but does not leak if time doesn't change" do 138 | expected_bucket = {bucket: 2, last_leak_time: @time.to_i, over_limit: false} 139 | Prop::Limiter.cache.write(@key, { bucket: 5, last_leak_time: @time.to_i }) 140 | decrement_test_helper(false, expected_bucket, 3, interval: 1, threshold: 10, burst_rate: 15) 141 | end 142 | 143 | it "decrements and leaks bucket" do 144 | # bucket updated 5 seconds ago with leak rate of 1/second 145 | expected_bucket = {bucket: 4, last_leak_time: @time.to_i, over_limit: false} 146 | Prop::Limiter.cache.write(@key, { bucket: 10, last_leak_time: @time.to_i - 5 }) 147 | decrement_test_helper(false, expected_bucket, 1, interval: 60, threshold: 60, burst_rate: 100) 148 | end 149 | 150 | it "decrements and does not leak bucket if interval is too short" do 151 | # leak_rate = (now - last_leak_time) / interval 152 | # leak_amount = leak_rate * threshold 153 | # in this case it should be leak_rate = 0.1, 1 second between updates / interval 154 | # leak_amount = 0.9, 0.1 * 9 which truncs to 0 so no leakage 155 | expected_bucket = {bucket: 9, last_leak_time: @time.to_i-1, over_limit: false} 156 | Prop::Limiter.cache.write(@key, { bucket: 10, last_leak_time: @time.to_i - 1 }) 157 | decrement_test_helper(false, expected_bucket, 1, interval: 10, threshold: 9, burst_rate: 100) 158 | end 159 | end 160 | 161 | describe "#reset" do 162 | before do 163 | Prop::Limiter.cache.write(@key, bucket: 100, last_leak_time: @time.to_i) 164 | end 165 | 166 | it "resets the bucket" do 167 | bucket_expected = { bucket: 0, last_leak_time: 0, over_limit: false } 168 | Prop::LeakyBucketStrategy.reset(@key) 169 | assert_equal(Prop::Limiter.cache.read(@key), bucket_expected) 170 | end 171 | end 172 | 173 | describe "#build" do 174 | it "returns a hexdigested key" do 175 | _(Prop::LeakyBucketStrategy.build(handle: :hello, key: [ "foo", 2, :bar ])).must_match(/prop\/leaky_bucket\/[a-f0-9]+/) 176 | end 177 | end 178 | 179 | describe "#validate_options!" do 180 | it "raise when burst rate is not valid" do 181 | @args = { threshold: 10, interval: 10, strategy: :leaky_bucket, burst_rate: 'five' } 182 | assert_raises(ArgumentError) { Prop::LeakyBucketStrategy.validate_options!(@args) } 183 | end 184 | 185 | describe "when :increment less than zero" do 186 | it "raises an exception" do 187 | @args = { threshold: 1, interval: 1, strategy: :leaky_bucket, burst_rate: 2, increment: -1} 188 | assert_raises(ArgumentError) { Prop::LeakyBucketStrategy.validate_options!(@args) } 189 | end 190 | end 191 | end 192 | end 193 | -------------------------------------------------------------------------------- /test/test_limiter.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | describe Prop::Limiter do 5 | context "both strategies are being used simultaneously" do 6 | before do 7 | cache = setup_fake_store 8 | def cache.increment(_, _, _) 9 | sleep 0.0001 10 | end 11 | 12 | Prop.configure(:interval, threshold: 1, interval: 1.minute) 13 | Prop.configure(:leaky_bucket, strategy: :leaky_bucket, burst_rate: 2, threshold: 1, interval: 1.minute) 14 | end 15 | 16 | def self.it_handles_concurrency(method) 17 | it "handles concurrency" do 18 | calculated_value = 0 19 | 20 | interval_thread = Thread.new do 21 | Prop.send(method, :interval) { calculated_value += 1 } 22 | end 23 | leaky_bucket_strategy_thread = Thread.new do 24 | Prop.send(method, :leaky_bucket) { calculated_value += 1 } 25 | end 26 | interval_thread.join 27 | leaky_bucket_strategy_thread.join 28 | 29 | calculated_value.must_equal 2 30 | end 31 | end 32 | 33 | describe "#throttle" do 34 | it_handles_concurrency('throttle') 35 | end 36 | 37 | describe "#throttle!" do 38 | it_handles_concurrency('throttle!') 39 | end 40 | end 41 | end 42 | 43 | describe Prop::Limiter do 44 | before do 45 | @cache_key = "cache_key" 46 | setup_fake_store 47 | freeze_time 48 | end 49 | 50 | describe Prop::IntervalStrategy do 51 | before do 52 | Prop::Limiter.configure(:something, threshold: 10, interval: 10) 53 | Prop::IntervalStrategy.stubs(:build).returns(@cache_key) 54 | Prop.reset(:something) 55 | end 56 | 57 | describe "#throttle" do 58 | it "returns false when disabled" do 59 | Prop::Limiter.disabled { Prop.throttle(:something) }.must_equal false 60 | end 61 | 62 | it "supports decrement and can below 0" do 63 | Prop.throttle(:something) 64 | Prop.count(:something).must_equal 1 65 | 66 | Prop.throttle(:something, nil, decrement: 5) 67 | Prop.count(:something).must_equal -4 68 | end 69 | 70 | describe "when there is a after_evaluated callback" do 71 | it "invokes the callback" do 72 | Prop.after_evaluated do |handle, counter, options| 73 | @handle = handle 74 | @counter = counter 75 | @options = options 76 | end 77 | 78 | Prop.throttle(:something) 79 | 80 | @handle.must_equal :something 81 | @counter.must_equal 1 82 | @options.must_equal({ threshold: 10, interval: 10, key: "", strategy: Prop::IntervalStrategy }) 83 | end 84 | end 85 | 86 | describe "and the threshold has been reached" do 87 | before { Prop::IntervalStrategy.stubs(:compare_threshold?).returns(true) } 88 | 89 | it "returns true" do 90 | assert Prop.throttle(:something) 91 | end 92 | it "increments the throttle count" do 93 | Prop.throttle(:something) 94 | Prop.count(:something).must_equal 1 95 | end 96 | 97 | it "does not execute a block" do 98 | test_block_executed = false 99 | Prop.throttle(:something) { test_block_executed = true } 100 | refute test_block_executed 101 | end 102 | 103 | it "invokes before_throttle callback" do 104 | Prop.before_throttle do |handle, key, threshold, interval| 105 | @handle = handle 106 | @key = key 107 | @threshold = threshold 108 | @interval = interval 109 | end 110 | 111 | Prop.throttle(:something, [:extra]) 112 | 113 | @handle.must_equal :something 114 | @key.must_equal [:extra] 115 | @threshold.must_equal 10 116 | @interval.must_equal 10 117 | end 118 | end 119 | 120 | describe "and the threshold has not been reached" do 121 | before { Prop::IntervalStrategy.stubs(:compare_threshold?).returns(false) } 122 | 123 | it "returns false" do 124 | refute Prop.throttle(:something) 125 | end 126 | 127 | it "increments the throttle count by one" do 128 | Prop.throttle(:something) 129 | 130 | Prop.count(:something).must_equal 1 131 | end 132 | 133 | it "increments the throttle count by the specified number when provided" do 134 | Prop.throttle(:something, nil, increment: 5) 135 | 136 | Prop.count(:something).must_equal 5 137 | end 138 | 139 | it "executes a block" do 140 | calls = [] 141 | Prop.throttle(:something) { calls << true }.must_equal [true] 142 | end 143 | end 144 | end 145 | 146 | describe "#throttle!" do 147 | it "throttles the given handle/key combination" do 148 | Prop::Limiter.expects(:_throttle).with( 149 | Prop::IntervalStrategy, 150 | :something, 151 | :key, 152 | 'cache_key', 153 | { 154 | threshold: 10, 155 | interval: 10, 156 | key: 'key', 157 | strategy: Prop::IntervalStrategy, 158 | options: true 159 | } 160 | ) 161 | 162 | Prop.throttle!(:something, :key, options: true) 163 | end 164 | 165 | describe "when the threshold has been reached" do 166 | before { Prop::Limiter.stubs(:_throttle).returns([true]) } 167 | 168 | it "raises a rate-limited exception" do 169 | assert_raises(Prop::RateLimited) { Prop.throttle!(:something) } 170 | end 171 | 172 | it "does not executes a block" do 173 | test_block_executed = false 174 | assert_raises Prop::RateLimited do 175 | Prop.throttle!(:something) { test_block_executed = true } 176 | end 177 | refute test_block_executed 178 | end 179 | end 180 | 181 | describe "when the threshold has not been reached" do 182 | it "returns the counter value" do 183 | Prop.throttle!(:something).must_equal Prop.count(:something) 184 | end 185 | 186 | it "returns the return value of a block" do 187 | calls = [] 188 | Prop.throttle!(:something) { calls << 'block_value' }.must_equal ['block_value'] 189 | end 190 | end 191 | end 192 | end 193 | 194 | describe Prop::LeakyBucketStrategy do 195 | before do 196 | Prop::Limiter.configure(:something, threshold: 10, interval: 60, burst_rate: 100, strategy: :leaky_bucket) 197 | Prop::LeakyBucketStrategy.stubs(:build).returns(@cache_key) 198 | end 199 | 200 | describe "#throttle" do 201 | describe "when the bucket is not full" do 202 | it "increments the count number and saves timestamp in the bucket" do 203 | refute Prop::Limiter.throttle(:something) 204 | Prop::Limiter.count(:something).must_equal( 205 | bucket: 1, last_leak_time: @time.to_i, over_limit: false 206 | ) 207 | end 208 | end 209 | 210 | describe "when there is a after_evaluated callback" do 211 | it "invokes the callback" do 212 | Prop.after_evaluated do |handle, counter, options| 213 | @handle = handle 214 | @counter = counter 215 | @options = options 216 | end 217 | 218 | Prop.throttle(:something) 219 | 220 | @handle.must_equal :something 221 | @counter.must_equal 1 222 | @options.keys.sort.must_equal [:bucket, :burst_rate, :interval, :key, :last_leak_time, :over_limit, :strategy, :threshold] 223 | end 224 | end 225 | end 226 | 227 | describe "#throttle!" do 228 | describe "when the bucket is full" do 229 | it "raises" do 230 | Prop::Limiter.expects(:_throttle).returns([true, 1]) 231 | assert_raises Prop::RateLimited do 232 | Prop::Limiter.throttle!(:something) 233 | end 234 | end 235 | end 236 | 237 | describe "when the bucket is not full" do 238 | it "returns the bucket" do 239 | expected_bucket = { bucket: 1, last_leak_time: @time.to_i, over_limit: false } 240 | Prop::Limiter.throttle!(:something).must_equal expected_bucket 241 | end 242 | 243 | it "throttles the given handle/key combination" do 244 | Prop::Limiter.expects(:_throttle).with( 245 | Prop::LeakyBucketStrategy, 246 | :something, 247 | :key, 248 | 'cache_key', 249 | { 250 | threshold: 10, 251 | interval: 60, 252 | key: 'key', 253 | burst_rate: 100, 254 | strategy: Prop::LeakyBucketStrategy, 255 | options: true 256 | } 257 | ) 258 | 259 | Prop::Limiter.throttle!(:something, :key, options: true) 260 | end 261 | end 262 | end 263 | end 264 | end 265 | -------------------------------------------------------------------------------- /test/test_middleware.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | require 'prop/middleware' 5 | require 'prop/rate_limited' 6 | 7 | describe Prop::Middleware do 8 | before do 9 | @app = stub() 10 | @env = {} 11 | @middleware = Prop::Middleware.new(@app) 12 | end 13 | 14 | it "return the response" do 15 | @app.expects(:call).with(@env).returns("response") 16 | @middleware.call(@env).must_equal "response" 17 | end 18 | 19 | describe "when throttled" do 20 | before do 21 | options = { 22 | handle: "foo", 23 | threshold: 10, 24 | interval: 60, 25 | cache_key: "wibble", 26 | description: "Boom!", 27 | first_throttled: false, 28 | strategy: Prop::IntervalStrategy 29 | } 30 | @app.expects(:call).with(@env).raises(Prop::RateLimited.new(options)) 31 | end 32 | 33 | it "return the rate limited message when throttled" do 34 | status, _, body = @middleware.call(@env) 35 | 36 | status.must_equal 429 37 | body.must_equal ["Boom!"] 38 | end 39 | 40 | it "allow setting a custom error handler" do 41 | @middleware = Prop::Middleware.new(@app, error_handler: Proc.new { |env, error| "Oops" }) 42 | @middleware.call(@env).must_equal "Oops" 43 | end 44 | end 45 | end 46 | -------------------------------------------------------------------------------- /test/test_options.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | describe Prop::Options do 5 | describe "#build" do 6 | before do 7 | @args = { key: "hello", params: { foo: "bif" }, defaults: { foo: "bar", baz: "moo", threshold: 10, interval: 5 }} 8 | end 9 | 10 | describe "when given valid input" do 11 | before do 12 | @options = Prop::Options.build(@args) 13 | end 14 | 15 | it "support defaults" do 16 | @options[:baz].must_equal "moo" 17 | end 18 | 19 | it "override defaults" do 20 | @options[:foo].must_equal "bif" 21 | end 22 | end 23 | 24 | describe "when given invalid input" do 25 | it "raise when not given an interval" do 26 | @args[:defaults].delete(:interval) 27 | assert_raises(ArgumentError) { Prop::Options.build(@args) } 28 | end 29 | 30 | it "raise when not given a threshold" do 31 | @args[:defaults].delete(:threshold) 32 | assert_raises(ArgumentError) { Prop::Options.build(@args) } 33 | end 34 | 35 | it "raise when not given a key" do 36 | @args.delete(:key) 37 | assert_raises KeyError do 38 | Prop::Options.build(@args) 39 | end 40 | end 41 | 42 | it "raise when increment is not an positive Integer" do 43 | @args[:defaults].merge!(increment: "one") 44 | assert_raises(ArgumentError) { Prop::Options.build(@args) } 45 | end 46 | end 47 | end 48 | end 49 | -------------------------------------------------------------------------------- /test/test_prop.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | # Integration level tests 5 | describe Prop do 6 | def self.with_each_strategy 7 | [{}, {strategy: :leaky_bucket, burst_rate: 2}].each do |options| 8 | describe "with #{options[:strategy] || :interval} strategy" do 9 | yield options 10 | end 11 | end 12 | end 13 | 14 | def count(counter) 15 | counter.is_a?(Hash) ? counter[:bucket] : counter 16 | end 17 | 18 | before do 19 | setup_fake_store 20 | freeze_time 21 | end 22 | 23 | describe "#defaults" do 24 | it "raise errors on invalid configuation" do 25 | assert_raises(ArgumentError) do 26 | Prop.configure :hello_there, threshold: 20, interval: 'hello' 27 | end 28 | 29 | assert_raises(ArgumentError) do 30 | Prop.configure :hello_there, threshold: 'wibble', interval: 100 31 | end 32 | end 33 | 34 | it "result in a default handle" do 35 | Prop.configure :hello_there, threshold: 4, interval: 10 36 | 4.times do |i| 37 | Prop.throttle!(:hello_there, 'some key').must_equal i + 1 38 | end 39 | 40 | assert_raises(Prop::RateLimited) { Prop.throttle!(:hello_there, 'some key') } 41 | Prop.throttle!(:hello_there, 'some key', threshold: 20).must_equal 6 42 | end 43 | 44 | it "create a handle accepts various cache key types" do 45 | Prop.configure :hello_there, threshold: 4, interval: 10 46 | Prop.throttle!(:hello_there, 5).must_equal 1 47 | Prop.throttle!(:hello_there, 5).must_equal 2 48 | Prop.throttle!(:hello_there, '6').must_equal 1 49 | Prop.throttle!(:hello_there, '6').must_equal 2 50 | Prop.throttle!(:hello_there, [ 5, '6' ]).must_equal 1 51 | Prop.throttle!(:hello_there, [ 5, '6' ]).must_equal 2 52 | end 53 | 54 | it "allow to create 0 limit threshold" do 55 | Prop.configure :hello_there, threshold: 0, interval: 10 56 | assert_raises(Prop::RateLimited) do 57 | Prop.throttle! :hello_there, 0 58 | end 59 | end 60 | end 61 | 62 | describe "#disable" do 63 | with_each_strategy do |options| 64 | it "does not increase the throttle" do 65 | Prop.configure :hello, options.merge(threshold: 2, interval: 10) 66 | count(Prop.throttle!(:hello)).must_equal 1 67 | Prop.disabled do 68 | count(Prop.throttle!(:hello)).must_equal 0 69 | end 70 | count(Prop.throttle!(:hello)).must_equal 2 71 | end 72 | 73 | it "does not increase the throttle with threshold 0" do 74 | Prop.configure :hello, options.merge(threshold: 0, interval: 10) 75 | Prop.disabled do 76 | count(Prop.throttle!(:hello)).must_equal 0 77 | end 78 | end 79 | end 80 | end 81 | 82 | describe "#reset" do 83 | describe "when use interval strategy" do 84 | before do 85 | Prop.configure :hello, threshold: 10, interval: 10 86 | 87 | 5.times do |i| 88 | Prop.throttle!(:hello).must_equal i + 1 89 | end 90 | end 91 | 92 | it "set the correct counter to 0" do 93 | Prop.throttle!(:hello, 'wibble') 94 | Prop.throttle!(:hello, 'wibble') 95 | 96 | Prop.reset(:hello) 97 | Prop.throttle!(:hello).must_equal 1 98 | 99 | Prop.throttle!(:hello, 'wibble').must_equal 3 100 | Prop.reset(:hello, 'wibble') 101 | Prop.throttle!(:hello, 'wibble').must_equal 1 102 | end 103 | end 104 | 105 | describe "when use leaky bucket strategy" do 106 | before do 107 | Prop.configure :leaky, threshold: 2, interval: 10, strategy: :leaky_bucket, burst_rate: 10 108 | 109 | 5.times do |i| 110 | Prop.throttle!(:leaky)[:bucket].must_equal i + 1 111 | end 112 | end 113 | 114 | it "set the correct counter to 0" do 115 | Prop.reset(:leaky) 116 | Prop.throttle!(:leaky)[:bucket].must_equal 1 117 | end 118 | end 119 | end 120 | 121 | describe "#throttled?" do 122 | with_each_strategy do |options| 123 | before do 124 | Prop.configure(:hello, options.merge(threshold: 2, interval: 10.seconds)) 125 | Prop.configure(:world, options.merge(threshold: 2, interval: 10.seconds)) 126 | end 127 | 128 | it "return true once it was throttled" do 129 | 2.times do 130 | refute Prop.throttle(:hello) 131 | refute Prop.throttled?(:hello) 132 | end 133 | 134 | assert Prop.throttle(:hello) 135 | assert Prop.throttled?(:hello) 136 | end 137 | 138 | it "counts different handles separately" do 139 | user_id = 42 140 | 2.times { Prop.throttle(:hello, user_id) } 141 | refute Prop.throttled?(:hello, user_id) 142 | assert Prop.throttle(:hello, user_id) 143 | refute Prop.throttled?(:world, user_id) 144 | end 145 | end 146 | end 147 | 148 | describe "#count" do 149 | before do 150 | Prop.configure(:hello, threshold: 20, interval: 20) 151 | Prop.throttle!(:hello) 152 | Prop.throttle!(:hello) 153 | end 154 | 155 | it "be aliased by #query" do 156 | Prop.query(:hello).must_equal 2 157 | end 158 | 159 | it "return the number of hits on a throttle" do 160 | Prop.count(:hello).must_equal 2 161 | end 162 | end 163 | 164 | describe "#throttle!" do 165 | describe "when use interval strategy" do 166 | it "increment counter correctly" do 167 | Prop.configure(:hello, threshold: 20, interval: 20) 168 | 3.times do |i| 169 | Prop.throttle!(:hello, nil, threshold: 10, interval: 10).must_equal i + 1 170 | end 171 | end 172 | 173 | it "reset counter when time window is passed" do 174 | Prop.configure(:hello, threshold: 20, interval: 20) 175 | 3.times do |i| 176 | Prop.throttle!(:hello, nil, threshold: 10, interval: 10).must_equal i + 1 177 | end 178 | 179 | Time.stubs(:now).returns(@time + 20) 180 | 181 | 3.times do |i| 182 | Prop.throttle!(:hello, nil, threshold: 10, interval: 10).must_equal i + 1 183 | end 184 | end 185 | 186 | it "increment the counter beyond the threshold" do 187 | Prop.configure(:hello, threshold: 5, interval: 1) 188 | 10.times do 189 | Prop.throttle!(:hello) rescue nil 190 | end 191 | 192 | Prop.query(:hello).must_equal 10 193 | end 194 | 195 | it "raise Prop::RateLimited when the threshold is exceeded" do 196 | Prop.configure(:hello, threshold: 5, interval: 10, description: "Boom!") 197 | 198 | 5.times do 199 | Prop.throttle!(:hello, nil) 200 | end 201 | 202 | assert_raises Prop::RateLimited do 203 | Prop.throttle!(:hello, nil) 204 | end 205 | 206 | e = assert_raises Prop::RateLimited do 207 | Prop.throttle!(:hello, nil) 208 | end 209 | 210 | e.handle.must_equal :hello 211 | e.message.must_include "5 tries per 10s exceeded for key" 212 | e.description.must_equal "Boom!" 213 | assert e.retry_after 214 | end 215 | 216 | it "support custom increments" do 217 | Prop.configure(:hello, threshold: 100, interval: 10) 218 | Prop.throttle!(:hello, nil, increment: 48) 219 | Prop.query(:hello).must_equal 48 220 | end 221 | end 222 | 223 | describe "when use leaky bucket strategy" do 224 | before do 225 | Prop.configure(:hello, threshold: 5, interval: 10, strategy: :leaky_bucket, burst_rate: 10, description: "Boom!") 226 | end 227 | 228 | it "increments counter correctly" do 229 | 3.times do |i| 230 | Prop.throttle!(:hello)[:bucket].must_equal i + 1 231 | end 232 | end 233 | 234 | it "leaks when time window is passed" do 235 | 3.times do |i| 236 | Prop.throttle!(:hello)[:bucket].must_equal i + 1 237 | end 238 | 239 | Time.stubs(:now).returns(@time + 10) 240 | 241 | 10.times do |i| 242 | Prop.throttle!(:hello)[:bucket].must_equal i + 1 243 | end 244 | Time.stubs(:now).returns(@time + 30) 245 | Prop.throttle!(:hello) # forces the leak 246 | Prop.query(:hello)[:bucket].must_equal 1 247 | end 248 | 249 | it "does not increment the counter beyond the burst rate" do 250 | 15.times do 251 | Prop.throttle!(:hello) rescue nil 252 | end 253 | 254 | Prop.query(:hello)[:bucket].must_equal 10 255 | end 256 | 257 | it "raises Prop::RateLimited when the bucket is full" do 258 | 10.times do 259 | Prop.throttle!(:hello, nil) 260 | end 261 | 262 | assert_raises Prop::RateLimited do 263 | Prop.throttle!(:hello, nil) 264 | end 265 | 266 | e = assert_raises Prop::RateLimited do 267 | Prop.throttle!(:hello, nil) 268 | end 269 | 270 | e.handle.must_equal :hello 271 | e.message.must_include "5 tries per 10s and burst rate 10 tries exceeded for key" 272 | e.description.must_equal "Boom!" 273 | assert e.retry_after 274 | end 275 | 276 | it "support custom increments" do 277 | Prop.configure(:hello, threshold: 100, interval: 10) 278 | Prop.throttle!(:hello, nil, increment: 48) 279 | Prop.query(:hello).must_equal 48 280 | end 281 | end 282 | 283 | it "raise a RuntimeError when a handle has not been configured" do 284 | assert_raises KeyError do 285 | Prop.throttle!(:no_such_handle, nil, threshold: 5, interval: 10) 286 | end 287 | end 288 | 289 | describe ":first_throttled" do 290 | before { Prop.configure(:hello, threshold: 10, interval: 10, first_throttled: true) } 291 | 292 | it "calls back on first throttling" do 293 | Prop.throttle(:hello, nil, increment: 10).must_equal false 294 | Prop.throttle(:hello, nil).must_equal :first_throttled 295 | Prop.throttle(:hello, nil).must_equal true 296 | end 297 | 298 | it "calls back on first throttling with increment" do 299 | Prop.throttle(:hello, nil, increment: 5).must_equal false 300 | Prop.throttle(:hello, nil, increment: 6).must_equal :first_throttled 301 | Prop.throttle(:hello, nil).must_equal true 302 | end 303 | 304 | it "is set on exceptions" do 305 | Prop.throttle!(:hello, nil, increment: 5) 306 | e = assert_raises Prop::RateLimited do 307 | Prop.throttle!(:hello, nil, increment: 6) 308 | end 309 | e.first_throttled.must_equal true 310 | 311 | e = assert_raises Prop::RateLimited do 312 | Prop.throttle!(:hello, nil, increment: 6) 313 | end 314 | e.first_throttled.must_equal false 315 | end 316 | end 317 | end 318 | 319 | describe "#configurations" do 320 | it "returns the configuration" do 321 | Prop.configure(:something, threshold: 100, interval: 30) 322 | config = Prop.configurations[:something] 323 | config[:threshold].must_equal 100 324 | config[:interval].must_equal 30 325 | end 326 | end 327 | end 328 | -------------------------------------------------------------------------------- /test/test_rate_limited.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | require_relative 'helper' 3 | 4 | describe Prop::RateLimited do 5 | before do 6 | freeze_time 1333685680 7 | 8 | Prop.configure :foo, threshold: 100, interval: 60, category: :api 9 | 10 | @error = Prop::RateLimited.new( 11 | handle: :foo, 12 | threshold: 10, 13 | interval: 60, 14 | cache_key: "wibble", 15 | description: "Boom!", 16 | strategy: Prop::IntervalStrategy, 17 | first_throttled: false 18 | ) 19 | end 20 | 21 | describe "#initialize" do 22 | it "returns an error instance" do 23 | @error.must_be_kind_of StandardError 24 | @error.must_be_kind_of Prop::RateLimited 25 | 26 | @error.handle.must_equal :foo 27 | @error.cache_key.must_equal "wibble" 28 | @error.description.must_equal "Boom!" 29 | @error.message.must_equal "foo threshold of 10 tries per 60s exceeded for key nil, hash wibble" 30 | @error.retry_after.must_equal 20 31 | @error.threshold.must_equal 10 32 | end 33 | end 34 | 35 | describe "#config" do 36 | it "returns the original configuration" do 37 | @error.config[:threshold].must_equal 100 38 | @error.config[:interval].must_equal 60 39 | @error.config[:category].must_equal :api 40 | end 41 | end 42 | end 43 | --------------------------------------------------------------------------------