mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-06 05:49:24 +08:00
b1eacb2ae8
issue: https://github.com/milvus-io/milvus/issues/25309 Signed-off-by: yiwangdr <yiwangdr@gmail.com>
1848 lines
54 KiB
Go
1848 lines
54 KiB
Go
package datacoord
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/samber/lo"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/stretchr/testify/suite"
|
|
"go.uber.org/zap"
|
|
"google.golang.org/grpc/codes"
|
|
grpcStatus "google.golang.org/grpc/status"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
"github.com/milvus-io/milvus/internal/metastore/mocks"
|
|
"github.com/milvus-io/milvus/internal/metastore/model"
|
|
mocks2 "github.com/milvus-io/milvus/internal/mocks"
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
|
"github.com/milvus-io/milvus/internal/proto/indexpb"
|
|
"github.com/milvus-io/milvus/internal/proto/internalpb"
|
|
"github.com/milvus-io/milvus/internal/types"
|
|
"github.com/milvus-io/milvus/pkg/log"
|
|
"github.com/milvus-io/milvus/pkg/mq/msgstream"
|
|
"github.com/milvus-io/milvus/pkg/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/util/metautil"
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/util/tsoutil"
|
|
)
|
|
|
|
type ServerSuite struct {
|
|
suite.Suite
|
|
|
|
testServer *Server
|
|
mockChMgr *MockChannelManager
|
|
}
|
|
|
|
func WithChannelManager(cm ChannelManager) Option {
|
|
return func(svr *Server) {
|
|
svr.channelManager = cm
|
|
}
|
|
}
|
|
|
|
func (s *ServerSuite) SetupTest() {
|
|
s.mockChMgr = NewMockChannelManager(s.T())
|
|
s.mockChMgr.EXPECT().Startup(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
|
|
s.mockChMgr.EXPECT().Close().Maybe()
|
|
|
|
s.testServer = newTestServer(s.T(), WithChannelManager(s.mockChMgr))
|
|
if s.testServer.channelManager != nil {
|
|
s.testServer.channelManager.Close()
|
|
}
|
|
}
|
|
|
|
func (s *ServerSuite) TearDownTest() {
|
|
if s.testServer != nil {
|
|
log.Info("ServerSuite tears down test", zap.String("name", s.T().Name()))
|
|
closeTestServer(s.T(), s.testServer)
|
|
}
|
|
}
|
|
|
|
func TestServerSuite(t *testing.T) {
|
|
suite.Run(t, new(ServerSuite))
|
|
}
|
|
|
|
func genMsg(msgType commonpb.MsgType, ch string, t Timestamp, sourceID int64) *msgstream.DataNodeTtMsg {
|
|
return &msgstream.DataNodeTtMsg{
|
|
BaseMsg: msgstream.BaseMsg{
|
|
HashValues: []uint32{0},
|
|
},
|
|
DataNodeTtMsg: msgpb.DataNodeTtMsg{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: msgType,
|
|
Timestamp: t,
|
|
SourceID: sourceID,
|
|
},
|
|
ChannelName: ch,
|
|
Timestamp: t,
|
|
SegmentsStats: []*commonpb.SegmentStats{{SegmentID: 2, NumRows: 100}},
|
|
},
|
|
}
|
|
}
|
|
|
|
func (s *ServerSuite) TestHandleDataNodeTtMsg() {
|
|
var (
|
|
chanName = "ch-1"
|
|
collID int64 = 100
|
|
sourceID int64 = 1
|
|
)
|
|
s.testServer.meta.AddCollection(&collectionInfo{
|
|
ID: collID,
|
|
Schema: newTestSchema(),
|
|
Partitions: []int64{10},
|
|
})
|
|
resp, err := s.testServer.AssignSegmentID(context.TODO(), &datapb.AssignSegmentIDRequest{
|
|
NodeID: sourceID,
|
|
SegmentIDRequests: []*datapb.SegmentIDRequest{
|
|
{
|
|
CollectionID: collID,
|
|
PartitionID: 10,
|
|
ChannelName: chanName,
|
|
Count: 100,
|
|
},
|
|
},
|
|
})
|
|
s.Require().NoError(err)
|
|
s.Require().True(merr.Ok(resp.GetStatus()))
|
|
s.Equal(1, len(resp.GetSegIDAssignments()))
|
|
assign := resp.GetSegIDAssignments()[0]
|
|
|
|
assignedSegmentID := resp.SegIDAssignments[0].SegID
|
|
segment := s.testServer.meta.GetHealthySegment(assignedSegmentID)
|
|
s.Require().NotNil(segment)
|
|
s.Equal(1, len(segment.allocations))
|
|
|
|
ts := tsoutil.AddPhysicalDurationOnTs(assign.ExpireTime, -3*time.Minute)
|
|
msg := genMsg(commonpb.MsgType_DataNodeTt, chanName, ts, sourceID)
|
|
msg.SegmentsStats = append(msg.SegmentsStats, &commonpb.SegmentStats{
|
|
SegmentID: assign.GetSegID(),
|
|
NumRows: 1,
|
|
})
|
|
mockCluster := NewMockCluster(s.T())
|
|
mockCluster.EXPECT().Close().Once()
|
|
mockCluster.EXPECT().Flush(mock.Anything, sourceID, chanName, mock.Anything).RunAndReturn(
|
|
func(ctx context.Context, nodeID int64, channel string, segments []*datapb.SegmentInfo) error {
|
|
s.EqualValues(chanName, channel)
|
|
s.EqualValues(sourceID, nodeID)
|
|
s.Equal(1, len(segments))
|
|
s.EqualValues(2, segments[0].GetID())
|
|
|
|
return fmt.Errorf("mock error")
|
|
}).Once()
|
|
s.testServer.cluster = mockCluster
|
|
s.mockChMgr.EXPECT().Match(sourceID, chanName).Return(true).Twice()
|
|
|
|
err = s.testServer.handleDataNodeTtMsg(context.TODO(), &msg.DataNodeTtMsg)
|
|
s.NoError(err)
|
|
|
|
tt := tsoutil.AddPhysicalDurationOnTs(assign.ExpireTime, 48*time.Hour)
|
|
msg = genMsg(commonpb.MsgType_DataNodeTt, chanName, tt, sourceID)
|
|
msg.SegmentsStats = append(msg.SegmentsStats, &commonpb.SegmentStats{
|
|
SegmentID: assign.GetSegID(),
|
|
NumRows: 1,
|
|
})
|
|
|
|
err = s.testServer.handleDataNodeTtMsg(context.TODO(), &msg.DataNodeTtMsg)
|
|
s.Error(err)
|
|
}
|
|
|
|
// restart the server for config DataNodeTimeTickByRPC=false
|
|
func (s *ServerSuite) initSuiteForTtChannel() {
|
|
s.testServer.serverLoopWg.Add(1)
|
|
s.testServer.startDataNodeTtLoop(s.testServer.serverLoopCtx)
|
|
|
|
s.testServer.meta.AddCollection(&collectionInfo{
|
|
ID: 1,
|
|
Schema: newTestSchema(),
|
|
Partitions: []int64{10},
|
|
})
|
|
}
|
|
|
|
func (s *ServerSuite) TestDataNodeTtChannel_ExpireAfterTt() {
|
|
s.initSuiteForTtChannel()
|
|
|
|
ctx := context.TODO()
|
|
ttMsgStream, err := s.testServer.factory.NewMsgStream(ctx)
|
|
s.Require().NoError(err)
|
|
|
|
ttMsgStream.AsProducer([]string{paramtable.Get().CommonCfg.DataCoordTimeTick.GetValue()})
|
|
defer ttMsgStream.Close()
|
|
|
|
var (
|
|
sourceID int64 = 9997
|
|
chanName = "ch-1"
|
|
signal = make(chan struct{})
|
|
collID int64 = 1
|
|
)
|
|
mockCluster := NewMockCluster(s.T())
|
|
mockCluster.EXPECT().Close().Once()
|
|
mockCluster.EXPECT().Flush(mock.Anything, sourceID, chanName, mock.Anything).RunAndReturn(
|
|
func(ctx context.Context, nodeID int64, channel string, segments []*datapb.SegmentInfo) error {
|
|
s.EqualValues(chanName, channel)
|
|
s.EqualValues(sourceID, nodeID)
|
|
s.Equal(1, len(segments))
|
|
s.EqualValues(2, segments[0].GetID())
|
|
|
|
signal <- struct{}{}
|
|
return nil
|
|
}).Once()
|
|
s.testServer.cluster = mockCluster
|
|
s.mockChMgr.EXPECT().Match(sourceID, chanName).Return(true).Once()
|
|
|
|
resp, err := s.testServer.AssignSegmentID(context.TODO(), &datapb.AssignSegmentIDRequest{
|
|
NodeID: sourceID,
|
|
SegmentIDRequests: []*datapb.SegmentIDRequest{
|
|
{
|
|
CollectionID: collID,
|
|
PartitionID: 10,
|
|
ChannelName: chanName,
|
|
Count: 100,
|
|
},
|
|
},
|
|
})
|
|
s.Require().NoError(err)
|
|
s.Require().True(merr.Ok(resp.GetStatus()))
|
|
s.Equal(1, len(resp.GetSegIDAssignments()))
|
|
|
|
assignedSegmentID := resp.SegIDAssignments[0].SegID
|
|
segment := s.testServer.meta.GetHealthySegment(assignedSegmentID)
|
|
s.Require().NotNil(segment)
|
|
s.Equal(1, len(segment.allocations))
|
|
|
|
msgPack := msgstream.MsgPack{}
|
|
tt := tsoutil.AddPhysicalDurationOnTs(resp.SegIDAssignments[0].ExpireTime, 48*time.Hour)
|
|
msg := genMsg(commonpb.MsgType_DataNodeTt, "ch-1", tt, sourceID)
|
|
msgPack.Msgs = append(msgPack.Msgs, msg)
|
|
err = ttMsgStream.Produce(&msgPack)
|
|
s.Require().NoError(err)
|
|
|
|
<-signal
|
|
segment = s.testServer.meta.GetHealthySegment(assignedSegmentID)
|
|
s.NotNil(segment)
|
|
s.Equal(0, len(segment.allocations))
|
|
}
|
|
|
|
func (s *ServerSuite) TestDataNodeTtChannel_FlushWithDiffChan() {
|
|
s.initSuiteForTtChannel()
|
|
|
|
ctx := context.TODO()
|
|
ttMsgStream, err := s.testServer.factory.NewMsgStream(ctx)
|
|
s.Require().NoError(err)
|
|
|
|
ttMsgStream.AsProducer([]string{paramtable.Get().CommonCfg.DataCoordTimeTick.GetValue()})
|
|
defer ttMsgStream.Close()
|
|
|
|
var (
|
|
sourceID int64 = 9998
|
|
chanName = "ch-1"
|
|
signal = make(chan struct{})
|
|
collID int64 = 1
|
|
)
|
|
|
|
mockCluster := NewMockCluster(s.T())
|
|
mockCluster.EXPECT().Close().Once()
|
|
mockCluster.EXPECT().Flush(mock.Anything, sourceID, chanName, mock.Anything).RunAndReturn(
|
|
func(ctx context.Context, nodeID int64, channel string, segments []*datapb.SegmentInfo) error {
|
|
s.EqualValues(chanName, channel)
|
|
s.EqualValues(sourceID, nodeID)
|
|
s.Equal(1, len(segments))
|
|
|
|
signal <- struct{}{}
|
|
return nil
|
|
}).Once()
|
|
mockCluster.EXPECT().FlushChannels(mock.Anything, sourceID, mock.Anything, []string{chanName}).Return(nil).Once()
|
|
s.testServer.cluster = mockCluster
|
|
|
|
s.mockChMgr.EXPECT().Match(sourceID, chanName).Return(true).Once()
|
|
s.mockChMgr.EXPECT().GetNodeChannelsByCollectionID(collID).Return(map[int64][]string{
|
|
sourceID: {chanName},
|
|
})
|
|
|
|
resp, err := s.testServer.AssignSegmentID(ctx, &datapb.AssignSegmentIDRequest{
|
|
NodeID: sourceID,
|
|
SegmentIDRequests: []*datapb.SegmentIDRequest{
|
|
{
|
|
CollectionID: collID,
|
|
PartitionID: 10,
|
|
ChannelName: chanName,
|
|
Count: 100,
|
|
},
|
|
{
|
|
CollectionID: collID,
|
|
PartitionID: 10,
|
|
ChannelName: "ch-2",
|
|
Count: 100,
|
|
},
|
|
},
|
|
})
|
|
s.Require().NoError(err)
|
|
s.Require().True(merr.Ok(resp.GetStatus()))
|
|
s.Equal(2, len(resp.GetSegIDAssignments()))
|
|
var assign *datapb.SegmentIDAssignment
|
|
for _, segment := range resp.SegIDAssignments {
|
|
if segment.GetChannelName() == chanName {
|
|
assign = segment
|
|
break
|
|
}
|
|
}
|
|
s.Require().NotNil(assign)
|
|
|
|
resp2, err := s.testServer.Flush(ctx, &datapb.FlushRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Flush,
|
|
SourceID: sourceID,
|
|
},
|
|
CollectionID: collID,
|
|
})
|
|
s.Require().NoError(err)
|
|
s.Require().True(merr.Ok(resp2.GetStatus()))
|
|
|
|
msgPack := msgstream.MsgPack{}
|
|
msg := genMsg(commonpb.MsgType_DataNodeTt, chanName, assign.ExpireTime, sourceID)
|
|
msg.SegmentsStats = append(msg.SegmentsStats, &commonpb.SegmentStats{
|
|
SegmentID: assign.GetSegID(),
|
|
NumRows: 1,
|
|
})
|
|
msgPack.Msgs = append(msgPack.Msgs, msg)
|
|
err = ttMsgStream.Produce(&msgPack)
|
|
s.NoError(err)
|
|
|
|
<-signal
|
|
}
|
|
|
|
func (s *ServerSuite) TestDataNodeTtChannel_SegmentFlushAfterTt() {
|
|
s.initSuiteForTtChannel()
|
|
|
|
var (
|
|
sourceID int64 = 9999
|
|
chanName = "ch-1"
|
|
signal = make(chan struct{})
|
|
collID int64 = 1
|
|
)
|
|
mockCluster := NewMockCluster(s.T())
|
|
mockCluster.EXPECT().Close().Once()
|
|
mockCluster.EXPECT().Flush(mock.Anything, sourceID, chanName, mock.Anything).RunAndReturn(
|
|
func(ctx context.Context, nodeID int64, channel string, segments []*datapb.SegmentInfo) error {
|
|
s.EqualValues(chanName, channel)
|
|
s.EqualValues(sourceID, nodeID)
|
|
s.Equal(1, len(segments))
|
|
|
|
signal <- struct{}{}
|
|
return nil
|
|
}).Once()
|
|
mockCluster.EXPECT().FlushChannels(mock.Anything, sourceID, mock.Anything, []string{chanName}).Return(nil).Once()
|
|
s.testServer.cluster = mockCluster
|
|
|
|
s.mockChMgr.EXPECT().Match(sourceID, chanName).Return(true).Once()
|
|
s.mockChMgr.EXPECT().GetNodeChannelsByCollectionID(collID).Return(map[int64][]string{
|
|
sourceID: {chanName},
|
|
})
|
|
|
|
ctx := context.TODO()
|
|
ttMsgStream, err := s.testServer.factory.NewMsgStream(ctx)
|
|
s.Require().NoError(err)
|
|
|
|
ttMsgStream.AsProducer([]string{paramtable.Get().CommonCfg.DataCoordTimeTick.GetValue()})
|
|
defer ttMsgStream.Close()
|
|
|
|
resp, err := s.testServer.AssignSegmentID(context.TODO(), &datapb.AssignSegmentIDRequest{
|
|
SegmentIDRequests: []*datapb.SegmentIDRequest{
|
|
{
|
|
CollectionID: 1,
|
|
PartitionID: 10,
|
|
ChannelName: chanName,
|
|
Count: 100,
|
|
},
|
|
},
|
|
})
|
|
s.Require().NoError(err)
|
|
s.Require().True(merr.Ok(resp.GetStatus()))
|
|
s.Require().Equal(1, len(resp.GetSegIDAssignments()))
|
|
|
|
assign := resp.GetSegIDAssignments()[0]
|
|
|
|
resp2, err := s.testServer.Flush(ctx, &datapb.FlushRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Flush,
|
|
},
|
|
CollectionID: 1,
|
|
})
|
|
s.Require().NoError(err)
|
|
s.Require().True(merr.Ok(resp2.GetStatus()))
|
|
|
|
msgPack := msgstream.MsgPack{}
|
|
msg := genMsg(commonpb.MsgType_DataNodeTt, "ch-1", assign.ExpireTime, 9999)
|
|
msg.SegmentsStats = append(msg.SegmentsStats, &commonpb.SegmentStats{
|
|
SegmentID: assign.GetSegID(),
|
|
NumRows: 1,
|
|
})
|
|
msgPack.Msgs = append(msgPack.Msgs, msg)
|
|
err = ttMsgStream.Produce(&msgPack)
|
|
s.Require().NoError(err)
|
|
|
|
<-signal
|
|
}
|
|
|
|
func (s *ServerSuite) TestGetFlushState_ByFlushTs() {
|
|
s.mockChMgr.EXPECT().GetChannelsByCollectionID(int64(0)).
|
|
Return([]RWChannel{&channelMeta{Name: "ch1", CollectionID: 0}}).Times(3)
|
|
|
|
s.mockChMgr.EXPECT().GetChannelsByCollectionID(int64(1)).Return(nil).Times(1)
|
|
tests := []struct {
|
|
description string
|
|
inTs Timestamp
|
|
|
|
expected bool
|
|
}{
|
|
{"channel cp > flush ts", 11, true},
|
|
{"channel cp = flush ts", 12, true},
|
|
{"channel cp < flush ts", 13, false},
|
|
}
|
|
|
|
err := s.testServer.meta.UpdateChannelCheckpoint("ch1", &msgpb.MsgPosition{
|
|
MsgID: []byte{1},
|
|
Timestamp: 12,
|
|
})
|
|
s.Require().NoError(err)
|
|
for _, test := range tests {
|
|
s.Run(test.description, func() {
|
|
resp, err := s.testServer.GetFlushState(context.TODO(), &datapb.GetFlushStateRequest{FlushTs: test.inTs})
|
|
s.NoError(err)
|
|
s.EqualValues(&milvuspb.GetFlushStateResponse{
|
|
Status: merr.Success(),
|
|
Flushed: test.expected,
|
|
}, resp)
|
|
})
|
|
}
|
|
|
|
resp, err := s.testServer.GetFlushState(context.TODO(), &datapb.GetFlushStateRequest{CollectionID: 1, FlushTs: 13})
|
|
s.NoError(err)
|
|
s.EqualValues(&milvuspb.GetFlushStateResponse{
|
|
Status: merr.Success(),
|
|
Flushed: true,
|
|
}, resp)
|
|
}
|
|
|
|
func (s *ServerSuite) TestGetFlushState_BySegment() {
|
|
s.mockChMgr.EXPECT().GetChannelsByCollectionID(mock.Anything).
|
|
Return([]RWChannel{&channelMeta{Name: "ch1", CollectionID: 0}}).Times(3)
|
|
|
|
tests := []struct {
|
|
description string
|
|
segID int64
|
|
state commonpb.SegmentState
|
|
|
|
expected bool
|
|
}{
|
|
{"flushed seg1", 1, commonpb.SegmentState_Flushed, true},
|
|
{"flushed seg2", 2, commonpb.SegmentState_Flushed, true},
|
|
{"sealed seg3", 3, commonpb.SegmentState_Sealed, false},
|
|
{"compacted/dropped seg4", 4, commonpb.SegmentState_Dropped, true},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
s.Run(test.description, func() {
|
|
err := s.testServer.meta.AddSegment(context.TODO(), &SegmentInfo{
|
|
SegmentInfo: &datapb.SegmentInfo{
|
|
ID: test.segID,
|
|
State: test.state,
|
|
},
|
|
})
|
|
|
|
s.Require().NoError(err)
|
|
err = s.testServer.meta.UpdateChannelCheckpoint("ch1", &msgpb.MsgPosition{
|
|
MsgID: []byte{1},
|
|
Timestamp: 12,
|
|
})
|
|
s.Require().NoError(err)
|
|
|
|
resp, err := s.testServer.GetFlushState(context.TODO(), &datapb.GetFlushStateRequest{SegmentIDs: []int64{test.segID}})
|
|
s.NoError(err)
|
|
s.EqualValues(&milvuspb.GetFlushStateResponse{
|
|
Status: merr.Success(),
|
|
Flushed: test.expected,
|
|
}, resp)
|
|
})
|
|
}
|
|
}
|
|
|
|
func (s *ServerSuite) TestSaveBinlogPath_ClosedServer() {
|
|
s.TearDownTest()
|
|
resp, err := s.testServer.SaveBinlogPaths(context.Background(), &datapb.SaveBinlogPathsRequest{
|
|
SegmentID: 1,
|
|
Channel: "test",
|
|
})
|
|
s.NoError(err)
|
|
s.ErrorIs(merr.Error(resp), merr.ErrServiceNotReady)
|
|
}
|
|
|
|
func (s *ServerSuite) TestSaveBinlogPath_ChannelNotMatch() {
|
|
s.mockChMgr.EXPECT().Match(mock.Anything, mock.Anything).Return(false)
|
|
resp, err := s.testServer.SaveBinlogPaths(context.Background(), &datapb.SaveBinlogPathsRequest{
|
|
SegmentID: 1,
|
|
Channel: "test",
|
|
})
|
|
s.NoError(err)
|
|
s.ErrorIs(merr.Error(resp), merr.ErrChannelNotFound)
|
|
}
|
|
|
|
func (s *ServerSuite) TestSaveBinlogPath_SaveUnhealthySegment() {
|
|
s.mockChMgr.EXPECT().Match(int64(0), "ch1").Return(true)
|
|
s.testServer.meta.AddCollection(&collectionInfo{ID: 0})
|
|
|
|
segments := map[int64]commonpb.SegmentState{
|
|
1: commonpb.SegmentState_NotExist,
|
|
2: commonpb.SegmentState_Dropped,
|
|
}
|
|
for segID, state := range segments {
|
|
info := &datapb.SegmentInfo{
|
|
ID: segID,
|
|
InsertChannel: "ch1",
|
|
State: state,
|
|
}
|
|
err := s.testServer.meta.AddSegment(context.TODO(), NewSegmentInfo(info))
|
|
s.Require().NoError(err)
|
|
}
|
|
|
|
tests := []struct {
|
|
description string
|
|
inSeg int64
|
|
|
|
expectedError error
|
|
}{
|
|
{"segment not exist", 1, merr.ErrSegmentNotFound},
|
|
{"segment dropped", 2, nil},
|
|
{"segment not in meta", 3, merr.ErrSegmentNotFound},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
s.Run(test.description, func() {
|
|
ctx := context.Background()
|
|
resp, err := s.testServer.SaveBinlogPaths(ctx, &datapb.SaveBinlogPathsRequest{
|
|
Base: &commonpb.MsgBase{
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
},
|
|
SegmentID: test.inSeg,
|
|
Channel: "ch1",
|
|
})
|
|
s.NoError(err)
|
|
s.ErrorIs(merr.Error(resp), test.expectedError)
|
|
})
|
|
}
|
|
}
|
|
|
|
func (s *ServerSuite) TestSaveBinlogPath_SaveDroppedSegment() {
|
|
s.mockChMgr.EXPECT().Match(int64(0), "ch1").Return(true)
|
|
s.testServer.meta.AddCollection(&collectionInfo{ID: 0})
|
|
|
|
segments := map[int64]commonpb.SegmentState{
|
|
0: commonpb.SegmentState_Flushed,
|
|
1: commonpb.SegmentState_Sealed,
|
|
}
|
|
for segID, state := range segments {
|
|
info := &datapb.SegmentInfo{
|
|
ID: segID,
|
|
InsertChannel: "ch1",
|
|
State: state,
|
|
Level: datapb.SegmentLevel_L1,
|
|
}
|
|
err := s.testServer.meta.AddSegment(context.TODO(), NewSegmentInfo(info))
|
|
s.Require().NoError(err)
|
|
}
|
|
|
|
tests := []struct {
|
|
description string
|
|
inSegID int64
|
|
inDropped bool
|
|
inFlushed bool
|
|
|
|
expectedState commonpb.SegmentState
|
|
}{
|
|
{"segID=0, flushed to dropped", 0, true, false, commonpb.SegmentState_Dropped},
|
|
{"segID=1, sealed to flushing", 1, false, true, commonpb.SegmentState_Flushing},
|
|
}
|
|
|
|
paramtable.Get().Save(paramtable.Get().DataCoordCfg.EnableAutoCompaction.Key, "False")
|
|
defer paramtable.Get().Reset(paramtable.Get().DataCoordCfg.EnableAutoCompaction.Key)
|
|
for _, test := range tests {
|
|
s.Run(test.description, func() {
|
|
ctx := context.Background()
|
|
resp, err := s.testServer.SaveBinlogPaths(ctx, &datapb.SaveBinlogPathsRequest{
|
|
Base: &commonpb.MsgBase{
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
},
|
|
SegmentID: test.inSegID,
|
|
Channel: "ch1",
|
|
Flushed: test.inFlushed,
|
|
Dropped: test.inDropped,
|
|
})
|
|
s.NoError(err)
|
|
s.EqualValues(resp.ErrorCode, commonpb.ErrorCode_Success)
|
|
|
|
segment := s.testServer.meta.GetSegment(test.inSegID)
|
|
s.NotNil(segment)
|
|
s.EqualValues(0, len(segment.GetBinlogs()))
|
|
s.EqualValues(segment.NumOfRows, 0)
|
|
|
|
flushing := []commonpb.SegmentState{commonpb.SegmentState_Flushed, commonpb.SegmentState_Flushing}
|
|
if lo.Contains(flushing, test.expectedState) {
|
|
s.True(lo.Contains(flushing, segment.GetState()))
|
|
} else {
|
|
s.Equal(test.expectedState, segment.GetState())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func (s *ServerSuite) TestSaveBinlogPath_L0Segment() {
|
|
s.mockChMgr.EXPECT().Match(int64(0), "ch1").Return(true)
|
|
s.testServer.meta.AddCollection(&collectionInfo{ID: 0})
|
|
|
|
segment := s.testServer.meta.GetHealthySegment(1)
|
|
s.Require().Nil(segment)
|
|
ctx := context.Background()
|
|
resp, err := s.testServer.SaveBinlogPaths(ctx, &datapb.SaveBinlogPathsRequest{
|
|
Base: &commonpb.MsgBase{
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
},
|
|
SegmentID: 1,
|
|
PartitionID: 1,
|
|
CollectionID: 0,
|
|
SegLevel: datapb.SegmentLevel_L0,
|
|
Channel: "ch1",
|
|
Deltalogs: []*datapb.FieldBinlog{
|
|
{
|
|
FieldID: 1,
|
|
Binlogs: []*datapb.Binlog{
|
|
{
|
|
LogPath: "/by-dev/test/0/1/1/1/1",
|
|
EntriesNum: 5,
|
|
},
|
|
{
|
|
LogPath: "/by-dev/test/0/1/1/1/2",
|
|
EntriesNum: 5,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
CheckPoints: []*datapb.CheckPoint{
|
|
{
|
|
SegmentID: 1,
|
|
Position: &msgpb.MsgPosition{
|
|
ChannelName: "ch1",
|
|
MsgID: []byte{1, 2, 3},
|
|
MsgGroup: "",
|
|
Timestamp: 0,
|
|
},
|
|
NumOfRows: 12,
|
|
},
|
|
},
|
|
Flushed: true,
|
|
})
|
|
s.NoError(err)
|
|
s.EqualValues(resp.ErrorCode, commonpb.ErrorCode_Success)
|
|
|
|
segment = s.testServer.meta.GetHealthySegment(1)
|
|
s.NotNil(segment)
|
|
s.EqualValues(datapb.SegmentLevel_L0, segment.GetLevel())
|
|
}
|
|
|
|
func (s *ServerSuite) TestSaveBinlogPath_NormalCase() {
|
|
s.mockChMgr.EXPECT().Match(int64(0), "ch1").Return(true)
|
|
s.testServer.meta.AddCollection(&collectionInfo{ID: 0})
|
|
|
|
segments := map[int64]int64{
|
|
0: 0,
|
|
1: 0,
|
|
}
|
|
for segID, collID := range segments {
|
|
info := &datapb.SegmentInfo{
|
|
ID: segID,
|
|
CollectionID: collID,
|
|
InsertChannel: "ch1",
|
|
State: commonpb.SegmentState_Growing,
|
|
}
|
|
err := s.testServer.meta.AddSegment(context.TODO(), NewSegmentInfo(info))
|
|
s.Require().NoError(err)
|
|
}
|
|
|
|
ctx := context.Background()
|
|
|
|
resp, err := s.testServer.SaveBinlogPaths(ctx, &datapb.SaveBinlogPathsRequest{
|
|
Base: &commonpb.MsgBase{
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
},
|
|
SegmentID: 1,
|
|
CollectionID: 0,
|
|
Channel: "ch1",
|
|
Field2BinlogPaths: []*datapb.FieldBinlog{
|
|
{
|
|
FieldID: 1,
|
|
Binlogs: []*datapb.Binlog{
|
|
{
|
|
LogPath: "/by-dev/test/0/1/1/1/1",
|
|
EntriesNum: 5,
|
|
},
|
|
{
|
|
LogPath: "/by-dev/test/0/1/1/1/2",
|
|
EntriesNum: 5,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
Field2StatslogPaths: []*datapb.FieldBinlog{
|
|
{
|
|
FieldID: 1,
|
|
Binlogs: []*datapb.Binlog{
|
|
{
|
|
LogPath: "/by-dev/test_stats/0/1/1/1/1",
|
|
EntriesNum: 5,
|
|
},
|
|
{
|
|
LogPath: "/by-dev/test_stats/0/1/1/1/2",
|
|
EntriesNum: 5,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
CheckPoints: []*datapb.CheckPoint{
|
|
{
|
|
SegmentID: 1,
|
|
Position: &msgpb.MsgPosition{
|
|
ChannelName: "ch1",
|
|
MsgID: []byte{1, 2, 3},
|
|
MsgGroup: "",
|
|
Timestamp: 0,
|
|
},
|
|
NumOfRows: 12,
|
|
},
|
|
},
|
|
Flushed: false,
|
|
})
|
|
s.NoError(err)
|
|
s.EqualValues(resp.ErrorCode, commonpb.ErrorCode_Success)
|
|
|
|
segment := s.testServer.meta.GetHealthySegment(1)
|
|
s.NotNil(segment)
|
|
binlogs := segment.GetBinlogs()
|
|
s.EqualValues(1, len(binlogs))
|
|
fieldBinlogs := binlogs[0]
|
|
s.NotNil(fieldBinlogs)
|
|
s.EqualValues(2, len(fieldBinlogs.GetBinlogs()))
|
|
s.EqualValues(1, fieldBinlogs.GetFieldID())
|
|
s.EqualValues("", fieldBinlogs.GetBinlogs()[0].GetLogPath())
|
|
s.EqualValues(int64(1), fieldBinlogs.GetBinlogs()[0].GetLogID())
|
|
s.EqualValues("", fieldBinlogs.GetBinlogs()[1].GetLogPath())
|
|
s.EqualValues(int64(2), fieldBinlogs.GetBinlogs()[1].GetLogID())
|
|
|
|
s.EqualValues(segment.DmlPosition.ChannelName, "ch1")
|
|
s.EqualValues(segment.DmlPosition.MsgID, []byte{1, 2, 3})
|
|
s.EqualValues(segment.NumOfRows, 10)
|
|
}
|
|
|
|
func (s *ServerSuite) TestFlush_NormalCase() {
|
|
req := &datapb.FlushRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Flush,
|
|
MsgID: 0,
|
|
Timestamp: 0,
|
|
SourceID: 0,
|
|
},
|
|
DbID: 0,
|
|
CollectionID: 0,
|
|
}
|
|
|
|
s.mockChMgr.EXPECT().GetNodeChannelsByCollectionID(mock.Anything).Return(map[int64][]string{
|
|
1: {"channel-1"},
|
|
})
|
|
|
|
mockCluster := NewMockCluster(s.T())
|
|
mockCluster.EXPECT().FlushChannels(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
|
|
Return(nil)
|
|
mockCluster.EXPECT().Close().Maybe()
|
|
s.testServer.cluster = mockCluster
|
|
|
|
schema := newTestSchema()
|
|
s.testServer.meta.AddCollection(&collectionInfo{ID: 0, Schema: schema, Partitions: []int64{}})
|
|
allocations, err := s.testServer.segmentManager.AllocSegment(context.TODO(), 0, 1, "channel-1", 1)
|
|
s.NoError(err)
|
|
s.EqualValues(1, len(allocations))
|
|
expireTs := allocations[0].ExpireTime
|
|
segID := allocations[0].SegmentID
|
|
|
|
resp, err := s.testServer.Flush(context.TODO(), req)
|
|
s.NoError(err)
|
|
s.EqualValues(commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
|
|
s.testServer.meta.SetCurrentRows(segID, 1)
|
|
ids, err := s.testServer.segmentManager.GetFlushableSegments(context.TODO(), "channel-1", expireTs)
|
|
s.NoError(err)
|
|
s.EqualValues(1, len(ids))
|
|
s.EqualValues(segID, ids[0])
|
|
}
|
|
|
|
func (s *ServerSuite) TestFlush_CollectionNotExist() {
|
|
req := &datapb.FlushRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Flush,
|
|
MsgID: 0,
|
|
Timestamp: 0,
|
|
SourceID: 0,
|
|
},
|
|
DbID: 0,
|
|
CollectionID: 0,
|
|
}
|
|
|
|
resp, err := s.testServer.Flush(context.TODO(), req)
|
|
s.NoError(err)
|
|
s.EqualValues(commonpb.ErrorCode_CollectionNotExists, resp.GetStatus().GetErrorCode())
|
|
|
|
mockHandler := NewNMockHandler(s.T())
|
|
mockHandler.EXPECT().GetCollection(mock.Anything, mock.Anything).
|
|
Return(nil, errors.New("mock error"))
|
|
s.testServer.handler = mockHandler
|
|
|
|
resp2, err2 := s.testServer.Flush(context.TODO(), req)
|
|
s.NoError(err2)
|
|
s.EqualValues(commonpb.ErrorCode_UnexpectedError, resp2.GetStatus().GetErrorCode())
|
|
}
|
|
|
|
func (s *ServerSuite) TestFlush_ClosedServer() {
|
|
s.TearDownTest()
|
|
req := &datapb.FlushRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Flush,
|
|
MsgID: 0,
|
|
Timestamp: 0,
|
|
SourceID: 0,
|
|
},
|
|
DbID: 0,
|
|
CollectionID: 0,
|
|
}
|
|
resp, err := s.testServer.Flush(context.Background(), req)
|
|
s.NoError(err)
|
|
s.ErrorIs(merr.Error(resp.GetStatus()), merr.ErrServiceNotReady)
|
|
}
|
|
|
|
func (s *ServerSuite) TestFlush_RollingUpgrade() {
|
|
req := &datapb.FlushRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Flush,
|
|
MsgID: 0,
|
|
Timestamp: 0,
|
|
SourceID: 0,
|
|
},
|
|
DbID: 0,
|
|
CollectionID: 0,
|
|
}
|
|
mockCluster := NewMockCluster(s.T())
|
|
mockCluster.EXPECT().FlushChannels(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
|
|
Return(merr.WrapErrServiceUnimplemented(grpcStatus.Error(codes.Unimplemented, "mock grpc unimplemented error")))
|
|
mockCluster.EXPECT().Close().Maybe()
|
|
s.testServer.cluster = mockCluster
|
|
s.testServer.meta.AddCollection(&collectionInfo{ID: 0})
|
|
s.mockChMgr.EXPECT().GetNodeChannelsByCollectionID(mock.Anything).Return(map[int64][]string{
|
|
1: {"channel-1"},
|
|
}).Once()
|
|
|
|
resp, err := s.testServer.Flush(context.TODO(), req)
|
|
s.NoError(err)
|
|
s.EqualValues(commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
s.EqualValues(0, resp.GetFlushTs())
|
|
}
|
|
|
|
func (s *ServerSuite) TestGetSegmentInfoChannel() {
|
|
resp, err := s.testServer.GetSegmentInfoChannel(context.TODO(), nil)
|
|
s.NoError(err)
|
|
s.EqualValues(commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
s.EqualValues(Params.CommonCfg.DataCoordSegmentInfo.GetValue(), resp.Value)
|
|
}
|
|
|
|
func (s *ServerSuite) TestGetSegmentInfo() {
|
|
testSegmentID := int64(1)
|
|
s.testServer.meta.AddSegment(context.TODO(), &SegmentInfo{
|
|
SegmentInfo: &datapb.SegmentInfo{
|
|
ID: 1,
|
|
Deltalogs: []*datapb.FieldBinlog{{FieldID: 100, Binlogs: []*datapb.Binlog{{LogID: 100}}}},
|
|
},
|
|
})
|
|
|
|
s.testServer.meta.AddSegment(context.TODO(), &SegmentInfo{
|
|
SegmentInfo: &datapb.SegmentInfo{
|
|
ID: 2,
|
|
Deltalogs: []*datapb.FieldBinlog{{FieldID: 100, Binlogs: []*datapb.Binlog{{LogID: 101}}}},
|
|
CompactionFrom: []int64{1},
|
|
},
|
|
})
|
|
|
|
resp, err := s.testServer.GetSegmentInfo(context.TODO(), &datapb.GetSegmentInfoRequest{
|
|
SegmentIDs: []int64{testSegmentID},
|
|
IncludeUnHealthy: true,
|
|
})
|
|
s.NoError(err)
|
|
s.EqualValues(2, len(resp.Infos[0].Deltalogs))
|
|
}
|
|
|
|
func (s *ServerSuite) TestAssignSegmentID() {
|
|
s.TearDownTest()
|
|
const collID = 100
|
|
const collIDInvalid = 101
|
|
const partID = 0
|
|
const channel0 = "channel0"
|
|
|
|
s.Run("assign segment normally", func() {
|
|
s.SetupTest()
|
|
defer s.TearDownTest()
|
|
|
|
schema := newTestSchema()
|
|
s.testServer.meta.AddCollection(&collectionInfo{
|
|
ID: collID,
|
|
Schema: schema,
|
|
Partitions: []int64{},
|
|
})
|
|
req := &datapb.SegmentIDRequest{
|
|
Count: 1000,
|
|
ChannelName: channel0,
|
|
CollectionID: collID,
|
|
PartitionID: partID,
|
|
}
|
|
|
|
resp, err := s.testServer.AssignSegmentID(context.TODO(), &datapb.AssignSegmentIDRequest{
|
|
NodeID: 0,
|
|
PeerRole: "",
|
|
SegmentIDRequests: []*datapb.SegmentIDRequest{req},
|
|
})
|
|
s.NoError(err)
|
|
s.EqualValues(1, len(resp.SegIDAssignments))
|
|
assign := resp.SegIDAssignments[0]
|
|
s.EqualValues(commonpb.ErrorCode_Success, assign.GetStatus().GetErrorCode())
|
|
s.EqualValues(collID, assign.CollectionID)
|
|
s.EqualValues(partID, assign.PartitionID)
|
|
s.EqualValues(channel0, assign.ChannelName)
|
|
s.EqualValues(1000, assign.Count)
|
|
})
|
|
|
|
s.Run("with closed server", func() {
|
|
s.SetupTest()
|
|
s.TearDownTest()
|
|
|
|
req := &datapb.SegmentIDRequest{
|
|
Count: 100,
|
|
ChannelName: channel0,
|
|
CollectionID: collID,
|
|
PartitionID: partID,
|
|
}
|
|
resp, err := s.testServer.AssignSegmentID(context.Background(), &datapb.AssignSegmentIDRequest{
|
|
NodeID: 0,
|
|
PeerRole: "",
|
|
SegmentIDRequests: []*datapb.SegmentIDRequest{req},
|
|
})
|
|
s.NoError(err)
|
|
s.ErrorIs(merr.Error(resp.GetStatus()), merr.ErrServiceNotReady)
|
|
})
|
|
|
|
s.Run("assign segment with invalid collection", func() {
|
|
s.SetupTest()
|
|
defer s.TearDownTest()
|
|
|
|
s.testServer.rootCoordClient = &mockRootCoord{
|
|
RootCoordClient: s.testServer.rootCoordClient,
|
|
collID: collID,
|
|
}
|
|
|
|
schema := newTestSchema()
|
|
s.testServer.meta.AddCollection(&collectionInfo{
|
|
ID: collID,
|
|
Schema: schema,
|
|
Partitions: []int64{},
|
|
})
|
|
req := &datapb.SegmentIDRequest{
|
|
Count: 1000,
|
|
ChannelName: channel0,
|
|
CollectionID: collIDInvalid,
|
|
PartitionID: partID,
|
|
}
|
|
|
|
resp, err := s.testServer.AssignSegmentID(context.TODO(), &datapb.AssignSegmentIDRequest{
|
|
NodeID: 0,
|
|
PeerRole: "",
|
|
SegmentIDRequests: []*datapb.SegmentIDRequest{req},
|
|
})
|
|
s.NoError(err)
|
|
s.EqualValues(1, len(resp.SegIDAssignments))
|
|
})
|
|
}
|
|
|
|
func TestBroadcastAlteredCollection(t *testing.T) {
|
|
t.Run("test server is closed", func(t *testing.T) {
|
|
s := &Server{}
|
|
s.stateCode.Store(commonpb.StateCode_Initializing)
|
|
ctx := context.Background()
|
|
resp, err := s.BroadcastAlteredCollection(ctx, nil)
|
|
assert.NotNil(t, resp.Reason)
|
|
assert.NoError(t, err)
|
|
})
|
|
|
|
t.Run("test meta non exist", func(t *testing.T) {
|
|
s := &Server{meta: &meta{collections: make(map[UniqueID]*collectionInfo, 1)}}
|
|
s.stateCode.Store(commonpb.StateCode_Healthy)
|
|
ctx := context.Background()
|
|
req := &datapb.AlterCollectionRequest{
|
|
CollectionID: 1,
|
|
PartitionIDs: []int64{1},
|
|
Properties: []*commonpb.KeyValuePair{{Key: "k", Value: "v"}},
|
|
}
|
|
resp, err := s.BroadcastAlteredCollection(ctx, req)
|
|
assert.NotNil(t, resp)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 1, len(s.meta.collections))
|
|
})
|
|
|
|
t.Run("test update meta", func(t *testing.T) {
|
|
s := &Server{meta: &meta{collections: map[UniqueID]*collectionInfo{
|
|
1: {ID: 1},
|
|
}}}
|
|
s.stateCode.Store(commonpb.StateCode_Healthy)
|
|
ctx := context.Background()
|
|
req := &datapb.AlterCollectionRequest{
|
|
CollectionID: 1,
|
|
PartitionIDs: []int64{1},
|
|
Properties: []*commonpb.KeyValuePair{{Key: "k", Value: "v"}},
|
|
}
|
|
|
|
assert.Nil(t, s.meta.collections[1].Properties)
|
|
resp, err := s.BroadcastAlteredCollection(ctx, req)
|
|
assert.NotNil(t, resp)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, s.meta.collections[1].Properties)
|
|
})
|
|
}
|
|
|
|
func TestServer_GcConfirm(t *testing.T) {
|
|
t.Run("closed server", func(t *testing.T) {
|
|
s := &Server{}
|
|
s.stateCode.Store(commonpb.StateCode_Initializing)
|
|
resp, err := s.GcConfirm(context.TODO(), &datapb.GcConfirmRequest{CollectionId: 100, PartitionId: 10000})
|
|
assert.NoError(t, err)
|
|
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
})
|
|
|
|
t.Run("normal case", func(t *testing.T) {
|
|
s := &Server{}
|
|
s.stateCode.Store(commonpb.StateCode_Healthy)
|
|
|
|
m := &meta{}
|
|
catalog := mocks.NewDataCoordCatalog(t)
|
|
m.catalog = catalog
|
|
|
|
catalog.On("GcConfirm",
|
|
mock.Anything,
|
|
mock.AnythingOfType("int64"),
|
|
mock.AnythingOfType("int64")).
|
|
Return(false)
|
|
|
|
s.meta = m
|
|
|
|
resp, err := s.GcConfirm(context.TODO(), &datapb.GcConfirmRequest{CollectionId: 100, PartitionId: 10000})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
assert.False(t, resp.GetGcFinished())
|
|
})
|
|
}
|
|
|
|
func TestGetRecoveryInfoV2(t *testing.T) {
|
|
t.Run("test get recovery info with no segments", func(t *testing.T) {
|
|
svr := newTestServer(t)
|
|
defer closeTestServer(t, svr)
|
|
|
|
svr.rootCoordClientCreator = func(ctx context.Context) (types.RootCoordClient, error) {
|
|
return newMockRootCoordClient(), nil
|
|
}
|
|
|
|
req := &datapb.GetRecoveryInfoRequestV2{
|
|
CollectionID: 0,
|
|
}
|
|
resp, err := svr.GetRecoveryInfoV2(context.TODO(), req)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
assert.EqualValues(t, 0, len(resp.GetSegments()))
|
|
assert.EqualValues(t, 0, len(resp.GetChannels()))
|
|
})
|
|
|
|
createSegment := func(id, collectionID, partitionID, numOfRows int64, posTs uint64,
|
|
channel string, state commonpb.SegmentState,
|
|
) *datapb.SegmentInfo {
|
|
return &datapb.SegmentInfo{
|
|
ID: id,
|
|
CollectionID: collectionID,
|
|
PartitionID: partitionID,
|
|
InsertChannel: channel,
|
|
NumOfRows: numOfRows,
|
|
State: state,
|
|
DmlPosition: &msgpb.MsgPosition{
|
|
ChannelName: channel,
|
|
MsgID: []byte{},
|
|
Timestamp: posTs,
|
|
},
|
|
StartPosition: &msgpb.MsgPosition{
|
|
ChannelName: "",
|
|
MsgID: []byte{},
|
|
MsgGroup: "",
|
|
Timestamp: 0,
|
|
},
|
|
}
|
|
}
|
|
|
|
t.Run("test get earliest position of flushed segments as seek position", func(t *testing.T) {
|
|
svr := newTestServer(t)
|
|
defer closeTestServer(t, svr)
|
|
|
|
svr.rootCoordClientCreator = func(ctx context.Context) (types.RootCoordClient, error) {
|
|
return newMockRootCoordClient(), nil
|
|
}
|
|
|
|
svr.meta.AddCollection(&collectionInfo{
|
|
Schema: newTestSchema(),
|
|
})
|
|
|
|
err := svr.meta.UpdateChannelCheckpoint("vchan1", &msgpb.MsgPosition{
|
|
ChannelName: "vchan1",
|
|
Timestamp: 10,
|
|
MsgID: []byte{0, 0, 0, 0, 0, 0, 0, 0},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
err = svr.meta.indexMeta.CreateIndex(&model.Index{
|
|
TenantID: "",
|
|
CollectionID: 0,
|
|
FieldID: 2,
|
|
IndexID: 0,
|
|
IndexName: "",
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
seg1 := createSegment(0, 0, 0, 100, 10, "vchan1", commonpb.SegmentState_Flushed)
|
|
seg1.Binlogs = []*datapb.FieldBinlog{
|
|
{
|
|
FieldID: 1,
|
|
Binlogs: []*datapb.Binlog{
|
|
{
|
|
EntriesNum: 20,
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 0, 0, 1, 901),
|
|
},
|
|
{
|
|
EntriesNum: 20,
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 0, 0, 1, 902),
|
|
},
|
|
{
|
|
EntriesNum: 20,
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 0, 0, 1, 903),
|
|
},
|
|
},
|
|
},
|
|
}
|
|
seg2 := createSegment(1, 0, 0, 100, 20, "vchan1", commonpb.SegmentState_Flushed)
|
|
seg2.Binlogs = []*datapb.FieldBinlog{
|
|
{
|
|
FieldID: 1,
|
|
Binlogs: []*datapb.Binlog{
|
|
{
|
|
EntriesNum: 30,
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 0, 1, 1, 801),
|
|
},
|
|
{
|
|
EntriesNum: 70,
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 0, 1, 1, 802),
|
|
},
|
|
},
|
|
},
|
|
}
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
|
|
assert.NoError(t, err)
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
|
|
assert.NoError(t, err)
|
|
err = svr.meta.indexMeta.AddSegmentIndex(&model.SegmentIndex{
|
|
SegmentID: seg1.ID,
|
|
BuildID: seg1.ID,
|
|
})
|
|
assert.NoError(t, err)
|
|
err = svr.meta.indexMeta.FinishTask(&indexpb.IndexTaskInfo{
|
|
BuildID: seg1.ID,
|
|
State: commonpb.IndexState_Finished,
|
|
})
|
|
assert.NoError(t, err)
|
|
err = svr.meta.indexMeta.AddSegmentIndex(&model.SegmentIndex{
|
|
SegmentID: seg2.ID,
|
|
BuildID: seg2.ID,
|
|
})
|
|
assert.NoError(t, err)
|
|
err = svr.meta.indexMeta.FinishTask(&indexpb.IndexTaskInfo{
|
|
BuildID: seg2.ID,
|
|
State: commonpb.IndexState_Finished,
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
ch := &channelMeta{Name: "vchan1", CollectionID: 0}
|
|
svr.channelManager.AddNode(0)
|
|
svr.channelManager.Watch(context.Background(), ch)
|
|
|
|
req := &datapb.GetRecoveryInfoRequestV2{
|
|
CollectionID: 0,
|
|
}
|
|
resp, err := svr.GetRecoveryInfoV2(context.TODO(), req)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
assert.EqualValues(t, 1, len(resp.GetChannels()))
|
|
assert.EqualValues(t, 0, len(resp.GetChannels()[0].GetUnflushedSegmentIds()))
|
|
assert.ElementsMatch(t, []int64{0, 1}, resp.GetChannels()[0].GetFlushedSegmentIds())
|
|
assert.EqualValues(t, 10, resp.GetChannels()[0].GetSeekPosition().GetTimestamp())
|
|
assert.EqualValues(t, 2, len(resp.GetSegments()))
|
|
// Row count corrected from 100 + 100 -> 100 + 60.
|
|
assert.EqualValues(t, 160, resp.GetSegments()[0].GetNumOfRows()+resp.GetSegments()[1].GetNumOfRows())
|
|
})
|
|
|
|
t.Run("test get recovery of unflushed segments ", func(t *testing.T) {
|
|
svr := newTestServer(t)
|
|
defer closeTestServer(t, svr)
|
|
|
|
svr.rootCoordClientCreator = func(ctx context.Context) (types.RootCoordClient, error) {
|
|
return newMockRootCoordClient(), nil
|
|
}
|
|
|
|
svr.meta.AddCollection(&collectionInfo{
|
|
ID: 0,
|
|
Schema: newTestSchema(),
|
|
})
|
|
|
|
err := svr.meta.UpdateChannelCheckpoint("vchan1", &msgpb.MsgPosition{
|
|
ChannelName: "vchan1",
|
|
Timestamp: 0,
|
|
MsgID: []byte{0, 0, 0, 0, 0, 0, 0, 0},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
seg1 := createSegment(3, 0, 0, 100, 30, "vchan1", commonpb.SegmentState_Growing)
|
|
seg1.Binlogs = []*datapb.FieldBinlog{
|
|
{
|
|
FieldID: 1,
|
|
Binlogs: []*datapb.Binlog{
|
|
{
|
|
EntriesNum: 20,
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 0, 3, 1, 901),
|
|
},
|
|
{
|
|
EntriesNum: 20,
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 0, 3, 1, 902),
|
|
},
|
|
{
|
|
EntriesNum: 20,
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 0, 3, 1, 903),
|
|
},
|
|
},
|
|
},
|
|
}
|
|
seg2 := createSegment(4, 0, 0, 100, 40, "vchan1", commonpb.SegmentState_Growing)
|
|
seg2.Binlogs = []*datapb.FieldBinlog{
|
|
{
|
|
FieldID: 1,
|
|
Binlogs: []*datapb.Binlog{
|
|
{
|
|
EntriesNum: 30,
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 0, 4, 1, 801),
|
|
},
|
|
{
|
|
EntriesNum: 70,
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 0, 4, 1, 802),
|
|
},
|
|
},
|
|
},
|
|
}
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
|
|
assert.NoError(t, err)
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
|
|
assert.NoError(t, err)
|
|
|
|
ch := &channelMeta{Name: "vchan1", CollectionID: 0}
|
|
svr.channelManager.AddNode(0)
|
|
svr.channelManager.Watch(context.Background(), ch)
|
|
|
|
req := &datapb.GetRecoveryInfoRequestV2{
|
|
CollectionID: 0,
|
|
}
|
|
resp, err := svr.GetRecoveryInfoV2(context.TODO(), req)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
assert.EqualValues(t, 0, len(resp.GetSegments()))
|
|
assert.EqualValues(t, 1, len(resp.GetChannels()))
|
|
assert.NotNil(t, resp.GetChannels()[0].SeekPosition)
|
|
assert.NotEqual(t, 0, resp.GetChannels()[0].GetSeekPosition().GetTimestamp())
|
|
})
|
|
|
|
t.Run("test get binlogs", func(t *testing.T) {
|
|
svr := newTestServer(t)
|
|
defer closeTestServer(t, svr)
|
|
|
|
svr.meta.AddCollection(&collectionInfo{
|
|
Schema: newTestSchema(),
|
|
})
|
|
|
|
svr.rootCoordClientCreator = func(ctx context.Context) (types.RootCoordClient, error) {
|
|
return newMockRootCoordClient(), nil
|
|
}
|
|
|
|
binlogReq := &datapb.SaveBinlogPathsRequest{
|
|
SegmentID: 0,
|
|
CollectionID: 0,
|
|
Field2BinlogPaths: []*datapb.FieldBinlog{
|
|
{
|
|
FieldID: 1,
|
|
Binlogs: []*datapb.Binlog{
|
|
{
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 100, 0, 1, 801),
|
|
},
|
|
{
|
|
LogPath: metautil.BuildInsertLogPath("a", 0, 100, 0, 1, 801),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
Field2StatslogPaths: []*datapb.FieldBinlog{
|
|
{
|
|
FieldID: 1,
|
|
Binlogs: []*datapb.Binlog{
|
|
{
|
|
LogPath: metautil.BuildStatsLogPath("a", 0, 100, 0, 1000, 10000),
|
|
},
|
|
{
|
|
LogPath: metautil.BuildStatsLogPath("a", 0, 100, 0, 1000, 10000),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
Deltalogs: []*datapb.FieldBinlog{
|
|
{
|
|
Binlogs: []*datapb.Binlog{
|
|
{
|
|
TimestampFrom: 0,
|
|
TimestampTo: 1,
|
|
LogPath: metautil.BuildDeltaLogPath("a", 0, 100, 0, 100000),
|
|
LogSize: 1,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
segment := createSegment(0, 0, 1, 100, 10, "vchan1", commonpb.SegmentState_Flushed)
|
|
err := svr.meta.AddSegment(context.TODO(), NewSegmentInfo(segment))
|
|
assert.NoError(t, err)
|
|
|
|
err = svr.meta.indexMeta.CreateIndex(&model.Index{
|
|
TenantID: "",
|
|
CollectionID: 0,
|
|
FieldID: 2,
|
|
IndexID: 0,
|
|
IndexName: "",
|
|
})
|
|
assert.NoError(t, err)
|
|
err = svr.meta.indexMeta.AddSegmentIndex(&model.SegmentIndex{
|
|
SegmentID: segment.ID,
|
|
BuildID: segment.ID,
|
|
})
|
|
assert.NoError(t, err)
|
|
err = svr.meta.indexMeta.FinishTask(&indexpb.IndexTaskInfo{
|
|
BuildID: segment.ID,
|
|
State: commonpb.IndexState_Finished,
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
err = svr.channelManager.AddNode(0)
|
|
assert.NoError(t, err)
|
|
err = svr.channelManager.Watch(context.Background(), &channelMeta{Name: "vchan1", CollectionID: 0})
|
|
assert.NoError(t, err)
|
|
|
|
sResp, err := svr.SaveBinlogPaths(context.TODO(), binlogReq)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, commonpb.ErrorCode_Success, sResp.ErrorCode)
|
|
|
|
req := &datapb.GetRecoveryInfoRequestV2{
|
|
CollectionID: 0,
|
|
PartitionIDs: []int64{1},
|
|
}
|
|
resp, err := svr.GetRecoveryInfoV2(context.TODO(), req)
|
|
assert.NoError(t, err)
|
|
assert.NoError(t, merr.Error(resp.Status))
|
|
assert.EqualValues(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
assert.EqualValues(t, 1, len(resp.GetSegments()))
|
|
assert.EqualValues(t, 0, resp.GetSegments()[0].GetID())
|
|
assert.EqualValues(t, 0, len(resp.GetSegments()[0].GetBinlogs()))
|
|
})
|
|
t.Run("with dropped segments", func(t *testing.T) {
|
|
svr := newTestServer(t)
|
|
defer closeTestServer(t, svr)
|
|
|
|
svr.rootCoordClientCreator = func(ctx context.Context) (types.RootCoordClient, error) {
|
|
return newMockRootCoordClient(), nil
|
|
}
|
|
|
|
svr.meta.AddCollection(&collectionInfo{
|
|
ID: 0,
|
|
Schema: newTestSchema(),
|
|
})
|
|
|
|
err := svr.meta.UpdateChannelCheckpoint("vchan1", &msgpb.MsgPosition{
|
|
ChannelName: "vchan1",
|
|
Timestamp: 0,
|
|
MsgID: []byte{0, 0, 0, 0, 0, 0, 0, 0},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
seg1 := createSegment(7, 0, 0, 100, 30, "vchan1", commonpb.SegmentState_Growing)
|
|
seg2 := createSegment(8, 0, 0, 100, 40, "vchan1", commonpb.SegmentState_Dropped)
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
|
|
assert.NoError(t, err)
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
|
|
assert.NoError(t, err)
|
|
|
|
ch := &channelMeta{Name: "vchan1", CollectionID: 0}
|
|
svr.channelManager.AddNode(0)
|
|
svr.channelManager.Watch(context.Background(), ch)
|
|
|
|
req := &datapb.GetRecoveryInfoRequestV2{
|
|
CollectionID: 0,
|
|
}
|
|
resp, err := svr.GetRecoveryInfoV2(context.TODO(), req)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
assert.EqualValues(t, 0, len(resp.GetSegments()))
|
|
assert.EqualValues(t, 1, len(resp.GetChannels()))
|
|
assert.NotNil(t, resp.GetChannels()[0].SeekPosition)
|
|
assert.NotEqual(t, 0, resp.GetChannels()[0].GetSeekPosition().GetTimestamp())
|
|
assert.Len(t, resp.GetChannels()[0].GetDroppedSegmentIds(), 1)
|
|
assert.Equal(t, UniqueID(8), resp.GetChannels()[0].GetDroppedSegmentIds()[0])
|
|
})
|
|
|
|
t.Run("with fake segments", func(t *testing.T) {
|
|
svr := newTestServer(t)
|
|
defer closeTestServer(t, svr)
|
|
|
|
svr.rootCoordClientCreator = func(ctx context.Context) (types.RootCoordClient, error) {
|
|
return newMockRootCoordClient(), nil
|
|
}
|
|
|
|
svr.meta.AddCollection(&collectionInfo{
|
|
ID: 0,
|
|
Schema: newTestSchema(),
|
|
})
|
|
|
|
err := svr.meta.UpdateChannelCheckpoint("vchan1", &msgpb.MsgPosition{
|
|
ChannelName: "vchan1",
|
|
Timestamp: 0,
|
|
MsgID: []byte{0, 0, 0, 0, 0, 0, 0, 0},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
seg1 := createSegment(7, 0, 0, 100, 30, "vchan1", commonpb.SegmentState_Growing)
|
|
seg2 := createSegment(8, 0, 0, 100, 40, "vchan1", commonpb.SegmentState_Flushed)
|
|
seg2.IsFake = true
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
|
|
assert.NoError(t, err)
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
|
|
assert.NoError(t, err)
|
|
|
|
ch := &channelMeta{Name: "vchan1", CollectionID: 0}
|
|
svr.channelManager.AddNode(0)
|
|
svr.channelManager.Watch(context.Background(), ch)
|
|
|
|
req := &datapb.GetRecoveryInfoRequestV2{
|
|
CollectionID: 0,
|
|
}
|
|
resp, err := svr.GetRecoveryInfoV2(context.TODO(), req)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
assert.EqualValues(t, 0, len(resp.GetSegments()))
|
|
assert.EqualValues(t, 1, len(resp.GetChannels()))
|
|
assert.NotNil(t, resp.GetChannels()[0].SeekPosition)
|
|
assert.NotEqual(t, 0, resp.GetChannels()[0].GetSeekPosition().GetTimestamp())
|
|
})
|
|
|
|
t.Run("with continuous compaction", func(t *testing.T) {
|
|
svr := newTestServer(t)
|
|
defer closeTestServer(t, svr)
|
|
|
|
svr.rootCoordClientCreator = func(ctx context.Context) (types.RootCoordClient, error) {
|
|
return newMockRootCoordClient(), nil
|
|
}
|
|
|
|
svr.meta.AddCollection(&collectionInfo{
|
|
ID: 0,
|
|
Schema: newTestSchema(),
|
|
})
|
|
|
|
err := svr.meta.UpdateChannelCheckpoint("vchan1", &msgpb.MsgPosition{
|
|
ChannelName: "vchan1",
|
|
Timestamp: 0,
|
|
MsgID: []byte{0, 0, 0, 0, 0, 0, 0, 0},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
seg1 := createSegment(9, 0, 0, 2048, 30, "vchan1", commonpb.SegmentState_Dropped)
|
|
seg2 := createSegment(10, 0, 0, 2048, 40, "vchan1", commonpb.SegmentState_Dropped)
|
|
seg3 := createSegment(11, 0, 0, 2048, 40, "vchan1", commonpb.SegmentState_Dropped)
|
|
seg3.CompactionFrom = []int64{9, 10}
|
|
seg4 := createSegment(12, 0, 0, 2048, 40, "vchan1", commonpb.SegmentState_Dropped)
|
|
seg5 := createSegment(13, 0, 0, 2048, 40, "vchan1", commonpb.SegmentState_Flushed)
|
|
seg5.CompactionFrom = []int64{11, 12}
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
|
|
assert.NoError(t, err)
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
|
|
assert.NoError(t, err)
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg3))
|
|
assert.NoError(t, err)
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg4))
|
|
assert.NoError(t, err)
|
|
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg5))
|
|
assert.NoError(t, err)
|
|
err = svr.meta.indexMeta.CreateIndex(&model.Index{
|
|
TenantID: "",
|
|
CollectionID: 0,
|
|
FieldID: 2,
|
|
IndexID: 0,
|
|
IndexName: "_default_idx_2",
|
|
IsDeleted: false,
|
|
CreateTime: 0,
|
|
TypeParams: nil,
|
|
IndexParams: nil,
|
|
IsAutoIndex: false,
|
|
UserIndexParams: nil,
|
|
})
|
|
assert.NoError(t, err)
|
|
svr.meta.indexMeta.updateSegmentIndex(&model.SegmentIndex{
|
|
SegmentID: seg4.ID,
|
|
CollectionID: 0,
|
|
PartitionID: 0,
|
|
NumRows: 100,
|
|
IndexID: 0,
|
|
BuildID: 0,
|
|
NodeID: 0,
|
|
IndexVersion: 1,
|
|
IndexState: commonpb.IndexState_Finished,
|
|
FailReason: "",
|
|
IsDeleted: false,
|
|
CreateTime: 0,
|
|
IndexFileKeys: nil,
|
|
IndexSize: 0,
|
|
})
|
|
|
|
ch := &channelMeta{Name: "vchan1", CollectionID: 0}
|
|
svr.channelManager.AddNode(0)
|
|
svr.channelManager.Watch(context.Background(), ch)
|
|
|
|
req := &datapb.GetRecoveryInfoRequestV2{
|
|
CollectionID: 0,
|
|
}
|
|
resp, err := svr.GetRecoveryInfoV2(context.TODO(), req)
|
|
assert.NoError(t, err)
|
|
assert.EqualValues(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
|
|
assert.NotNil(t, resp.GetChannels()[0].SeekPosition)
|
|
assert.NotEqual(t, 0, resp.GetChannels()[0].GetSeekPosition().GetTimestamp())
|
|
assert.Len(t, resp.GetChannels()[0].GetDroppedSegmentIds(), 0)
|
|
assert.ElementsMatch(t, []UniqueID{}, resp.GetChannels()[0].GetUnflushedSegmentIds())
|
|
assert.ElementsMatch(t, []UniqueID{9, 10, 12}, resp.GetChannels()[0].GetFlushedSegmentIds())
|
|
})
|
|
|
|
t.Run("with closed server", func(t *testing.T) {
|
|
svr := newTestServer(t)
|
|
closeTestServer(t, svr)
|
|
resp, err := svr.GetRecoveryInfoV2(context.TODO(), &datapb.GetRecoveryInfoRequestV2{})
|
|
assert.NoError(t, err)
|
|
err = merr.Error(resp.GetStatus())
|
|
assert.ErrorIs(t, err, merr.ErrServiceNotReady)
|
|
})
|
|
}
|
|
|
|
func TestImportV2(t *testing.T) {
|
|
ctx := context.Background()
|
|
mockErr := errors.New("mock err")
|
|
|
|
t.Run("ImportV2", func(t *testing.T) {
|
|
// server not healthy
|
|
s := &Server{}
|
|
s.stateCode.Store(commonpb.StateCode_Initializing)
|
|
resp, err := s.ImportV2(ctx, nil)
|
|
assert.NoError(t, err)
|
|
assert.NotEqual(t, int32(0), resp.GetStatus().GetCode())
|
|
s.stateCode.Store(commonpb.StateCode_Healthy)
|
|
|
|
// parse timeout failed
|
|
resp, err = s.ImportV2(ctx, &internalpb.ImportRequestInternal{
|
|
Options: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "timeout",
|
|
Value: "@$#$%#%$",
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.True(t, errors.Is(merr.Error(resp.GetStatus()), merr.ErrImportFailed))
|
|
|
|
// list binlog failed
|
|
cm := mocks2.NewChunkManager(t)
|
|
cm.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
|
|
s.meta = &meta{chunkManager: cm}
|
|
resp, err = s.ImportV2(ctx, &internalpb.ImportRequestInternal{
|
|
Files: []*internalpb.ImportFile{
|
|
{
|
|
Id: 1,
|
|
Paths: []string{"mock_insert_prefix"},
|
|
},
|
|
},
|
|
Options: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "backup",
|
|
Value: "true",
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.True(t, errors.Is(merr.Error(resp.GetStatus()), merr.ErrImportFailed))
|
|
|
|
// alloc failed
|
|
alloc := NewNMockAllocator(t)
|
|
alloc.EXPECT().allocN(mock.Anything).Return(0, 0, mockErr)
|
|
s.allocator = alloc
|
|
resp, err = s.ImportV2(ctx, &internalpb.ImportRequestInternal{})
|
|
assert.NoError(t, err)
|
|
assert.True(t, errors.Is(merr.Error(resp.GetStatus()), merr.ErrImportFailed))
|
|
alloc = NewNMockAllocator(t)
|
|
alloc.EXPECT().allocN(mock.Anything).Return(0, 0, nil)
|
|
s.allocator = alloc
|
|
|
|
// add job failed
|
|
catalog := mocks.NewDataCoordCatalog(t)
|
|
catalog.EXPECT().ListImportJobs().Return(nil, nil)
|
|
catalog.EXPECT().ListPreImportTasks().Return(nil, nil)
|
|
catalog.EXPECT().ListImportTasks().Return(nil, nil)
|
|
catalog.EXPECT().SaveImportJob(mock.Anything).Return(mockErr)
|
|
s.importMeta, err = NewImportMeta(catalog)
|
|
assert.NoError(t, err)
|
|
resp, err = s.ImportV2(ctx, &internalpb.ImportRequestInternal{
|
|
Files: []*internalpb.ImportFile{
|
|
{
|
|
Id: 1,
|
|
Paths: []string{"a.json"},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.True(t, errors.Is(merr.Error(resp.GetStatus()), merr.ErrImportFailed))
|
|
jobs := s.importMeta.GetJobBy()
|
|
assert.Equal(t, 0, len(jobs))
|
|
catalog.ExpectedCalls = lo.Filter(catalog.ExpectedCalls, func(call *mock.Call, _ int) bool {
|
|
return call.Method != "SaveImportJob"
|
|
})
|
|
catalog.EXPECT().SaveImportJob(mock.Anything).Return(nil)
|
|
|
|
// normal case
|
|
resp, err = s.ImportV2(ctx, &internalpb.ImportRequestInternal{
|
|
Files: []*internalpb.ImportFile{
|
|
{
|
|
Id: 1,
|
|
Paths: []string{"a.json"},
|
|
},
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, int32(0), resp.GetStatus().GetCode())
|
|
jobs = s.importMeta.GetJobBy()
|
|
assert.Equal(t, 1, len(jobs))
|
|
})
|
|
|
|
t.Run("GetImportProgress", func(t *testing.T) {
|
|
// server not healthy
|
|
s := &Server{}
|
|
s.stateCode.Store(commonpb.StateCode_Initializing)
|
|
resp, err := s.GetImportProgress(ctx, nil)
|
|
assert.NoError(t, err)
|
|
assert.NotEqual(t, int32(0), resp.GetStatus().GetCode())
|
|
s.stateCode.Store(commonpb.StateCode_Healthy)
|
|
|
|
// illegal jobID
|
|
resp, err = s.GetImportProgress(ctx, &internalpb.GetImportProgressRequest{
|
|
JobID: "@%$%$#%",
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.True(t, errors.Is(merr.Error(resp.GetStatus()), merr.ErrImportFailed))
|
|
|
|
// normal case
|
|
catalog := mocks.NewDataCoordCatalog(t)
|
|
catalog.EXPECT().ListImportJobs().Return(nil, nil)
|
|
catalog.EXPECT().ListPreImportTasks().Return(nil, nil)
|
|
catalog.EXPECT().ListImportTasks().Return(nil, nil)
|
|
catalog.EXPECT().SaveImportJob(mock.Anything).Return(nil)
|
|
s.importMeta, err = NewImportMeta(catalog)
|
|
assert.NoError(t, err)
|
|
var job ImportJob = &importJob{
|
|
ImportJob: &datapb.ImportJob{
|
|
JobID: 0,
|
|
Schema: &schemapb.CollectionSchema{},
|
|
State: internalpb.ImportJobState_Failed,
|
|
},
|
|
}
|
|
err = s.importMeta.AddJob(job)
|
|
assert.NoError(t, err)
|
|
resp, err = s.GetImportProgress(ctx, &internalpb.GetImportProgressRequest{
|
|
JobID: "0",
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, int32(0), resp.GetStatus().GetCode())
|
|
assert.Equal(t, int64(0), resp.GetProgress())
|
|
assert.Equal(t, internalpb.ImportJobState_Failed, resp.GetState())
|
|
})
|
|
|
|
t.Run("ListImports", func(t *testing.T) {
|
|
// server not healthy
|
|
s := &Server{}
|
|
s.stateCode.Store(commonpb.StateCode_Initializing)
|
|
resp, err := s.ListImports(ctx, nil)
|
|
assert.NoError(t, err)
|
|
assert.NotEqual(t, int32(0), resp.GetStatus().GetCode())
|
|
s.stateCode.Store(commonpb.StateCode_Healthy)
|
|
|
|
// normal case
|
|
catalog := mocks.NewDataCoordCatalog(t)
|
|
catalog.EXPECT().ListImportJobs().Return(nil, nil)
|
|
catalog.EXPECT().ListPreImportTasks().Return(nil, nil)
|
|
catalog.EXPECT().ListImportTasks().Return(nil, nil)
|
|
catalog.EXPECT().SaveImportJob(mock.Anything).Return(nil)
|
|
catalog.EXPECT().SavePreImportTask(mock.Anything).Return(nil)
|
|
s.importMeta, err = NewImportMeta(catalog)
|
|
assert.NoError(t, err)
|
|
var job ImportJob = &importJob{
|
|
ImportJob: &datapb.ImportJob{
|
|
JobID: 0,
|
|
CollectionID: 1,
|
|
Schema: &schemapb.CollectionSchema{},
|
|
},
|
|
}
|
|
err = s.importMeta.AddJob(job)
|
|
assert.NoError(t, err)
|
|
var task ImportTask = &preImportTask{
|
|
PreImportTask: &datapb.PreImportTask{
|
|
JobID: 0,
|
|
TaskID: 1,
|
|
State: datapb.ImportTaskStateV2_Failed,
|
|
},
|
|
}
|
|
err = s.importMeta.AddTask(task)
|
|
assert.NoError(t, err)
|
|
resp, err = s.ListImports(ctx, &internalpb.ListImportsRequestInternal{
|
|
CollectionID: 1,
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, int32(0), resp.GetStatus().GetCode())
|
|
assert.Equal(t, 1, len(resp.GetJobIDs()))
|
|
assert.Equal(t, 1, len(resp.GetStates()))
|
|
assert.Equal(t, 1, len(resp.GetReasons()))
|
|
assert.Equal(t, 1, len(resp.GetProgresses()))
|
|
})
|
|
}
|
|
|
|
type GcControlServiceSuite struct {
|
|
suite.Suite
|
|
|
|
server *Server
|
|
}
|
|
|
|
func (s *GcControlServiceSuite) SetupTest() {
|
|
s.server = newTestServer(s.T())
|
|
}
|
|
|
|
func (s *GcControlServiceSuite) TearDownTest() {
|
|
if s.server != nil {
|
|
closeTestServer(s.T(), s.server)
|
|
}
|
|
}
|
|
|
|
func (s *GcControlServiceSuite) TestClosedServer() {
|
|
closeTestServer(s.T(), s.server)
|
|
resp, err := s.server.GcControl(context.TODO(), &datapb.GcControlRequest{})
|
|
s.NoError(err)
|
|
s.False(merr.Ok(resp))
|
|
s.server = nil
|
|
}
|
|
|
|
func (s *GcControlServiceSuite) TestUnknownCmd() {
|
|
resp, err := s.server.GcControl(context.TODO(), &datapb.GcControlRequest{
|
|
Command: 0,
|
|
})
|
|
s.NoError(err)
|
|
s.False(merr.Ok(resp))
|
|
}
|
|
|
|
func (s *GcControlServiceSuite) TestPause() {
|
|
resp, err := s.server.GcControl(context.TODO(), &datapb.GcControlRequest{
|
|
Command: datapb.GcCommand_Pause,
|
|
})
|
|
s.Nil(err)
|
|
s.False(merr.Ok(resp))
|
|
|
|
resp, err = s.server.GcControl(context.TODO(), &datapb.GcControlRequest{
|
|
Command: datapb.GcCommand_Pause,
|
|
Params: []*commonpb.KeyValuePair{
|
|
{Key: "duration", Value: "not_int"},
|
|
},
|
|
})
|
|
s.Nil(err)
|
|
s.False(merr.Ok(resp))
|
|
|
|
resp, err = s.server.GcControl(context.TODO(), &datapb.GcControlRequest{
|
|
Command: datapb.GcCommand_Pause,
|
|
Params: []*commonpb.KeyValuePair{
|
|
{Key: "duration", Value: "60"},
|
|
},
|
|
})
|
|
s.Nil(err)
|
|
s.True(merr.Ok(resp))
|
|
}
|
|
|
|
func (s *GcControlServiceSuite) TestResume() {
|
|
resp, err := s.server.GcControl(context.TODO(), &datapb.GcControlRequest{
|
|
Command: datapb.GcCommand_Resume,
|
|
})
|
|
s.Nil(err)
|
|
s.True(merr.Ok(resp))
|
|
}
|
|
|
|
func (s *GcControlServiceSuite) TestTimeoutCtx() {
|
|
s.server.garbageCollector.close()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
resp, err := s.server.GcControl(ctx, &datapb.GcControlRequest{
|
|
Command: datapb.GcCommand_Resume,
|
|
})
|
|
s.Nil(err)
|
|
s.False(merr.Ok(resp))
|
|
|
|
resp, err = s.server.GcControl(ctx, &datapb.GcControlRequest{
|
|
Command: datapb.GcCommand_Pause,
|
|
Params: []*commonpb.KeyValuePair{
|
|
{Key: "duration", Value: "60"},
|
|
},
|
|
})
|
|
s.Nil(err)
|
|
s.False(merr.Ok(resp))
|
|
}
|
|
|
|
func TestGcControlService(t *testing.T) {
|
|
suite.Run(t, new(GcControlServiceSuite))
|
|
}
|