diff --git a/resourcemanager/BUILD.bazel b/resourcemanager/BUILD.bazel index 968a73b4a0f95..c6e7983df34be 100644 --- a/resourcemanager/BUILD.bazel +++ b/resourcemanager/BUILD.bazel @@ -13,6 +13,7 @@ go_library( "//resourcemanager/util", "//util", "//util/cpu", + "@com_github_google_uuid//:uuid", "@com_github_pingcap_log//:log", "@org_uber_go_zap//:zap", ], diff --git a/resourcemanager/rm.go b/resourcemanager/rm.go index aa2af6f949465..ed73cda8e1abf 100644 --- a/resourcemanager/rm.go +++ b/resourcemanager/rm.go @@ -17,6 +17,7 @@ package resourcemanager import ( "time" + "github.com/google/uuid" "github.com/pingcap/tidb/resourcemanager/scheduler" "github.com/pingcap/tidb/resourcemanager/util" tidbutil "github.com/pingcap/tidb/util" @@ -26,6 +27,11 @@ import ( // GlobalResourceManager is a global resource manager var GlobalResourceManager = NewResourceManger() +// RandomName is to get a random name for register pool. It is just for test. +func RandomName() string { + return uuid.New().String() +} + // ResourceManager is a resource manager type ResourceManager struct { poolMap *util.ShardPoolMap @@ -85,3 +91,8 @@ func (r *ResourceManager) registerPool(name string, pool *util.PoolContainer) er func (r *ResourceManager) Unregister(name string) { r.poolMap.Del(name) } + +// Reset is to Reset resource manager. it is just for test. +func (r *ResourceManager) Reset() { + r.poolMap = util.NewShardPoolMap() +} diff --git a/testkit/BUILD.bazel b/testkit/BUILD.bazel index 4e0e24091db27..c28ef0614eb04 100644 --- a/testkit/BUILD.bazel +++ b/testkit/BUILD.bazel @@ -21,6 +21,7 @@ go_library( "//parser/ast", "//parser/terror", "//planner/core", + "//resourcemanager", "//session", "//session/txninfo", "//sessionctx/variable", diff --git a/testkit/mockstore.go b/testkit/mockstore.go index 12afe0e0f2f68..9756d5bb65804 100644 --- a/testkit/mockstore.go +++ b/testkit/mockstore.go @@ -24,6 +24,7 @@ import ( "github.com/pingcap/tidb/ddl/schematracker" "github.com/pingcap/tidb/domain" "github.com/pingcap/tidb/kv" + "github.com/pingcap/tidb/resourcemanager" "github.com/pingcap/tidb/session" "github.com/pingcap/tidb/store/driver" "github.com/pingcap/tidb/store/mockstore" @@ -91,6 +92,7 @@ func bootstrap(t testing.TB, store kv.Storage, lease time.Duration) *domain.Doma err := store.Close() require.NoError(t, err) view.Stop() + resourcemanager.GlobalResourceManager.Reset() }) return dom }