support set

go.sum

@@ -79,9 +79,11 @@ google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzi
 google.golang.org/protobuf v1.23.0 h1:4MY060fB1DLGMB/7MBTLnwQUY6+F09GEiz6SsrNqyzM=
 google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gotest.tools v2.2.0+incompatible h1:VsBPFP1AI068pPrMxtb/S8Zkgf9xEmTLJjfM+P5UIEo=
 gotest.tools v2.2.0+incompatible/go.mod h1:DsYFclhRJ6vuDpmuTbkuFWG+y2sxOXAzmJt81HFBacw=
 honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=

loadbalance.go

@@ -20,6 +20,27 @@ type Aperture interface {
 	SetRemotePeers([]interface{})
 }
 
+// SetInfo contains region, zone and set
+type SetInfo struct {
+	// Name, app name defined as set
+	Name string
+	// Region, like `bj(beijing)` or `sh(shanghai)`
+	Region string
+	// UnitName, unit name defined as subsets
+	UnitName string
+}
+
+// Set supports divide remote peers into subsets
+// based on region, zone and set info
+type Set interface {
+	// Next returns next selected item.
+	Next() (interface{}, func(balancer.DoneInfo))
+	// Add a weighted item with set info.
+	Add(interface{}, float64, SetInfo)
+	// Reset this picker
+	Reset()
+}
+
 // Picker supports multiple algorithms for load balance,
 // uses the ideas behind the "power of 2 choices"
 // to select two nodes from the underlying vector.

set/set.go

@@ -0,0 +1,45 @@
+package set
+
+import (
+	"github.com/hnlq715/go-loadbalance"
+	"github.com/hnlq715/go-loadbalance/roundrobin"
+	"google.golang.org/grpc/balancer"
+)
+
+type Set struct {
+	info   loadbalance.SetInfo
+	picker loadbalance.Picker
+}
+
+func New(info loadbalance.SetInfo) loadbalance.Set {
+	return &Set{
+		info:   info,
+		picker: roundrobin.NewSmoothRoundrobin(),
+	}
+}
+
+func (s *Set) Next() (interface{}, func(balancer.DoneInfo)) {
+	return s.picker.Next()
+}
+
+func (s *Set) Add(item interface{}, weigth float64, info loadbalance.SetInfo) {
+	if info.Name != s.info.Name {
+		return
+	}
+
+	if info.Region != s.info.Region {
+		return
+	}
+
+	if s.info.UnitName != "*" {
+		if info.UnitName != s.info.UnitName {
+			return
+		}
+	}
+
+	s.picker.Add(item, weigth)
+}
+
+func (s *Set) Reset() {
+	s.picker.Reset()
+}

set/set_test.go

@@ -0,0 +1,110 @@
+package set_test
+
+import (
+	"testing"
+
+	"github.com/hnlq715/go-loadbalance"
+	"github.com/hnlq715/go-loadbalance/set"
+	"gotest.tools/assert"
+)
+
+func TestSet(t *testing.T) {
+	t.Run("same set", func(t *testing.T) {
+		s := set.New(loadbalance.SetInfo{
+			Name:     "app",
+			Region:   "bj",
+			UnitName: "01",
+		})
+
+		s.Add(1, 1, loadbalance.SetInfo{
+			Name:     "app",
+			Region:   "bj",
+			UnitName: "01",
+		})
+
+		item, _ := s.Next()
+		assert.Equal(t, 1, item)
+
+		s.Reset()
+
+		item, _ = s.Next()
+		assert.Equal(t, nil, item)
+	})
+
+	t.Run("different region", func(t *testing.T) {
+		s := set.New(loadbalance.SetInfo{
+			Name:     "app",
+			Region:   "bj",
+			UnitName: "01",
+		})
+
+		s.Add(1, 1, loadbalance.SetInfo{
+			Name:     "app",
+			Region:   "sh",
+			UnitName: "02",
+		})
+
+		item, _ := s.Next()
+		assert.Equal(t, nil, item)
+	})
+
+	t.Run("different name", func(t *testing.T) {
+		s := set.New(loadbalance.SetInfo{
+			Name:     "app01",
+			Region:   "bj",
+			UnitName: "01",
+		})
+
+		s.Add(1, 1, loadbalance.SetInfo{
+			Name:     "app02",
+			Region:   "bj",
+			UnitName: "01",
+		})
+
+		item, _ := s.Next()
+		assert.Equal(t, nil, item)
+	})
+
+	t.Run("different set", func(t *testing.T) {
+		s := set.New(loadbalance.SetInfo{
+			Name:     "app",
+			Region:   "bj",
+			UnitName: "01",
+		})
+
+		s.Add(1, 1, loadbalance.SetInfo{
+			Name:     "app",
+			Region:   "bj",
+			UnitName: "02",
+		})
+
+		item, _ := s.Next()
+		assert.Equal(t, nil, item)
+	})
+
+	t.Run("* set", func(t *testing.T) {
+		s := set.New(loadbalance.SetInfo{
+			Name:     "app",
+			Region:   "bj",
+			UnitName: "*",
+		})
+
+		s.Add(1, 1, loadbalance.SetInfo{
+			Name:     "app",
+			Region:   "bj",
+			UnitName: "*",
+		})
+
+		s.Add(2, 1, loadbalance.SetInfo{
+			Name:     "app",
+			Region:   "bj",
+			UnitName: "01",
+		})
+
+		item, _ := s.Next()
+		assert.Equal(t, 1, item)
+
+		item, _ = s.Next()
+		assert.Equal(t, 2, item)
+	})
+}