Smmoth_round_roubine:Performance improvement

client/selector.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"container/ring"
 	"context"
 	"math"
 	"math/rand"
@@ -111,42 +112,78 @@ 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.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
 				}
 			}
 		}

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)
+	}
+}