statistics: add for update into write to avoid inaccurate results

pkg/statistics/handle/ddl/BUILD.bazel

@@ -33,14 +33,17 @@ go_test(
     timeout = "short",
     srcs = ["ddl_test.go"],
     flaky = True,
-    shard_count = 18,
+    shard_count = 19,
     deps = [
+        ":ddl",
         "//pkg/ddl/notifier",
         "//pkg/meta/model",
         "//pkg/parser/model",
         "//pkg/planner/cardinality",
+        "//pkg/statistics/handle/storage",
         "//pkg/testkit",
         "//pkg/types",
+        "//pkg/util",
         "//pkg/util/mock",
         "@com_github_stretchr_testify//require",
     ],

pkg/statistics/handle/ddl/ddl.go

@@ -186,6 +186,15 @@ func (h *ddlHandlerImpl) HandleDDLEvent(s *notifier.SchemaChangeEvent) error {
 	return nil
 }
 
+// UpdateStatsWithCountDeltaAndModifyCountDeltaForTest updates the global stats with the given count delta and modify count delta.
+func UpdateStatsWithCountDeltaAndModifyCountDeltaForTest(
+	sctx sessionctx.Context,
+	tableID int64,
+	countDelta, modifyCountDelta int64,
+) error {
+	return updateStatsWithCountDeltaAndModifyCountDelta(sctx, tableID, countDelta, modifyCountDelta)
+}
+
 // updateStatsWithCountDeltaAndModifyCountDelta updates
 // the global stats with the given count delta and modify count delta.
 // Only used by some special DDLs, such as exchange partition.
@@ -226,7 +235,7 @@ func updateStatsWithCountDeltaAndModifyCountDelta(
 	}
 
 	// Because count can not be negative, so we need to get the current and calculate the delta.
-	count, modifyCount, isNull, err := storage.StatsMetaCountAndModifyCount(sctx, tableID)
+	count, modifyCount, isNull, err := storage.StatsMetaCountAndModifyCountForUpdate(sctx, tableID)
 	if err != nil {
 		return err
 	}

pkg/statistics/handle/ddl/ddl_test.go

@@ -23,8 +23,11 @@ import (
 	"github.com/pingcap/tidb/pkg/meta/model"
 	pmodel "github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/planner/cardinality"
+	"github.com/pingcap/tidb/pkg/statistics/handle/ddl"
+	"github.com/pingcap/tidb/pkg/statistics/handle/storage"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/types"
+	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/mock"
 	"github.com/stretchr/testify/require"
 )
@@ -1311,3 +1314,29 @@ func findEvent(eventCh <-chan *notifier.SchemaChangeEvent, eventType model.Actio
 		}
 	}
 }
+
+func TestExchangePartition(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("use test")
+	tk.MustExec("create table t (c1 int)")
+	is := dom.InfoSchema()
+	tbl, err := is.TableByName(context.Background(), pmodel.NewCIStr("test"), pmodel.NewCIStr("t"))
+	require.NoError(t, err)
+	var wg util.WaitGroupWrapper
+	for i := 0; i < 20; i++ {
+		tk1 := testkit.NewTestKit(t, store)
+		wg.Run(func() {
+			tk1.MustExec("begin")
+			ddl.UpdateStatsWithCountDeltaAndModifyCountDeltaForTest(tk1.Session(), tbl.Meta().ID, 10, 10)
+			tk1.MustExec("commit")
+		})
+	}
+	wg.Wait()
+	count, modifyCount, isNull, err := storage.StatsMetaCountAndModifyCount(tk.Session(), tbl.Meta().ID)
+	require.NoError(t, err)
+	require.False(t, isNull)
+	require.Equal(t, int64(200), count)
+	require.Equal(t, int64(200), modifyCount)
+}

pkg/statistics/handle/storage/read.go

@@ -43,7 +43,20 @@ import (
 
 // StatsMetaCountAndModifyCount reads count and modify_count for the given table from mysql.stats_meta.
 func StatsMetaCountAndModifyCount(sctx sessionctx.Context, tableID int64) (count, modifyCount int64, isNull bool, err error) {
-	rows, _, err := util.ExecRows(sctx, "select count, modify_count from mysql.stats_meta where table_id = %?", tableID)
+	return statsMetaCountAndModifyCount(sctx, tableID, false)
+}
+
+// StatsMetaCountAndModifyCountForUpdate reads count and modify_count for the given table from mysql.stats_meta with lock.
+func StatsMetaCountAndModifyCountForUpdate(sctx sessionctx.Context, tableID int64) (count, modifyCount int64, isNull bool, err error) {
+	return statsMetaCountAndModifyCount(sctx, tableID, false)
+}
+
+func statsMetaCountAndModifyCount(sctx sessionctx.Context, tableID int64, forUpdate bool) (count, modifyCount int64, isNull bool, err error) {
+	sql := "select count, modify_count from mysql.stats_meta where table_id = %?"
+	if forUpdate {
+		sql += " for update"
+	}
+	rows, _, err := util.ExecRows(sctx, sql, tableID)
 	if err != nil {
 		return 0, 0, false, err
 	}