kafe

% @hidden
-module(kafe_consumer_srv).
-compile([{parse_transform, lager_transform}]).
-behaviour(gen_server).

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

-callback init(Args :: list()) -> {ok, any()} | ignore.
-callback consume(Offset :: integer(),
                  Key :: binary(),
                  Value :: binary()) -> ok.

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

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

-record(state, {
          group_id,
          topics = [],
          fetchers = [],
          callback = undefined,
          fetch_interval = ?DEFAULT_CONSUMER_FETCH_INTERVAL,
          fetch_pids = [],
          max_bytes = ?DEFAULT_FETCH_MAX_BYTES,
          min_bytes = ?DEFAULT_FETCH_MIN_BYTES,
          max_wait_time = ?DEFAULT_FETCH_MAX_WAIT_TIME,
          commit = ?DEFAULT_CONSUMER_COMMIT,
          from_beginning = ?DEFAULT_CONSUMER_START_FROM_BEGINNING,
          allow_unordered_commit = ?DEFAULT_CONSUMER_ALLOW_UNORDERED_COMMIT,
          commits = #{},
          on_start_fetching = ?DEFAULT_CONSUMER_ON_START_FETCHING,
          on_stop_fetching = ?DEFAULT_CONSUMER_ON_STOP_FETCHING,
          on_assignment_change = ?DEFAULT_CONSUMER_ON_ASSIGNMENT_CHANGE,
          can_fetch = ?DEFAULT_CONSUMER_CAN_FETCH,
          errors_actions = ?DEFAULT_CONSUMER_ERRORS_ACTIONS
         }).

%% API.

% @hidden
-spec start_link(atom(), map()) -> {ok, pid()}.
start_link(GroupID, Options) ->
  gen_server:start_link(?MODULE, [GroupID, Options], []).

%% gen_server.

% @hidden
init([GroupID, Options]) ->
  erlang:process_flag(trap_exit, true),
  kafe_consumer_store:insert(GroupID, server_pid, self()),
  kafe_consumer_store:insert(GroupID, can_fetch, false),
  AllowUnorderedCommit = maps:get(allow_unordered_commit, Options, ?DEFAULT_CONSUMER_ALLOW_UNORDERED_COMMIT),
  kafe_consumer_store:insert(GroupID, allow_unordered_commit, AllowUnorderedCommit),
  FetchInterval = maps:get(fetch_interval, Options, ?DEFAULT_CONSUMER_FETCH_INTERVAL),
  MaxBytes = maps:get(max_bytes, Options, ?DEFAULT_FETCH_MAX_BYTES),
  MinBytes = maps:get(min_bytes, Options, ?DEFAULT_FETCH_MIN_BYTES),
  MaxWaitTime = maps:get(max_wait_time, Options, ?DEFAULT_FETCH_MAX_WAIT_TIME),
  Commit = maps:get(commit, Options, ?DEFAULT_CONSUMER_COMMIT),
  FromBeginning = maps:get(from_beginning, Options, ?DEFAULT_CONSUMER_START_FROM_BEGINNING),
  OnStartFetching = maps:get(on_start_fetching, Options, ?DEFAULT_CONSUMER_ON_START_FETCHING),
  OnStopFetching = maps:get(on_stop_fetching, Options, ?DEFAULT_CONSUMER_ON_STOP_FETCHING),
  OnAssignmentChange = maps:get(on_assignment_change, Options, ?DEFAULT_CONSUMER_ON_ASSIGNMENT_CHANGE),
  CanFetch = maps:get(can_fetch, Options, ?DEFAULT_CONSUMER_CAN_FETCH),
  ErrorsActions = maps:get(errors_actions, Options, ?DEFAULT_CONSUMER_ERRORS_ACTIONS),
  kafe_consumer_store:insert(GroupID, can_fetch_fun, CanFetch),
  {ok, #state{
          group_id = bucs:to_binary(GroupID),
          callback = maps:get(callback, Options),
          fetch_interval = FetchInterval,
          max_bytes = MaxBytes,
          min_bytes = MinBytes,
          max_wait_time = MaxWaitTime,
          commit = Commit,
          from_beginning = FromBeginning,
          allow_unordered_commit = AllowUnorderedCommit,
          on_start_fetching = OnStartFetching,
          on_stop_fetching = OnStopFetching,
          on_assignment_change = OnAssignmentChange,
          can_fetch = CanFetch,
          errors_actions = ErrorsActions
         }}.

% @hidden
handle_call({topics, Topics}, _From, #state{topics = CurrentTopics, group_id = GroupID} = State) ->
  if
    Topics == CurrentTopics ->
      {reply, ok, State};
    true ->
      kafe_consumer_store:insert(GroupID, topics, Topics),
      {reply, ok, update_fetchers(Topics, State#state{topics = Topics})}
  end;
handle_call(start_fetch, _From, #state{group_id = GroupID, on_start_fetching = OnStartFetching} = State) ->
  case kafe_consumer_store:lookup(GroupID, can_fetch) of
    {ok, true} ->
      ok;
    _ ->
      kafe_consumer_store:insert(GroupID, can_fetch, true),
      case OnStartFetching of
        Fun when is_function(Fun, 1) ->
          _ = erlang:spawn(fun() -> erlang:apply(Fun, [GroupID]) end);
        {Module, Function} when is_atom(Module),
                                is_atom(Function) ->
          case bucs:function_exists(Module, Function, 1) of
            true ->
              _ = erlang:spawn(fun() -> erlang:apply(Module, Function, [GroupID]) end);
            _ ->
              ok
          end;
        _ ->
          ok
      end
  end,
  {reply, ok, State};
handle_call(stop_fetch, _From, #state{group_id = GroupID, on_stop_fetching = OnStopFetching} = State) ->
  case kafe_consumer_store:lookup(GroupID, can_fetch) of
    {ok, true} ->
      kafe_consumer_store:insert(GroupID, can_fetch, false),
      case OnStopFetching of
        Fun when is_function(Fun, 1) ->
          _ = erlang:spawn(fun() -> erlang:apply(Fun, [GroupID]) end);
        {Module, Function} when is_atom(Module),
                                is_atom(Function) ->
          case bucs:function_exists(Module, Function, 1) of
            true ->
              _ = erlang:spawn(fun() -> erlang:apply(Module, Function, [GroupID]) end);
            _ ->
              ok
          end;
        _ ->
          ok
      end;
    _ ->
      ok
  end,
  {reply, ok, State};
handle_call(_Request, _From, State) ->
  {reply, ignored, State}.

% @hidden
handle_cast(_Msg, State) ->
  {noreply, State}.

% @hidden
handle_info({'DOWN', MonitorRef, Type, Object, Info}, State) ->
  lager:debug("DOWN ~p, ~p, ~p, ~p", [MonitorRef, Type, Object, Info]),
  {noreply, State};
handle_info(_Info, State) ->
  {noreply, State}.

% @hidden
terminate(Reason, #state{group_id = GroupID, fetchers = Fetchers} = State) ->
  lager:debug("Will stop fetchers : ~p~nStacktrace:~s", [Reason, lager:pr_stacktrace(erlang:get_stacktrace())]),
  stop_fetchers([TP || {TP, _, _} <- Fetchers], State),
  kafe_consumer_store:delete(GroupID, server_pid),
  ok.

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

% @hidden
update_fetchers(Topics, #state{fetchers = Fetchers,
                               group_id = GroupID,
                               on_assignment_change = OnAssignmentChange} = State) ->
  CurrentFetchers = [TP || {TP, _, _} <- Fetchers],
  NewFetchers = lists:foldl(fun({Topic, Partitions}, Acc) ->
                                lists:zip(
                                  lists:duplicate(length(Partitions), Topic),
                                  Partitions) ++ Acc
                            end, [], Topics),
  FetchersToStop = CurrentFetchers -- NewFetchers,
  FetchersToSart = NewFetchers -- CurrentFetchers,
  lager:debug("CurrentFetchers = ~p", [CurrentFetchers]),
  lager:debug("NewFetchers     = ~p", [NewFetchers]),
  lager:debug("Stop            = ~p", [FetchersToStop]),
  lager:debug("Start           = ~p", [FetchersToSart]),
  case OnAssignmentChange of
    Fun when is_function(Fun, 3) ->
      _ = erlang:spawn(fun() -> erlang:apply(Fun, [GroupID, FetchersToStop, FetchersToSart]) end);
    {Module, Function} when is_atom(Module),
                            is_atom(Function) ->
      case bucs:function_exists(Module, Function, 3) of
        true ->
          _ = erlang:spawn(fun() -> erlang:apply(Module, Function, [GroupID, FetchersToStop, FetchersToSart]) end);
        _ ->
          ok
      end;
    _ ->
      ok
  end,
  State1 = stop_fetchers(FetchersToStop, State),
  start_fetchers(FetchersToSart, State1).

stop_fetchers([], State) ->
  State;
stop_fetchers([TP|Rest], #state{fetchers = Fetchers, commits = Commits, group_id = GroupID} = State) ->
  case lists:keyfind(TP, 1, Fetchers) of
    {{Topic, Partition} = TP, Pid, MRef} ->
      kafe_metrics:delete_consumer_partition(GroupID, Topic, Partition),
      CommitStoreKey = erlang:term_to_binary(TP),
      _ = erlang:demonitor(MRef),
      try
        kafe_consumer_group_sup:stop_child(Pid)
      catch
        C:E ->
          lager:error("Can't terminate kafe_consumer_fetcher #~p: ~p:~p", [Pid, C, E])
      end,
      stop_fetchers(Rest, State#state{fetchers = lists:keydelete(TP, 1, Fetchers),
                                      commits = maps:remove(CommitStoreKey, Commits)});
    false ->
      stop_fetchers(Rest, State)
  end.

start_fetchers([], State) ->
  State;
start_fetchers([{Topic, Partition}|Rest], #state{fetchers = Fetchers,
                                                 fetch_interval = FetchInterval,
                                                 group_id = GroupID,
                                                 commit = Commit,
                                                 from_beginning = FromBeginning,
                                                 min_bytes = MinBytes,
                                                 max_bytes = MaxBytes,
                                                 max_wait_time = MaxWaitTime,
                                                 callback = Callback,
                                                 errors_actions = ErrorsActions} = State) ->
  kafe_metrics:init_consumer_partition(GroupID, Topic, Partition),
  case kafe_consumer_group_sup:start_child(Topic, Partition, FetchInterval,
                                           GroupID, Commit, FromBeginning,
                                           MinBytes, MaxBytes, MaxWaitTime,
                                           ErrorsActions, Callback) of
    {ok, Pid} ->
      MRef = erlang:monitor(process, Pid),
      start_fetchers(Rest, State#state{fetchers = [{{Topic, Partition}, Pid, MRef}|Fetchers]});
    {error, Error} ->
      lager:error("Faild to start fetcher for topic ~s, partition ~p: ~p", [Topic, Partition, Error]),
      start_fetchers(Rest, State)
  end.

-ifdef(TEST).
start_fetch_without_fun_test() ->
  kafe_consumer_store:new(<<"test_cg">>),
  ?assertNot(true == kafe_consumer_store:value(<<"test_cg">>, can_fetch)),
  State = #state{group_id = <<"test_cg">>, on_start_fetching = undefined},
  ?assertEqual({reply, ok, State}, handle_call(start_fetch, from, State)),
  ?assert(kafe_consumer_store:value(<<"test_cg">>, can_fetch)),
  kafe_consumer_store:delete(<<"test_cg">>).

start_fetch_with_fun_test() ->
  kafe_consumer_store:new(<<"test_cg">>),
  ?assertNot(true == kafe_consumer_store:value(<<"test_cg">>, can_fetch)),
  State = #state{group_id = <<"test_cg">>,
                 on_start_fetching = fun(G) ->
                                         ?assertEqual(<<"test_cg">>, G)
                                     end},
  ?assertEqual({reply, ok, State}, handle_call(start_fetch, from, State)),
  ?assert(kafe_consumer_store:value(<<"test_cg">>, can_fetch)),
  kafe_consumer_store:delete(<<"test_cg">>).

start_fetch_with_invalid_fun_test() ->
  kafe_consumer_store:new(<<"test_cg">>),
  ?assertNot(true == kafe_consumer_store:value(<<"test_cg">>, can_fetch)),
  State = #state{group_id = <<"test_cg">>,
                 on_start_fetching = fun(_, _) ->
                                         ?assert(false)
                                     end},
  ?assertEqual({reply, ok, State}, handle_call(start_fetch, from, State)),
  ?assert(kafe_consumer_store:value(<<"test_cg">>, can_fetch)),
  kafe_consumer_store:delete(<<"test_cg">>).

stop_fetch_without_fun_test() ->
  kafe_consumer_store:new(<<"test_cg">>),
  kafe_consumer_store:insert(<<"test_cg">>, can_fetch, true),
  State = #state{group_id = <<"test_cg">>,
                 on_stop_fetching = undefined},
  ?assertEqual({reply, ok, State}, handle_call(stop_fetch, from, State)),
  ?assertNot(kafe_consumer_store:value(<<"test_cg">>, can_fetch)),
  kafe_consumer_store:delete(<<"test_cg">>).

stop_fetch_with_fun_test() ->
  kafe_consumer_store:new(<<"test_cg">>),
  kafe_consumer_store:insert(<<"test_cg">>, can_fetch, true),
  State = #state{group_id = <<"test_cg">>,
                 on_stop_fetching = fun(G) ->
                                        ?assertEqual(<<"test_cg">>, G)
                                    end},
  ?assertEqual({reply, ok, State}, handle_call(stop_fetch, from, State)),
  ?assertNot(kafe_consumer_store:value(<<"test_cg">>, can_fetch)),
  kafe_consumer_store:delete(<<"test_cg">>).

stop_fetch_with_invalid_fun_test() ->
  kafe_consumer_store:new(<<"test_cg">>),
  kafe_consumer_store:insert(<<"test_cg">>, can_fetch, true),
  State = #state{group_id = <<"test_cg">>,
                 on_stop_fetching = fun(_, _) ->
                                        ?assert(false)
                                    end},
  ?assertEqual({reply, ok, State}, handle_call(stop_fetch, from, State)),
  ?assertNot(kafe_consumer_store:value(<<"test_cg">>, can_fetch)),
  kafe_consumer_store:delete(<<"test_cg">>).

update_fetchers_create_test() ->
  meck:new(kafe_consumer_group_sup, [passthrough]),
  meck:expect(kafe_consumer_group_sup, start_child, 11, {ok, c:pid(0, 0, 0)}),
  meck:expect(kafe_consumer_group_sup, stop_child, 1, ok),
  meck:new(kafe_metrics, [passthrough]),
  meck:expect(kafe_metrics, delete_consumer_partition, 3, ok),
  meck:expect(kafe_metrics, init_consumer_partition, 3, ok),
  kafe_consumer_store:new(<<"test_cg">>),
  State = #state{group_id = <<"test_cg">>},
  Topics = [{<<"topic0">>, [0, 1, 2]}, {<<"topic1">>, [0, 1]}],
  ?assertMatch({reply, ok,
                #state{group_id = <<"test_cg">>,
                       topics = Topics,
                       fetchers = [{{<<"topic0">>, 2}, _, _},
                                   {{<<"topic0">>, 1}, _, _},
                                   {{<<"topic0">>, 0}, _, _},
                                   {{<<"topic1">>, 1}, _, _},
                                   {{<<"topic1">>, 0}, _, _}]}},
               handle_call({topics, Topics}, from, State)),
  ?assertEqual([{<<"topic0">>, 0},
                {<<"topic0">>, 1},
                {<<"topic0">>, 2},
                {<<"topic1">>, 0},
                {<<"topic1">>, 1}],
               kafe_consumer:topics(<<"test_cg">>)),
  kafe_consumer_store:delete(<<"test_cg">>),
  meck:unload(kafe_metrics),
  meck:unload(kafe_consumer_group_sup).

update_fetchers_unchange_test() ->
  meck:new(kafe_consumer_group_sup, [passthrough]),
  meck:expect(kafe_consumer_group_sup, start_child, 11, {ok, c:pid(0, 0, 0)}),
  meck:expect(kafe_consumer_group_sup, stop_child, 1, ok),
%  meck:new(kafe_metrics, [passthrough]),
%  meck:expect(kafe_metrics, delete_consumer_partition, 3, ok),
%  meck:expect(kafe_metrics, init_consumer_partition, 3, ok),
  kafe_consumer_store:new(<<"test_cg">>),
  Topics = [{<<"topic0">>, [0, 1, 2]}, {<<"topic1">>, [0, 1]}],
  kafe_consumer_store:insert(<<"test_cg">>, topics, Topics),
  State = #state{group_id = <<"test_cg">>, topics = Topics},
  ?assertMatch({reply, ok, State},
               handle_call({topics, Topics}, from, State)),
  ?assertEqual([{<<"topic0">>, 0},
                {<<"topic0">>, 1},
                {<<"topic0">>, 2},
                {<<"topic1">>, 0},
                {<<"topic1">>, 1}],
               kafe_consumer:topics(<<"test_cg">>)),
  kafe_consumer_store:delete(<<"test_cg">>),
%  meck:unload(kafe_metrics),
  meck:unload(kafe_consumer_group_sup).

update_fetchers_add_test() ->
  meck:new(kafe_consumer_group_sup, [passthrough]),
  meck:expect(kafe_consumer_group_sup, start_child, 11, {ok, c:pid(0, 0, 0)}),
  meck:expect(kafe_consumer_group_sup, stop_child, 1, ok),
  meck:new(kafe_metrics, [passthrough]),
%  meck:expect(kafe_metrics, delete_consumer_partition, 3, ok),
  meck:expect(kafe_metrics, init_consumer_partition, 3, ok),
  kafe_consumer_store:new(<<"test_cg">>),
  Topics = [{<<"topic0">>, [0, 1, 2]}, {<<"topic1">>, [0, 1]}],
  kafe_consumer_store:insert(<<"test_cg">>, topics, Topics),
  State = #state{group_id = <<"test_cg">>, topics = Topics},
  NewTopics = [{<<"topic0">>, [0, 1, 2]}, {<<"topic1">>, [0, 1, 2]}, {<<"topic2">>, [0, 1, 2]}],
  ?assertMatch({reply, ok,
                #state{group_id = <<"test_cg">>,
                       topics = NewTopics,
                       fetchers = [{{<<"topic0">>, 2}, _, _},
                                   {{<<"topic0">>, 1}, _, _},
                                   {{<<"topic0">>, 0}, _, _},
                                   {{<<"topic1">>, 2}, _, _},
                                   {{<<"topic1">>, 1}, _, _},
                                   {{<<"topic1">>, 0}, _, _},
                                   {{<<"topic2">>, 2}, _, _},
                                   {{<<"topic2">>, 1}, _, _},
                                   {{<<"topic2">>, 0}, _, _}]}},
               handle_call({topics, NewTopics}, from, State)),
  ?assertEqual([{<<"topic0">>, 0},
                {<<"topic0">>, 1},
                {<<"topic0">>, 2},
                {<<"topic1">>, 0},
                {<<"topic1">>, 1},
                {<<"topic1">>, 2},
                {<<"topic2">>, 0},
                {<<"topic2">>, 1},
                {<<"topic2">>, 2}],
               kafe_consumer:topics(<<"test_cg">>)),
  kafe_consumer_store:delete(<<"test_cg">>),
  meck:unload(kafe_metrics),
  meck:unload(kafe_consumer_group_sup).

update_fetchers_delete_test() ->
  meck:new(kafe_consumer_group_sup, [passthrough]),
  meck:expect(kafe_consumer_group_sup, start_child, 11, {ok, c:pid(0, 0, 0)}),
  meck:expect(kafe_consumer_group_sup, stop_child, 1, ok),
  meck:new(kafe_metrics, [passthrough]),
  meck:expect(kafe_metrics, delete_consumer_partition, 3, ok),
  meck:expect(kafe_metrics, init_consumer_partition, 3, ok),
  kafe_consumer_store:new(<<"test_cg">>),
  Topics = [{<<"topic0">>, [0, 1, 2]}, {<<"topic1">>, [0, 1, 2]}, {<<"topic2">>, [0, 1, 2]}],
  kafe_consumer_store:insert(<<"test_cg">>, topics, Topics),
  State = #state{group_id = <<"test_cg">>, topics = Topics},
  NewTopics = [{<<"topic0">>, [0, 1, 2]}, {<<"topic1">>, [0, 1]}],
  ?assertMatch({reply, ok,
                #state{group_id = <<"test_cg">>,
                       topics = NewTopics,
                       fetchers = [{{<<"topic0">>, 2}, _, _},
                                   {{<<"topic0">>, 1}, _, _},
                                   {{<<"topic0">>, 0}, _, _},
                                   {{<<"topic1">>, 1}, _, _},
                                   {{<<"topic1">>, 0}, _, _}]}},
               handle_call({topics, NewTopics}, from, State)),
  ?assertEqual([{<<"topic0">>, 0},
                {<<"topic0">>, 1},
                {<<"topic0">>, 2},
                {<<"topic1">>, 0},
                {<<"topic1">>, 1}],
               kafe_consumer:topics(<<"test_cg">>)),
  kafe_consumer_store:delete(<<"test_cg">>),
  meck:unload(kafe_metrics),
  meck:unload(kafe_consumer_group_sup).

update_fetchers_update_test() ->
  meck:new(kafe_consumer_group_sup, [passthrough]),
  meck:expect(kafe_consumer_group_sup, start_child, 11, {ok, c:pid(0, 0, 0)}),
  meck:expect(kafe_consumer_group_sup, stop_child, 1, ok),
  meck:new(kafe_metrics, [passthrough]),
  meck:expect(kafe_metrics, delete_consumer_partition, 3, ok),
  meck:expect(kafe_metrics, init_consumer_partition, 3, ok),
  kafe_consumer_store:new(<<"test_cg">>),
  Topics = [{<<"topic0">>, [0, 1, 2]}, {<<"topic1">>, [0, 1]}],
  kafe_consumer_store:insert(<<"test_cg">>, topics, Topics),
  State = #state{group_id = <<"test_cg">>, topics = Topics},
  NewTopics = [{<<"topic1">>, [0, 1, 2]}, {<<"topic2">>, [0, 1, 2]}],
  ?assertMatch({reply, ok,
                #state{group_id = <<"test_cg">>,
                       topics = NewTopics,
                       fetchers = [{{<<"topic1">>, 2}, _, _},
                                   {{<<"topic1">>, 1}, _, _},
                                   {{<<"topic1">>, 0}, _, _},
                                   {{<<"topic2">>, 2}, _, _},
                                   {{<<"topic2">>, 1}, _, _},
                                   {{<<"topic2">>, 0}, _, _}]}},
               handle_call({topics, NewTopics}, from, State)),
  ?assertEqual([{<<"topic1">>, 0},
                {<<"topic1">>, 1},
                {<<"topic1">>, 2},
                {<<"topic2">>, 0},
                {<<"topic2">>, 1},
                {<<"topic2">>, 2}],
               kafe_consumer:topics(<<"test_cg">>)),
  kafe_consumer_store:delete(<<"test_cg">>),
  meck:unload(kafe_metrics),
  meck:unload(kafe_consumer_group_sup).
-endif.