From 86c8928223d0aec1880640ef1711e14c94cc13e3 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 9a3633db6e3b..59500c9dd09d 100644
--- a/pkg/raft/rafttest/BUILD.bazel
+++ b/pkg/raft/rafttest/BUILD.bazel
@@ -52,6 +52,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 b0345c2a8892..4f94a1c1c3c5 100644
--- a/pkg/raft/rafttest/node_test.go
+++ b/pkg/raft/rafttest/node_test.go
@@ -23,7 +23,6 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/raft"
-	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 )
 
 func TestBasicProgress(t *testing.T) {
@@ -53,9 +52,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)
 
@@ -135,15 +131,14 @@ func TestPause(t *testing.T) {
 }
 
 func waitLeader(ns []*node) int {
-	var l map[uint64]struct{}
-	var lindex int
-
+	l := make(map[uint64]struct{})
 	for {
-		l = make(map[uint64]struct{})
+		clear(l)
+		lindex := -1
 
 		for i, n := range ns {
 			lead := n.Status().SoftState.Lead
-			if lead != 0 {
+			if lead != raft.None {
 				l[lead] = struct{}{}
 				if n.id == lead {
 					lindex = i
@@ -151,7 +146,7 @@ func waitLeader(ns []*node) int {
 			}
 		}
 
-		if len(l) == 1 {
+		if len(l) == 1 && lindex != -1 {
 			return lindex
 		}
 	}