micro/store/etcd/etcd.go
Asim Aslam 078dd4eb9b
fix etcd (#1340)
* fix etcd

* update go mod
2020-03-13 08:55:23 +00:00

269 lines
6.7 KiB
Go

// Package etcd implements a go-micro/v2/store with etcd
package etcd
import (
"bytes"
"context"
"encoding/gob"
"math"
"strings"
"time"
"github.com/micro/go-micro/v2/store"
"github.com/pkg/errors"
"github.com/coreos/etcd/clientv3"
"github.com/coreos/etcd/clientv3/namespace"
)
type etcdStore struct {
options store.Options
client *clientv3.Client
config clientv3.Config
}
// NewStore returns a new etcd store
func NewStore(opts ...store.Option) store.Store {
e := &etcdStore{}
for _, o := range opts {
o(&e.options)
}
e.init()
return e
}
func (e *etcdStore) Init(opts ...store.Option) error {
for _, o := range opts {
o(&e.options)
}
return e.init()
}
func (e *etcdStore) init() error {
// ensure context is non-nil
e.options.Context = context.Background()
// set up config
e.config = clientv3.Config{}
e.applyConfig(&e.config)
if len(e.options.Nodes) == 0 {
e.config.Endpoints = []string{"http://127.0.0.1:2379"}
} else {
e.config.Endpoints = make([]string, len(e.options.Nodes))
copy(e.config.Endpoints, e.options.Nodes)
}
if e.client != nil {
e.client.Close()
}
client, err := clientv3.New(e.config)
if err != nil {
return err
}
e.client = client
ns := ""
if len(e.options.Prefix) > 0 {
ns = e.options.Prefix
}
if len(e.options.Namespace) > 0 {
ns = e.options.Namespace + "/" + ns
}
if len(ns) > 0 {
e.client.KV = namespace.NewKV(e.client.KV, ns)
e.client.Watcher = namespace.NewWatcher(e.client.Watcher, ns)
e.client.Lease = namespace.NewLease(e.client.Lease, ns)
}
return nil
}
func (e *etcdStore) Options() store.Options {
return e.options
}
func (e *etcdStore) String() string {
return "etcd"
}
func (e *etcdStore) Read(key string, opts ...store.ReadOption) ([]*store.Record, error) {
readOpts := store.ReadOptions{}
for _, o := range opts {
o(&readOpts)
}
if readOpts.Suffix {
return e.readSuffix(key, readOpts)
}
var etcdOpts []clientv3.OpOption
if readOpts.Prefix {
etcdOpts = append(etcdOpts, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend))
}
resp, err := e.client.KV.Get(context.Background(), key, etcdOpts...)
if err != nil {
return nil, err
}
if resp.Count == 0 && !(readOpts.Prefix || readOpts.Suffix) {
return nil, store.ErrNotFound
}
var records []*store.Record
for _, kv := range resp.Kvs {
ir := internalRecord{}
if err := gob.NewDecoder(bytes.NewReader(kv.Value)).Decode(&ir); err != nil {
return records, errors.Wrapf(err, "couldn't decode %s into internalRecord", err.Error())
}
r := store.Record{
Key: ir.Key,
Value: ir.Value,
}
if !ir.ExpiresAt.IsZero() {
r.Expiry = time.Until(ir.ExpiresAt)
}
records = append(records, &r)
}
if readOpts.Limit > 0 || readOpts.Offset > 0 {
return records[readOpts.Offset:min(readOpts.Limit, uint(len(records)))], nil
}
return records, nil
}
func (e *etcdStore) readSuffix(key string, readOpts store.ReadOptions) ([]*store.Record, error) {
opts := []store.ListOption{store.ListSuffix(key)}
if readOpts.Prefix {
opts = append(opts, store.ListPrefix(key))
}
keys, err := e.List(opts...)
if err != nil {
return nil, errors.Wrapf(err, "Couldn't list with suffix %s", key)
}
var records []*store.Record
for _, k := range keys {
resp, err := e.client.KV.Get(context.Background(), k)
if err != nil {
return nil, errors.Wrapf(err, "Couldn't get key %s", k)
}
ir := internalRecord{}
if err := gob.NewDecoder(bytes.NewReader(resp.Kvs[0].Value)).Decode(&ir); err != nil {
return records, errors.Wrapf(err, "couldn't decode %s into internalRecord", err.Error())
}
r := store.Record{
Key: ir.Key,
Value: ir.Value,
}
if !ir.ExpiresAt.IsZero() {
r.Expiry = time.Until(ir.ExpiresAt)
}
records = append(records, &r)
}
if readOpts.Limit > 0 || readOpts.Offset > 0 {
return records[readOpts.Offset:min(readOpts.Limit, uint(len(records)))], nil
}
return records, nil
}
func (e *etcdStore) Write(r *store.Record, opts ...store.WriteOption) error {
options := store.WriteOptions{}
for _, o := range opts {
o(&options)
}
if len(opts) > 0 {
// Copy the record before applying options, or the incoming record will be mutated
newRecord := store.Record{}
newRecord.Key = r.Key
newRecord.Value = make([]byte, len(r.Value))
copy(newRecord.Value, r.Value)
newRecord.Expiry = r.Expiry
if !options.Expiry.IsZero() {
newRecord.Expiry = time.Until(options.Expiry)
}
if options.TTL != 0 {
newRecord.Expiry = options.TTL
}
return e.write(&newRecord)
}
return e.write(r)
}
func (e *etcdStore) write(r *store.Record) error {
var putOpts []clientv3.OpOption
ir := &internalRecord{}
ir.Key = r.Key
ir.Value = make([]byte, len(r.Value))
copy(ir.Value, r.Value)
if r.Expiry != 0 {
ir.ExpiresAt = time.Now().Add(r.Expiry)
var leasexpiry int64
if r.Expiry.Seconds() < 5.0 {
// minimum etcd lease is 5 seconds
leasexpiry = 5
} else {
leasexpiry = int64(math.Ceil(r.Expiry.Seconds()))
}
lr, err := e.client.Lease.Grant(context.Background(), leasexpiry)
if err != nil {
return errors.Wrapf(err, "couldn't grant an etcd lease for %s", r.Key)
}
putOpts = append(putOpts, clientv3.WithLease(lr.ID))
}
b := &bytes.Buffer{}
if err := gob.NewEncoder(b).Encode(ir); err != nil {
return errors.Wrapf(err, "couldn't encode %s", r.Key)
}
_, err := e.client.KV.Put(context.Background(), ir.Key, string(b.Bytes()), putOpts...)
return errors.Wrapf(err, "couldn't put key %s in to etcd", err)
}
func (e *etcdStore) Delete(key string, opts ...store.DeleteOption) error {
options := store.DeleteOptions{}
for _, o := range opts {
o(&options)
}
_, err := e.client.KV.Delete(context.Background(), key)
return errors.Wrapf(err, "couldn't delete key %s", key)
}
func (e *etcdStore) List(opts ...store.ListOption) ([]string, error) {
options := store.ListOptions{}
for _, o := range opts {
o(&options)
}
searchPrefix := ""
if len(options.Prefix) > 0 {
searchPrefix = options.Prefix
}
resp, err := e.client.KV.Get(context.Background(), searchPrefix, clientv3.WithPrefix(), clientv3.WithKeysOnly(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend))
if err != nil {
return nil, errors.Wrap(err, "couldn't list, etcd get failed")
}
if len(options.Suffix) == 0 {
keys := make([]string, resp.Count)
for i, kv := range resp.Kvs {
keys[i] = string(kv.Key)
}
return keys, nil
}
keys := []string{}
for _, kv := range resp.Kvs {
if strings.HasSuffix(string(kv.Key), options.Suffix) {
keys = append(keys, string(kv.Key))
}
}
if options.Limit > 0 || options.Offset > 0 {
return keys[options.Offset:min(options.Limit, uint(len(keys)))], nil
}
return keys, nil
}
type internalRecord struct {
Key string
Value []byte
ExpiresAt time.Time
}
func min(i, j uint) uint {
if i < j {
return i
}
return j
}