github.com/lingyao2333/mo-zero@v1.4.1/core/discov/publisher.go (about)

     1  package discov
     2  
     3  import (
     4  	"github.com/lingyao2333/mo-zero/core/discov/internal"
     5  	"github.com/lingyao2333/mo-zero/core/lang"
     6  	"github.com/lingyao2333/mo-zero/core/logx"
     7  	"github.com/lingyao2333/mo-zero/core/proc"
     8  	"github.com/lingyao2333/mo-zero/core/syncx"
     9  	"github.com/lingyao2333/mo-zero/core/threading"
    10  	clientv3 "go.etcd.io/etcd/client/v3"
    11  )
    12  
    13  type (
    14  	// PubOption defines the method to customize a Publisher.
    15  	PubOption func(client *Publisher)
    16  
    17  	// A Publisher can be used to publish the value to an etcd cluster on the given key.
    18  	Publisher struct {
    19  		endpoints  []string
    20  		key        string
    21  		fullKey    string
    22  		id         int64
    23  		value      string
    24  		lease      clientv3.LeaseID
    25  		quit       *syncx.DoneChan
    26  		pauseChan  chan lang.PlaceholderType
    27  		resumeChan chan lang.PlaceholderType
    28  	}
    29  )
    30  
    31  // NewPublisher returns a Publisher.
    32  // endpoints is the hosts of the etcd cluster.
    33  // key:value are a pair to be published.
    34  // opts are used to customize the Publisher.
    35  func NewPublisher(endpoints []string, key, value string, opts ...PubOption) *Publisher {
    36  	publisher := &Publisher{
    37  		endpoints:  endpoints,
    38  		key:        key,
    39  		value:      value,
    40  		quit:       syncx.NewDoneChan(),
    41  		pauseChan:  make(chan lang.PlaceholderType),
    42  		resumeChan: make(chan lang.PlaceholderType),
    43  	}
    44  
    45  	for _, opt := range opts {
    46  		opt(publisher)
    47  	}
    48  
    49  	return publisher
    50  }
    51  
    52  // KeepAlive keeps key:value alive.
    53  func (p *Publisher) KeepAlive() error {
    54  	cli, err := internal.GetRegistry().GetConn(p.endpoints)
    55  	if err != nil {
    56  		return err
    57  	}
    58  
    59  	p.lease, err = p.register(cli)
    60  	if err != nil {
    61  		return err
    62  	}
    63  
    64  	proc.AddWrapUpListener(func() {
    65  		p.Stop()
    66  	})
    67  
    68  	return p.keepAliveAsync(cli)
    69  }
    70  
    71  // Pause pauses the renewing of key:value.
    72  func (p *Publisher) Pause() {
    73  	p.pauseChan <- lang.Placeholder
    74  }
    75  
    76  // Resume resumes the renewing of key:value.
    77  func (p *Publisher) Resume() {
    78  	p.resumeChan <- lang.Placeholder
    79  }
    80  
    81  // Stop stops the renewing and revokes the registration.
    82  func (p *Publisher) Stop() {
    83  	p.quit.Close()
    84  }
    85  
    86  func (p *Publisher) keepAliveAsync(cli internal.EtcdClient) error {
    87  	ch, err := cli.KeepAlive(cli.Ctx(), p.lease)
    88  	if err != nil {
    89  		return err
    90  	}
    91  
    92  	threading.GoSafe(func() {
    93  		for {
    94  			select {
    95  			case _, ok := <-ch:
    96  				if !ok {
    97  					p.revoke(cli)
    98  					if err := p.KeepAlive(); err != nil {
    99  						logx.Errorf("KeepAlive: %s", err.Error())
   100  					}
   101  					return
   102  				}
   103  			case <-p.pauseChan:
   104  				logx.Infof("paused etcd renew, key: %s, value: %s", p.key, p.value)
   105  				p.revoke(cli)
   106  				select {
   107  				case <-p.resumeChan:
   108  					if err := p.KeepAlive(); err != nil {
   109  						logx.Errorf("KeepAlive: %s", err.Error())
   110  					}
   111  					return
   112  				case <-p.quit.Done():
   113  					return
   114  				}
   115  			case <-p.quit.Done():
   116  				p.revoke(cli)
   117  				return
   118  			}
   119  		}
   120  	})
   121  
   122  	return nil
   123  }
   124  
   125  func (p *Publisher) register(client internal.EtcdClient) (clientv3.LeaseID, error) {
   126  	resp, err := client.Grant(client.Ctx(), TimeToLive)
   127  	if err != nil {
   128  		return clientv3.NoLease, err
   129  	}
   130  
   131  	lease := resp.ID
   132  	if p.id > 0 {
   133  		p.fullKey = makeEtcdKey(p.key, p.id)
   134  	} else {
   135  		p.fullKey = makeEtcdKey(p.key, int64(lease))
   136  	}
   137  	_, err = client.Put(client.Ctx(), p.fullKey, p.value, clientv3.WithLease(lease))
   138  
   139  	return lease, err
   140  }
   141  
   142  func (p *Publisher) revoke(cli internal.EtcdClient) {
   143  	if _, err := cli.Revoke(cli.Ctx(), p.lease); err != nil {
   144  		logx.Error(err)
   145  	}
   146  }
   147  
   148  // WithId customizes a Publisher with the id.
   149  func WithId(id int64) PubOption {
   150  	return func(publisher *Publisher) {
   151  		publisher.id = id
   152  	}
   153  }
   154  
   155  // WithPubEtcdAccount provides the etcd username/password.
   156  func WithPubEtcdAccount(user, pass string) PubOption {
   157  	return func(pub *Publisher) {
   158  		RegisterAccount(pub.endpoints, user, pass)
   159  	}
   160  }
   161  
   162  // WithPubEtcdTLS provides the etcd CertFile/CertKeyFile/CACertFile.
   163  func WithPubEtcdTLS(certFile, certKeyFile, caFile string, insecureSkipVerify bool) PubOption {
   164  	return func(pub *Publisher) {
   165  		logx.Must(RegisterTLS(pub.endpoints, certFile, certKeyFile, caFile, insecureSkipVerify))
   166  	}
   167  }