├── .gitignore ├── .travis.yml ├── LICENSE ├── Makefile ├── README.md ├── conf └── test.config ├── rebar.config ├── rebar.config.script ├── rebar.lock ├── rebar3 ├── src ├── drivers │ ├── throttle_driver.erl │ ├── throttle_ets.erl │ └── throttle_mnesia.erl ├── throttle.app.src ├── throttle.erl ├── throttle_app.erl ├── throttle_sup.erl └── throttle_time.erl └── test ├── profile.erl ├── throttle_distributed_SUITE.erl └── throttle_test_SUITE.erl /.gitignore: -------------------------------------------------------------------------------- 1 | .rebar3 2 | _* 3 | .eunit 4 | *.o 5 | *.beam 6 | *.plt 7 | *.swp 8 | *.swo 9 | .erlang.cookie 10 | ebin 11 | log 12 | erl_crash.dump 13 | .rebar 14 | logs 15 | _build 16 | .idea 17 | rebar3.crashdump 18 | ct_log 19 | .DS_Store 20 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: erlang 2 | otp_release: 3 | - 21.3 4 | - 22.3 5 | - 23.1.2 6 | 7 | 8 | script: 9 | - make travis_test 10 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 LambdaClass 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 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: dev test 2 | 3 | dev: 4 | ./rebar3 compile && ./rebar3 shell 5 | 6 | test: 7 | ./rebar3 ct --name node1@127.0.0.1 ; ./rebar3 dialyzer 8 | 9 | travis_test: 10 | ./rebar3 ct --name node1@127.0.0.1 ; ./rebar3 as test coveralls send ; ./rebar3 dialyzer 11 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # throttle 2 | [![Hex.pm](https://img.shields.io/hexpm/v/lambda_throttle.svg)](https://hex.pm/packages/lambda_throttle) 3 | [![Build Status](https://travis-ci.org/lambdaclass/throttle.svg?branch=master)](https://travis-ci.org/lambdaclass/throttle) 4 | [![Coverage Status](https://coveralls.io/repos/github/lambdaclass/throttle/badge.svg?branch=master)](https://coveralls.io/github/lambdaclass/throttle?branch=master) 5 | 6 | An OTP application to implement throttling/rate limiting of resources. 7 | 8 | ## Rebar3 dependency 9 | 10 | ```erl 11 | {throttle, "0.3.0", {pkg, lambda_throttle}} 12 | ``` 13 | 14 | ## Build 15 | 16 | $ rebar3 compile 17 | 18 | ## Usage 19 | 20 | The application allows to limit different resources (scopes) at different rates. 21 | 22 | * `throttle:setup(Scope, RateLimit, RatePeriod)`: setup a rate limit 23 | for a given `Scope`, allowing at most `RateLimit` requests per 24 | `RatePeriod`. Allowed rate periods are `per_second`, `per_minute`, 25 | `per_hour` and `per_day`. 26 | 27 | Rates can also be set via application environment instead of 28 | calling `setup`: 29 | 30 | ```erlang 31 | {throttle, [{rates, [{my_global_scope, 10, per_second} 32 | {my_expensive_endpoint, 2, per_minute}]}]} 33 | ``` 34 | 35 | * `throttle:check(Scope, Key)`: attempt to request `Scope` with a 36 | given `Key` (e.g. user token, IP). The result will be `{ok, 37 | RemainingAttempts, TimeToReset}` if there are attempts left or 38 | `{limit_exceeded, 0, TimeToReset}` if there aren't. 39 | 40 | * `throttle:peek(Scope, Key)`: returns the same result as `check` 41 | without increasing the requests count. 42 | 43 | ### Distributed support 44 | 45 | By default, throttle keeps the attempt counters on ETS tables, and 46 | therefore those are local to the Erlang node. Mnesia can be used 47 | instead to enfore access limits across all connected nodes, by setting 48 | the `driver` configuration parameter to `throttle_mnesia`: 49 | 50 | ``` erlang 51 | {throttle, [{driver, throttle_mnesia}, 52 | {rates, [{my_global_scope, 10, per_second}]}]} 53 | ``` 54 | 55 | When using the Mnesia driver, `throttle_mnesia:setup()` needs to be 56 | called after the cluster is connected (the tables have to be shared across 57 | nodes, so the nodes must be visible before intialization): 58 | 59 | ``` erlang 60 | (n1@127.0.0.1)1> application:set_env(throttle, driver, throttle_mnesia). 61 | ok 62 | (n1@127.0.0.1)2> application:ensure_all_started(throttle). 63 | {ok,[throttle]} 64 | (n1@127.0.0.1)3> net_kernel:connect('n2@127.0.0.1'). 65 | true 66 | (n1@127.0.0.1)4> throttle_mnesia:setup(). 67 | ok 68 | ``` 69 | 70 | When checking for a Key to access a given Scope, an access counter is 71 | incremented in Mnesia. The 72 | [activity access context](http://learnyousomeerlang.com/mnesia#access-and-context) 73 | for that operation can be configured with the `access_context` 74 | parameter: 75 | 76 | ``` erlang 77 | {throttle, [{driver, throttle_mnesia}, 78 | {access_context, sync_transaction}]}. 79 | ``` 80 | 81 | By default, the `async_dirty` context is used, which prioritizes speed 82 | over consistency when propagating the counter increment. This means 83 | there's a chance of two nodes getting access to a resource when there 84 | is one attempt left. Depending the application, it may make more 85 | sense to choose a different context (like `sync_transaction`) to 86 | reduce the chances of allowing accesses above the limit. 87 | 88 | ## Examples 89 | 90 | ### Shell 91 | ``` erlang 92 | 1> application:ensure_all_started(throttle). 93 | {ok,[throttle]} 94 | 2> throttle:setup(my_api_endpoint, 3, per_minute). 95 | ok 96 | 3> throttle:check(my_api_endpoint, my_token_or_ip). 97 | {ok,2,30362} 98 | 4> throttle:check(my_api_endpoint, my_token_or_ip). 99 | {ok,1,29114} 100 | 5> throttle:check(my_api_endpoint, my_token_or_ip). 101 | {ok,0,27978} 102 | 6> throttle:check(my_api_endpoint, my_token_or_ip). 103 | {limit_exceeded,0,26722} 104 | ``` 105 | 106 | ### Cowboy 2.0 limit by IP 107 | 108 | Middleware module: 109 | 110 | ``` erlang 111 | -module(throttling_middleware). 112 | 113 | -behavior(cowboy_middleware). 114 | 115 | -export([execute/2]). 116 | 117 | execute(Req, Env) -> 118 | {{IP, _}, Req2} = cowboy_req:peer(Req), 119 | 120 | case throttle:check(my_api_rate, IP) of 121 | {limit_exceeded, _, _} -> 122 | lager:warning("IP ~p exceeded api limit", [IP]), 123 | Req3 = cowboy_req:reply(429, Req2), 124 | {stop, Req3}; 125 | _ -> 126 | {ok, Req2, Env} 127 | end. 128 | ``` 129 | 130 | Using it: 131 | 132 | ``` erlang 133 | cowboy:start_clear(my_http_listener, [{port, 8080}], #{ 134 | env => #{dispatch => Dispatch}, 135 | middlewares => [cowboy_router, throttling_middleware, cowboy_handler] 136 | }), 137 | ``` 138 | 139 | ### Cowboy 2.0 limit by Authorization header 140 | 141 | ``` erlang 142 | -module(throttling_middleware). 143 | 144 | -behavior(cowboy_middleware). 145 | 146 | -export([execute/2]). 147 | 148 | execute(Req, Env) -> 149 | Authorization = cowboy_req:header(<<"authorization">>, Req), 150 | 151 | case throttle:check(my_api_rate, Authorization) of 152 | {limit_exceeded, _, _} -> 153 | lager:warning("Auth ~p exceeded api limit", [Authorization]), 154 | Req3 = cowboy_req:reply(429, Req), 155 | {stop, Req2}; 156 | _ -> 157 | {ok, Req, Env} 158 | end. 159 | ``` 160 | 161 | Note that assumes all requests have an authorization header. A more 162 | realistic approach would be to fallback to an IP limit when 163 | Authorization is not present. 164 | 165 | ### Cowboy 1.0 limit by IP 166 | 167 | Middleware module: 168 | 169 | ``` erlang 170 | -module(throttling_middleware). 171 | 172 | -behavior(cowboy_middleware). 173 | 174 | -export([execute/2]). 175 | 176 | execute(Req, Env) -> 177 | {{IP, _}, Req2} = cowboy_req:peer(Req), 178 | 179 | case throttle:check(my_api_rate, IP) of 180 | {limit_exceeded, _, _} -> 181 | lager:warning("IP ~p exceeded api limit", [IP]), 182 | {error, 429, Req2}; 183 | _ -> 184 | {ok, Req2, Env} 185 | end. 186 | ``` 187 | 188 | Using it: 189 | 190 | ``` erlang 191 | cowboy:start_http(my_http_listener, 100, [{port, 8080}], 192 | [{env, [{dispatch, Dispatch}]}, 193 | {middlewares, [cowboy_router, throttling_middleware, cowboy_handler]}] 194 | ), 195 | ``` 196 | 197 | A more detailed example, choosing the rate based on the path, can be found [here](https://github.com/lambdaclass/holiday_ping/blob/26a3d83faaad6977c936a40fe273cd45954d9259/src/throttling_middleware.erl). 198 | -------------------------------------------------------------------------------- /conf/test.config: -------------------------------------------------------------------------------- 1 | [ 2 | {lager, [ 3 | {handlers, [ 4 | {lager_console_backend, [{level, debug}]}, 5 | {lager_file_backend, [{file, "log/error.log"}, {level, error}]}, 6 | {lager_file_backend, [{file, "log/console.log"}, {level, debug}]}]} 7 | ]} 8 | ]. 9 | -------------------------------------------------------------------------------- /rebar.config: -------------------------------------------------------------------------------- 1 | {erl_opts, [debug_info]}. 2 | {deps, []}. 3 | {plugins, [coveralls]}. 4 | 5 | {ct_opts, [{sys_config, "conf/test.config"}]}. 6 | 7 | {profiles, [{test, [{erl_opts, [debug_info, {parse_transform, lager_transform}]}, 8 | {cover_enabled , true}, 9 | {cover_export_enabled , true}, 10 | {coveralls_coverdata , "_build/test/cover/ct.coverdata"}, 11 | {coveralls_service_name , "travis-ci"}, 12 | {deps, [{lager, "3.5.1"}]}]}]}. 13 | -------------------------------------------------------------------------------- /rebar.config.script: -------------------------------------------------------------------------------- 1 | case os:getenv("TRAVIS") of 2 | "true" -> 3 | JobId = os:getenv("TRAVIS_JOB_ID"), 4 | lists:keystore(coveralls_service_job_id, 1, CONFIG, {coveralls_service_job_id, JobId}); 5 | _ -> 6 | CONFIG 7 | end. 8 | -------------------------------------------------------------------------------- /rebar.lock: -------------------------------------------------------------------------------- 1 | []. 2 | -------------------------------------------------------------------------------- /rebar3: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lambdaclass/throttle/4f5fb17c9d4a86ba016e7011648ae5dfe539ac01/rebar3 -------------------------------------------------------------------------------- /src/drivers/throttle_driver.erl: -------------------------------------------------------------------------------- 1 | %% 2 | %% Interface to interact with the different stores that can track scope/key 3 | %% access counts. 4 | %% 5 | -module(throttle_driver). 6 | 7 | -export([setup/0, 8 | initialize/3, 9 | reset/2, 10 | update/2, 11 | lookup/2]). 12 | 13 | %% Performs any global (non scope-specific) setup required by the driver. 14 | -callback setup() -> ok. 15 | 16 | %% Performs scope-specific initialization of a scope. 17 | -callback initialize(throttle:scope(), throttle:rate_limit(), NextReset :: integer()) -> ok. 18 | 19 | %% Resets all the key counters of the scope back to zero. 20 | -callback reset(throttle:scope(), NextReset :: integer()) -> ok. 21 | 22 | %% Increase the access count for the scope/key and return its current value. 23 | -callback update(throttle:scope(), Key :: term()) -> 24 | {Count :: pos_integer(), throttle:rate_limit(), NextReset :: integer()} | rate_not_set. 25 | 26 | %% Retrieve the current access count for the scope/key without increasing it. 27 | -callback lookup(throttle:scope(), Key :: term()) -> 28 | {Count :: pos_integer(), throttle:rate_limit(), NextReset :: integer()} | rate_not_set. 29 | 30 | setup() -> 31 | Module = callback_module(), 32 | Module:setup(). 33 | 34 | initialize(Scope, Limit, NextReset) -> 35 | Module = callback_module(), 36 | Module:initialize(Scope, Limit, NextReset). 37 | 38 | reset(Scope, NextReset) -> 39 | Module = callback_module(), 40 | Module:reset(Scope, NextReset). 41 | 42 | update(Scope, Key) -> 43 | Module = callback_module(), 44 | Module:update(Scope, Key). 45 | 46 | lookup(Scope, Key) -> 47 | Module = callback_module(), 48 | Module:lookup(Scope, Key). 49 | 50 | 51 | %%% Internal 52 | callback_module() -> 53 | application:get_env(throttle, driver, throttle_ets). 54 | -------------------------------------------------------------------------------- /src/drivers/throttle_ets.erl: -------------------------------------------------------------------------------- 1 | -module(throttle_ets). 2 | 3 | -behavior(throttle_driver). 4 | 5 | -export([setup/0, 6 | initialize/3, 7 | reset/2, 8 | update/2, 9 | lookup/2]). 10 | 11 | -define(STATE_TABLE, throttle_state_table). 12 | 13 | %%% intialize the index table that tracks state for all the scopes 14 | setup() -> 15 | ets:new(?STATE_TABLE, [set, named_table, public]), 16 | ok. 17 | 18 | initialize(Scope, Limit, NextReset) -> 19 | TableId = ets:new(scope_counters, [set, public]), 20 | %% add + 1 to allow up to (including) that number 21 | ets:insert(?STATE_TABLE, {Scope, TableId, Limit + 1, NextReset}), 22 | ok. 23 | 24 | reset(Scope, NextReset) -> 25 | [{Scope, TableId, Limit, _PreviousReset}] = ets:lookup(?STATE_TABLE, Scope), 26 | true = ets:delete_all_objects(TableId), 27 | true = ets:insert(?STATE_TABLE, {Scope, TableId, Limit, NextReset}), 28 | ok. 29 | 30 | update(Scope, Key) -> 31 | case ets:lookup(?STATE_TABLE, Scope) of 32 | [{Scope, TableId, Limit, NextReset}] -> 33 | 34 | %% add 1 to counter in position 2, if it's less or equal than Limit, default counter to 0 35 | Count = ets:update_counter(TableId, Key, {2, 1, Limit, Limit}, {Key, 0}), 36 | 37 | {Count, Limit, NextReset}; 38 | [] -> 39 | rate_not_set 40 | end. 41 | 42 | lookup(Scope, Key) -> 43 | case ets:lookup(?STATE_TABLE, Scope) of 44 | [{Scope, TableId, Limit, NextReset}] -> 45 | case ets:lookup(TableId, Key) of 46 | [{Key, Count}] -> 47 | {Count, Limit, NextReset}; 48 | [] -> 49 | {0, Limit, NextReset} 50 | end; 51 | [] -> 52 | rate_not_set 53 | end. 54 | -------------------------------------------------------------------------------- /src/drivers/throttle_mnesia.erl: -------------------------------------------------------------------------------- 1 | -module(throttle_mnesia). 2 | 3 | -behavior(throttle_driver). 4 | 5 | -export([ 6 | setup/0, 7 | initialize/3, 8 | reset/2, 9 | update/2, 10 | lookup/2 11 | ]). 12 | 13 | -record(scope_state, {scope, limit, next_reset, access_context}). 14 | 15 | setup() -> 16 | %% intially this will be called by the sup and every node will get its local 17 | %% version of the tables. After connecting, call this manually so a single 18 | %% copy of the table is shared across nodes. 19 | 20 | AllNodes = [node() | nodes()], 21 | OtherNodes = nodes(), 22 | 23 | case application:start(mnesia) of 24 | ok -> 25 | %% first time, create state table 26 | {atomic, ok} = mnesia:create_table(scope_state, 27 | [{attributes, record_info(fields, scope_state)}, 28 | {ram_copies, AllNodes}, 29 | {type, set}]); 30 | {error, {already_started, mnesia}} -> 31 | %% not the first time, add new nodes 32 | rpc:multicall(OtherNodes, application, stop, [mnesia]), 33 | rpc:multicall(OtherNodes, application, start, [mnesia]), 34 | mnesia:change_config(extra_db_nodes, AllNodes) 35 | end, 36 | ok. 37 | 38 | initialize(Scope, Limit, NextReset) -> 39 | %% counter table uses the default key/value 40 | {atomic, ok} = mnesia:create_table(Scope, [{ram_copies, [node() | nodes()]}, 41 | {type, set}]), 42 | 43 | AccessContext = application:get_env(throttle, access_context, async_dirty), 44 | AddScope = fun() -> 45 | ok = mnesia:write(#scope_state{scope=Scope, 46 | limit=Limit + 1, %% add + 1 to allow up to (including) that number 47 | next_reset=NextReset, 48 | access_context=AccessContext}) 49 | end, 50 | 51 | mnesia:activity(transaction, AddScope). 52 | 53 | reset(Scope, NextReset) -> 54 | {atomic, ok} = mnesia:clear_table(Scope), 55 | 56 | %% update last reset timestamp 57 | [State] = mnesia:dirty_read(scope_state, Scope), 58 | NewState = State#scope_state{next_reset=NextReset}, 59 | ok = mnesia:dirty_write(scope_state, NewState). 60 | 61 | update(Scope, Key) -> 62 | case mnesia:dirty_read(scope_state, Scope) of 63 | [#scope_state{limit=Limit, 64 | next_reset=NextReset, 65 | access_context=AccessContext}] -> 66 | 67 | UpdateCounter = fun() -> 68 | mnesia:dirty_update_counter(Scope, Key, 1) 69 | end, 70 | 71 | Count = mnesia:activity(AccessContext, UpdateCounter), 72 | LimitedCount = min(Limit, Count), 73 | 74 | {LimitedCount, Limit, NextReset}; 75 | [] -> 76 | rate_not_set 77 | end. 78 | 79 | lookup(Scope, Key) -> 80 | case mnesia:dirty_read(scope_state, Scope) of 81 | [#scope_state{limit=Limit, next_reset=NextReset}] -> 82 | 83 | case mnesia:dirty_read(Scope, Key) of 84 | [{Scope, Key, Count}] -> 85 | {Count, Limit, NextReset}; 86 | [] -> 87 | {0, Limit, NextReset} 88 | end; 89 | [] -> 90 | rate_not_set 91 | end. 92 | -------------------------------------------------------------------------------- /src/throttle.app.src: -------------------------------------------------------------------------------- 1 | {application, throttle, 2 | [{description, "Erlang/OTP application to throttle/rate limit resource access"}, 3 | {pkg_name, lambda_throttle}, 4 | {vsn, "0.3.0"}, 5 | {registered, []}, 6 | {mod, { throttle_app, []}}, 7 | {applications, 8 | [kernel, 9 | stdlib 10 | ]}, 11 | {env,[]}, 12 | {modules, []}, 13 | 14 | {maintainers, []}, 15 | {licenses, ["MIT"]}, 16 | {links, [{"Github", "https://github.com/lambdaclass/throttle"}]} 17 | ]}. 18 | -------------------------------------------------------------------------------- /src/throttle.erl: -------------------------------------------------------------------------------- 1 | -module(throttle). 2 | 3 | -behaviour(gen_server). 4 | 5 | -export([setup/3, 6 | check/2, 7 | peek/2, 8 | 9 | start_link/3, 10 | init/1, 11 | handle_call/3, 12 | handle_cast/2, 13 | handle_info/2]). 14 | 15 | -type scope() :: atom(). 16 | -type rate_limit() :: pos_integer(). 17 | 18 | -export_type([scope/0, rate_limit/0]). 19 | 20 | %% API functions 21 | 22 | %%% setup throttling for a specific scope 23 | setup(Scope, RateLimit, RatePeriod) -> 24 | {ok, _Pid} = supervisor:start_child(throttle_sup, [Scope, RateLimit, RatePeriod]), 25 | ok. 26 | 27 | check(Scope, Key) -> 28 | Result = throttle_driver:update(Scope, Key), 29 | count_result(Result). 30 | 31 | peek(Scope, Key) -> 32 | Result = throttle_driver:lookup(Scope, Key), 33 | count_result(Result). 34 | 35 | %% Gen server callbacks 36 | start_link(Scope, Limit, Period) -> 37 | gen_server:start_link(?MODULE, {Scope, Limit, Period}, []). 38 | 39 | init({Scope, Limit, Period} = State) -> 40 | NextReset = throttle_time:schedule_reset(Period), 41 | throttle_driver:initialize(Scope, Limit, NextReset), 42 | {ok, _} = timer:send_interval(throttle_time:interval(Period), reset_counters), 43 | {ok, State}. 44 | 45 | handle_info(reset_counters, {Scope, _Limit, Period} = State) -> 46 | NextReset = throttle_time:schedule_reset(Period), 47 | throttle_driver:reset(Scope, NextReset), 48 | {noreply, State}. 49 | 50 | handle_call(_Request, _From, State) -> 51 | {noreply, State}. 52 | 53 | handle_cast(_Request, State) -> 54 | {noreply, State}. 55 | 56 | %%% Internal functions 57 | count_result({Count, Limit, NextReset}) when Count == Limit -> 58 | LeftToReset = throttle_time:left_til(NextReset), 59 | {limit_exceeded, 0, LeftToReset}; 60 | count_result({Count, Limit, NextReset}) -> 61 | LeftToReset = throttle_time:left_til(NextReset), 62 | {ok, Limit - Count - 1, LeftToReset}; 63 | count_result(rate_not_set) -> 64 | rate_not_set. 65 | -------------------------------------------------------------------------------- /src/throttle_app.erl: -------------------------------------------------------------------------------- 1 | -module(throttle_app). 2 | 3 | -behaviour(application). 4 | 5 | %% Application callbacks 6 | -export([start/2, stop/1]). 7 | 8 | start(_StartType, _StartArgs) -> 9 | %% start supervisor first so its available when calling throttle:setup 10 | {ok, Pid} = throttle_sup:start_link(), 11 | 12 | case application:get_env(throttle, rates) of 13 | {ok, Rates} -> 14 | lists:foreach(fun({Scope, Limit, Period}) -> 15 | throttle:setup(Scope, Limit, Period) 16 | end, Rates); 17 | _ -> 18 | ok 19 | end, 20 | {ok, Pid}. 21 | 22 | stop(_State) -> 23 | ok. 24 | -------------------------------------------------------------------------------- /src/throttle_sup.erl: -------------------------------------------------------------------------------- 1 | %%%------------------------------------------------------------------- 2 | %% @doc throttle top level supervisor. 3 | %% @end 4 | %%%------------------------------------------------------------------- 5 | 6 | -module(throttle_sup). 7 | 8 | -behaviour(supervisor). 9 | 10 | %% API 11 | -export([start_link/0]). 12 | 13 | %% Supervisor callbacks 14 | -export([init/1]). 15 | 16 | -define(SERVER, ?MODULE). 17 | 18 | %% API functions 19 | 20 | start_link() -> 21 | supervisor:start_link({local, ?SERVER}, ?MODULE, []). 22 | 23 | %% Supervisor callbacks 24 | 25 | init([]) -> 26 | throttle_driver:setup(), 27 | 28 | {ok, { #{ strategy => simple_one_for_one, intensity => 5, period => 1 }, 29 | [#{ 30 | id => throttle, 31 | start => {throttle, start_link, []}, 32 | restart => transient, 33 | shutdown => 5000, 34 | type => worker, 35 | modules => [throttle] 36 | }] 37 | }}. 38 | -------------------------------------------------------------------------------- /src/throttle_time.erl: -------------------------------------------------------------------------------- 1 | -module(throttle_time). 2 | 3 | -export([now/0, 4 | interval/1, 5 | schedule_reset/1, 6 | left_til/1]). 7 | 8 | -type interval() :: per_day | per_hour | per_minute | per_second | pos_integer(). 9 | -export_type([interval/0]). 10 | 11 | now() -> 12 | erlang:system_time(millisecond). 13 | 14 | interval(per_day) -> 15 | 1000 * 60 * 60 * 24; 16 | interval(per_hour) -> 17 | 1000 * 60 * 60; 18 | interval(per_minute) -> 19 | 1000 * 60; 20 | interval(per_second) -> 21 | 1000; 22 | interval(CustomMs) when is_integer(CustomMs) -> 23 | CustomMs. 24 | 25 | schedule_reset(Period) -> 26 | throttle_time:now() + interval(Period). 27 | 28 | left_til(NextReset) -> 29 | NextReset - throttle_time:now(). 30 | -------------------------------------------------------------------------------- /test/profile.erl: -------------------------------------------------------------------------------- 1 | -module(profile). 2 | 3 | -export([stress/4, 4 | profile/1]). 5 | 6 | %% TODO shorter arity with sane defaults 7 | 8 | stress(Driver, Scopes, Limit, Processes) -> 9 | init(Driver, Scopes, Limit, Processes), 10 | timer:sleep(500), 11 | 12 | %% start profiler 13 | {ok, _} = eprof:profile(fun() -> 14 | throttle:check(1, self()) 15 | end), 16 | 17 | eprof:stop_profiling(), 18 | eprof:analyze(). 19 | 20 | %% test the standalone reset_counters function 21 | %% which is what changes between ets implementation 22 | profile(Driver) -> 23 | Scope = profile_test, 24 | Driver:init(), 25 | Driver:init_counters(Scope, 1000, per_second), 26 | 27 | %% setup some counters 28 | lists:foreach(fun(N) -> 29 | Driver:update_counter(Scope, N) 30 | end, lists:seq(0, 1000)), 31 | 32 | %% time reset counters 33 | Result = timer:tc(Driver, reset_counters, [Scope]), 34 | 35 | %% cleanup ets 36 | ets:delete(throttle_state_table), 37 | case Driver of 38 | throttle_ets_match -> ets:delete(throttle_counter_table); 39 | _ -> ok 40 | end, 41 | Result. 42 | 43 | %% internal 44 | init(Driver, Scopes, Limit, Processes) -> 45 | application:set_env(throttle, driver, Driver), 46 | application:set_env(throttle, rates, [{N, Limit, per_second} || N <- lists:seq(0, Scopes)]), 47 | {ok, _Started} = application:ensure_all_started(throttle), 48 | spawn_processes(Scopes, Processes). 49 | 50 | spawn_processes(Scopes, Processes) -> 51 | lists:foreach(fun (N) -> 52 | Scope = N rem Scopes, 53 | %% spawn process that constantly queries throttle 54 | spawn(fun() -> loop(Scope) end) 55 | end, lists:seq(0, Processes)). 56 | 57 | loop(Scope) -> 58 | case throttle:check(Scope, self()) of 59 | {ok, _, _} -> 60 | dumb_operation(), 61 | loop(Scope); 62 | {limit_exceeded, _, _} -> 63 | loop(Scope) 64 | end. 65 | 66 | dumb_operation() -> 67 | %% does this make sense? 68 | 1 + 1. 69 | -------------------------------------------------------------------------------- /test/throttle_distributed_SUITE.erl: -------------------------------------------------------------------------------- 1 | -module(throttle_distributed_SUITE). 2 | 3 | -include_lib("common_test/include/ct.hrl"). 4 | 5 | -compile(export_all). 6 | 7 | groups() -> 8 | %% repeat to avoid "lucky" successes 9 | [{async_dirty, [{repeat, 5}], [test_limit]}, 10 | {transaction, [{repeat, 5}], [test_limit]}, 11 | {sync_transaction, [{repeat, 5}], [test_limit]} 12 | ]. 13 | 14 | all() -> 15 | [{group, async_dirty}, 16 | {group, transaction}, 17 | {group, sync_transaction}]. 18 | 19 | init_per_suite(Config) -> 20 | {ok, Slave} = start_slave(), 21 | [{slave, Slave} | Config]. 22 | 23 | end_per_suite(Config) -> 24 | Slave = proplists:get_value(slave, Config), 25 | ct_slave:stop(Slave), 26 | ok. 27 | 28 | init_per_group(AccessContext, Config) -> 29 | Slave = proplists:get_value(slave, Config), 30 | ok = rpc:call(Slave, throttle_distributed_SUITE, start_throttle, [AccessContext]), 31 | start_throttle(AccessContext), 32 | 33 | Sleep = case AccessContext of 34 | %% async takes a bit to propagate the changes after the operation returns 35 | async_dirty -> 50; 36 | _ -> 0 37 | end, 38 | [{sleep, Sleep} | Config]. 39 | 40 | end_per_group(_Context, Config) -> 41 | application:stop(throttle), 42 | ok. 43 | 44 | start_slave() -> 45 | Slave = 'node2@127.0.0.1', 46 | 47 | %% make the same binaries available to the other node 48 | CodePath = code:get_path(), 49 | PathFlag = "-pa " ++ lists:concat(lists:join(" ", CodePath)), 50 | {ok, _} = ct_slave:start(Slave, [{erl_flags, PathFlag}]), 51 | {ok, Slave}. 52 | 53 | start_throttle(Context) -> 54 | application:ensure_all_started(lager), 55 | ok = application:set_env(throttle, driver, throttle_mnesia), 56 | ok = application:set_env(throttle, access_context, Context), 57 | {ok, _Apps} = application:ensure_all_started(throttle), 58 | throttle_mnesia:setup(), 59 | ok. 60 | 61 | test_limit(Config) -> 62 | Slave = proplists:get_value(slave, Config), 63 | Sleep = proplists:get_value(sleep, Config), 64 | Scope = binary_to_atom(integer_to_binary(rand:uniform(10000)), latin1), 65 | throttle:setup(Scope, 4, per_second), 66 | 67 | %% check both here and in the other node 68 | {ok, 3, _} = throttle:check(Scope, <<"john">>), 69 | timer:sleep(Sleep), 70 | {ok, 2, _} = rpc:call(Slave, throttle, check, [Scope, <<"john">>]), 71 | timer:sleep(Sleep), 72 | {ok, 1, _} = throttle:check(Scope, <<"john">>), 73 | timer:sleep(Sleep), 74 | {ok, 0, _} = rpc:call(Slave, throttle, check, [Scope, <<"john">>]), 75 | timer:sleep(Sleep), 76 | {limit_exceeded, 0, _} = throttle:check(Scope, <<"john">>), 77 | 78 | timer:sleep(1100), 79 | 80 | {ok, 3, _} = throttle:check(Scope, <<"john">>), 81 | 82 | ok. 83 | -------------------------------------------------------------------------------- /test/throttle_test_SUITE.erl: -------------------------------------------------------------------------------- 1 | -module(throttle_test_SUITE). 2 | -include_lib("common_test/include/ct.hrl"). 3 | 4 | -compile(export_all). 5 | 6 | -define(ALL_TESTS, [test_limit, 7 | test_peek, 8 | test_scopes_not_mixed, 9 | test_keys_not_mixed, 10 | test_minute, 11 | test_hour, 12 | test_day, 13 | test_custom_ms, 14 | rate_not_set]). 15 | 16 | %% we want to repeat the same suit with the different drivers 17 | groups() -> 18 | [{throttle_ets, [], ?ALL_TESTS}, 19 | {throttle_mnesia, [], ?ALL_TESTS}]. 20 | 21 | all() -> 22 | [{group, throttle_ets}, 23 | {group, throttle_mnesia}]. 24 | 25 | init_per_suite(Config) -> 26 | application:ensure_all_started(lager), 27 | Config. 28 | 29 | end_per_suite(_Config) -> 30 | ok. 31 | 32 | init_per_group(Driver, Config) -> 33 | ok = application:set_env(throttle, driver, Driver), 34 | {ok, _Apps} = application:ensure_all_started(throttle), 35 | Config. 36 | 37 | end_per_group(_Driver, _Config) -> 38 | ok = application:stop(throttle), 39 | ok = application:unload(throttle). 40 | 41 | test_limit(_Config) -> 42 | throttle:setup(test_rate, 3, per_second), 43 | 44 | {ok, 2, _} = throttle:check(test_rate, <<"john">>), 45 | {ok, 1, _} = throttle:check(test_rate, <<"john">>), 46 | {ok, 0, _} = throttle:check(test_rate, <<"john">>), 47 | {limit_exceeded, 0, _} = throttle:check(test_rate, <<"john">>), 48 | 49 | timer:sleep(1100), 50 | 51 | {ok, 2, _} = throttle:check(test_rate, <<"john">>), 52 | 53 | ok. 54 | 55 | test_peek(_Config) -> 56 | throttle:setup(test_rate2, 2, per_second), 57 | 58 | {ok, 2, _} = throttle:peek(test_rate2, <<"john">>), 59 | {ok, 2, _} = throttle:peek(test_rate2, <<"john">>), 60 | {ok, 2, _} = throttle:peek(test_rate2, <<"john">>), 61 | 62 | ok. 63 | 64 | test_scopes_not_mixed(_Config) -> 65 | throttle:setup(test_rate3, 3, per_second), 66 | throttle:setup(test_rate4, 3, per_second), 67 | 68 | {ok, 2, _} = throttle:check(test_rate3, <<"john">>), 69 | {ok, 1, _} = throttle:check(test_rate3, <<"john">>), 70 | {ok, 0, _} = throttle:check(test_rate3, <<"john">>), 71 | {limit_exceeded, 0, _} = throttle:check(test_rate3, <<"john">>), 72 | 73 | {ok, 2, _} = throttle:check(test_rate4, <<"john">>), 74 | 75 | ok. 76 | 77 | test_keys_not_mixed(_Config) -> 78 | throttle:setup(test_rate5, 2, per_second), 79 | 80 | {ok, 1, _} = throttle:check(test_rate5, <<"john">>), 81 | {ok, 0, _} = throttle:check(test_rate5, <<"john">>), 82 | {limit_exceeded, 0, _} = throttle:check(test_rate5, <<"john">>), 83 | 84 | {ok, 1, _} = throttle:check(test_rate5, <<"mary">>), 85 | 86 | ok. 87 | 88 | test_minute(_Config) -> 89 | throttle:setup(test_rate6, 3, per_minute), 90 | 91 | {ok, 2, TimeToReset} = throttle:check(test_rate6, <<"john">>), 92 | 93 | Diff = 1000 * 60 - TimeToReset, 94 | true = Diff < 1000, 95 | 96 | ok. 97 | 98 | test_hour(_Config) -> 99 | throttle:setup(test_rate7, 3, per_hour), 100 | 101 | {ok, 2, TimeToReset} = throttle:check(test_rate7, <<"john">>), 102 | 103 | Diff = 1000 * 60 * 60 - TimeToReset, 104 | true = Diff < 1000, 105 | 106 | ok. 107 | 108 | test_day(_Config) -> 109 | throttle:setup(test_rate8, 3, per_day), 110 | 111 | {ok, 2, TimeToReset} = throttle:check(test_rate8, <<"john">>), 112 | 113 | Diff = 1000 * 60 * 60 * 24 - TimeToReset, 114 | true = Diff < 1000, 115 | 116 | ok. 117 | 118 | test_custom_ms(_Config) -> 119 | throttle:setup(test_rate9, 3, 10000), 120 | 121 | {ok, 2, TimeToReset} = throttle:check(test_rate9, <<"john">>), 122 | 123 | Diff = 10000 - TimeToReset, 124 | true = Diff < 1000, 125 | 126 | ok. 127 | 128 | rate_not_set(_Config) -> 129 | rate_not_set = throttle:check(didnt_set, <<"john">>), 130 | 131 | ok. 132 | --------------------------------------------------------------------------------