statistics: release session met error to avoid possible leak

Author: winoros

pkg/domain/infosync/info.go

@@ -1339,6 +1339,19 @@ func DeleteInternalSession(se any) {
 	sm.DeleteInternalSession(se)
 }
 
+// HasInternalSessionForTest is the entry function for check whether an internal session is stored in SessionManager.
+func HasInternalSessionForTest(se any) bool {
+	is, err := getGlobalInfoSyncer()
+	if err != nil {
+		return false
+	}
+	sm := is.GetSessionManager()
+	if sm == nil {
+		return false
+	}
+	return sm.HasInternalSessionForTest(se)
+}
+
 // SetEtcdClient is only used for test.
 func SetEtcdClient(etcdCli *clientv3.Client) {
 	is, err := getGlobalInfoSyncer()

pkg/server/server.go

@@ -1064,6 +1064,14 @@ func (s *Server) StoreInternalSession(se any) {
 	s.sessionMapMutex.Unlock()
 }
 
+// HasInternalSessionForTest implements SessionManager interface.
+func (s *Server) HasInternalSessionForTest(se any) bool {
+	s.sessionMapMutex.Lock()
+	_, ok := s.internalSessions[se]
+	s.sessionMapMutex.Unlock()
+	return ok
+}
+
 // DeleteInternalSession implements SessionManager interface.
 // @param addr	The address of a session.session struct variable
 func (s *Server) DeleteInternalSession(se any) {

pkg/statistics/handle/util/BUILD.bazel

@@ -12,6 +12,7 @@ go_library(
     importpath = "github.com/pingcap/tidb/pkg/statistics/handle/util",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/domain/infosync",
         "//pkg/infoschema",
         "//pkg/kv",
         "//pkg/meta/model",
@@ -45,7 +46,10 @@ go_test(
     flaky = True,
     deps = [
         ":util",
+        "//pkg/domain/infosync",
+        "//pkg/sessionctx",
         "//pkg/testkit",
+        "@com_github_pingcap_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/statistics/handle/util/util.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/tidb/pkg/domain/infosync"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/pingcap/tidb/pkg/metrics"
@@ -85,6 +86,8 @@ func CallWithSCtx(pool util.SessionPool, f func(sctx sessionctx.Context) error,
 	defer func() {
 		if err == nil { // only recycle when no error
 			pool.Put(se)
+		} else {
+			infosync.DeleteInternalSession(se)
 		}
 	}()
 	sctx := se.(sessionctx.Context)

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")
+	})
+	notReleased := infosync.HasInternalSessionForTest(sctxWithFailure)
+	require.Error(t, err)
+	require.Equal(t, "simulated error", err.Error())
+	require.False(t, notReleased)
+}

pkg/testkit/mocksessionmanager.go

@@ -137,6 +137,17 @@ func (msm *MockSessionManager) StoreInternalSession(s any) {
 	msm.mu.Unlock()
 }
 
+// HasInternalSessionForTest is to check whether the internal session is stored in the SessionManager
+func (msm *MockSessionManager) HasInternalSessionForTest(se any) bool {
+	msm.mu.Lock()
+	defer msm.mu.Unlock()
+	if msm.internalSessions == nil {
+		return false
+	}
+	_, ok := msm.internalSessions[se]
+	return ok
+}
+
 // DeleteInternalSession is to delete the internal session pointer from the map in the SessionManager
 func (msm *MockSessionManager) DeleteInternalSession(s any) {
 	msm.mu.Lock()

pkg/util/processinfo.go

@@ -218,6 +218,8 @@ type SessionManager interface {
 	ServerID() uint64
 	// StoreInternalSession puts the internal session pointer to the map in the SessionManager.
 	StoreInternalSession(se any)
+	// HasInternalSessionForTest checks whether the internal session is stored in the SessionManager.
+	HasInternalSessionForTest(se any) bool
 	// DeleteInternalSession deletes the internal session pointer from the map in the SessionManager.
 	DeleteInternalSession(se any)
 	// GetInternalSessionStartTSList gets all startTS of every transactions running in the current internal sessions.