Merge pull request #836 from blight19/master

Smmoth_round_roubine:Performance improvement
smallnest · Jan 10, 2024 · af9bc96 · af9bc96
2 parents 24009f7 + 98b076c
commit af9bc96
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 17 deletions.
diff --git a/client/selector.go b/client/selector.go
@@ -1,6 +1,7 @@
 package client
 
 import (
+	"container/ring"
 	"context"
 	"math"
 	"math/rand"
@@ -55,7 +56,7 @@ func newRandomSelector(servers map[string]string) Selector {
 	return &randomSelector{servers: ss}
 }
 
-func (s randomSelector) Select(ctx context.Context, servicePath, serviceMethod string, args interface{}) string {
+func (s *randomSelector) Select(ctx context.Context, servicePath, serviceMethod string, args interface{}) string {
 	ss := s.servers
 	if len(ss) == 0 {
 		return ""
@@ -111,42 +112,79 @@ func (s *roundRobinSelector) UpdateServer(servers map[string]string) {
 
 // weightedRoundRobinSelector selects servers with weighted.
 type weightedRoundRobinSelector struct {
-	servers []*Weighted
+	servers     []*Weighted
+	totalWeight int
+	rr          *ring.Ring
 }
 
 func newWeightedRoundRobinSelector(servers map[string]string) Selector {
 	ss := createWeighted(servers)
-	return &weightedRoundRobinSelector{servers: ss}
+	s := &weightedRoundRobinSelector{servers: ss}
+	s.buildRing()
+	return s
 }
 
 func (s *weightedRoundRobinSelector) Select(ctx context.Context, servicePath, serviceMethod string, args interface{}) string {
 	ss := s.servers
 	if len(ss) == 0 {
 		return ""
 	}
-	w := nextWeighted(ss)
-	if w == nil {
-		return ""
-	}
-	return w.Server
+	val := s.rr.Value
+	s.rr = s.rr.Next()
+	return val.(*Weighted).Server
+
 }
 
 func (s *weightedRoundRobinSelector) UpdateServer(servers map[string]string) {
 	ss := createWeighted(servers)
+	s.buildRing()
 	s.servers = ss
 }
-
+func (s *weightedRoundRobinSelector) buildRing() {
+	s.totalWeight = 0
+	for _, w := range s.servers {
+		s.totalWeight += w.Weight
+	}
+	s.rr = ring.New(s.totalWeight)
+	for i := 0; i < s.totalWeight; i++ {
+		n := s.next()
+		s.rr.Value = n
+		s.rr = s.rr.Next()
+	}
+}
+func (s *weightedRoundRobinSelector) next() *Weighted {
+	if len(s.servers) == 0 {
+		return nil
+	}
+	n := len(s.servers)
+	if n == 0 {
+		return nil
+	}
+	if n == 1 {
+		return s.servers[0]
+	}
+	flag := 0
+	m := 0
+	for i := 0; i < n; i++ {
+		s.servers[i].CurrentWeight += s.servers[i].Weight
+		if s.servers[i].CurrentWeight > m {
+			m = s.servers[i].CurrentWeight
+			flag = i
+		}
+	}
+	s.servers[flag].CurrentWeight -= s.totalWeight
+	return s.servers[flag]
+}
 func createWeighted(servers map[string]string) []*Weighted {
 	ss := make([]*Weighted, 0, len(servers))
 	for k, metadata := range servers {
-		w := &Weighted{Server: k, Weight: 1, EffectiveWeight: 1}
+		w := &Weighted{Server: k, Weight: 1}
 
 		if v, err := url.ParseQuery(metadata); err == nil {
 			ww := v.Get("weight")
 			if ww != "" {
 				if weight, err := strconv.Atoi(ww); err == nil {
 					w.Weight = weight
-					w.EffectiveWeight = weight
 				}
 			}
 		}
@@ -177,19 +215,19 @@ func newGeoSelector(servers map[string]string, latitude, longitude float64) Sele
 	return &geoSelector{servers: ss, Latitude: latitude, Longitude: longitude, r: r}
 }
 
-func (s geoSelector) Select(ctx context.Context, servicePath, serviceMethod string, args interface{}) string {
+func (s *geoSelector) Select(ctx context.Context, servicePath, serviceMethod string, args interface{}) string {
 	if len(s.servers) == 0 {
 		return ""
 	}
 
 	var server []string
-	min := math.MaxFloat64
+	minNum := math.MaxFloat64
 	for _, gs := range s.servers {
 		d := getDistanceFrom(s.Latitude, s.Longitude, gs.Latitude, gs.Longitude)
-		if d < min {
+		if d < minNum {
 			server = []string{gs.Server}
-			min = d
-		} else if d == min {
+			minNum = d
+		} else if d == minNum {
 			server = append(server, gs.Server)
 		}
 	}
@@ -253,7 +291,7 @@ func newConsistentHashSelector(servers map[string]string) Selector {
 	return &consistentHashSelector{servers: ss, h: h}
 }
 
-func (s consistentHashSelector) Select(ctx context.Context, servicePath, serviceMethod string, args interface{}) string {
+func (s *consistentHashSelector) Select(ctx context.Context, servicePath, serviceMethod string, args interface{}) string {
 	ss := s.servers
 	if len(ss) == 0 {
 		return ""

diff --git a/client/selector_test.go b/client/selector_test.go
@@ -1,6 +1,7 @@
 package client
 
 import (
+	"context"
 	"testing"
 )
 
@@ -36,3 +37,35 @@ func Test_consistentHashSelector_UpdateServer(t *testing.T) {
 		t.Errorf("UpdateServer: expected %d server but got %d", len(servers), len(s.h.All()))
 	}
 }
+
+func TestWeightedRoundRobinSelector_Select(t *testing.T) {
+	// a b a c a b a a b a c a b a
+	sers := []string{"ServerA", "ServerB", "ServerA", "ServerC", "ServerA", "ServerB", "ServerA",
+		"ServerA", "ServerB", "ServerA", "ServerC", "ServerA", "ServerB", "ServerA"}
+	servers := make(map[string]string)
+	servers["ServerA"] = "weight=4"
+	servers["ServerB"] = "weight=2"
+	servers["ServerC"] = "weight=1"
+	ctx := context.Background()
+	weightSelector := newWeightedRoundRobinSelector(servers).(*weightedRoundRobinSelector)
+
+	for i := 0; i < 14; i++ {
+		s := weightSelector.Select(ctx, "", "", nil)
+		if s != sers[i] {
+			t.Errorf("expected %s but got %s", sers[i], s)
+		}
+	}
+}
+
+func BenchmarkWeightedRoundRobinSelector_Select(b *testing.B) {
+	servers := make(map[string]string)
+	servers["ServerA"] = "weight=4"
+	servers["ServerB"] = "weight=2"
+	servers["ServerC"] = "weight=1"
+	ctx := context.Background()
+	weightSelector := newWeightedRoundRobinSelector(servers).(*weightedRoundRobinSelector)
+
+	for i := 0; i < b.N; i++ {
+		weightSelector.Select(ctx, "", "", nil)
+	}
+}