raft: clean up IsLocalMsg and IsResponseMsg logic

Use array indexing to clean up the code and make it constant time. Also, add a test for IsResponseMsg. Signed-off-by: Nathan VanBenschoten <[email protected]>
etcd-io · Oct 25, 2022 · 498aa12 · 498aa12
1 parent 49c2a62
commit 498aa12
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 3 deletions.
diff --git a/raft/util.go b/raft/util.go
@@ -40,13 +40,34 @@ func max(a, b uint64) uint64 {
 	return b
 }
 
+var isLocalMsg = [...]bool{
+	pb.MsgHup:         true,
+	pb.MsgBeat:        true,
+	pb.MsgUnreachable: true,
+	pb.MsgSnapStatus:  true,
+	pb.MsgCheckQuorum: true,
+}
+
+var isResponseMsg = [...]bool{
+	pb.MsgAppResp:       true,
+	pb.MsgVoteResp:      true,
+	pb.MsgHeartbeatResp: true,
+	pb.MsgUnreachable:   true,
+	pb.MsgReadIndexResp: true,
+	pb.MsgPreVoteResp:   true,
+}
+
+func isMsgInArray(msgt pb.MessageType, arr []bool) bool {
+	i := int(msgt)
+	return i < len(arr) && arr[i]
+}
+
 func IsLocalMsg(msgt pb.MessageType) bool {
-	return msgt == pb.MsgHup || msgt == pb.MsgBeat || msgt == pb.MsgUnreachable ||
-		msgt == pb.MsgSnapStatus || msgt == pb.MsgCheckQuorum
+	return isMsgInArray(msgt, isLocalMsg[:])
 }
 
 func IsResponseMsg(msgt pb.MessageType) bool {
-	return msgt == pb.MsgAppResp || msgt == pb.MsgVoteResp || msgt == pb.MsgHeartbeatResp || msgt == pb.MsgUnreachable || msgt == pb.MsgPreVoteResp
+	return isMsgInArray(msgt, isResponseMsg[:])
 }
 
 // voteResponseType maps vote and prevote message types to their corresponding responses.

diff --git a/raft/util_test.go b/raft/util_test.go
@@ -104,3 +104,37 @@ func TestIsLocalMsg(t *testing.T) {
 		}
 	}
 }
+
+func TestIsResponseMsg(t *testing.T) {
+	tests := []struct {
+		msgt       pb.MessageType
+		isResponse bool
+	}{
+		{pb.MsgHup, false},
+		{pb.MsgBeat, false},
+		{pb.MsgUnreachable, true},
+		{pb.MsgSnapStatus, false},
+		{pb.MsgCheckQuorum, false},
+		{pb.MsgTransferLeader, false},
+		{pb.MsgProp, false},
+		{pb.MsgApp, false},
+		{pb.MsgAppResp, true},
+		{pb.MsgVote, false},
+		{pb.MsgVoteResp, true},
+		{pb.MsgSnap, false},
+		{pb.MsgHeartbeat, false},
+		{pb.MsgHeartbeatResp, true},
+		{pb.MsgTimeoutNow, false},
+		{pb.MsgReadIndex, false},
+		{pb.MsgReadIndexResp, true},
+		{pb.MsgPreVote, false},
+		{pb.MsgPreVoteResp, true},
+	}
+
+	for i, tt := range tests {
+		got := IsResponseMsg(tt.msgt)
+		if got != tt.isResponse {
+			t.Errorf("#%d: got %v, want %v", i, got, tt.isResponse)
+		}
+	}
+}