Merge #130304

130304: raft: use maps and slices packages r=nvanbenschoten a=pav-kv

This PR replaces ad-hoc map keys slices with `maps.Keys` calls, and uses `slices` package to sort slices.

Epic: none
Release note: none

Co-authored-by: Pavel Kalinnikov <[email protected]>

pkg/raft/BUILD.bazel

@@ -26,6 +26,7 @@ go_library(
         "//pkg/raft/raftstoreliveness",
         "//pkg/raft/tracker",
         "//pkg/util/hlc",
+        "@org_golang_x_exp//maps",
     ],
 )
 

pkg/raft/confchange/restore_test.go

@@ -20,7 +20,7 @@ package confchange
 import (
 	"math/rand"
 	"reflect"
-	"sort"
+	"slices"
 	"testing"
 	"testing/quick"
 
@@ -108,7 +108,7 @@ func TestRestore(t *testing.T) {
 			cs.VotersOutgoing,
 			cs.LearnersNext,
 		} {
-			sort.Slice(sl, func(i, j int) bool { return sl[i] < sl[j] })
+			slices.Sort(sl)
 		}
 
 		cs2 := chg.Config.ConfState()

pkg/raft/quorum/BUILD.bazel

@@ -15,6 +15,7 @@ go_library(
     deps = [
         "//pkg/raft/raftpb",
         "//pkg/util/hlc",
+        "@org_golang_x_exp//maps",
     ],
 )
 

pkg/raft/quorum/config.go

@@ -12,6 +12,7 @@ package quorum
 
 import (
 	"fmt"
+	"maps"
 	"strings"
 
 	pb "github.com/cockroachdb/cockroach/pkg/raft/raftpb"
@@ -101,20 +102,10 @@ func (c Config) String() string {
 
 // Clone returns a copy of the Config that shares no memory with the original.
 func (c *Config) Clone() Config {
-	clone := func(m map[pb.PeerID]struct{}) map[pb.PeerID]struct{} {
-		if m == nil {
-			return nil
-		}
-		mm := make(map[pb.PeerID]struct{}, len(m))
-		for k := range m {
-			mm[k] = struct{}{}
-		}
-		return mm
-	}
 	return Config{
-		Voters:       JointConfig{clone(c.Voters[0]), clone(c.Voters[1])},
-		Learners:     clone(c.Learners),
-		LearnersNext: clone(c.LearnersNext),
+		Voters:       JointConfig{maps.Clone(c.Voters[0]), maps.Clone(c.Voters[1])},
+		Learners:     maps.Clone(c.Learners),
+		LearnersNext: maps.Clone(c.LearnersNext),
 	}
 }
 

pkg/raft/quorum/joint.go

@@ -36,7 +36,7 @@ func (c JointConfig) String() string {
 // IDs returns a newly initialized map representing the set of voters present
 // in the joint configuration.
 func (c JointConfig) IDs() map[pb.PeerID]struct{} {
-	m := map[pb.PeerID]struct{}{}
+	m := make(map[pb.PeerID]struct{}, len(c[0])+len(c[1]))
 	for _, cc := range c {
 		for id := range cc {
 			m[id] = struct{}{}
@@ -55,12 +55,7 @@ func (c JointConfig) Describe(l AckedIndexer) string {
 // quorum. An index is jointly committed if it is committed in both constituent
 // majorities.
 func (c JointConfig) CommittedIndex(l AckedIndexer) Index {
-	idx0 := c[0].CommittedIndex(l)
-	idx1 := c[1].CommittedIndex(l)
-	if idx0 < idx1 {
-		return idx0
-	}
-	return idx1
+	return min(c[0].CommittedIndex(l), c[1].CommittedIndex(l))
 }
 
 // VoteResult takes a mapping of voters to yes/no (true/false) votes and returns

pkg/raft/quorum/majority.go

@@ -18,32 +18,28 @@
 package quorum
 
 import (
+	"cmp"
 	"fmt"
 	"math"
 	"slices"
-	"sort"
 	"strings"
 
 	pb "github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"golang.org/x/exp/maps"
 )
 
 // MajorityConfig is a set of IDs that uses majority quorums to make decisions.
 type MajorityConfig map[pb.PeerID]struct{}
 
 func (c MajorityConfig) String() string {
-	sl := make([]pb.PeerID, 0, len(c))
-	for id := range c {
-		sl = append(sl, id)
-	}
-	sort.Slice(sl, func(i, j int) bool { return sl[i] < sl[j] })
 	var buf strings.Builder
 	buf.WriteByte('(')
-	for i := range sl {
+	for i, id := range c.Slice() {
 		if i > 0 {
 			buf.WriteByte(' ')
 		}
-		fmt.Fprint(&buf, sl[i])
+		fmt.Fprint(&buf, id)
 	}
 	buf.WriteByte(')')
 	return buf.String()
@@ -72,51 +68,42 @@ func (c MajorityConfig) Describe(l AckedIndexer) string {
 		idx, ok := l.AckedIndex(id)
 		info = append(info, tup{id: id, idx: idx, ok: ok})
 	}
-
-	// Sort by index
-	sort.Slice(info, func(i, j int) bool {
-		if info[i].idx == info[j].idx {
-			return info[i].id < info[j].id
-		}
-		return info[i].idx < info[j].idx
+	// Sort by (index, ID).
+	slices.SortFunc(info, func(a, b tup) int {
+		return cmp.Or(cmp.Compare(a.idx, b.idx), cmp.Compare(a.id, b.id))
 	})
-
 	// Populate .bar.
 	for i := range info {
 		if i > 0 && info[i-1].idx < info[i].idx {
 			info[i].bar = i
 		}
 	}
-
 	// Sort by ID.
-	sort.Slice(info, func(i, j int) bool {
-		return info[i].id < info[j].id
-	})
+	slices.SortFunc(info, func(a, b tup) int { return cmp.Compare(a.id, b.id) })
 
 	var buf strings.Builder
 
 	// Print.
 	fmt.Fprint(&buf, strings.Repeat(" ", n)+"    idx\n")
-	for i := range info {
-		bar := info[i].bar
-		if !info[i].ok {
+	for _, t := range info {
+		if !t.ok {
 			fmt.Fprint(&buf, "?"+strings.Repeat(" ", n))
 		} else {
-			fmt.Fprint(&buf, strings.Repeat("x", bar)+">"+strings.Repeat(" ", n-bar))
+			fmt.Fprint(&buf, strings.Repeat("x", t.bar)+">"+strings.Repeat(" ", n-t.bar))
 		}
-		fmt.Fprintf(&buf, " %5d    (id=%d)\n", info[i].idx, info[i].id)
+		fmt.Fprintf(&buf, " %5d    (id=%d)\n", t.idx, t.id)
 	}
 	return buf.String()
 }
 
 // Slice returns the MajorityConfig as a sorted slice.
 func (c MajorityConfig) Slice() []pb.PeerID {
-	var sl []pb.PeerID
-	for id := range c {
-		sl = append(sl, id)
+	if len(c) == 0 {
+		return nil
 	}
-	sort.Slice(sl, func(i, j int) bool { return sl[i] < sl[j] })
-	return sl
+	peers := maps.Keys(c)
+	slices.Sort(peers)
+	return peers
 }
 
 // NB: A lot of logic in CommittedIndex, VoteResult, and LeadSupportExpiration

pkg/raft/raft.go

@@ -25,7 +25,7 @@ import (
 	"fmt"
 	"math"
 	"math/big"
-	"sort"
+	"slices"
 	"strings"
 	"sync"
 
@@ -35,6 +35,7 @@ import (
 	pb "github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftstoreliveness"
 	"github.com/cockroachdb/cockroach/pkg/raft/tracker"
+	"golang.org/x/exp/maps"
 )
 
 const (
@@ -1228,15 +1229,8 @@ func (r *raft) campaign(t CampaignType) {
 		voteMsg = pb.MsgVote
 		term = r.Term
 	}
-	var ids []pb.PeerID
-	{
-		idMap := r.config.Voters.IDs()
-		ids = make([]pb.PeerID, 0, len(idMap))
-		for id := range idMap {
-			ids = append(ids, id)
-		}
-		sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
-	}
+	ids := maps.Keys(r.config.Voters.IDs())
+	slices.Sort(ids)
 	for _, id := range ids {
 		if id == r.id {
 			// The candidate votes for itself and should account for this self

pkg/raft/raft_paper_test.go

@@ -31,7 +31,8 @@ package raft
 
 import (
 	"fmt"
-	"sort"
+	"slices"
+	"strings"
 	"testing"
 
 	pb "github.com/cockroachdb/cockroach/pkg/raft/raftpb"
@@ -129,7 +130,7 @@ func TestLeaderBcastBeat(t *testing.T) {
 			}
 
 			msgs := r.readMessages()
-			sort.Sort(messageSlice(msgs))
+			slices.SortFunc(msgs, cmpMessages)
 			if storeLivenessEnabled {
 				assert.Equal(t, []pb.Message{
 					{From: 1, To: 2, Term: 1, Type: pb.MsgFortifyLeader},
@@ -184,7 +185,7 @@ func testNonleaderStartElection(t *testing.T, state StateType) {
 	assert.True(t, r.electionTracker.TestingGetVotes()[r.id])
 
 	msgs := r.readMessages()
-	sort.Sort(messageSlice(msgs))
+	slices.SortFunc(msgs, cmpMessages)
 	assert.Equal(t, []pb.Message{
 		{From: 1, To: 2, Term: 2, Type: pb.MsgVote},
 		{From: 1, To: 3, Term: 2, Type: pb.MsgVote},
@@ -399,7 +400,7 @@ func TestLeaderStartReplication(t *testing.T) {
 	assert.Equal(t, li+1, r.raftLog.lastIndex())
 	assert.Equal(t, li, r.raftLog.committed)
 	msgs := r.readMessages()
-	sort.Sort(messageSlice(msgs))
+	slices.SortFunc(msgs, cmpMessages)
 	wents := []pb.Entry{{Index: li + 1, Term: 1, Data: []byte("some data")}}
 	assert.Equal(t, []pb.Message{
 		{From: 1, To: 2, Term: 1, Type: pb.MsgApp, Index: li, LogTerm: 1, Entries: wents, Commit: li, Match: li},
@@ -435,7 +436,7 @@ func TestLeaderCommitEntry(t *testing.T) {
 		{Index: li + 1, Term: 1, Data: []byte("some data")},
 	}, r.raftLog.nextCommittedEnts(true))
 	msgs := r.readMessages()
-	sort.Sort(messageSlice(msgs))
+	slices.SortFunc(msgs, cmpMessages)
 	for i, m := range msgs {
 		assert.Equal(t, pb.PeerID(i+2), m.To)
 		assert.Equal(t, pb.MsgApp, m.Type)
@@ -711,7 +712,7 @@ func TestVoteRequest(t *testing.T) {
 		}
 
 		msgs := r.readMessages()
-		sort.Sort(messageSlice(msgs))
+		slices.SortFunc(msgs, cmpMessages)
 		require.Len(t, msgs, 2, "#%d", j)
 		for i, m := range msgs {
 			assert.Equal(t, pb.MsgVote, m.Type, "#%d.%d", j, i)
@@ -797,11 +798,9 @@ func TestLeaderOnlyCommitsLogFromCurrentTerm(t *testing.T) {
 	}
 }
 
-type messageSlice []pb.Message
-
-func (s messageSlice) Len() int           { return len(s) }
-func (s messageSlice) Less(i, j int) bool { return fmt.Sprint(s[i]) < fmt.Sprint(s[j]) }
-func (s messageSlice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
+func cmpMessages(a, b pb.Message) int {
+	return strings.Compare(fmt.Sprint(a), fmt.Sprint(b))
+}
 
 func commitNoopEntry(r *raft, s *MemoryStorage) {
 	if r.state != StateLeader {

pkg/raft/raftpb/confstate.go

@@ -20,7 +20,7 @@ package raftpb
 import (
 	"fmt"
 	"reflect"
-	"sort"
+	"slices"
 )
 
 // Equivalent returns a nil error if the inputs describe the same configuration.
@@ -30,7 +30,7 @@ func (cs ConfState) Equivalent(cs2 ConfState) error {
 	orig1, orig2 := cs1, cs2
 	s := func(sl *[]PeerID) {
 		*sl = append([]PeerID(nil), *sl...)
-		sort.Slice(*sl, func(i, j int) bool { return (*sl)[i] < (*sl)[j] })
+		slices.Sort(*sl)
 	}
 
 	for _, cs := range []*ConfState{&cs1, &cs2} {

pkg/raft/tracker/BUILD.bazel

@@ -17,6 +17,7 @@ go_library(
         "//pkg/raft/raftpb",
         "//pkg/raft/raftstoreliveness",
         "//pkg/util/hlc",
+        "@org_golang_x_exp//maps",
     ],
 )
 

pkg/raft/tracker/progress.go

@@ -19,10 +19,11 @@ package tracker
 
 import (
 	"fmt"
-	"sort"
+	"slices"
 	"strings"
 
 	pb "github.com/cockroachdb/cockroach/pkg/raft/raftpb"
+	"golang.org/x/exp/maps"
 )
 
 // Progress represents a follower’s progress in the view of the leader. Leader
@@ -414,13 +415,8 @@ func MakeEmptyProgressMap() ProgressMap {
 
 // String prints the ProgressMap in sorted key order, one Progress per line.
 func (m ProgressMap) String() string {
-	ids := make([]pb.PeerID, 0, len(m))
-	for k := range m {
-		ids = append(ids, k)
-	}
-	sort.Slice(ids, func(i, j int) bool {
-		return ids[i] < ids[j]
-	})
+	ids := maps.Keys(m)
+	slices.Sort(ids)
 	var buf strings.Builder
 	for _, id := range ids {
 		fmt.Fprintf(&buf, "%d: %s\n", id, m[id])

pkg/raft/tracker/progresstracker.go

@@ -19,10 +19,10 @@ package tracker
 
 import (
 	"slices"
-	"sort"
 
 	"github.com/cockroachdb/cockroach/pkg/raft/quorum"
 	pb "github.com/cockroachdb/cockroach/pkg/raft/raftpb"
+	"golang.org/x/exp/maps"
 )
 
 // ProgressTracker tracks the progress made by each peer in the currently active
@@ -123,12 +123,8 @@ func (p *ProgressTracker) QuorumActive() bool {
 
 // VoterNodes returns a sorted slice of voters.
 func (p *ProgressTracker) VoterNodes() []pb.PeerID {
-	m := p.config.Voters.IDs()
-	nodes := make([]pb.PeerID, 0, len(m))
-	for id := range m {
-		nodes = append(nodes, id)
-	}
-	sort.Slice(nodes, func(i, j int) bool { return nodes[i] < nodes[j] })
+	nodes := maps.Keys(p.config.Voters.IDs())
+	slices.Sort(nodes)
 	return nodes
 }
 
@@ -137,10 +133,7 @@ func (p *ProgressTracker) LearnerNodes() []pb.PeerID {
 	if len(p.config.Learners) == 0 {
 		return nil
 	}
-	nodes := make([]pb.PeerID, 0, len(p.config.Learners))
-	for id := range p.config.Learners {
-		nodes = append(nodes, id)
-	}
-	sort.Slice(nodes, func(i, j int) bool { return nodes[i] < nodes[j] })
+	nodes := maps.Keys(p.config.Learners)
+	slices.Sort(nodes)
 	return nodes
 }