feat: consistent hash balancer support (#5246)

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

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