Make snapshots with a pre phase restartable

kjnilsson · kjnilsson · commit 588d066ee482 · 2025-07-21T16:38:30.000+01:00
diff --git a/src/ra.hrl b/src/ra.hrl
@@ -83,7 +83,7 @@
 %% represent a unique entry in the ra log
 -type log_entry() :: {ra_index(), ra_term(), term()}.
 
--type chunk_flag() :: pre | next | last.
+-type chunk_flag() :: init | pre | next | last.
 
 -type consistent_query_ref() :: {From :: term(), Query :: ra:query_fun(), ConmmitIndex :: ra_index()}.
 
@@ -169,7 +169,7 @@
         {term :: ra_term(), % the leader's term
          leader_id :: ra_server_id(),
          meta :: snapshot_meta(),
-         chunk_state = {0, pre} :: {non_neg_integer(), chunk_flag()},
+         chunk_state :: {non_neg_integer(), chunk_flag()},
          data :: term()
         }).
 
diff --git a/src/ra_log.erl b/src/ra_log.erl
@@ -875,6 +875,7 @@ handle_event(major_compaction, #?MODULE{reader = Reader0,
     end;
 handle_event({snapshot_written, {SnapIdx, _} = Snap, LiveIndexes, SnapKind},
              #?MODULE{cfg = #cfg{uid = UId,
+                                 log_id = LogId,
                                  names = Names} = Cfg,
                       range = {FstIdx, _} = Range,
                       mem_table = Mt0,
@@ -886,6 +887,8 @@ handle_event({snapshot_written, {SnapIdx, _} = Snap, LiveIndexes, SnapKind},
     % ?assert(ra_snapshot:pending(SnapState0) =/= undefined),
     SnapState1 = ra_snapshot:complete_snapshot(Snap, SnapKind, LiveIndexes,
                                                SnapState0),
+    ?DEBUG("~ts: ra_log: ~s written at index ~b with ~b live indexes",
+           [LogId, SnapKind, SnapIdx, ra_seq:length(LiveIndexes)]),
     case SnapKind of
         snapshot ->
             put_counter(Cfg, ?C_RA_SVR_METRIC_SNAPSHOT_INDEX, SnapIdx),
diff --git a/src/ra_log_wal.erl b/src/ra_log_wal.erl
@@ -526,7 +526,7 @@ handle_msg({append, {UId, Pid} = Id, MtTid, ExpectedPrevIdx, Idx, Term, Entry},
         {in_seq, PrevIdx} ->
             % writer was in seq but has sent an out of seq entry
             % notify writer
-            ?DEBUG("WAL in ~ts: requesting resend from `~w`, "
+            ?DEBUG("WAL in ~ts: requesting resend for `~s`, "
                    "last idx ~b idx received (~b,~b)",
                    [Conf#conf.system, UId, PrevIdx, ExpectedPrevIdx, Idx]),
             Pid ! {ra_log_event, {resend_write, PrevIdx + 1}},
diff --git a/src/ra_server.erl b/src/ra_server.erl
@@ -99,7 +99,8 @@
       queries_waiting_heartbeats := queue:queue({non_neg_integer(),
                                                  consistent_query_ref()}),
       pending_consistent_queries := [consistent_query_ref()],
-      commit_latency => option(non_neg_integer())
+      commit_latency => option(non_neg_integer()),
+      snapshot_phase => chunk_flag()
      }.
 
 -type state() :: ra_server_state().
@@ -1447,25 +1448,27 @@ handle_follower(#install_snapshot_rpc{term = Term,
        SnapIdx > LastApplied andalso
        %% only install snapshot if the machine version is understood
        MacVer >= SnapMacVer andalso
-       Num =< 1 ->
+       Num =< 1 andalso
+       ChunkFlag /= pre ->
     %% only begin snapshot procedure if Idx is higher than the last_applied
     %% index.
-    ?DEBUG("~ts: begin_accept snapshot at index ~b in term ~b",
-           [LogId, SnapIdx, Term]),
+    ?DEBUG("~ts: begin_accept snapshot at index ~b in term ~b, phase ~s",
+           [LogId, SnapIdx, Term, ChunkFlag]),
     SnapState0 = ra_log:snapshot_state(Log0),
     {ok, SS} = ra_snapshot:begin_accept(Meta, SnapState0),
     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 ->
+              init ->
                   {ok, L} = ra_log:set_last_index(LastApplied, Log1),
                   L;
               _ ->
                   Log1
           end,
     {receive_snapshot, update_term(Term, State0#{log => Log,
+                                                 snapshot_phase => ChunkFlag,
                                                  leader_id => LeaderId}),
      [{next_event, Rpc}, {record_leader_msg, LeaderId}]};
 handle_follower(#install_snapshot_rpc{term = Term,
@@ -1540,7 +1543,7 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
                                                        cluster := ClusterIds,
                                                        term := SnapTerm} = SnapMeta,
                                               chunk_state = {Num, ChunkFlag},
-                                              data = ChunkOrEntries},
+                                              data = ChunkOrEntries} = Rpc,
                         #{cfg := #cfg{id = Id,
                                       log_id = LogId,
                                       effective_machine_version = CurEffMacVer,
@@ -1550,15 +1553,33 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
                           cluster := Cluster,
                           current_term := CurTerm,
                           last_applied := LastApplied,
-                          machine_state := OldMacState} = State0)
+                          machine_state := OldMacState,
+                          snapshot_phase := SnapPhase} = State0)
   when Term >= CurTerm ->
-    ?DEBUG("~ts: receiving snapshot chunk: ~b / ~w, index ~b, term ~b",
-           [LogId, Num, ChunkFlag, SnapIndex, SnapTerm]),
     Reply = #install_snapshot_result{term = CurTerm,
                                      last_term = SnapTerm,
                                      last_index = SnapIndex},
     case ChunkFlag of
+        init when SnapPhase == init ->
+            %% this is ok, just reply
+            {receive_snapshot, State0, [{reply, Reply}]};
+        init  ->
+            ?DEBUG("~ts: receiving snapshot saw unexpected init phase at snapshot"
+                   " index term {~b, ~b}, current phase ~s restarting
+                   snapshot receive process",
+                   [LogId, SnapIndex, SnapTerm, SnapPhase]),
+            %% the snapshot sending must have been interrupted and restarted
+            %% during the init or pre-phase
+            %% abort the snapshot, and revert to follower
+            SnapState0 = ra_log:snapshot_state(Log00),
+            SnapState = ra_snapshot:abort_accept(SnapState0),
+            Log = ra_log:set_snapshot_state(SnapState, Log00),
+            {follower, maps:remove(snapshot_phase, State0#{log => Log}),
+             [{next_event, Rpc}]};
         pre when is_list(ChunkOrEntries) ->
+            [{_FstIdx, _, _} | _] = ChunkOrEntries,
+            % ?DEBUG("~ts: receiving snapshot chunk pre first index ~b snap index ~b, term ~b",
+            %        [LogId, FstIdx, SnapIndex, SnapTerm]),
             %% reset last index to last applied
             %% as we dont know for sure indexes after last applied
             %% are of the right term
@@ -1568,15 +1589,21 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
                                {ok, L} = ra_log:write_sparse(E, LstIdx, L0),
                                {L, I}
                        end, {Log00, LastIdx}, ChunkOrEntries),
-            State = update_term(Term, State0#{log => Log}),
+            State = update_term(Term, State0#{log => Log,
+                                              snapshot_phase => pre}),
             {receive_snapshot, State, [{reply, Reply}]};
         next ->
+            ?DEBUG("~ts: receiving snapshot chunk: ~b / ~w, index ~b, term ~b",
+                   [LogId, Num, ChunkFlag, SnapIndex, SnapTerm]),
             SnapState0 = ra_log:snapshot_state(Log00),
             SnapState = ra_snapshot:accept_chunk(ChunkOrEntries, Num, SnapState0),
             Log0 = ra_log:set_snapshot_state(SnapState, Log00),
-            State = update_term(Term, State0#{log => Log0}),
+            State = update_term(Term, State0#{log => Log0,
+                                              snapshot_phase => next}),
             {receive_snapshot, State, [{reply, Reply}]};
         last ->
+            ?DEBUG("~ts: receiving snapshot chunk: ~b / ~w, index ~b, term ~b",
+                   [LogId, Num, ChunkFlag, SnapIndex, SnapTerm]),
             SnapState0 = ra_log:snapshot_state(Log00),
             {SnapState, MacState, LiveIndexes, Effs0} =
                 ra_snapshot:complete_accept(ChunkOrEntries, Num, Machine,
@@ -1614,7 +1641,7 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
                                                               MacState,
                                                               OldMeta,
                                                               OldMacState),
-            State = update_term(Term,
+            State1 = update_term(Term,
                                 State0#{cfg => Cfg,
                                         log => Log,
                                         commit_index => SnapIndex,
@@ -1626,6 +1653,7 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
                                         membership =>
                                             get_membership(ClusterIds, State0),
                                         machine_state => MacState}),
+            State = maps:remove(snapshot_phase, State1),
             put_counter(Cfg, ?C_RA_SVR_METRIC_LAST_APPLIED, SnapIndex),
             %% it was the last snapshot chunk so we can revert back to
             %% follower status
@@ -1645,13 +1673,15 @@ handle_receive_snapshot(#append_entries_rpc{term = Term} = Msg,
     SnapState0 = ra_log:snapshot_state(Log0),
     SnapState = ra_snapshot:abort_accept(SnapState0),
     Log = ra_log:set_snapshot_state(SnapState, Log0),
-    {follower, update_term(Term, clear_leader_id(State#{log => Log})),
+    {follower, maps:remove(snapshot_phase,
+                           update_term(Term,
+                                       clear_leader_id(State#{log => Log}))),
      [{next_event, Msg}]};
 handle_receive_snapshot({ra_log_event, Evt},
-                        #{cfg := #cfg{log_id = LogId},
+                        #{cfg := #cfg{log_id = _LogId},
                           log := Log0} = State) ->
-    ?DEBUG("~ts: ~s ra_log_event received: ~w",
-          [LogId, ?FUNCTION_NAME, Evt]),
+    % ?DEBUG("~ts: ~s ra_log_event received: ~w",
+    %       [LogId, ?FUNCTION_NAME, Evt]),
     % simply forward all other events to ra_log
     % whilst the snapshot is being received
     {Log, Effects} = ra_log:handle_event(Evt, Log0),
@@ -1664,7 +1694,7 @@ handle_receive_snapshot(receive_snapshot_timeout,
     SnapState0 = ra_log:snapshot_state(Log0),
     SnapState = ra_snapshot:abort_accept(SnapState0),
     Log = ra_log:set_snapshot_state(SnapState, Log0),
-    {follower, State#{log => Log}, []};
+    {follower, maps:remove(snapshot_phase, State#{log => Log}), []};
 handle_receive_snapshot(#info_rpc{term = Term} = Msg,
                         #{current_term := CurTerm,
                           cfg := #cfg{log_id = LogId},
@@ -1677,7 +1707,8 @@ handle_receive_snapshot(#info_rpc{term = Term} = Msg,
     SnapState0 = ra_log:snapshot_state(Log0),
     SnapState = ra_snapshot:abort_accept(SnapState0),
     Log = ra_log:set_snapshot_state(SnapState, Log0),
-    {follower, update_term(Term, clear_leader_id(State#{log => Log})),
+    {follower, maps:remove(snapshot_phase,
+                           update_term(Term, clear_leader_id(State#{log => Log}))),
      [{next_event, Msg}]};
 handle_receive_snapshot(#info_rpc{} = InfoRpc, State) ->
     InfoReplyEffect = empty_info_reply_effect(State, InfoRpc),
@@ -1694,7 +1725,8 @@ handle_receive_snapshot(#info_reply{term = Term} = Msg,
     SnapState0 = ra_log:snapshot_state(Log0),
     SnapState = ra_snapshot:abort_accept(SnapState0),
     Log = ra_log:set_snapshot_state(SnapState, Log0),
-    {follower, update_term(Term, clear_leader_id(State#{log => Log})),
+    {follower, maps:remove(snapshot_phase,
+                           update_term(Term, clear_leader_id(State#{log => Log}))),
      [{next_event, Msg}]};
 handle_receive_snapshot(#info_reply{}, State) ->
     {receive_snapshot, State, []};
diff --git a/src/ra_server_proc.erl b/src/ra_server_proc.erl
@@ -1943,6 +1943,7 @@ send_snapshots(Id, Term, {_, ToNode} = To, ChunkSize,
                                                        [To, [last_applied]]),
             RPC = #install_snapshot_rpc{term = Term,
                                         leader_id = Id,
+                                        chunk_state = {0, init},
                                         meta = Meta},
             case ra_snapshot:indexes(
                    ra_snapshot:current_snapshot_dir(SnapState)) of
@@ -1951,6 +1952,9 @@ send_snapshots(Id, Term, {_, ToNode} = To, ChunkSize,
                     Indexes = ra_seq:floor(LastApplied + 1, Indexes0),
                     ?DEBUG("~ts: sending live indexes ~w to ~w ",
                            [LogId, ra_seq:range(Indexes), To]),
+                    %% first send the init phase
+                    _Res0 = gen_statem:call(To, RPC,
+                                            {dirty_timeout, InstallTimeout}),
                     %% there are live indexes to send before the snapshot
                     %% TODO: write ra_seq:list_chunk function to avoid expansion
                     Idxs = lists:reverse(ra_seq:expand(Indexes)),
@@ -1962,7 +1966,7 @@ send_snapshots(Id, Term, {_, ToNode} = To, ChunkSize,
                                      RPC1 = RPC#install_snapshot_rpc{chunk_state = {0, pre},
                                                                      data = Ents},
                                      _Res1 = gen_statem:call(To, RPC1,
-                                                            {dirty_timeout, InstallTimeout}),
+                                                             {dirty_timeout, InstallTimeout}),
                                      %% TODO: assert Res1 is successful
                                      F
                              end, undefined, ra_lib:lists_chunk(16, Idxs)),
diff --git a/test/ra_log_memory.erl b/test/ra_log_memory.erl
@@ -10,6 +10,7 @@
          close/1,
          append/2,
          write/2,
+         write_sparse/3,
          take/3,
          fold/5,
          last_index_term/1,
@@ -94,7 +95,7 @@ write([{FirstIdx, _, _} | _] = Entries,
                                            {Acc#{Idx => {Term, Data}}, Idx}
                                    end, {Log1, FirstIdx}, Entries),
     {ok, State#state{last_index = LastInIdx,
-                          entries = Log}};
+                     entries = Log}};
 write([{FirstIdx, _, _} | _] = Entries,
       #state{snapshot = {#{index := SnapIdx}, _}, entries = Log0} = State)
  when SnapIdx + 1 =:= FirstIdx ->
@@ -106,6 +107,12 @@ write([{FirstIdx, _, _} | _] = Entries,
 write(_Entries, _State) ->
     {error, {integrity_error, undefined}}.
 
+write_sparse({Idx, Term, Data}, PrevIdx, #state{last_index = PrevIdx,
+                                                entries = Log0} = State) ->
+    {ok, State#state{last_index = Idx,
+                     entries = Log0#{Idx => {Term, Data}}}};
+write_sparse(_, _, _) ->
+    {error, gap_detected}.
 
 take(Start, Num, #state{last_index = LastIdx, entries = Log} = State) ->
     Entries = sparse_take(Start, Log, Num, LastIdx, []),
diff --git a/test/ra_server_SUITE.erl b/test/ra_server_SUITE.erl
@@ -64,6 +64,7 @@ all() ->
      leader_appends_cluster_change_then_steps_before_applying_it,
      leader_receives_install_snapshot_rpc,
      follower_installs_snapshot,
+     follower_installs_snapshot_with_pre,
      follower_ignores_installs_snapshot_with_higher_machine_version,
      follower_receives_stale_snapshot,
      follower_receives_snapshot_lower_than_last_applied,
@@ -166,8 +167,8 @@ setup_log() ->
                         {ok, {Meta, Data}, []}
                 end),
     meck:expect(ra_snapshot, complete_accept,
-                fun(_Data, _Num, _Machine, {_Meta, MacSt} = State) ->
-                        {State, MacSt, [], []}
+                fun(Data, _Num, _Machine, {_Meta, MacSt} = State) ->
+                        {State, Data ++ MacSt, [], []}
                 end),
     meck:expect(ra_snapshot, abort_accept, fun(SS) -> SS end),
     meck:expect(ra_snapshot, accepting, fun(_SS) -> undefined end),
@@ -191,6 +192,7 @@ setup_log() ->
     meck:expect(ra_log, next_index, fun ra_log_memory:next_index/1),
     meck:expect(ra_log, append, fun ra_log_memory:append/2),
     meck:expect(ra_log, write, fun ra_log_memory:write/2),
+    meck:expect(ra_log, write_sparse, fun ra_log_memory:write_sparse/3),
     meck:expect(ra_log, handle_event, fun ra_log_memory:handle_event/2),
     meck:expect(ra_log, last_written, fun ra_log_memory:last_written/1),
     meck:expect(ra_log, last_index_term, fun ra_log_memory:last_index_term/1),
@@ -2288,6 +2290,55 @@ follower_installs_snapshot(_Config) ->
 
     ok.
 
+follower_installs_snapshot_with_pre(_Config) ->
+    N1 = ?N1, N2 = ?N2, N3 = ?N3,
+    #{N3 := {_, State = #{cluster := Config}, _}} =
+        init_servers([N1, N2, N3], {module, ra_queue, #{}}),
+    LastTerm = 1, % snapshot term
+    Term = 2, % leader term
+    Idx = 3,
+    ISRpcInit = #install_snapshot_rpc{term = Term, leader_id = N1,
+                                  meta = snap_meta(Idx, LastTerm, Config),
+                                  chunk_state = {0, init},
+                                  data = []},
+    %% the init message starts the process
+    {receive_snapshot, State1,
+     [{next_event, ISRpc}, {record_leader_msg, _}]} =
+        ra_server:handle_follower(ISRpcInit, State#{current_term => Term}),
+
+    %% actually process init message in the correct state
+    {receive_snapshot, State2, [{reply, _}]} =
+        ra_server:handle_receive_snapshot(ISRpc, State1),
+
+    %% now send a pre message
+    ISRpcPre = ISRpcInit#install_snapshot_rpc{chunk_state = {0, pre},
+                                           data = [{2, 1, <<"e1">>}]},
+    {receive_snapshot, State3, [{reply, _}]} =
+        ra_server:handle_receive_snapshot(ISRpcPre, State2),
+
+    %% test that init returns to follower and retries
+    {follower, State3b, [{next_event, ISRpcInit}]} =
+        ra_server:handle_receive_snapshot(ISRpcInit, State3),
+    ?assertNot(maps:is_key(snapshot_phase, State3b)),
+
+    meck:expect(ra_snapshot, complete_accept,
+                fun (Mac, _, _, S) ->
+                        {S, Mac, [], []}
+                end),
+
+    %% finally process the actual snapshot
+    ISRpc1 = ISRpc#install_snapshot_rpc{chunk_state = {1, last},
+                                        data = [2]},
+    {follower, #{current_term := Term,
+                 commit_index := Idx,
+                 last_applied := Idx,
+                 cluster := Config,
+                 machine_state := [2],
+                 leader_id := N1} = _State,
+     [{reply, #install_snapshot_result{}}]} =
+        ra_server:handle_receive_snapshot(ISRpc1, State3),
+    ok.
+
 follower_ignores_installs_snapshot_with_higher_machine_version(_Config) ->
     %% currently followers cannot correctly handle snapshots with a higher
     %% machine version so have to ignore them
