statistics: add Destroy method and handle session recycling (#59546)

0xPoe · web-flow · commit 9f5f53a645e0 · 2025-02-18T13:48:25.000Z
close #59524, close #59560
diff --git a/pkg/bindinfo/global_handle.go b/pkg/bindinfo/global_handle.go
@@ -88,8 +88,7 @@ type GlobalBindingHandle interface {
 
 // globalBindingHandle is used to handle all global sql bind operations.
 type globalBindingHandle struct {
-	sPool util.SessionPool
-
+	sPool        util.DestroyableSessionPool
 	bindingCache BindingCache
 
 	// lastTaskTime records the last update time for the global sql bind cache.
@@ -121,7 +120,7 @@ const (
 )
 
 // NewGlobalBindingHandle creates a new GlobalBindingHandle.
-func NewGlobalBindingHandle(sPool util.SessionPool) GlobalBindingHandle {
+func NewGlobalBindingHandle(sPool util.DestroyableSessionPool) GlobalBindingHandle {
 	h := &globalBindingHandle{sPool: sPool}
 	h.lastUpdateTime.Store(types.ZeroTimestamp)
 	h.bindingCache = newBindCache()
@@ -484,6 +483,9 @@ func (h *globalBindingHandle) callWithSCtx(wrapTxn bool, f func(sctx sessionctx.
 	defer func() {
 		if err == nil { // only recycle when no error
 			h.sPool.Put(resource)
+		} else {
+			// Note: Otherwise, the session will be leaked.
+			h.sPool.Destroy(resource)
 		}
 	}()
 	sctx := resource.(sessionctx.Context)
diff --git a/pkg/bindinfo/global_handle_test.go b/pkg/bindinfo/global_handle_test.go
@@ -607,6 +607,8 @@ func (p *mockSessionPool) Get() (pools.Resource, error) {
 
 func (p *mockSessionPool) Put(pools.Resource) {}
 
+func (p *mockSessionPool) Destroy(pools.Resource) {}
+
 func (p *mockSessionPool) Close() {}
 
 func TestShowBindingDigestField(t *testing.T) {
diff --git a/pkg/bindinfo/tests/cross_db_binding_test.go b/pkg/bindinfo/tests/cross_db_binding_test.go
@@ -342,4 +342,6 @@ func (p *mockSessionPool) Get() (pools.Resource, error) {
 
 func (p *mockSessionPool) Put(pools.Resource) {}
 
+func (p *mockSessionPool) Destroy(pools.Resource) {}
+
 func (p *mockSessionPool) Close() {}
diff --git a/pkg/ddl/notifier/testkit_test.go b/pkg/ddl/notifier/testkit_test.go
@@ -79,6 +79,7 @@ func TestBasicPubSub(t *testing.T) {
 		},
 		nil,
 		nil,
+		nil,
 	)
 
 	n := notifier.NewDDLNotifier(sessionPool, s, 50*time.Millisecond)
@@ -150,6 +151,7 @@ func TestDeliverOrderAndCleanup(t *testing.T) {
 		},
 		nil,
 		nil,
+		nil,
 	)
 	n := notifier.NewDDLNotifier(sessionPool, s, 50*time.Millisecond)
 
@@ -324,6 +326,7 @@ func Test2OwnerForAShortTime(t *testing.T) {
 		},
 		nil,
 		nil,
+		nil,
 	)
 
 	n := notifier.NewDDLNotifier(sessionPool, s, 50*time.Millisecond)
@@ -458,6 +461,7 @@ func TestBeginTwice(t *testing.T) {
 		},
 		nil,
 		nil,
+		nil,
 	)
 
 	n := notifier.NewDDLNotifier(sessionPool, s, 50*time.Millisecond)
@@ -508,6 +512,7 @@ func TestHandlersSeePessimisticTxnError(t *testing.T) {
 		},
 		nil,
 		nil,
+		nil,
 	)
 	n := notifier.NewDDLNotifier(sessionPool, s, 50*time.Millisecond)
 	// Always fails
@@ -561,6 +566,7 @@ func TestCommitFailed(t *testing.T) {
 		},
 		nil,
 		nil,
+		nil,
 	)
 	n := notifier.NewDDLNotifier(sessionPool, s, 50*time.Millisecond)
 	handler := func(_ context.Context, sctx sessionctx.Context, _ *notifier.SchemaChangeEvent) error {
diff --git a/pkg/domain/domain.go b/pkg/domain/domain.go
@@ -158,8 +158,10 @@ type Domain struct {
 	m               syncutil.Mutex
 	SchemaValidator SchemaValidator
 	schemaLease     time.Duration
-	sysSessionPool  util.SessionPool
-	exit            chan struct{}
+	// Note: If you no longer need the session, you must call Destroy to release it.
+	// Otherwise, the session will be leaked. Because there is a strong reference from the domain to the session.
+	sysSessionPool util.DestroyableSessionPool
+	exit           chan struct{}
 	// `etcdClient` must be used when keyspace is not set, or when the logic to each etcd path needs to be separated by keyspace.
 	etcdClient *clientv3.Client
 	// autoidClient is used when there are tables with AUTO_ID_CACHE=1, it is the client to the autoid service.
@@ -1345,6 +1347,10 @@ func NewDomainWithEtcdClient(store kv.Storage, schemaLease time.Duration, statsL
 				})
 				infosync.DeleteInternalSession(r)
 			},
+			func(r pools.Resource) {
+				intest.Assert(r != nil)
+				infosync.DeleteInternalSession(r)
+			},
 		),
 		statsLease:        statsLease,
 		schemaLease:       schemaLease,
@@ -1814,7 +1820,7 @@ func (do *Domain) distTaskFrameworkLoop(ctx context.Context, taskManager *storag
 }
 
 // SysSessionPool returns the system session pool.
-func (do *Domain) SysSessionPool() util.SessionPool {
+func (do *Domain) SysSessionPool() util.DestroyableSessionPool {
 	return do.sysSessionPool
 }
 
@@ -2550,7 +2556,7 @@ func (do *Domain) initStats(ctx context.Context) {
 	}
 	initstats.InitStatsPercentage.Store(100)
 	if err != nil {
-		logutil.BgLogger().Error("init stats info failed", zap.Bool("lite", liteInitStats), zap.Duration("take time", time.Since(t)), zap.Error(err))
+		logutil.ErrVerboseLogger().Error("init stats info failed", zap.Bool("lite", liteInitStats), zap.Duration("take time", time.Since(t)), zap.Error(err))
 	} else {
 		logutil.BgLogger().Info("init stats info time", zap.Bool("lite", liteInitStats), zap.Duration("take time", time.Since(t)))
 	}
@@ -2612,7 +2618,7 @@ func (do *Domain) asyncLoadHistogram() {
 		case <-cleanupTicker.C:
 			err = statsHandle.LoadNeededHistograms(do.InfoSchema())
 			if err != nil {
-				logutil.BgLogger().Warn("load histograms failed", zap.Error(err))
+				logutil.ErrVerboseLogger().Warn("load histograms failed", zap.Error(err))
 			}
 		case <-do.exit:
 			return
diff --git a/pkg/domain/infosync/info.go b/pkg/domain/infosync/info.go
@@ -1339,6 +1339,21 @@ func DeleteInternalSession(se any) {
 	sm.DeleteInternalSession(se)
 }
 
+// ContainsInternalSessionForTest is the entry function for check whether an internal session is in SessionManager.
+// It is only used for test.
+func ContainsInternalSessionForTest(se any) bool {
+	is, err := getGlobalInfoSyncer()
+	if err != nil {
+		return false
+	}
+	sm := is.GetSessionManager()
+	if sm == nil {
+		return false
+	}
+
+	return sm.ContainsInternalSession(se)
+}
+
 // SetEtcdClient is only used for test.
 func SetEtcdClient(etcdCli *clientv3.Client) {
 	is, err := getGlobalInfoSyncer()
diff --git a/pkg/planner/core/rule_collect_plan_stats.go b/pkg/planner/core/rule_collect_plan_stats.go
@@ -259,11 +259,11 @@ func RequestLoadStats(ctx base.PlanContext, neededHistItems []model.StatsLoadIte
 	if err != nil {
 		stmtCtx.IsSyncStatsFailed = true
 		if vardef.StatsLoadPseudoTimeout.Load() {
-			logutil.BgLogger().Warn("RequestLoadStats failed", zap.Error(err))
+			logutil.ErrVerboseLogger().Warn("RequestLoadStats failed", zap.Error(err))
 			stmtCtx.AppendWarning(err)
 			return nil
 		}
-		logutil.BgLogger().Warn("RequestLoadStats failed", zap.Error(err))
+		logutil.ErrVerboseLogger().Warn("RequestLoadStats failed", zap.Error(err))
 		return err
 	}
 	return nil
diff --git a/pkg/server/server.go b/pkg/server/server.go
@@ -1064,6 +1064,14 @@ func (s *Server) StoreInternalSession(se any) {
 	s.sessionMapMutex.Unlock()
 }
 
+// ContainsInternalSession implements SessionManager interface.
+func (s *Server) ContainsInternalSession(se any) bool {
+	s.sessionMapMutex.Lock()
+	defer s.sessionMapMutex.Unlock()
+	_, ok := s.internalSessions[se]
+	return ok
+}
+
 // DeleteInternalSession implements SessionManager interface.
 // @param addr	The address of a session.session struct variable
 func (s *Server) DeleteInternalSession(se any) {
diff --git a/pkg/statistics/handle/bootstrap.go b/pkg/statistics/handle/bootstrap.go
@@ -342,6 +342,9 @@ func (h *Handle) initStatsHistogramsByPaging(is infoschema.InfoSchema, cache sta
 	defer func() {
 		if err == nil { // only recycle when no error
 			h.Pool.SPool().Put(se)
+		} else {
+			// Note: Otherwise, the session will be leaked.
+			h.Pool.SPool().Destroy(se)
 		}
 	}()
 
@@ -468,6 +471,9 @@ func (h *Handle) initStatsTopNByPaging(cache statstypes.StatsCache, task initsta
 	defer func() {
 		if err == nil { // only recycle when no error
 			h.Pool.SPool().Put(se)
+		} else {
+			// Note: Otherwise, the session will be leaked.
+			h.Pool.SPool().Destroy(se)
 		}
 	}()
 	sctx := se.(sessionctx.Context)
@@ -667,6 +673,9 @@ func (h *Handle) initStatsBucketsByPaging(cache statstypes.StatsCache, task init
 	defer func() {
 		if err == nil { // only recycle when no error
 			h.Pool.SPool().Put(se)
+		} else {
+			// Note: Otherwise, the session will be leaked.
+			h.Pool.SPool().Destroy(se)
 		}
 	}()
 	sctx := se.(sessionctx.Context)
diff --git a/pkg/statistics/handle/handle.go b/pkg/statistics/handle/handle.go
@@ -115,7 +115,7 @@ func NewHandle(
 	initStatsCtx sessionctx.Context,
 	lease time.Duration,
 	is infoschema.InfoSchema,
-	pool pkgutil.SessionPool,
+	pool pkgutil.DestroyableSessionPool,
 	tracker sysproctrack.Tracker,
 	ddlNotifier *notifier.DDLNotifier,
 	autoAnalyzeProcIDGetter func() uint64,
diff --git a/pkg/statistics/handle/lockstats/lock_stats.go b/pkg/statistics/handle/lockstats/lock_stats.go
@@ -41,11 +41,11 @@ const (
 
 // statsLockImpl implements the util.StatsLock interface.
 type statsLockImpl struct {
-	pool pkgutil.SessionPool
+	pool pkgutil.DestroyableSessionPool
 }
 
 // NewStatsLock creates a new StatsLock.
-func NewStatsLock(pool pkgutil.SessionPool) types.StatsLock {
+func NewStatsLock(pool pkgutil.DestroyableSessionPool) types.StatsLock {
 	return &statsLockImpl{pool: pool}
 }
 
diff --git a/pkg/statistics/handle/storage/read.go b/pkg/statistics/handle/storage/read.go
@@ -36,6 +36,7 @@ import (
 	"github.com/pingcap/tidb/pkg/statistics/handle/util"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/chunk"
+	"github.com/pingcap/tidb/pkg/util/intest"
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/pingcap/tidb/pkg/util/memory"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
@@ -609,7 +610,16 @@ func LoadNeededHistograms(sctx sessionctx.Context, is infoschema.InfoSchema, sta
 			err = loadNeededIndexHistograms(sctx, is, statsHandle, item.TableItemID, loadFMSketch)
 		}
 		if err != nil {
-			return err
+			statslogutil.StatsLogger().Error("load needed histogram failed",
+				zap.Error(err),
+				zap.Int64("tableID", item.TableID),
+				zap.Int64("histID", item.ID),
+				zap.Bool("isIndex", item.IsIndex),
+				zap.Bool("IsSyncLoadFailed", item.IsSyncLoadFailed),
+				zap.Bool("fullLoad", item.FullLoad),
+			)
+			intest.Assert(false, "load needed histogram failed")
+			return errors.Trace(err)
 		}
 	}
 	return nil
diff --git a/pkg/statistics/handle/syncload/stats_syncload.go b/pkg/statistics/handle/syncload/stats_syncload.go
@@ -303,6 +303,9 @@ func (s *statsSyncLoad) handleOneItemTask(task *statstypes.NeededItemTask) (err
 		if err == nil { // only recycle when no error
 			sctx.GetSessionVars().StmtCtx.Priority = mysql.NoPriority
 			s.statsHandle.SPool().Put(se)
+		} else {
+			// Note: Otherwise, the session will be leaked.
+			s.statsHandle.SPool().Destroy(se)
 		}
 	}()
 	var skipTypes map[string]struct{}
diff --git a/pkg/statistics/handle/util/BUILD.bazel b/pkg/statistics/handle/util/BUILD.bazel
@@ -45,7 +45,10 @@ go_test(
     flaky = True,
     deps = [
         ":util",
+        "//pkg/domain/infosync",
+        "//pkg/sessionctx",
         "//pkg/testkit",
+        "@com_github_pingcap_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],
 )
diff --git a/pkg/statistics/handle/util/pool.go b/pkg/statistics/handle/util/pool.go
@@ -28,7 +28,7 @@ type Pool interface {
 	GPool() *gp.Pool
 
 	// SPool returns the session pool.
-	SPool() util.SessionPool
+	SPool() util.DestroyableSessionPool
 
 	// Close closes the goroutine pool.
 	Close()
@@ -39,11 +39,11 @@ var _ Pool = (*pool)(nil)
 type pool struct {
 	// This gpool is used to reuse goroutine in the mergeGlobalStatsTopN.
 	gpool *gp.Pool
-	pool  util.SessionPool
+	pool  util.DestroyableSessionPool
 }
 
 // NewPool creates a new Pool.
-func NewPool(p util.SessionPool) Pool {
+func NewPool(p util.DestroyableSessionPool) Pool {
 	return &pool{
 		gpool: gp.New(math.MaxInt16, time.Minute),
 		pool:  p,
@@ -56,7 +56,7 @@ func (p *pool) GPool() *gp.Pool {
 }
 
 // SPool returns the session pool.
-func (p *pool) SPool() util.SessionPool {
+func (p *pool) SPool() util.DestroyableSessionPool {
 	return p.pool
 }
 
diff --git a/pkg/statistics/handle/util/util.go b/pkg/statistics/handle/util/util.go
@@ -76,20 +76,23 @@ var (
 )
 
 // CallWithSCtx allocates a sctx from the pool and call the f().
-func CallWithSCtx(pool util.SessionPool, f func(sctx sessionctx.Context) error, flags ...int) (err error) {
+func CallWithSCtx(pool util.DestroyableSessionPool, f func(sctx sessionctx.Context) error, flags ...int) (err error) {
 	defer util.Recover(metrics.LabelStats, "CallWithSCtx", nil, false)
 	se, err := pool.Get()
 	if err != nil {
-		return err
+		return errors.Trace(err)
 	}
 	defer func() {
 		if err == nil { // only recycle when no error
 			pool.Put(se)
+		} else {
+			// Note: Otherwise, the session will be leaked.
+			pool.Destroy(se)
 		}
 	}()
 	sctx := se.(sessionctx.Context)
 	if err := UpdateSCtxVarsForStats(sctx); err != nil { // update stats variables automatically
-		return err
+		return errors.Trace(err)
 	}
 
 	wrapTxn := false
@@ -103,7 +106,7 @@ func CallWithSCtx(pool util.SessionPool, f func(sctx sessionctx.Context) error,
 	} else {
 		err = f(sctx)
 	}
-	return err
+	return errors.Trace(err)
 }
 
 // UpdateSCtxVarsForStats updates all necessary variables that may affect the behavior of statistics.
@@ -184,7 +187,7 @@ func UpdateSCtxVarsForStats(sctx sessionctx.Context) error {
 }
 
 // GetCurrentPruneMode returns the current latest partitioning table prune mode.
-func GetCurrentPruneMode(pool util.SessionPool) (mode string, err error) {
+func GetCurrentPruneMode(pool util.DestroyableSessionPool) (mode string, err error) {
 	err = CallWithSCtx(pool, func(sctx sessionctx.Context) error {
 		mode = sctx.GetSessionVars().PartitionPruneMode.Load()
 		return nil
diff --git a/pkg/statistics/handle/util/util_test.go b/pkg/statistics/handle/util/util_test.go
@@ -17,6 +17,9 @@ package util_test
 import (
 	"testing"
 
+	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/pkg/domain/infosync"
+	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/statistics/handle/util"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/stretchr/testify/require"
@@ -55,3 +58,17 @@ func TestIsSpecialGlobalIndex(t *testing.T) {
 	}
 	require.Equal(t, cnt, len(tblInfo.Indices))
 }
+
+func TestCallSCtxFailed(t *testing.T) {
+	_, dom := testkit.CreateMockStoreAndDomain(t)
+
+	var sctxWithFailure sessionctx.Context
+	err := util.CallWithSCtx(dom.StatsHandle().SPool(), func(sctx sessionctx.Context) error {
+		sctxWithFailure = sctx
+		return errors.New("simulated error")
+	})
+	require.Error(t, err)
+	require.Equal(t, "simulated error", err.Error())
+	notReleased := infosync.ContainsInternalSessionForTest(sctxWithFailure)
+	require.False(t, notReleased)
+}
diff --git a/pkg/testkit/mocksessionmanager.go b/pkg/testkit/mocksessionmanager.go
diff --git a/pkg/util/processinfo.go b/pkg/util/processinfo.go
diff --git a/pkg/util/session_pool.go b/pkg/util/session_pool.go
diff --git a/pkg/util/session_pool_test.go b/pkg/util/session_pool_test.go

Original file line number	Diff line number	Diff line change
`@@ -259,11 +259,11 @@ func RequestLoadStats(ctx base.PlanContext, neededHistItems []model.StatsLoadIte`
`259`	`259`	`if err != nil {`
`260`	`260`	`stmtCtx.IsSyncStatsFailed = true`
`261`	`261`	`if vardef.StatsLoadPseudoTimeout.Load() {`
`262`		`- logutil.BgLogger().Warn("RequestLoadStats failed", zap.Error(err))`
	`262`	`+ logutil.ErrVerboseLogger().Warn("RequestLoadStats failed", zap.Error(err))`
`263`	`263`	`stmtCtx.AppendWarning(err)`
`264`	`264`	`return nil`
`265`	`265`	`}`
`266`		`- logutil.BgLogger().Warn("RequestLoadStats failed", zap.Error(err))`
	`266`	`+ logutil.ErrVerboseLogger().Warn("RequestLoadStats failed", zap.Error(err))`
`267`	`267`	`return err`
`268`	`268`	`}`
`269`	`269`	`return nil`