feat: consistent hash balancer support (#5246)

Co-authored-by: 周曙光 <[email protected]>

Author: zhoushuguang

zrpc/client.go

@@ -1,12 +1,16 @@
 package zrpc
 
 import (
+	"context"
+	"fmt"
 	"time"
 
 	"github.com/zeromicro/go-zero/core/conf"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/zrpc/internal"
 	"github.com/zeromicro/go-zero/zrpc/internal/auth"
+	"github.com/zeromicro/go-zero/zrpc/internal/balancer/consistenthash"
+	"github.com/zeromicro/go-zero/zrpc/internal/balancer/p2c"
 	"github.com/zeromicro/go-zero/zrpc/internal/clientinterceptors"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/keepalive"
@@ -67,6 +71,9 @@ func NewClient(c RpcClientConf, options ...ClientOption) (Client, error) {
 		})))
 	}
 
+	svcCfg := makeLBServiceConfig(c.BalancerName)
+	opts = append(opts, WithDialOption(grpc.WithDefaultServiceConfig(svcCfg)))
+
 	opts = append(opts, options...)
 
 	target, err := c.BuildTarget()
@@ -111,7 +118,20 @@ func SetClientSlowThreshold(threshold time.Duration) {
 	clientinterceptors.SetSlowThreshold(threshold)
 }
 
+// SetHashKey sets the hash key into context.
+func SetHashKey(ctx context.Context, key string) context.Context {
+	return consistenthash.SetHashKey(ctx, key)
+}
+
 // WithCallTimeout return a call option with given timeout to make a method call.
 func WithCallTimeout(timeout time.Duration) grpc.CallOption {
 	return clientinterceptors.WithCallTimeout(timeout)
 }
+
+func makeLBServiceConfig(balancerName string) string {
+	if len(balancerName) == 0 {
+		balancerName = p2c.Name
+	}
+
+	return fmt.Sprintf(`{"loadBalancingPolicy":"%s"}`, balancerName)
+}

zrpc/client_test.go

@@ -12,6 +12,8 @@ import (
 	"github.com/zeromicro/go-zero/core/discov"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/internal/mock"
+	"github.com/zeromicro/go-zero/zrpc/internal/balancer/consistenthash"
+	"github.com/zeromicro/go-zero/zrpc/internal/balancer/p2c"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
@@ -245,3 +247,42 @@ func TestNewClientWithTarget(t *testing.T) {
 
 	assert.NotNil(t, err)
 }
+
+func TestMakeLBServiceConfig(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected string
+	}{
+		{
+			name:     "empty name uses default p2c",
+			input:    "",
+			expected: fmt.Sprintf(`{"loadBalancingPolicy":"%s"}`, p2c.Name),
+		},
+		{
+			name:     "custom balancer name",
+			input:    "consistent_hash",
+			expected: `{"loadBalancingPolicy":"consistent_hash"}`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := makeLBServiceConfig(tt.input)
+			if got != tt.expected {
+				t.Errorf("expected %q, got %q", tt.expected, got)
+			}
+		})
+	}
+}
+
+func TestSetHashKey(t *testing.T) {
+	ctx := context.Background()
+	key := "abc123"
+
+	ctx = SetHashKey(ctx, key)
+	got := consistenthash.GetHashKey(ctx)
+	assert.Equal(t, key, got)
+
+	assert.Empty(t, consistenthash.GetHashKey(context.Background()))
+}

zrpc/config.go

@@ -31,6 +31,7 @@ type (
 		Timeout       int64           `json:",default=2000"`
 		KeepaliveTime time.Duration   `json:",optional"`
 		Middlewares   ClientMiddlewaresConf
+		BalancerName  string `json:",default=p2c_ewma"`
 	}
 
 	// A RpcServerConf is a rpc server config.

zrpc/config_test.go

@@ -4,9 +4,11 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	zconf "github.com/zeromicro/go-zero/core/conf"
 	"github.com/zeromicro/go-zero/core/discov"
 	"github.com/zeromicro/go-zero/core/service"
 	"github.com/zeromicro/go-zero/core/stores/redis"
+	"github.com/zeromicro/go-zero/zrpc/internal/balancer/p2c"
 )
 
 func TestRpcClientConf(t *testing.T) {
@@ -39,6 +41,13 @@ func TestRpcClientConf(t *testing.T) {
 		_, err := conf.BuildTarget()
 		assert.Error(t, err)
 	})
+
+	t.Run("default balancer name", func(t *testing.T) {
+		var conf RpcClientConf
+		err := zconf.FillDefault(&conf)
+		assert.NoError(t, err)
+		assert.Equal(t, p2c.Name, conf.BalancerName)
+	})
 }
 
 func TestRpcServerConf(t *testing.T) {

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