raft: avoid allocating supportMap for every LeadSupportUntil calculation

This adds supportExpMap that hangs off the fortificationTracker. This is
useful to avoid allocating the map for every LeadSupportUntil
calculation. This is now possible since all the calls to
ComputeLeadSupportUntil are done with a lock held.

Note that this doesn't seem to improve the performance of
ComputeLeadSupportUntil when the number of voters are low. The reason
is that there is a go compiler optimization that seems to  make an
on-stack map allocation if the size is small.

References: cockroachdb#137264

Release note: None

pkg/raft/status.go

@@ -98,7 +98,6 @@ func getBasicStatus(r *raft) BasicStatus {
 	// NOTE: we assign to LeadSupportUntil even if RaftState is not currently
 	// StateLeader. The replica may have been the leader and stepped down to a
 	// follower before its lead support ran out.
-	//s.LeadSupportUntil = hlc.Timestamp{}
 	s.LeadSupportUntil = r.fortificationTracker.LeadSupportUntil(r.state)
 
 	assertTrue((s.RaftState == pb.StateLeader) == (s.Lead == r.id), "inconsistent lead / raft state")

pkg/raft/tracker/fortificationtracker.go

@@ -100,6 +100,12 @@ type FortificationTracker struct {
 	// LeadSupportUntil is called by every request trying to evaluate the lease's
 	// status.
 	computedLeadSupportUntil hlc.Timestamp
+
+	// supportExpMap is a map that hangs off the fortificationTracker to prevent
+	// allocations on every call to ComputeLeadSupportUntil. It stores the
+	// SupportFrom expiration timestamps for the voters, which are then used to
+	// calculate the LeadSupportUntil.
+	supportExpMap map[pb.PeerID]hlc.Timestamp
 }
 
 // NewFortificationTracker initializes a FortificationTracker.
@@ -112,6 +118,7 @@ func NewFortificationTracker(
 		fortification: map[pb.PeerID]pb.Epoch{},
 		votersSupport: map[pb.PeerID]bool{},
 		logger:        logger,
+		supportExpMap: map[pb.PeerID]hlc.Timestamp{},
 	}
 	return &st
 }
@@ -211,9 +218,7 @@ func (ft *FortificationTracker) ComputeLeadSupportUntil(state pb.StateType) hlc.
 		return ft.computedLeadSupportUntil // fast-path for no fortification
 	}
 
-	// TODO(ibrahim): avoid this map allocation as we're calling LeadSupportUntil
-	// on every tick, on every new fortification, and on config changes.
-	supportExpMap := make(map[pb.PeerID]hlc.Timestamp)
+	clear(ft.supportExpMap)
 	ft.config.Voters.Visit(func(id pb.PeerID) {
 		if supportEpoch, ok := ft.fortification[id]; ok {
 			curEpoch, curExp := ft.storeLiveness.SupportFrom(id)
@@ -223,11 +228,11 @@ func (ft *FortificationTracker) ComputeLeadSupportUntil(state pb.StateType) hlc.
 			// supporting the leader's store at the epoch in the MsgFortifyLeaderResp
 			// message.
 			if curEpoch == supportEpoch {
-				supportExpMap[id] = curExp
+				ft.supportExpMap[id] = curExp
 			}
 		}
 	})
-	ft.computedLeadSupportUntil = ft.config.Voters.LeadSupportExpiration(supportExpMap)
+	ft.computedLeadSupportUntil = ft.config.Voters.LeadSupportExpiration(ft.supportExpMap)
 	return ft.computedLeadSupportUntil
 }
 

pkg/raft/tracker/fortificationtracker_test.go

@@ -6,6 +6,7 @@
 package tracker
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/raft/quorum"
@@ -905,6 +906,44 @@ func TestConfigChangeSafe(t *testing.T) {
 	}
 }
 
+// BenchmarkComputeLeadSupportUntil keeps calling ComputeLeadSupportUntil() for
+// different number of members.
+func BenchmarkComputeLeadSupportUntil(b *testing.B) {
+	ts := func(ts int64) hlc.Timestamp {
+		return hlc.Timestamp{
+			WallTime: ts,
+		}
+	}
+
+	for _, members := range []int{1, 3, 5, 7, 100} {
+		b.Run(fmt.Sprintf("members=%d", members), func(b *testing.B) {
+			// Prepare the mock store liveness, and record fortifications.
+			livenessMap := map[pb.PeerID]mockLivenessEntry{}
+			for i := 1; i <= members; i++ {
+				livenessMap[pb.PeerID(i)] = makeMockLivenessEntry(10, ts(100))
+			}
+
+			mockLiveness := makeMockStoreLiveness(livenessMap)
+			cfg := quorum.MakeEmptyConfig()
+			for i := 1; i <= members; i++ {
+				cfg.Voters[0][pb.PeerID(i)] = struct{}{}
+			}
+
+			ft := NewFortificationTracker(&cfg, mockLiveness, raftlogger.DiscardLogger)
+			for i := 1; i <= members; i++ {
+				ft.RecordFortification(pb.PeerID(i), 10)
+			}
+
+			// The benchmark actually starts here.
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				ft.ComputeLeadSupportUntil(pb.StateLeader)
+			}
+		})
+	}
+}
+
 type mockLivenessEntry struct {
 	epoch pb.Epoch
 	ts    hlc.Timestamp