feat: consistent hash balancer support (#5246)

Co-authored-by: 周曙光 <zsg@zhoushuguangdeMacBook-Pro.local>

zrpc/internal/balancer/consistenthash/consistenthash.go

@@ -0,0 +1,97 @@
+package consistenthash
+
+import (
+	"context"
+
+	"github.com/zeromicro/go-zero/core/hash"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+const (
+	Name = "consistent_hash"
+
+	defaultReplicaCount = 100
+)
+
+var emptyPickResult balancer.PickResult
+
+func init() {
+	balancer.Register(newBuilder())
+}
+
+type (
+	// hashKey is the key type for consistent hash in context.
+	hashKey struct{}
+	// pickerBuilder is a builder for picker.
+	pickerBuilder struct{}
+	// picker is a picker that uses consistent hash to pick a sub connection.
+	picker struct {
+		hashRing *hash.ConsistentHash
+		conns    map[string]balancer.SubConn
+	}
+)
+
+func (b *pickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
+	readySCs := info.ReadySCs
+	if len(readySCs) == 0 {
+		return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
+	}
+
+	conns := make(map[string]balancer.SubConn, len(readySCs))
+	hashRing := hash.NewCustomConsistentHash(defaultReplicaCount, hash.Hash)
+	for conn, connInfo := range readySCs {
+		addr := connInfo.Address.Addr
+		conns[addr] = conn
+		hashRing.Add(addr)
+	}
+
+	return &picker{
+		hashRing: hashRing,
+		conns:    conns,
+	}
+}
+
+func newBuilder() balancer.Builder {
+	return base.NewBalancerBuilder(Name, &pickerBuilder{}, base.Config{HealthCheck: true})
+}
+
+func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	hashKey := GetHashKey(info.Ctx)
+	if len(hashKey) == 0 {
+		return emptyPickResult, status.Error(codes.InvalidArgument,
+			"[consistent_hash] missing hash key in context")
+	}
+
+	if addrAny, ok := p.hashRing.Get(hashKey); ok {
+		addr, ok := addrAny.(string)
+		if !ok {
+			return emptyPickResult, status.Error(codes.Internal,
+				"[consistent_hash] invalid addr type in consistent hash")
+		}
+
+		subConn, ok := p.conns[addr]
+		if !ok {
+			return emptyPickResult, status.Errorf(codes.Internal,
+				"[consistent_hash] no subConn for addr: %s", addr)
+		}
+
+		return balancer.PickResult{SubConn: subConn}, nil
+	}
+
+	return emptyPickResult, status.Errorf(codes.Unavailable,
+		"[consistent_hash] no matching conn for hashKey: %s", hashKey)
+}
+
+// SetHashKey sets the hash key into context.
+func SetHashKey(ctx context.Context, key string) context.Context {
+	return context.WithValue(ctx, hashKey{}, key)
+}
+
+// GetHashKey gets the hash key from context.
+func GetHashKey(ctx context.Context) string {
+	v, _ := ctx.Value(hashKey{}).(string)
+	return v
+}

zrpc/internal/balancer/consistenthash/consistenthash_test.go

@@ -0,0 +1,175 @@
+package consistenthash
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/hash"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/resolver"
+)
+
+type fakeSubConn struct{ id int }
+
+func (f *fakeSubConn) Connect()                             {}
+func (f *fakeSubConn) UpdateAddresses(_ []resolver.Address) {}
+func (f *fakeSubConn) Shutdown()                            {}
+func (f *fakeSubConn) GetOrBuildProducer(b balancer.ProducerBuilder) (balancer.Producer, func()) {
+	return nil, func() {}
+}
+
+func TestPickerBuilder_EmptyReadySCs(t *testing.T) {
+	b := &pickerBuilder{}
+	p := b.Build(base.PickerBuildInfo{ReadySCs: map[balancer.SubConn]base.SubConnInfo{}})
+
+	_, err := p.Pick(balancer.PickInfo{})
+	assert.Equal(t, balancer.ErrNoSubConnAvailable, err)
+}
+
+func TestPickerBuilder_BuildAndRing(t *testing.T) {
+	subConn1 := &fakeSubConn{id: 1}
+	subConn2 := &fakeSubConn{id: 2}
+	addr1 := "127.0.0.1:8080"
+	addr2 := "127.0.0.1:8081"
+
+	b := &pickerBuilder{}
+	info := base.PickerBuildInfo{
+		ReadySCs: map[balancer.SubConn]base.SubConnInfo{
+			subConn1: {Address: resolver.Address{Addr: addr1}},
+			subConn2: {Address: resolver.Address{Addr: addr2}},
+		},
+	}
+
+	p := b.Build(info).(*picker)
+	assert.NotNil(t, p.hashRing)
+	assert.Len(t, p.conns, 2)
+}
+
+func TestPicker_HashConsistency(t *testing.T) {
+	subConn1 := &fakeSubConn{id: 1}
+	subConn2 := &fakeSubConn{id: 2}
+
+	pb := &pickerBuilder{}
+	info := base.PickerBuildInfo{
+		ReadySCs: map[balancer.SubConn]base.SubConnInfo{
+			subConn1: {Address: resolver.Address{Addr: "127.0.0.1:8080"}},
+			subConn2: {Address: resolver.Address{Addr: "127.0.0.1:8081"}},
+		},
+	}
+	p := pb.Build(info).(*picker)
+	ctx := SetHashKey(context.Background(), "user_123")
+	res1, err := p.Pick(balancer.PickInfo{Ctx: ctx})
+	assert.NoError(t, err)
+	assert.NotNil(t, res1.SubConn)
+
+	// Multiple requests with the same key remain consistent
+	for i := 0; i < 5; i++ {
+		resN, err := p.Pick(balancer.PickInfo{Ctx: ctx})
+		assert.NoError(t, err)
+		assert.Equal(t, res1.SubConn, resN.SubConn)
+	}
+}
+
+func TestPicker_MissingKey(t *testing.T) {
+	subConn := &fakeSubConn{id: 1}
+
+	pb := &pickerBuilder{}
+	info := base.PickerBuildInfo{
+		ReadySCs: map[balancer.SubConn]base.SubConnInfo{
+			subConn: {Address: resolver.Address{Addr: "127.0.0.1:8080"}},
+		},
+	}
+	p := pb.Build(info).(*picker)
+
+	// No hash key in context
+	_, err := p.Pick(balancer.PickInfo{Ctx: context.Background()})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "[consistent_hash] missing hash key in context")
+}
+
+func TestPicker_NoMatchingConn(t *testing.T) {
+	emptyRing := newCustomRingForTest()
+	p := &picker{
+		hashRing: emptyRing,
+		conns:    map[string]balancer.SubConn{},
+	}
+
+	_, err := p.Pick(balancer.PickInfo{Ctx: SetHashKey(context.Background(), "someone")})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "[consistent_hash] no matching conn for hashKey: someone")
+}
+
+func TestPicker_InvalidAddrType(t *testing.T) {
+	ring := newCustomRingForTest()
+	ring.Add(12345)
+
+	subConn := &fakeSubConn{id: 1}
+	p := &picker{
+		hashRing: ring,
+		conns: map[string]balancer.SubConn{
+			"12345": subConn,
+		},
+	}
+
+	_, err := p.Pick(balancer.PickInfo{Ctx: SetHashKey(context.Background(), "anykey")})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "[consistent_hash] invalid addr type in consistent hash")
+}
+
+func TestPicker_NoSubConnForAddr(t *testing.T) {
+	ring := newCustomRingForTest()
+	ring.Add("ghost:9999")
+
+	exist := &fakeSubConn{id: 1}
+	p := &picker{
+		hashRing: ring,
+		conns: map[string]balancer.SubConn{
+			"real:8080": exist,
+		},
+	}
+
+	_, err := p.Pick(balancer.PickInfo{Ctx: SetHashKey(context.Background(), "anykey")})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "[consistent_hash] no subConn for addr: ghost:9999")
+}
+
+func TestSetAndGetHashKey(t *testing.T) {
+	ctx := context.Background()
+	key := "abc123"
+
+	ctx = SetHashKey(ctx, key)
+	got := GetHashKey(ctx)
+	assert.Equal(t, key, got)
+
+	assert.Empty(t, GetHashKey(context.Background()))
+}
+
+func BenchmarkPicker_HashConsistency(b *testing.B) {
+	subConn1 := &fakeSubConn{id: 1}
+	subConn2 := &fakeSubConn{id: 2}
+
+	pb := &pickerBuilder{}
+	info := base.PickerBuildInfo{
+		ReadySCs: map[balancer.SubConn]base.SubConnInfo{
+			subConn1: {Address: resolver.Address{Addr: "127.0.0.1:8080"}},
+			subConn2: {Address: resolver.Address{Addr: "127.0.0.1:8081"}},
+		},
+	}
+	p := pb.Build(info).(*picker)
+
+	ctx := SetHashKey(context.Background(), "hot_user_123")
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		res, err := p.Pick(balancer.PickInfo{Ctx: ctx})
+		if err != nil || res.SubConn == nil {
+			b.Fatalf("unexpected result: res=%v err=%v", res.SubConn, err)
+		}
+	}
+}
+
+func newCustomRingForTest() *hash.ConsistentHash {
+	return hash.NewCustomConsistentHash(defaultReplicaCount, hash.Hash)
+}