github.com/15mga/kiwi@v0.0.2-0.20240324021231-b95d5c3ac751/util/etd/etcd.go (about)

     1  package etd
     2  
     3  import (
     4  	"github.com/15mga/kiwi/util"
     5  	etcd "go.etcd.io/etcd/client/v3"
     6  	"go.etcd.io/etcd/client/v3/concurrency"
     7  )
     8  
     9  var (
    10  	_Etcd    *etcd.Client
    11  	_LockTtl = 10
    12  )
    13  
    14  func Client() *etcd.Client {
    15  	return _Etcd
    16  }
    17  
    18  func Conn(cfg etcd.Config) *util.Err {
    19  	client, e := etcd.New(cfg)
    20  	if e != nil {
    21  		return util.WrapErr(util.EcEtcdErr, e)
    22  	}
    23  	_Etcd = client
    24  	return nil
    25  }
    26  
    27  func Grant(ttl int64) (int64, *util.Err) {
    28  	res, e := _Etcd.Grant(util.Ctx(), ttl)
    29  	if e != nil {
    30  		return 0, util.WrapErr(util.EcEtcdErr, e)
    31  	}
    32  	id := res.ID
    33  	ch, e := _Etcd.KeepAlive(util.Ctx(), id)
    34  	if e != nil {
    35  		return 0, util.WrapErr(util.EcEtcdErr, e)
    36  	}
    37  	go func() {
    38  		for r := range ch {
    39  			if r != nil {
    40  				continue
    41  			}
    42  			return
    43  		}
    44  	}()
    45  	return int64(id), nil
    46  }
    47  
    48  func Revoke(id int64) *util.Err {
    49  	_, e := _Etcd.Revoke(util.Ctx(), etcd.LeaseID(id))
    50  	return util.WrapErr(util.EcEtcdErr, e)
    51  }
    52  
    53  func Del(key string, opts ...etcd.OpOption) *util.Err {
    54  	_, e := _Etcd.Delete(util.Ctx(), key, opts...)
    55  	if e != nil {
    56  		return util.WrapErr(util.EcEtcdErr, e)
    57  	}
    58  	return nil
    59  }
    60  
    61  func Put(key, val string) *util.Err {
    62  	_, e := _Etcd.Put(util.Ctx(), key, val)
    63  	if e != nil {
    64  		return util.WrapErr(util.EcEtcdErr, e)
    65  	}
    66  	return nil
    67  }
    68  
    69  func PutWithTtl(key, val string, ttl int64) (int64, *util.Err) {
    70  	if ttl <= 0 {
    71  		return 0, util.NewErr(util.EcParamsErr, util.M{
    72  			"ttl": ttl,
    73  		})
    74  	}
    75  	id, err := Grant(ttl)
    76  	if err != nil {
    77  		return 0, err
    78  	}
    79  	_, e := _Etcd.Put(util.Ctx(), key, val, etcd.WithLease(etcd.LeaseID(id)))
    80  	return id, util.WrapErr(util.EcEtcdErr, e)
    81  }
    82  
    83  func Lock(key string, fn util.ToErr) *util.Err {
    84  	if fn == nil {
    85  		return nil
    86  	}
    87  
    88  	s, e := concurrency.NewSession(_Etcd, concurrency.WithTTL(_LockTtl))
    89  	if e != nil {
    90  		return util.WrapErr(util.EcEtcdErr, e)
    91  	}
    92  	defer s.Close()
    93  	m := concurrency.NewMutex(s, key)
    94  
    95  	e = m.Lock(util.Ctx())
    96  	if e != nil {
    97  		return util.WrapErr(util.EcEtcdErr, e)
    98  	}
    99  	err := fn()
   100  	_ = m.Unlock(util.Ctx())
   101  	return err
   102  }
   103  
   104  func TryLock(key string, fn util.Fn) *util.Err {
   105  	if fn == nil {
   106  		return nil
   107  	}
   108  
   109  	s, e := concurrency.NewSession(_Etcd, concurrency.WithTTL(_LockTtl))
   110  	if e != nil {
   111  		return util.WrapErr(util.EcEtcdErr, e)
   112  	}
   113  	defer s.Close()
   114  	m := concurrency.NewMutex(s, key)
   115  
   116  	e = m.TryLock(util.Ctx())
   117  	if e != nil {
   118  		return util.WrapErr(util.EcEtcdErr, e)
   119  	}
   120  	fn()
   121  	_ = m.Unlock(util.Ctx())
   122  	return nil
   123  }
   124  
   125  func Get(key string, fn util.StrBytesToBool, opts ...etcd.OpOption) (bool, *util.Err) {
   126  	res, e := _Etcd.Get(util.Ctx(), key, opts...)
   127  	if e != nil {
   128  		return false, util.NewErr(util.EcEtcdErr, util.M{
   129  			"key": key,
   130  		})
   131  	}
   132  	if len(res.Kvs) == 0 {
   133  		return false, nil
   134  	}
   135  	for _, kv := range res.Kvs {
   136  		if ok := fn(string(kv.Key), kv.Value); !ok {
   137  			return false, nil
   138  		}
   139  	}
   140  	return true, nil
   141  }
   142  
   143  func GetOne(key string, opts ...etcd.OpOption) ([]byte, *util.Err) {
   144  	res, e := _Etcd.Get(util.Ctx(), key, opts...)
   145  	if e != nil {
   146  		return nil, util.NewErr(util.EcEtcdErr, util.M{
   147  			"key": key,
   148  		})
   149  	}
   150  	if len(res.Kvs) == 0 {
   151  		return nil, util.NewErr(util.EcNotExist, util.M{
   152  			"key": key,
   153  		})
   154  	}
   155  	return res.Kvs[0].Value, nil
   156  }
   157  
   158  func Has(key string, opts ...etcd.OpOption) (bool, *util.Err) {
   159  	exist := false
   160  	_, err := Get(key, func(s string, bytes []byte) bool {
   161  		exist = true
   162  		return true
   163  	}, opts...)
   164  	if err != nil {
   165  		return false, err
   166  	}
   167  	return exist, nil
   168  }