elib

%%% @doc

%%% 将数据库写入分摊，减少瞬时IO,将IO平摊

%%% @end

-module(slow_db).
-include_lib("kernel/include/logger.hrl").
-behaviour(gen_server).

%% API

-export([

  start_link/1,

  start_link/2,

  start_link/3,

  put/1,

  put/2,

  remove/1,

  remove/2

]).

%% gen_server callbacks

-export([

  init/1,

  handle_call/3,

  handle_cast/2,

  handle_info/2,

  terminate/2,

  code_change/3

]).


-type flush_config() :: #{

flush_interval => integer(),% 刷数据库时间间隔

flush_once => integer(), % 每次刷数据库条数

flush_error_drop => boolean() % 单条保存失败后是否取消保存

}.

-record(state, {

  name,

  queue = queue:new(),

  mark_map = #{},

  callback,

  flush_config :: flush_config(),

  flushing = false

}).

start_link(CB) ->
  start_link(?MODULE, CB, #{}).

start_link(CB, FlushCfg) ->
  start_link(?MODULE, CB, FlushCfg).

start_link(Name, CB, FlushCfg) ->
  gen_server:start_link({local, Name}, ?MODULE, [Name, CB, FlushCfg], []).

put(K) ->
  ?MODULE:put(?MODULE, K).

put(Name, K) ->
  gen_server:call(Name, {put, K}, 5000).

remove(K) ->
  remove(?MODULE, K).

remove(Name, K) ->
  gen_server:call(Name, {remove, K}, 5000).

init([Name, CB, FlushCfg0]) ->
  FlushCfg = maps:merge(default_flush_cfg(), FlushCfg0),
  {ok, #state{name = Name, queue = queue:new(), mark_map = #{}, callback = CB, flush_config = FlushCfg}}.

handle_call({put, K}, _From, #state{mark_map = M, queue = Q, flushing = Flushing, flush_config = #{flush_interval:= FI}} = S) ->
  {Q2, M2} = case maps:is_key(K, M) of
               true -> {Q, M};
               false -> {queue:in_r(K, Q), maps:put(K, 1, M)}
             end,
  S2 = S#state{mark_map = M2, queue = Q2},
  case Flushing of
    true -> pass;
    false -> erlang:send_after(FI, self(), flush)
  end,
  {reply, ok, S2#state{flushing = true}};
handle_call({remove, K}, _From, #state{mark_map = M} = State) ->
  M2 = maps:remove(K, M),
  {reply, ok, State#state{mark_map = M2}};
handle_call(_Request, _From, State) ->
  {reply, ok, State}.

handle_cast(_Request, State) ->
  {noreply, State}.

handle_info(flush, #state{queue = Q} = S) ->
  case queue:is_empty(Q) of
    true ->
      {noreply, S#state{flushing = false}};
    false ->
      Start = unix_time(),
      S2 = flush_i(S),
      End = unix_time(),
      #state{queue = Q2, flush_config = #{flush_interval:= FI}} = S2,
      case queue:is_empty(Q2) of
        true ->
          {noreply, S2#state{flushing = false}};
        false ->
          Cost = End - Start,
          case Cost > FI of
            true ->
              % flush once

              erlang:send(self(), flush);
            false ->
              erlang:send_after(FI - Cost, self(), flush)
          end,
          {noreply, S2#state{flushing = true}}
      end
  end;
handle_info(_Info, State) ->
  {noreply, State}.

terminate(_Reason, _State) ->
  ok.

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

%%%===================================================================

%%% Internal functions

%%%===================================================================


default_flush_cfg() ->
  % 1 sec flush 100 records

  #{flush_interval => 1000, flush_once => 100, flush_error_drop => false}.

unix_time() ->
  os:perf_counter(millisecond).

flush_i(#state{name = Name, queue = Q, flush_config = #{flush_once:= FO, flush_error_drop:= Drop}} = S) ->
  {Q2, Q3} = queue:split(FO, Q),
  L = queue:to_list(Q2),
  {Total, Fail, FailRec, S2} = flush_ii(L, S, 0, 0, []),
  ?LOG_DEBUG("~p flush total num:~p,fail num:~p", [Name, Total, Fail]),
  Q4 = case Drop of
         true -> Q3;
         false -> lists:foldl(fun(K, Acc) -> queue:in(K, Acc) end, Q3, FailRec)
       end,
  S2#state{queue = Q4}.

flush_ii([], State, Fail, Total, Acc) ->
  {Total, Fail, Acc, State};
flush_ii([K | L], #state{name = Name, mark_map = Map, callback = CB, flush_config = #{flush_error_drop:= Drop}} = S, Fail, Total, Acc) ->
  case maps:is_key(K, Map) of
    false ->
      flush_ii(L, S, Fail, Total, Acc);
    true ->
      CBReply = case CB of
                  {M, F} -> M:F(K);
                  _ when is_function(CB, 1) -> CB(K)
                end,
      case CBReply of
        ok ->
          Map2 = maps:remove(K, Map),
          flush_ii(L, S#state{mark_map = Map2}, Fail, Total + 1, Acc);
        {error, Reason} ->
          ?LOG_WARNING("~p flush ~p failed with reason ~p", [Name, K, Reason]),
          Map2 = case Drop of
                   true -> maps:remove(K, Map);
                   false -> Map
                 end,
          flush_ii(L, S#state{mark_map = Map2}, Fail + 1, Total + 1, Acc)
      end
  end.