importinto/lightning: do remote checksum via sql (#44803)

br/pkg/checksum/executor.go

@@ -28,7 +28,8 @@ type ExecutorBuilder struct {
 
 	oldTable *metautil.Table
 
-	concurrency uint
+	concurrency   uint
+	backoffWeight int
 
 	oldKeyspace []byte
 	newKeyspace []byte
@@ -56,6 +57,12 @@ func (builder *ExecutorBuilder) SetConcurrency(conc uint) *ExecutorBuilder {
 	return builder
 }
 
+// SetBackoffWeight set the backoffWeight of the checksum executing.
+func (builder *ExecutorBuilder) SetBackoffWeight(backoffWeight int) *ExecutorBuilder {
+	builder.backoffWeight = backoffWeight
+	return builder
+}
+
 func (builder *ExecutorBuilder) SetOldKeyspace(keyspace []byte) *ExecutorBuilder {
 	builder.oldKeyspace = keyspace
 	return builder
@@ -79,7 +86,7 @@ func (builder *ExecutorBuilder) Build() (*Executor, error) {
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	return &Executor{reqs: reqs}, nil
+	return &Executor{reqs: reqs, backoffWeight: builder.backoffWeight}, nil
 }
 
 func buildChecksumRequest(
@@ -294,7 +301,8 @@ func updateChecksumResponse(resp, update *tipb.ChecksumResponse) {
 
 // Executor is a checksum executor.
 type Executor struct {
-	reqs []*kv.Request
+	reqs          []*kv.Request
+	backoffWeight int
 }
 
 // Len returns the total number of checksum requests.
@@ -347,7 +355,11 @@ func (exec *Executor) Execute(
 			err  error
 		)
 		err = utils.WithRetry(ctx, func() error {
-			resp, err = sendChecksumRequest(ctx, client, req, kv.NewVariables(&killed))
+			vars := kv.NewVariables(&killed)
+			if exec.backoffWeight > 0 {
+				vars.BackOffWeight = exec.backoffWeight
+			}
+			resp, err = sendChecksumRequest(ctx, client, req, vars)
 			failpoint.Inject("checksumRetryErr", func(val failpoint.Value) {
 				// first time reach here. return error
 				if val.(bool) {

br/pkg/lightning/backend/local/BUILD.bazel

@@ -48,6 +48,7 @@ go_library(
         "//kv",
         "//parser/model",
         "//parser/mysql",
+        "//sessionctx/variable",
         "//store/pdtypes",
         "//table",
         "//tablecodec",
@@ -73,6 +74,7 @@ go_library(
         "@com_github_pingcap_kvproto//pkg/pdpb",
         "@com_github_pingcap_tipb//go-tipb",
         "@com_github_tikv_client_go_v2//error",
+        "@com_github_tikv_client_go_v2//kv",
         "@com_github_tikv_client_go_v2//oracle",
         "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_pd_client//:client",

br/pkg/lightning/backend/local/checksum.go

@@ -32,8 +32,10 @@ import (
 	"github.com/pingcap/tidb/br/pkg/lightning/metric"
 	"github.com/pingcap/tidb/br/pkg/lightning/verification"
 	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/util/mathutil"
 	"github.com/pingcap/tipb/go-tipb"
+	tikvstore "github.com/tikv/client-go/v2/kv"
 	"github.com/tikv/client-go/v2/oracle"
 	pd "github.com/tikv/pd/client"
 	"go.uber.org/atomic"
@@ -49,7 +51,14 @@ const (
 var (
 	serviceSafePointTTL int64 = 10 * 60 // 10 min in seconds
 
-	minDistSQLScanConcurrency = 4
+	// MinDistSQLScanConcurrency is the minimum value of tidb_distsql_scan_concurrency.
+	MinDistSQLScanConcurrency = 4
+
+	// DefaultBackoffWeight is the default value of tidb_backoff_weight for checksum.
+	// when TiKV client encounters an error of "region not leader", it will keep retrying every 500 ms.
+	// If it still fails after 2 * 20 = 40 seconds, it will return "region unavailable".
+	// If we increase the BackOffWeight to 6, then the TiKV client will keep retrying for 120 seconds.
+	DefaultBackoffWeight = 3 * tikvstore.DefBackOffWeight
 )
 
 // RemoteChecksum represents a checksum result got from tidb.
@@ -102,16 +111,39 @@ func (e *tidbChecksumExecutor) Checksum(ctx context.Context, tableInfo *checkpoi
 
 	task := log.FromContext(ctx).With(zap.String("table", tableName)).Begin(zap.InfoLevel, "remote checksum")
 
+	conn, err := e.db.Conn(ctx)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	defer func() {
+		if err := conn.Close(); err != nil {
+			task.Warn("close connection failed", zap.Error(err))
+		}
+	}()
 	// ADMIN CHECKSUM TABLE <table>,<table>  example.
 	// 	mysql> admin checksum table test.t;
 	// +---------+------------+---------------------+-----------+-------------+
 	// | Db_name | Table_name | Checksum_crc64_xor  | Total_kvs | Total_bytes |
 	// +---------+------------+---------------------+-----------+-------------+
 	// | test    | t          | 8520875019404689597 |   7296873 |   357601387 |
 	// +---------+------------+---------------------+-----------+-------------+
+	backoffWeight, err := common.GetBackoffWeightFromDB(ctx, e.db)
+	if err == nil && backoffWeight < DefaultBackoffWeight {
+		task.Info("increase tidb_backoff_weight", zap.Int("original", backoffWeight), zap.Int("new", DefaultBackoffWeight))
+		// increase backoff weight
+		if _, err := conn.ExecContext(ctx, fmt.Sprintf("SET SESSION %s = '%d';", variable.TiDBBackOffWeight, DefaultBackoffWeight)); err != nil {
+			task.Warn("set tidb_backoff_weight failed", zap.Error(err))
+		} else {
+			defer func() {
+				if _, err := conn.ExecContext(ctx, fmt.Sprintf("SET SESSION %s = '%d';", variable.TiDBBackOffWeight, backoffWeight)); err != nil {
+					task.Warn("recover tidb_backoff_weight failed", zap.Error(err))
+				}
+			}()
+		}
+	}
 
 	cs := RemoteChecksum{}
-	err = common.SQLWithRetry{DB: e.db, Logger: task.Logger}.QueryRow(ctx, "compute remote checksum",
+	err = common.SQLWithRetry{DB: conn, Logger: task.Logger}.QueryRow(ctx, "compute remote checksum",
 		"ADMIN CHECKSUM TABLE "+tableName, &cs.Schema, &cs.Table, &cs.Checksum, &cs.TotalKVs, &cs.TotalBytes,
 	)
 	dur := task.End(zap.ErrorLevel, err)
@@ -239,22 +271,25 @@ type TiKVChecksumManager struct {
 	client                 kv.Client
 	manager                gcTTLManager
 	distSQLScanConcurrency uint
+	backoffWeight          int
 }
 
 var _ ChecksumManager = &TiKVChecksumManager{}
 
 // NewTiKVChecksumManager return a new tikv checksum manager
-func NewTiKVChecksumManager(client kv.Client, pdClient pd.Client, distSQLScanConcurrency uint) *TiKVChecksumManager {
+func NewTiKVChecksumManager(client kv.Client, pdClient pd.Client, distSQLScanConcurrency uint, backoffWeight int) *TiKVChecksumManager {
 	return &TiKVChecksumManager{
 		client:                 client,
 		manager:                newGCTTLManager(pdClient),
 		distSQLScanConcurrency: distSQLScanConcurrency,
+		backoffWeight:          backoffWeight,
 	}
 }
 
 func (e *TiKVChecksumManager) checksumDB(ctx context.Context, tableInfo *checkpoints.TidbTableInfo, ts uint64) (*RemoteChecksum, error) {
 	executor, err := checksum.NewExecutorBuilder(tableInfo.Core, ts).
 		SetConcurrency(e.distSQLScanConcurrency).
+		SetBackoffWeight(e.backoffWeight).
 		Build()
 	if err != nil {
 		return nil, errors.Trace(err)
@@ -286,8 +321,8 @@ func (e *TiKVChecksumManager) checksumDB(ctx context.Context, tableInfo *checkpo
 		if !common.IsRetryableError(err) {
 			break
 		}
-		if distSQLScanConcurrency > minDistSQLScanConcurrency {
-			distSQLScanConcurrency = mathutil.Max(distSQLScanConcurrency/2, minDistSQLScanConcurrency)
+		if distSQLScanConcurrency > MinDistSQLScanConcurrency {
+			distSQLScanConcurrency = mathutil.Max(distSQLScanConcurrency/2, MinDistSQLScanConcurrency)
 		}
 	}
 

br/pkg/lightning/backend/local/checksum_test.go

@@ -49,6 +49,7 @@ func TestDoChecksum(t *testing.T) {
 		WithArgs("10m").
 		WillReturnResult(sqlmock.NewResult(2, 1))
 	mock.ExpectClose()
+	mock.ExpectClose()
 
 	manager := NewTiDBChecksumExecutor(db)
 	checksum, err := manager.Checksum(context.Background(), &TidbTableInfo{DB: "test", Name: "t"})
@@ -228,6 +229,7 @@ func TestDoChecksumWithErrorAndLongOriginalLifetime(t *testing.T) {
 		WithArgs("300h").
 		WillReturnResult(sqlmock.NewResult(1, 1))
 	mock.ExpectClose()
+	mock.ExpectClose()
 
 	manager := NewTiDBChecksumExecutor(db)
 	_, err = manager.Checksum(context.Background(), &TidbTableInfo{DB: "test", Name: "t"})

br/pkg/lightning/common/BUILD.bazel

@@ -27,6 +27,7 @@ go_library(
         "//meta/autoid",
         "//parser/model",
         "//parser/mysql",
+        "//sessionctx/variable",
         "//store/driver/error",
         "//table/tables",
         "//types",

br/pkg/lightning/common/common.go

@@ -39,6 +39,7 @@ var DefaultImportantVariables = map[string]string{
 	"default_week_format":     "0",
 	"block_encryption_mode":   "aes-128-ecb",
 	"group_concat_max_len":    "1024",
+	"tidb_backoff_weight":     "6",
 }
 
 // DefaultImportVariablesTiDB is used in ObtainImportantVariables to retrieve the system

br/pkg/lightning/common/util.go

@@ -39,6 +39,7 @@ import (
 	"github.com/pingcap/tidb/errno"
 	"github.com/pingcap/tidb/parser/model"
 	tmysql "github.com/pingcap/tidb/parser/mysql"
+	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/table/tables"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/codec"
@@ -605,3 +606,45 @@ func IsDupKeyError(err error) bool {
 	}
 	return false
 }
+
+// GetBackoffWeightFromDB gets the backoff weight from database.
+func GetBackoffWeightFromDB(ctx context.Context, db *sql.DB) (int, error) {
+	val, err := getSessionVariable(ctx, db, variable.TiDBBackOffWeight)
+	if err != nil {
+		return 0, err
+	}
+	return strconv.Atoi(val)
+}
+
+// copy from dbutil to avoid import cycle
+func getSessionVariable(ctx context.Context, db *sql.DB, variable string) (value string, err error) {
+	query := fmt.Sprintf("SHOW VARIABLES LIKE '%s'", variable)
+	rows, err := db.QueryContext(ctx, query)
+
+	if err != nil {
+		return "", errors.Trace(err)
+	}
+	defer rows.Close()
+
+	// Show an example.
+	/*
+		mysql> SHOW VARIABLES LIKE "binlog_format";
+		+---------------+-------+
+		| Variable_name | Value |
+		+---------------+-------+
+		| binlog_format | ROW   |
+		+---------------+-------+
+	*/
+
+	for rows.Next() {
+		if err = rows.Scan(&variable, &value); err != nil {
+			return "", errors.Trace(err)
+		}
+	}
+
+	if err := rows.Err(); err != nil {
+		return "", errors.Trace(err)
+	}
+
+	return value, nil
+}

br/pkg/lightning/config/config.go

@@ -599,6 +599,7 @@ type PostRestore struct {
 	Level1Compact     bool        `toml:"level-1-compact" json:"level-1-compact"`
 	PostProcessAtLast bool        `toml:"post-process-at-last" json:"post-process-at-last"`
 	Compact           bool        `toml:"compact" json:"compact"`
+	ChecksumViaSQL    bool        `toml:"checksum-via-sql" json:"checksum-via-sql"`
 }
 
 // StringOrStringSlice can unmarshal a TOML string as string slice with one element.
@@ -974,6 +975,7 @@ func NewConfig() *Config {
 			Checksum:          OpLevelRequired,
 			Analyze:           OpLevelOptional,
 			PostProcessAtLast: true,
+			ChecksumViaSQL:    true,
 		},
 	}
 }

br/pkg/lightning/importer/checksum_helper.go

@@ -20,6 +20,7 @@ import (
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/br/pkg/lightning/backend/local"
 	"github.com/pingcap/tidb/br/pkg/lightning/checkpoints"
+	"github.com/pingcap/tidb/br/pkg/lightning/common"
 	"github.com/pingcap/tidb/br/pkg/lightning/config"
 	"github.com/pingcap/tidb/br/pkg/lightning/log"
 	"github.com/pingcap/tidb/br/pkg/lightning/metric"
@@ -44,14 +45,22 @@ func NewChecksumManager(ctx context.Context, rc *Controller, store kv.Storage) (
 
 	// for v4.0.0 or upper, we can use the gc ttl api
 	var manager local.ChecksumManager
-	if pdVersion.Major >= 4 {
+	if pdVersion.Major >= 4 && !rc.cfg.PostRestore.ChecksumViaSQL {
 		tlsOpt := rc.tls.ToPDSecurityOption()
 		pdCli, err := pd.NewClientWithContext(ctx, []string{pdAddr}, tlsOpt)
 		if err != nil {
 			return nil, errors.Trace(err)
 		}
 
-		manager = local.NewTiKVChecksumManager(store.GetClient(), pdCli, uint(rc.cfg.TiDB.DistSQLScanConcurrency))
+		backoffWeight, err := common.GetBackoffWeightFromDB(ctx, rc.db)
+		// only set backoff weight when it's smaller than default value
+		if err == nil && backoffWeight >= local.DefaultBackoffWeight {
+			log.FromContext(ctx).Info("get tidb_backoff_weight", zap.Int("backoff_weight", backoffWeight))
+		} else {
+			log.FromContext(ctx).Info("set tidb_backoff_weight to default", zap.Int("backoff_weight", local.DefaultBackoffWeight))
+			backoffWeight = local.DefaultBackoffWeight
+		}
+		manager = local.NewTiKVChecksumManager(store.GetClient(), pdCli, uint(rc.cfg.TiDB.DistSQLScanConcurrency), backoffWeight)
 	} else {
 		manager = local.NewTiDBChecksumExecutor(rc.db)
 	}

br/pkg/lightning/importer/table_import_test.go

@@ -802,6 +802,7 @@ func (s *tableRestoreSuite) TestCompareChecksumSuccess() {
 		WithArgs("10m").
 		WillReturnResult(sqlmock.NewResult(2, 1))
 	mock.ExpectClose()
+	mock.ExpectClose()
 
 	ctx := MockDoChecksumCtx(db)
 	remoteChecksum, err := DoChecksum(ctx, s.tr.tableInfo)
@@ -832,7 +833,7 @@ func (s *tableRestoreSuite) TestCompareChecksumFailure() {
 		WithArgs("10m").
 		WillReturnResult(sqlmock.NewResult(2, 1))
 	mock.ExpectClose()
-
+	mock.ExpectClose()
 	ctx := MockDoChecksumCtx(db)
 	remoteChecksum, err := DoChecksum(ctx, s.tr.tableInfo)
 	require.NoError(s.T(), err)

br/pkg/lightning/importer/tidb_test.go

@@ -337,6 +337,7 @@ func TestObtainRowFormatVersionSucceed(t *testing.T) {
 
 	sysVars := ObtainImportantVariables(ctx, s.db, true)
 	require.Equal(t, map[string]string{
+		"tidb_backoff_weight":     "6",
 		"tidb_row_format_version": "2",
 		"max_allowed_packet":      "1073741824",
 		"div_precision_increment": "10",
@@ -360,6 +361,7 @@ func TestObtainRowFormatVersionFailure(t *testing.T) {
 
 	sysVars := ObtainImportantVariables(ctx, s.db, true)
 	require.Equal(t, map[string]string{
+		"tidb_backoff_weight":     "6",
 		"tidb_row_format_version": "1",
 		"max_allowed_packet":      "67108864",
 		"div_precision_increment": "4",

br/tests/lightning_add_index/config1.toml

@@ -1,3 +1,6 @@
 [tikv-importer]
 backend = 'local'
 add-index-by-sql = false
+
+[post-restore]
+checksum-via-sql = false

disttask/framework/dispatcher/dispatcher.go

@@ -502,6 +502,10 @@ func (d *dispatcher) WithNewSession(fn func(se sessionctx.Context) error) error
 	return d.taskMgr.WithNewSession(fn)
 }
 
+func (d *dispatcher) WithNewTxn(ctx context.Context, fn func(se sessionctx.Context) error) error {
+	return d.taskMgr.WithNewTxn(ctx, fn)
+}
+
 func (*dispatcher) checkConcurrencyOverflow(cnt int) bool {
 	if cnt >= DefaultDispatchConcurrency {
 		logutil.BgLogger().Info("dispatch task loop, running GTask cnt is more than concurrency",