From 4971fd3eafbcd9956f81c97af8f23a83d2dcfb91 Mon Sep 17 00:00:00 2001
From: Andrew Mayorov <encube.ul@gmail.com>
Date: Fri, 2 Aug 2024 09:52:55 +0200
Subject: [PATCH] chore(dsraft): make shard info pterms saner and more visible

---
 .../src/emqx_ds_replication_layer_shard.erl   | 27 +++++++++++++++----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/apps/emqx_ds_builtin_raft/src/emqx_ds_replication_layer_shard.erl b/apps/emqx_ds_builtin_raft/src/emqx_ds_replication_layer_shard.erl
index 36dc654b4..bb0a9c99d 100644
--- a/apps/emqx_ds_builtin_raft/src/emqx_ds_replication_layer_shard.erl
+++ b/apps/emqx_ds_builtin_raft/src/emqx_ds_replication_layer_shard.erl
@@ -55,6 +55,8 @@
 
 -define(MEMBERSHIP_CHANGE_TIMEOUT, 30_000).
 
+-define(PTERM(DB, SHARD, KEY), {?MODULE, DB, SHARD, KEY}).
+
 %%
 
 start_link(DB, Shard, Opts) ->
@@ -164,7 +166,7 @@ local_site() ->
 
 -spec shard_info(emqx_ds:db(), emqx_ds_replication_layer:shard_id(), _Info) -> _Value.
 shard_info(DB, Shard, ready) ->
-    persistent_term:get({?MODULE, DB, Shard}, false).
+    get_shard_info(DB, Shard, ready, false).
 
 %%
 
@@ -372,8 +374,8 @@ handle_cast(_Msg, State) ->
     {noreply, State}.
 
 terminate(_Reason, {DB, Shard}) ->
-    %% FIXME
-    persistent_term:erase({?MODULE, DB, Shard}),
+    %% NOTE: Mark as not ready right away.
+    ok = erase_shard_info(DB, Shard),
     %% NOTE: Timeouts are ignored, it's a best effort attempt.
     catch prep_stop_server(DB, Shard),
     LocalServer = get_local_server(DB, Shard),
@@ -457,8 +459,7 @@ trigger_election(Server) ->
     end.
 
 announce_shard_ready(DB, Shard) ->
-    %% FIXME
-    persistent_term:put({?MODULE, DB, Shard}, true).
+    set_shard_info(DB, Shard, ready, true).
 
 server_uid(_DB, Shard) ->
     %% NOTE
@@ -472,6 +473,22 @@ server_uid(_DB, Shard) ->
 
 %%
 
+get_shard_info(DB, Shard, K, Default) ->
+    persistent_term:get(?PTERM(DB, Shard, K), Default).
+
+set_shard_info(DB, Shard, K, V) ->
+    persistent_term:put(?PTERM(DB, Shard, K), V).
+
+erase_shard_info(DB, Shard) ->
+    lists:foreach(fun(K) -> erase_shard_info(DB, Shard, K) end, [
+        ready
+    ]).
+
+erase_shard_info(DB, Shard, K) ->
+    persistent_term:erase(?PTERM(DB, Shard, K)).
+
+%%
+
 prep_stop_server(DB, Shard) ->
     prep_stop_server(DB, Shard, 5_000).