Kitex 一致性哈希算法源码解读

290 阅读9分钟

一、背景

Kitex 是字节跳动内部的 Golang 微服务 RPC 框架,具有高性能强可扩展的特点,在字节内部已广泛使用。同样,Kitex 在服务治理方面也十分全面,支持服务的负载均衡。
负载均衡是一种分配网络流量到多个服务器的技术,以确保没有一个服务器过载,从而提高应用程序的性能和可靠性。负载均衡器根据预定的算法决定将每个请求分配到哪个服务器。
目前 Kitex 支持的负载均衡算法有:

  • WeightedRoundRobin
  • InterleavedWeightedRoundRobin(kitex >= v0.7.0)
  • WeightedRandom
  • Alias Method(kitex >= v0.9.0)
  • ConsistentHash
  • Tagging Based

接下来将通过阅读算法的源码来介绍 Kitex 是怎样实现 ConsistentHash 算法

二、一致性哈希算法介绍

介绍

一致性哈希的基本原理

  • 哈希环
    • 一致性哈希将整个哈希空间(通常是一个大整数空间)组织成一个虚拟的环。
  • 节点映射
    • 每个服务器节点通过哈希函数映射到环上的某个点。这个点被称为该节点的“位置”。
    • Kitex 使用 xxhash3 包下的哈希函数将服务节点进行哈希,得到一个整数值,然后将这个值映射到环上。
  • 数据映射
    • 数据项(如缓存中的键)也通过哈希函数映射到环上的某个点。
    • 通过一致性哈希算法,数据项将被分配到顺时针方向遇到的第一个节点上。这保证了每个数据项有唯一的主节点。
  • 虚拟节点
    • 为了解决数据倾斜问题(即某些节点承载的数据量过多),一致性哈希引入了虚拟节点(Virtual Nodes)。
    • 每个物理节点对应多个虚拟节点,并且这些虚拟节点通过不同的哈希值分布在环上,从而实现更均匀的数据分布。

下图为一个带虚拟节点的哈希环,其中 A,B,C为真实节点,而每个真实节点对应有三个虚拟节点,分别为:
真实节点 A 对应的虚拟节点: A-01, A-02, A-03
真实节点 B 对应的虚拟节点: B-01, B-02, B-03
真实节点 C 对应的虚拟节点: C-01, C-02, C-03
一致性哈希步骤:对每个服务节点进行哈希 -> 映射到虚拟节点 -> 根据虚拟节点获取真实节点。
这样可以使得服务节点分布更加均匀。

dbb57b8d6071d011d05eeadd93269e13.webp

主要结构

负载均衡器接口

Picker 接口的作用是选择一个实例。
Loadbalancer 接口的作用是拿到一个 Picker 进行后续实例的选择。
Rebalancer 接口的作用是当服务节点发生变化时,可以将余下服务再均衡。

// Picker picks an instance for next RPC call.
type Picker interface {
	Next(ctx context.Context, request interface{}) discovery.Instance
}

// Loadbalancer generates pickers for the given service discovery result.
type Loadbalancer interface {
	GetPicker(discovery.Result) Picker
	Name() string // unique key
}

// Rebalancer is a kind of Loadbalancer that performs rebalancing when the result of service discovery changes.
type Rebalancer interface {
	Rebalance(discovery.Change)
	Delete(discovery.Change)
}

配置项

type ConsistentHashOption struct {
    
    // 这是一个函数类型,用于从输入数据中提取键值。该键值用于哈希计算。
	GetKey KeyFunc

    // 表示副本的数量。当连接到主节点失败时,可以使用副本进行连接。这会增加内存和CPU开销。
	Replica uint32

    // 每个真实节点对应的虚拟节点数量。虚拟节点数量越多,内存和计算成本越高,但负载分布会更均衡。
	VirtualFactor uint32

    // 是否按照权重进行负载均衡。如果设置为 false,则每个实例生成相同数量的虚拟节点;如果设置为 true,则根据权重生成虚拟节点数量。
	Weighted bool

    // 是否进行过期处理。实现将缓存所有键值。如果不设置过期,可能会导致内存不断增长,最终导致内存不足。设置过期会带来额外的性能开销。当前实现每分钟扫描一次进行删除,并在实例更改时重建一次。
	ExpireDuration time.Duration
}

一致性哈希结构

// 虚拟节点,N 虚拟节点 --> 1 真实节点
type virtualNode struct {
    hash     uint64
    RealNode *realNode
}

// 真实节点
type realNode struct {
    Ins discovery.Instance
}

/**
 * 一致性哈希结果
 *	Primary: 首选节点
 *	Replicas: 副本节点数组
 *	Touch: 记录最近一次的访问时间
 */
type consistResult struct {
    Primary  discovery.Instance
    Replicas []discovery.Instance
    Touch    atomic.Value
}

/**
 * 一致性哈希信息
 *	cachedConsistResult: map[hashKey]*consistResult
 *	sfg: 防止多个请求同时请求同一个数据项的哈希结果时,重复进行哈希计算。
 *	realNodes: 真实节点数组
 *	virtualNodes: 虚拟节点数组
 */
type consistInfo struct {
    cachedConsistResult sync.Map
    sfg                 singleflight.Group // To prevent multiple builds on the first request for the same key

    realNodes    []realNode
    virtualNodes []virtualNode
}

/**
 * 一致性哈希选择器, 实现 Picker 的接口
 *	cb: 一致性哈希均衡器(指针)
 *	info: 一致性哈希信息(指针)
 *	index: 索引
 *	result: 一致性哈希结果(指针)
 */
type consistPicker struct {
	cb     *consistBalancer
	info   *consistInfo
	index  int
	result *consistResult
}

/**
 * 一致性哈希均衡器, 实现 Loadbalancer 的接口
 *	cachedConsistInfo: map[discovery.CacheKey]*consistInfo
 *	info: 一致性哈希信息(指针)
 *	index: 索引
 *	result: 一致性哈希结果(指针)
 */
type consistBalancer struct {
	cachedConsistInfo sync.Map
	// The main purpose of this lock is to improve performance and prevent Change from being performed while expire
	// which may cause Change to do a lot of extra computation and memory allocation
	updateLock sync.Mutex
	opt        ConsistentHashOption
	sfg        singleflight.Group
}

获取主节点的实现细节

实例化 一致性哈希 Balancer

NewConsistBalancer 根据所给配置项实例化一个负载均衡器。
若在配置项中设置了过期时间,则会将该均衡器添加到一个守护协程中。在这个守护协程中,将每两分钟轮询所有的均衡器,若这些均衡器到达超时时间,则将删除缓存中的键值。

func NewConsistBalancer(opt ConsistentHashOption) Loadbalancer {
	if opt.GetKey == nil {
		panic("loadbalancer: new consistBalancer failed, getKey func cannot be nil")
	}
	if opt.VirtualFactor == 0 {
		panic("loadbalancer: new consistBalancer failed, virtual factor must > 0")
	}
	cb := &consistBalancer{
		opt: opt,
	}
	if cb.opt.ExpireDuration > 0 {
		cb.AddToDaemon()
	}
	return cb
}

获取 一致性哈希 Picker

负载均衡 cb 通过 GetPicker 方法法 根据获取一个 Picker。
Picker 可以通过 用户自定义实现的 getKey 函数获取到一个key值,再通过对 key 进行哈希函数映射到虚拟节点上。最后通过虚拟节点获取到真实节点并返回。

// GetPicker implements the Loadbalancer interface.
func (cb *consistBalancer) GetPicker(e discovery.Result) Picker {
	var ci *consistInfo
	if e.Cacheable {
		cii, ok := cb.cachedConsistInfo.Load(e.CacheKey)
		if !ok {
			cii, _, _ = cb.sfg.Do(e.CacheKey, func() (interface{}, error) {
				return cb.newConsistInfo(e), nil
			})
			cb.cachedConsistInfo.Store(e.CacheKey, cii)
		}
		ci = cii.(*consistInfo)
	} else {
		ci = cb.newConsistInfo(e)
	}
	picker := consistPickerPool.Get().(*consistPicker)
	picker.cb = cb
	picker.info = ci
	return picker
}

获取 realNode 的过程

通过 Picker.Next() 方法 获取到 realNode

  • 利用 xxhash3 包下的 hashString 方法对 key 进行哈希函数计算,获取到哈希值。
  • 根据 哈希值 调用 consistPicker 下的 getConsistResult 方法 获取到 realNode。
func (cp *consistPicker) Next(ctx context.Context, request interface{}) discovery.Instance {
    if len(cp.info.realNodes) == 0 {
        return nil
    }
    if cp.result == nil {
        key := cp.cb.opt.GetKey(ctx, request)
        if key == "" {
            return nil
        }
        /**
         * xxhash3.HashString(key) 对key进行哈希函数
         * getConsistResult: 根据 hash 值获取到结果
         */
        cp.result = cp.getConsistResult(xxhash3.HashString(key))
        cp.index = 0
        return cp.result.Primary
    }
    if cp.index < len(cp.result.Replicas) {
        cp.index++
        return cp.result.Replicas[cp.index-1]
    }
    return nil
}

通过 getConsistResult 获取一致性哈希结果

  • 根据 哈希值 从 一致性哈希信息(consistInfo) 缓存中获取到结果。
    • 若缓存中不存在,则调用 buildConsistResult 方法构造结果。
    • 若缓存中存在,则直接返回结果。
  • 并且结合 Touch 和 option 中的过期时间,将缓存中的数据进行更新。
/**
 * 根据 哈希值 从 一致性哈希信息(consistInfo) 缓存中获取到结果。
 * 若缓存中不存在,则调用 buildConsistResult 方法构造结果。
 * 若缓存中存在,则之间返回结果。
 * 并且结合 Touch 和 option 中的过期时间,将缓存中的数据进行更新。
 */
func (cp *consistPicker) getConsistResult(key uint64) *consistResult {
	var cr *consistResult
	cri, ok := cp.info.cachedConsistResult.Load(key)
	if !ok {
		cri, _, _ = cp.info.sfg.Do(strconv.FormatUint(key, 10), func() (interface{}, error) {
			cr := buildConsistResult(cp.cb, cp.info, key)
			if cp.cb.opt.ExpireDuration > 0 {
				cr.Touch.Store(time.Now())
			}
			return cr, nil
		})
		cp.info.cachedConsistResult.Store(key, cri)
	}
	cr = cri.(*consistResult)
	if cp.cb.opt.ExpireDuration > 0 {
		cr.Touch.Store(time.Now())
	}
	return cr
}

通过 buildConsistResult 构建一致性哈希结果

func buildConsistResult(cb *consistBalancer, info *consistInfo, key uint64) *consistResult {
	cr := &consistResult{}
    // 在虚拟节点中查找哈希值大于 key 的第一个节点,如果找不到,则返回数组的长度
	index := sort.Search(len(info.virtualNodes), func(i int) bool {
		return info.virtualNodes[i].hash > key
	})
	// 如果索引等于虚拟节点数组的长度,说明需要回到哈希环的起点。
	if index == len(info.virtualNodes) {
		index = 0
	}
    // 确定主节点
	cr.Primary = info.virtualNodes[index].RealNode.Ins
    // 确定副本数量
	replicas := int(cb.opt.Replica)
	// 根据配置的副本数量和真实节点的数量,调整副本数量,确保不会超过实际的节点数量。
	if len(info.realNodes)-1 < replicas {
		replicas = len(info.realNodes) - 1
	}
    /** 
     * 处理主节点(真实节点)和副本节点(虚拟节点)之间的映射关系
     * 一个主节点 映射 多个副本节点
     */
	if replicas > 0 {
        // 使用一个 map 记录已使用的节点,确保不会选择重复的节点作为副本。
		used := make(map[discovery.Instance]struct{}, replicas) // should be 1 + replicas - 1
		used[cr.Primary] = struct{}{}
		cr.Replicas = make([]discovery.Instance, replicas)
		for i := 0; i < replicas; i++ {
            // 循环查找下一个未使用的节点,并将其添加到副本列表中。
            // 如果索引到达数组末尾,则回到起点,继续查找。
			for {
				index++
				if index == len(info.virtualNodes) {
					index = 0
				}
				ins := info.virtualNodes[index].RealNode.Ins
				if _, ok := used[ins]; !ok {
					used[ins] = struct{}{}
					cr.Replicas[i] = ins
					break
				}
			}
		}
	}
	return cr
}

通过上述方法链的调用,我们成功通过 一致性哈希 算法获取到了服务节点,后续可以通过访问该服务节点处理相应的请求。

更新节点的实现细节

通过 Balancer.ReBalance() 方法进行更新

只有 一致性哈希结果 允许缓存的情况下可以更新节点。
通过 updateLock 实现了线程安全。在资源调动区执行 updateConsistInfo 方法执行实际的更新逻辑。

func (cb *consistBalancer) Rebalance(change discovery.Change) {
    if !change.Result.Cacheable {
        return
    }
    cb.updateLock.Lock()
    cb.updateConsistInfo(change.Result)
    cb.updateLock.Unlock()
}

通过 updateConsistInfo 方法进行更新

func (cb *consistBalancer) updateConsistInfo(e discovery.Result) {
    // 创建新的 consistInfo 对象
    newInfo := cb.newConsistInfo(e)
    // 使用 LoadOrStore 方法尝试将新的 consistInfo 对象存储到缓存中,如果缓存中已存在对应的 CacheKey,则返回已存在的对象。
    infoI, loaded := cb.cachedConsistInfo.LoadOrStore(e.CacheKey, newInfo)
    // 如果缓存中不存在对应的 CacheKey,则直接返回,不进行后续操作。
    if !loaded {
        return
    }
    // 遍历旧的 cachedConsistResult,为每个条目重新计算哈希结果并存储到新的 cachedConsistResult 中。如果配置了过期时间,还需要更新 Touch 时间。
    info := infoI.(*consistInfo)
    info.cachedConsistResult.Range(func(key, value interface{}) bool {
        cr := buildConsistResult(cb, newInfo, key.(uint64))
        if cb.opt.ExpireDuration > 0 {
            t := value.(*consistResult).Touch.Load().(time.Time)
            if time.Now().After(t.Add(cb.opt.ExpireDuration)) {
                return true
            }
            cr.Touch.Store(t)
        }
        // 新节点 存储结果缓存
        newInfo.cachedConsistResult.Store(key, cr)
        return true
    })
    // 更新缓存
    cb.cachedConsistInfo.Store(e.CacheKey, newInfo)
}

删除节点的实现细节

通过 Balancer.Delete() 方法进行删除

同样,只有当结果允许采用缓存存储时才能进行删除。
通过 updateLock 实现了线程安全。删除逻辑非常简单,在资源调动区执行 将一致性结果缓存中相应 key 值的map进行删除即可。

func (cb *consistBalancer) Delete(change discovery.Change) {
    if (!change.Result.Cacheable) {
        return
    }
    cb.updateLock.Lock()
    cb.cachedConsistInfo.Delete(change.Result.CacheKey)
    cb.updateLock.Unlock()
}

三、使用用例

// 创建 client 的时候初始化一致性哈希策略:
cli, err := echo.NewClient(
    "echo",
    client.WithLoadBalancer(loadbalance.NewConsistBalancer(loadbalance.NewConsistentHashOption(func(ctx context.Context, request interface{}) string {
		// 根据请求上下文信息设置 key
        return "key"
    }))),
)

参考:

  1. www.xiaolincoding.com/os/8_networ…

  2. www.cloudwego.io/zh/docs/kit…