memstore

%% Copyright (c) 2017-present. Benoit Chesneau
%%
%% This source code is licensed under the MIT license found in the
%% LICENSE file in the root directory of this source tree.

-module(memstore_writer).
-author("benoitc").

%% API
-export([
  start_link/4,
  write_batch/2
]).


%% process callbacks
-export([
  init_writer/5,
  writer_loop/1
]).


%% system callbacks
-export([
  system_continue/3,
  system_code_change/4,
  system_terminate/4
]).

-define(TIMEOUT, 5000).


-spec start_link(Writer :: atom(), Name :: atom(), Tab :: atom(), WriterBufferSize :: non_neg_integer())
                -> {ok, pid()} |{error, term()}.
start_link(Writer, Name, Tab, WriterBufferSize) ->
  proc_lib:start_link(?MODULE, init_writer, [self(), Writer, Name, Tab, WriterBufferSize]).


write_batch(Writer, Batch) -> write_batch(Writer, Batch, ?TIMEOUT).

write_batch(Writer, Batch, Timeout) when is_list(Batch) ->
  WriterPid = whereis(Writer),
  Ref = erlang:make_ref(),
  MRef = erlang:monitor(process, WriterPid),
  WriterPid ! {write_batch, {self(), Ref}, Batch, erlang:external_size(Batch)},
  receive
    {write_batch, Ref, ok} ->
      erlang:demonitor(MRef, [flush]),
      ok;
    {'DOWN', MRef, _, _, Reason} ->
      {error, Reason}
  after Timeout ->
    erlang:demonitor(MRef, [flush]),
    {error, timeout}
  end.

init_writer(Parent, Writer, Name, Tab, WriterBufferSize) ->
  proc_lib:init_ack(Parent, {ok, self()}),
  erlang:register(Writer, self()),
  _ = ets:new(Tab, [ordered_set, named_table, public, {read_concurrency, true}]),
  writer_loop(#{ parent => Parent, name => Name, tab => Tab, max_buffer_size => WriterBufferSize}).

writer_loop(DbState = #{ parent := Parent }) ->
  receive
    {write_batch, From, Batch, BatchSize} ->
      update_loop([Batch], [From], BatchSize, DbState);
    {'EXIT', Parent, Reason} ->
      erlang:exit(Reason);
    {system, From, Request} ->
      sys:handle_system_msg(
        Request, From, Parent, ?MODULE, [],
        {writer_loop, DbState});
    Other ->
      _ = error_logger:error_msg(
        "~s: got unexpected message: ~p~n",
        [?MODULE_STRING, Other]
      ),
      exit({unexpected_message, Other})
  end.

update_loop(Batches, Writers, BufferSize, DbState = #{ max_buffer_size := Max}) when BufferSize >= Max ->
  process_ops(Batches, Writers, DbState);
update_loop(Batches, Writers, BufferSize, DbState) ->
  receive
    {write_batch, From, Batch, BatchSize} ->
      update_loop([Batch | Batches], [From | Writers], BatchSize + BufferSize, DbState)
  after 0 ->
    process_ops(Batches, Writers, DbState)
  end.

process_ops(Batches, Writers, DbState = #{ name := Name}) ->
  CurrentVersion = ets:update_counter(Name, last_seq, {2, 0}),
  process_ops(Batches, Writers, #{}, CurrentVersion + 1, DbState).

process_ops([Batch | Rest], Writers, KeyMap, Version, DbState) ->
  KeyMap2 = lists:foldl(
    fun
      ({put, Key, Value}, KeyMap1) ->
        KeyMap1#{ Key => {{Key, Version}, {ok, Value}} };
      ({delete, Key}, KeyMap1) ->
        KeyMap1#{ Key => {{Key, Version}, deleted} }
    end,
    KeyMap,
    Batch
  ),
  process_ops(Rest, Writers, KeyMap2, Version, DbState);
process_ops([], Writers, KeyMap2, _Version, DbState = #{ name := Name, tab := Tab }) ->
  _ = ets:update_counter(Name, last_seq, {2, 1}),
  Batch = maps:values(KeyMap2),
  true = ets:insert(Tab, Batch),
  _ = [reply(Writer, ok) || Writer <- Writers],
  writer_loop(DbState).

reply({Pid, Tag}, Msg) ->
  Pid ! {write_batch, Tag, Msg}.

system_continue(_, _, {writer_loop, State}) ->
  writer_loop(State).

-spec system_terminate(any(), _, _, _) -> no_return().
system_terminate(Reason, _, _, {writer_loop, State}) ->
  error_logger:info_msg(
    "sytem terminate ~p~n", [State]
  ),
  exit(Reason).

system_code_change(Misc, _, _, _) ->
  {ok, Misc}.