Merge pull request #10309 from tbg/fix/raft-status-allocs

raft: add (*RawNode).WithProgress

raft/rawnode.go

@@ -236,6 +236,39 @@ func (rn *RawNode) Status() *Status {
 	return &status
 }
 
+// StatusWithoutProgress returns a Status without populating the Progress field
+// (and returns the Status as a value to avoid forcing it onto the heap). This
+// is more performant if the Progress is not required. See WithProgress for an
+// allocation-free way to introspect the Progress.
+func (rn *RawNode) StatusWithoutProgress() Status {
+	return getStatusWithoutProgress(rn.raft)
+}
+
+// ProgressType indicates the type of replica a Progress corresponds to.
+type ProgressType byte
+
+const (
+	// ProgressTypePeer accompanies a Progress for a regular peer replica.
+	ProgressTypePeer ProgressType = iota
+	// ProgressTypeLearner accompanies a Progress for a learner replica.
+	ProgressTypeLearner
+)
+
+// WithProgress is a helper to introspect the Progress for this node and its
+// peers.
+func (rn *RawNode) WithProgress(visitor func(id uint64, typ ProgressType, pr Progress)) {
+	for id, pr := range rn.raft.prs {
+		pr := *pr
+		pr.ins = nil
+		visitor(id, ProgressTypePeer, pr)
+	}
+	for id, pr := range rn.raft.learnerPrs {
+		pr := *pr
+		pr.ins = nil
+		visitor(id, ProgressTypeLearner, pr)
+	}
+}
+
 // ReportUnreachable reports the given node is not reachable for the last send.
 func (rn *RawNode) ReportUnreachable(id uint64) {
 	_ = rn.raft.Step(pb.Message{Type: pb.MsgUnreachable, From: id})

raft/rawnode_test.go

@@ -16,6 +16,7 @@ package raft
 
 import (
 	"bytes"
+	"fmt"
 	"reflect"
 	"testing"
 
@@ -545,3 +546,73 @@ func TestRawNodeBoundedLogGrowthWithPartition(t *testing.T) {
 	rawNode.Advance(rd)
 	checkUncommitted(0)
 }
+
+func BenchmarkStatusProgress(b *testing.B) {
+	setup := func(members int) *RawNode {
+		peers := make([]uint64, members)
+		for i := range peers {
+			peers[i] = uint64(i + 1)
+		}
+		cfg := newTestConfig(1, peers, 3, 1, NewMemoryStorage())
+		cfg.Logger = discardLogger
+		r := newRaft(cfg)
+		r.becomeFollower(1, 1)
+		r.becomeCandidate()
+		r.becomeLeader()
+		return &RawNode{raft: r}
+	}
+
+	for _, members := range []int{1, 3, 5, 100} {
+		b.Run(fmt.Sprintf("members=%d", members), func(b *testing.B) {
+			// NB: call getStatus through rn.Status because that incurs an additional
+			// allocation.
+			rn := setup(members)
+
+			b.Run("Status", func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					_ = rn.Status()
+				}
+			})
+
+			b.Run("Status-example", func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					s := rn.Status()
+					var n uint64
+					for _, pr := range s.Progress {
+						n += pr.Match
+					}
+					_ = n
+				}
+			})
+
+			b.Run("StatusWithoutProgress", func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					_ = rn.StatusWithoutProgress()
+				}
+			})
+
+			b.Run("WithProgress", func(b *testing.B) {
+				b.ReportAllocs()
+				visit := func(uint64, ProgressType, Progress) {}
+
+				for i := 0; i < b.N; i++ {
+					rn.WithProgress(visit)
+				}
+			})
+			b.Run("WithProgress-example", func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					var n uint64
+					visit := func(_ uint64, _ ProgressType, pr Progress) {
+						n += pr.Match
+					}
+					rn.WithProgress(visit)
+					_ = n
+				}
+			})
+		})
+	}
+}

raft/status.go

@@ -32,29 +32,35 @@ type Status struct {
 	LeadTransferee uint64
 }
 
-// getStatus gets a copy of the current raft status.
-func getStatus(r *raft) Status {
+func getProgressCopy(r *raft) map[uint64]Progress {
+	prs := make(map[uint64]Progress)
+	for id, p := range r.prs {
+		prs[id] = *p
+	}
+
+	for id, p := range r.learnerPrs {
+		prs[id] = *p
+	}
+	return prs
+}
+
+func getStatusWithoutProgress(r *raft) Status {
 	s := Status{
 		ID:             r.id,
 		LeadTransferee: r.leadTransferee,
 	}
-
 	s.HardState = r.hardState()
 	s.SoftState = *r.softState()
-
 	s.Applied = r.raftLog.applied
+	return s
+}
 
+// getStatus gets a copy of the current raft status.
+func getStatus(r *raft) Status {
+	s := getStatusWithoutProgress(r)
 	if s.RaftState == StateLeader {
-		s.Progress = make(map[uint64]Progress)
-		for id, p := range r.prs {
-			s.Progress[id] = *p
-		}
-
-		for id, p := range r.learnerPrs {
-			s.Progress[id] = *p
-		}
+		s.Progress = getProgressCopy(r)
 	}
-
 	return s
 }