From e90826302ed48fc14a33c7dea22af2c440abbd9d Mon Sep 17 00:00:00 2001
From: vinllen <cvinllen@gmail.com>
Date: Thu, 4 Apr 2019 14:19:42 +0800
Subject: [PATCH] use restore to recovery key

---
 src/pkg/rdb/loader.go           |  8 ++++++--
 src/redis-shake/common/utils.go | 21 ++++++++++++++++-----
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/pkg/rdb/loader.go b/src/pkg/rdb/loader.go
index 7ca4faf..c7ec8d4 100644
--- a/src/pkg/rdb/loader.go
+++ b/src/pkg/rdb/loader.go
@@ -63,6 +63,8 @@ type BinEntry struct {
 	ExpireAt        uint64
 	RealMemberCount uint32
 	NeedReadLen     byte
+	IdleTime        uint32
+	Freq            uint8
 }
 
 func (e *BinEntry) ObjEntry() (*ObjEntry, error) {
@@ -169,16 +171,18 @@ func (l *Loader) NextBinEntry() (*BinEntry, error) {
 			}
 		case rdbFlagIdle:
 			// ignore idle because target redis doesn't support this for given key
-			_, err := l.ReadLength()
+			idle, err := l.ReadLength()
 			if err != nil {
 				return nil, err
 			}
+			entry.IdleTime = idle
 		case rdbFlagFreq:
 			// ignore freq because target redis doesn't support this for given key
-			_, err := l.readUint8()
+			freq, err := l.readUint8()
 			if err != nil {
 				return nil, err
 			}
+			entry.Freq = freq
 		default:
 			var key []byte
 			if l.remainMember == 0 {
diff --git a/src/redis-shake/common/utils.go b/src/redis-shake/common/utils.go
index f8ef081..e4b30df 100644
--- a/src/redis-shake/common/utils.go
+++ b/src/redis-shake/common/utils.go
@@ -711,12 +711,22 @@ func RestoreRdbEntry(c redigo.Conn, e *rdb.BinEntry) {
 		return
 	}
 
-	// fmt.Printf("kkkey: %v, value: %v\n", string(e.Key), e.Value)
-	s, err := redigo.String(c.Do("restore", e.Key, ttlms, e.Value))
+	params := []interface{}{e.Key, ttlms, e.Value}
+	if e.IdleTime != 0 {
+		params = append(params, "IDLETIME")
+		params = append(params, e.IdleTime)
+	}
+	if e.Freq != 0 {
+		params = append(params, "FREQ")
+		params = append(params, e.Freq)
+	}
+	// fmt.Printf("key: %v, value: %v params: %v\n", string(e.Key), e.Value, params)
+	s, err := redigo.String(c.Do("restore", params...))
 	if err != nil {
 		/*The reply value of busykey in 2.8 kernel is "target key name is busy",
 		  but in 4.0 kernel is "BUSYKEY Target key name already exists"*/
-		if strings.Contains(err.Error(), "Target key name is busy") || strings.Contains(err.Error(), "BUSYKEY Target key name already exists") {
+		if strings.Contains(err.Error(), "Target key name is busy") ||
+			strings.Contains(err.Error(), "BUSYKEY Target key name already exists") {
 			if conf.Options.Rewrite {
 				if !conf.Options.Metric {
 					log.Infof("warning, rewrite key: %v", string(e.Key))
@@ -724,10 +734,11 @@ func RestoreRdbEntry(c redigo.Conn, e *rdb.BinEntry) {
 				var s2 string
 				var rerr error
 				if conf.Options.TargetReplace {
-					s2, rerr = redigo.String(c.Do("restore", e.Key, ttlms, e.Value, "replace"))
+					params = append(params, "REPLACE")
+					s2, rerr = redigo.String(c.Do("restore", params...))
 				} else {
 					_, _ = redigo.String(c.Do("del", e.Key))
-					s2, rerr = redigo.String(c.Do("restore", e.Key, ttlms, e.Value))
+					s2, rerr = redigo.String(c.Do("restore", params...))
 				}
 				if rerr != nil {
 					log.Info(s2, rerr, "key ", string(e.Key))