sonet/pkg/grpc/balancer/consistent_hash.go

package balancer

import (
	"encoding/json"
	"errors"
	"fmt"
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/base"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/resolver"
	"sonet/pkg/utils/logger"
	"strconv"
)

const ConsistentHash = "consistent_hash_x"

var ConsistentHashKey = "consistent-hash"

func InitConsistentHashBuilder() {
	balancer.Register(newConsistentHashBuilder())
}

// newConsistentHashBuilder creates a new ConsistentHash balancer builder.
func newConsistentHashBuilder() balancer.Builder {
	return base.NewBalancerBuilder(
		ConsistentHash,
		&consistentHashPickerBuilder{},
		base.Config{HealthCheck: true},
	)
}

type consistentHashPickerBuilder struct{}

func (b *consistentHashPickerBuilder) Build(buildInfo base.PickerBuildInfo) balancer.Picker {
	// grpclog.Infof("consistentHashPicker: newPicker called with buildInfo: %v", buildInfo)
	if len(buildInfo.ReadySCs) == 0 {
		return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
	}

	picker := &consistentHashPicker{
		subConns: make(map[string]balancer.SubConn),
		hash:     NewKetama(DefaultReplicas, nil),
	}

	for sc, conInfo := range buildInfo.ReadySCs {
		weight := GetWeight(conInfo.Address)
		for i := 0; i < weight; i++ {
			node := wrapAddr(conInfo.Address.Addr, i)
			picker.hash.Add(node)
			picker.subConns[node] = sc
		}
	}
	return picker
}

type consistentHashPicker struct {
	subConns map[string]balancer.SubConn
	hash     *Ketama
}

func (p *consistentHashPicker) Pick(info balancer.PickInfo) (ret balancer.PickResult, err error) {
	key, ok := info.Ctx.Value(ConsistentHashKey).(string)
	if !ok || key == "" {
		//key = strconv.Itoa(rand.Intn(65536))
		//grpclog.Warning("empty consistent hash key")
		panic(errors.New("empty consistent hash key"))
	}
	targetAddr, ok := p.hash.Get(key)
	if ok {
		ret.SubConn = p.subConns[targetAddr]
	}
	return
}

func wrapAddr(addr string, idx int) string {
	return fmt.Sprintf("%s-%d", addr, idx)
}

func GetWeight(addr resolver.Address) (weight int) {
	weight = DefaultWeight
	var val any = nil

	// from metadata...
	if marshal, ok := addr.Metadata.(string); ok {
		m := make(map[string]any)
		err := json.Unmarshal([]byte(marshal), &m)
		if err != nil {
			logger.Error("unmarshal metadata error: ", err)
		} else {
			val = m[WeightKey]
		}
	}

	// from attributes...
	if addr.Attributes != nil {
		val = addr.Attributes.Value(WeightKey)
	}

	if val == nil {
		return
	}

	switch val.(type) {
	case int:
		weight = val.(int)
	case string:
		w, err := strconv.Atoi(val.(string))
		if err != nil {
			grpclog.Errorf("instance weight format error: %v\n", val)
			return
		}
		weight = w
		//default:
		//	grpclog.Errorf("instance weight value type not string: %v\n", val)
	}
	return
}