From bd9edcb3a8f4b708ca20a2b3a9af2417ec8ac7a1 Mon Sep 17 00:00:00 2001
From: Nathan VanBenschoten <nvanbenschoten@gmail.com>
Date: Wed, 4 Sep 2024 11:15:32 -0400
Subject: [PATCH] raft: fix flaky leader index in waitLeader function

Fixes #127413.

This commit bypasses the larger rebase in #122133 to pick up the test
flake fix in https://github.com/etcd-io/raft/pull/188. There was some
discussion in https://github.com/etcd-io/raft/issues/181 about alternatives
for fixing this test. For now, we stick with a direct cherry-pick.

Release note: None
---
 pkg/raft/rafttest/BUILD.bazel  |  1 -
 pkg/raft/rafttest/node_test.go | 15 +++++----------
 2 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/pkg/raft/rafttest/BUILD.bazel b/pkg/raft/rafttest/BUILD.bazel
index 3a25f95c50b2..d7ebb0bacae1 100644
--- a/pkg/raft/rafttest/BUILD.bazel
+++ b/pkg/raft/rafttest/BUILD.bazel
@@ -59,6 +59,5 @@ go_test(
     deps = [
         "//pkg/raft",
         "//pkg/raft/raftpb",
-        "//pkg/testutils/skip",
     ],
 )
diff --git a/pkg/raft/rafttest/node_test.go b/pkg/raft/rafttest/node_test.go
index 4323b76005b9..4a6ed4ec1e11 100644
--- a/pkg/raft/rafttest/node_test.go
+++ b/pkg/raft/rafttest/node_test.go
@@ -24,7 +24,6 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/raft"
 	pb "github.com/cockroachdb/cockroach/pkg/raft/raftpb"
-	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 )
 
 func TestBasicProgress(t *testing.T) {
@@ -54,9 +53,6 @@ func TestBasicProgress(t *testing.T) {
 }
 
 func TestRestart(t *testing.T) {
-	// TODO(pav-kv): de-flake it. See https://github.com/etcd-io/raft/issues/181.
-	skip.UnderStress(t, "the test is flaky")
-
 	peers := []raft.Peer{{ID: 1, Context: nil}, {ID: 2, Context: nil}, {ID: 3, Context: nil}, {ID: 4, Context: nil}, {ID: 5, Context: nil}}
 	nt := newRaftNetwork(1, 2, 3, 4, 5)
 
@@ -136,15 +132,14 @@ func TestPause(t *testing.T) {
 }
 
 func waitLeader(ns []*node) int {
-	var l map[pb.PeerID]struct{}
-	var lindex int
-
+	l := make(map[pb.PeerID]struct{})
 	for {
-		l = make(map[pb.PeerID]struct{})
+		clear(l)
+		lindex := -1
 
 		for i, n := range ns {
 			lead := n.Status().HardState.Lead
-			if lead != 0 {
+			if lead != raft.None {
 				l[lead] = struct{}{}
 				if n.id == lead {
 					lindex = i
@@ -152,7 +147,7 @@ func waitLeader(ns []*node) int {
 			}
 		}
 
-		if len(l) == 1 {
+		if len(l) == 1 && lindex != -1 {
 			return lindex
 		}
 	}