添加scan时候的一次扫描次数控制配置，防止源端cpu高

tair-opensource · Mar 20, 2024 · f286da4 · f286da4
1 parent 36915f0
commit f286da4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/internal/client/redis.go b/internal/client/redis.go
@@ -148,8 +148,8 @@ func (r *Redis) Close() {
 
 /* Commands */
 
-func (r *Redis) Scan(cursor uint64) (newCursor uint64, keys []string) {
-	r.Send("scan", strconv.FormatUint(cursor, 10), "count", "2048")
+func (r *Redis) Scan(cursor uint64, batch string) (newCursor uint64, keys []string) {
+	r.Send("scan", strconv.FormatUint(cursor, 10), "count", batch)
 	reply, err := r.Receive()
 	if err != nil {
 		log.Panicf(err.Error())

diff --git a/internal/reader/scan_standalone_reader.go b/internal/reader/scan_standalone_reader.go
@@ -26,6 +26,7 @@ type ScanReaderOptions struct {
 	KSN           bool   `mapstructure:"ksn" default:"false"`
 	DBS           []int  `mapstructure:"dbs"`
 	PreferReplica bool   `mapstructure:"prefer_replica" default:"false"`
+	Batch         int    `mapstructure:"batch" default:"2048"`
 }
 
 type dbKey struct {
@@ -131,9 +132,10 @@ func (r *scanStandaloneReader) scan() {
 		}
 
 		var cursor uint64 = 0
+		batch := strconv.Itoa(r.opts.Batch)
 		for {
 			var keys []string
-			cursor, keys = c.Scan(cursor)
+			cursor, keys = c.Scan(cursor, batch)
 			for _, key := range keys {
 				r.keyQueue.Put(dbKey{dbId, key}) // pass value not pointer
 			}