ttl: reschedule task to other instances when shriking worker

Author: lcwangchao

pkg/ttl/ttlworker/del.go

@@ -22,6 +22,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/pingcap/log"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/ttl/cache"
 	"github.com/pingcap/tidb/pkg/ttl/metrics"
@@ -94,12 +95,12 @@ func (t *ttlDeleteTask) doDelete(ctx context.Context, rawSe session.Session) (re
 	leftRows := t.rows
 	defer func() {
 		if len(leftRows) > 0 {
-			t.statistics.IncErrorRows(len(leftRows))
+			retryRows = append(retryRows, leftRows...)
 		}
 	}()
 
 	se := newTableSession(rawSe, t.tbl, t.expire)
-	for len(leftRows) > 0 {
+	for len(leftRows) > 0 && ctx.Err() == nil {
 		maxBatch := variable.TTLDeleteBatchSize.Load()
 		var delBatch [][]types.Datum
 		if int64(len(leftRows)) < maxBatch {
@@ -133,7 +134,6 @@ func (t *ttlDeleteTask) doDelete(ctx context.Context, rawSe session.Session) (re
 		sqlInterval := time.Since(sqlStart)
 		if err != nil {
 			metrics.DeleteErrorDuration.Observe(sqlInterval.Seconds())
-			needRetry = needRetry && ctx.Err() == nil
 			logutil.BgLogger().Warn(
 				"delete SQL in TTL failed",
 				zap.Error(err),
@@ -214,6 +214,11 @@ func (b *ttlDelRetryBuffer) DoRetry(do func(*ttlDeleteTask) [][]types.Datum) tim
 	return b.retryInterval
 }
 
+// SetRetryInterval sets the retry interval of the buffer.
+func (b *ttlDelRetryBuffer) SetRetryInterval(interval time.Duration) {
+	b.retryInterval = interval
+}
+
 // Drain drains a retry buffer.
 func (b *ttlDelRetryBuffer) Drain() {
 	for ele := b.list.Front(); ele != nil; ele = ele.Next() {
@@ -296,8 +301,36 @@ func (w *ttlDeleteWorker) loop() error {
 	timer := time.NewTimer(w.retryBuffer.retryInterval)
 	defer timer.Stop()
 
-	// drain retry buffer to make sure the statistics are correct
-	defer w.retryBuffer.Drain()
+	defer func() {
+		// have a final try for all rows in retry buffer before the worker stops
+		if w.retryBuffer.Len() > 0 {
+			start := time.Now()
+			log.Info(
+				"try to delete TTL rows in del worker buffer immediately because the worker is going to stop",
+				zap.Int("bufferLen", w.retryBuffer.Len()),
+			)
+			retryCtx, cancel := context.WithTimeout(context.Background(), time.Second*10)
+			defer cancel()
+			w.retryBuffer.SetRetryInterval(0)
+			w.retryBuffer.DoRetry(func(task *ttlDeleteTask) [][]types.Datum {
+				return task.doDelete(retryCtx, se)
+			})
+			log.Info(
+				"delete TTL rows in del worker buffer finished",
+				zap.Duration("duration", time.Since(start)),
+			)
+		}
+
+		// drain retry buffer to make sure the statistics are correct
+		if w.retryBuffer.Len() > 0 {
+			log.Warn(
+				"some TTL rows are still in the buffer while the worker is going to stop, mark them as error",
+				zap.Int("bufferLen", w.retryBuffer.Len()),
+			)
+			w.retryBuffer.Drain()
+		}
+	}()
+
 	for w.Status() == workerStatusRunning {
 		tracer.EnterPhase(metrics.PhaseIdle)
 		select {

pkg/ttl/ttlworker/scan.go

@@ -82,9 +82,10 @@ type ttlScanTask struct {
 }
 
 type ttlScanTaskExecResult struct {
-	time time.Time
-	task *ttlScanTask
-	err  error
+	time       time.Time
+	task       *ttlScanTask
+	reschedule bool
+	err        error
 }
 
 func (t *ttlScanTask) result(err error) *ttlScanTaskExecResult {

pkg/ttl/ttlworker/task_manager.go

@@ -146,6 +146,7 @@ func (m *taskManager) resizeScanWorkers(count int) error {
 		var scanErr error
 		result := s.PollTaskResult()
 		if result != nil {
+			result.reschedule = true
 			jobID = result.task.JobID
 			scanID = result.task.ScanID
 
@@ -470,11 +471,13 @@ func (m *taskManager) checkFinishedTask(se session.Session, now time.Time) {
 		}
 		// we should cancel task to release inner context and avoid memory leak
 		task.cancel()
-		err := m.reportTaskFinished(se, now, task)
-		if err != nil {
-			logutil.Logger(m.ctx).Error("fail to report finished task", zap.Error(err))
-			stillRunningTasks = append(stillRunningTasks, task)
-			continue
+		if !task.result.reschedule {
+			err := m.reportTaskFinished(se, now, task)
+			if err != nil {
+				logutil.Logger(m.ctx).Error("fail to report finished task", zap.Error(err))
+				stillRunningTasks = append(stillRunningTasks, task)
+				continue
+			}
 		}
 	}
 
@@ -613,6 +616,11 @@ func (t *runningScanTask) finished(logger *zap.Logger) bool {
 		zap.String("table", t.tbl.Name.O),
 	)
 
+	if t.result.reschedule {
+		logger.Info("task should be rescheduled in other instances")
+		return true
+	}
+
 	totalRows := t.statistics.TotalRows.Load()
 	errRows := t.statistics.ErrorRows.Load()
 	successRows := t.statistics.SuccessRows.Load()