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

     1  package discovery
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  	"time"
     7  
     8  	"github.com/weedge/lib/log"
     9  	"github.com/weedge/lib/runtimer"
    10  
    11  	"github.com/pkg/errors"
    12  	"go.etcd.io/etcd/client/v3"
    13  )
    14  
    15  type Register struct {
    16  	*EtcdClient
    17  	serviceName string
    18  	selfIpPort  string
    19  	cancel      context.CancelFunc
    20  }
    21  
    22  func NewRegister(endpoints []string, dialTimeout time.Duration, serviceName, selfIpPort string) (m *Register) {
    23  	if len(endpoints) == 0 {
    24  		endpoints = []string{"0.0.0.0:2379"}
    25  	}
    26  	if dialTimeout <= 0 {
    27  		dialTimeout = 3 * time.Second
    28  	}
    29  
    30  	m = &Register{
    31  		serviceName: serviceName,
    32  		selfIpPort:  selfIpPort,
    33  		EtcdClient: &EtcdClient{
    34  			endpoints:   endpoints,
    35  			dialTimeout: dialTimeout,
    36  		},
    37  	}
    38  
    39  	return
    40  }
    41  
    42  func (m *Register) Do() (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  	var ctx context.Context
    50  	ctx, m.cancel = context.WithCancel(context.Background())
    51  
    52  	runtimer.GoSafely(nil, false, func() {
    53  		err = m.register(ctx)
    54  		if err != nil {
    55  			log.Error(err)
    56  			return
    57  		}
    58  	}, nil, nil)
    59  
    60  	return
    61  }
    62  
    63  func (m *Register) register(ctx context.Context) error {
    64  	resp, err := m.client.Grant(ctx, 5)
    65  	if err != nil {
    66  		return errors.Wrap(err, "etcd grant")
    67  	}
    68  	_, err = m.client.Put(ctx, strings.Join([]string{m.serviceName, m.selfIpPort}, "/"), m.selfIpPort, clientv3.WithLease(resp.ID))
    69  	if err != nil {
    70  		return errors.Wrap(err, "etcd put")
    71  	}
    72  	log.Infof("etcd put ok; path:%s selfIpPort:%s leaseId:%d",
    73  		strings.Join([]string{m.serviceName, m.selfIpPort}, "/"), m.selfIpPort, resp.ID)
    74  	respCh, err := m.client.KeepAlive(ctx, resp.ID)
    75  	if err != nil {
    76  		return errors.Wrap(err, "etcd keep alive")
    77  	}
    78  
    79  	for {
    80  		select {
    81  		case <-ctx.Done():
    82  			return nil
    83  		case res, ok := <-respCh:
    84  			if ok {
    85  				if DEBUG {
    86  					log.Infof("etcd keepalive resp:%+v", res)
    87  				}
    88  			}
    89  		}
    90  	}
    91  }
    92  
    93  func (m *Register) Close() {
    94  
    95  	m.cancel()
    96  }