ttl: rollback the scan trasnaction even if the BEGIN failed. (pingcap#58943)

close pingcap#58900

Author: YangKeao

pkg/ttl/session/session.go

@@ -111,6 +111,20 @@ func (s *session) ExecuteSQL(ctx context.Context, sql string, args ...any) ([]ch
 
 // RunInTxn executes the specified function in a txn
 func (s *session) RunInTxn(ctx context.Context, fn func() error, txnMode TxnMode) (err error) {
+	success := false
+	defer func() {
+		// Always try to `ROLLBACK` the transaction even if only the `BEGIN` fails. If the `BEGIN` is killed
+		// after it runs the first `Next`, the transaction is already active and needs to be `ROLLBACK`ed.
+		if !success {
+			// For now, the "ROLLBACK" can execute successfully even when the context has already been cancelled.
+			// Using another timeout context to avoid that this behavior will be changed in the future.
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			_, rollbackErr := s.ExecuteSQL(ctx, "ROLLBACK")
+			terror.Log(rollbackErr)
+			cancel()
+		}
+	}()
+
 	tracer := metrics.PhaseTracerFromCtx(ctx)
 	defer tracer.EnterPhase(tracer.Phase())
 
@@ -129,14 +143,6 @@ func (s *session) RunInTxn(ctx context.Context, fn func() error, txnMode TxnMode
 	}
 	tracer.EnterPhase(metrics.PhaseOther)
 
-	success := false
-	defer func() {
-		if !success {
-			_, rollbackErr := s.ExecuteSQL(ctx, "ROLLBACK")
-			terror.Log(rollbackErr)
-		}
-	}()
-
 	if err = fn(); err != nil {
 		return err
 	}

pkg/ttl/ttlworker/BUILD.bazel

@@ -61,6 +61,7 @@ go_test(
         "del_test.go",
         "job_manager_integration_test.go",
         "job_manager_test.go",
+        "scan_integration_test.go",
         "scan_test.go",
         "session_integration_test.go",
         "session_test.go",

pkg/ttl/ttlworker/scan_integration_test.go

@@ -0,0 +1,99 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ttlworker_test
+
+import (
+	"context"
+	"fmt"
+	"math/rand/v2"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/pingcap/tidb/pkg/parser/ast"
+	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/pingcap/tidb/pkg/testkit/testflag"
+	"github.com/pingcap/tidb/pkg/ttl/cache"
+	"github.com/pingcap/tidb/pkg/ttl/ttlworker"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCancelWhileScan(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("create table test.t (id int, created_at datetime) TTL= created_at + interval 1 hour")
+	testTable, err := dom.InfoSchema().TableByName(context.Background(), ast.NewCIStr("test"), ast.NewCIStr("t"))
+	require.NoError(t, err)
+	for i := 0; i < 10000; i++ {
+		tk.MustExec(fmt.Sprintf("insert into test.t values (%d, NOW() - INTERVAL 24 HOUR)", i))
+	}
+	testPhysicalTableCache, err := cache.NewPhysicalTable(ast.NewCIStr("test"), testTable.Meta(), ast.NewCIStr(""))
+	require.NoError(t, err)
+
+	delCh := make(chan *ttlworker.TTLDeleteTask)
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for range delCh {
+			// do nothing
+		}
+	}()
+
+	testStart := time.Now()
+	testDuration := time.Second
+	if testflag.Long() {
+		testDuration = time.Minute
+	}
+	for time.Since(testStart) < testDuration {
+		ctx, cancel := context.WithCancel(context.Background())
+		ttlTask := ttlworker.NewTTLScanTask(ctx, testPhysicalTableCache, &cache.TTLTask{
+			JobID:            "test",
+			TableID:          1,
+			ScanID:           1,
+			ScanRangeStart:   nil,
+			ScanRangeEnd:     nil,
+			ExpireTime:       time.Now().Add(-12 * time.Hour),
+			OwnerID:          "test",
+			OwnerAddr:        "test",
+			OwnerHBTime:      time.Now(),
+			Status:           cache.TaskStatusRunning,
+			StatusUpdateTime: time.Now(),
+			State:            &cache.TTLTaskState{},
+			CreatedTime:      time.Now(),
+		})
+
+		wg := &sync.WaitGroup{}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			ttlTask.DoScan(ctx, delCh, dom.SysSessionPool())
+		}()
+
+		// randomly sleep for a while and cancel the scan
+		time.Sleep(time.Duration(rand.Int() % int(time.Millisecond*10)))
+		startCancel := time.Now()
+		cancel()
+
+		wg.Wait()
+		// make sure the scan is canceled in time
+		require.Less(t, time.Since(startCancel), time.Second)
+	}
+
+	close(delCh)
+	wg.Wait()
+}

pkg/ttl/ttlworker/scan_test.go

@@ -554,3 +554,21 @@ func TestScanTaskCancelStmt(t *testing.T) {
 	task.ctx, cancel = context.WithCancel(context.Background())
 	testCancel(context.Background(), cancel)
 }
+
+// NewTTLScanTask creates a new TTL scan task for test.
+func NewTTLScanTask(ctx context.Context, tbl *cache.PhysicalTable, ttlTask *cache.TTLTask) *ttlScanTask {
+	return &ttlScanTask{
+		ctx:        ctx,
+		tbl:        tbl,
+		TTLTask:    ttlTask,
+		statistics: &ttlStatistics{},
+	}
+}
+
+// DoScan is an exported version of `doScan` for test.
+func (t *ttlScanTask) DoScan(ctx context.Context, delCh chan<- *TTLDeleteTask, sessPool util.SessionPool) *ttlScanTaskExecResult {
+	return t.doScan(ctx, delCh, sessPool)
+}
+
+// TTLDeleteTask is an exported version of `ttlDeleteTask` for test.
+type TTLDeleteTask = ttlDeleteTask