cloudpods/pkg/cloudcommon/etcd/etcd.go

// Copyright 2019 Yunion
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package etcd

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"time"

	"go.etcd.io/etcd/clientv3"

	"yunion.io/x/log"

	"yunion.io/x/onecloud/pkg/util/seclib2"
)

var (
	ErrNoSuchKey = errors.New("No such key")
)

type SEtcdClient struct {
	client          *clientv3.Client
	requestTimeout  time.Duration
	leaseTtlTimeout int

	namespace string

	leaseId clientv3.LeaseID

	watchers map[string]*SEtcdWatcher
}

func NewEtcdClient(opt *SEtcdOptions) (*SEtcdClient, error) {
	var err error
	var tlsConfig *tls.Config

	if opt.EtcdEnabldSsl {
		tlsConfig, err = seclib2.InitTLSConfig(opt.EtcdSslCertfile, opt.EtcdSslKeyfile)
		if err != nil {
			log.Errorf("init tls config fail %s", err)
			return nil, err
		}
	}

	etcdClient := &SEtcdClient{}

	timeoutSeconds := opt.EtcdTimeoutSeconds
	if timeoutSeconds == 0 {
		timeoutSeconds = 5
	}

	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   opt.EtcdEndpoint,
		DialTimeout: time.Duration(timeoutSeconds) * time.Second,
		Username:    opt.EtcdUsername,
		Password:    opt.EtcdPassword,
		TLS:         tlsConfig,
	})
	if err != nil {
		return nil, err
	}

	etcdClient.client = cli

	timeoutSeconds = opt.EtcdRequestTimeoutSeconds
	if timeoutSeconds == 0 {
		timeoutSeconds = 2
	}
	etcdClient.requestTimeout = time.Duration(timeoutSeconds) * time.Second

	timeoutSeconds = opt.EtcdLeaseExpireSeconds
	if timeoutSeconds == 0 {
		timeoutSeconds = 5
	}

	etcdClient.leaseTtlTimeout = timeoutSeconds

	etcdClient.watchers = make(map[string]*SEtcdWatcher)

	etcdClient.namespace = opt.EtcdNamspace

	err = etcdClient.startSession()
	if err != nil {
		etcdClient.Close()
		return nil, err
	}
	return etcdClient, nil
}

func (cli *SEtcdClient) Close() error {
	if cli.client != nil {
		err := cli.client.Close()
		if err != nil {
			return err
		}
		cli.client = nil
	}
	return nil
}

func (cli *SEtcdClient) GetClient() *clientv3.Client {
	return cli.client
}

func (cli *SEtcdClient) startSession() error {
	ctx := context.TODO()

	resp, err := cli.client.Grant(ctx, int64(cli.leaseTtlTimeout))
	if err != nil {
		return err
	}
	cli.leaseId = resp.ID

	ch, err := cli.client.KeepAlive(ctx, cli.leaseId)
	if err != nil {
		return err
	}

	go func() {
		for {
			ka := <-ch
			if ka == nil {
				log.Fatalf("fail to keepalive")
			} else {
				log.Debugf("etcd session %d keepalive ttl: %d", ka.ID, ka.TTL)
			}
		}
	}()

	return nil
}

func (cli *SEtcdClient) getKey(key string) string {
	if len(cli.namespace) > 0 {
		return fmt.Sprintf("%s%s", cli.namespace, key)
	} else {
		return key
	}
}

func (cli *SEtcdClient) Put(ctx context.Context, key string, val string) error {
	return cli.put(ctx, key, val, false)
}

func (cli *SEtcdClient) PutSession(ctx context.Context, key string, val string) error {
	return cli.put(ctx, key, val, true)
}

func (cli *SEtcdClient) put(ctx context.Context, key string, val string, session bool) error {
	nctx, cancel := context.WithTimeout(ctx, cli.requestTimeout)
	defer cancel()

	key = cli.getKey(key)
	if session {
		_, err := cli.client.Put(nctx, key, val, clientv3.WithLease(cli.leaseId))
		return err
	} else {
		_, err := cli.client.Put(nctx, key, val)
		return err
	}
}

func (cli *SEtcdClient) Get(ctx context.Context, key string) ([]byte, error) {
	nctx, cancel := context.WithTimeout(ctx, cli.requestTimeout)
	defer cancel()

	key = cli.getKey(key)

	resp, err := cli.client.Get(nctx, key)
	if err != nil {
		return nil, err
	}
	if len(resp.Kvs) == 0 {
		return nil, ErrNoSuchKey
	}

	return resp.Kvs[0].Value, nil
}

type SEtcdKeyValue struct {
	Key   string
	Value []byte
}

func (cli *SEtcdClient) List(ctx context.Context, prefix string) ([]SEtcdKeyValue, error) {
	nctx, cancel := context.WithTimeout(ctx, cli.requestTimeout)
	defer cancel()

	prefix = cli.getKey(prefix)

	resp, err := cli.client.Get(nctx, prefix, clientv3.WithPrefix(),
		clientv3.WithSort(clientv3.SortByKey, clientv3.SortDescend))
	if err != nil {
		return nil, err
	}
	ret := make([]SEtcdKeyValue, len(resp.Kvs))
	for i := 0; i < len(resp.Kvs); i += 1 {
		ret[i] = SEtcdKeyValue{
			Key:   string(resp.Kvs[i].Key[len(cli.namespace):]),
			Value: resp.Kvs[i].Value,
		}
	}
	return ret, nil
}

type TEtcdCreateEventFunc func(key, value []byte)
type TEtcdModifyEventFunc func(key, oldvalue, value []byte)

type SEtcdWatcher struct {
	watcher clientv3.Watcher
	cancel  context.CancelFunc
}

func (w *SEtcdWatcher) Cancel() {
	w.watcher.Close()
	w.cancel()
}

func (cli *SEtcdClient) Watch(ctx context.Context, prefix string, onCreate TEtcdCreateEventFunc, onModify TEtcdModifyEventFunc) {
	_, ok := cli.watchers[prefix]
	if ok {
		return
	}

	watcher := clientv3.NewWatcher(cli.client)
	nctx, cancel := context.WithCancel(ctx)

	cli.watchers[prefix] = &SEtcdWatcher{
		watcher: watcher,
		cancel:  cancel,
	}

	prefix = cli.getKey(prefix)

	rch := watcher.Watch(nctx, prefix, clientv3.WithPrefix(), clientv3.WithPrevKV())
	go func() {
		for wresp := range rch {
			for _, ev := range wresp.Events {
				if ev.PrevKv == nil {
					onCreate(ev.Kv.Key[len(cli.namespace):], ev.Kv.Value)
				} else {
					onModify(ev.Kv.Key[len(cli.namespace):], ev.PrevKv.Value, ev.Kv.Value)
				}
			}
		}
		log.Infof("stop watching %s", prefix)
	}()
}

func (cli *SEtcdClient) Unwatch(prefix string) {
	watcher, ok := cli.watchers[prefix]
	if ok {
		log.Debugf("unwatch %s", prefix)
		watcher.Cancel()
		delete(cli.watchers, prefix)
	} else {
		log.Debugf("prefix %s not watched!!", prefix)
	}
}

func (cli *SEtcdClient) Delete(ctx context.Context, key string) ([]byte, error) {
	nctx, cancel := context.WithTimeout(ctx, cli.requestTimeout)
	defer cancel()

	key = cli.getKey(key)

	dresp, err := cli.client.Delete(nctx, key, clientv3.WithPrevKV())
	if err != nil {
		return nil, err
	}

	if dresp.Deleted == 1 {
		return dresp.PrevKvs[0].Value, nil
	} else {
		return nil, nil
	}
}