Add picker interface

Author: sysulq

aperture/aperture.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/hnlq715/go-loadbalance"
 	"github.com/hnlq715/go-loadbalance/p2c"
+	"github.com/hnlq715/go-loadbalance/roundrobin"
 	"google.golang.org/grpc/balancer"
 )
 
@@ -18,7 +19,7 @@ type Aperture struct {
 	remotePeers     []interface{}
 	logicalAperture int
 
-	p2c           loadbalance.P2C
+	picker        loadbalance.Picker
 	apertureIdxes []int
 }
 
@@ -35,7 +36,7 @@ func NewLeastLoadedApeture() loadbalance.Aperture {
 		localPeers:      make([]string, 0),
 		localPeersMap:   make(map[string]int),
 		remotePeers:     make([]interface{}, 0),
-		p2c:             p2c.NewLeastLoaded(),
+		picker:          p2c.NewLeastLoaded(),
 	}
 }
 
@@ -46,7 +47,18 @@ func NewPeakEwmaAperture() loadbalance.Aperture {
 		localPeers:      make([]string, 0),
 		localPeersMap:   make(map[string]int),
 		remotePeers:     make([]interface{}, 0),
-		p2c:             p2c.NewPeakEwma(),
+		picker:          p2c.NewPeakEwma(),
+	}
+}
+
+// NewSmoothRoundrobin returns an Apeture interface with smooth roundrobin
+func NewSmoothRoundrobin() loadbalance.Aperture {
+	return &Aperture{
+		logicalAperture: defaultLogicalAperture,
+		localPeers:      make([]string, 0),
+		localPeersMap:   make(map[string]int),
+		remotePeers:     make([]interface{}, 0),
+		picker:          roundrobin.NewSmoothRoundrobin(),
 	}
 }
 
@@ -82,7 +94,7 @@ func (a *Aperture) SetRemotePeers(remotePeers []interface{}) {
 
 // Next returns the next selected item
 func (a *Aperture) Next() (interface{}, func(balancer.DoneInfo)) {
-	return a.p2c.Next()
+	return a.picker.Next()
 }
 
 // List returns the remote peers for the local peer id
@@ -119,10 +131,10 @@ func (a *Aperture) rebuild() {
 	ring := NewRing(len(a.remotePeers))
 	a.apertureIdxes = ring.Slice(offset, apertureWidth)
 
-	a.p2c = p2c.NewLeastLoaded()
+	a.picker.Reset()
 	for _, apertureIdx := range a.apertureIdxes {
 		weight := ring.Weight(apertureIdx, offset, apertureWidth)
-		a.p2c.Add(a.remotePeers[apertureIdx], weight)
+		a.picker.Add(a.remotePeers[apertureIdx], weight)
 	}
 }
 

aperture/aperture_test.go

@@ -18,6 +18,18 @@ func TestAperture(t *testing.T) {
 		assert.Nil(t, item)
 	})
 
+	t.Run("1 client 1 server", func(t *testing.T) {
+		ll := aperture.NewSmoothRoundrobin()
+		ll.SetLocalPeers(nil)
+		ll.SetLocalPeers([]string{"1"})
+		ll.SetRemotePeers([]interface{}{"8"})
+		ll.SetLocalPeerID("1")
+
+		item, done := ll.Next()
+		done(balancer.DoneInfo{})
+		assert.Equal(t, "8", item)
+	})
+
 	t.Run("1 client 1 server", func(t *testing.T) {
 		ll := aperture.NewLeastLoadedApeture()
 		ll.SetLocalPeers(nil)

loadbalance.go

@@ -20,20 +20,14 @@ type Aperture interface {
 	SetRemotePeers([]interface{})
 }
 
-// P2C support p2c algorithm for load balance,
+// Picker supports multiple algorithms for load balance,
 // uses the ideas behind the "power of 2 choices"
 // to select two nodes from the underlying vector.
-type P2C interface {
+type Picker interface {
 	// Next returns next selected item.
 	Next() (interface{}, func(balancer.DoneInfo))
 	// Add a weighted item.
 	Add(interface{}, float64)
-}
-
-// RoundRobin support roundrobin algorithm for load balance
-type RoundRobin interface {
-	// Next returns next selected item.
-	Next() interface{}
-	// Add a weighted item.
-	Add(interface{}, int64)
+	// Reset this picker
+	Reset()
 }

p2c/least_loaded.go

@@ -22,7 +22,7 @@ type leastLoaded struct {
 	rand  *rand.Rand
 }
 
-func NewLeastLoaded() loadbalance.P2C {
+func NewLeastLoaded() loadbalance.Picker {
 	return &leastLoaded{
 		items: make([]*leastLoadedNode, 0),
 		rand:  rand.New(rand.NewSource(time.Now().Unix())),
@@ -33,6 +33,10 @@ func (p *leastLoaded) Add(item interface{}, weight float64) {
 	p.items = append(p.items, &leastLoadedNode{item: item, weight: weight})
 }
 
+func (p *leastLoaded) Reset() {
+	p.items = p.items[:0]
+}
+
 func (p *leastLoaded) Next() (interface{}, func(balancer.DoneInfo)) {
 	var sc, backsc *leastLoadedNode
 

p2c/pewma.go

@@ -64,7 +64,7 @@ type pewma struct {
 	rand  *rand.Rand
 }
 
-func NewPeakEwma() loadbalance.P2C {
+func NewPeakEwma() loadbalance.Picker {
 	return &pewma{
 		items: make([]*peakEwmaNode, 0),
 		rand:  rand.New(rand.NewSource(time.Now().Unix())),
@@ -75,6 +75,13 @@ func (p *pewma) Add(item interface{}, weight float64) {
 	p.items = append(p.items, &peakEwmaNode{item: item, latency: newPEWMA(), weight: weight})
 }
 
+func (p *pewma) Reset() {
+	*p = pewma{
+		items: make([]*peakEwmaNode, 0),
+		rand:  rand.New(rand.NewSource(time.Now().Unix())),
+	}
+}
+
 func (p *pewma) Next() (interface{}, func(balancer.DoneInfo)) {
 	var sc, backsc *peakEwmaNode
 	begin := time.Now().UnixNano()

roundrobin/roundrobin.go

@@ -1,6 +1,11 @@
 package roundrobin
 
-import "github.com/hnlq715/go-loadbalance"
+import (
+	"math"
+
+	"github.com/hnlq715/go-loadbalance"
+	"google.golang.org/grpc/balancer"
+)
 
 // smoothRoundrobinNode is a wrapped weighted item.
 type smoothRoundrobinNode struct {
@@ -26,28 +31,34 @@ type smoothRoundrobin struct {
 // among peers.
 // In case of { 5, 1, 1 } weights this gives the following sequence of
 // current_weight's: (a, a, b, a, c, a, a)
-func NewSmoothRoundrobin() loadbalance.RoundRobin {
+func NewSmoothRoundrobin() loadbalance.Picker {
 	return &smoothRoundrobin{}
 }
 
 // Add a weighted server.
-func (w *smoothRoundrobin) Add(item interface{}, weight int64) {
-	weighted := &smoothRoundrobinNode{Item: item, Weight: weight, EffectiveWeight: weight}
+func (w *smoothRoundrobin) Add(item interface{}, weight float64) {
+	wt := int64(math.Floor(weight))
+	weighted := &smoothRoundrobinNode{Item: item, Weight: wt, EffectiveWeight: wt}
 	w.items = append(w.items, weighted)
 	w.n++
 }
 
+func (w *smoothRoundrobin) Reset() {
+	w.items = w.items[:0]
+	w.n = 0
+}
+
 // Next returns next selected server.
-func (w *smoothRoundrobin) Next() interface{} {
+func (w *smoothRoundrobin) Next() (interface{}, func(balancer.DoneInfo)) {
 	if w.n == 0 {
-		return nil
+		return nil, func(balancer.DoneInfo) {}
 	}
 
 	if w.n == 1 {
-		return w.items[0].Item
+		return w.items[0].Item, func(balancer.DoneInfo) {}
 	}
 
-	return nextSmoothWeighted(w.items).Item
+	return nextSmoothWeighted(w.items).Item, func(balancer.DoneInfo) {}
 }
 
 // nextSmoothWeighted selects the best node through the smooth weighted roundrobin .

roundrobin/roundrobin_test.go

@@ -9,19 +9,29 @@ import (
 func TestSW_Next(t *testing.T) {
 	w := NewSmoothRoundrobin()
 
-	assert.Nil(t, w.Next())
+	s, _ := w.Next()
+	assert.Nil(t, s)
 
 	w.Add("server1", 5)
-	assert.Equal(t, "server1", w.Next().(string))
+	s, _ = w.Next()
+	assert.Equal(t, "server1", s.(string))
+
+	w.Reset()
+	s, _ = w.Next()
+	assert.Nil(t, s)
+
+	w.Add("server1", 5)
+	s, _ = w.Next()
+	assert.Equal(t, "server1", s.(string))
 
 	w.Add("server2", 2)
 	w.Add("server3", 3)
 
 	results := make(map[string]int)
 
 	for i := 0; i < 1000; i++ {
-		s := w.Next().(string)
-		results[s]++
+		s, _ := w.Next()
+		results[s.(string)]++
 	}
 
 	if results["server1"] != 500 || results["server2"] != 200 || results["server3"] != 300 {