util

%%------------------------------------------------------------------------------
%% Copyright (c) 2011 Serge Aleynikov
%%
%% Permission is hereby granted, free of charge, to any person
%% obtaining a copy of this software and associated documentation
%% files (the "Software"), to deal in the Software without restriction,
%% including without limitation the rights to use, copy, modify, merge,
%% publish, distribute, sublicense, and/or sell copies of the Software,
%% and to permit persons to whom the Software is furnished to do
%% so, subject to the following conditions:
%%
%% The above copyright notice and this permission notice shall be included
%% in all copies or substantial portions of the Software.
%%
%% THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
%% EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
%% MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
%% IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
%% CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
%% TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
%% SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
%%------------------------------------------------------------------------------
-module(throttle).
-moduledoc """
Throttle given rate over a number of seconds.

Implementation uses time spacing reservation algorithm where each allocation of
samples reserves a fraction of space in the throttling window. The reservation
gets freed as the time goes by.  No more than the `Rate` number of samples are
allowed to fit in the milliseconds `Window`.

This is an Erlang implementation of the throttling algorithm from the utxx
library found at this URL:
[https://github.com/saleyn/utxx/blob/master/include/utxx/rate_throttler.hpp]

Author: Serge Aleynikov <saleyn at gmail dot com>
""".
-export([new/1, new/2, new/3, available/1, available/2, used/1, used/2]).
-export([next_timeout/1, call/2, call/3, call/4, call/5]).
-export([now/0, reset/1, reset/2, add/1, add/2, add/3, curr_rps/1, curr_rps/2]).

-compile({no_auto_import,[now/0]}).

-record(throttle, {
    rate
  , window  :: integer()    %% In microseconds
  , step    :: integer()    %% In microseconds
  , next_ts :: integer()
}).

-type throttle() :: #throttle{}.
-type time()     :: non_neg_integer().

-doc """
`retries`     - number of retries. `retry_delay` - delay in milliseconds between
successive retries. `blocking`    - instructs to block the call if throttled.
""".
-type throttle_opts() :: #{
  retries     => integer(),
  retry_delay => integer(),
  blocking    => boolean()
}.

-type throttle_result() ::
  {ok, any()} | {error, throttled | {Reason::any(), StackTrace::list()}}.

-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
-endif.

-doc "Create a new throttle given the `Rate` per second.".
-spec new(non_neg_integer()) -> throttle().
new(Rate) ->
  new(Rate, 1000).

%% @see new/3
-spec new(non_neg_integer(), non_neg_integer()) -> throttle().
new(Rate, Window) ->
  new(Rate, Window, now()).

-doc """
Create a new throttle given the `Rate` per `Window` milliseconds. `Now` is
expressesed in microseconds since epoch using `now()`.
""".
-spec new(non_neg_integer(), non_neg_integer(), time()) -> throttle().
new(Rate, Window, Now) when is_integer(Rate), is_integer(Window), is_integer(Now) ->
  Win  = Window * 1000,
  Step = if Rate == 0 -> 0; true -> Win div Rate end,
  #throttle{rate = Rate, window = Win, step = Step, next_ts = Now}.

%% @see reset/2
reset(#throttle{} = T) ->
  reset(T, now()).

-doc "Reset the throttle request counter".
reset(#throttle{} = T, Now) when is_integer(Now) ->
  T#throttle{next_ts = Now}.

-doc """
Call the lambda `F`, ensuring that it's not called more frequently than the
throttle would allow.

Example:
```
1> T = throttle:new(10, 1000). 2> lists:foldl(fun(_,{T1,A}) -> {T2,R} =
throttle:call(T1, fun() -> http:get("google.com") end), {T2, [R|A]} end, {T,[]},
lists:seq(1, 100)).
```
""".
call(T, F) ->
  call2(T, F, #{}, now()).

-doc """
Call the lambda `F`, ensuring that it's not called more often then the throttle
would allow. `Opts` are a map of options. When `{retries, R}` option is given
and `R` is greater than 0, the throttler will call the function `F()` up to `R`
times if the `F()` raises an exception.  The delay between retries is controlled
by the `{retry_delay, D}` options, expressed in milliseconds (default: `1`)
between successive executions of `F()`. If `F()` still raises an exception after
the R's retry, that exception would be reraised and it would be the
responsibility of the caller to handle it.
""".
-spec call(#throttle{}, fun(() -> any()), throttle_opts()) ->
  {#throttle{}, throttle_result()}.
call(#throttle{} = T, F, Opts) when is_function(F, 0), is_map(Opts) ->
  call2(T, F, Opts, now()).

-doc """
Call M,F,A, ensuring that it's not called more frequently than the throttle
would allow.

Example:
```
1> T = throttle:new(10, 1000). 2> lists:foldl(fun(_,{T1,A}) -> {T2,R} =
throttle:call(T1, http, get, ["google.com"]), {T2, [R|A]} end, {T,[]},
lists:seq(1, 100)).
```
""".
call(T, M,F,A) when is_atom(M), is_atom(F), is_list(A) ->
  call(T, M,F,A, #{}, now());
call(T, F, Opts, Now) when is_function(F, 0), is_map(Opts), is_integer(Now) ->
  call2(T, F, Opts, Now).

-doc """
Call M,F,A, ensuring that it's not called more frequently than the throttle
would allow.
""".
call(T, M,F,A, Now) when is_integer(Now) ->
  call(T, M,F,A, #{}, Now);
call(T, M,F,A, Opts) when is_map(Opts) ->
  call(T, M,F,A, Opts, now()).

-doc """
Call M,F,A, ensuring that it's not called more frequently than the throttle
would allow.
""".
-spec call(#throttle{}, atom(), atom(), [any()], non_neg_integer(), throttle_opts()) ->
  {#throttle{}, throttle_result()}.
call(#throttle{} = T, M,F,A, Opts, Now) when is_atom(M), is_atom(F), is_list(A) ->
  call2(T, fun() -> apply(M,F,A) end, Opts, Now).

-spec call2(#throttle{}, fun(() -> any()), throttle_opts(), non_neg_integer()) ->
  {#throttle{}, throttle_result()}.
call2(#throttle{} = T, F, Opts, Now) when is_integer(Now), is_function(F, 0), is_map(Opts) ->
  Retries = maps:get(retries,     Opts, 0),
  DelayMS = maps:get(retry_delay, Opts, 1),
  Block   = maps:get(blocking,    Opts, true),
  call3(T, F, Now, Block, Retries, DelayMS).

call3(T, F, Now, Block, Retries, DelayMS) ->
  case next_timeout(T, Now) of
    0 ->
      {1, T1} = add(T, 1, Now),
      Result  = try {ok, F()} catch _:R:Trace -> {retry, {R, Trace}} end,
      case Result of
        {ok, _} ->
          {T1, Result};
        {retry, _} when Retries > 0 ->
          receive after DelayMS -> ok end,
          call3(T, F, now(), Block, Retries-1, DelayMS);
        {retry, Error} ->
          {error, Error}
      end;
    _ when not Block ->
      {error, throttled};
    WaitMS ->
      receive after WaitMS -> ok end,
      call3(T, F, now(), Block, Retries, DelayMS)
  end.

-doc "Add one sample to the throttle".
add(T)          -> add(T, 1).

-doc "Add `Samples` to the throttle".
add(T, Samples) -> add(T, Samples, now()).

-doc """
Add `Samples` to the throtlle. Return `{FitSamples, State}`, where `FitSamples`
are the number of samples that fit in the throttling window. 0 means that the
throttler is fully congested, and more time needs to elapse before the throttles
gets reset to accept more samples.
""".
-spec add(throttle(), integer(), time()) -> {integer(), throttle()}.
add(#throttle{rate = 0}, Samples, _Now) ->
  Samples;
add(#throttle{next_ts = TS, step = Step, window = Win} = T, Samples, Now) ->
  NextTS    = TS     + Samples * Step,
  NowNextTS = Now    + Win,
  Diff      = NextTS - NowNextTS,
  if
    Diff < -Win ->
      {Samples, T#throttle{next_ts = Now + Step}};
    Diff < 0 ->
      {Samples, T#throttle{next_ts = NextTS}};
    true ->
      N = max(0, Samples - (Diff div Step)),
      {N, T#throttle{next_ts = TS + N * Step}}
  end.

%% @see available/2
available(T) -> available(T, now()).

-doc "Return the number of available samples given `Now` current time.".
available(#throttle{rate=0}=T,_Now) -> T#throttle.window;
available(#throttle{}      =T, Now) -> calc_available(T, Now).

%% @see used/2
used(T) -> used(T, now()).

-doc "Return the number of used samples given `a_now` current time.".
used(#throttle{rate = 0},     _Now) -> 0;
used(#throttle{rate = R} = T,  Now) -> R-calc_available(T, Now).

-doc """
Return the number of milliseconds to wait until the throttling threshold is
satisfied to fit another sample.
""".
next_timeout(T) -> next_timeout(T, now()).
next_timeout(#throttle{next_ts = TS, step = Step, window = Win}, Now) ->
  NextTS    = TS     + Step,
  NowNextTS = Now    + Win,
  Diff      = NextTS - NowNextTS,
  if
    Diff =< 0 -> 0;
    true      -> ceil(Diff / 1000)
  end.
  
%% @see curr_rps/2
curr_rps(T) -> curr_rps(T, now()).

-doc "Return currently used rate per second.".
curr_rps(#throttle{rate=0},   _Now) -> 0;
curr_rps(#throttle{rate=R}=T,  Now) ->
  (R-calc_available(T, Now))*1000000/T#throttle.window.

now() ->
  erlang:system_time(microsecond).

%%------------------------------------------------------------------------------
%% Internal functions
%%------------------------------------------------------------------------------

%% Return the number of available samples given `Now' current time.
calc_available(#throttle{rate=R, window=W, step=S} = T, Now) ->
  Diff = Now - T#throttle.next_ts,
  if Diff >= 0 -> R;
     true      -> min(R, max(0, (W+Diff) div S))
  end.

%%------------------------------------------------------------------------------
%% Unit testing
%%------------------------------------------------------------------------------

-ifdef(EUNIT).
time(TS, US) -> erlang:universaltime_to_posixtime(TS) * 1000000 + US.

all_test() ->

  Now = time({{2015, 6, 1}, {11,59,58}}, 900000),
  Thr = throttle:new(10, 1000, Now),  %% Throttle 10 samples / sec
  ?assertEqual(100000, Thr#throttle.step),
  ?assertEqual(10,     throttle:available(Thr, Now)),
  {N, T1} = throttle:add(Thr, 1, Now),
  ?assertEqual(1, N),
  ?assertEqual(9, throttle:available(T1, Now)),
  ?assertEqual(1, throttle:used(T1, Now)),

  Now1 = time({{2015, 6, 1}, {11,59,58}}, 999999),
  ?assertEqual(9, throttle:available(T1, Now1)),
  ?assertEqual(1, throttle:used(T1, Now1)),

  Now2 = time({{2015, 6, 1}, {11,59,59}}, 0),
  ?assertEqual(10,throttle:available(T1, Now2)),
  ?assertEqual(0, throttle:used(T1, Now2)),

  Now3 = time({{2015, 6, 1}, {12,0,0}}, 0),
  % 1 second elapsed, the throttler's interval is reset, and 10 samples are available
  ?assertEqual(10,throttle:available(T1, Now3)),
  ?assertEqual(0, throttle:used(T1, Now2)),

  {N1, T2} = throttle:add(T1, 1, Now3),
  ?assertEqual(1, N1),
  ?assertEqual(9, throttle:available(T2, Now3)),
  ?assertEqual(1, throttle:used(T2, Now3)),

  TT = throttle:new(5, 1000),
  Now4 = now(),
  {_,RR} = lists:foldl(fun(_,{TT1,A}) ->
    {TT2, R} = throttle:call(TT1, erlang, system_time, [millisecond]),
    {TT2, [R|A]}
  end, {TT, []}, lists:seq(1, 15)),
  Now5 = now(),

  ?assertEqual(15, length(RR)),
  ?assert(Now5 - Now4 >= 2000000),

  TT2 = throttle:new(5, 1000),
  {_,RR1} = lists:foldl(fun(_,{TT3,A}) ->
    {TT4, R} = throttle:call(TT3, fun() -> erlang:system_time(millisecond) end),
    {TT4, [R|A]}
  end, {TT2, []}, lists:seq(1, 15)),
  Now6 = now(),

  ?assertEqual(15, length(RR1)),
  ?assert(Now6 - Now5 >= 2000000),

  ok.

retry_ok_test() ->
  Now  = now(),
  Opts = #{retries => 3, retry_delay => 100},
  TT   = throttle:new(10, 1000),
  Inc  = fun(undefined) -> 1; (N) -> N+1 end,
  F    = fun() ->
    case get(count) of
      2 -> success;
      N -> put(count, Inc(N)), erlang:error(exception)
    end
  end,
  {_, R} = throttle:call(TT, F, Opts, Now),
  erase(count),
  ?assertEqual({ok, success}, R).

retry_fail_test() ->
  Now  = now(),
  Opts = #{retries => 3, retry_delay => 100},
  TT   = throttle:new(10, 1000),
  {error, {exception, _}} = throttle:call(TT, fun() -> erlang:error(exception) end, Opts, Now),
  Diff = now() - Now,
  %% Expected delay should be around 300ms
  ?assert(Diff >= 300000),
  ?assert(Diff <  400000).

retry_fail_no_block_test() ->
  Now  = now(),
  Opts = #{retries => 3, retry_delay => 0},
  TT   = throttle:new(10, 1000),
  {error, {exception, _}} = throttle:call(TT, fun() -> erlang:error(exception) end, Opts, Now),
  Diff = now() - Now,
  %% Expected delay should be around 0ms
  ?assert(Diff >   0),
  ?assert(Diff < 100).

-endif.