gRPC Go 服务发现与负载均衡(更新版)

// m is a map from scheme to resolver builder.
var  m = make(map[string]Builder)

type Target struct {
  Scheme    string
  Authority string
  Endpoint  string
}

// Builder creates a resolver that will be used to watch name resolution updates.
type Builder interface {
  // Build creates a new resolver for the given target.
  //
  // gRPC dial calls Build synchronously, and fails if the returned error is
  // not nil.
  Build(target Target, cc ClientConn, opts BuildOptions) (Resolver, error)
  // Scheme returns the scheme supported by this resolver.
  // Scheme is defined at https://github.com/grpc/grpc/blob/master/doc/naming.md.
  Scheme() string
}

// State contains the current Resolver state relevant to the ClientConn.
type State struct {
  // Addresses is the latest set of resolved addresses for the target.
  Addresses []Address

  // ServiceConfig contains the result from parsing the latest service
  // config.  If it is nil, it indicates no service config is present or the
  // resolver does not provide service configs.
  ServiceConfig *serviceconfig.ParseResult

  // Attributes contains arbitrary data about the resolver intended for
  // consumption by the load balancing policy.
  Attributes *attributes.Attributes
}

// Resolver watches for the updates on the specified target.
// Updates include address updates and service config updates.
type Resolver interface {
  // ResolveNow will be called by gRPC to try to resolve the target name
  // again. It's just a hint, resolver can ignore this if it's not necessary.
  //
  // It could be called multiple times concurrently.
  ResolveNow(ResolveNowOptions)
  // Close closes the resolver.
  Close()
}

// 定义 Scheme 名称
const exampleScheme = "example"

type exampleResolverBuilder struct {
  addrsStore map[string][]string
}

func NewExampleResolverBuilder(addrsStore map[string][]string) *exampleResolverBuilder {
  return &exampleResolverBuilder{addrsStore: addrsStore}
}

func (e *exampleResolverBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
  // 初始化 resolver, 将 addrsStore 传递进去
  r := &exampleResolver{
    target:     target,
    cc:         cc,
    addrsStore: e.addrsStore,
  }
  // 调用 start 初始化地址
  r.start()
  return r, nil
}
func (e *exampleResolverBuilder) Scheme() string { return exampleScheme }

type exampleResolver struct {
  target     resolver.Target
  cc         resolver.ClientConn
  addrsStore map[string][]string
}

func (r *exampleResolver) start() {
  // 在静态路由表中查询此 Endpoint 对应 addrs
  addrStrs := r.addrsStore[r.target.Endpoint]
  addrs := make([]resolver.Address, len(addrStrs))
  for i, s := range addrStrs {
    addrs[i] = resolver.Address{Addr: s}
  }
  // addrs 列表转化为 state, 调用 cc.UpdateState 更新地址
  r.cc.UpdateState(resolver.State{Addresses: addrs})
}
func (*exampleResolver) ResolveNow(o resolver.ResolveNowOptions) {}
func (*exampleResolver) Close()                                  {}

// 注册我们的 resolver
resolver.Register(NewExampleResolverBuilder(map[string][]string{
  "test": []string{"localhost:8080", "localhost:8081"},
}))

// 建立对应 scheme 的连接, 并且配置负载均衡
conn, err := grpc.Dial("example:///test", grpc.WithDefaultServiceConfig(`{"loadBalancingPolicy":"round_robin"}`))

func Register(ctx context.Context, client *clientv3.Client, service, self string) error {
  resp, err := client.Grant(ctx, 2)
  if err != nil {
    return errors.Wrap(err, "etcd grant")
  }
  _, err = client.Put(ctx, strings.Join([]string{service, self}, "/"), self, clientv3.WithLease(resp.ID))
  if err != nil {
    return errors.Wrap(err, "etcd put")
  }
  // respCh 需要消耗, 不然会有 warning
  respCh, err := client.KeepAlive(ctx, resp.ID)
  if err != nil {
    return errors.Wrap(err, "etcd keep alive")
  }

  for {
    select {
    case <-ctx.Done():
      return nil
    case <-respCh:

    }
  }
}

const (
  // etcd resolver 负责的 scheme 类型
  Scheme      = "etcd"
  defaultFreq = time.Minute * 30
)

type Builder struct {
  client *clientv3.Client
  // 全局路由表快照, 非必要
  store  map[string]map[string]struct{}
}

func NewBuilder(client *clientv3.Client) *Builder {
  return &Builder{
    client: client,
    store:  make(map[string]map[string]struct{}),
  }
}

func (b *Builder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
  b.store[target.Endpoint] = make(map[string]struct{})

  // 初始化 etcd resolver
  r := &etcdResolver{
    client: b.client,
    target: target,
    cc:     cc,
    store:  b.store[target.Endpoint],
    stopCh: make(chan struct{}, 1),
    rn:     make(chan struct{}, 1),
    t:      time.NewTicker(defaultFreq),
  }

  // 开启后台更新 goroutine
  go r.start(context.Background())
  // 全量更新服务地址
  r.ResolveNow(resolver.ResolveNowOptions{})

  return r, nil
}

func (b *Builder) Scheme() string {
  return Scheme
}

type etcdResolver struct {
  client *clientv3.Client
  target resolver.Target
  cc     resolver.ClientConn
  store  map[string]struct{}
  stopCh chan struct{}
  // rn channel is used by ResolveNow() to force an immediate resolution of the target.
  rn chan struct{}
  t  *time.Ticker
}

func (r *etcdResolver) start(ctx context.Context) {
  target := r.target.Endpoint

  w := clientv3.NewWatcher(r.client)
  rch := w.Watch(ctx, target+"/", clientv3.WithPrefix())
  for {
    select {
    case <-r.rn:
      r.resolveNow()
    case <-r.t.C:
      r.ResolveNow(resolver.ResolveNowOptions{})
    case <-r.stopCh:
      w.Close()
      return
    case wresp := <-rch:
      for _, ev := range wresp.Events {
        switch ev.Type {
        case mvccpb.PUT:
          r.store[string(ev.Kv.Value)] = struct{}{}
        case mvccpb.DELETE:
          delete(r.store, strings.Replace(string(ev.Kv.Key), target+"/", "", 1))
        }
      }
      r.updateTargetState()
    }
  }
}

func (r *etcdResolver) resolveNow() {
  target := r.target.Endpoint
  resp, err := r.client.Get(context.Background(), target+"/", clientv3.WithPrefix())
  if err != nil {
    r.cc.ReportError(errors.Wrap(err, "get init endpoints"))
    return
  }

  for _, kv := range resp.Kvs {
    r.store[string(kv.Value)] = struct{}{}
  }

  r.updateTargetState()
}

func (r *etcdResolver) updateTargetState() {
  addrs := make([]resolver.Address, len(r.store))
  i := 0
  for k := range r.store {
    addrs[i] = resolver.Address{Addr: k}
    i++
  }
  r.cc.UpdateState(resolver.State{Addresses: addrs})
}

// 会并发调用, 所以这里防止同时多次全量刷新
func (r *etcdResolver) ResolveNow(o resolver.ResolveNowOptions) {
  select {
  case r.rn <- struct{}{}:
  default:

  }
}

func (r *etcdResolver) Close() {
  r.t.Stop()
  close(r.stopCh)
}

type SubConn interface {
  UpdateAddresses([]resolver.Address)
  Connect()
}

type Balancer interface {
  // UpdateClientConnState 会在 ClientConnState 改变时被调用
  // 例如: resolver 更新地址列表
  UpdateClientConnState(ClientConnState) error
  ResolverError(error)
  // UpdateSubConnState 会在 SubConn 自身状态改变时被调用
  // 例如: 连接断了
  UpdateSubConnState(SubConn, SubConnState)
  Close()
}

// Picker 会在每次 rpc 调用时被调用, 他需要负责连接选取
type Picker interface {
  Pick(info PickInfo) (PickResult, error)
}

type PickerBuilder interface {
  Build(info PickerBuildInfo) balancer.Picker
}

type PickerBuildInfo struct {
  ReadySCs map[balancer.SubConn]SubConnInfo
}

type rrPickerBuilder struct{}

func (*rrPickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
  if len(info.ReadySCs) == 0 {
    return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
  }
  scs := make([]balancer.SubConn, 0, len(info.ReadySCs))
  for sc := range info.ReadySCs {
    scs = append(scs, sc)
  }
  // 将参数传递进来的最新健康连接保存, 并且生成一个随机数作为初始 pick index
  return &rrPicker{
    subConns: scs,
    next: grpcrand.Intn(len(scs)),
  }
}

type rrPicker struct {
  subConns []balancer.SubConn

  mu   sync.Mutex
  next int
}

// 每次 rpc 调用时, 返回 next 对应的连接
// 并且将 next 计数器加一
func (p *rrPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
  p.mu.Lock()
  sc := p.subConns[p.next]
  p.next = (p.next + 1) % len(p.subConns)
  p.mu.Unlock()
  return balancer.PickResult{SubConn: sc}, nil
}

// 去除掉细节字段
type baseBalancer struct {
  // gRPC 连接, baseBalancer 需要维护连接状态
  // 在连接状态变动时生成新的 Picker,
  // 并调用 cc.UpdateState 通知变动
  cc            balancer.ClientConn
  // pickerBuilder 需要用户提供
  pickerBuilder PickerBuilder
  // 维护 addr 和 SubConn 的映射
  subConns *resolver.AddressMap
  // 维护 SubConn 和状态的映射
  scStates map[balancer.SubConn]connectivity.State
}

// UpdateClientConnState 会在连接基础信息发生改变时被调用
// 例如: resolver 更新地址列表信息
func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
  b.resolverErr = nil

  // 遍历 s.ResolverState.Addresses 地址列表
  // 和现有的 subConns 进行比对
  // 如果发现是新的 addr, 则调用 NewSubConn 创建新连接
  addrsSet := resolver.NewAddressMap()
  for _, a := range s.ResolverState.Addresses {
    addrsSet.Set(a, nil)
    if _, ok := b.subConns.Get(a); !ok {
      // a is a new address (not existing in b.subConns).
      sc, err := b.cc.NewSubConn([]resolver.Address{a}, balancer.NewSubConnOptions{HealthCheckEnabled: b.config.HealthCheck})
      if err != nil {
        logger.Warningf("base.baseBalancer: failed to create new SubConn: %v", err)
        continue
      }
      b.subConns.Set(a, sc)
      b.scStates[sc] = connectivity.Idle
      b.csEvltr.RecordTransition(connectivity.Shutdown, connectivity.Idle)
      sc.Connect()
    }
  }
  // 将不存在于 s.ResolverState.Addresses 地址列表中的连接关闭
  for _, a := range b.subConns.Keys() {
    sci, _ := b.subConns.Get(a)
    sc := sci.(balancer.SubConn)
    // a was removed by resolver.
    if _, ok := addrsSet.Get(a); !ok {
      b.cc.RemoveSubConn(sc)
      b.subConns.Delete(a)
    }
  }

  if len(s.ResolverState.Addresses) == 0 {
    b.ResolverError(errors.New("produced zero addresses"))
    return balancer.ErrBadResolverState
  }
  return nil
}

// 从当前最新的 subConns 中筛选出状态为 Ready 的连接
// 调用用户的 pickerBuilder.Build 方法创建新的 Picker
func (b *baseBalancer) regeneratePicker() {
  if b.state == connectivity.TransientFailure {
    b.picker = NewErrPicker(b.mergeErrors())
    return
  }
  readySCs := make(map[balancer.SubConn]SubConnInfo)

  // Filter out all ready SCs from full subConn map.
  for _, addr := range b.subConns.Keys() {
    sci, _ := b.subConns.Get(addr)
    sc := sci.(balancer.SubConn)
    if st, ok := b.scStates[sc]; ok && st == connectivity.Ready {
      readySCs[sc] = SubConnInfo{Address: addr}
    }
  }
  b.picker = b.pickerBuilder.Build(PickerBuildInfo{ReadySCs: readySCs})
}

// UpdateSubConnState 会在 SubConn 自身连接状态发生改变时被调用
func (b *baseBalancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
  // 改变后的新状态
  s := state.ConnectivityState
  // 老状态
  oldS, ok := b.scStates[sc]
  if !ok {
    return
  }
  if oldS == connectivity.TransientFailure &&
    (s == connectivity.Connecting || s == connectivity.Idle) {
    // 从暂时失败转换为空闲状态时, 调用 Connect 发起连接
    if s == connectivity.Idle {
      sc.Connect()
    }
    // 正在连接状态, 不做处理
    return
  }
  // 更新状态为新状态
  b.scStates[sc] = s
  switch s {
  case connectivity.Idle:
    sc.Connect()
  case connectivity.Shutdown:
    // 连接断开时, 删除他的本地状态映射
    delete(b.scStates, sc)
  case connectivity.TransientFailure:
    // 保存错误
    b.connErr = state.ConnectionError
  }

  // 更新聚合状态
  b.state = b.csEvltr.RecordTransition(oldS, s)

  // 在下面两种情况时重新生成 Picker
  //  - 当前 SubConn 状态由其他变为 Ready 或者由 Ready 变为其他
  //  - 聚合状态转变为临时性错误时
  if (s == connectivity.Ready) != (oldS == connectivity.Ready) ||
    b.state == connectivity.TransientFailure {
    b.regeneratePicker()
  }
  // 调用 cc.UpdateState 更新聚合 state 和最新 Picker
  b.cc.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
}