|
@@ -1,5 +1,12 @@
|
|
|
-module(token_bucket).
|
|
|
|
|
|
+-behaviour(gen_server).
|
|
|
+
|
|
|
+%% callbacks
|
|
|
+-export([start_link/0]).
|
|
|
+-export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]).
|
|
|
+
|
|
|
+%% API
|
|
|
-export([
|
|
|
is_limit_reached/1,
|
|
|
is_limit_reached/2
|
|
@@ -15,5 +22,69 @@ is_limit_reached(UserId) ->
|
|
|
is_limit_reached(UserId, ?DEFAULT_RPS).
|
|
|
|
|
|
-spec is_limit_reached(user_id(), max_rps()) -> boolean().
|
|
|
-is_limit_reached(_UserId, _MaxRps) ->
|
|
|
- {error, not_implemented}.
|
|
|
+is_limit_reached(_UserId, infinity) -> false; %% infinity means no limit
|
|
|
+is_limit_reached(UserId, MaxRps) ->
|
|
|
+ Timestamp = get_timestamp_now(),
|
|
|
+ Key = {Timestamp, UserId},
|
|
|
+ {ok, RPS_Count} = gen_server:call(?MODULE, {get_rps_count, Key}),
|
|
|
+ RPS_Count2 = RPS_Count + 1,
|
|
|
+ case RPS_Count2 > MaxRps of
|
|
|
+ true ->
|
|
|
+ %% limit is reached
|
|
|
+ %% ok = gen_server:call(?MODULE, {set_rps_count, {Key, RPS_Count2}}), %% uncomment this for update RPS_Count even when count in upper limitation
|
|
|
+ true;
|
|
|
+ false ->
|
|
|
+ %% limit not reached
|
|
|
+ ok = gen_server:call(?MODULE, {set_rps_count, {Key, RPS_Count2}}),
|
|
|
+ false
|
|
|
+ end.
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+start_link() ->
|
|
|
+ gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).
|
|
|
+
|
|
|
+
|
|
|
+init([]) ->
|
|
|
+ ets:new(users_rps_counters_table, [set, named_table, { keypos, 1 }, private]), % Key = {Timestamp, User_Id}
|
|
|
+
|
|
|
+ State = [],
|
|
|
+ {ok, State}.
|
|
|
+
|
|
|
+
|
|
|
+handle_call({get_rps_count, {Timestamp, User_Id} = Key}, _From, State) ->
|
|
|
+ Count = case ets:lookup(users_rps_counters_table, Key) of
|
|
|
+ [{{Timestamp, User_Id}, RPS_Count}] -> RPS_Count;
|
|
|
+ _ -> 0
|
|
|
+ end,
|
|
|
+ {reply, {ok, Count}, State};
|
|
|
+
|
|
|
+
|
|
|
+handle_call({set_rps_count, {{_Timestamp, _User_Id}, _Count} = V}, _From, State) ->
|
|
|
+ true = ets:insert(users_rps_counters_table, V),
|
|
|
+ {reply, ok, State};
|
|
|
+
|
|
|
+
|
|
|
+handle_call(_Req, _From, State) ->
|
|
|
+ {reply, not_handled, State}.
|
|
|
+
|
|
|
+
|
|
|
+handle_cast(_Req, State) ->
|
|
|
+ {noreply, State}.
|
|
|
+
|
|
|
+
|
|
|
+handle_info(_Request, State) ->
|
|
|
+ {noreply, State}.
|
|
|
+
|
|
|
+
|
|
|
+terminate(_Reason, _State) ->
|
|
|
+ ok.
|
|
|
+
|
|
|
+
|
|
|
+code_change(_OldVsn, State, _Extra) ->
|
|
|
+ {ok, State}.
|
|
|
+
|
|
|
+
|
|
|
+get_timestamp_now() ->
|
|
|
+ erlang:system_time(second).
|
|
|
+
|