Fix snapshot live indexes replication bug

Author: kjnilsson

src/ra_log.erl

@@ -706,10 +706,21 @@ last_written(#?MODULE{last_written_index_term = LWTI}) ->
     {ok, state()} | {not_found, state()}.
 set_last_index(Idx, #?MODULE{cfg = Cfg,
                              range = Range,
+                             snapshot_state = SnapState,
                              last_written_index_term = {LWIdx0, _}} = State0) ->
+    Cur = ra_snapshot:current(SnapState),
     case fetch_term(Idx, State0) of
-        {undefined, State} ->
+        {undefined, State} when element(1, Cur) =/= Idx ->
+            %% not found and Idx isn't equal to latest snapshot index
             {not_found, State};
+        {_, State} when element(1, Cur) =:= Idx ->
+            {_, SnapTerm} = Cur,
+            %% Idx is equal to the current snapshot
+            put_counter(Cfg, ?C_RA_SVR_METRIC_LAST_INDEX, Idx),
+            put_counter(Cfg, ?C_RA_SVR_METRIC_LAST_WRITTEN_INDEX, Idx),
+            {ok, State#?MODULE{range = ra_range:limit(Idx + 1, Range),
+                               last_term = SnapTerm,
+                               last_written_index_term = Cur}};
         {Term, State1} ->
             LWIdx = min(Idx, LWIdx0),
             {LWTerm, State2} = fetch_term(LWIdx, State1),

src/ra_log_segment.erl

@@ -456,6 +456,7 @@ segref(Filename) ->
     SegRef.
 
 -type infos() :: #{size => non_neg_integer(),
+                   index_size => non_neg_integer(),
                    max_count => non_neg_integer(),
                    file_type => regular | symlink,
                    ctime => integer(),

src/ra_mt.erl

@@ -120,7 +120,10 @@ insert({Idx, _, _} = _Entry,
     end.
 
 -spec insert_sparse(log_entry(), undefined | ra:index(), state()) ->
-    {ok, state()} | {error, overwriting | gap_detected | limit_reached}.
+    {ok, state()} | {error,
+                     overwriting |
+                     gap_detected |
+                     limit_reached}.
 insert_sparse({Idx, _, _} = Entry, LastIdx,
               #?MODULE{tid = Tid,
                        indexes = Seq} = State) ->

src/ra_server.erl

@@ -1427,7 +1427,7 @@ handle_follower(#install_snapshot_rpc{term = Term,
                                       meta = #{index := SnapIdx,
                                                machine_version := SnapMacVer} = Meta,
                                       leader_id = LeaderId,
-                                      chunk_state = {Num, _ChunkFlag}} = Rpc,
+                                      chunk_state = {Num, ChunkFlag}} = Rpc,
                 #{cfg := #cfg{log_id = LogId,
                               machine_version = MacVer}, log := Log0,
                   last_applied := LastApplied,
@@ -1443,7 +1443,17 @@ handle_follower(#install_snapshot_rpc{term = Term,
            [LogId, SnapIdx, Term]),
     SnapState0 = ra_log:snapshot_state(Log0),
     {ok, SS} = ra_snapshot:begin_accept(Meta, SnapState0),
-    Log = ra_log:set_snapshot_state(SS, Log0),
+    Log1 = ra_log:set_snapshot_state(SS, Log0),
+
+    %% if the snaphost includes pre entries (live entries) then we need
+    %% to reset the log to the last applied index to avoid issues
+    Log = case ChunkFlag of
+              pre ->
+                  {ok, L} = ra_log:set_last_index(LastApplied, Log1),
+                  L;
+              _ ->
+                  Log1
+          end,
     {receive_snapshot, update_term(Term, State0#{log => Log,
                                                  leader_id => LeaderId}),
      [{next_event, Rpc}, {record_leader_msg, LeaderId}]};
@@ -1538,17 +1548,18 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
                                      last_index = SnapIndex},
     case ChunkFlag of
         pre when is_list(ChunkOrEntries) ->
-            %% TODO: we may need to reset the log here to
-            %% the last applied index as we
-            %% dont know for sure indexes after last applied
+            %% reset last index to last applied
+            %% as we dont know for sure indexes after last applied
             %% are of the right term
             {LastIndex, _} = ra_log:last_index_term(Log00),
-            {Log0, _} = lists:foldl(
+            {Log, _} = lists:foldl(
                        fun ({I, _, _} = E, {L0, LstIdx}) ->
                                {ok, L} = ra_log:write_sparse(E, LstIdx, L0),
                                {L, I}
                        end, {Log00, LastIndex}, ChunkOrEntries),
-            State = update_term(Term, State0#{log => Log0}),
+            ?DEBUG("~ts: receiving snapshot log last index ~p",
+                   [LogId, ra_log:last_index_term(Log)]),
+            State = update_term(Term, State0#{log => Log}),
             {receive_snapshot, State, [{reply, Reply}]};
         next ->
             SnapState0 = ra_log:snapshot_state(Log00),
@@ -1606,6 +1617,7 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
                                         membership =>
                                             get_membership(ClusterIds, State0),
                                         machine_state => MacState}),
+            put_counter(Cfg, ?C_RA_SVR_METRIC_LAST_APPLIED, SnapIndex),
             %% it was the last snapshot chunk so we can revert back to
             %% follower status
             {follower, persist_last_applied(State), [{reply, Reply} |

src/ra_server_proc.erl

@@ -1535,7 +1535,9 @@ handle_effect(leader, {send_snapshot, {_, ToNode} = To, {SnapState, _Id, Term}},
               #state{server_state = SS0,
                      monitors = Monitors,
                      conf = #conf{snapshot_chunk_size = ChunkSize,
-                     install_snap_rpc_timeout = InstallSnapTimeout} = Conf} = State0,
+                                  log_id = LogId,
+                                  install_snap_rpc_timeout = InstallSnapTimeout} = Conf}
+              = State0,
               Actions) ->
     case lists:member(ToNode, [node() | nodes()]) of
         true ->
@@ -1546,7 +1548,7 @@ handle_effect(leader, {send_snapshot, {_, ToNode} = To, {SnapState, _Id, Term}},
             Pid = spawn(fun () ->
                                 try send_snapshots(Id, Term, To,
                                                    ChunkSize, InstallSnapTimeout,
-                                                   SnapState, Machine) of
+                                                   SnapState, Machine, LogId) of
                                     _ -> ok
                                 catch
                                     C:timeout:S ->
@@ -1906,29 +1908,39 @@ read_entries0(From, Idxs, #state{server_state = #{log := Log}} = State) ->
     {keep_state, State, [{reply, From, {ok, ReadState}}]}.
 
 send_snapshots(Id, Term, {_, ToNode} = To, ChunkSize,
-               InstallTimeout, SnapState, Machine) ->
+               InstallTimeout, SnapState, Machine, LogId) ->
     Context = ra_snapshot:context(SnapState, ToNode),
     {ok, #{machine_version := SnapMacVer} = Meta, ReadState} =
         ra_snapshot:begin_read(SnapState, Context),
 
-    %% only send the snapshot if the target server can accept it
-    %% TODO: grab the last_applied index also and use this to floor
-    %% the live indexes
+    %% TODO: consolidate getting the context, machinve version and last
+    %% applied index in one rpc, and handle errors
     TheirMacVer = erpc:call(ToNode, ra_machine, version, [Machine]),
 
-    %% rpc the check what their
+    %% only send the snapshot if the target server can accept it
     case SnapMacVer > TheirMacVer of
         true ->
+            ?DEBUG("~ts: not sending snapshot to ~w as their machine version ~b "
+                   "is lower than snapshot machine version ~b",
+                   [LogId, To, TheirMacVer, SnapMacVer]),
             ok;
         false ->
+            #{last_applied := LastApplied} = erpc:call(ToNode,
+                                                       ra_counters,
+                                                       counters,
+                                                       [To, [last_applied]]),
             RPC = #install_snapshot_rpc{term = Term,
                                         leader_id = Id,
                                         meta = Meta},
-            case ra_snapshot:indexes(ra_snapshot:current_snapshot_dir(SnapState)) of
-                {ok, [_|_] = Indexes} ->
+            case ra_snapshot:indexes(
+                   ra_snapshot:current_snapshot_dir(SnapState)) of
+                {ok, [_|_] = Indexes0} ->
+                    %% remove all indexes lower than the target's last applied
+                    Indexes = ra_seq:floor(LastApplied + 1, Indexes0),
+                    ?DEBUG("~ts: sending live indexes ~w to ~w ",
+                           [LogId, ra_seq:range(Indexes), To]),
                     %% there are live indexes to send before the snapshot
-                    %% %% TODO: only send live indexes higher than the follower's
-                    %% last_applied index
+                    %% TODO: write ra_seq:list_chunk function to avoid expansion
                     Idxs = lists:reverse(ra_seq:expand(Indexes)),
                     Flru = lists:foldl(
                              fun (Is, F0) ->

test/ra_kv_SUITE.erl

@@ -20,7 +20,8 @@ all() ->
 
 all_tests() ->
     [
-     basics
+     basics,
+     snapshot_replication
     ].
 
 groups() ->
@@ -56,6 +57,60 @@ end_per_testcase(_TestCase, _Config) ->
 %%%===================================================================
 
 
+snapshot_replication(_Config) ->
+    Members = [{kv1, node()}, {kv2, node()}],
+    KvId = hd(Members),
+
+    {ok, _, _} = ra_kv:start_cluster(?SYS, ?FUNCTION_NAME,
+                                     #{members => Members}),
+    ra:transfer_leadership(KvId, KvId),
+    {ok, #{}} = ra_kv:put(KvId, <<"k1">>, <<"k1-value01">>, 5000),
+    %% write 10k entries of the same key
+    [{ok, #{}} = ra_kv:put(KvId, integer_to_binary(I), I, 5000)
+     || I <- lists:seq(1, 5000)],
+
+    ct:pal("kv get ~p", [ra_kv:get(KvId, <<"k1">>, 5000)]),
+    ct:pal("leaderboard ~p", [ets:tab2list(ra_leaderboard)]),
+
+    ?assertMatch({ok, #{machine := #{num_keys := _}}, KvId},
+                 ra:member_overview(KvId)),
+    ra_log_wal:force_roll_over(ra_log_wal),
+    %% wait for rollover processing
+    ra_log_wal:last_writer_seq(ra_log_wal, <<>>),
+    %% wait for segment writer to process
+    ra_log_segment_writer:await(ra_log_segment_writer),
+    %% promt ra_kv to take a snapshot
+    ok = ra:aux_command(KvId, take_snapshot),
+    ok = ra_lib:retry(
+           fun () ->
+                   {ok, #{log := #{snapshot_index := SnapIdx,
+                                   last_index := LastIdx}}, _} =
+                       ra:member_overview(KvId),
+                   SnapIdx == LastIdx
+           end, 100, 100),
+
+    KvId3 = {kv3, node()},
+    ok = ra_kv:add_member(?SYS, KvId3, KvId),
+    KvId3Pid = whereis(kv3),
+    ?assert(is_pid(KvId3Pid)),
+    {ok, #{}} = ra_kv:put(KvId, <<"k3">>, <<"k3-value01">>, 5000),
+    {ok, #{}} = ra_kv:put(KvId, <<"k4">>, <<"k3-value01">>, 5000),
+    ok = ra:aux_command(KvId, take_snapshot),
+    % timer:sleep(1000),
+    {ok, #{log := #{last_index := Kv1LastIndex  }}, _} =
+        ra:member_overview(KvId),
+    ok = ra_lib:retry(
+           fun () ->
+                   {ok, #{log := #{last_index := LastIdx}}, _} =
+                       ra:member_overview(KvId3),
+                   Kv1LastIndex == LastIdx
+           end, 100, 100),
+    ct:pal("counters ~p", [ra_counters:counters(KvId3, [last_applied])]),
+    %% ensure kv3 did not crash during snapshot replication
+    ?assertEqual(KvId3Pid, whereis(kv3)),
+    ra:delete_cluster([KvId, {kv2, node()}, KvId3]),
+    ok.
+
 basics(_Config) ->
     Members = [{kv1, node()}],
     KvId = hd(Members),
@@ -136,4 +191,5 @@ basics(_Config) ->
                                                     undefined, 1000),
     ?assertEqual(Reads4, Reads5),
     ct:pal("counters ~p", [ra_counters:overview(KvId)]),
+    ra:delete_cluster([KvId, KvId2]),
     ok.