github.com/vicanso/pike@v1.0.1-0.20210630235453-9099e041f6ec/config/etcd_client.go (about)

     1  // MIT License
     2  
     3  // Copyright (c) 2020 Tree Xie
     4  
     5  // Permission is hereby granted, free of charge, to any person obtaining a copy
     6  // of this software and associated documentation files (the "Software"), to deal
     7  // in the Software without restriction, including without limitation the rights
     8  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     9  // copies of the Software, and to permit persons to whom the Software is
    10  // furnished to do so, subject to the following conditions:
    11  
    12  // The above copyright notice and this permission notice shall be included in all
    13  // copies or substantial portions of the Software.
    14  
    15  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    16  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    17  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    18  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    19  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    20  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    21  // SOFTWARE.
    22  
    23  // etcd client for config
    24  
    25  package config
    26  
    27  import (
    28  	"context"
    29  	"crypto/tls"
    30  	"net/url"
    31  	"strings"
    32  	"time"
    33  
    34  	"github.com/coreos/etcd/clientv3"
    35  )
    36  
    37  // etcdClient etcd client
    38  type etcdClient struct {
    39  	c       *clientv3.Client
    40  	key     string
    41  	Timeout time.Duration
    42  }
    43  
    44  const (
    45  	defaultEtcdTimeout = 5 * time.Second
    46  )
    47  
    48  // NewEtcdClient create a new etcd client
    49  func NewEtcdClient(uri string) (client *etcdClient, err error) {
    50  	u, err := url.Parse(uri)
    51  	if err != nil {
    52  		return
    53  	}
    54  	conf := clientv3.Config{
    55  		Endpoints:   strings.Split(u.Host, ","),
    56  		DialTimeout: defaultEtcdTimeout,
    57  	}
    58  	if u.User != nil {
    59  		conf.Username = u.User.Username()
    60  		conf.Password, _ = u.User.Password()
    61  	}
    62  	cert := u.Query().Get("cert")
    63  	key := u.Query().Get("key")
    64  	if cert != "" && key != "" {
    65  		tlsConfig := tls.Config{}
    66  		tlsConfig.Certificates = make([]tls.Certificate, 1)
    67  		// TODO 支持更多种形式的拉取证书，如HTTP的方式
    68  		tlsConfig.Certificates[0], err = tls.LoadX509KeyPair(cert, key)
    69  		if err != nil {
    70  			return nil, err
    71  		}
    72  		conf.TLS = &tlsConfig
    73  	}
    74  	// TODO 后续支持从querystring中配置更多的参数
    75  	c, err := clientv3.New(conf)
    76  	if err != nil {
    77  		return
    78  	}
    79  	client = &etcdClient{
    80  		c:   c,
    81  		key: u.Path,
    82  	}
    83  	return
    84  }
    85  
    86  func (ec *etcdClient) context() (context.Context, context.CancelFunc) {
    87  	d := ec.Timeout
    88  	if d == 0 {
    89  		d = defaultEtcdTimeout
    90  	}
    91  	return context.WithTimeout(context.Background(), d)
    92  }
    93  
    94  // Get get data from etcd
    95  func (ec *etcdClient) Get() (data []byte, err error) {
    96  	ctx, cancel := ec.context()
    97  	defer cancel()
    98  	resp, err := ec.c.Get(ctx, ec.key)
    99  	if err != nil {
   100  		return
   101  	}
   102  	kvs := resp.Kvs
   103  	if len(kvs) == 0 {
   104  		return
   105  	}
   106  	data = kvs[0].Value
   107  	return
   108  }
   109  
   110  // Set set data to etcd
   111  func (ec *etcdClient) Set(data []byte) (err error) {
   112  	ctx, cancel := ec.context()
   113  	defer cancel()
   114  	_, err = ec.c.Put(ctx, ec.key, string(data))
   115  	return
   116  }
   117  
   118  // Watch watch config change
   119  func (ec *etcdClient) Watch(onChange OnChange) {
   120  	ch := ec.c.Watch(context.Background(), ec.key)
   121  	// 只监听有变化则可
   122  	for range ch {
   123  		onChange()
   124  	}
   125  }
   126  
   127  // Close close etcd client
   128  func (ec *etcdClient) Close() error {
   129  	return ec.c.Close()
   130  }