Merge cockroachdb#138374

138374: raft: avoid allocating supportMap for every LeadSupportUntil calculation r=iskettaneh a=iskettaneh 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 Co-authored-by: Ibrahim Kettaneh <[email protected]>
mgartner · Jan 14, 2025 · 2729738 · 2729738
2 parents 4f224db + a785188
commit 2729738
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 6 deletions.
diff --git a/pkg/raft/status.go b/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")

diff --git a/pkg/raft/tracker/fortificationtracker.go b/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
 }
 

diff --git a/pkg/raft/tracker/fortificationtracker_test.go b/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