diff --git a/DEPS.bzl b/DEPS.bzl index 381bd2e9a9f28..30682100fc3cc 100644 --- a/DEPS.bzl +++ b/DEPS.bzl @@ -4207,8 +4207,8 @@ def go_deps(): build_tags = ["nextgen", "intest"], build_file_proto_mode = "disable_global", importpath = "github.com/tikv/pd/client", - sum = "h1:OoBvgoeWmdNEXtS+eOlhysz/OvhA4GS0OdPVhTXteGA=", - version = "v0.0.0-20260708075407-4e05b9d2c2d3", + sum = "h1:BiPn4oLXVqB9ptPNmyo3rh/2IH1dTouXrrnneH86Z3M=", + version = "v0.0.0-20260720043438-0b37df9a48ed", ) go_repository( name = "com_github_timakin_bodyclose", diff --git a/cmd/tidb-server/main.go b/cmd/tidb-server/main.go index e117e3e138629..98891d4aec575 100644 --- a/cmd/tidb-server/main.go +++ b/cmd/tidb-server/main.go @@ -710,6 +710,9 @@ func overrideConfig(cfg *config.Config, fset *flag.FlagSet) { fset.Visit(func(f *flag.Flag) { actualFlags[f.Name] = true }) + if actualFlags[nmStarterParams] && cfg.DeployMode == deploymode.Starter { + terror.MustNil(applyStarterAdditionalParams(cfg, getStarterAdditionalParams())) + } // Base if actualFlags[nmHost] { @@ -1352,6 +1355,7 @@ type starterParams struct { podName string podIP string podNamespace string + enableRGFallback bool } func parseStarterAdditionalParams(raw string) (starterParams, error) { @@ -1394,6 +1398,12 @@ func parseStarterAdditionalParams(raw string) (starterParams, error) { params.podIP = value case "pod-namespace": params.podNamespace = value + case "enable-rg-fallback": + enable, err := strconv.ParseBool(value) + if err != nil { + return params, fmt.Errorf("starter additional param %q must be a bool: %w", key, err) + } + params.enableRGFallback = enable default: return params, fmt.Errorf("unknown starter additional param %q", key) } @@ -1401,6 +1411,15 @@ func parseStarterAdditionalParams(raw string) (starterParams, error) { return params, nil } +func applyStarterAdditionalParams(cfg *config.Config, raw string) error { + params, err := parseStarterAdditionalParams(raw) + if err != nil { + return err + } + cfg.StarterParams.EnableRGFallback = params.enableRGFallback + return nil +} + func getStarterAdditionalParams() string { if starterAdditionalParams == nil { return "" diff --git a/cmd/tidb-server/main_test.go b/cmd/tidb-server/main_test.go index 5808ccc14bdf9..307444e36eb9e 100644 --- a/cmd/tidb-server/main_test.go +++ b/cmd/tidb-server/main_test.go @@ -85,13 +85,15 @@ func TestOverrideConfigKeyspaceActivateMode(t *testing.T) { fset := initFlagSet() require.NoError(t, fset.Parse([]string{ "--keyspace-activate=true", - "--starter-additional-params=pod-name=pod-1,pod-ip=10.0.0.1,pod-namespace=ns-1", + "--starter-additional-params=pod-name=pod-1,pod-ip=10.0.0.1,pod-namespace=ns-1,enable-rg-fallback=true", })) cfg := config.NewConfig() + cfg.DeployMode = deploymode.Starter overrideConfig(cfg, fset) require.True(t, cfg.KeyspaceActivateMode) - require.Equal(t, "pod-name=pod-1,pod-ip=10.0.0.1,pod-namespace=ns-1", *starterAdditionalParams) + require.True(t, cfg.StarterParams.EnableRGFallback) + require.Equal(t, "pod-name=pod-1,pod-ip=10.0.0.1,pod-namespace=ns-1,enable-rg-fallback=true", *starterAdditionalParams) } func TestSetGlobalVars(t *testing.T) { @@ -217,6 +219,11 @@ func TestCreateMgrClientRequiresPodIdentityInStarter(t *testing.T) { _, err = createMgrClientForStarter() require.ErrorContains(t, err, `unknown starter additional param "unknown"`) + invalidBoolParam := "enable-rg-fallback=definitely" + starterAdditionalParams = &invalidBoolParam + _, err = createMgrClientForStarter() + require.ErrorContains(t, err, `starter additional param "enable-rg-fallback" must be a bool`) + config.UpdateGlobal(func(conf *config.Config) { conf.StarterParams.ManagerAddr = "" }) diff --git a/go.mod b/go.mod index e590aed6c5c42..179607ad82d4e 100644 --- a/go.mod +++ b/go.mod @@ -123,7 +123,7 @@ require ( github.com/stretchr/testify v1.11.1 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2 github.com/tikv/client-go/v2 v2.0.8-0.20260708122311-01bd8f99f4da - github.com/tikv/pd/client v0.0.0-20260708075407-4e05b9d2c2d3 + github.com/tikv/pd/client v0.0.0-20260720043438-0b37df9a48ed github.com/timakin/bodyclose v0.0.0-20241222091800-1db5c5ca4d67 github.com/twmb/murmur3 v1.1.6 github.com/uber/jaeger-client-go v2.22.1+incompatible diff --git a/go.sum b/go.sum index a69dcfc6c875c..0440b05504932 100644 --- a/go.sum +++ b/go.sum @@ -902,8 +902,8 @@ github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= github.com/tikv/client-go/v2 v2.0.8-0.20260708122311-01bd8f99f4da h1:Ju9uUKu3M5gPl7Z90e9dcGSe8LtXxGqvSetbPcr6tyc= github.com/tikv/client-go/v2 v2.0.8-0.20260708122311-01bd8f99f4da/go.mod h1:MRhIujZMMkYcI49Euif4+A3+SHSGwTpyUUadNz48R4g= -github.com/tikv/pd/client v0.0.0-20260708075407-4e05b9d2c2d3 h1:OoBvgoeWmdNEXtS+eOlhysz/OvhA4GS0OdPVhTXteGA= -github.com/tikv/pd/client v0.0.0-20260708075407-4e05b9d2c2d3/go.mod h1:3/Bu91CJONgkDA+Y0v/cnbROSJnu5tQ09vv7JGybUBA= +github.com/tikv/pd/client v0.0.0-20260720043438-0b37df9a48ed h1:BiPn4oLXVqB9ptPNmyo3rh/2IH1dTouXrrnneH86Z3M= +github.com/tikv/pd/client v0.0.0-20260720043438-0b37df9a48ed/go.mod h1:3/Bu91CJONgkDA+Y0v/cnbROSJnu5tQ09vv7JGybUBA= github.com/timakin/bodyclose v0.0.0-20241222091800-1db5c5ca4d67 h1:9LPGD+jzxMlnk5r6+hJnar67cgpDIz/iyD+rfl5r2Vk= github.com/timakin/bodyclose v0.0.0-20241222091800-1db5c5ca4d67/go.mod h1:mkjARE7Yr8qU23YcGMSALbIxTQ9r9QBVahQOBRfU460= github.com/tjfoc/gmsm v1.3.2/go.mod h1:HaUcFuY0auTiaHB9MHFGCPx5IaLhTUd2atbCFBQXn9w= diff --git a/pkg/config/config.go b/pkg/config/config.go index a7cbf232abf5f..07e0c58d7c460 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -1094,6 +1094,9 @@ type StarterParams struct { // ManagerAddr is the TiDB manager address used by the shutdown notifier. // When empty and EnableManagerNotifier is true, the Starter path derives the service address from starter additional params. ManagerAddr string `toml:"manager-addr" json:"manager-addr,omitempty"` + // EnableRGFallback enables resource group lookup fallback for resource control. + // It is populated from --starter-additional-params and is not file-backed config. + EnableRGFallback bool `toml:"-" json:"-"` // MaxImportDataSize is the maximum total real source data size allowed for IMPORT INTO. // Zero means unlimited. MaxImportDataSize configtypes.ByteSize `toml:"max-import-data-size" json:"max-import-data-size,omitempty"` diff --git a/pkg/domain/BUILD.bazel b/pkg/domain/BUILD.bazel index 3b18f8d2dce0f..f6804d4b65a2d 100644 --- a/pkg/domain/BUILD.bazel +++ b/pkg/domain/BUILD.bazel @@ -11,6 +11,7 @@ go_library( "optimize_trace.go", "plan_replayer.go", "plan_replayer_dump.go", + "resource_group_controller_options.go", "ru_stats.go", "runaway.go", "schema_checker.go", @@ -116,6 +117,7 @@ go_library( "@com_github_pingcap_failpoint//:failpoint", "@com_github_pingcap_kvproto//pkg/metapb", "@com_github_pingcap_kvproto//pkg/pdpb", + "@com_github_pingcap_kvproto//pkg/resource_manager", "@com_github_pingcap_log//:log", "@com_github_pingcap_metering_sdk//config", "@com_github_stretchr_testify//require", @@ -149,14 +151,17 @@ go_test( "plan_replayer_slow_log_test.go", "plan_replayer_test.go", "ru_stats_test.go", + "runaway_test.go", "schema_checker_test.go", "topn_slow_query_test.go", ], embed = [":domain"], flaky = True, - shard_count = 30, + shard_count = 33, deps = [ "//pkg/config", + "//pkg/config/deploymode", + "//pkg/config/kerneltype", "//pkg/ddl", "//pkg/domain/infosync", "//pkg/domain/serverinfo", @@ -174,6 +179,7 @@ go_test( "//pkg/parser/mysql", "//pkg/parser/terror", "//pkg/planner/extstore", + "//pkg/resourcegroup/runaway", "//pkg/server", "//pkg/session", "//pkg/sessionctx/vardef", @@ -191,15 +197,21 @@ go_test( "@com_github_ngaut_pools//:pools", "@com_github_pingcap_errors//:errors", "@com_github_pingcap_failpoint//:failpoint", + "@com_github_pingcap_kvproto//pkg/kvrpcpb", + "@com_github_pingcap_kvproto//pkg/meta_storagepb", "@com_github_pingcap_kvproto//pkg/metapb", "@com_github_pingcap_kvproto//pkg/resource_manager", "@com_github_prometheus_client_model//go", "@com_github_stretchr_testify//require", + "@com_github_tikv_client_go_v2//tikvrpc", "@com_github_tikv_client_go_v2//txnkv/transaction", "@com_github_tikv_pd_client//:client", + "@com_github_tikv_pd_client//errs", "@com_github_tikv_pd_client//opt", "@com_github_tikv_pd_client//resource_group/controller", "@io_etcd_go_etcd_tests_v3//integration", + "@org_golang_google_grpc//codes", + "@org_golang_google_grpc//status", "@org_uber_go_goleak//:goleak", ], ) diff --git a/pkg/domain/infosync/resource_manager_client.go b/pkg/domain/infosync/resource_manager_client.go index 391507d3fb38f..25348059687e5 100644 --- a/pkg/domain/infosync/resource_manager_client.go +++ b/pkg/domain/infosync/resource_manager_client.go @@ -148,6 +148,14 @@ func (*mockResourceManagerClient) LoadResourceGroups(context.Context) ([]*rmpb.R return nil, 0, nil } +func (*mockResourceManagerClient) Get(context.Context, []byte, ...opt.MetaStorageOption) (*meta_storagepb.GetResponse, error) { + return &meta_storagepb.GetResponse{Header: &meta_storagepb.ResponseHeader{}}, nil +} + +func (*mockResourceManagerClient) Put(context.Context, []byte, []byte, ...opt.MetaStorageOption) (*meta_storagepb.PutResponse, error) { + return &meta_storagepb.PutResponse{Header: &meta_storagepb.ResponseHeader{}}, nil +} + func (m *mockResourceManagerClient) Watch(_ context.Context, key []byte, _ ...opt.MetaStorageOption) (chan *metastorage.WatchResponse, error) { if bytes.Equal(pd.GroupSettingsPathPrefixBytes(m.keyspaceID), key) { return m.eventCh, nil diff --git a/pkg/domain/resource_group_controller_options.go b/pkg/domain/resource_group_controller_options.go new file mode 100644 index 0000000000000..abb9c0f1ef749 --- /dev/null +++ b/pkg/domain/resource_group_controller_options.go @@ -0,0 +1,66 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package domain + +import ( + "time" + + rmpb "github.com/pingcap/kvproto/pkg/resource_manager" + "github.com/pingcap/tidb/pkg/config" + "github.com/pingcap/tidb/pkg/config/deploymode" + "github.com/pingcap/tidb/pkg/resourcegroup/runaway" + rmclient "github.com/tikv/pd/client/resource_group/controller" +) + +const ( + defaultDegradedRUFillRate = 2_000_000 + defaultDegradedRUBurstLimit = 50_000_000_000 + defaultDegradedModeWaitTimeout = 3 * time.Second / 2 + tokenWaitRetryInterval = 100 * time.Millisecond + tokenWaitRetryTimes = 20 +) + +func newDefaultDegradedRUSettings() *rmpb.GroupRequestUnitSettings { + return &rmpb.GroupRequestUnitSettings{ + RU: &rmpb.TokenBucket{ + Settings: &rmpb.TokenLimitSettings{ + FillRate: defaultDegradedRUFillRate, + BurstLimit: defaultDegradedRUBurstLimit, + }, + }, + } +} + +func newResourceGroupsControllerOptions() []rmclient.ResourceControlCreateOption { + opts := []rmclient.ResourceControlCreateOption{ + rmclient.WithMaxWaitDuration(runaway.MaxWaitDuration), + } + if deploymode.IsStarter() && config.GetGlobalConfig().StarterParams.EnableRGFallback { + opts = append(opts, + // This Starter-only fallback path is a best-effort UX fallback for + // temporary GetResourceGroup failures. It provides a permissive + // group so user requests do not fail immediately; it is not intended + // to define a precise cross-RPC RU limit or response-side accounting + // contract while resource manager is unavailable. Keep synthesis + // inside the controller so degraded groups are not inserted into the + // normal metadata cache. + rmclient.WithDegradedRUSettings(newDefaultDegradedRUSettings()), + rmclient.WithDegradedModeWaitDuration(defaultDegradedModeWaitTimeout), + rmclient.WithWaitRetryInterval(tokenWaitRetryInterval), + rmclient.WithWaitRetryTimes(tokenWaitRetryTimes), + ) + } + return opts +} diff --git a/pkg/domain/runaway.go b/pkg/domain/runaway.go index ae1f2d6d233c2..439369bc2536f 100644 --- a/pkg/domain/runaway.go +++ b/pkg/domain/runaway.go @@ -39,7 +39,7 @@ func (do *Domain) initResourceGroupsController(ctx context.Context, pdClient pd. if codec := do.Store().GetCodec(); codec != nil { keyspaceID = uint32(codec.GetKeyspaceID()) } - control, err := rmclient.NewResourceGroupController(ctx, uniqueID, pdClient, nil, keyspaceID, rmclient.WithMaxWaitDuration(runaway.MaxWaitDuration)) + control, err := rmclient.NewResourceGroupController(ctx, uniqueID, pdClient, nil, keyspaceID, newResourceGroupsControllerOptions()...) if err != nil { return err } diff --git a/pkg/domain/runaway_test.go b/pkg/domain/runaway_test.go new file mode 100644 index 0000000000000..dcec631574306 --- /dev/null +++ b/pkg/domain/runaway_test.go @@ -0,0 +1,290 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package domain + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/pingcap/kvproto/pkg/kvrpcpb" + "github.com/pingcap/kvproto/pkg/meta_storagepb" + rmpb "github.com/pingcap/kvproto/pkg/resource_manager" + "github.com/pingcap/tidb/pkg/config" + "github.com/pingcap/tidb/pkg/config/deploymode" + "github.com/pingcap/tidb/pkg/config/kerneltype" + "github.com/pingcap/tidb/pkg/domain/infosync" + "github.com/pingcap/tidb/pkg/resourcegroup/runaway" + "github.com/stretchr/testify/require" + "github.com/tikv/client-go/v2/tikvrpc" + pd "github.com/tikv/pd/client" + pderr "github.com/tikv/pd/client/errs" + "github.com/tikv/pd/client/opt" + rmclient "github.com/tikv/pd/client/resource_group/controller" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +type resourceGroupProviderStub struct { + rmclient.ResourceGroupProvider + resourceGroup *rmpb.ResourceGroup + resourceErr error + controllerConfig *rmclient.Config +} + +func newResourceGroupProviderStub(t *testing.T, resourceGroup *rmpb.ResourceGroup, resourceErr error) *resourceGroupProviderStub { + t.Helper() + baseProvider, ok := infosync.NewMockResourceManagerClient(0).(rmclient.ResourceGroupProvider) + require.True(t, ok) + return &resourceGroupProviderStub{ + ResourceGroupProvider: baseProvider, + resourceGroup: resourceGroup, + resourceErr: resourceErr, + } +} + +// GetResourceGroup returns both the mocked resource group and the mocked error. +// This lets the test verify whether the controller uses the degraded fallback +// only for the editions that enable it. +func (s *resourceGroupProviderStub) GetResourceGroup(context.Context, string, ...pd.GetResourceGroupOption) (*rmpb.ResourceGroup, error) { + return s.resourceGroup, s.resourceErr +} + +func (s *resourceGroupProviderStub) Get(ctx context.Context, key []byte, opts ...opt.MetaStorageOption) (*meta_storagepb.GetResponse, error) { + if s.controllerConfig == nil { + return s.ResourceGroupProvider.Get(ctx, key, opts...) + } + value, err := json.Marshal(s.controllerConfig) + if err != nil { + return nil, err + } + return &meta_storagepb.GetResponse{ + Kvs: []*meta_storagepb.KeyValue{{ + Key: key, + Value: value, + }}, + }, nil +} + +func newStarterControllerForTest(t *testing.T, provider rmclient.ResourceGroupProvider) *rmclient.ResourceGroupsController { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + require.NoError(t, deploymode.Set(deploymode.Starter)) + controller, err := rmclient.NewResourceGroupController( + ctx, + 1, + provider, + nil, + 0, + newResourceGroupsControllerOptions()..., + ) + require.NoError(t, err) + controller.Start(ctx) + t.Cleanup(func() { + require.NoError(t, controller.Stop()) + }) + return controller +} + +func requireDegradedResourceGroup(t *testing.T, group *rmpb.ResourceGroup, name string) { + t.Helper() + require.NotNil(t, group) + require.Equal(t, name, group.Name) + require.Equal(t, rmpb.GroupMode_RUMode, group.Mode) + require.NotNil(t, group.RUSettings) + require.NotNil(t, group.RUSettings.RU) + require.NotNil(t, group.RUSettings.RU.Settings) + require.EqualValues(t, defaultDegradedRUFillRate, group.RUSettings.RU.Settings.FillRate) + require.EqualValues(t, defaultDegradedRUBurstLimit, group.RUSettings.RU.Settings.BurstLimit) +} + +func newTransientGetResourceGroupErr(name string) error { + err := status.Error(codes.Unavailable, "resource manager unavailable") + return &pderr.ErrClientGetResourceGroup{ + ResourceGroupName: name, + Cause: err.Error(), + Err: err, + } +} + +func restoreResourceGroupControllerTestState(t *testing.T) { + t.Helper() + restoreConfig := config.RestoreFunc() + t.Cleanup(restoreConfig) + if !kerneltype.IsNextGen() { + return + } + // Preserve the process-wide deploy mode because deploymode.IsStarter reads + // it directly when newResourceGroupsControllerOptions builds controller options. + originalDeployMode := deploymode.Get() + t.Cleanup(func() { + require.NoError(t, deploymode.Set(originalDeployMode)) + }) +} + +func newTestResourceGroup(name string) *rmpb.ResourceGroup { + return &rmpb.ResourceGroup{ + Name: name, + Mode: rmpb.GroupMode_RUMode, + RUSettings: &rmpb.GroupRequestUnitSettings{ + RU: &rmpb.TokenBucket{ + Settings: &rmpb.TokenLimitSettings{FillRate: 1}, + }, + }, + } +} + +func TestStarterDegradedResourceGroup(t *testing.T) { + if !kerneltype.IsNextGen() { + t.Skip("Starter deploy mode is only available in NextGen builds") + } + + t.Run("fallback", func(t *testing.T) { + restoreResourceGroupControllerTestState(t) + config.UpdateGlobal(func(conf *config.Config) { + conf.StarterParams.EnableRGFallback = true + }) + + provider := newResourceGroupProviderStub(t, nil, newTransientGetResourceGroupErr("test-group")) + controller := newStarterControllerForTest(t, provider) + group, err := controller.GetResourceGroup("test-group") + require.NoError(t, err) + requireDegradedResourceGroup(t, group, "test-group") + }) + + t.Run("recovery does not cache degraded group", func(t *testing.T) { + restoreResourceGroupControllerTestState(t) + config.UpdateGlobal(func(conf *config.Config) { + conf.StarterParams.EnableRGFallback = true + }) + + provider := newResourceGroupProviderStub(t, nil, newTransientGetResourceGroupErr("test-group")) + controller := newStarterControllerForTest(t, provider) + group, err := controller.GetResourceGroup("test-group") + require.NoError(t, err) + requireDegradedResourceGroup(t, group, "test-group") + + provider.resourceGroup = newTestResourceGroup("test-group") + provider.resourceErr = nil + group, err = controller.GetResourceGroup("test-group") + require.NoError(t, err) + require.Equal(t, provider.resourceGroup, group) + }) +} + +func TestResourceGroupsControllerOptions(t *testing.T) { + if !kerneltype.IsNextGen() { + t.Skip("Starter deploy mode is only available in NextGen builds") + } + + newController := func(t *testing.T) *rmclient.ResourceGroupsController { + t.Helper() + provider := newResourceGroupProviderStub(t, nil, nil) + provider.controllerConfig = rmclient.DefaultConfig() + provider.controllerConfig.WaitRetryInterval = rmclient.NewDuration(250 * time.Millisecond) + provider.controllerConfig.WaitRetryTimes = 4 + provider.controllerConfig.LTBTokenRPCMaxDelay = rmclient.NewDuration(time.Second) + + controller, err := rmclient.NewResourceGroupController( + context.Background(), + 1, + provider, + nil, + 0, + newResourceGroupsControllerOptions()..., + ) + require.NoError(t, err) + return controller + } + + t.Run("starter enables degraded mode explicitly", func(t *testing.T) { + restoreResourceGroupControllerTestState(t) + require.NoError(t, deploymode.Set(deploymode.Starter)) + config.UpdateGlobal(func(conf *config.Config) { + conf.StarterParams.EnableRGFallback = true + }) + + ruConfig := newController(t).GetConfig() + require.Equal(t, tokenWaitRetryInterval, ruConfig.WaitRetryInterval) + require.Equal(t, tokenWaitRetryTimes, ruConfig.WaitRetryTimes) + require.Equal(t, defaultDegradedModeWaitTimeout, ruConfig.DegradedModeWaitDuration) + }) + + t.Run("starter without degraded flag keeps default retry settings", func(t *testing.T) { + restoreResourceGroupControllerTestState(t) + require.NoError(t, deploymode.Set(deploymode.Starter)) + config.UpdateGlobal(func(conf *config.Config) { + conf.StarterParams.EnableRGFallback = false + }) + + ruConfig := newController(t).GetConfig() + require.Equal(t, 250*time.Millisecond, ruConfig.WaitRetryInterval) + require.Equal(t, 4, ruConfig.WaitRetryTimes) + require.Zero(t, ruConfig.DegradedModeWaitDuration) + }) + + t.Run("non starter ignores degraded flag", func(t *testing.T) { + restoreResourceGroupControllerTestState(t) + require.NoError(t, deploymode.Set(deploymode.Premium)) + config.UpdateGlobal(func(conf *config.Config) { + conf.StarterParams.EnableRGFallback = true + }) + + ruConfig := newController(t).GetConfig() + require.Equal(t, 250*time.Millisecond, ruConfig.WaitRetryInterval) + require.Equal(t, 4, ruConfig.WaitRetryTimes) + require.Zero(t, ruConfig.DegradedModeWaitDuration) + }) +} + +func TestStarterRunawaySwitchGroup(t *testing.T) { + if !kerneltype.IsNextGen() { + t.Skip("Starter deploy mode is only available in NextGen builds") + } + restoreResourceGroupControllerTestState(t) + + config.UpdateGlobal(func(conf *config.Config) { + conf.StarterParams.EnableRGFallback = true + }) + provider := newResourceGroupProviderStub(t, nil, newTransientGetResourceGroupErr("target-switch-group")) + controller := newStarterControllerForTest(t, provider) + manager := runaway.NewRunawayManager(controller, "127.0.0.1:4000", nil, make(chan struct{}), nil, nil) + t.Cleanup(manager.Stop) + checker := runaway.NewChecker( + manager, + "source-group", + &rmpb.RunawaySettings{ + Action: rmpb.RunawayAction_SwitchGroup, + SwitchGroupName: "target-switch-group", + Rule: &rmpb.RunawayRule{ProcessedKeys: 1}, + }, + "SELECT 1", + "sql_digest", + "plan_digest", + time.Now(), + ) + + require.NoError(t, checker.CheckThresholds(nil, 10, nil)) + req := &tikvrpc.Request{ + Context: kvrpcpb.Context{ + ResourceControlContext: &kvrpcpb.ResourceControlContext{}, + }, + } + require.NoError(t, checker.BeforeCopRequest(req)) + require.Equal(t, "target-switch-group", req.GetResourceControlContext().GetResourceGroupName()) +} diff --git a/pkg/executor/BUILD.bazel b/pkg/executor/BUILD.bazel index 0c50820197f26..f433305294238 100644 --- a/pkg/executor/BUILD.bazel +++ b/pkg/executor/BUILD.bazel @@ -564,8 +564,6 @@ go_test( "@com_github_tikv_client_go_v2//tikv", "@com_github_tikv_client_go_v2//tikvrpc", "@com_github_tikv_client_go_v2//util", - "@com_github_tikv_pd_client//:client", - "@com_github_tikv_pd_client//clients/metastorage", "@com_github_tikv_pd_client//http", "@com_github_tikv_pd_client//opt", "@com_github_tikv_pd_client//resource_group/controller", diff --git a/pkg/executor/adapter_internal_test.go b/pkg/executor/adapter_internal_test.go index c1a85252735cd..5a5d1263f73ac 100644 --- a/pkg/executor/adapter_internal_test.go +++ b/pkg/executor/adapter_internal_test.go @@ -22,8 +22,8 @@ import ( "time" "github.com/pingcap/kvproto/pkg/meta_storagepb" - rmpb "github.com/pingcap/kvproto/pkg/resource_manager" "github.com/pingcap/tidb/pkg/domain" + "github.com/pingcap/tidb/pkg/domain/infosync" "github.com/pingcap/tidb/pkg/expression" "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/parser" @@ -40,8 +40,6 @@ import ( "github.com/pingcap/tidb/pkg/util/topsql/stmtstats" "github.com/stretchr/testify/require" "github.com/tikv/client-go/v2/util" - pd "github.com/tikv/pd/client" - metastorage "github.com/tikv/pd/client/clients/metastorage" "github.com/tikv/pd/client/opt" rmclient "github.com/tikv/pd/client/resource_group/controller" ) @@ -452,9 +450,20 @@ func TestObserveStmtFinishedOnTopProfilingIgnores(t *testing.T) { } type mockResourceGroupProvider struct { + rmclient.ResourceGroupProvider config *rmclient.Config } +func newMockResourceGroupProvider(t *testing.T, config *rmclient.Config) *mockResourceGroupProvider { + t.Helper() + baseProvider, ok := infosync.NewMockResourceManagerClient(1).(rmclient.ResourceGroupProvider) + require.True(t, ok) + return &mockResourceGroupProvider{ + ResourceGroupProvider: baseProvider, + config: config, + } +} + func newMockDomainWithRUVersion(t *testing.T, version rmclient.RUVersion) *domain.Domain { t.Helper() ctx, cancel := context.WithCancel(context.Background()) @@ -462,7 +471,7 @@ func newMockDomainWithRUVersion(t *testing.T, version rmclient.RUVersion) *domai cfg := rmclient.DefaultConfig() cfg.RUVersionPolicy = &rmclient.RUVersionPolicy{Default: version} - provider := &mockResourceGroupProvider{config: cfg} + provider := newMockResourceGroupProvider(t, cfg) controller, err := rmclient.NewResourceGroupController(ctx, 1, provider, nil, 1) require.NoError(t, err) @@ -481,48 +490,4 @@ func (m *mockResourceGroupProvider) Get(ctx context.Context, key []byte, opts .. }, nil } -func (*mockResourceGroupProvider) Watch(ctx context.Context, key []byte, opts ...opt.MetaStorageOption) (chan *metastorage.WatchResponse, error) { - ch := make(chan *metastorage.WatchResponse) - go func() { - <-ctx.Done() - close(ch) - }() - return ch, nil -} - -func (*mockResourceGroupProvider) Put(context.Context, []byte, []byte, ...opt.MetaStorageOption) (*meta_storagepb.PutResponse, error) { - return &meta_storagepb.PutResponse{}, nil -} - -func (*mockResourceGroupProvider) GetResourceGroup(context.Context, string, ...pd.GetResourceGroupOption) (*rmpb.ResourceGroup, error) { - return nil, nil -} - -func (*mockResourceGroupProvider) ListResourceGroups(context.Context, ...pd.GetResourceGroupOption) ([]*rmpb.ResourceGroup, error) { - return nil, nil -} - -func (*mockResourceGroupProvider) AddResourceGroup(context.Context, *rmpb.ResourceGroup) (string, error) { - return "", nil -} - -func (*mockResourceGroupProvider) ModifyResourceGroup(context.Context, *rmpb.ResourceGroup) (string, error) { - return "", nil -} - -func (*mockResourceGroupProvider) DeleteResourceGroup(context.Context, string) (string, error) { - return "", nil -} - -func (*mockResourceGroupProvider) AcquireTokenBuckets(context.Context, *rmpb.TokenBucketsRequest) ([]*rmpb.TokenBucketResponse, error) { - return nil, nil -} - -func (*mockResourceGroupProvider) LoadResourceGroups(context.Context) ([]*rmpb.ResourceGroup, int64, error) { - return nil, 0, nil -} - -var ( - _ metastorage.Client = (*mockResourceGroupProvider)(nil) - _ rmclient.ResourceGroupProvider = (*mockResourceGroupProvider)(nil) -) +var _ rmclient.ResourceGroupProvider = (*mockResourceGroupProvider)(nil)