Repository: pingcap/tidb (go)

Issue:
dxf: the `Task Status` pannel in `Dist Execute Framework` display incorrectly
## Bug Report

Please answer these questions before submitting your issue. Thanks!

### 1. Minimal reproduce step (Required)
Use sysbench prepare 100 table with 10000 rows each.
<!-- a step by step guide for reproducing the bug. -->

### 2. What did you expect to see? (Required)
After all the add index  jobs finished, the `backfill-completed` task count should be 100 and the `backfill-running` task count should be 0.
### 3. What did you see instead (Required)
![img_v3_02gc_401b2848-5df2-4087-a8c9-0df8d7d5bbcg](https://github.com/user-attachments/assets/4fc439f7-93f7-4e42-b90f-98891c782cc6)

### 4. What is your TiDB version? (Required)
master
<!-- Paste the output of SELECT tidb_version() -->



Pull request description:
dxf: fix the task num count incorrectly (#57127)
This is an automated cherry-pick of #57127

<!--

Thank you for contributing to TiDB!

PR Title Format:
1. pkg [, pkg2, pkg3]: what's changed
2. *: what's changed

-->

### What problem does this PR solve?
<!--

Please create an issue first to describe the problem.

There MUST be one line starting with "Issue Number:  " and
linking the relevant issues via the "close" or "ref".

For more info, check https://pingcap.github.io/tidb-dev-guide/contribute-to-tidb/contribute-code.html#referring-to-an-issue.

-->

Issue Number: close #57172

Problem Summary:

### What changed and how does it work?
The task count metric is updating in scheduler loop,  but the task maybe cleaned up in `doCleanupTask()`, the scheduler will not handle the finished/failed... state, so the task count metric will not be reduced.
### Check List

Tests <!-- At least one of them must be included. -->

- [ ] Unit test
- [ ] Integration test
- [x] Manual test (add detailed scripts or steps below)
Use sysbench prepare 100 tables with 10000 rows. Watch the pannel:
<img width="929" alt="lnAQTlHyy1" src="https://github.com/user-attachments/assets/ccd2d2db-ad26-4938-aa81-199c19597977">

- [ ] No need to test
  > - [ ] I checked and no code files have been changed.
  > <!-- Or your custom  "No need to test" reasons -->

Side effects

- [ ] Performance regression: Consumes more CPU
- [ ] Performance regression: Consumes more Memory
- [ ] Breaking backward compatibility

Documentation

- [ ] Affects user behaviors
- [ ] Contains syntax changes
- [ ] Contains variable changes
- [ ] Contains experimental features
- [ ] Changes MySQL compatibility

### Release note

<!-- compatibility change, improvement, bugfix, and new feature need a release note -->

Please refer to [Release Notes Language Style Guide](https://pingcap.github.io/tidb-dev-guide/contribute-to-tidb/release-notes-style-guide.html) to write a quality release note.

```release-note
None
```


Patch:
diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go
index 132ecf498a..838205e585 100644
--- a/pkg/ddl/index.go
+++ b/pkg/ddl/index.go
@@ -2266,7 +2266,7 @@ func (w *worker) updateJobRowCount(taskKey string, jobID int64) {
 		logutil.DDLLogger().Warn("cannot get task manager", zap.String("task_key", taskKey), zap.Error(err))
 		return
 	}
-	task, err := taskMgr.GetTaskByKey(w.ctx, taskKey)
+	task, err := taskMgr.GetTaskByKeyWithHistory(w.ctx, taskKey)
 	if err != nil {
 		logutil.DDLLogger().Warn("cannot get task", zap.String("task_key", taskKey), zap.Error(err))
 		return
diff --git a/pkg/disttask/framework/scheduler/scheduler.go b/pkg/disttask/framework/scheduler/scheduler.go
index 3ac08f3191..eae669409d 100644
--- a/pkg/disttask/framework/scheduler/scheduler.go
+++ b/pkg/disttask/framework/scheduler/scheduler.go
@@ -29,7 +29,6 @@ import (
 	"github.com/pingcap/tidb/pkg/disttask/framework/storage"
 	"github.com/pingcap/tidb/pkg/domain/infosync"
 	"github.com/pingcap/tidb/pkg/kv"
-	"github.com/pingcap/tidb/pkg/metrics"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/util/backoff"
 	disttaskutil "github.com/pingcap/tidb/pkg/util/disttask"
@@ -173,6 +172,8 @@ func (s *BaseScheduler) scheduleTask() {
 				if errors.Cause(err) == storage.ErrTaskNotFound {
 					// this can happen when task is reverted/succeed, but before
 					// we reach here, cleanup routine move it to history.
+					s.logger.Debug("task not found, might be reverted/succeed/failed", zap.Int64("task_id", s.GetTask().ID),
+						zap.String("task_key", s.GetTask().Key))
 					return
 				}
 				s.logger.Error("refresh task failed", zap.Error(err))
@@ -412,7 +413,6 @@ func (s *BaseScheduler) onRunning() error {
 
 func (s *BaseScheduler) onFinished() {
 	task := s.GetTask()
-	metrics.UpdateMetricsForFinishTask(task)
 	s.logger.Debug("schedule task, task is finished", zap.Stringer("state", task.State))
 }
 
diff --git a/pkg/disttask/framework/scheduler/scheduler_manager.go b/pkg/disttask/framework/scheduler/scheduler_manager.go
index a2f8f8789a..df28c84462 100644
--- a/pkg/disttask/framework/scheduler/scheduler_manager.go
+++ b/pkg/disttask/framework/scheduler/scheduler_manager.go
@@ -285,8 +285,7 @@ func (sm *Manager) startSchedulers(schedulableTasks []*proto.TaskBase) error {
 				zap.Int64("task-id", task.ID), zap.Stringer("state", task.State))
 		}
 
-		metrics.DistTaskGauge.WithLabelValues(task.Type.String(), metrics.SchedulingStatus).Inc()
-		metrics.UpdateMetricsForScheduleTask(task.ID, task.Type)
+		metrics.UpdateMetricsForScheduleTask(task)
 		sm.startScheduler(task, allocateSlots, reservedExecID)
 	}
 	return nil
@@ -437,6 +436,7 @@ func (sm *Manager) cleanupFinishedTasks(tasks []*proto.Task) error {
 			// if task doesn't register cleanup function, mark it as cleaned.
 			cleanedTasks = append(cleanedTasks, task)
 		}
+		metrics.UpdateMetricsForFinishTask(task)
 	}
 	if firstErr != nil {
 		sm.logger.Warn("cleanup routine failed", zap.Error(errors.Trace(firstErr)))
diff --git a/pkg/disttask/framework/storage/task_table.go b/pkg/disttask/framework/storage/task_table.go
index aa8f82262f..aaeee4b446 100644
--- a/pkg/disttask/framework/storage/task_table.go
+++ b/pkg/disttask/framework/storage/task_table.go
@@ -480,8 +480,12 @@ func (mgr *TaskManager) GetAllSubtasksByStepAndState(ctx context.Context, taskID
 func (mgr *TaskManager) GetSubtaskRowCount(ctx context.Context, taskID int64, step proto.Step) (int64, error) {
 	rs, err := mgr.ExecuteSQLWithNewSession(ctx, `select
     	cast(sum(json_extract(summary, '$.row_count')) as signed) as row_count
-		from mysql.tidb_background_subtask where task_key = %? and step = %?`,
-		taskID, step)
+		from (
+			select summary from mysql.tidb_background_subtask where task_key = %? and step = %?
+			union all
+			select summary from mysql.tidb_background_subtask_history where task_key = %? and step = %?
+		) as combined`,
+		taskID, step, taskID, step)
 	if err != nil {
 		return 0, err
 	}
diff --git a/pkg/executor/aggregate/agg_hash_executor.go b/pkg/executor/aggregate/agg_hash_executor.go
index 4ffe6e7063..f2b0925f29 100644
--- a/pkg/executor/aggregate/agg_hash_executor.go
+++ b/pkg/executor/aggregate/agg_hash_executor.go
@@ -154,6 +154,8 @@ type HashAggExec struct {
 	spillHelper *parallelHashAggSpillHelper
 	// isChildDrained indicates whether the all data from child has been taken out.
 	isChildDrained bool
+
+	invalidMemoryUsageForTrackingTest bool
 }
 
 // Close implements the Executor Close interface.
@@ -204,7 +206,11 @@ func (e *HashAggExec) Close() error {
 		channel.Clear(e.finalOutputCh)
 		e.executed.Store(false)
 		if e.memTracker != nil {
-			e.memTracker.ReplaceBytesUsed(0)
+			if e.memTracker.BytesConsumed() < 0 {
+				e.invalidMemoryUsageForTrackingTest = true
+			} else {
+				e.memTracker.ReplaceBytesUsed(0)
+			}
 		}
 		e.parallelExecValid = false
 		if e.parallelAggSpillAction != nil {
@@ -285,6 +291,8 @@ func (e *HashAggExec) initForUnparallelExec() {
 }
 
 func (e *HashAggExec) initPartialWorkers(partialConcurrency int, finalConcurrency int, ctx sessionctx.Context) {
+	memUsage := int64(0)
+
 	for i := 0; i < partialConcurrency; i++ {
 		partialResultsMap := make([]aggfuncs.AggPartialResultMapper, finalConcurrency)
 		for i := 0; i < finalConcurrency; i++ {
@@ -311,6 +319,8 @@ func (e *HashAggExec) initPartialWorkers(partialConcurrency int, finalConcurrenc
 			inflightChunkSync:    e.inflightChunkSync,
 		}
 
+		memUsage += e.partialWorkers[i].chk.MemoryUsage()
+
 		e.partialWorkers[i].partialResultNumInRow = e.partialWorkers[i].getPartialResultSliceLenConsiderByteAlign()
 		for j := 0; j < finalConcurrency; j++ {
 			e.partialWorkers[i].BInMaps[j] = 0
@@ -328,9 +338,11 @@ func (e *HashAggExec) initPartialWorkers(partialConcurrency int, finalConcurrenc
 			chk:        exec.NewFirstChunk(e.Children(0)),
 			giveBackCh: e.partialWorkers[i].inputCh,
 		}
-		e.memTracker.Consume(input.chk.MemoryUsage())
+		memUsage += input.chk.MemoryUsage()
 		e.inputCh <- input
 	}
+
+	e.memTracker.Consume(memUsage)
 }
 
 func (e *HashAggExec) initFinalWorkers(finalConcurrency int) {
@@ -442,6 +454,7 @@ func (e *HashAggExec) fetchChildData(ctx context.Context, waitGroup *sync.WaitGr
 		ok    bool
 		err   error
 	)
+
 	defer func() {
 		if r := recover(); r != nil {
 			recoveryHashAgg(e.finalOutputCh, r)
@@ -494,6 +507,7 @@ func (e *HashAggExec) fetchChildData(ctx context.Context, waitGroup *sync.WaitGr
 		input.giveBackCh <- chk
 
 		if hasError := e.spillIfNeed(); hasError {
+			e.memTracker.Consume(-mSize)
 			return
 		}
 	}
@@ -857,3 +871,8 @@ func (e *HashAggExec) IsSpillTriggeredForTest() bool {
 	}
 	return false
 }
+
+// IsInvalidMemoryUsageTrackingForTest is for test
+func (e *HashAggExec) IsInvalidMemoryUsageTrackingForTest() bool {
+	return e.invalidMemoryUsageForTrackingTest
+}
diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go
index 9cbd8b0f99..c0467e97b0 100644
--- a/pkg/executor/importer/import.go
+++ b/pkg/executor/importer/import.go
@@ -79,20 +79,30 @@ const (
 	// 0 means no limit
 	unlimitedWriteSpeed = config.ByteSize(0)
 
-	characterSetOption          = "character_set"
-	fieldsTerminatedByOption    = "fields_terminated_by"
-	fieldsEnclosedByOption      = "fields_enclosed_by"
-	fieldsEscapedByOption       = "fields_escaped_by"
-	fieldsDefinedNullByOption   = "fields_defined_null_by"
-	linesTerminatedByOption     = "lines_terminated_by"
-	skipRowsOption              = "skip_rows"
-	splitFileOption             = "split_file"
-	diskQuotaOption             = "disk_quota"
-	threadOption                = "thread"
-	maxWriteSpeedOption         = "max_write_speed"
-	checksumTableOption         = "checksum_table"
-	recordErrorsOption          = "record_errors"
-	detachedOption              = "detached"
+	characterSetOption        = "character_set"
+	fieldsTerminatedByOption  = "fields_terminated_by"
+	fieldsEnclosedByOption    = "fields_enclosed_by"
+	fieldsEscapedByOption     = "fields_escaped_by"
+	fieldsDefinedNullByOption = "fields_defined_null_by"
+	linesTerminatedByOption   = "lines_terminated_by"
+	skipRowsOption            = "skip_rows"
+	splitFileOption           = "split_file"
+	diskQuotaOption           = "disk_quota"
+	threadOption              = "thread"
+	maxWriteSpeedOption       = "max_write_speed"
+	checksumTableOption       = "checksum_table"
+	recordErrorsOption        = "record_errors"
+	detachedOption            = "detached"
+	// if 'import mode' enabled, TiKV will:
+	//  - set level0_stop_writes_trigger = max(old, 1 << 30)
+	//  - set level0_slowdown_writes_trigger = max(old, 1 << 30)
+	//  - set soft_pending_compaction_bytes_limit = 0,
+	//  - set hard_pending_compaction_bytes_limit = 0,
+	//  - will not trigger flow control when SST count in L0 is large
+	//  - will not trigger region split, it might cause some region became
+	//    very large and be a hotspot, might cause latency spike.
+	//
+	// default false for local sort, true for global sort.
 	disableTiKVImportModeOption = "disable_tikv_import_mode"
 	cloudStorageURIOption       = "cloud_storage_uri"
 	disablePrecheckOption       = "disable_precheck"
@@ -780,6 +790,9 @@ func (p *Plan) adjustOptions(targetNodeCPUCnt int) {
 			zap.Int("before", p.ThreadCnt), zap.Int("after", limit))
 		p.ThreadCnt = limit
 	}
+	if p.IsGlobalSort() {
+		p.DisableTiKVImportMode = true
+	}
 }
 
 func (p *Plan) initParameters(plan *plannercore.ImportInto) error {
diff --git a/pkg/metrics/disttask.go b/pkg/metrics/disttask.go
index ac9b8d767a..c27073bddf 100644
--- a/pkg/metrics/disttask.go
+++ b/pkg/metrics/disttask.go
@@ -80,10 +80,11 @@ func UpdateMetricsForAddTask(task *proto.TaskBase) {
 }
 
 // UpdateMetricsForScheduleTask update metrics when a task is added
-func UpdateMetricsForScheduleTask(id int64, taskType proto.TaskType) {
-	DistTaskGauge.WithLabelValues(taskType.String(), WaitingStatus).Dec()
-	DistTaskStartTimeGauge.DeleteLabelValues(taskType.String(), WaitingStatus, fmt.Sprint(id))
-	DistTaskStartTimeGauge.WithLabelValues(taskType.String(), SchedulingStatus, fmt.Sprint(id)).SetToCurrentTime()
+func UpdateMetricsForScheduleTask(task *proto.TaskBase) {
+	DistTaskGauge.WithLabelValues(task.Type.String(), WaitingStatus).Dec()
+	DistTaskGauge.WithLabelValues(task.Type.String(), SchedulingStatus).Inc()
+	DistTaskStartTimeGauge.DeleteLabelValues(task.Type.String(), WaitingStatus, fmt.Sprint(task.ID))
+	DistTaskStartTimeGauge.WithLabelValues(task.Type.String(), SchedulingStatus, fmt.Sprint(task.ID)).SetToCurrentTime()
 }
 
 // UpdateMetricsForRunTask update metrics when a task starts running
diff --git a/pkg/metrics/grafana/tidb.json b/pkg/metrics/grafana/tidb.json
index 6af04afdf7..c1ec5ade91 100644
--- a/pkg/metrics/grafana/tidb.json
+++ b/pkg/metrics/grafana/tidb.json
@@ -15769,7 +15769,7 @@
           "targets": [
             {
               "exemplar": true,
-              "expr": "sum(tidb_disttask_subtasks{k8s_cluster=\"$k8s_cluster\", tidb_cluster=\"$tidb_cluster\", instance=~\"$instance\", status=~\"pending|running|reverting|revert_pending|paused\"}) by (exec_id)",
+              "expr": "sum(tidb_disttask_subtasks{k8s_cluster=\"$k8s_cluster\", tidb_cluster=\"$tidb_cluster\", instance=~\"$instance\", status=~\"pending|running|failed|canceled|paused\"}) by (exec_id)",
               "interval": "",
               "legendFormat": "",
               "queryType": "randomWalk",


Test patch:
diff --git a/pkg/executor/aggregate/agg_spill_test.go b/pkg/executor/aggregate/agg_spill_test.go
index 253856249c..600e735a75 100644
--- a/pkg/executor/aggregate/agg_spill_test.go
+++ b/pkg/executor/aggregate/agg_spill_test.go
@@ -150,6 +150,7 @@ func generateResult(t *testing.T, ctx *mock.Context, dataSource *testutil.MockDa
 			resultRows = append(resultRows, chk.GetRow(i))
 		}
 	}
+	require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 	aggExec.Close()
 
 	require.False(t, aggExec.IsSpillTriggeredForTest())
@@ -315,6 +316,7 @@ func executeCorrecResultTest(t *testing.T, ctx *mock.Context, aggExec *aggregate
 			resultRows = append(resultRows, chk.GetRow(i))
 		}
 	}
+	require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 	aggExec.Close()
 
 	require.True(t, aggExec.IsSpillTriggeredForTest())
@@ -351,6 +353,7 @@ func fallBackActionTest(t *testing.T) {
 		}
 		chk.Reset()
 	}
+	require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 	aggExec.Close()
 	require.Less(t, 0, newRootExceedAction.GetTriggeredNum())
 }
@@ -373,6 +376,7 @@ func randomFailTest(t *testing.T, ctx *mock.Context, aggExec *aggregate.HashAggE
 	go func() {
 		time.Sleep(time.Duration(rand.Int31n(300)) * time.Millisecond)
 		once.Do(func() {
+			require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 			aggExec.Close()
 		})
 		goRoutineWaiter.Done()
@@ -382,6 +386,7 @@ func randomFailTest(t *testing.T, ctx *mock.Context, aggExec *aggregate.HashAggE
 		err := aggExec.Next(tmpCtx, chk)
 		if err != nil {
 			once.Do(func() {
+				require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 				err = aggExec.Close()
 				require.Equal(t, nil, err)
 			})
@@ -393,6 +398,7 @@ func randomFailTest(t *testing.T, ctx *mock.Context, aggExec *aggregate.HashAggE
 		chk.Reset()
 	}
 	once.Do(func() {
+		require.False(t, aggExec.IsInvalidMemoryUsageTrackingForTest())
 		aggExec.Close()
 	})
 }
diff --git a/pkg/executor/importer/import_test.go b/pkg/executor/importer/import_test.go
index 9ebd90a276..014a001c4d 100644
--- a/pkg/executor/importer/import_test.go
+++ b/pkg/executor/importer/import_test.go
@@ -189,11 +189,17 @@ func TestAdjustOptions(t *testing.T) {
 	plan.adjustOptions(16)
 	require.Equal(t, 16, plan.ThreadCnt)
 	require.Equal(t, config.ByteSize(10), plan.MaxWriteSpeed) // not adjusted
+	require.False(t, plan.DisableTiKVImportMode)
 
 	plan.ThreadCnt = 100000000
 	plan.DataSourceType = DataSourceTypeQuery
 	plan.adjustOptions(16)
 	require.Equal(t, 32, plan.ThreadCnt)
+	require.False(t, plan.DisableTiKVImportMode)
+
+	plan.CloudStorageURI = "s3://bucket/path"
+	plan.adjustOptions(16)
+	require.True(t, plan.DisableTiKVImportMode)
 }
 
 func TestAdjustDiskQuota(t *testing.T) {
diff --git a/tests/realtikvtest/importintotest3/from_server_test.go b/tests/realtikvtest/importintotest3/from_server_test.go
index d41e44afe2..544446e9d7 100644
--- a/tests/realtikvtest/importintotest3/from_server_test.go
+++ b/tests/realtikvtest/importintotest3/from_server_test.go
@@ -73,4 +73,5 @@ func (s *mockGCSSuite) TestImportFromServer() {
 	var taskMeta importinto.TaskMeta
 	require.NoError(s.T(), json.Unmarshal(task.Meta, &taskMeta))
 	require.Len(s.T(), taskMeta.ChunkMap, 2)
+	require.False(s.T(), taskMeta.Plan.DisableTiKVImportMode)
 }
diff --git a/tests/realtikvtest/importintotest4/global_sort_test.go b/tests/realtikvtest/importintotest4/global_sort_test.go
index 7f9a8c6982..f6b03c8cab 100644
--- a/tests/realtikvtest/importintotest4/global_sort_test.go
+++ b/tests/realtikvtest/importintotest4/global_sort_test.go
@@ -96,6 +96,7 @@ func (s *mockGCSSuite) TestGlobalSortBasic() {
 	taskMeta := importinto.TaskMeta{}
 	s.NoError(json.Unmarshal(task.Meta, &taskMeta))
 	urlEqual(s.T(), redactedSortStorageURI, taskMeta.Plan.CloudStorageURI)
+	require.True(s.T(), taskMeta.Plan.DisableTiKVImportMode)
 
 	// merge-sort data kv
 	s.tk.MustExec("truncate table t")