statistics: optimize stats delta dumping with batch processing (#58791)

close #57869

pkg/statistics/handle/handletest/statstest/BUILD.bazel

@@ -9,7 +9,7 @@ go_test(
     ],
     flaky = True,
     race = "on",
-    shard_count = 13,
+    shard_count = 14,
     deps = [
         "//pkg/config",
         "//pkg/parser/ast",

pkg/statistics/handle/handletest/statstest/stats_test.go

@@ -436,3 +436,32 @@ func TestInitStatsIssue41938(t *testing.T) {
 	require.NoError(t, h.InitStats(context.Background(), dom.InfoSchema()))
 	h.SetLease(0)
 }
+
+func TestDumpStatsDeltaInBatch(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	testKit := testkit.NewTestKit(t, store)
+	testKit.MustExec("use test")
+	testKit.MustExec("create table t1 (c1 int, c2 int)")
+	testKit.MustExec("insert into t1 values (1, 1), (2, 2), (3, 3)")
+	testKit.MustExec("create table t2 (c1 int, c2 int)")
+	testKit.MustExec("insert into t2 values (1, 1), (2, 2), (3, 3)")
+
+	// Dump stats delta in one batch.
+	handle := dom.StatsHandle()
+	require.NoError(t, handle.DumpStatsDeltaToKV(true))
+
+	// Check the mysql.stats_meta table.
+	rows := testKit.MustQuery("select modify_count, count, version from mysql.stats_meta order by table_id").Rows()
+	require.Len(t, rows, 2)
+
+	require.Equal(t, "3", rows[0][0])
+	require.Equal(t, "3", rows[0][1])
+	require.Equal(t, "3", rows[1][0])
+	require.Equal(t, "3", rows[1][1])
+	require.Equal(
+		t,
+		rows[0][2],
+		rows[1][2],
+		"The version of two tables should be the same because they are dumped in the same transaction.",
+	)
+}

pkg/statistics/handle/storage/update.go

@@ -74,6 +74,7 @@ func NewDeltaUpdate(tableID int64, delta variable.TableDelta, isLocked bool) *De
 
 // UpdateStatsMeta updates the stats meta for multiple tables.
 // It uses the INSERT INTO ... ON DUPLICATE KEY UPDATE syntax to fill the missing records.
+// Note: Make sure call this function in a transaction.
 func UpdateStatsMeta(
 	ctx context.Context,
 	sctx sessionctx.Context,
@@ -85,14 +86,22 @@ func UpdateStatsMeta(
 	}
 
 	// Separate locked and unlocked updates
-	var lockedValues, unlockedPosValues, unlockedNegValues []string
-	var cacheInvalidateIDs []int64
+	// In most cases, the number of locked tables is small.
+	lockedTableIDs := make([]string, 0, 20)
+	lockedValues := make([]string, 0, 20)
+	// In most cases, the number of unlocked tables is large.
+	unlockedTableIDs := make([]string, 0, len(updates))
+	unlockedPosValues := make([]string, 0, max(len(updates)/2, 1))
+	unlockedNegValues := make([]string, 0, max(len(updates)/2, 1))
+	cacheInvalidateIDs := make([]int64, 0, len(updates))
 
 	for _, update := range updates {
 		if update.IsLocked {
+			lockedTableIDs = append(lockedTableIDs, fmt.Sprintf("%d", update.TableID))
 			lockedValues = append(lockedValues, fmt.Sprintf("(%d, %d, %d, %d)",
 				startTS, update.TableID, update.Delta.Count, update.Delta.Delta))
 		} else {
+			unlockedTableIDs = append(unlockedTableIDs, fmt.Sprintf("%d", update.TableID))
 			if update.Delta.Delta < 0 {
 				unlockedNegValues = append(unlockedNegValues, fmt.Sprintf("(%d, %d, %d, %d)",
 					startTS, update.TableID, update.Delta.Count, -update.Delta.Delta))
@@ -104,6 +113,22 @@ func UpdateStatsMeta(
 		}
 	}
 
+	// Lock the stats_meta and stats_table_locked tables using SELECT FOR UPDATE to prevent write conflicts.
+	// This ensures that we acquire the necessary locks before attempting to update the tables, reducing the likelihood
+	// of encountering lock conflicts during the update process.
+	lockedTableIDsStr := strings.Join(lockedTableIDs, ",")
+	if lockedTableIDsStr != "" {
+		if _, err = statsutil.ExecWithCtx(ctx, sctx, fmt.Sprintf("select * from mysql.stats_table_locked where table_id in (%s) for update", lockedTableIDsStr)); err != nil {
+			return err
+		}
+	}
+
+	unlockedTableIDsStr := strings.Join(unlockedTableIDs, ",")
+	if unlockedTableIDsStr != "" {
+		if _, err = statsutil.ExecWithCtx(ctx, sctx, fmt.Sprintf("select * from mysql.stats_meta where table_id in (%s) for update", unlockedTableIDsStr)); err != nil {
+			return err
+		}
+	}
 	// Execute locked updates
 	if len(lockedValues) > 0 {
 		sql := fmt.Sprintf("insert into mysql.stats_table_locked (version, table_id, modify_count, count) values %s "+

pkg/statistics/handle/usage/BUILD.bazel

@@ -15,6 +15,7 @@ go_library(
         "//pkg/metrics",
         "//pkg/sessionctx",
         "//pkg/sessionctx/variable",
+        "//pkg/statistics/handle/logutil",
         "//pkg/statistics/handle/storage",
         "//pkg/statistics/handle/types",
         "//pkg/statistics/handle/usage/indexusage",
@@ -26,6 +27,7 @@ go_library(
         "//pkg/util/sqlescape",
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",
+        "@org_uber_go_zap//:zap",
     ],
 )