%%% @copyright (C) 2016, AdRoll
%%% @doc
%%%
%%%     Generic record processor with support for batching and checkpointing.
%%%
%%%     We use an underlying 'erlmld_flusher' module to maintain the current
%%%     batch. When the current batch is full (according to the flusher's
%%%     definition of fullness), or when 'flush_interval_ms' milliseconds have
%%%     elapsed, we ask the flusher to flush the current batch and start a new
%%%     one.
%%%
%%%     We checkpoint every 'checkpoint_interval_ms' milliseconds. After
%%%     'watchdog_timeout_ms' milliseconds with no activity (after any initial
%%%     activity), the current process is killed. Checkpointing works by
%%%     keeping track of the latest record for which itself and all
%%%     predecessors have been flushed. For example, suppose 'process_record'
%%%     has been called with the following records, and the flusher has already
%%%     flushed the ones marked with '*':
%%%
%%%         R1 R2 R3 R4 R5 R6 R7 R8 R9 R10
%%%         *  *  *     *  *     *
%%%
%%%     Then the stream will be checkpointed at (the sequence number of) R3.
%%%     For the internal mechanics, see the comment in 'note_success'.
%%%
%%% @end
%%% Created : 19 Dec 2016 by Constantin Berzan <constantin.berzan@adroll.com>
%%% based on original code by Mike Watters <mike.watters@adroll.com>
-module(erlmld_batch_processor).

-behaviour(erlmld_worker).

-export([initialize/3, ready/1, process_record/2, checkpointed/3, shutdown/2]).

-include("erlmld.hrl").

-record(state,
        {%% Handler module implementing the flusher behavior.
         flusher_mod :: atom(),
         %% Flusher state (opaque to this module) passed to flusher_mod callbacks.
         flusher_state :: flusher_state(),
         %% Optional callback to call each time process_record returns a checkpoint.
         on_checkpoint :: fun((term(), shard_id()) -> term()),
         %% Optional, false by default. Tells whether to log or not every successful
         %% checkpoint from the KCL worker.
         log_checkpoints :: boolean(),
         description :: term(),
         shard_id :: shard_id(),
         count = 0, % non-ignored records seen
         checkpoints = 0, % checkpoints requested
         next_counter_checkpoint = 0, % counter value at which we can next checkpoint
         last_flush_time = os:timestamp(),
         last_checkpoint_time = os:timestamp(),
         checkpointable = gb_trees:empty(), % {Count, SequenceNo}
         flush_interval_ms,
         checkpoint_interval_ms,
         watchdog_timeout_ms,
         watchdog,
         enable_subsequence_checkpoints = false :: boolean()}).

%%%===================================================================
%%% API
%%%===================================================================

initialize(Opts, ShardId, ISN) ->
    Defaults =
        #{on_checkpoint => fun(_, _) -> ok end,
          log_checkpoints => false,
          description => undefined,
          enable_subsequence_checkpoints => false},
    #{description := Description,
      flusher_mod := FlusherMod,
      flusher_mod_data := FlusherModData,
      on_checkpoint := OnCheckpoint,
      log_checkpoints := LogCheckpoints,
      flush_interval_ms := FlushIntervalMs,
      checkpoint_interval_ms := CheckpointIntervalMs,
      watchdog_timeout_ms := WatchdogTimeoutMs,
      enable_subsequence_checkpoints := EnableSubCP} =
        maps:merge(Defaults, Opts),
    State =
        #state{flusher_mod = FlusherMod,
               flusher_state = erlmld_flusher:init(FlusherMod, ShardId, FlusherModData),
               on_checkpoint = OnCheckpoint,
               log_checkpoints = LogCheckpoints,
               description = Description,
               shard_id = ShardId,
               flush_interval_ms = FlushIntervalMs,
               checkpoint_interval_ms = CheckpointIntervalMs,
               watchdog_timeout_ms = WatchdogTimeoutMs,
               enable_subsequence_checkpoints = EnableSubCP},
    error_logger:info_msg("~p initialized for shard {~p,~p} at ~p~n",
                          [self(), Description, ShardId, ISN]),
    {ok, update_watchdog(State)}.

ready(#state{flusher_mod = FMod, flusher_state = FState} = State) ->
    {ok, NFState, Tokens} = erlmld_flusher:heartbeat(FMod, FState),
    NState = flusher_state(State, NFState),
    NNState =
        case Tokens of
            [] ->
                NState;
            _ ->
                note_success(note_flush(NState), Tokens)
        end,
    maybe_checkpoint(update_watchdog(NNState)).

process_record(#state{last_flush_time = LastFlush, flush_interval_ms = FlushInterval} =
                   State,
               Record) ->
    {ok, NState} =
        case {add_record(State, Record), elapsed_ms(LastFlush)} of
            {{ok, State1}, E} when E >= FlushInterval ->
                {ok, flush(State1)};
            {{ok, State1}, _} ->
                {ok, State1};
            {{error, full}, _} ->
                add_record(flush(State), Record)
        end,
    maybe_checkpoint(update_watchdog(NState)).

checkpointed(#state{log_checkpoints = LogCheckpoints} = State,
             SequenceNumber,
             Checkpoint) ->
    case LogCheckpoints of
        true ->
            error_logger:info_msg("~p checkpointed at ~p (~p)~n",
                                  [State, Checkpoint, SequenceNumber]);
        false ->
            ok
    end,
    {ok, State}.

shutdown(#state{description = Description,
                shard_id = ShardId,
                count = Count} =
             State,
         Reason) ->
    error_logger:info_msg("{~p,~p} (~p) shutting down, reason: ~p~n",
                          [Description, ShardId, Count, Reason]),
    case Reason of
        zombie ->
            %% we lost our lease, nothing else to do now.
            ok;
        terminate ->
            %% the shard is closing (e.g., as a result of a merge or split).  we should
            %% flush all outstanding data and checkpoint.
            NState = flush_full(State),
            case next_checkpoint(NState) of
                %% we flushed all items and are able to checkpoint.  there shouldn't be
                %% any more items left to checkpoint:
                {CPState, CPSN} ->
                    error_logger:info_msg("{~p,~p} will checkpoint at ~p~n",
                                          [Description, ShardId, CPSN]),
                    Remaining = checkpointable(CPState),
                    Remaining = gb_trees:empty(),
                    {ok, #checkpoint{sequence_number = CPSN}};
                %% nothing was checkpointable.  there shouldn't be any items left to
                %% checkpoint:
                undefined ->
                    error_logger:info_msg("{~p,~p} has nothing to checkpoint, will use latest~n",
                                          [Description, ShardId]),
                    Remaining = checkpointable(NState),
                    Remaining = gb_trees:empty(),
                    %% we checkpoint with an empty sequence number.  this will cause the
                    %% KCL to checkpoint at the most recent record it has seen on the
                    %% shard:
                    {ok, #checkpoint{}}
            end
    end.

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

incr_count(#state{count = Count} = State) ->
    State#state{count = Count + 1}.

incr_checkpoints(#state{checkpoints = Count} = State) ->
    State#state{checkpoints = Count + 1}.

checkpointable(#state{checkpointable = CPT}) ->
    CPT.

checkpointable(State, CPT) ->
    State#state{checkpointable = CPT}.

next_counter_checkpoint(State, N) ->
    State#state{next_counter_checkpoint = N}.

flusher_state(State, FState) ->
    State#state{flusher_state = FState}.

add_record(#state{count = Count,
                  flusher_mod = FMod,
                  flusher_state = FState} =
               State,
           #stream_record{sequence_number = SN} = Record) ->
    case erlmld_flusher:add_record(FMod, FState, Record, {Count, SN}) of
        {ok, NFState} ->
            {ok, incr_count(flusher_state(State, NFState))};
        {ignored, NFState} ->
            {ok, flusher_state(State, NFState)};
        {error, full} ->
            {error, full}
    end.

flush(State) ->
    flush(State, partial).

flush_full(State) ->
    flush(State, full).

flush(#state{flusher_mod = FMod, flusher_state = FState} = State, Kind) ->
    {ok, NFState, Tokens} = erlmld_flusher:flush(FMod, FState, Kind),
    NState = flusher_state(State, NFState),
    note_success(note_flush(NState), Tokens).

note_success(#state{checkpointable = CPT} = State, Tokens) ->
    %% The items with the given Tokens have been successfully processed.
    %% Update the 'checkpointable' gb_tree, which we are using as a priority
    %% queue. When the time comes to checkpoint, we pop items until the next
    %% smallest item has a discontiguous id, then checkpoint at the sequence
    %% number corresponding to the most recently-popped item. This way, we can
    %% do time-based checkpointing without forcing a flush of the current
    %% pending batch.
    %%
    %% fixme; when adding a new element, also pop all smaller elements which have
    %% contiguous keys with the new element (only need to track gaps).
    NCPT =
        lists:foldl(fun({Counter, SeqNum}, Acc) -> gb_trees:insert(Counter, SeqNum, Acc) end,
                    CPT,
                    Tokens),
    checkpointable(State, NCPT).

maybe_checkpoint(#state{last_checkpoint_time = LastCheckpoint,
                        checkpoint_interval_ms = CheckpointInterval} =
                     State) ->
    case elapsed_ms(LastCheckpoint) of
        N when N >= CheckpointInterval ->
            case next_checkpoint(State) of
                {NState, SequenceNumber} ->
                    {ok, note_checkpoint(NState), #checkpoint{sequence_number = SequenceNumber}};
                undefined ->
                    {ok, State}
            end;
        _ ->
            {ok, State}
    end.

%% given a state, return undefined if nothing is ready to be checkpointed.  otherwise,
%% return {NState, SeqNo}, where SeqNo is a #sequence_number{} at which we can checkpoint
%% (i.e., all items below that SN have been successfully processed), and NState represents
%% the state after checkpointing.
next_checkpoint(State) ->
    CPT = checkpointable(State),
    case gb_trees:is_empty(CPT) of
        true ->
            undefined;
        false ->
            next_checkpoint(State, gb_trees:take_smallest(CPT), undefined)
    end.

next_checkpoint(State, {SmallestCount, _, _} = SmallestItem, LatestFinished) ->
    next_checkpoint(State, SmallestCount, SmallestItem, LatestFinished).

%% in order to checkpoint, the smallest tree element must have a key whose value is equal
%% to the `next_counter_checkpoint` field; otherwise, we can't checkpoint yet (we have an
%% initial gap).
next_checkpoint(#state{next_counter_checkpoint = N}, FirstSmallestCount, _, _)
    when FirstSmallestCount > N ->
    undefined;
%% if subsequence checkpointing is disabled, we can only checkpoint at a record which has
%% been fully processed.  while advancing through contiguous items, we keep track of the
%% highest-numbered item corresponding to a completely-processed record (which may be
%% undefined).
next_checkpoint(#state{enable_subsequence_checkpoints = EnableSubCP} = State,
                FirstSmallestCount,
                {SmallestCount, SmallestSN, CPT},
                LatestFinished) ->
    %% can checkpoint at SmallestSN if subsequence checkpointing is enabled, or if it
    %% corresponds to the last subrecord in an aggregate record, or if it corresponds to a
    %% non-aggregate record:
    CanCheckpoint = EnableSubCP orelse not is_sub_record(SmallestSN),
    %% if we can checkpoint at the current SN, use it as the latest 'finished' value:
    NLatestFinished =
        case CanCheckpoint of
            true ->
                {next_counter_checkpoint(checkpointable(State, CPT), SmallestCount + 1),
                 SmallestSN};
            false ->
                LatestFinished
        end,
    case gb_trees:is_empty(CPT) of
        true ->
            NLatestFinished;
        false ->
            NextContiguous = SmallestCount + 1,
            case gb_trees:take_smallest(CPT) of
                {NextContiguous, _, _} = Next ->
                    next_checkpoint(State, FirstSmallestCount, Next, NLatestFinished);
                _ ->
                    NLatestFinished
            end
    end.

note_checkpoint(#state{description = Description,
                       shard_id = ShardId,
                       on_checkpoint = OnCheckpoint} =
                    State) ->
    OnCheckpoint(Description, ShardId),
    incr_checkpoints(State#state{last_checkpoint_time = os:timestamp()}).

note_flush(State) ->
    State#state{last_flush_time = os:timestamp()}.

%% given an erlang timestamp, return the elapsed duration in milliseconds.
elapsed_ms(When) ->
    timer:now_diff(
        os:timestamp(), When)
    div 1000.

%% start a watchdog timer, cancelling any which is outstanding.  if the timer fires, it
%% will result in the current process exiting with a reason of 'watchdog_timeout'
update_watchdog(#state{watchdog_timeout_ms = undefined} = State) ->
    State;
update_watchdog(#state{watchdog_timeout_ms = WatchdogTimeout, watchdog = Watchdog} =
                    State) ->
    case Watchdog of
        undefined ->
            ok;
        _ ->
            timer:cancel(Watchdog)
    end,
    {ok, Ref} = timer:exit_after(WatchdogTimeout, watchdog_timeout),
    State#state{watchdog = Ref}.

%% We use user_sub here instead of sub: for KPL records, those values will be
%% the same. For our own KPL-like protocol user_sub will be filled in, and
%% in other use cases it will be undefined. In both of these last cases sub
%% will be the original KCL sent (0, possible) so it can't be really used to
%% know if this is a subrecord or not.
is_sub_record(#sequence_number{user_sub = Sub, user_total = Total})
    when is_integer(Sub) andalso is_integer(Total) andalso Sub < Total - 1 ->
    true;
is_sub_record(_) ->
    false.

%%%===================================================================
%%% TESTS
%%%===================================================================

-ifdef(TEST).

-include_lib("eunit/include/eunit.hrl").

equal_cpt(A, B) ->
    maps:from_list(
        gb_trees:to_list(checkpointable(A)))
    == maps:from_list(
           gb_trees:to_list(checkpointable(B))).

checkpointing_test() ->
    State = #state{},
    ?assertEqual(undefined, next_checkpoint(State)),

    %% items 0, 1, and 3 completed and checkpointable:
    CPT0 =
        gb_trees:insert(0, 0, gb_trees:insert(1, 1, gb_trees:insert(3, 3, gb_trees:empty()))),
    Expected1 = checkpointable(State, CPT0),
    State1 = note_success(State, [{0, 0}, {1, 1}, {3, 3}]),
    ?assert(equal_cpt(Expected1, State1)),
    %% the initial checkpoint should only cover items 0-1 due to the gap (2 not
    %% checkpointable yet), so item 3 should still be 'checkpointable':
    CPT1 = gb_trees:delete(0, gb_trees:delete(1, CPT0)),

    AfterCheckpoint1 = next_counter_checkpoint(checkpointable(State, CPT1), 2),
    ?assertEqual({AfterCheckpoint1, 1}, next_checkpoint(State1)),

    %% after completing item 2, the next checkpoint should be at sequence number 3,
    %% because the gap between successfully-processed items has been closed:
    CPT2 = gb_trees:insert(2, 2, CPT1),
    Expected2 = checkpointable(AfterCheckpoint1, CPT2),
    State2 = note_success(AfterCheckpoint1, [{2, 2}]),
    ?assert(equal_cpt(Expected2, State2)),
    AfterCheckpoint2 =
        next_counter_checkpoint(checkpointable(AfterCheckpoint1, gb_trees:empty()), 4),
    ?assertEqual({AfterCheckpoint2, 3}, next_checkpoint(State2)).

checkpointing_subrecord_test() ->
    State = #state{enable_subsequence_checkpoints = false},
    ?assertEqual(undefined, next_checkpoint(State)),

    SN0 = #sequence_number{user_sub = 0, user_total = 3},
    SN1 = #sequence_number{user_sub = 1, user_total = 3},
    SN2 = #sequence_number{user_sub = 2, user_total = 3},
    SN3 = #sequence_number{},

    %% items 0 and 2 completed and checkpointable, but not all subrecords have completed,
    %% so we can't checkpoint yet:
    State1 = note_success(State, [{0, SN0}, {2, SN2}]),
    ?assertEqual(undefined, next_checkpoint(State1)),

    %% item 1 is the last pending subrecord in an aggregate record, and so item 2 should
    %% be the next checkpoint:
    State2 = note_success(State1, [{1, SN1}]),
    {State3, SN2} = next_checkpoint(State2),

    %% item 3 is not a subrecord, so is checkpointable:
    State4 = note_success(State3, [{3, SN3}]),
    {_State5, SN3} = next_checkpoint(State4),

    %% if subsequence checkpointing is enabled, we can checkpoint at subrecords:
    State6 = State#state{enable_subsequence_checkpoints = true},
    State7 = note_success(State6, [{0, SN0}, {2, SN2}]),
    {State8, SN0} = next_checkpoint(State7),

    State9 = note_success(State8, [{1, SN1}]),
    {_State10, SN2} = next_checkpoint(State9).

watchdog_test() ->
    process_flag(trap_exit, true),
    State = update_watchdog(#state{watchdog_timeout_ms = 200}),
    receive
        {'EXIT', _, watchdog_timeout} ->
            error("unexpected watchdog trigger")
    after 100 ->
        ok
    end,
    update_watchdog(State),
    receive
        {'EXIT', _, watchdog_timeout} ->
            ok
    after 400 ->
        error("watchdog failed to trigger")
    end.

is_sub_record_test() ->
    ?assertEqual(true, is_sub_record(#sequence_number{user_sub = 0, user_total = 2})),
    ?assertEqual(false, is_sub_record(#sequence_number{user_sub = 1, user_total = 2})),
    ?assertEqual(false,
                 is_sub_record(#sequence_number{user_sub = undefined, user_total = undefined})).

-endif.