lightning: mark background and resource group name when in local mode

DEPS.bzl

@@ -1711,13 +1711,13 @@ def go_deps():
         name = "com_github_elastic_gosigar",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/elastic/gosigar",
-        sha256 = "ba59c159fd8b4c64f34dd0c915df96d39f7f5d5e21f9ae2d9639793828a67e18",
-        strip_prefix = "github.com/elastic/gosigar@v0.14.2",
+        sha256 = "21ddafe5b9912113983e4210aa48f12a87d3a88a0a6946601e65725b736ca852",
+        strip_prefix = "github.com/elastic/gosigar@v0.9.0",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/elastic/gosigar/com_github_elastic_gosigar-v0.14.2.zip",
-            "http://ats.apps.svc/gomod/github.com/elastic/gosigar/com_github_elastic_gosigar-v0.14.2.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/elastic/gosigar/com_github_elastic_gosigar-v0.14.2.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/elastic/gosigar/com_github_elastic_gosigar-v0.14.2.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/elastic/gosigar/com_github_elastic_gosigar-v0.9.0.zip",
+            "http://ats.apps.svc/gomod/github.com/elastic/gosigar/com_github_elastic_gosigar-v0.9.0.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/elastic/gosigar/com_github_elastic_gosigar-v0.9.0.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/elastic/gosigar/com_github_elastic_gosigar-v0.9.0.zip",
         ],
     )
     go_repository(
@@ -6924,26 +6924,26 @@ def go_deps():
         name = "com_github_tikv_client_go_v2",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/tikv/client-go/v2",
-        sha256 = "fda64379a1d2f8bc263aa1e9faac235a59c085bc06269a46a88f3dd60e37d372",
-        strip_prefix = "github.com/tikv/client-go/v2@v2.0.8-0.20230714052714-85fc8f337565",
+        sha256 = "2dcd0fa5258d9cf2dedf7a5d6066c9c221707c269f895824b05a5be4af721263",
+        strip_prefix = "github.com/tikv/client-go/v2@v2.0.8-0.20230726071354-76d6d93e9856",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20230714052714-85fc8f337565.zip",
-            "http://ats.apps.svc/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20230714052714-85fc8f337565.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20230714052714-85fc8f337565.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20230714052714-85fc8f337565.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20230726071354-76d6d93e9856.zip",
+            "http://ats.apps.svc/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20230726071354-76d6d93e9856.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20230726071354-76d6d93e9856.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20230726071354-76d6d93e9856.zip",
         ],
     )
     go_repository(
         name = "com_github_tikv_pd_client",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/tikv/pd/client",
-        sha256 = "6cdab562dd442f33d5c4fca3999c8e7fafca1ae28a55cbd597ede42a9adef249",
-        strip_prefix = "github.com/tikv/pd/client@v0.0.0-20230613052906-7158cb319935",
+        sha256 = "401dddb69309b07fa9d4464adeaf5d98a456dda622114fd479d01d18ab622cbc",
+        strip_prefix = "github.com/tikv/pd/client@v0.0.0-20230724080549-de985b8e0afc",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20230613052906-7158cb319935.zip",
-            "http://ats.apps.svc/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20230613052906-7158cb319935.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20230613052906-7158cb319935.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20230613052906-7158cb319935.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20230724080549-de985b8e0afc.zip",
+            "http://ats.apps.svc/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20230724080549-de985b8e0afc.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20230724080549-de985b8e0afc.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/tikv/pd/client/com_github_tikv_pd_client-v0.0.0-20230724080549-de985b8e0afc.zip",
         ],
     )
     go_repository(
@@ -10035,13 +10035,13 @@ def go_deps():
         name = "org_golang_x_exp",
         build_file_proto_mode = "disable_global",
         importpath = "golang.org/x/exp",
-        sha256 = "1bac2a4c8fb2891f3db33d4beeb04b4fcc3b1e00733f67a58c3d5942e21b4ff2",
-        strip_prefix = "golang.org/x/exp@v0.0.0-20230519143937-03e91628a987",
+        sha256 = "d4b6a3cc6bf072a05030324328169b5f878f7be012508fff618e251cccb3d0aa",
+        strip_prefix = "golang.org/x/exp@v0.0.0-20230711005742-c3f37128e5a4",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/golang.org/x/exp/org_golang_x_exp-v0.0.0-20230519143937-03e91628a987.zip",
-            "http://ats.apps.svc/gomod/golang.org/x/exp/org_golang_x_exp-v0.0.0-20230519143937-03e91628a987.zip",
-            "https://cache.hawkingrei.com/gomod/golang.org/x/exp/org_golang_x_exp-v0.0.0-20230519143937-03e91628a987.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/golang.org/x/exp/org_golang_x_exp-v0.0.0-20230519143937-03e91628a987.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/golang.org/x/exp/org_golang_x_exp-v0.0.0-20230711005742-c3f37128e5a4.zip",
+            "http://ats.apps.svc/gomod/golang.org/x/exp/org_golang_x_exp-v0.0.0-20230711005742-c3f37128e5a4.zip",
+            "https://cache.hawkingrei.com/gomod/golang.org/x/exp/org_golang_x_exp-v0.0.0-20230711005742-c3f37128e5a4.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/golang.org/x/exp/org_golang_x_exp-v0.0.0-20230711005742-c3f37128e5a4.zip",
         ],
     )
     go_repository(

br/cmd/tidb-lightning-ctl/main.go

@@ -128,7 +128,7 @@
 		tls.WithHost(cfg.TiDB.PdAddr),
 		tikv.StoreStateDisconnected,
 		func(c context.Context, store *tikv.Store) error {
-			return tikv.Compact(c, tls, store.Address, importer.FullLevelCompact)
+			return tikv.Compact(c, tls, store.Address, importer.FullLevelCompact, "")
 		},
 	)
 }

br/pkg/checksum/executor.go

@@ -33,6 +33,9 @@
 
 	oldKeyspace []byte
 	newKeyspace []byte
+
+	resourceGroupName         string
+	explicitRequestSourceType string
 }
 
 // NewExecutorBuilder returns a new executor builder.
@@ -73,6 +76,16 @@
 	return builder
 }
 
+func (builder *ExecutorBuilder) SetResourceGroupName(name string) *ExecutorBuilder {
+	builder.resourceGroupName = name
+	return builder
+}
+
+func (builder *ExecutorBuilder) SetExplicitRequestSourceType(name string) *ExecutorBuilder {
+	builder.explicitRequestSourceType = name
+	return builder
+}
+
 // Build builds a checksum executor.
 func (builder *ExecutorBuilder) Build() (*Executor, error) {
 	reqs, err := buildChecksumRequest(
@@ -82,6 +95,8 @@
 		builder.concurrency,
 		builder.oldKeyspace,
 		builder.newKeyspace,
+		builder.resourceGroupName,
+		builder.explicitRequestSourceType,
 	)
 	if err != nil {
 		return nil, errors.Trace(err)
@@ -96,6 +111,7 @@
 	concurrency uint,
 	oldKeyspace []byte,
 	newKeyspace []byte,
+	resourceGroupName, explicitRequestSourceType string,
 ) ([]*kv.Request, error) {
 	var partDefs []model.PartitionDefinition
 	if part := newTable.Partition; part != nil {
@@ -107,7 +123,8 @@
 	if oldTable != nil {
 		oldTableID = oldTable.Info.ID
 	}
-	rs, err := buildRequest(newTable, newTable.ID, oldTable, oldTableID, startTS, concurrency, oldKeyspace, newKeyspace)
+	rs, err := buildRequest(newTable, newTable.ID, oldTable, oldTableID, startTS, concurrency,
+		oldKeyspace, newKeyspace, resourceGroupName, explicitRequestSourceType)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -122,7 +139,8 @@
 				}
 			}
 		}
-		rs, err := buildRequest(newTable, partDef.ID, oldTable, oldPartID, startTS, concurrency, oldKeyspace, newKeyspace)
+		rs, err := buildRequest(newTable, partDef.ID, oldTable, oldPartID, startTS, concurrency,
+			oldKeyspace, newKeyspace, resourceGroupName, explicitRequestSourceType)
 		if err != nil {
 			return nil, errors.Trace(err)
 		}
@@ -141,9 +159,11 @@
 	concurrency uint,
 	oldKeyspace []byte,
 	newKeyspace []byte,
+	resourceGroupName, explicitRequestSourceType string,
 ) ([]*kv.Request, error) {
 	reqs := make([]*kv.Request, 0)
-	req, err := buildTableRequest(tableInfo, tableID, oldTable, oldTableID, startTS, concurrency, oldKeyspace, newKeyspace)
+	req, err := buildTableRequest(tableInfo, tableID, oldTable, oldTableID, startTS, concurrency,
+		oldKeyspace, newKeyspace, resourceGroupName, explicitRequestSourceType)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -172,7 +192,8 @@
 			}
 		}
 		req, err = buildIndexRequest(
-			tableID, indexInfo, oldTableID, oldIndexInfo, startTS, concurrency, oldKeyspace, newKeyspace)
+			tableID, indexInfo, oldTableID, oldIndexInfo, startTS, concurrency,
+			oldKeyspace, newKeyspace, resourceGroupName, explicitRequestSourceType)
 		if err != nil {
 			return nil, errors.Trace(err)
 		}
@@ -191,6 +212,7 @@
 	concurrency uint,
 	oldKeyspace []byte,
 	newKeyspace []byte,
+	resourceGroupName, explicitRequestSourceType string,
 ) (*kv.Request, error) {
 	var rule *tipb.ChecksumRewriteRule
 	if oldTable != nil {
@@ -220,6 +242,8 @@
 		SetStartTS(startTS).
 		SetChecksumRequest(checksum).
 		SetConcurrency(int(concurrency)).
+		SetResourceGroupName(resourceGroupName).
+		SetExplicitRequestSourceType(explicitRequestSourceType).
 		Build()
 }
 
@@ -232,6 +256,7 @@
 	concurrency uint,
 	oldKeyspace []byte,
 	newKeyspace []byte,
+	resourceGroupName, ExplicitRequestSourceType string,
 ) (*kv.Request, error) {
 	var rule *tipb.ChecksumRewriteRule
 	if oldIndexInfo != nil {
@@ -257,6 +282,8 @@
 		SetStartTS(startTS).
 		SetChecksumRequest(checksum).
 		SetConcurrency(int(concurrency)).
+		SetResourceGroupName(resourceGroupName).
+		SetExplicitRequestSourceType(ExplicitRequestSourceType).
 		Build()
 }
 

br/pkg/gluetidb/glue.go

@@ -188,7 +188,7 @@ func (gs *tidbSession) ExecuteInternal(ctx context.Context, sql string, args ...
 		vars.TxnCtxMu.Unlock()
 	}()
 	// Some of SQLs (like ADMIN RECOVER INDEX) may lazily take effect
-	// when we polling the result set.
+	// when we are polling the result set.
 	// At least call `next` once for triggering theirs side effect.
 	// (Maybe we'd better drain all returned rows?)
 	if rs != nil {
@@ -379,7 +379,7 @@ func (s *mockSession) ExecuteInternal(ctx context.Context, sql string, args ...i
 		return err
 	}
 	// Some of SQLs (like ADMIN RECOVER INDEX) may lazily take effect
-	// when we polling the result set.
+	// when we are polling the result set.
 	// At least call `next` once for triggering theirs side effect.
 	// (Maybe we'd better drain all returned rows?)
 	if rs != nil {

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

@@ -77,6 +77,7 @@ go_library(
         "@com_github_tikv_client_go_v2//kv",
         "@com_github_tikv_client_go_v2//oracle",
         "@com_github_tikv_client_go_v2//tikv",
+        "@com_github_tikv_client_go_v2//util",
         "@com_github_tikv_pd_client//:client",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//backoff",

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

@@ -268,28 +268,34 @@
 
 // TiKVChecksumManager is a manager that can compute checksum of a table using TiKV.
 type TiKVChecksumManager struct {
-	client                 kv.Client
-	manager                gcTTLManager
-	distSQLScanConcurrency uint
-	backoffWeight          int
+	client                    kv.Client
+	manager                   gcTTLManager
+	distSQLScanConcurrency    uint
+	backoffWeight             int
+	resourceGroupName         string
+	explicitRequestSourceType string
 }
 
 var _ ChecksumManager = &TiKVChecksumManager{}
 
 // NewTiKVChecksumManager return a new tikv checksum manager
-func NewTiKVChecksumManager(client kv.Client, pdClient pd.Client, distSQLScanConcurrency uint, backoffWeight int) *TiKVChecksumManager {
+func NewTiKVChecksumManager(client kv.Client, pdClient pd.Client, distSQLScanConcurrency uint, backoffWeight int, resourceGroupName, explicitRequestSourceType string) *TiKVChecksumManager {
 	return &TiKVChecksumManager{
-		client:                 client,
-		manager:                newGCTTLManager(pdClient),
-		distSQLScanConcurrency: distSQLScanConcurrency,
-		backoffWeight:          backoffWeight,
+		client:                    client,
+		manager:                   newGCTTLManager(pdClient),
+		distSQLScanConcurrency:    distSQLScanConcurrency,
+		backoffWeight:             backoffWeight,
+		resourceGroupName:         resourceGroupName,
+		explicitRequestSourceType: explicitRequestSourceType,
 	}
 }
 
 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).
+		SetResourceGroupName(e.resourceGroupName).
+		SetExplicitRequestSourceType(e.explicitRequestSourceType).
 		Build()
 	if err != nil {
 		return nil, errors.Trace(err)

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

@@ -421,10 +421,11 @@
 	KeyspaceName string
 	// the scope when pause PD schedulers.
 	PausePDSchedulerScope config.PausePDSchedulerScope
+	ResourceGroupName     string
 }
 
 // NewBackendConfig creates a new BackendConfig.
-func NewBackendConfig(cfg *config.Config, maxOpenFiles int, keyspaceName string) BackendConfig {
+func NewBackendConfig(cfg *config.Config, maxOpenFiles int, keyspaceName, resourceGroupName string) BackendConfig {
 	return BackendConfig{
 		PDAddr:                  cfg.TiDB.PdAddr,
 		LocalStoreDir:           cfg.TikvImporter.SortedKVDir,
@@ -445,6 +446,7 @@
 		MaxOpenFiles:            maxOpenFiles,
 		KeyspaceName:            keyspaceName,
 		PausePDSchedulerScope:   cfg.TikvImporter.PausePDSchedulerScope,
+		ResourceGroupName:       resourceGroupName,
 	}
 }
 

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

@@ -36,8 +36,10 @@ import (
 	"github.com/pingcap/tidb/br/pkg/lightning/metric"
 	"github.com/pingcap/tidb/br/pkg/logutil"
 	"github.com/pingcap/tidb/br/pkg/restore/split"
+	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/util/codec"
 	"github.com/pingcap/tidb/util/mathutil"
+	"github.com/tikv/client-go/v2/util"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 )
@@ -230,6 +232,12 @@ func (local *Backend) writeToTiKV(ctx context.Context, j *regionJob) error {
 		Chunk: &sst.WriteRequest_Meta{
 			Meta: meta,
 		},
+		Context: &kvrpcpb.Context{
+			ResourceControlContext: &kvrpcpb.ResourceControlContext{
+				ResourceGroupName: local.ResourceGroupName,
+			},
+			RequestSource: util.BuildRequestSource(true, kv.InternalTxnLightning, util.ExplicitTypeLightning),
+		},
 	}
 	for _, peer := range region.GetPeers() {
 		cli, err := clientFactory.Create(ctx, peer.StoreId)
@@ -558,6 +566,10 @@ func (local *Backend) doIngest(ctx context.Context, j *regionJob) (*sst.IngestRe
 			RegionId:    j.region.Region.GetId(),
 			RegionEpoch: j.region.Region.GetRegionEpoch(),
 			Peer:        leader,
+			ResourceControlContext: &kvrpcpb.ResourceControlContext{
+				ResourceGroupName: local.ResourceGroupName,
+			},
+			RequestSource: util.BuildRequestSource(true, kv.InternalTxnLightning, util.ExplicitTypeLightning),
 		}
 
 		if supportMultiIngest {

br/pkg/lightning/common/util.go

@@ -616,6 +616,11 @@
 	return strconv.Atoi(val)
 }
 
+// GetExplicitRequestSourceTypeFromDB gets the explicit request source type from database.
+func GetExplicitRequestSourceTypeFromDB(ctx context.Context, db *sql.DB) (string, error) {
+	return getSessionVariable(ctx, db, variable.TiDBExplicitRequestSourceType)
+}
+
 // 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)
@@ -648,3 +653,10 @@
 
 	return value, nil
 }
+
+// IsFunctionNotExistErr checks if err is a function not exist error.
+func IsFunctionNotExistErr(err error, functionName string) bool {
+	return err != nil &&
+		(strings.Contains(err.Error(), "No database selected") ||
+			strings.Contains(err.Error(), fmt.Sprintf("%s does not exist", functionName)))
+}

br/pkg/lightning/importer/BUILD.bazel

@@ -58,6 +58,7 @@ go_library(
         "//parser/model",
         "//parser/mysql",
         "//planner/core",
+        "//sessionctx/variable",
         "//store/driver",
         "//store/driver/txn",
         "//store/pdtypes",
@@ -85,6 +86,7 @@ go_library(
         "@com_github_pingcap_kvproto//pkg/metapb",
         "@com_github_prometheus_client_golang//prometheus",
         "@com_github_tikv_client_go_v2//config",
+        "@com_github_tikv_client_go_v2//util",
         "@com_github_tikv_pd_client//:client",
         "@io_etcd_go_etcd_client_v3//:client",
         "@org_golang_google_grpc//:grpc",

br/pkg/lightning/importer/checksum_helper.go

@@ -60,7 +60,13 @@
 			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)
+
+		explicitRequestSourceType, err := common.GetExplicitRequestSourceTypeFromDB(ctx, rc.db)
+		if err != nil {
+			log.FromContext(ctx).Warn("get tidb_request_source_type failed", zap.Error(err), zap.String("tidb_request_source_type", explicitRequestSourceType))
+			return nil, errors.Trace(err)
+		}
+		manager = local.NewTiKVChecksumManager(store.GetClient(), pdCli, uint(rc.cfg.TiDB.DistSQLScanConcurrency), backoffWeight, rc.resourceGroupName, explicitRequestSourceType)
 	} else {
 		manager = local.NewTiDBChecksumExecutor(rc.db)
 	}