Fix snapshot live indexes replication bug

kjnilsson · kjnilsson · commit 873db7704666 · 2025-06-24T16:55:32.000+01:00
diff --git a/src/ra_log.erl b/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),
@@ -1107,12 +1118,11 @@ assert(#?MODULE{cfg = #cfg{log_id = LogId},
                 range = Range,
                 snapshot_state = SnapState,
                 current_snapshot = CurrSnap,
-                live_indexes = LiveIndexes,
-                mem_table = _Mt
+                live_indexes = LiveIndexes
                } = State) ->
     %% TODO: remove this at some point?
-    ?DEBUG("~ts: ra_log: asserting Range ~p Snapshot ~p LiveIndexes ~p",
-           [LogId, Range, CurrSnap, LiveIndexes]),
+    ?DEBUG("~ts: ra_log: asserting Range ~p Snapshot ~p",
+           [LogId, Range, CurrSnap]),
     %% perform assertions to ensure log state is correct
     ?assert(CurrSnap =:= ra_snapshot:current(SnapState)),
     ?assert(Range == undefined orelse
diff --git a/src/ra_log_segment.erl b/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(),
@@ -493,8 +494,7 @@ info(Filename, Live0)
                                    {_, _, Sz, _} = maps:get(I, Index),
                                    Acc + Sz
                            end, 0, Live),
-    Info = #{
-             size => Seg#state.data_write_offset,
+    Info = #{size => Seg#state.data_write_offset,
              index_size => Seg#state.data_start,
              file_type => T,
              links => Links,
diff --git a/src/ra_log_segments.erl b/src/ra_log_segments.erl
@@ -726,32 +726,43 @@ compaction_groups([], Groups, _Conf) ->
     lists:reverse(Groups);
 compaction_groups(Infos, Groups, Conf) ->
     case take_group(Infos, Conf, []) of
+        {[], RemInfos} ->
+            compaction_groups(RemInfos, Groups, Conf);
         {Group, RemInfos} ->
             compaction_groups(RemInfos, [Group | Groups], Conf)
     end.
 
-%% TODO: try to take potential size into account
 take_group([], _, Acc) ->
     {lists:reverse(Acc), []};
 take_group([{#{num_entries := NumEnts,
-               live_size := LiveSize}, Live, {_, _}} = E | Rem] = All,
-           #{max_count := Mc,
+               index_size := IdxSz,
+               size := Sz,
+               live_size := LiveSz}, Live, {_, _}} = E | Rem] = All,
+           #{max_count := MaxCnt,
              max_size := MaxSz}, Acc) ->
-    Num = ra_seq:length(Live),
-    case Num / NumEnts < 0.5 of
+    NumLive = ra_seq:length(Live),
+    AllDataSz = Sz - IdxSz,
+    %% group on either num relaimable entries or data saved
+    case NumLive / NumEnts < 0.5 orelse
+         LiveSz / AllDataSz < 0.5 of
+        %% there are fewer than half live entries in the segment
         true ->
-            case Mc - Num < 0 orelse
-                 MaxSz - LiveSize < 0 of
+            %% check that adding this segment to the current group will no
+            %% exceed entry or size limits
+            case MaxCnt - NumLive < 0 orelse
+                 MaxSz - LiveSz < 0 of
                 true ->
+                    %% adding this segment to the group will exceed limits
+                    %% so returning current group
                     {lists:reverse(Acc), All};
                 false ->
-                    take_group(Rem, #{max_count => Mc - Num,
-                                      max_size => MaxSz - LiveSize},
+                    take_group(Rem, #{max_count => MaxCnt - NumLive,
+                                      max_size => MaxSz - LiveSz},
                                [E | Acc])
             end;
-            %% skip this secment
+        %% skip this segment
         false when Acc == [] ->
-            take_group(Rem, #{max_count => Mc,
+            take_group(Rem, #{max_count => MaxCnt,
                               max_size => MaxSz}, Acc);
         false ->
             {lists:reverse(Acc), Rem}
diff --git a/src/ra_mt.erl b/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) ->
@@ -299,6 +302,8 @@ get_items(Indexes, #?MODULE{} = State) ->
     non_neg_integer().
 delete(undefined) ->
     0;
+delete({indexes, _Tid, []}) ->
+    0;
 delete({indexes, Tid, Seq}) ->
     NumToDelete = ra_seq:length(Seq),
     Start = ra_seq:first(Seq),
diff --git a/src/ra_server.erl b/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} |
diff --git a/src/ra_server_proc.erl b/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) ->
diff --git a/test/ra_kv_SUITE.erl b/test/ra_kv_SUITE.erl
@@ -20,7 +20,8 @@ all() ->
 
 all_tests() ->
     [
-     basics
+     basics,
+     snapshot_replication
     ].
 
 groups() ->
@@ -56,6 +57,71 @@ 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)],
+
+    ?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-value">>, 5000),
+    {ok, #{}} = ra_kv:put(KvId, <<"k4">>, <<"k4-value">>, 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)),
+
+    ok = ra:stop_server(default, KvId3),
+
+    {ok, #{}} = ra_kv:put(KvId, <<"k5">>, <<"k5-value">>, 5000),
+    {ok, #{}} = ra_kv:put(KvId, <<"k6">>, <<"k6-value">>, 5000),
+    ok = ra:aux_command(KvId, take_snapshot),
+
+    ok = ra:restart_server(default, KvId3),
+    {ok, #{log := #{last_index := Kv1LastIndex2}}, _} = ra:member_overview(KvId),
+    ok = ra_lib:retry(
+           fun () ->
+                   {ok, #{log := #{last_index := LastIdx}}, _} =
+                       ra:member_overview(KvId3),
+                   Kv1LastIndex2 == LastIdx
+           end, 100, 100),
+
+    ra:delete_cluster([KvId, {kv2, node()}, KvId3]),
+    ok.
+
 basics(_Config) ->
     Members = [{kv1, node()}],
     KvId = hd(Members),
@@ -136,4 +202,5 @@ basics(_Config) ->
                                                     undefined, 1000),
     ?assertEqual(Reads4, Reads5),
     ct:pal("counters ~p", [ra_counters:overview(KvId)]),
+    ra:delete_cluster([KvId, KvId2]),
     ok.
diff --git a/test/ra_log_segments_SUITE.erl b/test/ra_log_segments_SUITE.erl
