statistics: add better API and better logs for lock stats (#47024)

ref #46351

executor/lockstats/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
  name = "lockstats",
@@ -19,3 +19,17 @@ go_library(
  "@com_github_pingcap_errors//:errors",
  ],
 )
+
+go_test(
+ name = "lockstats_test",
+ timeout = "short",
+ srcs = ["lock_stats_executor_test.go"],
+ embed = [":lockstats"],
+ flaky = True,
+ deps = [
+ "//infoschema",
+ "//parser/ast",
+ "//parser/model",
+ "@com_github_stretchr_testify//require",
+ ],
+)

executor/lockstats/lock_stats_executor.go

@@ -16,6 +16,7 @@ package lockstats
 
 import (
  "context"
+ "fmt"
 
  "github.com/pingcap/errors"
  "github.com/pingcap/tidb/domain"
@@ -51,12 +52,13 @@ func (e *LockExec) Next(_ context.Context, _ *chunk.Chunk) error {
  is := do.InfoSchema()
 
  if e.onlyLockPartitions() {
- tableName := e.Tables[0]
- tid, pidNames, err := populatePartitionIDAndNames(tableName, tableName.PartitionNames, is)
+ table := e.Tables[0]
+ tid, pidNames, err := populatePartitionIDAndNames(table, table.PartitionNames, is)
  if err != nil {
  return err
  }
 
+ tableName := fmt.Sprintf("%s.%s", table.Schema.L, table.Name.L)
  msg, err := h.LockPartitions(tid, tableName, pidNames)
  if err != nil {
  return err
@@ -65,12 +67,12 @@ func (e *LockExec) Next(_ context.Context, _ *chunk.Chunk) error {
  e.Ctx().GetSessionVars().StmtCtx.AppendWarning(errors.New(msg))
  }
  } else {
- tids, pids, err := populateTableAndPartitionIDs(e.Tables, is)
+ tidAndNames, pidAndNames, err := populateTableAndPartitionIDs(e.Tables, is)
  if err != nil {
  return err
  }
 
- msg, err := h.LockTables(tids, pids, e.Tables)
+ msg, err := h.LockTables(tidAndNames, pidAndNames)
  if err != nil {
  return err
  }
@@ -97,15 +99,23 @@ func (*LockExec) Open(context.Context) error {
 }
 
 // populatePartitionIDAndNames returns the table ID and partition IDs for the given table name and partition names.
-func populatePartitionIDAndNames(tableName *ast.TableName, partitionNames []model.CIStr, is infoschema.InfoSchema) (int64, map[int64]string, error) {
- tbl, err := is.TableByName(tableName.Schema, tableName.Name)
+func populatePartitionIDAndNames(
+ table *ast.TableName,
+ partitionNames []model.CIStr,
+ is infoschema.InfoSchema,
+) (int64, map[int64]string, error) {
+ if len(partitionNames) == 0 {
+ return 0, nil, errors.New("partition list should not be empty")
+ }
+ tbl, err := is.TableByName(table.Schema, table.Name)
  if err != nil {
  return 0, nil, err
  }
 
  pi := tbl.Meta().GetPartitionInfo()
  if pi == nil {
- return 0, nil, errors.Errorf("table %s is not a partition table", tableName.Name)
+ return 0, nil, errors.Errorf("table %s is not a partition table",
+ fmt.Sprintf("%s.%s", table.Schema.L, table.Name.L))
  }
 
  pidNames := make(map[int64]string, len(partitionNames))
@@ -121,25 +131,36 @@ func populatePartitionIDAndNames(tableName *ast.TableName, partitionNames []mode
 }
 
 // populateTableAndPartitionIDs returns table IDs and partition IDs for the given table names.
-func populateTableAndPartitionIDs(tables []*ast.TableName, is infoschema.InfoSchema) ([]int64, []int64, error) {
- tids := make([]int64, 0, len(tables))
- pids := make([]int64, 0)
+func populateTableAndPartitionIDs(
+ tables []*ast.TableName,
+ is infoschema.InfoSchema,
+) (map[int64]string, map[int64]string, error) {
+ if len(tables) == 0 {
+ return nil, nil, errors.New("table list should not be empty")
+ }
+
+ tidAndNames := make(map[int64]string, len(tables))
+ pidAndNames := make(map[int64]string, len(tables))
 
  for _, table := range tables {
  tbl, err := is.TableByName(table.Schema, table.Name)
  if err != nil {
  return nil, nil, err
  }
- tids = append(tids, tbl.Meta().ID)
+ tidAndNames[tbl.Meta().ID] = fmt.Sprintf("%s.%s", table.Schema.L, table.Name.L)
 
  pi := tbl.Meta().GetPartitionInfo()
  if pi == nil {
  continue
  }
  for _, p := range pi.Definitions {
- pids = append(pids, p.ID)
+ pidAndNames[p.ID] = genFullPartitionName(table, p.Name.L)
  }
  }
 
- return tids, pids, nil
+ return tidAndNames, pidAndNames, nil
+}
+
+func genFullPartitionName(table *ast.TableName, partitionName string) string {
+ return fmt.Sprintf("%s.%s partition (%s)", table.Schema.L, table.Name.L, partitionName)
 }

executor/lockstats/lock_stats_executor_test.go

@@ -0,0 +1,108 @@
+// 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 lockstats
+
+import (
+ "testing"
+
+ "github.com/pingcap/tidb/infoschema"
+ "github.com/pingcap/tidb/parser/ast"
+ "github.com/pingcap/tidb/parser/model"
+ "github.com/stretchr/testify/require"
+)
+
+func TestPopulatePartitionIDAndNames(t *testing.T) {
+ fakeInfo := infoschema.MockInfoSchema([]*model.TableInfo{
+ tInfo(1, "t1", "p1", "p2"),
+ })
+
+ table := &ast.TableName{
+ Schema: model.NewCIStr("test"),
+ Name: model.NewCIStr("t1"),
+ PartitionNames: []model.CIStr{
+ model.NewCIStr("p1"),
+ model.NewCIStr("p2"),
+ },
+ }
+
+ gotTID, gotPIDNames, err := populatePartitionIDAndNames(table, table.PartitionNames, fakeInfo)
+ require.NoError(t, err)
+ require.Equal(t, int64(1), gotTID)
+ require.Equal(t, map[int64]string{
+ 2: "p1",
+ 3: "p2",
+ }, gotPIDNames)
+
+ // Empty partition names.
+ _, _, err = populatePartitionIDAndNames(nil, nil, fakeInfo)
+ require.Error(t, err)
+}
+
+func TestPopulateTableAndPartitionIDs(t *testing.T) {
+ fakeInfo := infoschema.MockInfoSchema([]*model.TableInfo{
+ tInfo(1, "t1", "p1", "p2"),
+ tInfo(4, "t2"),
+ })
+
+ tables := []*ast.TableName{
+ {
+ Schema: model.NewCIStr("test"),
+ Name: model.NewCIStr("t1"),
+ PartitionNames: []model.CIStr{
+ model.NewCIStr("p1"),
+ model.NewCIStr("p2"),
+ },
+ },
+ {
+ Schema: model.NewCIStr("test"),
+ Name: model.NewCIStr("t2"),
+ },
+ }
+
+ gotTIDAndNames, gotPIDAndNames, err := populateTableAndPartitionIDs(tables, fakeInfo)
+ require.NoError(t, err)
+ require.Equal(t, map[int64]string{
+ 1: "test.t1",
+ 4: "test.t2",
+ }, gotTIDAndNames)
+ require.Equal(t, map[int64]string{
+ 2: "test.t1 partition (p1)",
+ 3: "test.t1 partition (p2)",
+ }, gotPIDAndNames)
+
+ // Empty table list.
+ _, _, err = populateTableAndPartitionIDs(nil, fakeInfo)
+ require.Error(t, err)
+}
+
+func tInfo(id int, tableName string, partitionNames ...string) *model.TableInfo {
+ tbl := &model.TableInfo{
+ ID: int64(id),
+ Name: model.NewCIStr(tableName),
+ }
+ if len(partitionNames) > 0 {
+ tbl.Partition = &model.PartitionInfo{
+ Enable: true,
+ }
+ for i, partitionName := range partitionNames {
+ tbl.Partition.Definitions = append(tbl.Partition.Definitions, model.PartitionDefinition{
+ ID: int64(id + 1 + i),
+ Name: model.NewCIStr(partitionName),
+ })
+ }
+ }
+
+ return tbl
+}

executor/lockstats/unlock_stats_executor.go

@@ -16,6 +16,7 @@ package lockstats
 
 import (
  "context"
+ "fmt"
 
  "github.com/pingcap/errors"
  "github.com/pingcap/tidb/domain"
@@ -48,11 +49,12 @@ func (e *UnlockExec) Next(context.Context, *chunk.Chunk) error {
  is := do.InfoSchema()
 
  if e.onlyUnlockPartitions() {
- tableName := e.Tables[0]
- tid, pidNames, err := populatePartitionIDAndNames(tableName, tableName.PartitionNames, is)
+ table := e.Tables[0]
+ tid, pidNames, err := populatePartitionIDAndNames(table, table.PartitionNames, is)
  if err != nil {
  return err
  }
+ tableName := fmt.Sprintf("%s.%s", table.Schema.O, table.Name.O)
  msg, err := h.RemoveLockedPartitions(tid, tableName, pidNames)
  if err != nil {
  return err
@@ -61,11 +63,11 @@ func (e *UnlockExec) Next(context.Context, *chunk.Chunk) error {
  e.Ctx().GetSessionVars().StmtCtx.AppendWarning(errors.New(msg))
  }
  } else {
- tids, pids, err := populateTableAndPartitionIDs(e.Tables, is)
+ tidAndNames, pidAndNames, err := populateTableAndPartitionIDs(e.Tables, is)
  if err != nil {
  return err
  }
- msg, err := h.RemoveLockedTables(tids, pids, e.Tables)
+ msg, err := h.RemoveLockedTables(tidAndNames, pidAndNames)
  if err != nil {
  return err
  }

statistics/handle/handletest/lockstats/BUILD.bazel

@@ -9,7 +9,7 @@ go_test(
  "main_test.go",
  ],
  flaky = True,
- shard_count = 12,
+ shard_count = 14,
  deps = [
  "//config",
  "//domain",

statistics/handle/handletest/lockstats/lock_partition_stats_test.go

@@ -148,7 +148,7 @@ func TestLockAndUnlockPartitionStatsRepeatedly(t *testing.T) {
  // Lock the partition again and check the warning.
  tk.MustExec("lock stats t partition p0")
  tk.MustQuery("show warnings").Check(testkit.Rows(
- "Warning 1105 skip locking locked table: test.t partition (p0)",
+ "Warning 1105 skip locking locked partition of table test.t: p0",
  ))
 
  // Unlock the partition.
@@ -160,7 +160,7 @@ func TestLockAndUnlockPartitionStatsRepeatedly(t *testing.T) {
  // Unlock the partition again and check the warning.
  tk.MustExec("unlock stats t partition p0")
  tk.MustQuery("show warnings").Check(testkit.Rows(
- "Warning 1105 skip unlocking unlocked table: test.t partition (p0)",
+ "Warning 1105 skip unlocking unlocked partition of table test.t: p0",
  ))
 }
 
@@ -260,6 +260,28 @@ func TestUnlockTheWholeTableWouldUnlockLockedPartitionsAndGenerateWarning(t *tes
  require.Equal(t, 0, num)
 }
 
+func TestSkipLockALotOfPartitions(t *testing.T) {
+ store, _ := testkit.CreateMockStoreAndDomain(t)
+ tk := testkit.NewTestKit(t, store)
+ tk.MustExec("set @@tidb_analyze_version = 1")
+ tk.MustExec("use test")
+ tk.MustExec("drop table if exists t")
+ tk.MustExec("create table t(a int, b varchar(10), index idx_b (b)) partition by range(a) " +
+ "(partition p0 values less than (10), partition p1 values less than (20), " +
+ "partition a values less than (30), " +
+ "partition b values less than (40), " +
+ "partition g values less than (90), " +
+ "partition h values less than (100))")
+
+ tk.MustExec("lock stats t partition p0, p1, a, b, g, h")
+
+ // Skip locking a lot of partitions.
+ tk.MustExec("lock stats t partition p0, p1, a, b, g, h")
+ tk.MustQuery("show warnings").Check(testkit.Rows(
+ "Warning 1105 skip locking locked partitions of table test.t: a, b, g, h, p0, p1",
+ ))
+}
+
 func setupTestEnvironmentWithPartitionedTableT(t *testing.T) (kv.Storage, *testkit.TestKit, *model.TableInfo) {
  store, dom := testkit.CreateMockStoreAndDomain(t)
  tk := testkit.NewTestKit(t, store)

statistics/handle/handletest/lockstats/lock_table_stats_test.go

@@ -316,6 +316,26 @@ func TestShowStatsLockedTablePrivilege(t *testing.T) {
  require.Len(t, rows, 1)
 }
 
+func TestSkipLockALotOfTables(t *testing.T) {
+ store, _ := testkit.CreateMockStoreAndDomain(t)
+ tk := testkit.NewTestKit(t, store)
+ tk.MustExec("set @@tidb_analyze_version = 1")
+ tk.MustExec("use test")
+ tk.MustExec("drop table if exists t; drop table if exists a; drop table if exists x; drop table if exists y; drop table if exists z")
+ tk.MustExec("create table t(a int, b varchar(10), index idx_b (b)); " +
+ "create table a(a int, b varchar(10), index idx_b (b)); " +
+ "create table x(a int, b varchar(10), index idx_b (b)); " +
+ "create table y(a int, b varchar(10), index idx_b (b)); " +
+ "create table z(a int, b varchar(10), index idx_b (b))")
+ tk.MustExec("lock stats test.t, test.a, test.x, test.y, test.z")
+
+ // Skip locking a lot of tables.
+ tk.MustExec("lock stats test.t, test.a, test.x, test.y, test.z")
+ tk.MustQuery("show warnings").Check(testkit.Rows(
+ "Warning 1105 skip locking locked tables: test.a, test.t, test.x, test.y, test.z",
+ ))
+}
+
 func setupTestEnvironmentWithTableT(t *testing.T) (kv.Storage, *testkit.TestKit, *model.TableInfo) {
  store, dom := testkit.CreateMockStoreAndDomain(t)
  tk := testkit.NewTestKit(t, store)