Update server queue length without requiring staleness change

Summary:
Introduce a separate info method for updating the server queue length and call
it unconditionally whenever we are reevaluating the liveness and staleness of the
RAFT server to improve the accuracy of the value.

This fixes an issue where the queue length value can get stuck if the RAFT server
stops receiving RPCs (for example if a custom distribution implementation withholds
RPCs when the queue is too large). Also fixes an issue where the queue length was
always being stored with the default name of the RAFT partition.

Reviewed By: jaher, SarahDesouky

Differential Revision:
D86316661

Privacy Context Container: L1384697

fbshipit-source-id: c1c91dea7dd5a1914565803bbfb01a168a104469

Author: hsun324

src/wa_raft_info.erl

@@ -30,7 +30,9 @@
     set_membership/3,
     set_live/3,
     set_stale/3,
-    set_state/3
+    set_state/3,
+    set_message_queue_length/1,
+    set_message_queue_length/2
 ]).
 
 %% Local RAFT server's current FSM state
@@ -132,10 +134,17 @@ set_live(Table, Partition, Live) ->
 
 -spec set_stale(wa_raft:table(), wa_raft:partition(), boolean()) -> true.
 set_stale(Table, Partition, Stale) ->
-    set(?RAFT_STALE_KEY(Table, Partition), Stale),
-    {message_queue_len, MsgQLen} = process_info(self(), message_queue_len),
-    set(?RAFT_MSG_QUEUE_LENGTH_KEY(wa_raft_server:default_name(Table, Partition)), MsgQLen).
+    set(?RAFT_STALE_KEY(Table, Partition), Stale).
 
 -spec set_membership(wa_raft:table(), wa_raft:partition(), wa_raft_server:membership()) -> true.
 set_membership(Table, Partition, Membership) ->
     set(?RAFT_MEMBERSHIP_KEY(Table, Partition), Membership).
+
+-spec set_message_queue_length(Name :: atom()) -> true.
+set_message_queue_length(Name) ->
+    {message_queue_len, Length} = process_info(self(), message_queue_len),
+    set_message_queue_length(Name, Length).
+
+-spec set_message_queue_length(Name :: atom(), Length :: non_neg_integer()) -> true.
+set_message_queue_length(Name, Length) ->
+    set(?RAFT_MSG_QUEUE_LENGTH_KEY(Name), Length).

src/wa_raft_server.erl

@@ -688,12 +688,13 @@ callback_mode() ->
     [state_functions, state_enter].
 
 -spec terminate(Reason :: term(), State :: state(), Data :: #raft_state{}) -> ok.
-terminate(Reason, State, #raft_state{table = Table, partition = Partition} = Data) ->
+terminate(Reason, State, #raft_state{name = Name, table = Table, partition = Partition} = Data) ->
     ?SERVER_LOG_NOTICE(State, Data, "terminating due to ~0P", [Reason, 20]),
     wa_raft_durable_state:sync(Data),
     wa_raft_info:delete_state(Table, Partition),
     wa_raft_info:set_live(Table, Partition, false),
     wa_raft_info:set_stale(Table, Partition, true),
+    wa_raft_info:set_message_queue_length(Name, 0),
     ok.
 
 %%------------------------------------------------------------------------------
@@ -2657,14 +2658,16 @@ check_stale_upon_entry(witness, #raft_state{table = Table, partition = Partition
     wa_raft_info:set_stale(Table, Partition, true),
     ok;
 %% Leaders are always live and never stale upon entry.
-check_stale_upon_entry(leader, #raft_state{table = Table, partition = Partition}) ->
+check_stale_upon_entry(leader, #raft_state{name = Name, table = Table, partition = Partition}) ->
     wa_raft_info:set_live(Table, Partition, true),
     wa_raft_info:set_stale(Table, Partition, false),
+    wa_raft_info:set_message_queue_length(Name),
     ok;
 %% Stalled and disabled servers are never live and always stale.
-check_stale_upon_entry(_, #raft_state{table = Table, partition = Partition}) ->
+check_stale_upon_entry(_, #raft_state{name = Name, table = Table, partition = Partition}) ->
     wa_raft_info:set_live(Table, Partition, false),
     wa_raft_info:set_stale(Table, Partition, true),
+    wa_raft_info:set_message_queue_length(Name),
     ok.
 
 %% Set a new current term and voted-for peer and clear any state that is associated with the previous term.
@@ -3405,6 +3408,7 @@ check_follower_liveness(
     State,
     #raft_state{
         application = App,
+        name = Name,
         table = Table,
         partition = Partition,
         leader_heartbeat_ts = LeaderHeartbeatTs
@@ -3428,6 +3432,7 @@ check_follower_liveness(
                 wa_raft_info:set_stale(Table, Partition, true)
             end
     end,
+    wa_raft_info:set_message_queue_length(Name),
     ok.
 
 %% Check the state of a follower or witness's last applied log entry versus the leader's
@@ -3439,6 +3444,7 @@ check_follower_lagging(
     LeaderCommit,
     #raft_state{
         application = App,
+        name = Name,
         table = Table,
         partition = Partition,
         last_applied = LastApplied
@@ -3462,6 +3468,7 @@ check_follower_lagging(
                 wa_raft_info:set_stale(Table, Partition, false)
         end
     end,
+    wa_raft_info:set_message_queue_length(Name),
 
     ok.
 
@@ -3472,6 +3479,7 @@ check_follower_lagging(
 check_leader_liveness(
     #raft_state{
         application = App,
+        name = Name,
         table = Table,
         partition = Partition,
         heartbeat_response_ts = HeartbeatResponse
@@ -3502,7 +3510,10 @@ check_leader_liveness(
         false ->
             ?SERVER_LOG_NOTICE(leader, State, "is no longer stale after heartbeat quorum age drops to ~0p ms < ~0p ms max", [QuorumAge, MaxAge]),
             wa_raft_info:set_stale(Table, Partition, ShouldBeStale)
-    end.
+    end,
+
+    % Update message queue length
+    wa_raft_info:set_message_queue_length(Name).
 
 %% Based on information that the leader has available as a result of heartbeat replies, attempt
 %% to discern what the best subsequent replication mode would be for this follower.