From 99ae81bf6f19441d9b33e5902488e4844660f237 Mon Sep 17 00:00:00 2001
From: ethanvc <xht_xht@qq.com>
Date: Fri, 2 Sep 2022 14:19:31 +0800
Subject: [PATCH] roundrobin: optimization of the roundrobin implementation.
 (#5607)

* optimization of the roundrobin implementation.
---
 balancer/roundrobin/roundrobin.go | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/balancer/roundrobin/roundrobin.go b/balancer/roundrobin/roundrobin.go
index 274eb2f85..f7031ad22 100644
--- a/balancer/roundrobin/roundrobin.go
+++ b/balancer/roundrobin/roundrobin.go
@@ -22,7 +22,7 @@
 package roundrobin
 
 import (
-	"sync"
+	"sync/atomic"
 
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/balancer/base"
@@ -60,7 +60,7 @@ func (*rrPickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
 		// Start at a random index, as the same RR balancer rebuilds a new
 		// picker when SubConn states change, and we don't want to apply excess
 		// load to the first server in the list.
-		next: grpcrand.Intn(len(scs)),
+		next: uint32(grpcrand.Intn(len(scs))),
 	}
 }
 
@@ -69,15 +69,13 @@ type rrPicker struct {
 	// created. The slice is immutable. Each Get() will do a round robin
 	// selection from it and return the selected SubConn.
 	subConns []balancer.SubConn
-
-	mu   sync.Mutex
-	next int
+	next     uint32
 }
 
 func (p *rrPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
-	p.mu.Lock()
-	sc := p.subConns[p.next]
-	p.next = (p.next + 1) % len(p.subConns)
-	p.mu.Unlock()
+	subConnsLen := uint32(len(p.subConns))
+	nextIndex := atomic.AddUint32(&p.next, 1)
+
+	sc := p.subConns[nextIndex%subConnsLen]
 	return balancer.PickResult{SubConn: sc}, nil
 }