dxf: add api to manually tune the calculation of resource related params

pkg/ddl/reorg_util.go

@@ -95,9 +95,22 @@ func initJobReorgMetaFromVariables(ctx context.Context, job *model.Job, tbl tabl
 		}
 	})
 
+	var (
+		autoConc, autoMaxNode int
+		factorField           = zap.Skip()
+	)
+	if shouldCalResource {
+		factors, err := dxfhandle.GetScheduleTuneFactors(ctx, sctx.GetStore().GetKeyspace())
+		if err != nil {
+			return err
+		}
+		calc := scheduler.NewRCCalcForAddIndex(tableSizeInBytes, cpuNum, factors)
+		autoConc = calc.CalcConcurrency()
+		autoMaxNode = calc.CalcMaxNodeCountForAddIndex()
+		factorField = zap.Float64("amplifyFactor", factors.AmplifyFactor)
+	}
 	if setReorgParam {
 		if shouldCalResource && setDistTaskParam {
-			autoConc := scheduler.CalcConcurrencyByDataSize(tableSizeInBytes, cpuNum)
 			m.SetConcurrency(autoConc)
 		} else {
 			if sv, ok := sessVars.GetSystemVar(vardef.TiDBDDLReorgWorkerCount); ok {
@@ -115,7 +128,7 @@ func initJobReorgMetaFromVariables(ctx context.Context, job *model.Job, tbl tabl
 		m.IsFastReorg = vardef.EnableFastReorg.Load()
 		m.TargetScope = dxfhandle.GetTargetScope()
 		if shouldCalResource {
-			m.MaxNodeCount = scheduler.CalcMaxNodeCountByTableSize(tableSizeInBytes, cpuNum)
+			m.MaxNodeCount = autoMaxNode
 		} else {
 			if sv, ok := sessVars.GetSystemVar(vardef.TiDBMaxDistTaskNodes); ok {
 				m.MaxNodeCount = variable.TidbOptInt(sv, 0)
@@ -152,6 +165,7 @@ func initJobReorgMetaFromVariables(ctx context.Context, job *model.Job, tbl tabl
 		zap.String("tableSizeInBytes", units.BytesSize(float64(tableSizeInBytes))),
 		zap.Int("concurrency", m.GetConcurrency()),
 		zap.Int("batchSize", m.GetBatchSize()),
+		factorField,
 	)
 	return nil
 }

pkg/disttask/framework/handle/BUILD.bazel

@@ -16,6 +16,7 @@ go_library(
         "//pkg/disttask/framework/storage",
         "//pkg/domain/infosync",
         "//pkg/kv",
+        "//pkg/meta",
         "//pkg/meta/tidbvar",
         "//pkg/metrics",
         "//pkg/sessionctx",
@@ -27,6 +28,7 @@ go_library(
         "//pkg/util/sqlexec",
         "@com_github_docker_go_units//:go-units",
         "@com_github_pingcap_errors//:errors",
+        "@com_github_tikv_client_go_v2//util",
         "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
     ],
@@ -42,12 +44,14 @@ go_test(
     ],
     embed = [":handle"],
     flaky = True,
-    shard_count = 8,
+    shard_count = 9,
     deps = [
         "//pkg/config/kerneltype",
         "//pkg/disttask/framework/proto",
         "//pkg/disttask/framework/schstatus",
         "//pkg/disttask/framework/storage",
+        "//pkg/kv",
+        "//pkg/meta",
         "//pkg/sessionctx",
         "//pkg/sessionctx/vardef",
         "//pkg/testkit",

pkg/disttask/framework/handle/handle.go

@@ -30,13 +30,15 @@ import (
 	"github.com/pingcap/tidb/pkg/disttask/framework/schstatus"
 	"github.com/pingcap/tidb/pkg/disttask/framework/storage"
 	"github.com/pingcap/tidb/pkg/kv"
+	"github.com/pingcap/tidb/pkg/meta"
 	"github.com/pingcap/tidb/pkg/meta/tidbvar"
 	"github.com/pingcap/tidb/pkg/metrics"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/sessionctx/vardef"
 	"github.com/pingcap/tidb/pkg/util/backoff"
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
+	"github.com/tikv/client-go/v2/util"
 	"go.uber.org/atomic"
 	"go.uber.org/zap"
 )
@@ -298,6 +300,7 @@ func GetCloudStorageURI(ctx context.Context, store kv.Storage) string {
 
 // UpdatePauseScaleInFlag updates the pause scale-in flag.
 func UpdatePauseScaleInFlag(ctx context.Context, flag *schstatus.TTLFlag) error {
+	ctx = util.WithInternalSourceType(ctx, kv.InternalDistTask)
 	manager, err := storage.GetTaskManager()
 	if err != nil {
 		return err
@@ -314,6 +317,34 @@ func UpdatePauseScaleInFlag(ctx context.Context, flag *schstatus.TTLFlag) error
 	})
 }
 
+// GetScheduleTuneFactors gets the schedule tune factors for a keyspace.
+// if not set or expired, it returns the default tune factors.
+func GetScheduleTuneFactors(ctx context.Context, keyspace string) (*schstatus.TuneFactors, error) {
+	ctx = util.WithInternalSourceType(ctx, kv.InternalDistTask)
+	mgr, err := storage.GetDXFSvcTaskMgr()
+	if err != nil {
+		return nil, err
+	}
+	var factors *schstatus.TTLTuneFactors
+	if err = mgr.WithNewSession(func(se sessionctx.Context) error {
+		return kv.RunInNewTxn(ctx, se.GetStore(), true, func(_ context.Context, txn kv.Transaction) error {
+			mutator := meta.NewMutator(txn)
+			var err2 error
+			factors, err2 = mutator.GetDXFScheduleTuneFactors(keyspace)
+			if err2 != nil {
+				return err2
+			}
+			return nil
+		})
+	}); err != nil {
+		return nil, err
+	}
+	if factors == nil || factors.ExpireTime.Before(time.Now()) {
+		return schstatus.GetDefaultTuneFactors(), nil
+	}
+	return &factors.TuneFactors, nil
+}
+
 func init() {
 	// domain will init this var at runtime, we store it here for test, as some
 	// test might not start domain.

pkg/disttask/framework/handle/status.go

@@ -24,16 +24,19 @@ import (
 	"github.com/pingcap/tidb/pkg/disttask/framework/schstatus"
 	"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/meta/tidbvar"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/util/cpu"
 	disttaskutil "github.com/pingcap/tidb/pkg/util/disttask"
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
+	"github.com/tikv/client-go/v2/util"
 )
 
 // GetScheduleStatus returns the schedule status.
 func GetScheduleStatus(ctx context.Context) (*schstatus.Status, error) {
+	ctx = util.WithInternalSourceType(ctx, kv.InternalDistTask)
 	manager, err := storage.GetTaskManager()
 	if err != nil {
 		return nil, err
@@ -178,6 +181,7 @@ func GetScheduleFlags(ctx context.Context, manager *storage.TaskManager) (map[sc
 
 func getPauseScaleInFlag(ctx context.Context, manager *storage.TaskManager) (*schstatus.TTLFlag, error) {
 	flag := &schstatus.TTLFlag{}
+	ctx = util.WithInternalSourceType(ctx, kv.InternalDistTask)
 	if err := manager.WithNewSession(func(se sessionctx.Context) error {
 		rs, err2 := sqlexec.ExecSQL(ctx, se.GetSQLExecutor(), `SELECT VARIABLE_VALUE from mysql.tidb WHERE VARIABLE_NAME = %?`,
 			tidbvar.DXFSchedulePauseScaleIn)

pkg/disttask/framework/handle/status_testkit_test.go

@@ -26,6 +26,8 @@ import (
 	"github.com/pingcap/tidb/pkg/disttask/framework/proto"
 	"github.com/pingcap/tidb/pkg/disttask/framework/schstatus"
 	"github.com/pingcap/tidb/pkg/disttask/framework/storage"
+	"github.com/pingcap/tidb/pkg/kv"
+	"github.com/pingcap/tidb/pkg/meta"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
@@ -43,7 +45,7 @@ func TestGetScheduleStatus(t *testing.T) {
 	}, 1, 1, time.Second)
 	defer pool.Close()
 	storage.SetTaskManager(storage.NewTaskManager(pool))
-	ctx := util.WithInternalSourceType(context.Background(), "handle_test")
+	ctx := context.Background()
 	status, err := handle.GetScheduleStatus(ctx)
 	require.NoError(t, err)
 	require.Equal(t, &schstatus.Status{
@@ -105,7 +107,7 @@ func TestNodeInfoAndBusyNodes(t *testing.T) {
 
 func TestScheduleFlag(t *testing.T) {
 	testkit.CreateMockStore(t)
-	ctx := util.WithInternalSourceType(context.Background(), "handle_test")
+	ctx := context.Background()
 	taskMgr, err := storage.GetTaskManager()
 	require.NoError(t, err)
 	// not set
@@ -116,7 +118,7 @@ func TestScheduleFlag(t *testing.T) {
 	// in CI, the stored and unmarshalled time might have different time zone,
 	// so unify to UTC.
 	expireTime := time.Unix(time.Now().Add(time.Hour).Unix(), 0).In(time.UTC)
-	flag := schstatus.TTLFlag{Enabled: true, TTL: time.Hour, ExpireTime: expireTime}
+	flag := schstatus.TTLFlag{Enabled: true, TTLInfo: schstatus.TTLInfo{TTL: time.Hour, ExpireTime: expireTime}}
 	require.NoError(t, handle.UpdatePauseScaleInFlag(ctx, &flag))
 	flags, err = handle.GetScheduleFlags(ctx, taskMgr)
 	require.NoError(t, err)
@@ -137,3 +139,37 @@ func TestScheduleFlag(t *testing.T) {
 	require.NoError(t, err)
 	require.Empty(t, flags)
 }
+
+func TestGetScheduleTuneFactors(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	ctx := context.Background()
+	// factors not set
+	factors, err := handle.GetScheduleTuneFactors(ctx, store.GetKeyspace())
+	require.NoError(t, err)
+	require.EqualValues(t, schstatus.GetDefaultTuneFactors(), factors)
+	// non expired factors
+	ctx2 := util.WithInternalSourceType(ctx, kv.InternalDistTask)
+	// in CI, the stored and unmarshalled time might have different time zone,
+	// so unify to UTC.
+	expireTime := time.Unix(time.Now().Add(time.Hour).Unix(), 0).In(time.UTC)
+	ttlFactors := schstatus.TTLTuneFactors{
+		TTLInfo:     schstatus.TTLInfo{TTL: time.Hour, ExpireTime: expireTime},
+		TuneFactors: schstatus.TuneFactors{AmplifyFactor: 1.5},
+	}
+	require.NoError(t, kv.RunInNewTxn(ctx2, store, true, func(ctx context.Context, txn kv.Transaction) error {
+		m := meta.NewMutator(txn)
+		return m.SetDXFScheduleTuneFactors(store.GetKeyspace(), &ttlFactors)
+	}))
+	factors, err = handle.GetScheduleTuneFactors(ctx, store.GetKeyspace())
+	require.NoError(t, err)
+	require.EqualValues(t, &ttlFactors.TuneFactors, factors)
+	// expired factors
+	ttlFactors.ExpireTime = time.Now().Add(-time.Hour)
+	require.NoError(t, kv.RunInNewTxn(ctx2, store, true, func(ctx context.Context, txn kv.Transaction) error {
+		m := meta.NewMutator(txn)
+		return m.SetDXFScheduleTuneFactors(store.GetKeyspace(), &ttlFactors)
+	}))
+	factors, err = handle.GetScheduleTuneFactors(ctx, store.GetKeyspace())
+	require.NoError(t, err)
+	require.EqualValues(t, schstatus.GetDefaultTuneFactors(), factors)
+}

pkg/disttask/framework/scheduler/BUILD.bazel

@@ -19,6 +19,7 @@ go_library(
         "//pkg/config/kerneltype",
         "//pkg/disttask/framework/handle",
         "//pkg/disttask/framework/proto",
+        "//pkg/disttask/framework/schstatus",
         "//pkg/disttask/framework/storage",
         "//pkg/disttask/framework/taskexecutor/execute",
         "//pkg/domain/infosync",
@@ -60,14 +61,15 @@ go_test(
     embed = [":scheduler"],
     flaky = True,
     race = "off",
-    shard_count = 39,
+    shard_count = 41,
     deps = [
         "//pkg/config",
         "//pkg/config/kerneltype",
         "//pkg/disttask/framework/handle",
         "//pkg/disttask/framework/mock",
         "//pkg/disttask/framework/proto",
         "//pkg/disttask/framework/scheduler/mock",
+        "//pkg/disttask/framework/schstatus",
         "//pkg/disttask/framework/storage",
         "//pkg/disttask/framework/testutil",
         "//pkg/domain/infosync",

pkg/disttask/framework/scheduler/autoscaler.go

@@ -22,6 +22,7 @@ import (
 	"github.com/docker/go-units"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/disttask/framework/handle"
+	"github.com/pingcap/tidb/pkg/disttask/framework/schstatus"
 	"github.com/pingcap/tidb/pkg/disttask/framework/storage"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/sessionctx/vardef"
@@ -38,33 +39,80 @@ const (
 	// Each node should handle at least 2 subtasks, each 100GiB data.
 	// For every additional 200 GiB of data, add 1 node.
 	baseDataSize = 200 * units.GiB
-	// To improve performance for small tasks, we assume that on an 8c machine,
+	// To improve performance for small tasks, we assume that on a 8c machine,
 	// importing 200 GiB of data requires full utilization of a single node’s resources.
 	// Therefore, for every additional 25 GiB, add 1 concurrency unit as an estimate for task concurrency.
-	baseSizePerConc                = 25 * units.GiB
+	baseSizePerConc = 25 * units.GiB
+	// The maximum number of nodes that can be used for add-index.
+	maxNodeCountLimitForAddIndex = 30
+	// The maximum number of nodes that can be used for import-into.
+	// this value is based on previous performance test, for a quite common scenario,
+	// to import 100TiB data within 24 hours, we need about 32 8c nodes.
+	// Note: import speed is affected by many factors, such as row length, index
+	// count, table schema complexity, etc.
 	maxNodeCountLimitForImportInto = 32
 	// this value is calculated by 256/8, we have test on a 8c machine with 256
 	// concurrency, it's fast enough for checksum. we can tune this later if needed.
 	maxDistSQLConcurrencyPerCore = 32
 )
 
-// CalcMaxNodeCountByTableSize calculates the maximum number of nodes to execute DXF based on the table size.
-func CalcMaxNodeCountByTableSize(size int64, coresPerNode int) int {
-	return calcMaxNodeCountBySize(size, coresPerNode, 30)
+// ResourceCalc is used to calculate the resource required for a DXF task
+// in nextgen.
+type ResourceCalc struct {
+	// dataSize is the input data size in bytes.
+	// for import-into, it's the size of the source data.
+	// for add-index, it's the size of the table to be indexed.
+	dataSize int64
+	// nodeCPU is the number of CPU cores of each execution node.
+	nodeCPU int
+	// indexSizeRatio is the ratio of index KV size to data KV size, we get it
+	// through sampling and estimation.
+	// it's used by import-into, for add-index, it's always 0.
+	indexSizeRatio float64
+	// TuneFactors is the tuning factors for resource calculation.
+	factors schstatus.TuneFactors
 }
 
-// CalcMaxNodeCountByDataSize calculates the maximum number of nodes to execute DXF based on the data size.
-func CalcMaxNodeCountByDataSize(size int64, coresPerNode int) int {
-	return calcMaxNodeCountBySize(size, coresPerNode, maxNodeCountLimitForImportInto)
+// NewRCCalcForAddIndex creates a new ResourceCalc for add-index task.
+func NewRCCalcForAddIndex(dataSize int64, nodeCPU int, factors *schstatus.TuneFactors) *ResourceCalc {
+	return NewRCCalc(dataSize, nodeCPU, 0, factors)
 }
 
-func calcMaxNodeCountBySize(size int64, coresPerNode int, factor float64) int {
-	if coresPerNode <= 0 {
+// NewRCCalc creates a new ResourceCalc.
+func NewRCCalc(dataSize int64, nodeCPU int, indexSizeRatio float64, factors *schstatus.TuneFactors) *ResourceCalc {
+	return &ResourceCalc{
+		dataSize:       dataSize,
+		nodeCPU:        nodeCPU,
+		indexSizeRatio: indexSizeRatio,
+		factors:        *factors,
+	}
+}
+
+// CalcMaxNodeCountForAddIndex calculates the maximum number of nodes to execute add-index.
+func (rc *ResourceCalc) CalcMaxNodeCountForAddIndex() int {
+	size := rc.getAmplifiedDataSize()
+	limit := rc.factors.AmplifyFactor * maxNodeCountLimitForAddIndex
+	return rc.calcMaxNodeCountBySize(size, limit)
+}
+
+// CalcMaxNodeCountForImportInto calculates the maximum number of nodes to execute import-into.
+func (rc *ResourceCalc) CalcMaxNodeCountForImportInto() int {
+	size := rc.getAmplifiedDataSize()
+	limit := rc.factors.AmplifyFactor * maxNodeCountLimitForImportInto
+	return rc.calcMaxNodeCountBySize(size, limit)
+}
+
+func (rc *ResourceCalc) getAmplifiedDataSize() int64 {
+	return int64(rc.factors.AmplifyFactor * (1 + rc.indexSizeRatio) * float64(rc.dataSize))
+}
+
+func (rc *ResourceCalc) calcMaxNodeCountBySize(size int64, limit float64) int {
+	if rc.nodeCPU <= 0 {
 		return 0
 	}
-	r := baseCores / float64(coresPerNode)
+	r := baseCores / float64(rc.nodeCPU)
 	nodeCnt := float64(size) * r / baseDataSize
-	nodeCnt = min(nodeCnt, factor*r)
+	nodeCnt = min(nodeCnt, limit*r)
 	nodeCnt = max(nodeCnt, 1)
 	return int(math.Round(nodeCnt))
 }
@@ -91,13 +139,14 @@ func CalcMaxNodeCountByStoresNum(ctx context.Context, store kv.Storage) int {
 	return max(3, len(stores)/3)
 }
 
-// CalcConcurrencyByDataSize calculates the concurrency based on the data size.
-func CalcConcurrencyByDataSize(size int64, coresPerNode int) int {
+// CalcConcurrency calculates the concurrency based on the data size.
+func (rc *ResourceCalc) CalcConcurrency() int {
+	size := rc.getAmplifiedDataSize()
 	if size <= 0 {
 		return 4
 	}
 	concurrency := float64(size) / baseSizePerConc
-	concurrency = min(concurrency, float64(coresPerNode))
+	concurrency = min(concurrency, float64(rc.nodeCPU))
 	concurrency = max(concurrency, 1)
 	return int(math.Round(concurrency))
 }