From a82339821ccff507b08bb73338b6f8d4256476c5 Mon Sep 17 00:00:00 2001
From: ti-srebot <66930949+ti-srebot@users.noreply.github.com>
Date: Mon, 19 Jul 2021 15:35:34 +0800
Subject: [PATCH] session: Fix txn.onStmtEnd is not called when some special
 error occurs in ExecutePreparedStmt (#26056) (#26080)

---
 session/session.go | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/session/session.go b/session/session.go
index 7ef59197fb3ca..610380b71151d 100644
--- a/session/session.go
+++ b/session/session.go
@@ -1908,6 +1908,7 @@ func (s *session) ExecutePreparedStmt(ctx context.Context, stmtID uint32, args [
 		return nil, err
 	}
 	s.txn.onStmtStart(preparedStmt.SQLDigest.String())
+	defer s.txn.onStmtEnd()
 	var is infoschema.InfoSchema
 	var snapshotTS uint64
 	if preparedStmt.ForUpdateRead {
@@ -1924,14 +1925,10 @@ func (s *session) ExecutePreparedStmt(ctx context.Context, stmtID uint32, args [
 	} else {
 		is = s.GetInfoSchema().(infoschema.InfoSchema)
 	}
-	var rs sqlexec.RecordSet
 	if ok {
-		rs, err = s.cachedPlanExec(ctx, is, snapshotTS, stmtID, preparedStmt, args)
-	} else {
-		rs, err = s.preparedStmtExec(ctx, is, snapshotTS, stmtID, preparedStmt, args)
+		return s.cachedPlanExec(ctx, is, snapshotTS, stmtID, preparedStmt, args)
 	}
-	s.txn.onStmtEnd()
-	return rs, err
+	return s.preparedStmtExec(ctx, is, snapshotTS, stmtID, preparedStmt, args)
 }
 
 func (s *session) DropPreparedStmt(stmtID uint32) error {