dxf: add DXF resource limit config (#68251)

ref #68181

Author: D3Hunter

pkg/config/config.go

@@ -97,6 +97,14 @@ const (
 	DefExpensiveTxnTimeThreshold = 600
 	// DefMemoryUsageAlarmRatio is the threshold triggering an alarm which the memory usage of tidb-server instance exceeds.
 	DefMemoryUsageAlarmRatio = 0.8
+	// DefDXFResourceLimit is the default resource percentage available for DXF.
+	DefDXFResourceLimit = 100
+	// MinDXFResourceLimit is the minimum resource percentage available for DXF.
+	// Keep it at 10 while mysql.dist_framework_meta.cpu_count changes from total
+	// node CPU to usable DXF CPU, so the stored value stays non-zero.
+	MinDXFResourceLimit = 10
+	// MaxDXFResourceLimit is the maximum resource percentage available for DXF.
+	MaxDXFResourceLimit = 100
 	// DefTempDir is the default temporary directory path for TiDB.
 	DefTempDir = "/tmp/tidb"
 	// DefPluginAuditLogBufferSize is the default buffer size for plugin audit log.
@@ -201,6 +209,7 @@ type Config struct {
 	TiDBEdition                string                  `toml:"tidb-edition" json:"tidb-edition"`
 	TiDBReleaseVersion         string                  `toml:"tidb-release-version" json:"tidb-release-version"`
 	DeployMode                 deploymode.Mode         `toml:"deploy-mode" json:"deploy-mode"`
+	DXFResourceLimit           int                     `toml:"dxf-resource-limit" json:"dxf-resource-limit"`
 	KeyspaceName               string                  `toml:"keyspace-name" json:"keyspace-name"`
 	TiKVWorkerURL              string                  `toml:"tikv-worker-url" json:"tikv-worker-url"`
 	Log                        Log                     `toml:"log" json:"log"`
@@ -1041,6 +1050,7 @@ var defaultConf = Config{
 	VersionComment:               "",
 	TiDBReleaseVersion:           "",
 	DeployMode:                   deploymode.Premium,
+	DXFResourceLimit:             DefDXFResourceLimit,
 	RUV2:                         DefaultRUV2Config(),
 	Log: Log{
 		Level:               "info",
@@ -1385,6 +1395,13 @@ func (c *Config) Load(confFile string) error {
 	if !kerneltype.IsNextGen() && metaData.IsDefined("deploy-mode") {
 		return fmt.Errorf("deploy-mode can only be configured for nextgen TiDB")
 	}
+	dxfResourceLimitDefined := metaData.IsDefined("dxf-resource-limit")
+	if !dxfResourceLimitDefined && c.DXFResourceLimit == 0 {
+		c.DXFResourceLimit = DefDXFResourceLimit
+	}
+	if dxfResourceLimitDefined && c.DeployMode != deploymode.PremiumReserved {
+		return fmt.Errorf("dxf-resource-limit can only be configured when deploy-mode is premium_reserved")
+	}
 	if c.TokenLimit == 0 {
 		c.TokenLimit = 1000
 	} else if c.TokenLimit > MaxTokenLimit {
@@ -1456,6 +1473,12 @@ func (c *Config) Valid() error {
 	if !kerneltype.IsNextGen() && c.DeployMode != deploymode.Premium {
 		return fmt.Errorf("deploy-mode can only be configured for nextgen TiDB")
 	}
+	if c.DXFResourceLimit < MinDXFResourceLimit || c.DXFResourceLimit > MaxDXFResourceLimit {
+		return fmt.Errorf("dxf-resource-limit should be between %d and %d", MinDXFResourceLimit, MaxDXFResourceLimit)
+	}
+	if c.DXFResourceLimit != DefDXFResourceLimit && c.DeployMode != deploymode.PremiumReserved {
+		return fmt.Errorf("dxf-resource-limit can only be configured when deploy-mode is premium_reserved")
+	}
 	if c.Store == StoreTypeMockTiKV && !c.Instance.TiDBEnableDDL.Load() {
 		return fmt.Errorf("can't disable DDL on mocktikv")
 	}

pkg/config/config.toml.nextgen.example

@@ -16,6 +16,10 @@ store = "unistore"
 # startup and cannot be changed after it is set.
 deploy-mode = "premium"
 
+# Limits the local TiDB resources available to DXF tasks, [10, 100].
+# This can only be configured when deploy-mode is "premium_reserved".
+# dxf-resource-limit = 100
+
 # TiDB storage path.
 path = "/tmp/tidb"
 

pkg/config/config_test.go

@@ -1052,6 +1052,7 @@ func TestTxnTotalSizeLimitValid(t *testing.T) {
 func TestDeployModeConfig(t *testing.T) {
 	conf := NewConfig()
 	require.Equal(t, deploymode.Premium, conf.DeployMode)
+	require.Equal(t, DefDXFResourceLimit, conf.DXFResourceLimit)
 	require.NoError(t, conf.Valid())
 	conf.DeployMode = deploymode.Mode(100)
 	require.ErrorContains(t, conf.Valid(), "invalid deploy-mode")
@@ -1061,6 +1062,10 @@ func TestDeployModeConfig(t *testing.T) {
 	configFile := filepath.Join(storeDir, "config.toml")
 
 	if kerneltype.IsClassic() {
+		require.NoError(t, os.WriteFile(configFile, []byte(`dxf-resource-limit = 30`), 0644))
+		conf = NewConfig()
+		require.ErrorContains(t, conf.Load(configFile), "dxf-resource-limit can only be configured when deploy-mode is premium_reserved")
+
 		require.NoError(t, os.WriteFile(configFile, []byte(`deploy-mode = "premium"`), 0644))
 		conf = NewConfig()
 		require.ErrorContains(t, conf.Load(configFile), "deploy-mode can only be configured for nextgen TiDB")
@@ -1076,8 +1081,34 @@ func TestDeployModeConfig(t *testing.T) {
 	conf = NewConfig()
 	require.NoError(t, conf.Load(configFile))
 	require.Equal(t, deploymode.PremiumReserved, conf.DeployMode)
+	require.Equal(t, DefDXFResourceLimit, conf.DXFResourceLimit)
+	require.NoError(t, conf.Valid())
+
+	require.NoError(t, os.WriteFile(configFile, []byte(`deploy-mode = "premium_reserved"
+dxf-resource-limit = 30`), 0644))
+	conf = NewConfig()
+	require.NoError(t, conf.Load(configFile))
+	require.Equal(t, deploymode.PremiumReserved, conf.DeployMode)
+	require.Equal(t, 30, conf.DXFResourceLimit)
 	require.NoError(t, conf.Valid())
 
+	require.NoError(t, os.WriteFile(configFile, []byte(`deploy-mode = "premium"
+dxf-resource-limit = 100`), 0644))
+	conf = NewConfig()
+	require.ErrorContains(t, conf.Load(configFile), "dxf-resource-limit can only be configured when deploy-mode is premium_reserved")
+
+	require.NoError(t, os.WriteFile(configFile, []byte(`deploy-mode = "premium_reserved"
+dxf-resource-limit = 9`), 0644))
+	conf = NewConfig()
+	require.NoError(t, conf.Load(configFile))
+	require.ErrorContains(t, conf.Valid(), "dxf-resource-limit should be between 10 and 100")
+
+	require.NoError(t, os.WriteFile(configFile, []byte(`deploy-mode = "premium_reserved"
+dxf-resource-limit = 101`), 0644))
+	conf = NewConfig()
+	require.NoError(t, conf.Load(configFile))
+	require.ErrorContains(t, conf.Valid(), "dxf-resource-limit should be between 10 and 100")
+
 	require.NoError(t, os.WriteFile(configFile, []byte(`deploy-mode = "starter"`), 0644))
 	conf = NewConfig()
 	require.NoError(t, conf.Load(configFile))

pkg/ddl/backfilling_dist_scheduler.go

@@ -774,7 +774,7 @@ func getRangeSplitter(
 			logger.Warn("fail to get region split keys and size", zap.Error(err))
 		}
 	}
-	nodeRc := handle.GetNodeResource()
+	nodeRc := diststorage.GetNodeResource()
 	rangeSize, rangeKeys := external.CalRangeSize(nodeRc.TotalMem/int64(nodeRc.TotalCPU), regionSplitSize, regionSplitKeys)
 	logutil.DDLIngestLogger().Info("split kv range with split size and keys",
 		zap.Int64("region-split-size", regionSplitSize),

pkg/domain/BUILD.bazel

@@ -36,7 +36,6 @@ go_library(
         "//pkg/domain/infosync",
         "//pkg/domain/metrics",
         "//pkg/domain/sqlsvrapi",
-        "//pkg/dxf/framework/handle",
         "//pkg/dxf/framework/metering",
         "//pkg/dxf/framework/proto",
         "//pkg/dxf/framework/scheduler",

pkg/domain/domain.go

@@ -50,7 +50,6 @@ import (
 	"github.com/pingcap/tidb/pkg/domain/globalconfigsync"
 	"github.com/pingcap/tidb/pkg/domain/infosync"
 	"github.com/pingcap/tidb/pkg/domain/sqlsvrapi"
-	disthandle "github.com/pingcap/tidb/pkg/dxf/framework/handle"
 	"github.com/pingcap/tidb/pkg/dxf/framework/metering"
 	"github.com/pingcap/tidb/pkg/dxf/framework/proto"
 	"github.com/pingcap/tidb/pkg/dxf/framework/scheduler"
@@ -1127,7 +1126,7 @@ func (do *Domain) InitDistTaskLoop() error {
 	if err != nil {
 		return err
 	}
-	disthandle.SetNodeResource(nodeRes)
+	storage.SetNodeResource(nodeRes)
 	executorManager, err := taskexecutor.NewManager(managerCtx, do.store, serverID, taskManager, nodeRes)
 	if err != nil {
 		return err
@@ -1175,11 +1174,16 @@ func calculateNodeResource() (*proto.NodeResource, error) {
 	} else {
 		totalDisk = sz.Capacity
 	}
+	nodeRes := proto.NewNodeResource(totalCPU, int64(totalMem), totalDisk)
+	dxfNodeRes := nodeRes.LimitDXFResource(cfg.DXFResourceLimit)
 	logger.Info("initialize node resource",
 		zap.Int("total-cpu", totalCPU),
 		zap.String("total-mem", units.BytesSize(float64(totalMem))),
+		zap.Int("dxf-resource-limit", cfg.DXFResourceLimit),
+		zap.Int("dxf-usable-cpu", dxfNodeRes.TotalCPU),
+		zap.String("dxf-usable-mem", units.BytesSize(float64(dxfNodeRes.TotalMem))),
 		zap.String("total-disk", units.BytesSize(float64(totalDisk))))
-	return proto.NewNodeResource(totalCPU, int64(totalMem), totalDisk), nil
+	return dxfNodeRes, nil
 }
 
 func (do *Domain) distTaskFrameworkLoop(ctx context.Context, taskManager *storage.TaskManager, executorManager *taskexecutor.Manager, serverID string, nodeRes *proto.NodeResource) {

pkg/dxf/framework/handle/BUILD.bazel

@@ -26,7 +26,6 @@ go_library(
         "//pkg/sessionctx",
         "//pkg/sessionctx/vardef",
         "//pkg/util/backoff",
-        "//pkg/util/cpu",
         "//pkg/util/disttask",
         "//pkg/util/logutil",
         "//pkg/util/sem/compat",
@@ -35,7 +34,6 @@ go_library(
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_tikv_client_go_v2//util",
-        "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
     ],
 )
@@ -56,6 +54,7 @@ go_test(
         "//pkg/dxf/framework/proto",
         "//pkg/dxf/framework/schstatus",
         "//pkg/dxf/framework/storage",
+        "//pkg/dxf/framework/testutil",
         "//pkg/kv",
         "//pkg/meta",
         "//pkg/sessionctx",

pkg/dxf/framework/handle/handle.go

@@ -46,7 +46,6 @@ import (
 	"github.com/pingcap/tidb/pkg/util/sem/compat"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
 	"github.com/tikv/client-go/v2/util"
-	"go.uber.org/atomic"
 	"go.uber.org/zap"
 )
 
@@ -292,18 +291,6 @@ func RunWithRetry(
 	return lastErr
 }
 
-var nodeResource atomic.Pointer[proto.NodeResource]
-
-// GetNodeResource gets the node resource.
-func GetNodeResource() *proto.NodeResource {
-	return nodeResource.Load()
-}
-
-// SetNodeResource gets the node resource.
-func SetNodeResource(rc *proto.NodeResource) {
-	nodeResource.Store(rc)
-}
-
 // GetDefaultRegionSplitConfig gets the default region split size and keys.
 func GetDefaultRegionSplitConfig() (splitSize, splitKeys int64) {
 	if kerneltype.IsNextGen() {
@@ -473,9 +460,3 @@ func SendRowAndSizeMeterData(ctx context.Context, task *proto.Task, rows int64,
 	failpoint.InjectCall("afterSendRowAndSizeMeterData", item)
 	return nil
 }
-
-func init() {
-	// domain will init this var at runtime, we store it here for test, as some
-	// test might not start domain.
-	nodeResource.Store(proto.NewNodeResource(8, 16*units.GiB, 100*units.GiB))
-}

pkg/dxf/framework/handle/status.go

@@ -27,11 +27,11 @@ import (
 	"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"
+	"go.uber.org/zap"
 )
 
 // GetScheduleStatus returns the schedule status.
@@ -115,8 +115,9 @@ func GetNodesInfo(ctx context.Context, manager *storage.TaskManager) (nodeCount
 	if len(nodes) == 0 {
 		// shouldn't happen normally as every node will register itself to the meta
 		// table.
-		logutil.BgLogger().Warn("no managed nodes found, use local node CPU count instead")
-		cpuCount = cpu.GetCPUCount()
+		cpuCount = storage.GetDXFCPUCount()
+		logutil.BgLogger().Warn("no managed nodes found, use local usable CPU count instead",
+			zap.Int("usableCPUCount", cpuCount))
 	} else {
 		cpuCount = nodes[0].CPUCount
 	}

pkg/dxf/framework/handle/status_testkit_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/pingcap/tidb/pkg/dxf/framework/proto"
 	"github.com/pingcap/tidb/pkg/dxf/framework/schstatus"
 	"github.com/pingcap/tidb/pkg/dxf/framework/storage"
+	"github.com/pingcap/tidb/pkg/dxf/framework/testutil"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta"
 	"github.com/pingcap/tidb/pkg/sessionctx"
@@ -59,7 +60,7 @@ func TestGetScheduleStatus(t *testing.T) {
 
 func TestNodeInfoAndBusyNodes(t *testing.T) {
 	testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/domain/MockDisableDistTask", "return(true)")
-	testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/util/cpu/mockNumCpu", "return(16)")
+	testutil.MockNodeResource(t, 16)
 	store := testkit.CreateMockStore(t)
 	gtk := testkit.NewTestKit(t, store)
 	pool := pools.NewResourcePool(func() (pools.Resource, error) {