feature: add weight_picker

Author: Stone-afk

balancer/types.go

@@ -0,0 +1,15 @@
+// Copyright 2023 ecodeclub
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package balancer

balancer/wrr/weight_picker.go

@@ -0,0 +1,121 @@
+package wrr
+
+import (
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"math"
+	"sync"
+)
+
+var (
+	_ balancer.Picker    = (*Picker)(nil)
+	_ base.PickerBuilder = (*PickerBuilder)(nil)
+)
+
+type GetWeightFunc func(ci base.SubConnInfo) uint32
+
+type Option func(b *PickerBuilder)
+
+func WithGetWeightFunc(fn GetWeightFunc) Option {
+	return func(b *PickerBuilder) {
+		b.getWeightFunc = fn
+	}
+}
+
+func defaultGetWeight(ci base.SubConnInfo) uint32 {
+	md, ok := ci.Address.Metadata.(map[string]any)
+	if !ok {
+		return 10
+	}
+	weightVal := md["weight"]
+	weight, _ := weightVal.(uint32)
+	return weight
+}
+
+type PickerBuilder struct {
+	getWeightFunc GetWeightFunc
+}
+
+func NewPickerBuilder(opts ...Option) *PickerBuilder {
+	res := &PickerBuilder{
+		getWeightFunc: defaultGetWeight,
+	}
+	for _, opt := range opts {
+		opt(res)
+	}
+	return res
+}
+
+func (b *PickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
+	connections := make([]*conn, 0, len(info.ReadySCs))
+	for con, conInfo := range info.ReadySCs {
+		weight := b.getWeightFunc(conInfo)
+		connections = append(connections, &conn{
+			SubConn:         con,
+			connInfo:        conInfo,
+			weight:          weight,
+			currentWeight:   weight,
+			efficientWeight: weight,
+			name:            conInfo.Address.Addr,
+		})
+	}
+	return &Picker{
+		connections: connections,
+	}
+}
+
+type Picker struct {
+	connections []*conn
+	mutex       sync.Mutex
+}
+
+func (p *Picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	if len(p.connections) == 0 {
+		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+	}
+	var totalWeight uint32
+	var maxConn *conn
+	for _, con := range p.connections {
+		con.mutex.Lock()
+		totalWeight += con.efficientWeight
+		con.currentWeight += con.efficientWeight
+		if maxConn == nil || con.currentWeight > maxConn.currentWeight {
+			maxConn = con
+		}
+		con.mutex.Unlock()
+	}
+	if maxConn == nil {
+		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
+	}
+	maxConn.mutex.Lock()
+	maxConn.currentWeight -= totalWeight
+	maxConn.mutex.Unlock()
+	return balancer.PickResult{
+		SubConn: maxConn,
+		Done: func(info balancer.DoneInfo) {
+			maxConn.mutex.Lock()
+			defer maxConn.mutex.Unlock()
+			if info.Err != nil && maxConn.weight == 0 {
+				return
+			}
+			if info.Err == nil && maxConn.efficientWeight == math.MaxUint32 {
+				return
+			}
+			if info.Err != nil {
+				maxConn.efficientWeight--
+			} else {
+				maxConn.efficientWeight++
+			}
+		},
+	}, nil
+}
+
+type conn struct {
+	balancer.SubConn
+	mutex           sync.Mutex
+	connInfo        base.SubConnInfo
+	name            string
+	weight          uint32
+	efficientWeight uint32
+	currentWeight   uint32
+}

balancer/wrr/weight_picker_test.go

@@ -0,0 +1,166 @@
+package wrr
+
+import (
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/resolver"
+	"testing"
+)
+
+func TestWeightBalancer_PickerBuilder(t *testing.T) {
+	testCases := []struct {
+		name        string
+		b           *PickerBuilder
+		buildInfoFn func() base.PickerBuildInfo
+		wantWeights map[string]uint32
+	}{
+		{
+			name: "default",
+			buildInfoFn: func() base.PickerBuildInfo {
+				readySCs := make(map[balancer.SubConn]base.SubConnInfo, 3)
+				k1 := &mockSubConn{id: 1}
+				k2 := &mockSubConn{id: 2}
+				k3 := &mockSubConn{id: 3}
+				readySCs[k1] = base.SubConnInfo{
+					Address: resolver.Address{
+						Addr:     "weight-1",
+						Metadata: map[string]any{"weight": uint32(15)},
+					},
+				}
+				readySCs[k2] = base.SubConnInfo{
+					Address: resolver.Address{
+						Addr:     "weight-2",
+						Metadata: map[string]any{"weight": uint32(20)},
+					},
+				}
+				readySCs[k3] = base.SubConnInfo{
+					Address: resolver.Address{
+						Addr:     "weight-3",
+						Metadata: map[string]any{"weight": uint32(25)},
+					},
+				}
+				return base.PickerBuildInfo{
+					ReadySCs: readySCs,
+				}
+			},
+			b:           NewPickerBuilder(),
+			wantWeights: map[string]uint32{"weight-1": 15, "weight-2": 20, "weight-3": 25},
+		},
+		{
+			name: "address attributes",
+			buildInfoFn: func() base.PickerBuildInfo {
+				readySCs := make(map[balancer.SubConn]base.SubConnInfo, 3)
+				k1 := &mockSubConn{id: 1}
+				k2 := &mockSubConn{id: 2}
+				k3 := &mockSubConn{id: 3}
+				readySCs[k1] = base.SubConnInfo{
+					Address: resolver.Address{
+						Addr:       "weight-1",
+						Attributes: attributes.New("weight", uint32(15)),
+					},
+				}
+				readySCs[k2] = base.SubConnInfo{
+					Address: resolver.Address{
+						Addr:       "weight-2",
+						Attributes: attributes.New("weight", uint32(20)),
+					},
+				}
+				readySCs[k3] = base.SubConnInfo{
+					Address: resolver.Address{
+						Addr:       "weight-3",
+						Attributes: attributes.New("weight", uint32(25)),
+					},
+				}
+				return base.PickerBuildInfo{
+					ReadySCs: readySCs,
+				}
+			},
+			b: NewPickerBuilder(WithGetWeightFunc(func(ci base.SubConnInfo) uint32 {
+				weight := ci.Address.Attributes.Value("weight").(uint32)
+				return weight
+			})),
+			wantWeights: map[string]uint32{"weight-1": 15, "weight-2": 20, "weight-3": 25},
+		},
+	}
+	for _, tc := range testCases {
+		tt := tc
+		t.Run(tt.name, func(t *testing.T) {
+			p := tt.b.Build(tt.buildInfoFn()).(*Picker)
+			targetWeights := make(map[string]uint32, len(p.connections))
+			for _, c := range p.connections {
+				targetWeights[c.name] = c.weight
+			}
+			assert.Equal(t, targetWeights, tt.wantWeights)
+		})
+	}
+}
+
+func TestWeightBalancer_Pick(t *testing.T) {
+	b := &Picker{
+		connections: []*conn{
+			{
+				name:            "weight-5",
+				weight:          5,
+				efficientWeight: 5,
+				currentWeight:   5,
+			},
+			{
+				name:            "weight-4",
+				weight:          4,
+				efficientWeight: 4,
+				currentWeight:   4,
+			},
+			{
+				name:            "weight-3",
+				weight:          3,
+				efficientWeight: 3,
+				currentWeight:   3,
+			},
+		},
+	}
+	pickRes, err := b.Pick(balancer.PickInfo{})
+	require.NoError(t, err)
+	assert.Equal(t, "weight-5", pickRes.SubConn.(*conn).name)
+
+	pickRes, err = b.Pick(balancer.PickInfo{})
+	require.NoError(t, err)
+	assert.Equal(t, "weight-4", pickRes.SubConn.(*conn).name)
+
+	pickRes, err = b.Pick(balancer.PickInfo{})
+	require.NoError(t, err)
+	assert.Equal(t, "weight-3", pickRes.SubConn.(*conn).name)
+
+	pickRes, err = b.Pick(balancer.PickInfo{})
+	require.NoError(t, err)
+	assert.Equal(t, "weight-5", pickRes.SubConn.(*conn).name)
+
+	pickRes, err = b.Pick(balancer.PickInfo{})
+	require.NoError(t, err)
+	assert.Equal(t, "weight-4", pickRes.SubConn.(*conn).name)
+
+	pickRes.Done(balancer.DoneInfo{})
+	// 断言这里面 efficient weight 是变化了的
+}
+
+var _ balancer.SubConn = (*mockSubConn)(nil)
+
+type mockSubConn struct{ id int }
+
+func (m *mockSubConn) UpdateAddresses(addresses []resolver.Address) {
+	return
+}
+
+func (m *mockSubConn) Connect() {
+	return
+}
+
+func (m *mockSubConn) GetOrBuildProducer(builder balancer.ProducerBuilder) (p balancer.Producer, close func()) {
+	return
+}
+
+func (m *mockSubConn) Shutdown() {
+	return
+}

go.mod

@@ -1,3 +1,16 @@
 module github.com/ecodeclub/grpcx
 
 go 1.21
+
+require (
+	github.com/stretchr/testify v1.9.0
+	google.golang.org/grpc v1.67.0
+)
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	golang.org/x/sys v0.24.0 // indirect
+	google.golang.org/protobuf v1.34.2 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

go.sum

@@ -0,0 +1,24 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+golang.org/x/net v0.28.0 h1:a9JDOJc5GMUJ0+UDqmLT86WiEy7iWyIhz8gz8E4e5hE=
+golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg=
+golang.org/x/sys v0.24.0 h1:Twjiwq9dn6R1fQcyiK+wQyHWfaz/BJB+YIpzU/Cv3Xg=
+golang.org/x/sys v0.24.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/text v0.17.0 h1:XtiM5bkSOt+ewxlOE/aE/AKEHibwj/6gvWMl9Rsh0Qc=
+golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 h1:e7S5W7MGGLaSu8j3YjdezkZ+m1/Nm0uRVRMEMGk26Xs=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142/go.mod h1:UqMtugtsSgubUsoxbuAoiCXvqvErP7Gf0so0mK9tHxU=
+google.golang.org/grpc v1.67.0 h1:IdH9y6PF5MPSdAntIcpjQ+tXO41pcQsfZV2RxtQgVcw=
+google.golang.org/grpc v1.67.0/go.mod h1:1gLDyUQU7CTLJI90u3nXZ9ekeghjeM7pTDZlqFNg2AA=
+google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
+google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=