raterl

%% -------------------------------------------------------------------
%%
%% Copyright (c) 2016 Luis Rascão.  All Rights Reserved.
%%
%% This file is provided to you under the Apache License,
%% Version 2.0 (the "License"); you may not use this file
%% except in compliance with the License.  You may obtain
%% a copy of the License at
%%
%%   http://www.apache.org/licenses/LICENSE-2.0
%%
%% Unless required by applicable law or agreed to in writing,
%% software distributed under the License is distributed on an
%% "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
%% KIND, either express or implied.  See the License for the
%% specific language governing permissions and limitations
%% under the License.
%%
%% -------------------------------------------------------------------
%%%-------------------------------------------------------------------
%% @doc raterl queue.
%% @end
%%%-------------------------------------------------------------------

-module(raterl_queue).

-behaviour(gen_server).

%% API
-export([start_link/1,
         new/1,
         info/1,
         modify_regulator/3,
         cancel_timer/1,
         restart_timer/1,
         stop/1]).

%% Gen_server callbacks
-export([init/1,
         handle_call/3,
         handle_cast/2,
         handle_info/2,
         terminate/2,
         code_change/3]).

-define(SERVER, ?MODULE).
-define(REFRESH_TIMEOUT, 1000).

-record(state, {
          name :: atom(),
          regulator :: proplists:proplist(),
          timer_ref :: reference()
         }).

%%====================================================================
%% API functions
%%====================================================================

start_link({Name, Opts}) ->
    RegName = raterl_utils:queue_name(Name), 
    gen_server:start_link({local, RegName}, ?MODULE, [{Name, Opts}], []).

new(Args) ->
    supervisor:start_child(raterl_queue_sup, [Args]).

stop(Name) ->
    ok = gen_server:cast(raterl_utils:queue_name(Name),
                         stop).
info(Name) ->
    gen_server:call(raterl_utils:queue_name(Name),
                    info).

modify_regulator(Name, RegName, Limit) ->
    gen_server:call(raterl_utils:queue_name(Name),
                    {modify_regulator, RegName, Limit}).

cancel_timer(Name) ->
    gen_server:call(raterl_utils:queue_name(Name), cancel_timer).

restart_timer(Name) ->
    gen_server:cast(raterl_utils:queue_name(Name), restart_timer).

%%====================================================================
%% Gen_server callbacks
%%====================================================================

init([{Name, Opts}]) ->
    Regulator = proplists:get_value(regulator,
                                     Opts),
    TimerRef = init_regulator(Name, Regulator),
    {ok, #state{name = Name,
                regulator = Regulator,
                timer_ref = TimerRef}}.

handle_call(info, _From, State) ->
    {reply, State, State};
handle_call({modify_regulator, RegName, Limit}, _From,
            #state{name = Name,
                   regulator = Regulator0} = State0) ->
    Table = raterl_utils:table_name(Name),
    %% we update the counter to the new limit immediately!
    true = ets:update_element(Table, RegName, {2, Limit}),

    Regulator = lists:keyreplace(limit, 1, Regulator0,
                                 {limit, Limit}),
    State = State0#state{regulator = Regulator},
    {reply, ok, State};
handle_call(cancel_timer, _From,
            #state{timer_ref = TimerRef} = State)
    when TimerRef =/= undefined ->
    Ret = erlang:cancel_timer(TimerRef),
    {reply, Ret, State#state{timer_ref = undefined}};
handle_call(_Msg, _From, State) ->
    {reply, error, State}.

handle_cast(stop, State) ->
    {stop, normal, State};
handle_cast(restart_timer, #state{name = QueueName,
                                  regulator = Regulator,
                                  timer_ref = undefined} = State) ->
    Table = raterl_utils:table_name(QueueName),
    Name = proplists:get_value(name, Regulator),
    TimerRef = set_refresh_timer(Table, Name),
    {noreply, State#state{timer_ref = TimerRef}};
handle_cast(restart_timer, #state{name = QueueName,
                                  timer_ref = TimerRef} = State)
    when TimerRef =/= undefined ->
    %% first cancel the timer and then cast to self to restart it!
    _ = erlang:cancel_timer(TimerRef),
    restart_timer(QueueName),
    {noreply, State#state{timer_ref = undefined}};
handle_cast(_, State) ->
    {noreply, State}.

handle_info({refresh_rate_limit, Table, Name},
            #state{regulator = Regulator} = State) ->
    %% if this is a rate regulator then we need to
    %% update the limit and start another timer
    TimerRef = case proplists:get_value(type, Regulator) of
                   rate ->
                       %% update the ets counter with the configured rate
                       %% to limit+1 since we'll be rate limiting at zero
                       %% instead of at negative 1
                       Limit = proplists:get_value(limit, Regulator),
                       true = ets:update_element(Table, Name, {2, Limit + 1}),
                       set_refresh_timer(Table, Name);
                   _ ->
                       undefined
               end,
    {noreply, State#state{timer_ref = TimerRef}};
handle_info(_Msg, State) ->
    {noreply, State}.

terminate(_Reason, _State) ->
    ok.

code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%%====================================================================
%% Internal functions
%%====================================================================

init_regulator(QueueName, Opts) ->
    Name = proplists:get_value(name, Opts),
    Type = proplists:get_value(type, Opts),
    %% save the queue configuration in ets so
    %% it's easily reached by every client
    ets:insert_new(raterl, {QueueName, Type, Name}),
    init_regulator(QueueName, Type, Name, Opts).

init_regulator(QueueName, rate, Name, Opts) ->
    Limit = proplists:get_value(limit, Opts),
    %% create the ets counter that will hold
    %% the limit
    Table = raterl_utils:table_name(QueueName),
    ets:new(Table,
            [public, set, named_table,
             {write_concurrency, true}]),
    ets:insert_new(Table, {Name, Limit + 1}),
    %% set a up a recurrent timer that sets the rate
    %% counter to the limit on every second
    set_refresh_timer(Table, Name);
init_regulator(QueueName, counter, Name, Opts) ->
    Limit = proplists:get_value(limit, Opts),
    %% create the ets counter that will hold
    %% the limit
    Table = raterl_utils:table_name(QueueName),
    ets:new(Table,
            [public, set, named_table,
             {write_concurrency, true}]),
    ets:insert_new(Table, {Name, Limit + 1}),
    undefined.

set_refresh_timer(Table, Name) ->
    erlang:send_after(?REFRESH_TIMEOUT, self(),
                      {refresh_rate_limit, Table, Name}).