erlperf

%%%-------------------------------------------------------------------
%%% @author Maxim Fedorov <maximfca@gmail.com>
%%% @copyright (C) 2020, Maxim Fedorov
%%% @doc
%%%   Rate-limited load generator.
%%% @end
-module(ep_load).
-author("maximfca@gmail.com").

%% API
-export([
    generate/3,
    stop/1,
    start_link/0,
    spawn_link/3,
    init/1,
    start/3
]).

-behaviour(supervisor).

-spec start_link() -> supervisor:startlink_ret().
start_link() ->
    supervisor:start_link({local, ?MODULE}, ?MODULE, []).

-spec init([]) -> {ok, {supervisor:sup_flags(), [supervisor:child_spec()]}}.
init([]) ->
    SupFlags = #{strategy => simple_one_for_one, intensity => 10, period => 60},
    ChildSpecs = [#{id => ?MODULE, start => {?MODULE, spawn_link, []}}],
    {ok, {SupFlags, ChildSpecs}}.

start(Fun, Count, PeriodMs) ->
    supervisor:start_child(?MODULE, [Fun, Count, PeriodMs]).

spawn_link(Fun, Count, PeriodMs) ->
    {ok, proc_lib:spawn_link(?MODULE, generate, [Fun, Count, PeriodMs])}.

generate(Fun, Count, PeriodMs) ->
    process_flag(priority, high),
    generate_impl(Fun, Count, PeriodMs, erlang:monotonic_time(millisecond) + PeriodMs).

stop(Pid) ->
    unlink(Pid),
    exit(Pid, kill).

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

generate_impl(Fun, Count, _Period, _NextTick) when not is_integer(Count); Count =< 0; not is_function(Fun, 0) ->
    error(badarg);
generate_impl(Fun, Count, Period, NextTick) ->
    case erlang:monotonic_time(millisecond) of
        Now when Now < NextTick ->
            erlang:send_after(NextTick, self(), timeout, [{abs, true}]),
            fun_impl(Fun, Count),
            receive
                timeout ->
                    generate_impl(Fun, Count, Period, NextTick + Period)
            end;
        _Now ->
            %% falling behind: no waiting
            fun_impl(Fun, Count),
            generate_impl(Fun, Count, Period, NextTick + Period)
    end.

fun_impl(Fun, 1) ->
    Fun();
fun_impl(Fun, Count) ->
    Fun(),
    fun_impl(Fun, Count - 1).