基于golang实现常用的负载均衡算法

原文链接：Golang实现四种负载均衡算法

接口

package algo

type LoadBalance interface {
   Add(...string) error
   Get(string) (string, error)
}

random

package algo

import (
   "errors"
   "math/rand"
)

type RandomBalance struct {
   curIdx int
   rss []string
}

func (r *RandomBalance) Add(params ...string) error {
   if len(params) == 0 {
      return errors.New("params len at least 1")
   }
   addr := params[0]
   r.rss = append(r.rss, addr)

   return nil
}

func (r *RandomBalance) Next() string {
   if len(r.rss) == 0 {
      return ""
   }

   r.curIdx = rand.Intn(len(r.rss))
   return r.rss[r.curIdx]
}

func (r *RandomBalance) Get(string2 string) (string, error) {
   return r.Next(), nil
}

round-robin

package algo

import "errors"

type RoundRobinBalance struct {
   curIdx int
   rss []string
}

func (r *RoundRobinBalance) Add(params ...string) error {
   if len(params) == 0 {
      return errors.New("params len at least 1")
   }

   addr := params[0]
   r.rss = append(r.rss, addr)
   return nil
}

func (r *RoundRobinBalance) Get(string2 string) (string, error) {
   return r.Next(), nil
}

func (r *RoundRobinBalance) Next() string {
   if len(r.rss) == 0 {
      return ""
   }

   length := len(r.rss)
   // over size, reset to zero index
   if r.curIdx >= length {
      r.curIdx = 0
   }

   curAddr := r.rss[r.curIdx]
   r.curIdx = (r.curIdx + 1) % length
   return curAddr
}

weight-round-robin

package algo

import (
   "errors"
   "strconv"
)

type WeightRoundRobinBalance struct {
   curIdx int
   rss    []*WeightNode
   rsw    []int
}

type WeightNode struct {
   addr            string
   Weight          int // 初始化时堆节点的约定权重
   curWeight       int // 节点临时权重，每轮都会变化
   effectiveWeight int // 有效权重, 默认与weight相同 , totalWeight = sum(effectiveWeight)  //出现故障就-1
}

//1, currentWeight = currentWeight + effectiveWeight
//2, 选中最大的currentWeight节点为选中节点
//3, currentWeight = currentWeight - totalWeight

func (r *WeightRoundRobinBalance) Add(params ...string) error {
   if len(params) != 2 {
      return errors.New("params len must be 2.")
   }

   parseInt, err := strconv.ParseInt(params[1], 10, 64)
   if err != nil {
      return err
   }

   node := &WeightNode{
      addr: params[0],
      Weight: int(parseInt),
   }

   node.effectiveWeight = node.Weight
   r.rss = append(r.rss, node)
   return nil
}

func (r *WeightRoundRobinBalance) Next() string {
   var choosen *WeightNode
   total := 0
   for i := 0; i < len(r.rss); i++ {
      w := r.rss[i]
      // calculate weight sum
      total += w.effectiveWeight
      // modify cur node tmp weight
      w.curWeight += w.effectiveWeight
      // valid weight == weight default, if communication error, -1, else +1, until reseat to weight
      if w.effectiveWeight < w.Weight {
         w.effectiveWeight++
      }
      // choose the largest tmp weight
      if choosen == nil || w.curWeight > choosen.curWeight {
         choosen = w
      }
   }

   if choosen == nil {
      return ""
   }

   // modify tmp weight to tmp - total
   choosen.curWeight -= total
   return choosen.addr
}

func (r *WeightRoundRobinBalance) Get(string2 string) (string, error) {
   return r.Next(), nil
}

consistent-hash

package algo

import (
   "errors"
   "hash/crc32"
   "sort"
   "strconv"
   "sync"
)

// 1. 单调性（唯一） 2. 平衡性 (数据 目标元素均衡) 3. 分散性(散列)
type Hash func(data []byte) uint32

type Uint32Slice []uint32

func (s Uint32Slice) Len() int {
   return len(s)
}

func (s Uint32Slice) Less(i, j int) bool {
   return s[i] < s[j]
}

func (s Uint32Slice) Swap(i, j int) {
   s[i], s[j] = s[j], s[i]
}

type ConsistentHashBalance struct {
   mux      sync.RWMutex // protect hashmap
   hash     Hash
   replicas int // copy vars
   keys     Uint32Slice // sorted node hash slice
   hashMap  map[uint32]string // node hash key , key-hash, value-node's key
}

func NewConsistentHashBalance(replicas int, fn Hash) *ConsistentHashBalance {
   m := &ConsistentHashBalance{
      replicas: replicas,
      hash: fn,
      hashMap: make(map[uint32]string),
   }

   if m.hash == nil {
      m.hash = crc32.ChecksumIEEE
   }

   return m
}

func (c *ConsistentHashBalance) IsEmpty() bool {
   return len(c.keys) == 0
}

func (c *ConsistentHashBalance) Add(params ...string) error {
   if len(params) == 0 {
      return errors.New("params len at least 1")
   }

   addr := params[0]
   c.mux.Lock()
   defer c.mux.Unlock()

   // combine with replicas to calculate virtual node's hash value, save to m.keys，and save to m.hashmap
   for i := 0; i < c.replicas; i++ {
      hash := c.hash([]byte(strconv.Itoa(i) + addr))
      c.keys = append(c.keys, hash)
      c.hashMap[hash] = addr
   }

   // sort, and binary search
   sort.Sort(c.keys)
   return nil
}

// Get 方法根据给定的对象获取最靠近它的那个节点
func (c *ConsistentHashBalance) Get(key string) (string, error) {
   if c.IsEmpty() {
      return "", errors.New("node is empty")
   }
   hash := c.hash([]byte(key))

   // 通过二分查找获取最优节点，第一个"服务器hash"值大于"数据hash"值的就是最优"服务器节点"
   idx := sort.Search(len(c.keys), func(i int) bool { return c.keys[i] >= hash })

   // 如果查找结果 大于 服务器节点哈希数组的最大索引，表示此时该对象哈希值位于最后一个节点之后，那么放入第一个节点中
   if idx == len(c.keys) {
      idx = 0
   }
   c.mux.RLock()
   defer c.mux.RUnlock()
   return c.hashMap[c.keys[idx]], nil
}

main

package algo

type LbType int

const (
   LbRandom LbType = iota
   LbRoundRobin
   LbWeightRoundRobin
   LbConsistentHash
)

func LoadBalanceFactory(lbType LbType) LoadBalance {
   switch lbType {
   case LbRandom:
      return &RandomBalance{}
   case LbConsistentHash:
      return NewConsistentHashBalance(10, nil)
   case LbRoundRobin:
      return &RoundRobinBalance{}
   case LbWeightRoundRobin:
      return &WeightRoundRobinBalance{}
   default:
      return &RandomBalance{}
   }
}