ttl: force to kill SQL in scan task when canceling TTL job/task (#56518) (#56578)

close #56511

Author: ti-chi-bot

ttl/session/BUILD.bazel

@@ -28,12 +28,13 @@ go_test(
         "sysvar_test.go",
     ],
     flaky = True,
-    shard_count = 6,
+    shard_count = 7,
     deps = [
         ":session",
         "//sessionctx/variable",
         "//testkit",
         "//testkit/testsetup",
+        "//util",
         "@com_github_pingcap_errors//:errors",
         "@com_github_stretchr_testify//require",
         "@org_uber_go_goleak//:goleak",

ttl/session/session.go

@@ -16,6 +16,7 @@ package session
 
 import (
 	"context"
+	"sync/atomic"
 	"time"
 
 	"github.com/pingcap/errors"
@@ -51,6 +52,8 @@ type Session interface {
 	RunInTxn(ctx context.Context, fn func() error, mode TxnMode) (err error)
 	// ResetWithGlobalTimeZone resets the session time zone to global time zone
 	ResetWithGlobalTimeZone(ctx context.Context) error
+	// KillStmt kills the current statement execution
+	KillStmt()
 	// Close closes the session
 	Close()
 	// Now returns the current time in location specified by session var
@@ -168,6 +171,11 @@ func (s *session) ResetWithGlobalTimeZone(ctx context.Context) error {
 	return err
 }
 
+// KillStmt kills the current statement execution
+func (s *session) KillStmt() {
+	atomic.StoreUint32(&s.GetSessionVars().Killed, 1)
+}
+
 // Close closes the session
 func (s *session) Close() {
 	if s.closeFn != nil {

ttl/session/session_test.go

@@ -17,10 +17,12 @@ package session_test
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/testkit"
 	"github.com/pingcap/tidb/ttl/session"
+	"github.com/pingcap/tidb/util"
 	"github.com/stretchr/testify/require"
 )
 
@@ -64,3 +66,28 @@ func TestSessionResetTimeZone(t *testing.T) {
 	require.NoError(t, se.ResetWithGlobalTimeZone(context.TODO()))
 	tk.MustQuery("select @@time_zone").Check(testkit.Rows("UTC"))
 }
+
+func TestSessionKill(t *testing.T) {
+	store, do := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	se := session.NewSession(tk.Session(), tk.Session(), nil)
+	sleepStmt := "select sleep(123)"
+	wg := util.WaitGroupWrapper{}
+	wg.Run(func() {
+		start := time.Now()
+		for time.Since(start) < 10*time.Second {
+			time.Sleep(10 * time.Millisecond)
+			processes := do.InfoSyncer().GetSessionManager().ShowProcessList()
+			for _, proc := range processes {
+				if proc.Info == sleepStmt {
+					se.KillStmt()
+					return
+				}
+			}
+		}
+		require.FailNow(t, "wait sleep stmt timeout")
+	})
+	// the killed sleep stmt will return "1"
+	tk.MustQuery(sleepStmt).Check(testkit.Rows("1"))
+	wg.Wait()
+}

ttl/ttlworker/BUILD.bazel

@@ -59,7 +59,7 @@ go_test(
     embed = [":ttlworker"],
     flaky = True,
     race = "on",
-    shard_count = 38,
+    shard_count = 39,
     deps = [
         "//domain",
         "//infoschema",
@@ -78,6 +78,7 @@ go_test(
         "//ttl/metrics",
         "//ttl/session",
         "//types",
+        "//util",
         "//util/chunk",
         "//util/logutil",
         "@com_github_google_uuid//:uuid",

ttl/ttlworker/scan.go

@@ -28,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/ttl/metrics"
 	"github.com/pingcap/tidb/ttl/sqlbuilder"
 	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/pingcap/tidb/util/logutil"
 	"go.uber.org/zap"
@@ -110,7 +111,45 @@ func (t *ttlScanTask) doScan(ctx context.Context, delCh chan<- *ttlDeleteTask, s
 	if err != nil {
 		return t.result(err)
 	}
-	defer rawSess.Close()
+
+	doScanFinished, setDoScanFinished := context.WithCancel(context.Background())
+	wg := util.WaitGroupWrapper{}
+	wg.Run(func() {
+		select {
+		case <-taskCtx.Done():
+		case <-ctx.Done():
+		case <-doScanFinished.Done():
+			return
+		}
+		logger := logutil.BgLogger().With(
+			zap.Int64("tableID", t.TableID),
+			zap.String("table", t.tbl.Name.O),
+			zap.String("partition", t.tbl.Partition.O),
+			zap.String("jobID", t.JobID),
+			zap.Int64("scanID", t.ScanID),
+		)
+		logger.Info("kill the running statement in scan task because the task or worker cancelled")
+		rawSess.KillStmt()
+		ticker := time.NewTicker(time.Minute)
+		defer ticker.Stop()
+		for {
+			// Have a small probability that the kill signal will be lost when the session is idle.
+			// So wait for a while and send the kill signal again if the scan is still running.
+			select {
+			case <-doScanFinished.Done():
+				return
+			case <-ticker.C:
+				logger.Warn("scan task is still running after the kill signal sent, kill it again")
+				rawSess.KillStmt()
+			}
+		}
+	})
+
+	defer func() {
+		setDoScanFinished()
+		wg.Wait()
+		rawSess.Close()
+	}()
 
 	safeExpire, err := t.tbl.EvalExpireTime(taskCtx, rawSess, time.Now())
 	if err != nil {
@@ -181,7 +220,7 @@ func (t *ttlScanTask) doScan(ctx context.Context, delCh chan<- *ttlDeleteTask, s
 		selectInterval := time.Since(sqlStart)
 		if sqlErr != nil {
 			metrics.SelectErrorDuration.Observe(selectInterval.Seconds())
-			needRetry := retryable && retryTimes < scanTaskExecuteSQLMaxRetry && ctx.Err() == nil
+			needRetry := retryable && retryTimes < scanTaskExecuteSQLMaxRetry && ctx.Err() == nil && t.ctx.Err() == nil
 			logutil.BgLogger().Error("execute query for ttl scan task failed",
 				zap.String("SQL", sql),
 				zap.Int("retryTimes", retryTimes),

ttl/ttlworker/scan_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/ttl/cache"
 	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/stretchr/testify/require"
 )
@@ -445,3 +446,51 @@ func TestScanTaskCheck(t *testing.T) {
 	require.Equal(t, 1, len(ch))
 	require.Equal(t, "Total Rows: 1, Success Rows: 0, Error Rows: 0", task.statistics.String())
 }
+
+func TestScanTaskCancelStmt(t *testing.T) {
+	task := &ttlScanTask{
+		ctx: context.Background(),
+		tbl: newMockTTLTbl(t, "t1"),
+		TTLTask: &cache.TTLTask{
+			ExpireTime:     time.UnixMilli(0),
+			ScanRangeStart: []types.Datum{types.NewIntDatum(0)},
+		},
+		statistics: &ttlStatistics{},
+	}
+
+	testCancel := func(ctx context.Context, doCancel func()) {
+		mockPool := newMockSessionPool(t)
+		startExec := make(chan struct{})
+		mockPool.se.sessionInfoSchema = newMockInfoSchema(task.tbl.TableInfo)
+		mockPool.se.executeSQL = func(_ context.Context, _ string, _ ...any) ([]chunk.Row, error) {
+			close(startExec)
+			select {
+			case <-mockPool.se.killed:
+				return nil, errors.New("killed")
+			case <-time.After(10 * time.Second):
+				return nil, errors.New("timeout")
+			}
+		}
+		wg := util.WaitGroupWrapper{}
+		wg.Run(func() {
+			select {
+			case <-startExec:
+			case <-time.After(10 * time.Second):
+				require.FailNow(t, "timeout")
+			}
+			doCancel()
+		})
+		r := task.doScan(ctx, nil, mockPool)
+		require.NotNil(t, r)
+		require.EqualError(t, r.err, "killed")
+		wg.Wait()
+	}
+
+	// test cancel with input context
+	ctx, cancel := context.WithCancel(context.Background())
+	testCancel(ctx, cancel)
+
+	// test cancel with task context
+	task.ctx, cancel = context.WithCancel(context.Background())
+	testCancel(context.Background(), cancel)
+}

ttl/ttlworker/session_test.go

@@ -140,6 +140,7 @@ type mockSession struct {
 	resetTimeZoneCalls int
 	closed             bool
 	commitErr          error
+	killed             chan struct{}
 }
 
 func newMockSession(t *testing.T, tbl ...*cache.PhysicalTable) *mockSession {
@@ -154,6 +155,7 @@ func newMockSession(t *testing.T, tbl ...*cache.PhysicalTable) *mockSession {
 		sessionInfoSchema: newMockInfoSchema(tbls...),
 		evalExpire:        time.Now(),
 		sessionVars:       sessVars,
+		killed:            make(chan struct{}),
 	}
 }
 
@@ -201,6 +203,11 @@ func (s *mockSession) ResetWithGlobalTimeZone(_ context.Context) (err error) {
 	return nil
 }
 
+// KillStmt kills the current statement execution
+func (s *mockSession) KillStmt() {
+	close(s.killed)
+}
+
 func (s *mockSession) Close() {
 	s.closed = true
 }