增加带次数的租约续期

middleware/etcd/lease.go

@@ -9,6 +9,7 @@ package etcd
 
 import (
 	"context"
+	"time"
 
 	"github.com/pkg/errors"
 	"go.etcd.io/etcd/clientv3"
@@ -72,11 +73,101 @@ func LeaseKeepAliveForever(ctx context.Context, key string, val string, ttl int6
 	if respChan, err = 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(key, ka)
+			keepAliveHandler(leaseData)
 		}
 	}
 	return nil
 }
+
+// LeaseKeepAliveWithDuration 设置最大支持续期的时间, 中途可随时取消
+//
+// Author : go_developer@163.com<白茶清欢>
+//
+// Date : 8:05 下午 2021/11/23
+func 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 = Client.Grant(ctx, ttl); err != nil {
+		return nil, errors.New("lease grant error : " + err.Error())
+	}
+	leaseData.StartTime = time.Now().Unix()
+
+	// ttl 秒钟之后,  这个key就会被移除
+	if _, err = 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 = 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
+}