session: reset affected rows in retry. (#2949)

session.go

@@ -353,6 +353,7 @@ func (s *session) retry(maxCnt int) error {
 				log.Warnf("[%d] Retry [%d] query [%d]", connID, retryCnt, i)
 			}
 			s.sessionVars.StmtCtx = sr.stmtCtx
+			s.sessionVars.StmtCtx.ResetForRetry()
 			_, err = st.Exec(s)
 			if err != nil {
 				break

session_test.go

@@ -2692,3 +2692,22 @@ func (s *testSessionSuite) TestRetryCleanTxn(c *C) {
 	c.Assert(se.Txn(), IsNil)
 	c.Assert(se.sessionVars.InTxn(), IsFalse)
 }
+
+func (s *testSessionSuite) TestRetryResetStmtCtx(c *C) {
+	defer testleak.AfterTest(c)()
+	dbName := "test_retry_reset_stmtctx"
+	se := newSession(c, s.store, dbName).(*session)
+	se.Execute("create table retrytxn (a int unique, b int)")
+	_, err := se.Execute("insert retrytxn values (1, 1)")
+	c.Assert(err, IsNil)
+	se.Execute("begin")
+	se.Execute("update retrytxn set b = b + 1 where a = 1")
+
+	// Make retryable error.
+	se2 := newSession(c, s.store, dbName)
+	se2.Execute("update retrytxn set b = b + 1 where a = 1")
+
+	err = se.CommitTxn()
+	c.Assert(err, IsNil)
+	c.Assert(se.AffectedRows(), Equals, uint64(1))
+}

sessionctx/variable/session.go

@@ -368,3 +368,12 @@ func (sc *StatementContext) HandleTruncate(err error) error {
 	}
 	return err
 }
+
+// ResetForRetry resets the changed states during execution.
+func (sc *StatementContext) ResetForRetry() {
+	sc.mu.Lock()
+	sc.mu.affectedRows = 0
+	sc.mu.foundRows = 0
+	sc.mu.warnings = nil
+	sc.mu.Unlock()
+}

sessionctx/variable/session_test.go

@@ -44,4 +44,9 @@ func (*testSessionSuite) TestSession(c *C) {
 	// For last insert id
 	ctx.GetSessionVars().SetLastInsertID(1)
 	c.Assert(ctx.GetSessionVars().LastInsertID, Equals, uint64(1))
+
+	ss.ResetForRetry()
+	c.Assert(ss.AffectedRows(), Equals, uint64(0))
+	c.Assert(ss.FoundRows(), Equals, uint64(0))
+	c.Assert(ss.WarningCount(), Equals, uint16(0))
 }