etcd/lease.go

// Package etcd ...
//
// Description : 租约相关操作
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2021-11-23 6:03 下午
package etcd

import (
	"context"
	"time"

	"errors"
	clientv3 "go.etcd.io/etcd/client/v3"
)

// LeaseOnce 申请一个一次性租约
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 6:06 下午 2021/11/23
func (wc *WrapperClient) LeaseOnce(ctx context.Context, key string, val string, ttl int64) error {
	if ttl <= 0 {
		return errors.New("lease time must be more than 0")
	}

	if nil == ctx {
		ctx = context.Background()
	}

	var (
		resp       *clientv3.LeaseGrantResponse
		err        error
		cancelFunc context.CancelFunc
	)
	ctx, cancelFunc = context.WithCancel(ctx)
	defer cancelFunc()
	// 创建一个5秒的租约
	if resp, err = wc.client.Grant(ctx, ttl); err != nil {
		return errors.New("lease grant error : " + err.Error())
	}

	// ttl 秒钟之后,  这个key就会被移除
	if _, err = wc.client.Put(ctx, key, val, clientv3.WithLease(resp.ID)); err != nil {
		return errors.New("lease key put fail : " + err.Error())
	}
	_, err = wc.client.KeepAliveOnce(ctx, resp.ID)
	return err
}

// LeaseKeepAliveForever 无限续租一个key
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 7:40 下午 2021/11/23
func (wc *WrapperClient) LeaseKeepAliveForever(ctx context.Context, key string, val string, ttl int64, keepAliveHandler LeaseKeepALiveHandler) error {
	if ttl <= 0 {
		return errors.New("lease time must be more than 0")
	}
	if nil == ctx {
		ctx = context.TODO()
	}
	var (
		resp     *clientv3.LeaseGrantResponse
		respChan <-chan *clientv3.LeaseKeepAliveResponse
		err      error
	)
	// 创建一个5秒的租约
	if resp, err = wc.client.Grant(ctx, ttl); err != nil {
		return errors.New("lease grant error : " + err.Error())
	}

	// ttl 秒钟之后,  这个key就会被移除
	if _, err = wc.client.Put(context.TODO(), key, val, clientv3.WithLease(resp.ID)); err != nil {
		return errors.New("lease key put fail : " + err.Error())
	}
	// the key will be kept forever
	if respChan, err = wc.client.KeepAlive(ctx, resp.ID); nil != err {
		return errors.New("lease keep alive fail : " + err.Error())
	}
	leaseData := &LeaseKeepAliveData{
		Key:             key,
		StartTime:       time.Now().Unix(),
		LastLeaseTime:   0,
		LeaseCnt:        0,
		HasFinish:       false,
		LeaseFinishType: "",
		LeaseFinishTime: 0,
		LeaseDetail:     nil,
		Data:            make(map[string]interface{}),
	}
	// 监听 chan
	for ka := range respChan {
		leaseData.LeaseCnt++
		leaseData.LeaseDetail = ka
		leaseData.LastLeaseTime = time.Now().Unix()
		if nil != keepAliveHandler {
			keepAliveHandler(leaseData)
		}
	}
	return nil
}

// LeaseKeepAliveWithDuration 设置最大支持续期的时间, 中途可随时取消
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 8:05 下午 2021/11/23
func (wc *WrapperClient) LeaseKeepAliveWithDuration(ctx context.Context, key string, val string, ttl int64, keepAliveHandler LeaseKeepALiveHandler, cancelLeaseChan chan *LeaseKeepAliveData, maxCnt int64) (*LeaseKeepAliveData, error) {
	if ttl <= 0 {
		return nil, errors.New("lease time must be more than 0")
	}

	if nil == cancelLeaseChan {
		cancelLeaseChan = make(chan *LeaseKeepAliveData, 1)
	}
	var cancelFunc context.CancelFunc
	if nil == ctx {
		ctx, cancelFunc = context.WithCancel(context.Background())
	} else {
		ctx, cancelFunc = context.WithCancel(ctx)
	}
	defer cancelFunc()
	var (
		resp     *clientv3.LeaseGrantResponse
		respChan <-chan *clientv3.LeaseKeepAliveResponse
		err      error
	)
	leaseData := &LeaseKeepAliveData{
		Key:             key,
		StartTime:       0,
		LastLeaseTime:   0,
		LeaseCnt:        0,
		LeaseFinishType: "",
		LeaseFinishTime: 0,
		Data:            make(map[string]interface{}),
		HasFinish:       false,
	}
	// 创建一个 ttl 秒的租约
	if resp, err = wc.client.Grant(ctx, ttl); err != nil {
		return nil, errors.New("lease grant error : " + err.Error())
	}
	leaseData.StartTime = time.Now().Unix()

	// ttl 秒钟之后,  这个key就会被移除
	if _, err = wc.client.Put(context.TODO(), key, val, clientv3.WithLease(resp.ID)); err != nil {
		return nil, errors.New("lease key put fail : " + err.Error())
	}
	// the key will be kept forever
	if respChan, err = wc.client.KeepAlive(ctx, resp.ID); nil != err {
		return nil, errors.New("lease keep alive fail : " + err.Error())
	}

	for {
		select {
		case <-cancelLeaseChan:
			leaseData.HasFinish = true
			leaseData.LeaseFinishTime = time.Now().Unix()
			leaseData.LeaseFinishType = "SIGNAL_CANCEL"
		case leaseResp := <-respChan:
			leaseData.LeaseCnt++
			leaseData.LastLeaseTime = time.Now().Unix()
			leaseData.LeaseDetail = leaseResp
			if nil != keepAliveHandler {
				keepAliveHandler(leaseData)
			}
			if leaseData.LeaseCnt >= maxCnt {
				leaseData.HasFinish = true
				leaseData.LeaseFinishType = "OVER_MAX_CNT"
				leaseData.LeaseFinishTime = time.Now().Unix()
			}
		}
		if leaseData.HasFinish {
			break
		}
	}
	return leaseData, nil
}