Add ability to update throttle by custom amount

rebar.config

@@ -9,4 +9,5 @@
                     {cover_export_enabled   , true},
                     {coveralls_coverdata    , "_build/test/cover/ct.coverdata"},
                     {coveralls_service_name , "travis-ci"},
+                    {erl_opts, [nowarn_export_all]},
                     {deps, [{lager, "3.5.1"}]}]}]}.

src/drivers/throttle_driver.erl

@@ -7,7 +7,7 @@
 -export([setup/0,
          initialize/3,
          reset/2,
-         update/2,
+         update/3,
          lookup/2]).
 
 %% Performs any global (non scope-specific) setup required by the driver.
@@ -20,7 +20,7 @@
 -callback reset(throttle:scope(), NextReset :: integer()) -> ok.
 
 %% Increase the access count for the scope/key and return its current value.
--callback update(throttle:scope(), Key :: term()) ->
+-callback update(throttle:scope(), Key :: term(), Value :: pos_integer()) ->
     {Count :: pos_integer(), throttle:rate_limit(), NextReset :: integer()} | rate_not_set.
 
 %% Retrieve the current access count for the scope/key without increasing it.
@@ -39,9 +39,9 @@ reset(Scope, NextReset) ->
     Module = callback_module(),
     Module:reset(Scope, NextReset).
 
-update(Scope, Key) ->
+update(Scope, Key, Value) ->
     Module = callback_module(),
-    Module:update(Scope, Key).
+    Module:update(Scope, Key, Value).
 
 lookup(Scope, Key) ->
     Module = callback_module(),

src/drivers/throttle_ets.erl

@@ -5,7 +5,7 @@
 -export([setup/0,
          initialize/3,
          reset/2,
-         update/2,
+         update/3,
          lookup/2]).
 
 -define(STATE_TABLE, throttle_state_table).
@@ -27,12 +27,12 @@ reset(Scope, NextReset) ->
   true = ets:insert(?STATE_TABLE, {Scope, TableId, Limit, NextReset}),
   ok.
 
-update(Scope, Key) ->
+update(Scope, Key, Increment) ->
   case ets:lookup(?STATE_TABLE, Scope) of
     [{Scope, TableId, Limit, NextReset}] ->
 
-      %% add 1 to counter in position 2, if it's less or equal than Limit, default counter to 0
-      Count = ets:update_counter(TableId, Key, {2, 1, Limit, Limit}, {Key, 0}),
+      %% add `Increment' to counter in position 2, if it's less or equal than Limit, default counter to 0
+      Count = ets:update_counter(TableId, Key, {2, Increment, Limit, Limit}, {Key, 0}),
 
       {Count, Limit, NextReset};
     [] ->

src/drivers/throttle_mnesia.erl

@@ -6,7 +6,7 @@
          setup/0,
          initialize/3,
          reset/2,
-         update/2,
+         update/3,
          lookup/2
         ]).
 
@@ -58,14 +58,14 @@ reset(Scope, NextReset) ->
   NewState = State#scope_state{next_reset=NextReset},
   ok = mnesia:dirty_write(scope_state, NewState).
 
-update(Scope, Key) ->
+update(Scope, Key, Increment) ->
   case mnesia:dirty_read(scope_state, Scope) of
     [#scope_state{limit=Limit,
                   next_reset=NextReset,
                   access_context=AccessContext}] ->
 
       UpdateCounter = fun() ->
-                          mnesia:dirty_update_counter(Scope, Key, 1)
+                          mnesia:dirty_update_counter(Scope, Key, Increment)
                       end,
 
       Count = mnesia:activity(AccessContext, UpdateCounter),

src/throttle.erl

@@ -4,6 +4,7 @@
 
 -export([setup/3,
          check/2,
+         update/3,
          peek/2,
 
          start_link/3,
@@ -25,7 +26,11 @@ setup(Scope, RateLimit, RatePeriod) ->
   ok.
 
 check(Scope, Key) ->
-  Result = throttle_driver:update(Scope, Key),
+  Result = throttle_driver:update(Scope, Key, 1),
+  count_result(Result).
+
+update(Scope, Key, Value) when is_integer(Value), Value > 0 ->
+  Result = throttle_driver:update(Scope, Key, Value),
   count_result(Result).
 
 peek(Scope, Key) ->

test/throttle_distributed_SUITE.erl

@@ -37,7 +37,7 @@ init_per_group(AccessContext, Config) ->
           end,
   [{sleep, Sleep} | Config].
 
-end_per_group(_Context, Config) ->
+end_per_group(_Context, _Config) ->
   application:stop(throttle),
   ok.
 

test/throttle_test_SUITE.erl

@@ -46,6 +46,10 @@ test_limit(_Config) ->
   {ok, 0, _} = throttle:check(test_rate, <<"john">>),
   {limit_exceeded, 0, _} = throttle:check(test_rate, <<"john">>),
 
+  {ok, 1, _} = throttle:update(test_rate, <<"paul">>, 2),
+  {ok, 0, _} = throttle:update(test_rate, <<"paul">>, 1),
+  {limit_exceeded, 0, _} = throttle:update(test_rate, <<"paul">>, 2),
+
   timer:sleep(1100),
 
   {ok, 2, _} = throttle:check(test_rate, <<"john">>),