Merge pull request #83 from ahmczsy/master

add insert limiter
zeromicro · Jun 20, 2021 · d429399 · d429399
2 parents 7ac386d + 6e321e9
commit d429399
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 17 deletions.
diff --git a/go.mod b/go.mod
@@ -36,6 +36,7 @@ require (
 	golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9
 	golang.org/x/sys v0.0.0-20201119102817-f84b799fce68 // indirect
 	golang.org/x/text v0.3.4 // indirect
+	golang.org/x/time v0.0.0-20191024005414-555d28b269f0
 	golang.org/x/tools v0.0.0-20201125231158-b5590deeca9b // indirect
 	gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b // indirect
 	gopkg.in/mgo.v2 v2.0.0-20190816093944-a6b53ec6cb22

diff --git a/pkg/ckgroup/exec.go b/pkg/ckgroup/exec.go
@@ -2,6 +2,7 @@ package ckgroup
 
 import (
 	"errors"
+	"math/rand"
 	"sort"
 	"strings"
 	"sync"
@@ -154,25 +155,21 @@ func (g *dbGroup) ExecAll(query string, args [][]interface{}) error {
 }
 
 func (g *dbGroup) exec(idx int, query string, rows []rowValue) error {
+	shardConns := g.GetAllShard()[idx].GetAllConn()
+	execOrder := rand.Perm(len(shardConns))
 	var err error
 	for attempt := 1; attempt <= g.opt.RetryNum; attempt++ {
-		err = saveData(g.ShardNodes[idx].GetShardConn().GetRawConn(), query, rows)
-		if err != nil {
-			logx.Infof("[attempt %d/%d] Node[%d] primary node execute error:%v, will switch to replica node", attempt, g.opt.RetryNum, idx, err)
-		} else {
-			return nil
-		}
-		for i, replicaNode := range g.ShardNodes[idx].GetReplicaConn() {
-			err = saveData(replicaNode.GetRawConn(), query, rows)
-			if err != nil {
-				logx.Infof("[attempt %d/%d] Node[%d] replica[%d] execute error:%v, will switch to next replica node", attempt, g.opt.RetryNum, idx, i, err)
-			} else {
+		for _, order := range execOrder {
+			err = saveData(shardConns[order].GetRawConn(), query, rows)
+			if err == nil {
 				return nil
 			}
+			logx.Errorf("[attempt %d/%d] shard[%d] node[%d] insert error:%s, will switch to next node", attempt, g.opt.RetryNum, idx+1, order+1, err.Error())
+			continue
 		}
 	}
 	if err != nil {
-		logx.Errorf("All node exec failed. Retry num:%d. Last fail reason: %v, query: %s", g.opt.RetryNum, err, query)
+		logx.Errorf("shard[%d] all node exec failed. Retry num:%d. Last fail reason: %v, query: %s", idx+1, g.opt.RetryNum, err, query)
 	}
 	return err
 }
diff --git a/pkg/ckgroup/insert.go b/pkg/ckgroup/insert.go
@@ -1,6 +1,7 @@
 package ckgroup
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"reflect"
@@ -20,6 +21,11 @@ func (g *dbGroup) InsertAuto(query string, hashTag string, sliceData interface{}
 		return err
 	}
 
+	err = g.opt.GroupInsertLimiter.Wait(context.Background())
+	if err != nil {
+		logx.Error(err)
+	}
+
 	var eg errgroup.Group
 	for i, shardConn := range g.ShardNodes {
 
@@ -54,6 +60,12 @@ func (g *dbGroup) InsertAutoDetail(query string, hashTag string, sliceData inter
 	if err != nil {
 		return nil, err
 	}
+
+	err = g.opt.GroupInsertLimiter.Wait(context.Background())
+	if err != nil {
+		logx.Error(err)
+	}
+
 	waitGroup := sync.WaitGroup{}
 	waitGroup.Add(len(g.ShardNodes))
 	ch := make(chan InsertErrDetail, len(g.GetAllShard()))

diff --git a/pkg/ckgroup/insert_err_detail_test.go b/pkg/ckgroup/insert_err_detail_test.go
@@ -6,14 +6,15 @@ import (
 	"testing"
 
 	"github.com/tal-tech/cds/pkg/ckgroup/dbtesttool/dbtool"
+	"golang.org/x/time/rate"
 )
 
 func Test_dbGroup_InsertAutoDetail(t *testing.T) {
 	dataSet := dbtool.GenerateDataSet(10000)
 
 	c1 := dbGroup{
 		ShardNodes: []ShardConn{&fakeShardConn{true}, &fakeShardConn{true}, &fakeShardConn{true}},
-		opt:        option{RetryNum: 3},
+		opt:        option{RetryNum: 3, GroupInsertLimiter: rate.NewLimiter(rate.Inf, 0)},
 	}
 	errDetail1, err := c1.InsertAutoDetail(``, "pk", dataSet)
 	if err != nil {
@@ -29,7 +30,7 @@ func Test_dbGroup_InsertAutoDetail(t *testing.T) {
 	}
 	c2 := dbGroup{
 		ShardNodes: []ShardConn{&fakeShardConn{false}, &fakeShardConn{false}, &fakeShardConn{false}},
-		opt:        option{RetryNum: 3},
+		opt:        option{RetryNum: 3, GroupInsertLimiter: rate.NewLimiter(rate.Inf, 0)},
 	}
 	errDetail2, err := c2.InsertAutoDetail(``, "pk", dataSet)
 	if err != nil {
@@ -41,7 +42,7 @@ func Test_dbGroup_InsertAutoDetail(t *testing.T) {
 
 	c3 := dbGroup{
 		ShardNodes: []ShardConn{&fakeShardConn{false}, &fakeShardConn{true}, &fakeShardConn{false}},
-		opt:        option{RetryNum: 3},
+		opt:        option{RetryNum: 3, GroupInsertLimiter: rate.NewLimiter(rate.Inf, 0)},
 	}
 	errDetail3, err := c3.InsertAutoDetail(``, "pk", dataSet)
 	if err != nil {

diff --git a/pkg/ckgroup/option.go b/pkg/ckgroup/option.go
@@ -1,13 +1,19 @@
 package ckgroup
 
+import (
+	"golang.org/x/time/rate"
+)
+
 type option struct {
-	RetryNum int
+	RetryNum           int
+	GroupInsertLimiter *rate.Limiter
 }
 type OptionFunc func(*option)
 
 func newOptions(opts ...OptionFunc) option {
 	opt := option{
-		RetryNum: 1,
+		RetryNum:           1,
+		GroupInsertLimiter: rate.NewLimiter(rate.Inf, 0),
 	}
 
 	for _, o := range opts {
@@ -22,3 +28,9 @@ func WithRetryNum(retryNum int) OptionFunc {
 		o.RetryNum = retryNum
 	}
 }
+
+func WithGroupInsertLimiter(limit rate.Limit, burst int) OptionFunc {
+	return func(o *option) {
+		o.GroupInsertLimiter = rate.NewLimiter(limit, burst)
+	}
+}
diff --git a/pkg/ckgroup/option_test.go b/pkg/ckgroup/option_test.go
@@ -0,0 +1,28 @@
+package ckgroup
+
+import (
+	"testing"
+	"time"
+
+	"github.com/tal-tech/cds/pkg/ckgroup/dbtesttool/dbtool"
+	"golang.org/x/time/rate"
+)
+
+func TestWithGroupInsertLimiter(t *testing.T) {
+	db := dbGroup{
+		ShardNodes: []ShardConn{&fakeShardConn{true}},
+		opt:        option{GroupInsertLimiter: rate.NewLimiter(rate.Every(time.Millisecond*500), 1)},
+	}
+	dataSet := dbtool.GenerateDataSet(10000)
+	start := time.Now()
+	for i := 1; i <= 3; i++ {
+		err := db.InsertAuto(``, `pk`, dataSet)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	if time.Since(start) <= time.Millisecond*900 {
+		t.Fatal(`insert limiter did not meet expectations`)
+	}
+}