store/copr: set concurrency to partition number for limit statements

distsql/request_builder.go

@@ -158,7 +158,11 @@ func (builder *RequestBuilder) SetDAGRequest(dag *tipb.DAGRequest) *RequestBuild
  if len(dag.Executors) == 2 && dag.Executors[1].GetLimit() != nil {
  limit := dag.Executors[1].GetLimit()
  if limit != nil && limit.Limit < estimatedRegionRowCount {
- builder.Request.Concurrency = 1
+ if kr := builder.Request.KeyRanges; kr != nil {
+ builder.Request.Concurrency = kr.PartitionNum()
+ } else {
+ builder.Request.Concurrency = 1
+ }
  }
  builder.Request.LimitSize = limit.GetLimit()
  }
@@ -265,9 +269,13 @@ func (*RequestBuilder) getKVPriority(sv *variable.SessionVars) int {
 // SetFromSessionVars sets the following fields for "kv.Request" from session variables:
 // "Concurrency", "IsolationLevel", "NotFillCache", "TaskID", "Priority", "ReplicaRead", "ResourceGroupTagger".
 func (builder *RequestBuilder) SetFromSessionVars(sv *variable.SessionVars) *RequestBuilder {
+ distsqlConcurrency := sv.DistSQLScanConcurrency()
  if builder.Request.Concurrency == 0 {
- // Concurrency may be set to 1 by SetDAGRequest
- builder.Request.Concurrency = sv.DistSQLScanConcurrency()
+ // Concurrency unset.
+ builder.Request.Concurrency = distsqlConcurrency
+ } else if builder.Request.Concurrency > distsqlConcurrency {
+ // Concurrency is set in SetDAGRequest, check the upper limit.
+ builder.Request.Concurrency = distsqlConcurrency
  }
  replicaReadType := sv.GetReplicaRead()
  if sv.StmtCtx.WeakConsistency {

executor/distsqltest/BUILD.bazel

@@ -0,0 +1,18 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_test")
+
+go_test(
+ name = "distsqltest_test",
+ srcs = [
+ "distsql_test.go",
+ "main_test.go",
+ ],
+ deps = [
+ "//config",
+ "//kv",
+ "//meta/autoid",
+ "//testkit",
+ "@com_github_stretchr_testify//require",
+ "@com_github_tikv_client_go_v2//tikv",
+ "@org_uber_go_goleak//:goleak",
+ ],
+)

executor/distsqltest/distsql_test.go

@@ -0,0 +1,78 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package distsql_test
+
+import (
+ "context"
+ "fmt"
+ "strings"
+ "testing"
+
+ "github.com/pingcap/tidb/kv"
+ "github.com/pingcap/tidb/testkit"
+ "github.com/stretchr/testify/require"
+)
+
+func TestDistsqlPartitionTableConcurrency(t *testing.T) {
+ store := testkit.CreateMockStore(t)
+ tk := testkit.NewTestKit(t, store)
+ tk.MustExec("use test")
+ tk.MustExec("drop table if exists t1, t2, t3")
+ tk.MustExec("create table t1(id int primary key , val int)")
+ partitions := make([]string, 0, 20)
+ for i := 0; i < 20; i++ {
+ pid := i + 1
+ partitions = append(partitions, fmt.Sprintf("PARTITION p%d VALUES LESS THAN (%d00)", pid, pid))
+ }
+ tk.MustExec("create table t2(id int primary key, val int)" +
+ "partition by range(id)" +
+ "(" + strings.Join(partitions[:10], ",") + ")")
+ tk.MustExec("create table t3(id int primary key, val int)" +
+ "partition by range(id)" +
+ "(" + strings.Join(partitions, ",") + ")")
+ for i := 0; i < 20; i++ {
+ for _, tbl := range []string{"t1", "t2", "t3"} {
+ tk.MustExec(fmt.Sprintf("insert into %s values(%d, %d)", tbl, i*50, i*50))
+ }
+ }
+ tk.MustExec("analyze table t1, t2, t3")
+ // non-partitioned table checker
+ ctx1 := context.WithValue(context.Background(), "CheckSelectRequestHook", func(req *kv.Request) {
+ require.Equal(t, req.KeyRanges.PartitionNum(), 1)
+ require.Equal(t, req.Concurrency, 1)
+ })
+ // 10-ranges-partitioned table checker
+ ctx2 := context.WithValue(context.Background(), "CheckSelectRequestHook", func(req *kv.Request) {
+ require.Equal(t, req.KeyRanges.PartitionNum(), 10)
+ require.Equal(t, req.Concurrency, 10)
+ })
+ // 20-ranges-partitioned table checker
+ ctx3 := context.WithValue(context.Background(), "CheckSelectRequestHook", func(req *kv.Request) {
+ require.Equal(t, req.KeyRanges.PartitionNum(), 20)
+ require.Equal(t, req.Concurrency, 15)
+ })
+ ctxs := []context.Context{ctx1, ctx2, ctx3}
+ for i, tbl := range []string{"t1", "t2", "t3"} {
+ ctx := ctxs[i]
+ tk.MustQueryWithContext(ctx, fmt.Sprintf("select * from %s order by id asc limit 1", tbl)).
+ Check(testkit.Rows("0 0"))
+ tk.MustQueryWithContext(ctx, fmt.Sprintf("select * from %s order by id asc limit 5", tbl)).
+ Check(testkit.Rows("0 0", "50 50", "100 100", "150 150", "200 200"))
+ tk.MustQueryWithContext(ctx, fmt.Sprintf("select * from %s order by id desc limit 1", tbl)).
+ Check(testkit.Rows("950 950"))
+ tk.MustQueryWithContext(ctx, fmt.Sprintf("select * from %s order by id desc limit 5", tbl)).
+ Check(testkit.Rows("950 950", "900 900", "850 850", "800 800", "750 750"))
+ }
+}

executor/distsqltest/main_test.go

@@ -0,0 +1,44 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package distsql_test
+
+import (
+ "testing"
+
+ "github.com/pingcap/tidb/config"
+ "github.com/pingcap/tidb/meta/autoid"
+ "github.com/tikv/client-go/v2/tikv"
+ "go.uber.org/goleak"
+)
+
+func TestMain(m *testing.M) {
+ autoid.SetStep(5000)
+ config.UpdateGlobal(func(conf *config.Config) {
+ conf.Log.SlowThreshold = 30000 // 30s
+ conf.TiKVClient.AsyncCommit.SafeWindow = 0
+ conf.TiKVClient.AsyncCommit.AllowedClockDrift = 0
+ conf.Experimental.AllowsExpressionIndex = true
+ })
+ tikv.EnableFailpoints()
+
+ opts := []goleak.Option{
+ goleak.IgnoreTopFunction("github.com/golang/glog.(*loggingT).flushDaemon"),
+ goleak.IgnoreTopFunction("github.com/lestrrat-go/httprc.runFetchWorker"),
+ goleak.IgnoreTopFunction("go.etcd.io/etcd/client/pkg/v3/logutil.(*MergeLogger).outputLoop"),
+ goleak.IgnoreTopFunction("gopkg.in/natefinch/lumberjack%2ev2.(*Logger).millRun"),
+ goleak.IgnoreTopFunction("github.com/tikv/client-go/v2/txnkv/transaction.keepAlive"),
+ }
+ goleak.VerifyTestMain(m, opts...)
+}

planner/core/partition_pruner_test.go

@@ -545,7 +545,9 @@ func TestListColumnsPartitionPrunerRandom(t *testing.T) {
 
 func TestIssue22635(t *testing.T) {
  failpoint.Enable("github.com/pingcap/tidb/planner/core/forceDynamicPrune", `return(true)`)
- defer failpoint.Disable("github.com/pingcap/tidb/planner/core/forceDynamicPrune")
+ defer func() {
+ failpoint.Disable("github.com/pingcap/tidb/planner/core/forceDynamicPrune")
+ }()
  store := testkit.CreateMockStore(t)
  tk := testkit.NewTestKit(t, store)
  tk.MustExec("USE test;")
@@ -557,19 +559,19 @@ CREATE TABLE t1 (
 ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin
 PARTITION BY HASH( a )
 PARTITIONS 4`)
- tk.MustQuery("SELECT (SELECT tt.a FROM t1  tt LIMIT 1) aa, COUNT(DISTINCT b) FROM t1  GROUP BY aa").Check(testkit.Rows()) // work fine without any error
+ tk.MustQuery("SELECT (SELECT tt.a FROM t1 tt ORDER BY a ASC LIMIT 1) aa, COUNT(DISTINCT b) FROM t1 GROUP BY aa").Check(testkit.Rows()) // work fine without any error
 
  tk.MustExec("insert into t1 values (1, 1)")
- tk.MustQuery("SELECT (SELECT tt.a FROM t1  tt LIMIT 1) aa, COUNT(DISTINCT b) FROM t1  GROUP BY aa").Check(testkit.Rows("1 1"))
+ tk.MustQuery("SELECT (SELECT tt.a FROM t1 tt ORDER BY a ASC LIMIT 1) aa, COUNT(DISTINCT b) FROM t1 GROUP BY aa").Check(testkit.Rows("1 1"))
 
  tk.MustExec("insert into t1 values (2, 2), (2, 2)")
- tk.MustQuery("SELECT (SELECT tt.a FROM t1  tt LIMIT 1) aa, COUNT(DISTINCT b) FROM t1  GROUP BY aa").Check(testkit.Rows("1 2"))
+ tk.MustQuery("SELECT (SELECT tt.a FROM t1 tt ORDER BY a ASC LIMIT 1) aa, COUNT(DISTINCT b) FROM t1 GROUP BY aa").Check(testkit.Rows("1 2"))
 
  tk.MustExec("insert into t1 values (3, 3), (3, 3), (3, 3)")
- tk.MustQuery("SELECT (SELECT tt.a FROM t1  tt LIMIT 1) aa, COUNT(DISTINCT b) FROM t1  GROUP BY aa").Check(testkit.Rows("1 3"))
+ tk.MustQuery("SELECT (SELECT tt.a FROM t1 tt ORDER BY a ASC LIMIT 1) aa, COUNT(DISTINCT b) FROM t1 GROUP BY aa").Check(testkit.Rows("1 3"))
 
  tk.MustExec("insert into t1 values (4, 4), (4, 4), (4, 4), (4, 4)")
- tk.MustQuery("SELECT (SELECT tt.a FROM t1  tt LIMIT 1) aa, COUNT(DISTINCT b) FROM t1  GROUP BY aa").Check(testkit.Rows("4 4"))
+ tk.MustQuery("SELECT (SELECT tt.a FROM t1 tt ORDER BY a DESC LIMIT 1) aa, COUNT(DISTINCT b) FROM t1 GROUP BY aa").Check(testkit.Rows("4 4"))
 }
 
 func TestIssue22898(t *testing.T) {

store/copr/coprocessor.go

@@ -217,7 +217,14 @@ func (c *CopClient) BuildCopIterator(ctx context.Context, req *kv.Request, vars
  // For ReqTypeAnalyze, we keep its concurrency to avoid slow analyze(see https://github.com/pingcap/tidb/issues/40162 for details).
  if it.concurrency > 2 && it.req.Tp != kv.ReqTypeAnalyze {
  oldConcurrency := it.concurrency
+ partitionNum := req.KeyRanges.PartitionNum()
+ if partitionNum > it.concurrency {
+ partitionNum = it.concurrency
+ }
  it.concurrency = 2
+ if it.concurrency < partitionNum {
+ it.concurrency = partitionNum
+ }
 
  failpoint.Inject("testRateLimitActionMockConsumeAndAssert", func(val failpoint.Value) {
  if val.(bool) {