From fbc3c7e58a8bd077d005627ffce2c313518b00a9 Mon Sep 17 00:00:00 2001
From: Weizhen Wang <wangweizhen@pingcap.com>
Date: Tue, 27 Feb 2024 18:47:08 +0800
Subject: [PATCH] statistics: fix panic when to enable force-init-stats
 (#51357)

close pingcap/tidb#51358
---
 pkg/statistics/handle/bootstrap.go            |  2 +-
 pkg/statistics/handle/cache/BUILD.bazel       |  1 +
 pkg/statistics/handle/cache/statscache.go     |  4 +++
 .../handle/handletest/statstest/BUILD.bazel   |  3 +-
 .../handle/handletest/statstest/stats_test.go | 28 +++++++++++++++++++
 5 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/pkg/statistics/handle/bootstrap.go b/pkg/statistics/handle/bootstrap.go
index 676ae05f70..f05ce9b972 100644
--- a/pkg/statistics/handle/bootstrap.go
+++ b/pkg/statistics/handle/bootstrap.go
@@ -165,10 +165,10 @@ func (h *Handle) initStatsHistograms4Chunk(is infoschema.InfoSchema, cache stats
 	for row := iter.Begin(); row != iter.End(); row = iter.Next() {
 		tblID, statsVer := row.GetInt64(0), row.GetInt64(8)
 		table, ok := cache.Get(tblID)
-		table = table.Copy()
 		if !ok {
 			continue
 		}
+		table = table.Copy()
 		id, ndv, nullCount, version, totColSize := row.GetInt64(2), row.GetInt64(3), row.GetInt64(5), row.GetUint64(4), row.GetInt64(7)
 		lastAnalyzePos := row.GetDatum(11, types.NewFieldType(mysql.TypeBlob))
 		tbl, _ := h.TableInfoByID(is, table.PhysicalID)
diff --git a/pkg/statistics/handle/cache/BUILD.bazel b/pkg/statistics/handle/cache/BUILD.bazel
index 1a3d418bae..ee1c0d9a03 100644
--- a/pkg/statistics/handle/cache/BUILD.bazel
+++ b/pkg/statistics/handle/cache/BUILD.bazel
@@ -30,6 +30,7 @@ go_library(
         "//pkg/util/logutil",
         "//pkg/util/syncutil",
         "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_failpoint//:failpoint",
         "@org_uber_go_zap//:zap",
     ],
 )
diff --git a/pkg/statistics/handle/cache/statscache.go b/pkg/statistics/handle/cache/statscache.go
index ed4804c92b..b8a1b1e993 100644
--- a/pkg/statistics/handle/cache/statscache.go
+++ b/pkg/statistics/handle/cache/statscache.go
@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/infoschema"
 	tidbmetrics "github.com/pingcap/tidb/pkg/metrics"
@@ -204,6 +205,9 @@ func (s *StatsCacheImpl) MemConsumed() (size int64) {
 
 // Get returns the specified table's stats.
 func (s *StatsCacheImpl) Get(tableID int64) (*statistics.Table, bool) {
+	failpoint.Inject("StatsCacheGetNil", func() {
+		failpoint.Return(nil, false)
+	})
 	return s.Load().Get(tableID)
 }
 
diff --git a/pkg/statistics/handle/handletest/statstest/BUILD.bazel b/pkg/statistics/handle/handletest/statstest/BUILD.bazel
index 9f74a3d998..e14d55030f 100644
--- a/pkg/statistics/handle/handletest/statstest/BUILD.bazel
+++ b/pkg/statistics/handle/handletest/statstest/BUILD.bazel
@@ -9,13 +9,14 @@ go_test(
     ],
     flaky = True,
     race = "on",
-    shard_count = 8,
+    shard_count = 9,
     deps = [
         "//pkg/config",
         "//pkg/parser/model",
         "//pkg/statistics/handle/internal",
         "//pkg/testkit",
         "//pkg/testkit/testsetup",
+        "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_stretchr_testify//require",
         "@org_uber_go_goleak//:goleak",
     ],
diff --git a/pkg/statistics/handle/handletest/statstest/stats_test.go b/pkg/statistics/handle/handletest/statstest/stats_test.go
index caa0167589..7c710becad 100644
--- a/pkg/statistics/handle/handletest/statstest/stats_test.go
+++ b/pkg/statistics/handle/handletest/statstest/stats_test.go
@@ -19,6 +19,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/statistics/handle/internal"
@@ -237,6 +238,33 @@ func TestInitStats(t *testing.T) {
 	h.SetLease(0)
 }
 
+func TestInitStats51358(t *testing.T) {
+	originValue := config.GetGlobalConfig().Performance.LiteInitStats
+	defer func() {
+		config.GetGlobalConfig().Performance.LiteInitStats = originValue
+	}()
+	config.GetGlobalConfig().Performance.LiteInitStats = false
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	testKit := testkit.NewTestKit(t, store)
+	testKit.MustExec("use test")
+	testKit.MustExec("set @@session.tidb_analyze_version = 1")
+	testKit.MustExec("create table t(a int, b int, c int, primary key(a), key idx(b))")
+	testKit.MustExec("insert into t values (1,1,1),(2,2,2),(3,3,3),(4,4,4),(5,5,5),(6,7,8)")
+	testKit.MustExec("analyze table t")
+	h := dom.StatsHandle()
+	is := dom.InfoSchema()
+	// `Update` will not use load by need strategy when `Lease` is 0, and `InitStats` is only called when
+	// `Lease` is not 0, so here we just change it.
+	h.SetLease(time.Millisecond)
+
+	h.Clear()
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/statistics/handle/cache/StatsCacheGetNil", "return()"))
+	defer func() {
+		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/statistics/handle/cache/StatsCacheGetNil"))
+	}()
+	require.NoError(t, h.InitStats(is))
+}
+
 func TestInitStatsVer2(t *testing.T) {
 	originValue := config.GetGlobalConfig().Performance.LiteInitStats
 	defer func() {