ttl: disable paging in TTL (pingcap#58759)

close pingcap#58342

Author: lcwangchao

pkg/ttl/ttlworker/scan.go

@@ -17,13 +17,11 @@ package ttlworker
 import (
 	"context"
 	"fmt"
-	"strconv"
 	"sync/atomic"
 	"time"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/parser/ast"
-	"github.com/pingcap/tidb/pkg/parser/terror"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/ttl/cache"
 	"github.com/pingcap/tidb/pkg/ttl/metrics"
@@ -199,17 +197,12 @@ func (t *ttlScanTask) doScan(ctx context.Context, delCh chan<- *ttlDeleteTask, s
 		))
 	}
 
-	origConcurrency := rawSess.GetSessionVars().DistSQLScanConcurrency()
-	if _, err = rawSess.ExecuteSQL(ctx, "set @@tidb_distsql_scan_concurrency=1"); err != nil {
+	sess, restoreSession, err := NewScanSession(rawSess, t.tbl, t.ExpireTime)
+	if err != nil {
 		return t.result(err)
 	}
+	defer restoreSession()
 
-	defer func() {
-		_, err = rawSess.ExecuteSQL(ctx, "set @@tidb_distsql_scan_concurrency="+strconv.Itoa(origConcurrency))
-		terror.Log(err)
-	}()
-
-	sess := newTableSession(rawSess, t.tbl, t.ExpireTime)
 	generator, err := sqlbuilder.NewScanQueryGenerator(t.tbl, t.ExpireTime, t.ScanRangeStart, t.ScanRangeEnd)
 	if err != nil {
 		return t.result(err)

pkg/ttl/ttlworker/session.go

@@ -208,6 +208,36 @@ func newTableSession(se session.Session, tbl *cache.PhysicalTable, expire time.T
 	}
 }
 
+// NewScanSession creates a session for scan
+func NewScanSession(se session.Session, tbl *cache.PhysicalTable, expire time.Time) (*ttlTableSession, func(), error) {
+	origConcurrency := se.GetSessionVars().DistSQLScanConcurrency()
+	origPaging := se.GetSessionVars().EnablePaging
+
+	restore := func() {
+		_, err := se.ExecuteSQL(context.Background(), "set @@tidb_distsql_scan_concurrency=%?", origConcurrency)
+		terror.Log(err)
+		_, err = se.ExecuteSQL(context.Background(), "set @@tidb_enable_paging=%?", origPaging)
+		terror.Log(err)
+	}
+
+	// Set the distsql scan concurrency to 1 to reduce the number of cop tasks in TTL scan.
+	if _, err := se.ExecuteSQL(context.Background(), "set @@tidb_distsql_scan_concurrency=1"); err != nil {
+		restore()
+		return nil, nil, err
+	}
+
+	// Disable tidb_enable_paging because we have already had a `LIMIT` in the SQL to limit the result set.
+	// If `tidb_enable_paging` is enabled, it may have multiple cop tasks even in one region that makes some extra
+	// processed keys in TiKV side, see issue: https://github.com/pingcap/tidb/issues/58342.
+	// Disable it to make the scan more efficient.
+	if _, err := se.ExecuteSQL(context.Background(), "set @@tidb_enable_paging=OFF"); err != nil {
+		restore()
+		return nil, nil, err
+	}
+
+	return newTableSession(se, tbl, expire), restore, nil
+}
+
 type ttlTableSession struct {
 	session.Session
 	tbl    *cache.PhysicalTable

pkg/ttl/ttlworker/session_integration_test.go

@@ -20,11 +20,13 @@ import (
 	"strings"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/ngaut/pools"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/pingcap/tidb/pkg/ttl/cache"
 	"github.com/pingcap/tidb/pkg/ttl/ttlworker"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/logutil"
@@ -185,3 +187,64 @@ func TestGetSessionWithFault(t *testing.T) {
 		require.True(t, se != nil || err != nil)
 	}
 }
+
+func TestNewScanSession(t *testing.T) {
+	_, dom := testkit.CreateMockStoreAndDomain(t)
+	pool := newFaultSessionPool(dom.SysSessionPool())
+	pool.setFault(newFaultWithFilter(func(s string) bool { return false }, newFaultAfterCount(0)))
+	se, err := ttlworker.GetSessionForTest(pool)
+	require.NoError(t, err)
+
+	_, err = se.ExecuteSQL(context.Background(), "set @@tidb_distsql_scan_concurrency=123")
+	require.NoError(t, err)
+	require.Equal(t, 123, se.GetSessionVars().DistSQLScanConcurrency())
+
+	_, err = se.ExecuteSQL(context.Background(), "set @@tidb_enable_paging=ON")
+	require.NoError(t, err)
+	require.True(t, se.GetSessionVars().EnablePaging)
+
+	for _, errSQL := range []string{
+		"",
+		"set @@tidb_distsql_scan_concurrency=1",
+		"set @@tidb_enable_paging=OFF",
+	} {
+		t.Run("test err in SQL: "+errSQL, func(t *testing.T) {
+			var faultCnt atomic.Int64
+			pool.setFault(newFaultWithFilter(func(s string) bool {
+				if s == errSQL && s != "" {
+					faultCnt.Add(1)
+					return true
+				}
+				return false
+			}, newFaultAfterCount(0)))
+			tblSe, restore, err := ttlworker.NewScanSession(se, &cache.PhysicalTable{}, time.Now())
+			if errSQL == "" {
+				// success case
+				require.NoError(t, err)
+				require.NotNil(t, tblSe)
+				require.NotNil(t, restore)
+				require.Same(t, se, tblSe.Session)
+				require.Equal(t, int64(0), faultCnt.Load())
+
+				// NewScanSession should override @@dist_sql_scan_concurrency and @@tidb_enable_paging
+				require.Equal(t, 1, se.GetSessionVars().DistSQLScanConcurrency())
+				require.False(t, se.GetSessionVars().EnablePaging)
+
+				// restore should restore the session variables
+				restore()
+				require.Equal(t, 123, se.GetSessionVars().DistSQLScanConcurrency())
+				require.True(t, se.GetSessionVars().EnablePaging)
+			} else {
+				// error case
+				require.Equal(t, int64(1), faultCnt.Load())
+				require.EqualError(t, err, "fault in test")
+				require.Nil(t, tblSe)
+				require.Nil(t, restore)
+
+				// NewScanSession should not change session state if error occurs
+				require.Equal(t, 123, se.GetSessionVars().DistSQLScanConcurrency())
+				require.True(t, se.GetSessionVars().EnablePaging)
+			}
+		})
+	}
+}