add optional storage callback for creating witness snapshots

Reviewed By: hsun324

Differential Revision: D72425041

fbshipit-source-id: e34c5378a659698efb67500d18009246c5547f19

Author: Sarah Hassan

include/wa_raft.hrl

@@ -51,6 +51,10 @@
 -define(SNAPSHOT_PREFIX, "snapshot").
 %% Snapshot name
 -define(SNAPSHOT_NAME(Index, Term), (?SNAPSHOT_PREFIX "." ++ integer_to_list(Index) ++ "." ++ integer_to_list(Term))).
+
+%% Witness Snapshot name
+-define(WITNESS_SNAPSHOT_NAME(Index, Term), (?SNAPSHOT_PREFIX "." ++ integer_to_list(Index) ++ "." ++ integer_to_list(Term) ++ ".witness")).
+
 %% Location of a snapshot
 -define(RAFT_SNAPSHOT_PATH(Path, Name), (filename:join(Path, Name))).
 -define(RAFT_SNAPSHOT_PATH(Table, Partition, Name), ?RAFT_SNAPSHOT_PATH(?RAFT_PARTITION_PATH(Table, Partition), Name)).

src/wa_raft_storage.erl

@@ -31,6 +31,8 @@
     open_snapshot/3,
     create_snapshot/1,
     create_snapshot/2,
+    create_witness_snapshot/1,
+    create_witness_snapshot/2,
     make_empty_snapshot/5,
     delete_snapshot/2
 ]).
@@ -272,6 +274,14 @@
 -callback storage_make_empty_snapshot(Name :: atom(), Identifier :: #raft_identifier{}, Path :: file:filename(), Position :: wa_raft_log:log_pos(), Config :: wa_raft_server:config(), Data :: dynamic()) -> ok | error().
 -optional_callback([storage_make_empty_snapshot/6]).
 
+%% Create a new witness snapshot at the provided path, containing only the current
+%% position in storage, configuration, and virtual partition range information.
+%% The snapshot will be empty (without actual storage data) but will retain all
+%% necessary metadata. When loaded, this witness snapshot will reflect the exact
+%% position state of the original storage without the storage contents.
+-callback storage_create_witness_snapshot(Path :: file:filename(), Handle :: storage_handle()) -> ok | error().
+-optional_callback([storage_create_witness_snapshot/6]).
+
 %%-----------------------------------------------------------------------------
 %% RAFT Storage - Types
 %%-----------------------------------------------------------------------------
@@ -358,6 +368,14 @@ create_snapshot(ServiceRef) ->
 create_snapshot(ServiceRef, Name) ->
     gen_server:call(ServiceRef, {snapshot_create, Name}, ?RAFT_STORAGE_CALL_TIMEOUT()).
 
+-spec create_witness_snapshot(ServiceRef :: pid() | atom()) -> {ok, Pos :: wa_raft_log:log_pos()} | error().
+create_witness_snapshot(ServiceRef) ->
+    gen_server:call(ServiceRef, snapshot_create_witness, ?RAFT_STORAGE_CALL_TIMEOUT()).
+
+-spec create_witness_snapshot(ServiceRef :: pid() | atom(), Name :: string()) -> {ok, Pos :: wa_raft_log:log_pos()} | error().
+create_witness_snapshot(ServiceRef, Name) ->
+    gen_server:call(ServiceRef, {snapshot_create_witness, Name}, ?RAFT_STORAGE_CALL_TIMEOUT()).
+
 -spec make_empty_snapshot(ServiceRef :: pid() | atom(), Path :: file:filename(), Position :: wa_raft_log:log_pos(), Config :: wa_raft_server:config(), Data :: term()) -> ok | error().
 make_empty_snapshot(ServiceRef, Path, Position, Config, Data) ->
     gen_server:call(ServiceRef, {make_empty_snapshot, Path, Position, Config, Data}).
@@ -450,8 +468,10 @@ init(#raft_options{table = Table, partition = Partition, identifier = Identifier
         open |
         {read, Op :: wa_raft_acceptor:command()} |
         snapshot_create |
+        snapshot_create_witness |
         status |
         {snapshot_create, Name :: string()} |
+        {snapshot_create_witness, Name :: string()} |
         {snapshot_open, Path :: file:filename(), LastAppliedPos :: wa_raft_log:log_pos()} |
         {make_empty_snapshot, Path :: file:filename(), Position :: wa_raft_log:log_pos(), Config :: wa_raft_server:config(), Data :: term()} |
         position |
@@ -475,6 +495,18 @@ handle_call({snapshot_create, Name}, _From, #state{last_applied = #raft_log_pos{
             {reply, Error, State}
     end;
 
+handle_call(snapshot_create_witness, From, #state{last_applied = #raft_log_pos{index = LastIndex, term = LastTerm}} = State) ->
+    Name = ?WITNESS_SNAPSHOT_NAME(LastIndex, LastTerm),
+    handle_call({snapshot_create, Name}, From, State);
+
+handle_call({snapshot_create_witness, Name}, _From, #state{last_applied = LastApplied} = State) ->
+    case create_witness_snapshot_impl(Name, State) of
+        ok ->
+            {reply, {ok, LastApplied}, State};
+        {error, _} = Error ->
+            {reply, Error, State}
+    end;
+
 handle_call({snapshot_open, SnapshotPath, LogPos}, _From, #state{name = Name, module = Module, handle = Handle, last_applied = LastApplied} = State) ->
     ?LOG_NOTICE("Storage[~0p] replacing storage at ~0p with snapshot at ~0p.", [Name, LastApplied, LogPos], #{domain => [whatsapp, wa_raft]}),
     case Module:storage_open_snapshot(SnapshotPath, LogPos, Handle) of
@@ -637,6 +669,22 @@ create_snapshot_impl(SnapName, #state{name = Name, root_dir = RootDir, module =
             Module:storage_create_snapshot(SnapshotPath, Handle)
     end.
 
+-spec create_witness_snapshot_impl(SnapName :: string(), Storage :: #state{}) -> ok | error().
+create_witness_snapshot_impl(SnapName, #state{name = Name, root_dir = RootDir, module = Module, handle = Handle} = State) ->
+    SnapshotPath = filename:join(RootDir, SnapName),
+    case filelib:is_dir(SnapshotPath) of
+        true ->
+            ?LOG_NOTICE("Snapshot ~s for ~p already exists. Skipping witness snapshot creation.", [SnapName, Name], #{domain => [whatsapp, wa_raft]}),
+            ok;
+        false ->
+            cleanup_snapshots(State),
+            ?LOG_NOTICE("Create witness snapshot ~s for ~p.", [SnapName, Name], #{domain => [whatsapp, wa_raft]}),
+            case erlang:function_exported(Module, storage_create_witness_snapshot, 2) of
+                true -> Module:storage_create_witness_snapshot(SnapshotPath, Handle);
+                false -> {error, not_supported}
+            end
+    end.
+
 -define(MAX_RETAINED_SNAPSHOT, 1).
 
 -spec cleanup_snapshots(#state{}) -> ok.

src/wa_raft_storage_ets.erl

@@ -22,6 +22,7 @@
     storage_apply_config/3,
     storage_read/3,
     storage_create_snapshot/2,
+    storage_create_witness_snapshot/2,
     storage_open_snapshot/3,
     storage_make_empty_snapshot/6
 ]).
@@ -98,8 +99,17 @@ storage_apply({delete, _Table, Key}, Position, #state{storage = Storage} = State
     LogPos :: wa_raft_log:log_pos(),
     State :: #state{}
 ) -> {ok | wa_raft_storage:error(), #state{}}.
-storage_apply_config(Config, LogPos, #state{storage = Storage} = State) ->
-    true = ets:insert(Storage, [{{?METADATA_TAG, config}, {LogPos, Config}}, {?POSITION_TAG, LogPos}]),
+storage_apply_config(Config, LogPos, State) ->
+    storage_apply_config(Config, LogPos, LogPos, State).
+
+-spec storage_apply_config(
+    Config :: wa_raft_server:config(),
+    ConfigPos :: wa_raft_log:log_pos(),
+    LogPos :: wa_raft_log:log_pos(),
+    State :: #state{}
+) -> {ok | wa_raft_storage:error(), #state{}}.
+storage_apply_config(Config, ConfigPos, LogPos, #state{storage = Storage} = State) ->
+    true = ets:insert(Storage, [{{?METADATA_TAG, config}, {ConfigPos, Config}}, {?POSITION_TAG, LogPos}]),
     {ok, State}.
 
 -spec storage_read(Command :: wa_raft_acceptor:command(), Position :: wa_raft_log:log_pos(), State :: #state{}) -> ok | {ok, Value :: dynamic()} | not_found.
@@ -118,6 +128,12 @@ storage_create_snapshot(SnapshotPath, #state{storage = Storage}) ->
         {error, Reason} -> {error, Reason}
     end.
 
+-spec storage_create_witness_snapshot(file:filename(), #state{}) -> ok | wa_raft_storage:error().
+storage_create_witness_snapshot(SnapshotPath, #state{name = Name, table = Table, partition = Partition} = State) ->
+    {ok, ConfigPosition, Config} = storage_config(State),
+    SnapshotPosition = storage_position(State),
+    storage_make_empty_snapshot(Name, Table, Partition, SnapshotPath, SnapshotPosition, Config, ConfigPosition, #{}).
+
 -spec storage_open_snapshot(file:filename(), wa_raft_log:log_pos(), #state{}) -> {ok, #state{}} | wa_raft_storage:error().
 storage_open_snapshot(SnapshotPath, SnapshotPosition, #state{storage = Storage} = State) ->
     SnapshotData = filename:join(SnapshotPath, ?SNAPSHOT_FILENAME),
@@ -136,8 +152,12 @@ storage_open_snapshot(SnapshotPath, SnapshotPosition, #state{storage = Storage}
     end.
 
 -spec storage_make_empty_snapshot(atom(), #raft_identifier{}, file:filename(), wa_raft_log:log_pos(), wa_raft_server:config(), dynamic()) -> ok | wa_raft_storage:error().
-storage_make_empty_snapshot(Name, #raft_identifier{table = Table, partition = Partition}, SnapshotPath, SnapshotPosition, Config, _Data) ->
+storage_make_empty_snapshot(Name, #raft_identifier{table = Table, partition = Partition}, SnapshotPath, SnapshotPosition, Config, Data) ->
+    storage_make_empty_snapshot(Name, Table, Partition, SnapshotPath, SnapshotPosition, Config, SnapshotPosition, Data).
+
+-spec storage_make_empty_snapshot(atom(), wa_raft:table(), wa_raft:partition(), file:filename(), wa_raft_log:log_pos(), wa_raft_server:config(), wa_raft_log:log_pos(), dynamic()) -> ok | wa_raft_storage:error().
+storage_make_empty_snapshot(Name, Table, Partition, SnapshotPath, SnapshotPosition, Config, ConfigPosition, _Data) ->
     Storage = ets:new(Name, ?OPTIONS),
     State = #state{name = Name, table = Table, partition = Partition, storage = Storage},
-    storage_apply_config(Config, SnapshotPosition, State),
+    storage_apply_config(Config, ConfigPosition, SnapshotPosition, State),
     storage_create_snapshot(SnapshotPath, State).