sonet/pkg/grpc/discovery/resolver.go

package discovery

import (
	"context"
	"sonet/pkg/utils/logger"
	"time"

	clientv3 "go.etcd.io/etcd/client/v3"
	"google.golang.org/grpc/resolver"
)

const (
	schema = "etcd"
)

// Resolver for grpc client
type Resolver struct {
	schema      string
	DialTimeout int

	closeCh      chan struct{}
	watchCh      clientv3.WatchChan
	cli          *clientv3.Client
	keyPrifix    string
	srvAddrsList []resolver.Address

	cc resolver.ClientConn
}

// NewResolver create a new resolver.Builder base on etcd
func NewResolver(client *clientv3.Client) *Resolver {
	return &Resolver{
		cli:         client,
		schema:      schema,
		DialTimeout: 3,
	}
}

// Scheme returns the scheme supported by this resolver.
func (r *Resolver) Scheme() string {
	return r.schema
}

// Build creates a new resolver.Resolver for the given target
func (r *Resolver) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (rr resolver.Resolver, err error) {
	r.cc = cc
	r.keyPrifix = BuildPrefix(Server{Name: target.Endpoint()})
	if err = r.start(); err != nil {
		return nil, err
	}
	return r, nil
}

// ResolveNow resolver.Resolver interface
func (r *Resolver) ResolveNow(o resolver.ResolveNowOptions) {}

// Close resolver.Resolver interface
func (r *Resolver) Close() {
	if r.closeCh == nil {
		return
	}
	r.closeCh <- struct{}{}
	<-r.closeCh
}

// start
func (r *Resolver) start() error {
	var err error

	resolver.Register(r)

	if r.closeCh == nil {
		r.closeCh = make(chan struct{})
	}

	if err = r.sync(); err != nil {
		return err
	}

	go r.watch()

	return nil
}

// watch update events
func (r *Resolver) watch() {
	ticker := time.NewTicker(time.Minute)
	r.watchCh = r.cli.Watch(context.Background(), r.keyPrifix, clientv3.WithPrefix())

	for {
		select {
		case <-r.closeCh:
			r.closeCh <- struct{}{}
			return
		case res, ok := <-r.watchCh:
			if ok {
				r.update(res.Events)
			}
		case <-ticker.C:
			if err := r.sync(); err != nil {
				logger.Error("resolver sync failed: ", err)
			}
		}
	}
}

// update
func (r *Resolver) update(events []*clientv3.Event) {
	for _, ev := range events {
		var info Server
		var err error

		switch ev.Type {
		case clientv3.EventTypePut:
			info, err = ParseValue(ev.Kv.Value)
			if err != nil {
				continue
			}
			addr := resolver.Address{Addr: info.Addr}
			for k, v := range info.Attrs {
				addr.Attributes = addr.Attributes.WithValue(k, v)
			}
			if !Exist(r.srvAddrsList, addr) {
				r.srvAddrsList = append(r.srvAddrsList, addr)
				err = r.cc.UpdateState(resolver.State{Addresses: r.srvAddrsList})
				if err != nil {
					logger.Error("resolver conn update put state err: ", err)
				}
			}
		case clientv3.EventTypeDelete:
			info, err = SplitPath(string(ev.Kv.Key))
			if err != nil {
				continue
			}
			addr := resolver.Address{Addr: info.Addr}
			if s, ok := Remove(r.srvAddrsList, addr); ok {
				r.srvAddrsList = s
				err = r.cc.UpdateState(resolver.State{Addresses: r.srvAddrsList})
				if err != nil {
					logger.Error("resolver conn update delete state err: ", err)
				}
			}
		}
	}
}

// sync 同步获取所有地址信息
func (r *Resolver) sync() (err error) {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	res, err := r.cli.Get(ctx, r.keyPrifix, clientv3.WithPrefix())
	if err != nil {
		return
	}
	r.srvAddrsList = []resolver.Address{}

	for _, v := range res.Kvs {
		info, err := ParseValue(v.Value)
		if err != nil {
			continue
		}
		addr := resolver.Address{Addr: info.Addr}
		for k, v := range info.Attrs {
			addr.Attributes = addr.Attributes.WithValue(k, v)
		}
		r.srvAddrsList = append(r.srvAddrsList, addr)
	}
	err = r.cc.UpdateState(resolver.State{Addresses: r.srvAddrsList})
	return
}