This is an automated cherry-pick of pingcap#56518

Signed-off-by: ti-chi-bot <[email protected]>

Author: lcwangchao

ttl/session/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
     importpath = "github.com/pingcap/tidb/ttl/session",
     visibility = ["//visibility:public"],
     deps = [
+<<<<<<< HEAD:ttl/session/BUILD.bazel
         "//infoschema",
         "//kv",
         "//parser/terror",
@@ -15,6 +16,19 @@ go_library(
         "//ttl/metrics",
         "//util/chunk",
         "//util/sqlexec",
+=======
+        "//pkg/infoschema",
+        "//pkg/kv",
+        "//pkg/parser/terror",
+        "//pkg/sessionctx",
+        "//pkg/sessionctx/variable",
+        "//pkg/sessiontxn",
+        "//pkg/ttl/metrics",
+        "//pkg/util/chunk",
+        "//pkg/util/sqlexec",
+        "//pkg/util/sqlkiller",
+        "//pkg/util/timeutil",
+>>>>>>> e68c26a0e67 (ttl: force to kill SQL in scan task when canceling TTL job/task (#56518)):pkg/ttl/session/BUILD.bazel
         "@com_github_pingcap_errors//:errors",
     ],
 )
@@ -28,12 +42,19 @@ go_test(
         "sysvar_test.go",
     ],
     flaky = True,
-    shard_count = 6,
+    shard_count = 7,
     deps = [
         ":session",
+<<<<<<< HEAD:ttl/session/BUILD.bazel
         "//sessionctx/variable",
         "//testkit",
         "//testkit/testsetup",
+=======
+        "//pkg/sessionctx/variable",
+        "//pkg/testkit",
+        "//pkg/testkit/testsetup",
+        "//pkg/util",
+>>>>>>> e68c26a0e67 (ttl: force to kill SQL in scan task when canceling TTL job/task (#56518)):pkg/ttl/session/BUILD.bazel
         "@com_github_pingcap_errors//:errors",
         "@com_github_stretchr_testify//require",
         "@org_uber_go_goleak//:goleak",

ttl/session/session.go

@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+<<<<<<< HEAD:ttl/session/session.go
 	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/parser/terror"
@@ -28,6 +29,19 @@ import (
 	"github.com/pingcap/tidb/ttl/metrics"
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/pingcap/tidb/util/sqlexec"
+=======
+	"github.com/pingcap/tidb/pkg/infoschema"
+	"github.com/pingcap/tidb/pkg/kv"
+	"github.com/pingcap/tidb/pkg/parser/terror"
+	"github.com/pingcap/tidb/pkg/sessionctx"
+	"github.com/pingcap/tidb/pkg/sessionctx/variable"
+	"github.com/pingcap/tidb/pkg/sessiontxn"
+	"github.com/pingcap/tidb/pkg/ttl/metrics"
+	"github.com/pingcap/tidb/pkg/util/chunk"
+	"github.com/pingcap/tidb/pkg/util/sqlexec"
+	"github.com/pingcap/tidb/pkg/util/sqlkiller"
+	"github.com/pingcap/tidb/pkg/util/timeutil"
+>>>>>>> e68c26a0e67 (ttl: force to kill SQL in scan task when canceling TTL job/task (#56518)):pkg/ttl/session/session.go
 )
 
 // TxnMode represents using optimistic or pessimistic mode in the transaction
@@ -51,6 +65,13 @@ 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
+<<<<<<< HEAD:ttl/session/session.go
+=======
+	// GlobalTimeZone returns the global timezone. It is used to compute expire time for TTL
+	GlobalTimeZone(ctx context.Context) (*time.Location, error)
+	// KillStmt kills the current statement execution
+	KillStmt()
+>>>>>>> e68c26a0e67 (ttl: force to kill SQL in scan task when canceling TTL job/task (#56518)):pkg/ttl/session/session.go
 	// Close closes the session
 	Close()
 	// Now returns the current time in location specified by session var
@@ -168,6 +189,23 @@ func (s *session) ResetWithGlobalTimeZone(ctx context.Context) error {
 	return err
 }
 
+<<<<<<< HEAD:ttl/session/session.go
+=======
+// GlobalTimeZone returns the global timezone
+func (s *session) GlobalTimeZone(ctx context.Context) (*time.Location, error) {
+	str, err := s.GetSessionVars().GetGlobalSystemVar(ctx, "time_zone")
+	if err != nil {
+		return nil, err
+	}
+	return timeutil.ParseTimeZone(str)
+}
+
+// KillStmt kills the current statement execution
+func (s *session) KillStmt() {
+	s.GetSessionVars().SQLKiller.SendKillSignal(sqlkiller.QueryInterrupted)
+}
+
+>>>>>>> e68c26a0e67 (ttl: force to kill SQL in scan task when canceling TTL job/task (#56518)):pkg/ttl/session/session.go
 // Close closes the session
 func (s *session) Close() {
 	if s.closeFn != nil {

ttl/session/session_test.go

@@ -17,10 +17,17 @@ package session_test
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/pingcap/errors"
+<<<<<<< HEAD:ttl/session/session_test.go
 	"github.com/pingcap/tidb/testkit"
 	"github.com/pingcap/tidb/ttl/session"
+=======
+	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/pingcap/tidb/pkg/ttl/session"
+	"github.com/pingcap/tidb/pkg/util"
+>>>>>>> e68c26a0e67 (ttl: force to kill SQL in scan task when canceling TTL job/task (#56518)):pkg/ttl/session/session_test.go
 	"github.com/stretchr/testify/require"
 )
 
@@ -64,3 +71,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/scan.go

@@ -110,7 +110,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 +219,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

@@ -21,11 +21,20 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+<<<<<<< HEAD:ttl/ttlworker/scan_test.go
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/ttl/cache"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/chunk"
+=======
+	"github.com/pingcap/tidb/pkg/parser/mysql"
+	"github.com/pingcap/tidb/pkg/sessionctx/variable"
+	"github.com/pingcap/tidb/pkg/ttl/cache"
+	"github.com/pingcap/tidb/pkg/types"
+	"github.com/pingcap/tidb/pkg/util"
+	"github.com/pingcap/tidb/pkg/util/chunk"
+>>>>>>> e68c26a0e67 (ttl: force to kill SQL in scan task when canceling TTL job/task (#56518)):pkg/ttl/ttlworker/scan_test.go
 	"github.com/stretchr/testify/require"
 )
 
@@ -445,3 +454,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,19 @@ func (s *mockSession) ResetWithGlobalTimeZone(_ context.Context) (err error) {
 	return nil
 }
 
+<<<<<<< HEAD:ttl/ttlworker/session_test.go
+=======
+// GlobalTimeZone returns the global timezone
+func (s *mockSession) GlobalTimeZone(_ context.Context) (*time.Location, error) {
+	return time.Local, nil
+}
+
+// KillStmt kills the current statement execution
+func (s *mockSession) KillStmt() {
+	close(s.killed)
+}
+
+>>>>>>> e68c26a0e67 (ttl: force to kill SQL in scan task when canceling TTL job/task (#56518)):pkg/ttl/ttlworker/session_test.go
 func (s *mockSession) Close() {
 	s.closed = true
 }