Browse Source

token_bucket

221V 2 years ago
parent
commit
a17ddf82be
5 changed files with 98 additions and 3 deletions
  1. 6 0
      Makefile
  2. 73 2
      src/token_bucket.erl
  3. 7 1
      src/token_bucket_sup.erl
  4. 3 0
      sys.config
  5. 9 0
      vm.args

+ 6 - 0
Makefile

@@ -4,3 +4,9 @@ PROJECT_VERSION = 0.1.0
 
 include erlang.mk
 
+run:
+	erl -pa ebin +pc unicode -args_file vm.args -config sys.config -eval 'application:start(token_bucket)'
+# add -pa ebin in erl .. command because no deps with nested structure
+#   and no ERL_LIBS="apps:deps" before erl .. command
+
+.PHONY: run

+ 73 - 2
src/token_bucket.erl

@@ -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).
+

+ 7 - 1
src/token_bucket_sup.erl

@@ -9,5 +9,11 @@ start_link() ->
   supervisor:start_link({local, ?MODULE}, ?MODULE, []).
 
 init([]) ->
-  Procs = [],
+  Procs = [
+    #{id => worker_1,
+      start => {token_bucket, start_link, []},
+      restart => permanent,
+      type => worker,
+      modules => [token_bucket]}
+  ],
   {ok, { #{strategy => one_for_one, intensity => 10, period => 100}, Procs}}.

+ 3 - 0
sys.config

@@ -0,0 +1,3 @@
+[
+].
+

+ 9 - 0
vm.args

@@ -0,0 +1,9 @@
+-name token_bucket@127.0.0.1
+-setcookie D23uthURSyk6kRdP
++pc unicode
++K true
++A 5
+-env ERL_MAX_PORTS 4096
+-env ERL_FULLSWEEP_AFTER 10
+-env ERL_EPMD_ADDRESS 127.0.0.1
+-kernel inet_dist_use_interface {127,0,0,1}