executor: fix plan replayer for sql file input wrongly fetch startTS …

…after `OnTxnEnd` (#46201)

close #46197

executor/plan_replayer.go

@@ -64,6 +64,7 @@ type PlanReplayerDumpInfo struct {
 	ExecStmts         []ast.StmtNode
 	Analyze           bool
 	HistoricalStatsTS uint64
+	StartTS           uint64
 	Path              string
 	File              *os.File
 	FileName          string
@@ -86,6 +87,15 @@ func (e *PlanReplayerExec) Next(ctx context.Context, req *chunk.Chunk) error {
 	if err != nil {
 		return err
 	}
+	// Note:
+	// For the dumping for SQL file case (len(e.DumpInfo.Path) > 0), the DumpInfo.dump() is called in
+	// handleFileTransInConn(), which is after TxnManager.OnTxnEnd(), where we can't access the TxnManager anymore.
+	// So we must fetch the startTS now.
+	startTS, err := sessiontxn.GetTxnManager(e.Ctx()).GetStmtReadTS()
+	if err != nil {
+		return err
+	}
+	e.DumpInfo.StartTS = startTS
 	if len(e.DumpInfo.Path) > 0 {
 		err = e.prepare()
 		if err != nil {
@@ -165,12 +175,8 @@ func (e *PlanReplayerExec) createFile() error {
 func (e *PlanReplayerDumpInfo) dump(ctx context.Context) (err error) {
 	fileName := e.FileName
 	zf := e.File
-	startTS, err := sessiontxn.GetTxnManager(e.ctx).GetStmtReadTS()
-	if err != nil {
-		return err
-	}
 	task := &domain.PlanReplayerDumpTask{
-		StartTS:           startTS,
+		StartTS:           e.StartTS,
 		FileName:          fileName,
 		Zf:                zf,
 		SessionVars:       e.ctx.GetSessionVars(),

server/BUILD.bazel

@@ -137,7 +137,7 @@ go_test(
     data = glob(["testdata/**"]),
     embed = [":server"],
     flaky = True,
-    shard_count = 46,
+    shard_count = 47,
     deps = [
         "//config",
         "//domain",

server/server_test.go

@@ -24,7 +24,9 @@ import (
 	"testing"
 
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/server/internal"
+	"github.com/pingcap/tidb/server/internal/testutil"
 	"github.com/pingcap/tidb/server/internal/util"
 	"github.com/pingcap/tidb/testkit"
 	"github.com/pingcap/tidb/testkit/testdata"
@@ -103,3 +105,48 @@ func TestOptimizerDebugTrace(t *testing.T) {
 		require.NoError(t, os.Remove(filepath.Join(replayer.GetPlanReplayerDirName(), task.FileName)))
 	}
 }
+
+func TestIssue46197(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tidbdrv := NewTiDBDriver(store)
+	cfg := util.NewTestConfig()
+	cfg.Port, cfg.Status.StatusPort = 0, 0
+	cfg.Status.ReportStatus = false
+	server, err := NewServer(cfg, tidbdrv)
+	require.NoError(t, err)
+	defer server.Close()
+
+	// Mock the content of the SQL file in PacketIO buffer.
+	// First 4 bytes are the header, followed by the actual content.
+	// This acts like we are sending "select * from t1;" from the client when tidb requests the "a.txt" file.
+	var inBuffer bytes.Buffer
+	_, err = inBuffer.Write([]byte{0x11, 0x00, 0x00, 0x01})
+	require.NoError(t, err)
+	_, err = inBuffer.Write([]byte("select * from t1;"))
+	require.NoError(t, err)
+
+	// clientConn setup
+	brc := util.NewBufferedReadConn(&testutil.BytesConn{Buffer: inBuffer})
+	pkt := internal.NewPacketIO(brc)
+	pkt.SetBufWriter(bufio.NewWriter(bytes.NewBuffer(nil)))
+	cc := &clientConn{
+		server:     server,
+		alloc:      arena.NewAllocator(1024),
+		chunkAlloc: chunk.NewAllocator(),
+		pkt:        pkt,
+		capability: mysql.ClientLocalFiles,
+	}
+	ctx := context.Background()
+	cc.SetCtx(&TiDBContext{Session: tk.Session(), stmts: make(map[int]*TiDBStatement)})
+
+	tk.MustExec("use test")
+	tk.MustExec("create table t1 (a int, b int)")
+
+	// 3 is mysql.ComQuery, followed by the SQL text.
+	require.NoError(t, cc.dispatch(ctx, []byte("\u0003plan replayer dump explain 'a.txt'")))
+
+	// clean up
+	path := testdata.ConvertRowsToStrings(tk.MustQuery("select @@tidb_last_plan_replayer_token").Rows())
+	require.NoError(t, os.Remove(filepath.Join(replayer.GetPlanReplayerDirName(), path[0])))
+}