ttl: fix scan workers shrink (#39860)

close #39859

ttl/ttlworker/job_manager.go

@@ -180,39 +180,73 @@ func (m *JobManager) reportMetrics() {
 
 func (m *JobManager) resizeScanWorkers(count int) error {
  var err error
- m.scanWorkers, err = m.resizeWorkers(m.scanWorkers, count, func() worker {
+ var canceledWorkers []worker
+ m.scanWorkers, canceledWorkers, err = m.resizeWorkers(m.scanWorkers, count, func() worker {
  return newScanWorker(m.delCh, m.notifyStateCh, m.sessPool)
  })
+ for _, w := range canceledWorkers {
+ s := w.(scanWorker)
+
+ var tableID int64
+ var scanErr error
+ result := s.PollTaskResult()
+ if result != nil {
+ tableID = result.task.tbl.ID
+ scanErr = result.err
+ } else {
+ // if the scan worker failed to poll the task, it's possible that the `WaitStopped` has timeout
+ // we still consider the scan task as finished
+ curTask := s.CurrentTask()
+ if curTask == nil {
+ continue
+ }
+ tableID = curTask.tbl.ID
+ scanErr = errors.New("timeout to cancel scan task")
+ }
+
+ job := findJobWithTableID(m.runningJobs, tableID)
+ if job == nil {
+ logutil.Logger(m.ctx).Warn("task state changed but job not found", zap.Int64("tableID", tableID))
+ continue
+ }
+ logutil.Logger(m.ctx).Debug("scan task finished", zap.String("jobID", job.id))
+ job.finishedScanTaskCounter += 1
+ job.scanTaskErr = multierr.Append(job.scanTaskErr, scanErr)
+ }
  return err
 }
 
 func (m *JobManager) resizeDelWorkers(count int) error {
  var err error
- m.delWorkers, err = m.resizeWorkers(m.delWorkers, count, func() worker {
+ m.delWorkers, _, err = m.resizeWorkers(m.delWorkers, count, func() worker {
  return newDeleteWorker(m.delCh, m.sessPool)
  })
  return err
 }
 
-func (m *JobManager) resizeWorkers(workers []worker, count int, factory func() worker) ([]worker, error) {
+// resizeWorkers scales the worker, and returns the full set of workers as the first return value. If there are workers
+// stopped, return the stopped worker in the second return value
+func (m *JobManager) resizeWorkers(workers []worker, count int, factory func() worker) ([]worker, []worker, error) {
  if count < len(workers) {
  logutil.Logger(m.ctx).Info("shrink ttl worker", zap.Int("originalCount", len(workers)), zap.Int("newCount", count))
 
  for _, w := range workers[count:] {
  w.Stop()
  }
+
  var errs error
+ ctx, cancel := context.WithTimeout(m.ctx, 30*time.Second)
  for _, w := range workers[count:] {
- err := w.WaitStopped(m.ctx, 30*time.Second)
+ err := w.WaitStopped(ctx, 30*time.Second)
  if err != nil {
  logutil.Logger(m.ctx).Warn("fail to stop ttl worker", zap.Error(err))
  errs = multierr.Append(errs, err)
  }
  }
+ cancel()
 
  // remove the existing workers, and keep the left workers
- workers = workers[:count]
- return workers, errs
+ return workers[:count], workers[count:], errs
  }
 
  if count > len(workers) {
@@ -223,29 +257,31 @@ func (m *JobManager) resizeWorkers(workers []worker, count int, factory func() w
  w.Start()
  workers = append(workers, w)
  }
- return workers, nil
+ return workers, nil, nil
  }
 
- return workers, nil
+ return workers, nil, nil
 }
 
 func (m *JobManager) updateTaskState() {
  results := m.pollScanWorkerResults()
  for _, result := range results {
  job := findJobWithTableID(m.runningJobs, result.task.tbl.ID)
- if job != nil {
- logutil.Logger(m.ctx).Debug("scan task state changed", zap.String("jobID", job.id))
-
- job.finishedScanTaskCounter += 1
- job.scanTaskErr = multierr.Append(job.scanTaskErr, result.err)
+ if job == nil {
+ logutil.Logger(m.ctx).Warn("task state changed but job not found", zap.Int64("tableID", result.task.tbl.ID))
+ continue
  }
+ logutil.Logger(m.ctx).Debug("scan task finished", zap.String("jobID", job.id))
+
+ job.finishedScanTaskCounter += 1
+ job.scanTaskErr = multierr.Append(job.scanTaskErr, result.err)
  }
 }
 
 func (m *JobManager) pollScanWorkerResults() []*ttlScanTaskExecResult {
  results := make([]*ttlScanTaskExecResult, 0, len(m.scanWorkers))
  for _, w := range m.scanWorkers {
- worker := w.(*ttlScanWorker)
+ worker := w.(scanWorker)
  result := worker.PollTaskResult()
  if result != nil {
  results = append(results, result)

ttl/ttlworker/job_manager_test.go

@@ -306,7 +306,7 @@ func TestResizeWorkers(t *testing.T) {
  m.SetScanWorkers4Test([]worker{
  scanWorker1,
  })
- newWorkers, err := m.resizeWorkers(m.scanWorkers, 2, func() worker {
+ newWorkers, _, err := m.resizeWorkers(m.scanWorkers, 2, func() worker {
  return scanWorker2
  })
  assert.NoError(t, err)
@@ -328,6 +328,24 @@ func TestResizeWorkers(t *testing.T) {
 
  assert.NoError(t, m.resizeScanWorkers(1))
  scanWorker2.checkWorkerStatus(workerStatusStopped, false, nil)
+
+ // shrink scan workers after job is run
+ scanWorker1 = newMockScanWorker(t)
+ scanWorker1.Start()
+ scanWorker2 = newMockScanWorker(t)
+ scanWorker2.Start()
+
+ m = NewJobManager("test-id", newMockSessionPool(t, tbl), nil)
+ m.SetScanWorkers4Test([]worker{
+ scanWorker1,
+ scanWorker2,
+ })
+ m.runningJobs = append(m.runningJobs, &ttlJob{tbl: tbl})
+
+ scanWorker2.curTaskResult = &ttlScanTaskExecResult{task: &ttlScanTask{tbl: tbl}}
+ assert.NoError(t, m.resizeScanWorkers(1))
+ scanWorker2.checkWorkerStatus(workerStatusStopped, false, nil)
+ assert.Equal(t, m.runningJobs[0].finishedScanTaskCounter, 1)
 }
 
 func TestLocalJobs(t *testing.T) {

ttl/ttlworker/scan.go

@@ -349,4 +349,6 @@ type scanWorker interface {
 
  Idle() bool
  Schedule(*ttlScanTask) error
+ PollTaskResult() *ttlScanTaskExecResult
+ CurrentTask() *ttlScanTask
 }

ttl/ttlworker/worker.go

@@ -96,6 +96,12 @@ func (w *baseWorker) Error() error {
 }
 
 func (w *baseWorker) WaitStopped(ctx context.Context, timeout time.Duration) error {
+ // consider the situation when the worker has stopped, but the context has also stopped. We should
+ // return without error
+ if w.Status() == workerStatusStopped {
+ return nil
+ }
+
  ctx, cancel := context.WithTimeout(ctx, timeout)
  go func() {
  w.wg.Wait()