mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-03 12:29:36 +08:00
856 lines
27 KiB
Go
856 lines
27 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 datacoord
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"testing"
|
|
|
|
"github.com/samber/lo"
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/suite"
|
|
"go.uber.org/zap"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
globalIDAllocator "github.com/milvus-io/milvus/internal/allocator"
|
|
kvmock "github.com/milvus-io/milvus/internal/kv/mocks"
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
|
"github.com/milvus-io/milvus/pkg/kv/predicates"
|
|
"github.com/milvus-io/milvus/pkg/log"
|
|
"github.com/milvus-io/milvus/pkg/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
|
)
|
|
|
|
func TestChannelManagerSuite(t *testing.T) {
|
|
suite.Run(t, new(ChannelManagerSuite))
|
|
}
|
|
|
|
type ChannelManagerSuite struct {
|
|
suite.Suite
|
|
|
|
mockKv *kvmock.MetaKv
|
|
mockCluster *MockSubCluster
|
|
mockAlloc *globalIDAllocator.MockGlobalIDAllocator
|
|
mockHandler *NMockHandler
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) prepareMeta(chNodes map[string]int64, state datapb.ChannelWatchState) {
|
|
s.SetupTest()
|
|
if chNodes == nil {
|
|
s.mockKv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Once()
|
|
return
|
|
}
|
|
var keys, values []string
|
|
for channel, nodeID := range chNodes {
|
|
keys = append(keys, fmt.Sprintf("channel_store/%d/%s", nodeID, channel))
|
|
info := generateWatchInfo(channel, state)
|
|
bs, err := proto.Marshal(info)
|
|
s.Require().NoError(err)
|
|
values = append(values, string(bs))
|
|
}
|
|
s.mockKv.EXPECT().LoadWithPrefix(mock.Anything).Return(keys, values, nil).Once()
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) checkAssignment(m *ChannelManagerImpl, nodeID int64, channel string, state ChannelState) {
|
|
rwChannel, found := m.GetChannel(nodeID, channel)
|
|
s.True(found)
|
|
s.NotNil(rwChannel)
|
|
s.Equal(channel, rwChannel.GetName())
|
|
sChannel, ok := rwChannel.(*StateChannel)
|
|
s.True(ok)
|
|
s.Equal(state, sChannel.currentState)
|
|
s.EqualValues(nodeID, sChannel.assignedNode)
|
|
s.True(m.Match(nodeID, channel))
|
|
|
|
if nodeID != bufferID {
|
|
gotNode, err := m.FindWatcher(channel)
|
|
s.NoError(err)
|
|
s.EqualValues(gotNode, nodeID)
|
|
}
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) checkNoAssignment(m *ChannelManagerImpl, nodeID int64, channel string) {
|
|
rwChannel, found := m.GetChannel(nodeID, channel)
|
|
s.False(found)
|
|
s.Nil(rwChannel)
|
|
s.False(m.Match(nodeID, channel))
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) SetupTest() {
|
|
s.mockKv = kvmock.NewMetaKv(s.T())
|
|
s.mockCluster = NewMockSubCluster(s.T())
|
|
s.mockHandler = NewNMockHandler(s.T())
|
|
s.mockHandler.EXPECT().GetDataVChanPositions(mock.Anything, mock.Anything).
|
|
RunAndReturn(func(ch RWChannel, partitionID UniqueID) *datapb.VchannelInfo {
|
|
return &datapb.VchannelInfo{
|
|
CollectionID: ch.GetCollectionID(),
|
|
ChannelName: ch.GetName(),
|
|
}
|
|
}).Maybe()
|
|
s.mockKv.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything).RunAndReturn(
|
|
func(save map[string]string, removals []string, preds ...predicates.Predicate) error {
|
|
log.Info("test save and remove", zap.Any("save", save), zap.Any("removals", removals))
|
|
return nil
|
|
}).Maybe()
|
|
s.mockAlloc = globalIDAllocator.NewMockGlobalIDAllocator(s.T())
|
|
s.mockAlloc.EXPECT().AllocOne().Return(1, nil).Maybe()
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) TearDownTest() {}
|
|
|
|
func (s *ChannelManagerSuite) TestAddNode() {
|
|
s.Run("AddNode with empty store", func() {
|
|
s.prepareMeta(nil, 0)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
var testNode int64 = 1
|
|
err = m.AddNode(testNode)
|
|
s.NoError(err)
|
|
|
|
info := m.store.GetNode(testNode)
|
|
s.NotNil(info)
|
|
s.Empty(info.Channels)
|
|
s.Equal(info.NodeID, testNode)
|
|
})
|
|
s.Run("AddNode with channel in bufferID", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": bufferID,
|
|
"ch2": bufferID,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToWatch)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
var (
|
|
testNodeID int64 = 1
|
|
testChannels = []string{"ch1", "ch2"}
|
|
)
|
|
lo.ForEach(testChannels, func(ch string, _ int) {
|
|
s.checkAssignment(m, bufferID, ch, Standby)
|
|
})
|
|
|
|
err = m.AddNode(testNodeID)
|
|
s.NoError(err)
|
|
|
|
lo.ForEach(testChannels, func(ch string, _ int) {
|
|
s.checkAssignment(m, testNodeID, ch, ToWatch)
|
|
})
|
|
})
|
|
s.Run("AddNode with channels evenly in other node", func() {
|
|
var (
|
|
testNodeID int64 = 100
|
|
storedNodeID int64 = 1
|
|
testChannel = "ch1"
|
|
)
|
|
|
|
chNodes := map[string]int64{testChannel: storedNodeID}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_WatchSuccess)
|
|
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
s.checkAssignment(m, storedNodeID, testChannel, Watched)
|
|
|
|
err = m.AddNode(testNodeID)
|
|
s.NoError(err)
|
|
s.ElementsMatch([]int64{100, 1}, m.store.GetNodes())
|
|
s.checkNoAssignment(m, testNodeID, testChannel)
|
|
|
|
testNodeID = 101
|
|
paramtable.Get().Save(paramtable.Get().DataCoordCfg.AutoBalance.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().DataCoordCfg.AutoBalance.Key)
|
|
|
|
err = m.AddNode(testNodeID)
|
|
s.NoError(err)
|
|
s.ElementsMatch([]int64{100, 101, 1}, m.store.GetNodes())
|
|
s.checkNoAssignment(m, testNodeID, testChannel)
|
|
})
|
|
s.Run("AddNode with channels unevenly in other node", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
"ch3": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_WatchSuccess)
|
|
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
var testNodeID int64 = 100
|
|
paramtable.Get().Save(paramtable.Get().DataCoordCfg.AutoBalance.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().DataCoordCfg.AutoBalance.Key)
|
|
|
|
err = m.AddNode(testNodeID)
|
|
s.NoError(err)
|
|
s.ElementsMatch([]int64{testNodeID, 1}, m.store.GetNodes())
|
|
})
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) TestWatch() {
|
|
s.Run("test Watch with empty store", func() {
|
|
s.prepareMeta(nil, 0)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
var testCh string = "ch1"
|
|
|
|
err = m.Watch(context.TODO(), getChannel(testCh, 1))
|
|
s.NoError(err)
|
|
|
|
s.checkAssignment(m, bufferID, testCh, Standby)
|
|
})
|
|
s.Run("test Watch with nodeID in store", func() {
|
|
s.prepareMeta(nil, 0)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
var (
|
|
testCh string = "ch1"
|
|
testNodeID int64 = 1
|
|
)
|
|
err = m.AddNode(testNodeID)
|
|
s.NoError(err)
|
|
s.checkNoAssignment(m, testNodeID, testCh)
|
|
|
|
err = m.Watch(context.TODO(), getChannel(testCh, 1))
|
|
s.NoError(err)
|
|
|
|
s.checkAssignment(m, testNodeID, testCh, ToWatch)
|
|
})
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) TestRelease() {
|
|
s.Run("release not exist nodeID and channel", func() {
|
|
s.prepareMeta(nil, 0)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
err = m.Release(1, "ch1")
|
|
s.Error(err)
|
|
log.Info("error", zap.String("msg", err.Error()))
|
|
|
|
m.AddNode(1)
|
|
err = m.Release(1, "ch1")
|
|
s.Error(err)
|
|
log.Info("error", zap.String("msg", err.Error()))
|
|
})
|
|
|
|
s.Run("release channel in bufferID", func() {
|
|
s.prepareMeta(nil, 0)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
m.Watch(context.TODO(), getChannel("ch1", 1))
|
|
s.checkAssignment(m, bufferID, "ch1", Standby)
|
|
|
|
err = m.Release(bufferID, "ch1")
|
|
s.NoError(err)
|
|
s.checkAssignment(m, bufferID, "ch1", Standby)
|
|
})
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) TestDeleteNode() {
|
|
s.Run("delete not exsit node", func() {
|
|
s.prepareMeta(nil, 0)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
info := m.store.GetNode(1)
|
|
s.Require().Nil(info)
|
|
|
|
err = m.DeleteNode(1)
|
|
s.NoError(err)
|
|
})
|
|
s.Run("delete bufferID", func() {
|
|
s.prepareMeta(nil, 0)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
info := m.store.GetNode(bufferID)
|
|
s.Require().NotNil(info)
|
|
|
|
err = m.DeleteNode(bufferID)
|
|
s.NoError(err)
|
|
})
|
|
|
|
s.Run("delete node without assigment", func() {
|
|
s.prepareMeta(nil, 0)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
err = m.AddNode(1)
|
|
s.NoError(err)
|
|
info := m.store.GetNode(bufferID)
|
|
s.Require().NotNil(info)
|
|
|
|
err = m.DeleteNode(1)
|
|
s.NoError(err)
|
|
info = m.store.GetNode(1)
|
|
s.Nil(info)
|
|
})
|
|
s.Run("delete node with channel", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
"ch3": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_WatchSuccess)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", Watched)
|
|
s.checkAssignment(m, 1, "ch2", Watched)
|
|
s.checkAssignment(m, 1, "ch3", Watched)
|
|
|
|
err = m.AddNode(2)
|
|
s.NoError(err)
|
|
|
|
err = m.DeleteNode(1)
|
|
s.NoError(err)
|
|
info := m.store.GetNode(bufferID)
|
|
s.NotNil(info)
|
|
|
|
s.Equal(3, len(info.Channels))
|
|
s.EqualValues(bufferID, info.NodeID)
|
|
s.checkAssignment(m, bufferID, "ch1", Standby)
|
|
s.checkAssignment(m, bufferID, "ch2", Standby)
|
|
s.checkAssignment(m, bufferID, "ch3", Standby)
|
|
|
|
info = m.store.GetNode(1)
|
|
s.Nil(info)
|
|
})
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) TestFindWatcher() {
|
|
chNodes := map[string]int64{
|
|
"ch1": bufferID,
|
|
"ch2": bufferID,
|
|
"ch3": 1,
|
|
"ch4": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_WatchSuccess)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
tests := []struct {
|
|
description string
|
|
testCh string
|
|
|
|
outNodeID int64
|
|
outError bool
|
|
}{
|
|
{"channel not exist", "ch-notexist", 0, true},
|
|
{"channel in bufferID", "ch1", bufferID, true},
|
|
{"channel in bufferID", "ch2", bufferID, true},
|
|
{"channel in nodeID=1", "ch3", 1, false},
|
|
{"channel in nodeID=1", "ch4", 1, false},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
s.Run(test.description, func() {
|
|
gotID, gotErr := m.FindWatcher(test.testCh)
|
|
s.EqualValues(test.outNodeID, gotID)
|
|
if test.outError {
|
|
s.Error(gotErr)
|
|
} else {
|
|
s.NoError(gotErr)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) TestAdvanceChannelState() {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
s.Run("advance statndby with no available nodes", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": bufferID,
|
|
"ch2": bufferID,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToWatch)
|
|
s.mockHandler.EXPECT().CheckShouldDropChannel(mock.Anything).Return(false)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, bufferID, "ch1", Standby)
|
|
s.checkAssignment(m, bufferID, "ch2", Standby)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, bufferID, "ch1", Standby)
|
|
s.checkAssignment(m, bufferID, "ch2", Standby)
|
|
})
|
|
|
|
s.Run("advance statndby with node 1", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": bufferID,
|
|
"ch2": bufferID,
|
|
"ch3": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_WatchSuccess)
|
|
s.mockHandler.EXPECT().CheckShouldDropChannel(mock.Anything).Return(false).Times(2)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, bufferID, "ch1", Standby)
|
|
s.checkAssignment(m, bufferID, "ch2", Standby)
|
|
s.checkAssignment(m, 1, "ch3", Watched)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
})
|
|
s.Run("advance towatch channels notify success check success", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToWatch)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Watching)
|
|
s.checkAssignment(m, 1, "ch2", Watching)
|
|
})
|
|
s.Run("advance watching channels check no progress", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToWatch)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Watching)
|
|
s.checkAssignment(m, 1, "ch2", Watching)
|
|
|
|
s.mockCluster.EXPECT().CheckChannelOperationProgress(mock.Anything, mock.Anything, mock.Anything).
|
|
Return(&datapb.ChannelOperationProgressResponse{State: datapb.ChannelWatchState_ToWatch}, nil).Twice()
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Watching)
|
|
s.checkAssignment(m, 1, "ch2", Watching)
|
|
})
|
|
s.Run("advance watching channels check ErrNodeNotFound", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToWatch)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Watching)
|
|
s.checkAssignment(m, 1, "ch2", Watching)
|
|
|
|
s.mockCluster.EXPECT().CheckChannelOperationProgress(mock.Anything, mock.Anything, mock.Anything).
|
|
Return(nil, merr.WrapErrNodeNotFound(1)).Twice()
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Standby)
|
|
s.checkAssignment(m, 1, "ch2", Standby)
|
|
})
|
|
|
|
s.Run("advance watching channels check watch success", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToWatch)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Watching)
|
|
s.checkAssignment(m, 1, "ch2", Watching)
|
|
|
|
s.mockCluster.EXPECT().CheckChannelOperationProgress(mock.Anything, mock.Anything, mock.Anything).
|
|
Return(&datapb.ChannelOperationProgressResponse{State: datapb.ChannelWatchState_WatchSuccess}, nil).Twice()
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Watched)
|
|
s.checkAssignment(m, 1, "ch2", Watched)
|
|
})
|
|
s.Run("advance watching channels check watch fail", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToWatch)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(2)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Watching)
|
|
s.checkAssignment(m, 1, "ch2", Watching)
|
|
|
|
s.mockCluster.EXPECT().CheckChannelOperationProgress(mock.Anything, mock.Anything, mock.Anything).
|
|
Return(&datapb.ChannelOperationProgressResponse{State: datapb.ChannelWatchState_WatchFailure}, nil).Twice()
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Standby)
|
|
s.checkAssignment(m, 1, "ch2", Standby)
|
|
|
|
s.mockHandler.EXPECT().CheckShouldDropChannel(mock.Anything).Return(false)
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
})
|
|
s.Run("advance releasing channels check release no progress", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToRelease)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToRelease)
|
|
s.checkAssignment(m, 1, "ch2", ToRelease)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Releasing)
|
|
s.checkAssignment(m, 1, "ch2", Releasing)
|
|
|
|
s.mockCluster.EXPECT().CheckChannelOperationProgress(mock.Anything, mock.Anything, mock.Anything).
|
|
Return(&datapb.ChannelOperationProgressResponse{State: datapb.ChannelWatchState_ToRelease}, nil).Twice()
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Releasing)
|
|
s.checkAssignment(m, 1, "ch2", Releasing)
|
|
})
|
|
s.Run("advance releasing channels check ErrNodeNotFound", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToRelease)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToRelease)
|
|
s.checkAssignment(m, 1, "ch2", ToRelease)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Releasing)
|
|
s.checkAssignment(m, 1, "ch2", Releasing)
|
|
|
|
s.mockCluster.EXPECT().CheckChannelOperationProgress(mock.Anything, mock.Anything, mock.Anything).
|
|
Return(nil, merr.WrapErrNodeNotFound(1)).Twice()
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Standby)
|
|
s.checkAssignment(m, 1, "ch2", Standby)
|
|
})
|
|
s.Run("advance releasing channels check release success", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToRelease)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToRelease)
|
|
s.checkAssignment(m, 1, "ch2", ToRelease)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Releasing)
|
|
s.checkAssignment(m, 1, "ch2", Releasing)
|
|
|
|
s.mockCluster.EXPECT().CheckChannelOperationProgress(mock.Anything, mock.Anything, mock.Anything).
|
|
Return(&datapb.ChannelOperationProgressResponse{State: datapb.ChannelWatchState_ReleaseSuccess}, nil).Twice()
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Standby)
|
|
s.checkAssignment(m, 1, "ch2", Standby)
|
|
|
|
s.mockHandler.EXPECT().CheckShouldDropChannel(mock.Anything).Return(false)
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
})
|
|
s.Run("advance releasing channels check release fail", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToRelease)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToRelease)
|
|
s.checkAssignment(m, 1, "ch2", ToRelease)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Releasing)
|
|
s.checkAssignment(m, 1, "ch2", Releasing)
|
|
|
|
s.mockCluster.EXPECT().CheckChannelOperationProgress(mock.Anything, mock.Anything, mock.Anything).
|
|
Return(&datapb.ChannelOperationProgressResponse{State: datapb.ChannelWatchState_ReleaseFailure}, nil).Twice()
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Standby)
|
|
s.checkAssignment(m, 1, "ch2", Standby)
|
|
|
|
s.mockHandler.EXPECT().CheckShouldDropChannel(mock.Anything).Return(false)
|
|
m.AdvanceChannelState(ctx)
|
|
// TODO, donot assign to abnormal nodes
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
})
|
|
s.Run("advance towatch channels notify fail", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToWatch)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).
|
|
Return(fmt.Errorf("mock error")).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", ToWatch)
|
|
s.checkAssignment(m, 1, "ch2", ToWatch)
|
|
})
|
|
s.Run("advance to release channels notify success", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToRelease)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToRelease)
|
|
s.checkAssignment(m, 1, "ch2", ToRelease)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", Releasing)
|
|
s.checkAssignment(m, 1, "ch2", Releasing)
|
|
})
|
|
s.Run("advance to release channels notify fail", func() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToRelease)
|
|
s.mockCluster.EXPECT().NotifyChannelOperation(mock.Anything, mock.Anything, mock.Anything).
|
|
Return(fmt.Errorf("mock error")).Twice()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 1, "ch1", ToRelease)
|
|
s.checkAssignment(m, 1, "ch2", ToRelease)
|
|
|
|
m.AdvanceChannelState(ctx)
|
|
s.checkAssignment(m, 1, "ch1", ToRelease)
|
|
s.checkAssignment(m, 1, "ch2", ToRelease)
|
|
})
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) TestStartup() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
"ch3": 3,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToRelease)
|
|
s.mockHandler.EXPECT().CheckShouldDropChannel(mock.Anything).Return(false)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
var (
|
|
legacyNodes = []int64{1}
|
|
allNodes = []int64{1}
|
|
)
|
|
err = m.Startup(context.TODO(), legacyNodes, allNodes)
|
|
s.NoError(err)
|
|
|
|
s.checkAssignment(m, 1, "ch1", Legacy)
|
|
s.checkAssignment(m, 1, "ch2", Legacy)
|
|
s.checkAssignment(m, bufferID, "ch3", Standby)
|
|
|
|
err = m.DeleteNode(1)
|
|
s.NoError(err)
|
|
|
|
s.checkAssignment(m, bufferID, "ch1", Standby)
|
|
s.checkAssignment(m, bufferID, "ch2", Standby)
|
|
|
|
err = m.AddNode(2)
|
|
s.NoError(err)
|
|
s.checkAssignment(m, 2, "ch1", ToWatch)
|
|
s.checkAssignment(m, 2, "ch2", ToWatch)
|
|
s.checkAssignment(m, 2, "ch3", ToWatch)
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) TestStartupNilSchema() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
"ch3": 3,
|
|
}
|
|
var keys, values []string
|
|
for channel, nodeID := range chNodes {
|
|
keys = append(keys, fmt.Sprintf("channel_store/%d/%s", nodeID, channel))
|
|
info := generateWatchInfo(channel, datapb.ChannelWatchState_ToRelease)
|
|
info.Schema = nil
|
|
bs, err := proto.Marshal(info)
|
|
s.Require().NoError(err)
|
|
values = append(values, string(bs))
|
|
}
|
|
s.mockKv.EXPECT().LoadWithPrefix(mock.Anything).Return(keys, values, nil).Once()
|
|
s.mockHandler.EXPECT().CheckShouldDropChannel(mock.Anything).Return(false)
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
err = m.Startup(context.TODO(), nil, []int64{1, 3})
|
|
s.Require().NoError(err)
|
|
|
|
for ch, node := range chNodes {
|
|
channel, got := m.GetChannel(node, ch)
|
|
s.Require().True(got)
|
|
s.Nil(channel.GetSchema())
|
|
s.Equal(ch, channel.GetName())
|
|
log.Info("Recovered nil schema channel", zap.Any("channel", channel))
|
|
}
|
|
|
|
s.mockHandler.EXPECT().GetCollection(mock.Anything, mock.Anything).Return(
|
|
&collectionInfo{ID: 111, Schema: &schemapb.CollectionSchema{Name: "coll111"}},
|
|
nil,
|
|
)
|
|
|
|
err = m.DeleteNode(1)
|
|
s.Require().NoError(err)
|
|
|
|
err = m.DeleteNode(3)
|
|
s.Require().NoError(err)
|
|
|
|
s.checkAssignment(m, bufferID, "ch1", Standby)
|
|
s.checkAssignment(m, bufferID, "ch2", Standby)
|
|
s.checkAssignment(m, bufferID, "ch3", Standby)
|
|
|
|
for ch := range chNodes {
|
|
channel, got := m.GetChannel(bufferID, ch)
|
|
s.Require().True(got)
|
|
s.NotNil(channel.GetSchema())
|
|
s.Equal(ch, channel.GetName())
|
|
|
|
s.NotNil(channel.GetWatchInfo())
|
|
s.NotNil(channel.GetWatchInfo().Schema)
|
|
log.Info("Recovered non-nil schema channel", zap.Any("channel", channel))
|
|
}
|
|
|
|
err = m.AddNode(7)
|
|
s.Require().NoError(err)
|
|
s.checkAssignment(m, 7, "ch1", ToWatch)
|
|
s.checkAssignment(m, 7, "ch2", ToWatch)
|
|
s.checkAssignment(m, 7, "ch3", ToWatch)
|
|
|
|
for ch := range chNodes {
|
|
channel, got := m.GetChannel(7, ch)
|
|
s.Require().True(got)
|
|
s.NotNil(channel.GetSchema())
|
|
s.Equal(ch, channel.GetName())
|
|
|
|
s.NotNil(channel.GetWatchInfo())
|
|
s.NotNil(channel.GetWatchInfo().Schema)
|
|
log.Info("non-nil schema channel", zap.Any("channel", channel))
|
|
}
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) TestStartupRootCoordFailed() {
|
|
chNodes := map[string]int64{
|
|
"ch1": 1,
|
|
"ch2": 1,
|
|
"ch3": 1,
|
|
"ch4": bufferID,
|
|
}
|
|
s.prepareMeta(chNodes, datapb.ChannelWatchState_ToWatch)
|
|
|
|
s.mockAlloc = globalIDAllocator.NewMockGlobalIDAllocator(s.T())
|
|
s.mockAlloc.EXPECT().AllocOne().Return(0, errors.New("mock error")).Maybe()
|
|
m, err := NewChannelManager(s.mockKv, s.mockHandler, s.mockCluster, s.mockAlloc)
|
|
s.Require().NoError(err)
|
|
|
|
err = m.Startup(context.TODO(), nil, []int64{2})
|
|
s.Error(err)
|
|
|
|
err = m.Startup(context.TODO(), nil, []int64{1, 2})
|
|
s.Error(err)
|
|
}
|
|
|
|
func (s *ChannelManagerSuite) TestCheckLoop() {}
|
|
func (s *ChannelManagerSuite) TestGet() {}
|
|
|
|
func (s *ChannelManagerSuite) TestGetChannelWatchInfos() {
|
|
store := NewMockRWChannelStore(s.T())
|
|
store.EXPECT().GetNodesChannels().Return([]*NodeChannelInfo{
|
|
{
|
|
NodeID: 1,
|
|
Channels: map[string]RWChannel{
|
|
"ch1": &channelMeta{
|
|
WatchInfo: &datapb.ChannelWatchInfo{
|
|
Vchan: &datapb.VchannelInfo{
|
|
ChannelName: "ch1",
|
|
},
|
|
StartTs: 100,
|
|
State: datapb.ChannelWatchState_ToWatch,
|
|
OpID: 1,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
NodeID: 2,
|
|
Channels: map[string]RWChannel{
|
|
"ch2": &channelMeta{
|
|
WatchInfo: &datapb.ChannelWatchInfo{
|
|
Vchan: &datapb.VchannelInfo{
|
|
ChannelName: "ch2",
|
|
},
|
|
StartTs: 10,
|
|
State: datapb.ChannelWatchState_WatchSuccess,
|
|
OpID: 1,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
|
|
cm := &ChannelManagerImpl{store: store}
|
|
infos := cm.GetChannelWatchInfos()
|
|
s.Equal(2, len(infos))
|
|
s.Equal("ch1", infos[1]["ch1"].GetVchan().ChannelName)
|
|
s.Equal("ch2", infos[2]["ch2"].GetVchan().ChannelName)
|
|
|
|
// test empty value
|
|
store.EXPECT().GetNodesChannels().Unset()
|
|
store.EXPECT().GetNodesChannels().Return([]*NodeChannelInfo{})
|
|
infos = cm.GetChannelWatchInfos()
|
|
s.Equal(0, len(infos))
|
|
}
|