Merge branch 'master' into cherry-pick-36089

server/driver_tidb.go

@@ -335,11 +335,11 @@ func (tc *TiDBContext) EncodeSessionStates(ctx context.Context, sctx sessionctx.
 		// Bound params are sent by CMD_STMT_SEND_LONG_DATA, the proxy can wait for COM_STMT_EXECUTE.
 		for _, boundParam := range stmt.BoundParams() {
 			if boundParam != nil {
-				return session.ErrCannotMigrateSession.GenWithStackByArgs("prepared statements have bound params")
+				return sessionstates.ErrCannotMigrateSession.GenWithStackByArgs("prepared statements have bound params")
 			}
 		}
 		if rs := stmt.GetResultSet(); rs != nil && !rs.IsClosed() {
-			return session.ErrCannotMigrateSession.GenWithStackByArgs("prepared statements have open result sets")
+			return sessionstates.ErrCannotMigrateSession.GenWithStackByArgs("prepared statements have open result sets")
 		}
 		preparedStmtInfo.ParamTypes = stmt.GetParamsType()
 	}

session/session.go

@@ -3367,23 +3367,23 @@ func (s *session) EncodeSessionStates(ctx context.Context, sctx sessionctx.Conte
 	valid := s.txn.Valid()
 	s.txn.mu.Unlock()
 	if valid {
-		return ErrCannotMigrateSession.GenWithStackByArgs("session has an active transaction")
+		return sessionstates.ErrCannotMigrateSession.GenWithStackByArgs("session has an active transaction")
 	}
 	// Data in local temporary tables is hard to encode, so we do not support it.
 	// Check temporary tables here to avoid circle dependency.
 	if s.sessionVars.LocalTemporaryTables != nil {
 		localTempTables := s.sessionVars.LocalTemporaryTables.(*infoschema.LocalTemporaryTables)
 		if localTempTables.Count() > 0 {
-			return ErrCannotMigrateSession.GenWithStackByArgs("session has local temporary tables")
+			return sessionstates.ErrCannotMigrateSession.GenWithStackByArgs("session has local temporary tables")
 		}
 	}
 	// The advisory locks will be released when the session is closed.
 	if len(s.advisoryLocks) > 0 {
-		return ErrCannotMigrateSession.GenWithStackByArgs("session has advisory locks")
+		return sessionstates.ErrCannotMigrateSession.GenWithStackByArgs("session has advisory locks")
 	}
 	// The TableInfo stores session ID and server ID, so the session cannot be migrated.
 	if len(s.lockedTables) > 0 {
-		return ErrCannotMigrateSession.GenWithStackByArgs("session has locked tables")
+		return sessionstates.ErrCannotMigrateSession.GenWithStackByArgs("session has locked tables")
 	}
 
 	if err := s.sessionVars.EncodeSessionStates(ctx, sessionStates); err != nil {

session/tidb.go

@@ -378,6 +378,5 @@ func ResultSetToStringSlice(ctx context.Context, s Session, rs sqlexec.RecordSet
 
 // Session errors.
 var (
-	ErrForUpdateCantRetry   = dbterror.ClassSession.NewStd(errno.ErrForUpdateCantRetry)
-	ErrCannotMigrateSession = dbterror.ClassSession.NewStd(errno.ErrCannotMigrateSession)
+	ErrForUpdateCantRetry = dbterror.ClassSession.NewStd(errno.ErrForUpdateCantRetry)
 )

sessionctx/sessionstates/BUILD.bazel

@@ -2,20 +2,33 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "sessionstates",
-    srcs = ["session_states.go"],
+    srcs = [
+        "session_states.go",
+        "session_token.go",
+    ],
     importpath = "github.com/pingcap/tidb/sessionctx/sessionstates",
     visibility = ["//visibility:public"],
     deps = [
+        "//errno",
         "//parser/types",
         "//sessionctx/stmtctx",
         "//types",
+        "//util/dbterror",
+        "//util/logutil",
+        "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_failpoint//:failpoint",
+        "@org_uber_go_zap//:zap",
     ],
 )
 
 go_test(
     name = "sessionstates_test",
     timeout = "short",
-    srcs = ["session_states_test.go"],
+    srcs = [
+        "session_states_test.go",
+        "session_token_test.go",
+    ],
+    embed = [":sessionstates"],
     deps = [
         "//config",
         "//errno",
@@ -25,8 +38,10 @@ go_test(
         "//sessionctx/variable",
         "//testkit",
         "//types",
+        "//util",
         "//util/sem",
         "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_stretchr_testify//require",
     ],
 )

sessionctx/sessionstates/session_states.go

@@ -17,14 +17,21 @@ package sessionstates
 import (
 	"time"
 
+	"github.com/pingcap/tidb/errno"
 	ptypes "github.com/pingcap/tidb/parser/types"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util/dbterror"
 )
 
 // SessionStateType is the type of session states.
 type SessionStateType int
 
+var (
+	// ErrCannotMigrateSession indicates the session cannot be migrated.
+	ErrCannotMigrateSession = dbterror.ClassSession.NewStd(errno.ErrCannotMigrateSession)
+)
+
 // These enums represents the types of session state handlers.
 const (
 	// StatePrepareStmt represents prepared statements.