Kitex 服务发现与负载均衡

// pkg/discovery/discovery.go
// Resolver resolves the target endpoint into a list of Instance.
type Resolver interface {
  // Target should return a description for the given target that is suitable for being a key for cache.
  Target(ctx context.Context, target rpcinfo.EndpointInfo) (description string)

  // Resolve returns a list of instances for the given description of a target.
  Resolve(ctx context.Context, desc string) (Result, error)

  // Diff computes the difference between two results.
  // When `next` is cacheable, the Change should be cacheable, too. And the `Result` field's CacheKey in
  // the return value should be set with the given cacheKey.
  Diff(cacheKey string, prev, next Result) (Change, bool)

  // Name returns the name of the resolver.
  Name() string
}

type Result struct {
  Cacheable bool
  CacheKey  string
  Instances []Instance
}

type Instance interface {
  Address() net.Addr
  Weight() int
  Tag(key string) (value string, exist bool)
}

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

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

func resolverBalancerMW(ctx context.Context, req, res interface{}) error {
  // 1. 获取目标地址
  destService := getDestService(ctx)
  // 2. 通过 resolver 拿到地址列表
  discoveryResult, _ := resolver.Resolve(ctx, destService)
  // 3. get picker
  picker := loadbalancer.GetPicker(discoveryResult)
  // 4. pick
  instance := picker.Next(ctx, req)
  // ... use instance handle rpc
}

type Balancer struct {
  b            *BalancerFactory
  target       string
  res          atomic.Value // 将最近一次 resolve 成功的 result 缓存
  expire       int32        // 单例缓存 balancer, 一定时间不使用会在 BalancerFactory cache 中删除当前实例
  sharedTicker *sharedTicker // 定时调用 resolve 的共享 ticker
}

// 此时的 GetPicker 的参数 discovery.Result 是直接拿的缓存结果
func (bl *Balancer) GetPicker() loadbalance.Picker {
  atomic.StoreInt32(&bl.expire, 0)
  res := bl.res.Load().(discovery.Result)
  return bl.b.balancer.GetPicker(res)
}

// GetPicker implements the Loadbalancer interface.
func (wb *weightedBalancer) GetPicker(e discovery.Result) Picker {
  var w *weightInfo
  // 如果支持缓存则尝试进行缓存
  if e.Cacheable {
    // 从缓存中尝试取值
    wi, ok := wb.cachedWeightInfo.Load(e.CacheKey)
    if !ok {
      // 没有值则初始化, 用 singleflight 防止并发重复计算
      wi, _, _ = wb.sfg.Do(e.CacheKey, func() (interface{}, error) {
        return wb.calcWeightInfo(e), nil
      })
      // 存入缓存
      wb.cachedWeightInfo.Store(e.CacheKey, wi)
    }
    w = wi.(*weightInfo)
  } else {
    w = wb.calcWeightInfo(e)
  }

  if w.weightSum == 0 {
    return new(DummyPicker)
  }

  // 优化操作, 如果所有 instance 权重一样, 直接退化到随机
  // 不进行权重计算
  if w.balance {
    picker := randomPickerPool.Get().(*randomPicker)
    picker.immutableInstances = w.instances
    picker.firstIndex = -1
    return picker
  }
  picker := weightedPickerPool.Get().(*weightedPicker)
  picker.immutableEntries = w.entries
  picker.weightSum = w.weightSum
  picker.immutableInstances = w.instances
  picker.firstIndex = -1
  return picker
}

type Rebalancer interface {
  // 地址列表发生变化时, 会调用此函数
  Rebalance(discovery.Change)
  // close 时通知
  Delete(discovery.Change)
}

// Rebalance implements the Rebalancer interface.
func (wb *weightedBalancer) Rebalance(change discovery.Change) {
  if !change.Result.Cacheable {
    return
  }
  // 感知到地址列表变化时, 重新计算并更新缓存
  wb.cachedWeightInfo.Store(change.Result.CacheKey, wb.calcWeightInfo(change.Result))
}

// Delete implements the Rebalancer interface.
func (wb *weightedBalancer) Delete(change discovery.Change) {
  if !change.Result.Cacheable {
    return
  }
  wb.cachedWeightInfo.Delete(change.Result.CacheKey)
}

func (bl *Balancer) refresh() {
  res, err := bl.b.resolver.Resolve(context.Background(), bl.target)
  if err != nil {
    klog.Warnf("KITEX: resolver refresh failed, key=%s error=%s", bl.target, err.Error())
    return
  }
  renameResultCacheKey(&res, bl.b.resolver.Name())
  prev := bl.res.Load().(discovery.Result)
  // 初始化时会判断 balancer 有没有实现 Rebalancer 接口
  // 有才会给 rebalancer 赋值
  if bl.b.rebalancer != nil {
    // 调用 resolver.Diff 判断是否发生变化
    if ch, ok := bl.b.resolver.Diff(res.CacheKey, prev, res); ok {
      // 变化时回调 Rebalance 通知
      bl.b.rebalancer.Rebalance(ch)
    }
  }
  // replace previous result
  bl.res.Store(res)
}

func resolverBalancerMW(ctx context.Context, req, res interface{}) error {
  // 1. 获取目标地址
  dest := getDestService(ctx)
  // 2. 从复杂均衡工厂方法里面根据 dest 信息拿到 lb 实例
  // 缓存单例优化性能
  lb, _ := lbf.Get(ctx, dest)
  // 3. get picker
  // resolver 通过轮询调用, 封装在了 GetPicker 方法里
  picker := lb.GetPicker()
  // 4. pick 连接并发起 rpc 调用
  for {
    select {
    case <-ctx.Done():
      return kerrors.ErrRPCTimeout
    default:
    }

    // 从 lb 拿到连接
    ins := picker.Next(ctx, request)
    if ins == nil {
      return kerrors.ErrNoMoreInstance.WithCause(fmt.Errorf("last error: %w", lastErr))
    }
    remote.SetInstance(ins)

    retryable := func(err error) bool {
      return errors.Is(err, kerrors.ErrGetConnection) || errors.Is(err, kerrors.ErrCircuitBreak)
    }

    // 执行中间件链式调用, 在 连接错误 或 熔断错误 发生时重试
    if err = next(ctx, request, response); err != nil && retryable(err) {
      lastErr = err
      continue
    }
    return err
  }
}