Merge #138228

138228: raft: properly setup BenchmarkStatus r=iskettaneh a=iskettaneh This commit properly sets up BenchmarkStatus so that the leader organically steps up and knows about all the followers. This way, calls like LeadSupportUntil will be actually calculated. References: #137264 Release note: None Co-authored-by: Ibrahim Kettaneh <[email protected]>
cockroachdb · Jan 6, 2025 · 7d91020 · 7d91020
2 parents 884b75a + 2f2d495
commit 7d91020
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 7 deletions.
diff --git a/pkg/raft/raft_test.go b/pkg/raft/raft_test.go
@@ -5217,6 +5217,7 @@ func SetRandomizedElectionTimeout(r *RawNode, v int64) {
 // testConfigModifiers allows callers to optionally modify newTestConfig.
 type testConfigModifiers struct {
 	testingStoreLiveness raftstoreliveness.StoreLiveness
+	testingLogger        raftlogger.Logger
 }
 
 // testConfigModifierOpt is the type of an optional parameter to newTestConfig
@@ -5235,6 +5236,13 @@ func withStoreLiveness(storeLiveness raftstoreliveness.StoreLiveness) testConfig
 	}
 }
 
+// withLogger explicitly uses the supplied raft logger.
+func withLogger(logger raftlogger.Logger) testConfigModifierOpt {
+	return func(modifier *testConfigModifiers) {
+		modifier.testingLogger = logger
+	}
+}
+
 func newTestConfig(
 	id pb.PeerID, election, heartbeat int64, storage Storage, opts ...testConfigModifierOpt,
 ) *Config {
@@ -5248,6 +5256,14 @@ func newTestConfig(
 	} else {
 		storeLiveness = raftstoreliveness.AlwaysLive{}
 	}
+
+	var logger raftlogger.Logger
+	if modifiers.testingLogger != nil {
+		logger = modifiers.testingLogger
+	} else {
+		logger = raftlogger.DefaultRaftLogger
+	}
+
 	return &Config{
 		ID:              id,
 		ElectionTick:    election,
@@ -5256,6 +5272,7 @@ func newTestConfig(
 		MaxSizePerMsg:   noLimit,
 		MaxInflightMsgs: 256,
 		StoreLiveness:   storeLiveness,
+		Logger:          logger,
 		CRDBVersion:     cluster.MakeTestingClusterSettings().Version,
 		Metrics:         NewMetrics(),
 	}

diff --git a/pkg/raft/rawnode_test.go b/pkg/raft/rawnode_test.go
@@ -887,13 +887,16 @@ func BenchmarkStatus(b *testing.B) {
 		for i := range peers {
 			peers[i] = pb.PeerID(i + 1)
 		}
-		cfg := newTestConfig(1, 3, 1, newTestMemoryStorage(withPeers(peers...)))
-		cfg.Logger = raftlogger.DiscardLogger
-		r := newRaft(cfg)
-		r.becomeFollower(1, 1)
-		r.becomeCandidate()
-		r.becomeLeader()
-		return &RawNode{raft: r}
+
+		raftPeers := make([]stateMachine, members)
+		for i := range raftPeers {
+			raftPeers[i] = newTestRaft(pb.PeerID(i+1), 10, 1, newTestMemoryStorage(withPeers(peers...)),
+				withLogger(raftlogger.DiscardLogger))
+		}
+
+		tt := newNetworkWithConfig(nil, raftPeers...)
+		tt.send(pb.Message{From: 1, To: 1, Type: pb.MsgHup})
+		return &RawNode{raft: raftPeers[0].(*raft)}
 	}
 
 	for _, members := range []int{1, 3, 5, 100} {