feat: introduce QueryCopStoreLimiter to manage TiKV cop request concurrency

- Added QueryCopStoreLimiter to limit the number of in-flight cop requests per TiKV store within a single query.
- Updated DistSQLContext to include QueryCopStoreLimiter.
- Modified Select function to apply the QueryCopStoreLimiter when sending requests to TiKV.
- Enhanced tests to validate the behavior of the new limiter, ensuring it correctly manages concurrent requests.
- Adjusted related variables and configurations to support the new functionality.

Signed-off-by: Zhigao TONG <tongzhigao@pingcap.com>

Author: solotzg

pkg/distsql/BUILD.bazel

@@ -68,7 +68,7 @@ go_test(
     embed = [":distsql"],
     flaky = True,
     race = "on",
-    shard_count = 33,
+    shard_count = 34,
     deps = [
         "//pkg/distsql/context",
         "//pkg/errctx",

pkg/distsql/context/context.go

@@ -63,6 +63,7 @@ type DistSQLContext struct {
 	TiFlashQuerySpillRatio               float64
 	TiFlashHashJoinVersion               string
 
+	QueryCopStoreLimiter          *kv.QueryCopStoreLimiter
 	DistSQLConcurrency            int
 	ReplicaReadType               kv.ReplicaReadType
 	WeakConsistency               bool

pkg/distsql/context/context_test.go

@@ -111,6 +111,7 @@ func TestContextDetach(t *testing.T) {
 			"$.RunawayChecker",
 			"$.RUConsumptionReporter",
 			"$.ExecDetails",
+			"$.QueryCopStoreLimiter",
 		}))
 
 	staticObj := obj.Detach()

pkg/distsql/distsql.go

@@ -60,6 +60,10 @@ func Select(ctx context.Context, dctx *distsqlctx.DistSQLContext, kvReq *kv.Requ
 	r, ctx := tracing.StartRegionEx(ctx, "distsql.Select")
 	defer r.End()
 
+	if kvReq.StoreType == kv.TiKV && dctx.QueryCopStoreLimiter != nil {
+		kvReq.QueryCopStoreLimiter = dctx.QueryCopStoreLimiter
+	}
+
 	// For testing purpose.
 	if hook := ctx.Value("CheckSelectRequestHook"); hook != nil {
 		hook.(func(*kv.Request))(kvReq)

pkg/distsql/distsql_test.go

@@ -103,6 +103,69 @@ func TestSelectWithRuntimeStats(t *testing.T) {
 	require.NoError(t, response.Close())
 }
 
+func TestSelectAppliesQueryCopStoreLimiter(t *testing.T) {
+	sctx := newMockSessionContext()
+	sctx.GetSessionVars().QueryCopStoreLimit = 3
+	dctx := sctx.GetDistSQLCtx()
+	require.NotNil(t, dctx.QueryCopStoreLimiter)
+	require.Equal(t, 3, dctx.QueryCopStoreLimiter.Capacity())
+
+	colTypes := []*types.FieldType{types.NewFieldType(mysql.TypeLonglong)}
+	buildRequest := func(storeType kv.StoreType) *kv.Request {
+		request, err := (&RequestBuilder{}).SetKeyRanges(nil).
+			SetDAGRequest(&tipb.DAGRequest{}).
+			SetStoreType(storeType).
+			SetFromSessionVars(DefaultDistSQLContext).
+			SetMemTracker(memory.NewTracker(-1, -1)).
+			Build()
+		require.NoError(t, err)
+		return request
+	}
+	checkRequest := func(check func(*kv.Request)) context.Context {
+		return context.WithValue(context.TODO(), "CheckSelectRequestHook", func(req *kv.Request) {
+			check(req)
+		})
+	}
+
+	request := buildRequest(kv.TiKV)
+	response, err := Select(checkRequest(func(req *kv.Request) {
+		require.Nil(t, req.CoprRequestRateLimit)
+		require.Same(t, dctx.QueryCopStoreLimiter, req.QueryCopStoreLimiter)
+	}), dctx, request, colTypes)
+	require.NoError(t, err)
+	require.NoError(t, response.Close())
+
+	request = buildRequest(kv.TiFlash)
+	response, err = Select(checkRequest(func(req *kv.Request) {
+		require.Nil(t, req.CoprRequestRateLimit)
+		require.Nil(t, req.QueryCopStoreLimiter)
+	}), dctx, request, colTypes)
+	require.NoError(t, err)
+	require.NoError(t, response.Close())
+
+	request = buildRequest(kv.TiKV)
+	explicitRateLimit := kv.NewCoprRequestRateLimit(7)
+	request.CoprRequestRateLimit = explicitRateLimit
+	response, err = Select(checkRequest(func(req *kv.Request) {
+		require.Same(t, explicitRateLimit, req.CoprRequestRateLimit)
+		require.Same(t, dctx.QueryCopStoreLimiter, req.QueryCopStoreLimiter)
+		require.False(t, req.CoprRequestRateLimit.Acquire(make(chan struct{})))
+		req.CoprRequestRateLimit.Release()
+	}), dctx, request, colTypes)
+	require.NoError(t, err)
+	require.NoError(t, response.Close())
+
+	dctx.QueryCopStoreLimiter = nil
+	request = buildRequest(kv.TiKV)
+	request.CoprRequestRateLimit = explicitRateLimit
+	response, err = Select(checkRequest(func(req *kv.Request) {
+		require.Same(t, explicitRateLimit, req.CoprRequestRateLimit)
+		require.Nil(t, req.QueryCopStoreLimiter)
+	}), dctx, request, colTypes)
+	require.NoError(t, err)
+	require.NoError(t, response.Close())
+}
+
 func TestSelectResultRuntimeStats(t *testing.T) {
 	stmtStats := execdetails.NewRuntimeStatsColl(nil)
 	basic := stmtStats.GetBasicRuntimeStats(1, true)

pkg/distsql/request_builder.go

@@ -391,7 +391,7 @@ func (builder *RequestBuilder) SetConcurrency(concurrency int) *RequestBuilder {
 }
 
 // SetCoprRequestRateLimit sets a shared in-flight cop request limiter for this request.
-func (builder *RequestBuilder) SetCoprRequestRateLimit(rateLimit *util.RateLimit) *RequestBuilder {
+func (builder *RequestBuilder) SetCoprRequestRateLimit(rateLimit kv.CoprRequestLimiter) *RequestBuilder {
 	builder.Request.CoprRequestRateLimit = rateLimit
 	return builder
 }

pkg/executor/distsql.go

@@ -65,7 +65,6 @@ import (
 	rangerctx "github.com/pingcap/tidb/pkg/util/ranger/context"
 	"github.com/pingcap/tidb/pkg/util/size"
 	"github.com/pingcap/tipb/go-tipb"
-	tikvutil "github.com/tikv/client-go/v2/util"
 	"go.uber.org/zap"
 )
 
@@ -1022,7 +1021,7 @@ func getSelectResultInFlightCost(result distsql.SelectResult) int {
 	return inFlightCost
 }
 
-func getMergeSortSharedCoprRequestRateLimit(needMerge bool, distSQLConcurrency int) *tikvutil.RateLimit {
+func getMergeSortSharedCoprRequestRateLimit(needMerge bool, distSQLConcurrency int) kv.CoprRequestLimiter {
 	if !needMerge {
 		return nil
 	}
@@ -1032,7 +1031,7 @@ func getMergeSortSharedCoprRequestRateLimit(needMerge bool, distSQLConcurrency i
 	if capacity < 1 {
 		capacity = 1
 	}
-	return tikvutil.NewRateLimit(2 * capacity)
+	return kv.NewCoprRequestRateLimit(2 * capacity)
 }
 
 func getMergeSortIndexScanConcurrency(needMerge bool, kvRangesCount int, distSQLConcurrency int) int {
@@ -1067,7 +1066,7 @@ func (e *IndexLookUpExecutor) buildIndexSelectResultForRange(
 	totalRanges int,
 	batchSize int,
 	indexScanConcurrency int,
-	sharedCoprRequestRateLimit *tikvutil.RateLimit,
+	sharedCoprRequestRateLimit kv.CoprRequestLimiter,
 ) (distsql.SelectResult, error) {
 	if tblScanIdxForRewritePartitionID >= 0 {
 		// We should set the TblScan's TableID to the partition physical ID to make sure

pkg/executor/set_test.go

@@ -1447,6 +1447,7 @@ func TestSetConcurrency(t *testing.T) {
 	tk.MustQuery("select @@tidb_streamagg_concurrency;").Check(testkit.Rows(strconv.Itoa(vardef.DefTiDBStreamAggConcurrency)))
 	tk.MustQuery("select @@tidb_projection_concurrency;").Check(testkit.Rows(strconv.Itoa(vardef.ConcurrencyUnset)))
 	tk.MustQuery("select @@tidb_distsql_scan_concurrency;").Check(testkit.Rows(strconv.Itoa(vardef.DefDistSQLScanConcurrency)))
+	tk.MustQuery("select @@tidb_query_cop_store_limit;").Check(testkit.Rows(strconv.Itoa(vardef.DefTiDBQueryCopStoreLimit)))
 
 	tk.MustQuery("select @@tidb_index_serial_scan_concurrency;").Check(testkit.Rows(strconv.Itoa(vardef.DefIndexSerialScanConcurrency)))
 
@@ -1461,6 +1462,7 @@ func TestSetConcurrency(t *testing.T) {
 	require.Equal(t, vardef.DefTiDBStreamAggConcurrency, vars.StreamAggConcurrency())
 	require.Equal(t, vardef.DefExecutorConcurrency, vars.ProjectionConcurrency())
 	require.Equal(t, vardef.DefDistSQLScanConcurrency, vars.DistSQLScanConcurrency())
+	require.Equal(t, vardef.DefTiDBQueryCopStoreLimit, vars.QueryCopStoreLimit)
 
 	// test setting deprecated variables
 	warnTpl := "Warning 1287 '%s' is deprecated and will be removed in a future release. Please use tidb_executor_concurrency instead"
@@ -1499,6 +1501,10 @@ func TestSetConcurrency(t *testing.T) {
 	tk.MustQuery(fmt.Sprintf("select @@%s;", vardef.TiDBDistSQLScanConcurrency)).Check(testkit.Rows("1"))
 	require.Equal(t, 1, vars.DistSQLScanConcurrency())
 
+	tk.MustExec(fmt.Sprintf("set @@%s=8;", vardef.TiDBQueryCopStoreLimit))
+	tk.MustQuery(fmt.Sprintf("select @@%s;", vardef.TiDBQueryCopStoreLimit)).Check(testkit.Rows("8"))
+	require.Equal(t, 8, vars.QueryCopStoreLimit)
+
 	tk.MustExec("set @@tidb_index_serial_scan_concurrency=4")
 	tk.MustQuery("show warnings").Check(testkit.Rows("Warning 1287 The 'tidb_index_serial_scan_concurrency' variable is deprecated. Sequential scans follow 'tidb_executor_concurrency', and index statistics collection uses 'tidb_analyze_distsql_scan_concurrency'."))
 	tk.MustQuery("select @@tidb_index_serial_scan_concurrency;").Check(testkit.Rows("4"))

pkg/kv/BUILD.bazel

@@ -85,7 +85,7 @@ go_test(
     ],
     embed = [":kv"],
     flaky = True,
-    shard_count = 27,
+    shard_count = 32,
     deps = [
         "//pkg/config/kerneltype",
         "//pkg/keyspace",