milvus/internal/querycoordv2/observers/replica_observer_test.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 observers

import (
	"testing"
	"time"

	"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/metastore/kv/querycoord"
	"github.com/milvus-io/milvus/internal/querycoordv2/meta"
	. "github.com/milvus-io/milvus/internal/querycoordv2/params"
	"github.com/milvus-io/milvus/internal/querycoordv2/session"
	"github.com/milvus-io/milvus/internal/querycoordv2/utils"
	"github.com/milvus-io/milvus/pkg/kv"
	"github.com/milvus-io/milvus/pkg/util/etcd"
	"github.com/milvus-io/milvus/pkg/util/paramtable"
	"github.com/milvus-io/milvus/pkg/util/typeutil"
)

type ReplicaObserverSuite struct {
	suite.Suite

	kv kv.MetaKv
	// dependency
	meta    *meta.Meta
	distMgr *meta.DistributionManager

	nodeMgr  *session.NodeManager
	observer *ReplicaObserver

	collectionID int64
	partitionID  int64
}

func (suite *ReplicaObserverSuite) SetupSuite() {
	paramtable.Init()
	paramtable.Get().Save(Params.QueryCoordCfg.CheckNodeInReplicaInterval.Key, "1")
}

func (suite *ReplicaObserverSuite) SetupTest() {
	var err error
	config := 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())

	// meta
	store := querycoord.NewCatalog(suite.kv)
	idAllocator := RandomIncrementIDAllocator()
	suite.nodeMgr = session.NewNodeManager()
	suite.meta = meta.NewMeta(idAllocator, store, suite.nodeMgr)

	suite.distMgr = meta.NewDistributionManager()
	suite.observer = NewReplicaObserver(suite.meta, suite.distMgr)
	suite.observer.Start()
	suite.collectionID = int64(1000)
	suite.partitionID = int64(100)
}

func (suite *ReplicaObserverSuite) TestCheckNodesInReplica() {
	suite.meta.ResourceManager.AddResourceGroup("rg1", &rgpb.ResourceGroupConfig{
		Requests: &rgpb.ResourceGroupLimit{NodeNum: 2},
		Limits:   &rgpb.ResourceGroupLimit{NodeNum: 2},
	})
	suite.meta.ResourceManager.AddResourceGroup("rg2", &rgpb.ResourceGroupConfig{
		Requests: &rgpb.ResourceGroupLimit{NodeNum: 2},
		Limits:   &rgpb.ResourceGroupLimit{NodeNum: 2},
	})
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   1,
		Address:  "localhost:8080",
		Hostname: "localhost",
	}))
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   2,
		Address:  "localhost:8080",
		Hostname: "localhost",
	}))
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   3,
		Address:  "localhost:8080",
		Hostname: "localhost",
	}))
	suite.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{
		NodeID:   4,
		Address:  "localhost:8080",
		Hostname: "localhost",
	}))
	suite.meta.ResourceManager.HandleNodeUp(1)
	suite.meta.ResourceManager.HandleNodeUp(2)
	suite.meta.ResourceManager.HandleNodeUp(3)
	suite.meta.ResourceManager.HandleNodeUp(4)

	err := suite.meta.CollectionManager.PutCollection(utils.CreateTestCollection(suite.collectionID, 2))
	suite.NoError(err)
	replicas, err := suite.meta.Spawn(suite.collectionID, map[string]int{
		"rg1": 1,
		"rg2": 1,
	}, nil)
	suite.NoError(err)
	suite.Equal(2, len(replicas))

	suite.Eventually(func() bool {
		availableNodes := typeutil.NewUniqueSet()
		for _, r := range replicas {
			replica := suite.meta.ReplicaManager.Get(r.GetID())
			suite.NotNil(replica)
			if replica.RWNodesCount() != 2 {
				return false
			}
			if replica.RONodesCount() != 0 {
				return false
			}
			availableNodes.Insert(replica.GetNodes()...)
		}
		return availableNodes.Len() == 4
	}, 6*time.Second, 2*time.Second)

	// Add some segment on nodes.
	for nodeID := int64(1); nodeID <= 4; nodeID++ {
		suite.distMgr.ChannelDistManager.Update(
			nodeID,
			utils.CreateTestChannel(suite.collectionID, nodeID, 1, "test-insert-channel1"))
		suite.distMgr.SegmentDistManager.Update(
			nodeID,
			utils.CreateTestSegment(suite.collectionID, suite.partitionID, 1, nodeID, 1, "test-insert-channel1"))
	}

	// Do a replica transfer.
	suite.meta.ReplicaManager.TransferReplica(suite.collectionID, "rg1", "rg2", 1)

	// All replica should in the rg2 but not rg1
	// And some nodes will become ro nodes before all segment and channel on it is cleaned.
	suite.Eventually(func() bool {
		for _, r := range replicas {
			replica := suite.meta.ReplicaManager.Get(r.GetID())
			suite.NotNil(replica)
			suite.Equal("rg2", replica.GetResourceGroup())
			// all replica should have ro nodes.
			// transferred replica should have 2 ro nodes.
			// not transferred replica should have 1 ro nodes for balancing.
			if !(replica.RONodesCount()+replica.RWNodesCount() == 2 && replica.RONodesCount() > 0) {
				return false
			}
		}
		return true
	}, 30*time.Second, 2*time.Second)

	// Add some segment on nodes.
	for nodeID := int64(1); nodeID <= 4; nodeID++ {
		suite.distMgr.ChannelDistManager.Update(nodeID)
		suite.distMgr.SegmentDistManager.Update(nodeID)
	}

	suite.Eventually(func() bool {
		for _, r := range replicas {
			replica := suite.meta.ReplicaManager.Get(r.GetID())
			suite.NotNil(replica)
			suite.Equal("rg2", replica.GetResourceGroup())
			if replica.RONodesCount() > 0 {
				return false
			}
			if replica.RWNodesCount() != 1 {
				return false
			}
		}
		return true
	}, 30*time.Second, 2*time.Second)
}

func (suite *ReplicaObserverSuite) TearDownSuite() {
	suite.kv.Close()
	suite.observer.Stop()
}

func TestReplicaObserver(t *testing.T) {
	suite.Run(t, new(ReplicaObserverSuite))
}