github.com/ppphp/yayagf@v0.0.1/pkg/etcd/etcd.go (about)

     1  package etcd
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"go.etcd.io/etcd/api/v3/mvccpb"
     9  	"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
    10  	clientv3 "go.etcd.io/etcd/client/v3"
    11  	"go.uber.org/atomic"
    12  )
    13  
    14  func RegisterService(context context.Context, etcdaddrs []string, serviceName string) error {
    15  	c, err := clientv3.New(clientv3.Config{
    16  		Endpoints:   etcdaddrs,
    17  		DialTimeout: 5 * time.Second,
    18  	})
    19  	if err != nil {
    20  		return err
    21  	}
    22  
    23  	kv := clientv3.NewKV(c)
    24  	lease := clientv3.NewLease(c)
    25  	var curLeaseId clientv3.LeaseID = 0
    26  	timer := time.NewTicker(1 * time.Second)
    27  
    28  	go func() {
    29  		for {
    30  			select {
    31  			case <-timer.C:
    32  				if curLeaseId == 0 {
    33  					leaseResp, err := lease.Grant(context, 10)
    34  					if err != nil {
    35  						panic(err)
    36  					}
    37  
    38  					key := "/services/" + serviceName + "/" + fmt.Sprintf("%d", leaseResp.ID)
    39  					if _, err := kv.Put(context, key, value, clientv3.WithLease(leaseResp.ID)); err != nil {
    40  						panic(err)
    41  					}
    42  					curLeaseId = leaseResp.ID
    43  				} else {
    44  					// 续约租约，如果租约已经过期将curLeaseId复位到0重新走创建租约的逻辑
    45  					if _, err := lease.KeepAliveOnce(context, curLeaseId); err == rpctypes.ErrLeaseNotFound {
    46  						curLeaseId = 0
    47  						continue
    48  					}
    49  				}
    50  			case <-context.Done():
    51  				return
    52  			}
    53  		}
    54  		defer c.Close()
    55  	}()
    56  
    57  	return nil
    58  }
    59  
    60  func GetService(context context.Context, etcdaddrs []string, serviceName string) *remoteService {
    61  	c, err := clientv3.New(clientv3.Config{
    62  		Endpoints:   etcdaddrs,
    63  		DialTimeout: 5 * time.Second,
    64  	})
    65  	if err != nil {
    66  		return err
    67  	}
    68  	kv := clientv3.NewKV(c)
    69  	rangeResp, err := kv.Get(context, "/services/"+serviceName, clientv3.WithPrefix())
    70  	if err != nil {
    71  		panic(err)
    72  	}
    73  
    74  	service.mutex.Lock()
    75  	for _, kv := range rangeResp.Kvs {
    76  		service.nodes[string(kv.Key)] = string(kv.Value)
    77  	}
    78  	service.mutex.Unlock()
    79  	watcher := clientv3.NewWatcher(client)
    80  	watchChan := watcher.Watch(context, service.name, clientv3.WithPrefix())
    81  
    82  	go func() {
    83  		for watchResp := range watchChan {
    84  			for _, event := range watchResp.Events {
    85  				service.mutex.Lock()
    86  				switch event.Type {
    87  				case mvccpb.PUT: //PUT事件，目录下有了新key
    88  					service.nodes[string(event.Kv.Key)] = string(event.Kv.Value)
    89  				case mvccpb.DELETE: //DELETE事件，目录中有key被删掉(Lease过期，key 也会被删掉)
    90  					delete(service.nodes, string(event.Kv.Key))
    91  				}
    92  				service.mutex.Unlock()
    93  			}
    94  		}
    95  	}()
    96  }
    97  
    98  type RoundRipper struct {
    99  	host  string
   100  	index atomic.Int64
   101  	ips   []string
   102  }
   103  
   104  func (r *RoundRipper) GetIP() string {
   105  	r1 := r.index.Add(1)
   106  	return r.ips[r1%len(r.ips)]
   107  }
   108  
   109  func NewRoundRipper(host string, ips []string) *RoundRipper {
   110  	return &RoundRipper{index: 0, ips: ips}
   111  }