mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 05:18:52 +08:00
622 lines
24 KiB
Go
622 lines
24 KiB
Go
// Licensed to the LF AI & Data foundation under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you 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 meta
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/rgpb"
|
|
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
|
|
"github.com/milvus-io/milvus/internal/kv/mocks"
|
|
"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
|
|
"github.com/milvus-io/milvus/internal/querycoordv2/params"
|
|
"github.com/milvus-io/milvus/internal/querycoordv2/session"
|
|
"github.com/milvus-io/milvus/pkg/kv"
|
|
"github.com/milvus-io/milvus/pkg/util/etcd"
|
|
"github.com/milvus-io/milvus/pkg/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
|
)
|
|
|
|
type ResourceManagerSuite struct {
|
|
suite.Suite
|
|
|
|
kv kv.MetaKv
|
|
manager *ResourceManager
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) SetupSuite() {
|
|
paramtable.Init()
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) SetupTest() {
|
|
config := params.GenerateEtcdConfig()
|
|
cli, err := etcd.GetEtcdClient(
|
|
config.UseEmbedEtcd.GetAsBool(),
|
|
config.EtcdUseSSL.GetAsBool(),
|
|
config.Endpoints.GetAsStrings(),
|
|
config.EtcdTLSCert.GetValue(),
|
|
config.EtcdTLSKey.GetValue(),
|
|
config.EtcdTLSCACert.GetValue(),
|
|
config.EtcdTLSMinVersion.GetValue())
|
|
suite.Require().NoError(err)
|
|
suite.kv = etcdkv.NewEtcdKV(cli, config.MetaRootPath.GetValue())
|
|
|
|
store := querycoord.NewCatalog(suite.kv)
|
|
suite.manager = NewResourceManager(store, session.NewNodeManager())
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) TearDownSuite() {
|
|
suite.kv.Close()
|
|
}
|
|
|
|
func TestResourceManager(t *testing.T) {
|
|
suite.Run(t, new(ResourceManagerSuite))
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) TestValidateConfiguration() {
|
|
err := suite.manager.validateResourceGroupConfig("rg1", newResourceGroupConfig(0, 0))
|
|
suite.NoError(err)
|
|
|
|
err = suite.manager.validateResourceGroupConfig("rg1", &rgpb.ResourceGroupConfig{})
|
|
suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)
|
|
|
|
err = suite.manager.validateResourceGroupConfig("rg1", newResourceGroupConfig(-1, 2))
|
|
suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)
|
|
|
|
err = suite.manager.validateResourceGroupConfig("rg1", newResourceGroupConfig(2, -1))
|
|
suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)
|
|
|
|
err = suite.manager.validateResourceGroupConfig("rg1", newResourceGroupConfig(3, 2))
|
|
suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)
|
|
|
|
cfg := newResourceGroupConfig(0, 0)
|
|
cfg.TransferFrom = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg1"}}
|
|
err = suite.manager.validateResourceGroupConfig("rg1", cfg)
|
|
suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)
|
|
|
|
cfg = newResourceGroupConfig(0, 0)
|
|
cfg.TransferFrom = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg2"}}
|
|
err = suite.manager.validateResourceGroupConfig("rg1", cfg)
|
|
suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)
|
|
|
|
cfg = newResourceGroupConfig(0, 0)
|
|
cfg.TransferTo = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg1"}}
|
|
err = suite.manager.validateResourceGroupConfig("rg1", cfg)
|
|
suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)
|
|
|
|
cfg = newResourceGroupConfig(0, 0)
|
|
cfg.TransferTo = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg2"}}
|
|
err = suite.manager.validateResourceGroupConfig("rg1", cfg)
|
|
suite.ErrorIs(err, merr.ErrResourceGroupIllegalConfig)
|
|
|
|
err = suite.manager.AddResourceGroup("rg2", newResourceGroupConfig(0, 0))
|
|
suite.NoError(err)
|
|
|
|
err = suite.manager.RemoveResourceGroup("rg2")
|
|
suite.NoError(err)
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) TestValidateDelete() {
|
|
// Non empty resource group can not be removed.
|
|
err := suite.manager.AddResourceGroup("rg1", newResourceGroupConfig(1, 1))
|
|
suite.NoError(err)
|
|
|
|
err = suite.manager.validateResourceGroupIsDeletable(DefaultResourceGroupName)
|
|
suite.ErrorIs(err, merr.ErrParameterInvalid)
|
|
|
|
err = suite.manager.validateResourceGroupIsDeletable("rg1")
|
|
suite.ErrorIs(err, merr.ErrParameterInvalid)
|
|
|
|
cfg := newResourceGroupConfig(0, 0)
|
|
cfg.TransferFrom = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg1"}}
|
|
suite.manager.AddResourceGroup("rg2", cfg)
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg1": newResourceGroupConfig(0, 0),
|
|
})
|
|
err = suite.manager.validateResourceGroupIsDeletable("rg1")
|
|
suite.ErrorIs(err, merr.ErrParameterInvalid)
|
|
|
|
cfg = newResourceGroupConfig(0, 0)
|
|
cfg.TransferTo = []*rgpb.ResourceGroupTransfer{{ResourceGroup: "rg1"}}
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg2": cfg,
|
|
})
|
|
err = suite.manager.validateResourceGroupIsDeletable("rg1")
|
|
suite.ErrorIs(err, merr.ErrParameterInvalid)
|
|
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg2": newResourceGroupConfig(0, 0),
|
|
})
|
|
err = suite.manager.validateResourceGroupIsDeletable("rg1")
|
|
suite.NoError(err)
|
|
|
|
err = suite.manager.RemoveResourceGroup("rg1")
|
|
suite.NoError(err)
|
|
err = suite.manager.RemoveResourceGroup("rg2")
|
|
suite.NoError(err)
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) TestManipulateResourceGroup() {
|
|
// test add rg
|
|
err := suite.manager.AddResourceGroup("rg1", newResourceGroupConfig(0, 0))
|
|
suite.NoError(err)
|
|
suite.True(suite.manager.ContainResourceGroup("rg1"))
|
|
suite.Len(suite.manager.ListResourceGroups(), 2)
|
|
|
|
// test add duplicate rg but same configuration is ok
|
|
err = suite.manager.AddResourceGroup("rg1", newResourceGroupConfig(0, 0))
|
|
suite.NoError(err)
|
|
|
|
err = suite.manager.AddResourceGroup("rg1", newResourceGroupConfig(1, 1))
|
|
suite.Error(err)
|
|
|
|
// test delete rg
|
|
err = suite.manager.RemoveResourceGroup("rg1")
|
|
suite.NoError(err)
|
|
|
|
// test delete rg which doesn't exist
|
|
err = suite.manager.RemoveResourceGroup("rg1")
|
|
suite.NoError(err)
|
|
// test delete default rg
|
|
err = suite.manager.RemoveResourceGroup(DefaultResourceGroupName)
|
|
suite.ErrorIs(err, merr.ErrParameterInvalid)
|
|
|
|
// test delete a rg not empty.
|
|
err = suite.manager.AddResourceGroup("rg2", newResourceGroupConfig(1, 1))
|
|
suite.NoError(err)
|
|
err = suite.manager.RemoveResourceGroup("rg2")
|
|
suite.ErrorIs(err, merr.ErrParameterInvalid)
|
|
|
|
// test delete a rg after update
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg2": newResourceGroupConfig(0, 0),
|
|
})
|
|
err = suite.manager.RemoveResourceGroup("rg2")
|
|
suite.NoError(err)
|
|
|
|
// assign a node to rg.
|
|
err = suite.manager.AddResourceGroup("rg2", newResourceGroupConfig(1, 1))
|
|
suite.NoError(err)
|
|
suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 1,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
defer suite.manager.nodeMgr.Remove(1)
|
|
suite.manager.HandleNodeUp(1)
|
|
err = suite.manager.RemoveResourceGroup("rg2")
|
|
suite.ErrorIs(err, merr.ErrParameterInvalid)
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg2": newResourceGroupConfig(0, 0),
|
|
})
|
|
// RemoveResourceGroup will remove all nodes from the resource group.
|
|
err = suite.manager.RemoveResourceGroup("rg2")
|
|
suite.NoError(err)
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) TestNodeUpAndDown() {
|
|
suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 1,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
err := suite.manager.AddResourceGroup("rg1", newResourceGroupConfig(1, 1))
|
|
suite.NoError(err)
|
|
// test add node to rg
|
|
suite.manager.HandleNodeUp(1)
|
|
suite.Equal(1, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
|
|
// test add non-exist node to rg
|
|
err = suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg1": newResourceGroupConfig(2, 3),
|
|
})
|
|
suite.NoError(err)
|
|
suite.manager.HandleNodeUp(2)
|
|
suite.Equal(1, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// teardown a non-exist node from rg.
|
|
suite.manager.HandleNodeDown(2)
|
|
suite.Equal(1, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// test add exist node to rg
|
|
suite.manager.HandleNodeUp(1)
|
|
suite.Equal(1, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// teardown a exist node from rg.
|
|
suite.manager.HandleNodeDown(1)
|
|
suite.Zero(suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// teardown a exist node from rg.
|
|
suite.manager.HandleNodeDown(1)
|
|
suite.Zero(suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
suite.manager.HandleNodeUp(1)
|
|
suite.Equal(1, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
err = suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg1": newResourceGroupConfig(4, 4),
|
|
})
|
|
suite.NoError(err)
|
|
suite.manager.AddResourceGroup("rg2", newResourceGroupConfig(1, 1))
|
|
suite.NoError(err)
|
|
|
|
suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 11,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 12,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 13,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 14,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.manager.HandleNodeUp(11)
|
|
suite.manager.HandleNodeUp(12)
|
|
suite.manager.HandleNodeUp(13)
|
|
suite.manager.HandleNodeUp(14)
|
|
|
|
suite.Equal(4, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(1, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
suite.manager.HandleNodeDown(11)
|
|
suite.manager.HandleNodeDown(12)
|
|
suite.manager.HandleNodeDown(13)
|
|
suite.manager.HandleNodeDown(14)
|
|
suite.Equal(1, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
suite.manager.HandleNodeDown(1)
|
|
suite.Zero(suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Zero(suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg1": newResourceGroupConfig(20, 30),
|
|
"rg2": newResourceGroupConfig(30, 40),
|
|
})
|
|
for i := 1; i <= 100; i++ {
|
|
suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: int64(i),
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.manager.HandleNodeUp(int64(i))
|
|
}
|
|
|
|
suite.Equal(20, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(30, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(50, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// down all nodes
|
|
for i := 1; i <= 100; i++ {
|
|
suite.manager.HandleNodeDown(int64(i))
|
|
suite.Equal(100-i, suite.manager.GetResourceGroup("rg1").NodeNum()+
|
|
suite.manager.GetResourceGroup("rg2").NodeNum()+
|
|
suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
}
|
|
|
|
// if there are all rgs reach limit, should be fall back to default rg.
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg1": newResourceGroupConfig(0, 0),
|
|
"rg2": newResourceGroupConfig(0, 0),
|
|
DefaultResourceGroupName: newResourceGroupConfig(0, 0),
|
|
})
|
|
|
|
for i := 1; i <= 100; i++ {
|
|
suite.manager.HandleNodeUp(int64(i))
|
|
suite.Equal(i, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
suite.Equal(0, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(0, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
}
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) TestAutoRecover() {
|
|
for i := 1; i <= 100; i++ {
|
|
suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: int64(i),
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.manager.HandleNodeUp(int64(i))
|
|
}
|
|
suite.Equal(100, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// Recover 10 nodes from default resource group
|
|
suite.manager.AddResourceGroup("rg1", newResourceGroupConfig(10, 30))
|
|
suite.Zero(suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(10, suite.manager.GetResourceGroup("rg1").MissingNumOfNodes())
|
|
suite.Equal(100, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
suite.manager.AutoRecoverResourceGroup("rg1")
|
|
suite.Equal(10, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(0, suite.manager.GetResourceGroup("rg1").MissingNumOfNodes())
|
|
suite.Equal(90, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// Recover 20 nodes from default resource group
|
|
suite.manager.AddResourceGroup("rg2", newResourceGroupConfig(20, 30))
|
|
suite.Zero(suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(20, suite.manager.GetResourceGroup("rg2").MissingNumOfNodes())
|
|
suite.Equal(10, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(90, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
suite.manager.AutoRecoverResourceGroup("rg2")
|
|
suite.Equal(20, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(10, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(70, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// Recover 5 redundant nodes from resource group
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg1": newResourceGroupConfig(5, 5),
|
|
})
|
|
suite.manager.AutoRecoverResourceGroup("rg1")
|
|
suite.Equal(20, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(5, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(75, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// Recover 10 redundant nodes from resource group 2 to resource group 1 and default resource group.
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg1": newResourceGroupConfig(10, 20),
|
|
"rg2": newResourceGroupConfig(5, 10),
|
|
})
|
|
|
|
suite.manager.AutoRecoverResourceGroup("rg2")
|
|
suite.Equal(10, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(10, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(80, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// recover redundant nodes from default resource group
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg1": newResourceGroupConfig(10, 20),
|
|
"rg2": newResourceGroupConfig(20, 30),
|
|
DefaultResourceGroupName: newResourceGroupConfig(10, 20),
|
|
})
|
|
suite.manager.AutoRecoverResourceGroup("rg1")
|
|
suite.manager.AutoRecoverResourceGroup("rg2")
|
|
suite.manager.AutoRecoverResourceGroup(DefaultResourceGroupName)
|
|
|
|
// Even though the default resource group has 20 nodes limits,
|
|
// all redundant nodes will be assign to default resource group.
|
|
suite.Equal(20, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(30, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(50, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// Test recover missing from high priority resource group by set `from`.
|
|
suite.manager.AddResourceGroup("rg3", &rgpb.ResourceGroupConfig{
|
|
Requests: &rgpb.ResourceGroupLimit{
|
|
NodeNum: 15,
|
|
},
|
|
Limits: &rgpb.ResourceGroupLimit{
|
|
NodeNum: 15,
|
|
},
|
|
TransferFrom: []*rgpb.ResourceGroupTransfer{{
|
|
ResourceGroup: "rg1",
|
|
}},
|
|
})
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
DefaultResourceGroupName: newResourceGroupConfig(30, 40),
|
|
})
|
|
|
|
suite.manager.AutoRecoverResourceGroup("rg1")
|
|
suite.manager.AutoRecoverResourceGroup("rg2")
|
|
suite.manager.AutoRecoverResourceGroup(DefaultResourceGroupName)
|
|
suite.manager.AutoRecoverResourceGroup("rg3")
|
|
|
|
// Get 10 from default group for redundant nodes, get 5 from rg1 for rg3 at high priority.
|
|
suite.Equal(15, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(30, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(15, suite.manager.GetResourceGroup("rg3").NodeNum())
|
|
suite.Equal(40, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// Test recover redundant to high priority resource group by set `to`.
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg3": {
|
|
Requests: &rgpb.ResourceGroupLimit{
|
|
NodeNum: 0,
|
|
},
|
|
Limits: &rgpb.ResourceGroupLimit{
|
|
NodeNum: 0,
|
|
},
|
|
TransferTo: []*rgpb.ResourceGroupTransfer{{
|
|
ResourceGroup: "rg2",
|
|
}},
|
|
},
|
|
"rg1": newResourceGroupConfig(15, 100),
|
|
"rg2": newResourceGroupConfig(15, 40),
|
|
})
|
|
|
|
suite.manager.AutoRecoverResourceGroup("rg1")
|
|
suite.manager.AutoRecoverResourceGroup("rg2")
|
|
suite.manager.AutoRecoverResourceGroup(DefaultResourceGroupName)
|
|
suite.manager.AutoRecoverResourceGroup("rg3")
|
|
|
|
// Recover rg3 by transfer 10 nodes to rg2 with high priority, 5 to rg1.
|
|
suite.Equal(20, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(40, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(0, suite.manager.GetResourceGroup("rg3").NodeNum())
|
|
suite.Equal(40, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
suite.testTransferNode()
|
|
|
|
// Test redundant nodes recover to default resource group.
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
DefaultResourceGroupName: newResourceGroupConfig(1, 1),
|
|
"rg3": newResourceGroupConfig(0, 0),
|
|
"rg2": newResourceGroupConfig(0, 0),
|
|
"rg1": newResourceGroupConfig(0, 0),
|
|
})
|
|
// Even default resource group has 1 node limit,
|
|
// all redundant nodes will be assign to default resource group if there's no resource group can hold.
|
|
suite.manager.AutoRecoverResourceGroup(DefaultResourceGroupName)
|
|
suite.manager.AutoRecoverResourceGroup("rg1")
|
|
suite.manager.AutoRecoverResourceGroup("rg2")
|
|
suite.manager.AutoRecoverResourceGroup("rg3")
|
|
suite.Equal(0, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(0, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(0, suite.manager.GetResourceGroup("rg3").NodeNum())
|
|
suite.Equal(100, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// Test redundant recover to missing nodes and missing nodes from redundant nodes.
|
|
// Initialize
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
DefaultResourceGroupName: newResourceGroupConfig(0, 0),
|
|
"rg3": newResourceGroupConfig(10, 10),
|
|
"rg2": newResourceGroupConfig(80, 80),
|
|
"rg1": newResourceGroupConfig(10, 10),
|
|
})
|
|
suite.manager.AutoRecoverResourceGroup(DefaultResourceGroupName)
|
|
suite.manager.AutoRecoverResourceGroup("rg1")
|
|
suite.manager.AutoRecoverResourceGroup("rg2")
|
|
suite.manager.AutoRecoverResourceGroup("rg3")
|
|
suite.Equal(10, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(80, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(10, suite.manager.GetResourceGroup("rg3").NodeNum())
|
|
suite.Equal(0, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
DefaultResourceGroupName: newResourceGroupConfig(0, 5),
|
|
"rg3": newResourceGroupConfig(5, 5),
|
|
"rg2": newResourceGroupConfig(80, 80),
|
|
"rg1": newResourceGroupConfig(20, 30),
|
|
})
|
|
suite.manager.AutoRecoverResourceGroup("rg3") // recover redundant to missing rg.
|
|
suite.Equal(15, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(80, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(5, suite.manager.GetResourceGroup("rg3").NodeNum())
|
|
suite.Equal(0, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
suite.manager.updateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
DefaultResourceGroupName: newResourceGroupConfig(5, 5),
|
|
"rg3": newResourceGroupConfig(5, 10),
|
|
"rg2": newResourceGroupConfig(80, 80),
|
|
"rg1": newResourceGroupConfig(10, 10),
|
|
})
|
|
suite.manager.AutoRecoverResourceGroup(DefaultResourceGroupName) // recover missing from redundant rg.
|
|
suite.Equal(10, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(80, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(5, suite.manager.GetResourceGroup("rg3").NodeNum())
|
|
suite.Equal(5, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) testTransferNode() {
|
|
// Test redundant nodes recover to default resource group.
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
DefaultResourceGroupName: newResourceGroupConfig(40, 40),
|
|
"rg3": newResourceGroupConfig(0, 0),
|
|
"rg2": newResourceGroupConfig(40, 40),
|
|
"rg1": newResourceGroupConfig(20, 20),
|
|
})
|
|
suite.manager.AutoRecoverResourceGroup("rg1")
|
|
suite.manager.AutoRecoverResourceGroup("rg2")
|
|
suite.manager.AutoRecoverResourceGroup(DefaultResourceGroupName)
|
|
suite.manager.AutoRecoverResourceGroup("rg3")
|
|
|
|
suite.Equal(20, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(40, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(0, suite.manager.GetResourceGroup("rg3").NodeNum())
|
|
suite.Equal(40, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
|
|
// Test TransferNode.
|
|
// param error.
|
|
err := suite.manager.TransferNode("rg1", "rg1", 1)
|
|
suite.Error(err)
|
|
|
|
err = suite.manager.TransferNode("rg1", "rg2", 0)
|
|
suite.Error(err)
|
|
|
|
err = suite.manager.TransferNode("rg3", "rg2", 1)
|
|
suite.Error(err)
|
|
|
|
err = suite.manager.TransferNode("rg1", "rg10086", 1)
|
|
suite.Error(err)
|
|
|
|
err = suite.manager.TransferNode("rg10086", "rg2", 1)
|
|
suite.Error(err)
|
|
|
|
// success
|
|
err = suite.manager.TransferNode("rg1", "rg3", 5)
|
|
suite.NoError(err)
|
|
|
|
suite.manager.AutoRecoverResourceGroup("rg1")
|
|
suite.manager.AutoRecoverResourceGroup("rg2")
|
|
suite.manager.AutoRecoverResourceGroup(DefaultResourceGroupName)
|
|
suite.manager.AutoRecoverResourceGroup("rg3")
|
|
|
|
suite.Equal(15, suite.manager.GetResourceGroup("rg1").NodeNum())
|
|
suite.Equal(40, suite.manager.GetResourceGroup("rg2").NodeNum())
|
|
suite.Equal(5, suite.manager.GetResourceGroup("rg3").NodeNum())
|
|
suite.Equal(40, suite.manager.GetResourceGroup(DefaultResourceGroupName).NodeNum())
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) TestIncomingNode() {
|
|
suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 1,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.manager.incomingNode.Insert(1)
|
|
|
|
suite.Equal(1, suite.manager.CheckIncomingNodeNum())
|
|
suite.manager.AssignPendingIncomingNode()
|
|
suite.Equal(0, suite.manager.CheckIncomingNodeNum())
|
|
nodes, err := suite.manager.GetNodes(DefaultResourceGroupName)
|
|
suite.NoError(err)
|
|
suite.Len(nodes, 1)
|
|
}
|
|
|
|
func (suite *ResourceManagerSuite) TestUnassignFail() {
|
|
// suite.man
|
|
mockKV := mocks.NewMetaKv(suite.T())
|
|
mockKV.EXPECT().MultiSave(mock.Anything).Return(nil).Once()
|
|
|
|
store := querycoord.NewCatalog(mockKV)
|
|
suite.manager = NewResourceManager(store, session.NewNodeManager())
|
|
|
|
suite.manager.UpdateResourceGroups(map[string]*rgpb.ResourceGroupConfig{
|
|
"rg1": newResourceGroupConfig(20, 30),
|
|
})
|
|
|
|
suite.manager.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
|
|
NodeID: 1,
|
|
Address: "localhost",
|
|
Hostname: "localhost",
|
|
}))
|
|
suite.manager.HandleNodeUp(1)
|
|
|
|
mockKV.EXPECT().MultiSave(mock.Anything).Return(merr.WrapErrServiceInternal("mocked")).Once()
|
|
|
|
suite.Panics(func() {
|
|
suite.manager.HandleNodeDown(1)
|
|
})
|
|
}
|