github.com/weedge/lib@v0.0.0-20230424045628-a36dcc1d90e4/client/etcd/discovery/grpc_resolver.go (about)

     1  package discovery
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strings"
     7  	"time"
     8  
     9  	"github.com/weedge/lib/runtimer"
    10  
    11  	"github.com/weedge/lib/log"
    12  
    13  	"github.com/pkg/errors"
    14  	"go.etcd.io/etcd/api/v3/mvccpb"
    15  	clientv3 "go.etcd.io/etcd/client/v3"
    16  	"google.golang.org/grpc/resolver"
    17  )
    18  
    19  type GrpcResolver struct {
    20  	*EtcdClient
    21  	store map[string]map[string]struct{}
    22  }
    23  
    24  func NewGrpcResolver(endpoints []string, dialTimeout time.Duration) (m *GrpcResolver) {
    25  	if len(endpoints) == 0 {
    26  		endpoints = []string{"0.0.0.0:2379"}
    27  	}
    28  	if dialTimeout <= 0 {
    29  		dialTimeout = 3 * time.Second
    30  	}
    31  	m = &GrpcResolver{
    32  		EtcdClient: &EtcdClient{
    33  			endpoints:   endpoints,
    34  			dialTimeout: dialTimeout,
    35  		},
    36  		store: make(map[string]map[string]struct{}),
    37  	}
    38  
    39  	return
    40  }
    41  
    42  func (m *GrpcResolver) InitGrpcResolver() (err error) {
    43  	m.client, err = clientv3.New(clientv3.Config{Endpoints: m.endpoints, DialTimeout: m.dialTimeout})
    44  	if err != nil {
    45  		log.Error(err)
    46  		return
    47  	}
    48  
    49  	resolver.Register(m)
    50  
    51  	return
    52  }
    53  
    54  func (b *GrpcResolver) DebugStore() {
    55  	fmt.Printf("store %+v\n", b.store)
    56  }
    57  
    58  // Build creates a new resolver for the given target.
    59  //
    60  // gRPC dial calls Build synchronously, and fails if the returned error is
    61  // not nil.
    62  func (b *GrpcResolver) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
    63  	fmt.Printf("call builder %s\n", target.Endpoint)
    64  	b.store[target.Endpoint()] = make(map[string]struct{})
    65  
    66  	r := &etcdResolver{
    67  		client: b.client,
    68  		target: target,
    69  		cc:     cc,
    70  		store:  b.store[target.Endpoint()],
    71  		stopCh: make(chan struct{}, 1),
    72  		rn:     make(chan struct{}, 1),
    73  		t:      time.NewTicker(defaultFreq),
    74  	}
    75  
    76  	runtimer.GoSafely(nil, false, func() {
    77  		r.start(context.Background())
    78  	}, nil, nil)
    79  	r.ResolveNow(resolver.ResolveNowOptions{})
    80  
    81  	return r, nil
    82  }
    83  
    84  // Scheme returns the scheme supported by this resolver.
    85  // Scheme is defined at https://github.com/grpc/grpc/blob/master/doc/naming.md.
    86  func (b *GrpcResolver) Scheme() string {
    87  	return Scheme
    88  }
    89  
    90  type etcdResolver struct {
    91  	client *clientv3.Client
    92  	target resolver.Target
    93  	cc     resolver.ClientConn
    94  	store  map[string]struct{}
    95  	stopCh chan struct{}
    96  	rn     chan struct{} // rn channel is used by ResolveNow() to force an immediate resolution of the target.
    97  	t      *time.Ticker
    98  }
    99  
   100  func (r *etcdResolver) start(ctx context.Context) {
   101  	target := r.target.Endpoint()
   102  
   103  	w := clientv3.NewWatcher(r.client)
   104  	rch := w.Watch(ctx, target+"/", clientv3.WithPrefix())
   105  	for {
   106  		select {
   107  		case <-r.rn:
   108  			r.resolveNow()
   109  		case <-r.t.C:
   110  			r.ResolveNow(resolver.ResolveNowOptions{})
   111  		case <-r.stopCh:
   112  			err := w.Close()
   113  			if err != nil {
   114  				log.Errorf("etcd watcher close err", err.Error())
   115  			}
   116  			return
   117  		case wresp := <-rch:
   118  			for _, ev := range wresp.Events {
   119  				switch ev.Type {
   120  				case mvccpb.PUT:
   121  					r.store[string(ev.Kv.Value)] = struct{}{}
   122  					log.Info("etcd watcher put kv", string(ev.Kv.Key), string(ev.Kv.Value))
   123  				case mvccpb.DELETE:
   124  					delete(r.store, strings.Replace(string(ev.Kv.Key), target+"/", "", 1))
   125  					log.Info("etcd watcher delete kv", string(ev.Kv.Key), string(ev.Kv.Value))
   126  				}
   127  			}
   128  			r.updateTargetState()
   129  		}
   130  	}
   131  }
   132  
   133  func (r *etcdResolver) resolveNow() {
   134  	target := r.target.Endpoint()
   135  	resp, err := r.client.Get(context.Background(), target+"/", clientv3.WithPrefix())
   136  	if err != nil {
   137  		r.cc.ReportError(errors.Wrap(err, "get init endpoints"))
   138  		return
   139  	}
   140  
   141  	for _, kv := range resp.Kvs {
   142  		r.store[string(kv.Value)] = struct{}{}
   143  	}
   144  
   145  	r.updateTargetState()
   146  }
   147  
   148  func (r *etcdResolver) updateTargetState() {
   149  	addrs := make([]resolver.Address, len(r.store))
   150  	i := 0
   151  	for k := range r.store {
   152  		addrs[i] = resolver.Address{Addr: k}
   153  		i++
   154  	}
   155  	r.cc.UpdateState(resolver.State{Addresses: addrs})
   156  }
   157  
   158  // ResolveNow will be called by gRPC to try to resolve the target name
   159  // again. It's just a hint, resolver can ignore this if it's not necessary.
   160  //
   161  // It could be called multiple times concurrently.
   162  func (r *etcdResolver) ResolveNow(o resolver.ResolveNowOptions) {
   163  	select {
   164  	case r.rn <- struct{}{}:
   165  	default:
   166  
   167  	}
   168  }
   169  
   170  // Close closes the resolver.
   171  func (r *etcdResolver) Close() {
   172  	r.t.Stop()
   173  	close(r.stopCh)
   174  }