Refactor snapshot receive to consolidate temporary state

- Remove snapshot_next_event field, store deferred event in
  snapshot_phase as {awaiting_pending, EventType, Rpc}
- Add ra_log:has_pending/1 for O(1) pending check
- Simplify abort_receive to reset log based on snapshot_has_live_indexes
  rather than checking individual phases

Author: kjnilsson

src/ra_log.erl

@@ -45,6 +45,7 @@
 
          can_write/1,
          exists/2,
+         has_pending/1,
          overview/1,
          %% config
          write_config/2,
@@ -1386,6 +1387,12 @@ exists({Idx, Term}, Log0) ->
     end.
 
 -spec overview(state()) -> overview().
+-spec has_pending(state()) -> boolean().
+has_pending(#?MODULE{pending = []}) ->
+    false;
+has_pending(#?MODULE{}) ->
+    true.
+
 overview(#?MODULE{range = Range,
                   last_term = LastTerm,
                   last_written_index_term = LWIT,

src/ra_log_segments.erl

@@ -520,8 +520,8 @@ segment_read_plan(#cfg{log_id = LogId} = Cfg,
         undefined ->
             %% not found, not good
             ?WARN("~ts: read plan request did not found all requested indexes"
-                  " missing ~w segrefs ~p",
-                  [LogId, Indexes, ra_lol:to_list(SegRefs)]),
+                  " missing ~w segrefs left ~p",
+                  [LogId, Indexes, SegRefs]),
             lists:reverse(Acc)
     end.
 

src/ra_server.erl

@@ -98,14 +98,17 @@
                                                  consistent_query_ref()}),
       pending_consistent_queries := [consistent_query_ref()],
       commit_latency => option(non_neg_integer()),
-      snapshot_phase => chunk_flag(),
+      %% snapshot_phase tracks the current phase of snapshot reception:
+      %% - chunk_flag() for normal phases (init, pre, next, last)
+      %% - {awaiting_pending, EventType, Rpc} when waiting for pending WAL
+      %%   writes to complete before finalizing snapshot installation
+      snapshot_phase => chunk_flag() |
+                        {awaiting_pending, term(), #install_snapshot_rpc{}},
       pending_release_cursor => {ra_index(), term(),
                                  [ra_machine:release_cursor_condition()]},
-      %% TODO: review these 3, specific to handle_receive_snapshot
-      %% could temporarily increase state size over small map (32)
+      %% temporary state for handle_receive_snapshot
       current_event_type => term(),
-      snapshot_has_live_indexes => boolean(),
-      snapshot_next_event => term()
+      snapshot_has_live_indexes => boolean()
      }.
 
 -type state() :: ra_server_state().
@@ -1621,7 +1624,7 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
     %% works as an assertion also
     {AcceptingSnapIdx, _} = ra_snapshot:accepting(SnapState0),
     SnapshotHasLiveIndexes = maps:get(snapshot_has_live_indexes, State0, false),
-    LogHasPendingIndexes = maps:get(num_pending, ra_log:overview(Log00)) > 0,
+    LogHasPendingIndexes = ra_log:has_pending(Log00),
     case ChunkFlag of
         init when SnapPhase == init andalso
                   SnapIndex == AcceptingSnapIdx ->
@@ -1673,16 +1676,14 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
           when SnapshotHasLiveIndexes andalso
                LogHasPendingIndexes ->
             %% we cannot yet complete the snapshot as there are pending
-            %% log indexes
+            %% log indexes, defer completion until they are written
             EventType = maps:get(current_event_type, State0),
             ?DEBUG("~ts: receiving snapshot chunk: ~b / ~w, index ~b, term ~b "
                    "cannot yet complete as log has pending indexes",
                    [LogId, Num, ChunkFlag, SnapIndex, SnapTerm]),
             {receive_snapshot,
-             State0#{snapshot_next_event => {EventType, Rpc}}, []};
+             State0#{snapshot_phase => {awaiting_pending, EventType, Rpc}}, []};
         last ->
-            %% TODO: we can't do this bit until all pending log entries have
-            %% been processed
             ?assert(SnapIndex == AcceptingSnapIdx),
             ?DEBUG("~ts: receiving snapshot chunk: ~b / ~w, index ~b, term ~b",
                    [LogId, Num, ChunkFlag, SnapIndex, SnapTerm]),
@@ -1737,7 +1738,7 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
                                          machine_state => MacState}),
             State = maps:without([snapshot_phase,
                                   snapshot_has_live_indexes,
-                                  snapshot_next_event], State1),
+                                  current_event_type], 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
@@ -1757,15 +1758,27 @@ handle_receive_snapshot(#append_entries_rpc{term = Term} = Msg,
     State = abort_receive(State0),
     {follower, update_term(Term, State), [{next_event, Msg}]};
 handle_receive_snapshot({ra_log_event, Evt},
-                        #{cfg := #cfg{log_id = _LogId},
+                        #{cfg := #cfg{log_id = LogId},
+                          snapshot_phase := Phase,
                           log := Log0} = State0) ->
-    % simply forward all other events to ra_log
-    % whilst the snapshot is being received
+    %% forward log events to ra_log whilst the snapshot is being received
     {Log, Effects} = ra_log:handle_event(Evt, Log0),
-    case maps:take(snapshot_next_event, State0) of
-        {{EventType, NextEvt}, State} ->
-            {receive_snapshot, State#{log => Log},
-             [{next_event, EventType, NextEvt} | Effects]};
+    case Phase of
+        {awaiting_pending, EventType, DeferredRpc} ->
+
+            case ra_log:has_pending(Log) of
+                false ->
+                    ?DEBUG("~ts: pending indexes cleared, completing snapshot",
+                           [LogId]),
+                    %% replay the deferred last chunk, set phase to next so
+                    %% the replayed event doesn't hit the awaiting_pending guard
+                    {receive_snapshot,
+                     State0#{log => Log,
+                             snapshot_phase => next},
+                     [{next_event, EventType, DeferredRpc} | Effects]};
+                true ->
+                    {receive_snapshot, State0#{log => Log}, Effects}
+            end;
         _ ->
             {receive_snapshot, State0#{log => Log}, Effects}
     end;
@@ -1824,23 +1837,24 @@ handle_receive_snapshot(Msg, State) ->
     %% TODO: work out what else to handle
     {receive_snapshot, State, [{reply, {error, {unsupported_call, Msg}}}]}.
 
-abort_receive(#{snapshot_phase := Phase,
+abort_receive(#{snapshot_phase := _Phase,
                 last_applied := LastApplied,
                 log := Log0} = State) ->
     SnapState0 = ra_log:snapshot_state(Log0),
     SnapState = ra_snapshot:abort_accept(SnapState0),
     Log1 = ra_log:set_snapshot_state(SnapState, Log0),
-    Log = case Phase of
-              pre ->
+    Log = case State of
+              #{snapshot_has_live_indexes := true} ->
+                  %% live indexes were written during pre phase,
+                  %% reset log index to undo them
                   {ok, Log2} = ra_log:set_last_index(LastApplied, Log1),
                   Log2;
               _ ->
                   Log1
           end,
     clear_leader_id(maps:without([snapshot_phase,
                                   snapshot_has_live_indexes,
-                                  current_event_type,
-                                  snapshot_next_event],
+                                  current_event_type],
                                  State#{log => Log})).
 
 -spec handle_await_condition(ra_msg(), ra_server_state()) ->

test/ra_log_2_SUITE.erl

@@ -73,6 +73,7 @@ all_tests() ->
      sparse_read,
      read_plan_modified,
      read_plan,
+     read_plan_missing_index,
      sparse_read_out_of_range,
      sparse_read_out_of_range_2,
      written_event_after_snapshot,
@@ -586,6 +587,28 @@ read_plan(Config) ->
     [] = Indexes -- [I || I <- maps:keys(EntriesOut)],
     ok.
 
+read_plan_missing_index(Config) ->
+    Log0 = ra_log_init(Config),
+    Log1 = write_and_roll(1, 500, 1, Log0, 50),
+    Log2 = deliver_all_log_events(Log1, 100),
+    {Log3, Effs} = ra_log:update_release_cursor(300, #{}, macctx(),
+                                                [1, 299], Log2),
+    ct:pal("Effs ~p", [Effs]),
+    run_effs(Effs),
+
+    Log4 = assert_log_events(Log3,
+                             fun (L) ->
+                                     #{snapshot_index := SnapIdx} = ra_log:overview(L),
+                                     SnapIdx == 300
+                             end),
+    timer:sleep(100),
+    flush(),
+    ReadPlan = ra_log:partial_read([128, 301], Log4, fun (_, _, Cmd) -> Cmd end),
+    ?assert(is_map(ra_log_read_plan:info(ReadPlan))),
+    ct:pal("ReadPlan ~p", [ReadPlan]),
+    {Entries, _} = ra_log_read_plan:execute(ReadPlan, undefined),
+    ok.
+
 written_event_after_snapshot(Config) ->
     Log0 = ra_log_init(Config, #{min_snapshot_interval => 1}),
     Log1 = ra_log:append({1, 1, <<"one">>}, Log0),

test/ra_server_SUITE.erl

@@ -212,6 +212,7 @@ setup_log() ->
                 end),
     meck:expect(ra_log, write_config, fun ra_log_memory:write_config/2),
     meck:expect(ra_log, next_index, fun ra_log_memory:next_index/1),
+    meck:expect(ra_log, has_pending, fun (_) -> false end),
     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),