From 0290e03b9c93ecb627646cc3ffbc56e211299c9f Mon Sep 17 00:00:00 2001
From: Michael Davis <mcarsondavis@gmail.com>
Date: Thu, 22 Aug 2024 15:23:55 -0400
Subject: [PATCH] Add a ra_machine:snapshot_installed/4 callback to pass old
 state

---
 src/ra_machine.erl          | 35 +++++++++++++++++++++++++++++++----
 src/ra_server.erl           |  5 ++++-
 test/coordination_SUITE.erl |  2 +-
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/src/ra_machine.erl b/src/ra_machine.erl
index 05882a60..e67dc23e 100644
--- a/src/ra_machine.erl
+++ b/src/ra_machine.erl
@@ -69,7 +69,7 @@
 -export([init/2,
          apply/4,
          tick/3,
-         snapshot_installed/3,
+         snapshot_installed/5,
          state_enter/3,
          overview/2,
          query/3,
@@ -210,6 +210,7 @@
 
 -optional_callbacks([tick/2,
                      snapshot_installed/2,
+                     snapshot_installed/4,
                      state_enter/2,
                      init_aux/1,
                      handle_aux/5,
@@ -243,6 +244,14 @@
 
 -callback snapshot_installed(ra_snapshot:meta(), state()) -> effects().
 
+-callback snapshot_installed(Meta, OldMacVer, OldState, NewState) -> Effects
+    when
+      Meta :: ra_snapshot:meta(),
+      OldMacVer :: version(),
+      OldState :: state(),
+      NewState :: state(),
+      Effects :: effects().
+
 -callback init_aux(Name :: atom()) -> AuxState :: term().
 
 -callback handle_aux(ra_server:ra_state(),
@@ -302,9 +311,27 @@ apply(Mod, Metadata, Cmd, State) ->
 tick(Mod, TimeMs, State) ->
     ?OPT_CALL(Mod:tick(TimeMs, State), []).
 
--spec snapshot_installed(module(), ra_snapshot:meta(), state()) -> effects().
-snapshot_installed(Mod, Meta, State) ->
-    ?OPT_CALL(Mod:snapshot_installed(Meta, State), []).
+-spec snapshot_installed(Module, Meta, OldMacVer, OldState, NewState) ->
+    Effects when
+      Module :: module(),
+      Meta :: ra_snapshot:meta(),
+      OldMacVer :: version(),
+      OldState :: state(),
+      NewState :: state(),
+      Effects :: effects().
+
+snapshot_installed(Mod, Meta, OldMacVer, OldState, NewState) ->
+    try
+        Mod:snapshot_installed(Meta, OldMacVer, OldState, NewState)
+    catch
+        error:undef ->
+            try
+                Mod:snapshot_installed(Meta, NewState)
+            catch
+                error:undef ->
+                    []
+            end
+    end.
 
 %% @doc called when the ra_server_proc enters a new state
 -spec state_enter(module(), ra_server:ra_state() | eol, state()) ->
diff --git a/src/ra_server.erl b/src/ra_server.erl
index 4dc3a701..6ae440ac 100644
--- a/src/ra_server.erl
+++ b/src/ra_server.erl
@@ -1342,7 +1342,8 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
                                       machine_versions = MachineVersions,
                                       machine = Machine} = Cfg0,
                           log := Log0,
-                          current_term := CurTerm} = State0)
+                          current_term := CurTerm,
+                          machine_state := OldMacState} = State0)
   when Term >= CurTerm ->
     ?DEBUG("~ts: receiving snapshot chunk: ~b / ~w, index ~b, term ~b",
            [LogId, Num, ChunkFlag, SnapIndex, SnapTerm]),
@@ -1377,6 +1378,8 @@ handle_receive_snapshot(#install_snapshot_rpc{term = Term,
 
             SnapInstalledEffs = ra_machine:snapshot_installed(EffMacMod,
                                                               SnapMeta,
+                                                              CurEffMacVer,
+                                                              OldMacState,
                                                               MacState),
 
             State = update_term(Term,
diff --git a/test/coordination_SUITE.erl b/test/coordination_SUITE.erl
index a97ee94d..70653ba0 100644
--- a/test/coordination_SUITE.erl
+++ b/test/coordination_SUITE.erl
@@ -1329,7 +1329,7 @@ apply(#{index := _Idx}, {segment_writer_or_wal_crash_follower, _}, State) ->
 apply(#{index := Idx}, _Cmd, State) ->
     {State, ok, [{release_cursor, Idx, State}]}.
 
-snapshot_installed(Meta, _State) ->
+snapshot_installed(Meta, _OldMacVer, _OldState, _NewState) ->
     case whereis(snapshot_installed_proc) of
         undefined ->
             [];