github.com/volts-dev/volts@v0.0.0-20240120094013-5e9c65924106/registry/etcd/watcher.go (about)

     1  package etcd
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"time"
     7  
     8  	"github.com/volts-dev/volts/registry"
     9  	clientv3 "go.etcd.io/etcd/client/v3"
    10  )
    11  
    12  type etcdWatcher struct {
    13  	stop    chan bool
    14  	w       clientv3.WatchChan
    15  	client  *clientv3.Client
    16  	timeout time.Duration
    17  }
    18  
    19  func newEtcdWatcher(r *etcdRegistry, timeout time.Duration, opts ...registry.WatchOptions) (registry.Watcher, error) {
    20  	var wo registry.WatchConfig
    21  	for _, o := range opts {
    22  		o(&wo)
    23  	}
    24  
    25  	ctx, cancel := context.WithCancel(context.Background())
    26  	stop := make(chan bool, 1)
    27  
    28  	go func() {
    29  		<-stop
    30  		cancel()
    31  	}()
    32  
    33  	watchPath := prefix
    34  	if len(wo.Service) > 0 {
    35  		watchPath = servicePath(wo.Service) + "/"
    36  	}
    37  
    38  	return &etcdWatcher{
    39  		stop:    stop,
    40  		w:       r.client.Watch(ctx, watchPath, clientv3.WithPrefix(), clientv3.WithPrevKV()),
    41  		client:  r.client,
    42  		timeout: timeout,
    43  	}, nil
    44  }
    45  
    46  func (ew *etcdWatcher) Next() (*registry.Result, error) {
    47  	for wresp := range ew.w {
    48  		if wresp.Err() != nil {
    49  			return nil, wresp.Err()
    50  		}
    51  		if wresp.Canceled {
    52  			return nil, errors.New("could not get next")
    53  		}
    54  
    55  		var service *registry.Service
    56  		for _, ev := range wresp.Events {
    57  			var action string
    58  
    59  			switch ev.Type {
    60  			case clientv3.EventTypePut:
    61  				service = decode(ev.Kv.Value)
    62  
    63  				if ev.IsCreate() {
    64  					action = "create"
    65  				} else if ev.IsModify() {
    66  					action = "update"
    67  				}
    68  			case clientv3.EventTypeDelete:
    69  				service = nil
    70  				action = "delete"
    71  
    72  				// get service from prevKv
    73  				service = decode(ev.PrevKv.Value)
    74  			}
    75  
    76  			if service == nil {
    77  				continue
    78  			}
    79  			return &registry.Result{
    80  				Action:  action,
    81  				Service: service,
    82  			}, nil
    83  		}
    84  	}
    85  	return nil, errors.New("could not get next")
    86  }
    87  
    88  func (ew *etcdWatcher) Stop() {
    89  	select {
    90  	case <-ew.stop:
    91  		return
    92  	default:
    93  		close(ew.stop)
    94  	}
    95  }