Repository: pingcap/tidb
Pull request #48258 (2023-11-03 02:59:58)

importinto: calculate subtask count for global sort (#48169)
This is an automated cherry-pick of #48169

<!--

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: ref https://github.com/pingcap/tidb/issues/47572

Problem Summary:

### What is changed and how it works?
for global sort, make sure subtask count is a multiple of execute nodes count
### Check List

Tests <!-- At least one of them must be included. -->

- [x] Unit test
- [ ] Integration test
- [ ] Manual test (add detailed scripts or steps below)
- [ ] 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
```


ti-chi-bot[bot]: This cherry pick PR is for a release branch and has not yet been approved by triage owners. 
Adding the `do-not-merge/cherry-pick-not-approved` label.

To merge this cherry pick:
1. It must be approved by the approvers firstly.
2. **AFTER** it has been approved by approvers, please wait for the cherry-pick merging approval from triage owners.


<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository.
</details>
ti-chi-bot[bot]: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by:
**Once this PR has been reviewed and has the lgtm label**, please assign [d3hunter](https://github.com/d3hunter) for approval. For more information see [the Kubernetes Code Review Process](https://book.prow.tidb.net/#/workflows/pr).

The full list of commands accepted by this bot can be found [here](https://prow.tidb.net/command-help?repo=pingcap%2Ftidb).

<details open>
Needs approval from an approver in each of these files:

- **[pkg/disttask/OWNERS](https://github.com/pingcap/tidb/blob/release-7.5/pkg/disttask/OWNERS)**
- **[pkg/executor/importer/OWNERS](https://github.com/pingcap/tidb/blob/release-7.5/pkg/executor/importer/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["d3hunter"]} -->
ti-chi-bot[bot]: @ti-chi-bot: The following test **failed**, say `/retest` to rerun all failed tests or `/retest-required` to rerun all mandatory failed tests:

Test name | Commit | Details | Required | Rerun command
--- | --- | --- | --- | ---
idc-jenkins-ci-tidb/unit-test | 7326fe6a7cfed88a90bf95c4a38d042be4f7532d | [link](https://do.pingcap.net/jenkins/job/pingcap/job/tidb/job/release-7.5/job/ghpr_unit_test/190/display/redirect) | true | `/test unit-test`

[Full PR test history](https://prow.tidb.net/pr-history?org=pingcap&repo=tidb&pr=48258). [Your PR dashboard](https://prow.tidb.net/pr?query=is%3Apr%20state%3Aopen%20author%3Ati-chi-bot).

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
ti-chi-bot[bot]: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/pingcap/tidb/pull/48258#pullrequestreview-1711794636" title="Approved">D3Hunter</a>*

The full list of commands accepted by this bot can be found [here](https://prow.tidb.net/command-help?repo=pingcap%2Ftidb).

The pull request process is described [here](https://book.prow.tidb.net/#/workflows/pr)

<details >
Needs approval from an approver in each of these files:

- ~~[pkg/disttask/OWNERS](https://github.com/pingcap/tidb/blob/release-7.5/pkg/disttask/OWNERS)~~ [D3Hunter]
- ~~[pkg/executor/importer/OWNERS](https://github.com/pingcap/tidb/blob/release-7.5/pkg/executor/importer/OWNERS)~~ [D3Hunter]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
ti-chi-bot[bot]: [LGTM Timeline notifier]
---
Timeline:

- `2023-11-03 03:53:18.980558127 +0000 UTC m=+3184396.567668273`: :ballot_box_with_check: agreed by [D3Hunter](https://github.com/D3Hunter).
ti-chi-bot[bot]: [LGTM Timeline notifier]
---
Timeline:

- `2023-11-03 03:53:18.980558127 +0000 UTC m=+3184396.567668273`: :ballot_box_with_check: agreed by [D3Hunter](https://github.com/D3Hunter).
- `2023-11-03 04:00:12.298084662 +0000 UTC m=+3184809.885194807`: :ballot_box_with_check: agreed by [GMHDBJD](https://github.com/GMHDBJD).
GMHDBJD: /retest

Patch:
diff --git a/pkg/disttask/framework/planner/plan.go b/pkg/disttask/framework/planner/plan.go
index fae9b2b9b9..e3d94a7c95 100644
--- a/pkg/disttask/framework/planner/plan.go
+++ b/pkg/disttask/framework/planner/plan.go
@@ -37,6 +37,7 @@ type PlanCtx struct {
 	PreviousSubtaskMetas map[proto.Step][][]byte
 	GlobalSort           bool
 	NextTaskStep         proto.Step
+	ExecuteNodesCnt      int
 }
 
 // LogicalPlan represents a logical plan in distribute framework.
diff --git a/pkg/disttask/importinto/dispatcher.go b/pkg/disttask/importinto/dispatcher.go
index b141f1e46c..ab8fd81bec 100644
--- a/pkg/disttask/importinto/dispatcher.go
+++ b/pkg/disttask/importinto/dispatcher.go
@@ -299,12 +299,18 @@ func (dsp *ImportDispatcherExt) OnNextSubtasksBatch(
 		return nil, errors.Errorf("unknown step %d", gTask.Step)
 	}
 
+	eligibleInstances, err := dsp.GetEligibleInstances(ctx, gTask)
+	if err != nil {
+		logger.Warn("failed to get eligible instances", zap.Error(err))
+	}
+
 	planCtx := planner.PlanCtx{
 		Ctx:                  ctx,
 		TaskID:               gTask.ID,
 		PreviousSubtaskMetas: previousSubtaskMetas,
 		GlobalSort:           dsp.GlobalSort,
 		NextTaskStep:         nextStep,
+		ExecuteNodesCnt:      len(eligibleInstances),
 	}
 	logicalPlan := &LogicalPlan{}
 	if err := logicalPlan.FromTaskMeta(gTask.Meta); err != nil {
diff --git a/pkg/disttask/importinto/planner.go b/pkg/disttask/importinto/planner.go
index 6ea23054b4..53ffe2e8c3 100644
--- a/pkg/disttask/importinto/planner.go
+++ b/pkg/disttask/importinto/planner.go
@@ -109,7 +109,7 @@ func (p *LogicalPlan) ToPhysicalPlan(planCtx planner.PlanCtx) (*planner.Physical
 	// we only generate needed plans for the next step.
 	switch planCtx.NextTaskStep {
 	case StepImport, StepEncodeAndSort:
-		specs, err := generateImportSpecs(planCtx.Ctx, p)
+		specs, err := generateImportSpecs(planCtx, p)
 		if err != nil {
 			return nil, err
 		}
@@ -249,7 +249,7 @@ func buildController(plan *importer.Plan, stmt string) (*importer.LoadDataContro
 	return controller, nil
 }
 
-func generateImportSpecs(ctx context.Context, p *LogicalPlan) ([]planner.PipelineSpec, error) {
+func generateImportSpecs(pCtx planner.PlanCtx, p *LogicalPlan) ([]planner.PipelineSpec, error) {
 	var chunkMap map[int32][]Chunk
 	if len(p.ChunkMap) > 0 {
 		chunkMap = p.ChunkMap
@@ -258,11 +258,12 @@ func generateImportSpecs(ctx context.Context, p *LogicalPlan) ([]planner.Pipelin
 		if err2 != nil {
 			return nil, err2
 		}
-		if err2 = controller.InitDataFiles(ctx); err2 != nil {
+		if err2 = controller.InitDataFiles(pCtx.Ctx); err2 != nil {
 			return nil, err2
 		}
 
-		engineCheckpoints, err2 := controller.PopulateChunks(ctx)
+		controller.SetExecuteNodeCnt(pCtx.ExecuteNodesCnt)
+		engineCheckpoints, err2 := controller.PopulateChunks(pCtx.Ctx)
 		if err2 != nil {
 			return nil, err2
 		}
diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel
index 07f6fb8622..5e1afc2d83 100644
--- a/pkg/executor/importer/BUILD.bazel
+++ b/pkg/executor/importer/BUILD.bazel
@@ -88,7 +88,7 @@ go_test(
     embed = [":importer"],
     flaky = True,
     race = "on",
-    shard_count = 19,
+    shard_count = 20,
     deps = [
         "//br/pkg/errors",
         "//br/pkg/lightning/backend/encode",
diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go
index 4d4cc844f8..2dd931e828 100644
--- a/pkg/executor/importer/import.go
+++ b/pkg/executor/importer/import.go
@@ -264,6 +264,8 @@ type LoadDataController struct {
 	dataFiles []*mydump.SourceFileMeta
 	// GlobalSortStore is used to store sorted data when using global sort.
 	GlobalSortStore storage.ExternalStorage
+	// ExecuteNodesCnt is the count of execute nodes.
+	ExecuteNodesCnt int
 }
 
 func getImportantSysVars(sctx sessionctx.Context) map[string]string {
@@ -452,10 +454,11 @@ func NewLoadDataController(plan *Plan, tbl table.Table, astArgs *ASTArgs) (*Load
 	fullTableName := tbl.Meta().Name.String()
 	logger := log.L().With(zap.String("table", fullTableName))
 	c := &LoadDataController{
-		Plan:    plan,
-		ASTArgs: astArgs,
-		Table:   tbl,
-		logger:  logger,
+		Plan:            plan,
+		ASTArgs:         astArgs,
+		Table:           tbl,
+		logger:          logger,
+		ExecuteNodesCnt: 1,
 	}
 	if err := c.checkFieldParams(); err != nil {
 		return nil, err
diff --git a/pkg/executor/importer/table_import.go b/pkg/executor/importer/table_import.go
index ace1df765e..016c688aeb 100644
--- a/pkg/executor/importer/table_import.go
+++ b/pkg/executor/importer/table_import.go
@@ -299,12 +299,11 @@ func (ti *TableImporter) getKVEncoder(chunk *checkpoints.ChunkCheckpoint) (KVEnc
 	return NewTableKVEncoder(cfg, ti)
 }
 
-func (e *LoadDataController) getAdjustedMaxEngineSize() int64 {
+func (e *LoadDataController) calculateSubtaskCnt() int {
 	// we want to split data files into subtask of size close to MaxEngineSize to reduce range overlap,
 	// and evenly distribute them to subtasks.
-	// so we adjust MaxEngineSize to make sure each subtask has a similar amount of data to import.
-	// we calculate subtask count first by round(TotalFileSize / maxEngineSize), then adjust maxEngineSize
-	//
+	// we calculate subtask count first by round(TotalFileSize / maxEngineSize)
+
 	// AllocateEngineIDs is using ceil() to calculate subtask count, engine size might be too small in some case,
 	// such as 501G data, maxEngineSize will be about 250G, so we don't relay on it.
 	// see https://github.com/pingcap/tidb/blob/b4183e1dc9bb01fb81d3aa79ca4b5b74387c6c2a/br/pkg/lightning/mydump/region.go#L109
@@ -315,13 +314,33 @@ func (e *LoadDataController) getAdjustedMaxEngineSize() int64 {
 	// [750, 1250)            2    [375, 625)
 	// [1250, 1750)           3    [416, 583)
 	// [1750, 2250)           4    [437, 562)
-	maxEngineSize := int64(e.MaxEngineSize)
+	var (
+		subtaskCount  float64
+		maxEngineSize = int64(e.MaxEngineSize)
+	)
 	if e.TotalFileSize <= maxEngineSize {
-		return e.TotalFileSize
+		subtaskCount = 1
+	} else {
+		subtaskCount = math.Round(float64(e.TotalFileSize) / float64(e.MaxEngineSize))
+	}
+
+	// for global sort task, since there is no overlap,
+	// we make sure subtask count is a multiple of execute nodes count
+	if e.IsGlobalSort() && e.ExecuteNodesCnt > 0 {
+		subtaskCount = math.Ceil(subtaskCount/float64(e.ExecuteNodesCnt)) * float64(e.ExecuteNodesCnt)
 	}
-	subtaskCount := math.Round(float64(e.TotalFileSize) / float64(maxEngineSize))
-	adjusted := math.Ceil(float64(e.TotalFileSize) / subtaskCount)
-	return int64(adjusted)
+	return int(subtaskCount)
+}
+
+func (e *LoadDataController) getAdjustedMaxEngineSize() int64 {
+	subtaskCount := e.calculateSubtaskCnt()
+	// we adjust MaxEngineSize to make sure each subtask has a similar amount of data to import.
+	return int64(math.Ceil(float64(e.TotalFileSize) / float64(subtaskCount)))
+}
+
+// SetExecuteNodeCnt sets the execute node count.
+func (e *LoadDataController) SetExecuteNodeCnt(cnt int) {
+	e.ExecuteNodesCnt = cnt
 }
 
 // PopulateChunks populates chunks from table regions.


Test patch:
diff --git a/pkg/executor/importer/table_import_test.go b/pkg/executor/importer/table_import_test.go
index 3fc5688a7d..e72b6c009f 100644
--- a/pkg/executor/importer/table_import_test.go
+++ b/pkg/executor/importer/table_import_test.go
@@ -80,29 +80,92 @@ func TestPrepareSortDir(t *testing.T) {
 	require.Nil(t, info)
 }
 
+func TestCalculateSubtaskCnt(t *testing.T) {
+	tests := []struct {
+		totalSize       int64
+		maxEngineSize   config.ByteSize
+		executeNodeCnt  int
+		cloudStorageURL string
+		want            int
+	}{
+		{1, 500, 0, "", 1},
+		{499, 500, 1, "", 1},
+		{500, 500, 2, "", 1},
+		{749, 500, 3, "", 1},
+		{750, 500, 4, "", 2},
+		{1249, 500, 5, "", 2},
+		{1250, 500, 6, "", 3},
+		{100, 30, 7, "", 3},
+
+		{1, 500, 0, "url", 1},
+		{499, 500, 1, "url", 1},
+		{500, 500, 2, "url", 2},
+		{749, 500, 3, "url", 3},
+		{750, 500, 4, "url", 4},
+		{1249, 500, 5, "url", 5},
+		{1250, 500, 6, "url", 6},
+		{100, 30, 2, "url", 4},
+		{400, 99, 3, "url", 6},
+		{500, 100, 5, "url", 5},
+		{500, 200, 5, "url", 5},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("%d/%d", tt.totalSize, tt.maxEngineSize), func(t *testing.T) {
+			e := &LoadDataController{
+				Plan: &Plan{
+					MaxEngineSize:   tt.maxEngineSize,
+					TotalFileSize:   tt.totalSize,
+					CloudStorageURI: tt.cloudStorageURL,
+				},
+				ExecuteNodesCnt: tt.executeNodeCnt,
+			}
+			if got := e.calculateSubtaskCnt(); got != tt.want {
+				t.Errorf("calculateSubtaskCnt() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
 func TestLoadDataControllerGetAdjustedMaxEngineSize(t *testing.T) {
 	tests := []struct {
-		totalSize     int64
-		maxEngineSize config.ByteSize
-		want          int64
+		totalSize       int64
+		maxEngineSize   config.ByteSize
+		executeNodeCnt  int
+		cloudStorageURL string
+		want            int64
 	}{
-		{1, 500, 1},
-		{499, 500, 499},
-		{500, 500, 500},
-		{749, 500, 749},
-		{750, 500, 375},
-		{1249, 500, 625},
-		{1250, 500, 417},
+		{1, 500, 0, "", 1},
+		{499, 500, 1, "", 499},
+		{500, 500, 2, "", 500},
+		{749, 500, 3, "", 749},
+		{750, 500, 4, "", 375},
+		{1249, 500, 5, "", 625},
+		{1250, 500, 6, "", 417},
 		// ceil(100/3)
-		{100, 30, 34},
+		{100, 30, 7, "", 34},
+
+		{1, 500, 0, "url", 1},
+		{499, 500, 1, "url", 499},
+		{500, 500, 2, "url", 250},
+		{749, 500, 3, "url", 250},
+		{750, 500, 4, "url", 188},
+		{1249, 500, 5, "url", 250},
+		{1250, 500, 6, "url", 209},
+		{100, 30, 2, "url", 25},
+		{400, 99, 3, "url", 67},
+		{500, 100, 5, "url", 100},
+		{500, 200, 5, "url", 100},
+		{500, 100, 1, "url", 100},
 	}
 	for _, tt := range tests {
 		t.Run(fmt.Sprintf("%d/%d", tt.totalSize, tt.maxEngineSize), func(t *testing.T) {
 			e := &LoadDataController{
 				Plan: &Plan{
-					MaxEngineSize: tt.maxEngineSize,
-					TotalFileSize: tt.totalSize,
+					MaxEngineSize:   tt.maxEngineSize,
+					TotalFileSize:   tt.totalSize,
+					CloudStorageURI: tt.cloudStorageURL,
 				},
+				ExecuteNodesCnt: tt.executeNodeCnt,
 			}
 			if got := e.getAdjustedMaxEngineSize(); got != tt.want {
 				t.Errorf("getAdjustedMaxEngineSize() = %v, want %v", got, tt.want)