mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-01 19:39:21 +08:00
424 lines
13 KiB
Go
424 lines
13 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 datanode
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/milvus-io/milvus/internal/msgstream"
|
|
"github.com/milvus-io/milvus/internal/proto/commonpb"
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
|
"github.com/milvus-io/milvus/internal/proto/internalpb"
|
|
"github.com/milvus-io/milvus/internal/util/flowgraph"
|
|
)
|
|
|
|
type mockFactory struct {
|
|
msgstream.Factory
|
|
}
|
|
|
|
func TestFlowGraph_DDNode_newDDNode(te *testing.T) {
|
|
tests := []struct {
|
|
inCollID UniqueID
|
|
|
|
inFlushedSegs []UniqueID
|
|
inFlushedChannelTs Timestamp
|
|
inUnFlushedSegID UniqueID
|
|
inUnFlushedChannelTs Timestamp
|
|
|
|
description string
|
|
}{
|
|
{UniqueID(1), []UniqueID{100, 101, 102}, 666666, 200, 666666,
|
|
"Input VchannelInfo with 3 flushed segs and 1 unflushed seg"},
|
|
{UniqueID(2), []UniqueID{103}, 666666, 200, 666666,
|
|
"Input VchannelInfo with 1 flushed seg and 1 unflushed seg"},
|
|
{UniqueID(3), []UniqueID{}, 666666, 200, 666666,
|
|
"Input VchannelInfo with 0 flushed segs and 1 unflushed seg"},
|
|
{UniqueID(3), []UniqueID{104}, 666666, 0, 0,
|
|
"Input VchannelInfo with 1 flushed seg and empty unflushed seg"},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
te.Run(test.description, func(t *testing.T) {
|
|
di := &datapb.SegmentInfo{}
|
|
|
|
if test.inUnFlushedSegID != 0 {
|
|
di.ID = test.inUnFlushedSegID
|
|
di.DmlPosition = &internalpb.MsgPosition{Timestamp: test.inUnFlushedChannelTs}
|
|
}
|
|
|
|
var fi []*datapb.SegmentInfo
|
|
for _, id := range test.inFlushedSegs {
|
|
s := &datapb.SegmentInfo{ID: id}
|
|
fi = append(fi, s)
|
|
}
|
|
|
|
mmf := &mockMsgStreamFactory{
|
|
true, true,
|
|
}
|
|
ddNode := newDDNode(
|
|
context.Background(),
|
|
test.inCollID,
|
|
&datapb.VchannelInfo{
|
|
FlushedSegments: fi,
|
|
UnflushedSegments: []*datapb.SegmentInfo{di},
|
|
ChannelName: "by-dev-rootcoord-dml-test",
|
|
},
|
|
mmf,
|
|
newCompactionExecutor(),
|
|
)
|
|
require.NotNil(t, ddNode)
|
|
|
|
var flushedSegIDs []UniqueID
|
|
for _, seg := range ddNode.flushedSegments {
|
|
flushedSegIDs = append(flushedSegIDs, seg.ID)
|
|
}
|
|
assert.Equal(t, "ddNode", ddNode.Name())
|
|
assert.Equal(t, test.inCollID, ddNode.collectionID)
|
|
assert.Equal(t, len(test.inFlushedSegs), len(ddNode.flushedSegments))
|
|
assert.ElementsMatch(t, test.inFlushedSegs, flushedSegIDs)
|
|
|
|
si, ok := ddNode.segID2SegInfo.Load(test.inUnFlushedSegID)
|
|
assert.True(t, ok)
|
|
assert.Equal(t, test.inUnFlushedSegID, si.(*datapb.SegmentInfo).GetID())
|
|
assert.Equal(t, test.inUnFlushedChannelTs, si.(*datapb.SegmentInfo).GetDmlPosition().GetTimestamp())
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFlowGraph_DDNode_Operate(to *testing.T) {
|
|
to.Run("Test DDNode Operate DropCollection Msg", func(te *testing.T) {
|
|
// invalid inputs
|
|
invalidInTests := []struct {
|
|
in []Msg
|
|
description string
|
|
}{
|
|
{[]Msg{},
|
|
"Invalid input length == 0"},
|
|
{[]Msg{&flowGraphMsg{}, &flowGraphMsg{}, &flowGraphMsg{}},
|
|
"Invalid input length == 3"},
|
|
{[]Msg{&flowGraphMsg{}},
|
|
"Invalid input length == 1 but input message is not msgStreamMsg"},
|
|
}
|
|
|
|
for _, test := range invalidInTests {
|
|
te.Run(test.description, func(t *testing.T) {
|
|
ddn := ddNode{}
|
|
rt := ddn.Operate(test.in)
|
|
assert.Empty(t, rt)
|
|
})
|
|
}
|
|
|
|
// valid inputs
|
|
tests := []struct {
|
|
ddnCollID UniqueID
|
|
|
|
msgCollID UniqueID
|
|
expectedChlen int
|
|
|
|
description string
|
|
}{
|
|
{1, 1, 1,
|
|
"DropCollectionMsg collID == ddNode collID"},
|
|
{1, 2, 0,
|
|
"DropCollectionMsg collID != ddNode collID"},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
te.Run(test.description, func(t *testing.T) {
|
|
factory := msgstream.NewPmsFactory()
|
|
deltaStream, err := factory.NewMsgStream(context.Background())
|
|
assert.Nil(t, err)
|
|
ddn := ddNode{
|
|
collectionID: test.ddnCollID,
|
|
deltaMsgStream: deltaStream,
|
|
vchannelName: "ddn_drop_msg",
|
|
compactionExecutor: newCompactionExecutor(),
|
|
}
|
|
|
|
var dropCollMsg msgstream.TsMsg = &msgstream.DropCollectionMsg{
|
|
DropCollectionRequest: internalpb.DropCollectionRequest{
|
|
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_DropCollection},
|
|
CollectionID: test.msgCollID,
|
|
},
|
|
}
|
|
tsMessages := []msgstream.TsMsg{dropCollMsg}
|
|
var msgStreamMsg Msg = flowgraph.GenerateMsgStreamMsg(tsMessages, 0, 0, nil, nil)
|
|
|
|
rt := ddn.Operate([]Msg{msgStreamMsg})
|
|
|
|
if test.ddnCollID == test.msgCollID {
|
|
assert.NotEmpty(t, rt)
|
|
assert.True(t, rt[0].(*flowGraphMsg).dropCollection)
|
|
} else {
|
|
assert.NotEmpty(t, rt)
|
|
}
|
|
})
|
|
}
|
|
})
|
|
|
|
to.Run("Test DDNode Operate Insert Msg", func(te *testing.T) {
|
|
tests := []struct {
|
|
ddnCollID UniqueID
|
|
inMsgCollID UniqueID
|
|
|
|
MsgEndTs Timestamp
|
|
threshold Timestamp
|
|
|
|
ddnFlushedSegment UniqueID
|
|
inMsgSegID UniqueID
|
|
|
|
expectedRtLen int
|
|
description string
|
|
}{
|
|
{1, 1, 2000, 3000, 100, 100, 0,
|
|
"MsgEndTs(2000) < threshold(3000), inMsgSegID(100) IN ddnFlushedSeg {100}"},
|
|
{1, 1, 2000, 3000, 100, 200, 1,
|
|
"MsgEndTs(2000) < threshold(3000), inMsgSegID(200) NOT IN ddnFlushedSeg {100}"},
|
|
{1, 1, 4000, 3000, 100, 101, 1,
|
|
"Seg 101, MsgEndTs(4000) > FilterThreshold(3000)"},
|
|
{1, 1, 4000, 3000, 100, 200, 1,
|
|
"Seg 200, MsgEndTs(4000) > FilterThreshold(3000)"},
|
|
{1, 2, 4000, 3000, 100, 100, 0,
|
|
"inMsgCollID(2) != ddnCollID"},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
te.Run(test.description, func(t *testing.T) {
|
|
fs := &datapb.SegmentInfo{ID: test.ddnFlushedSegment}
|
|
factory := msgstream.NewPmsFactory()
|
|
deltaStream, err := factory.NewMsgStream(context.Background())
|
|
assert.Nil(t, err)
|
|
// Prepare ddNode states
|
|
ddn := ddNode{
|
|
flushedSegments: []*datapb.SegmentInfo{fs},
|
|
collectionID: test.ddnCollID,
|
|
deltaMsgStream: deltaStream,
|
|
}
|
|
FilterThreshold = test.threshold
|
|
|
|
// Prepare insert messages
|
|
var iMsg msgstream.TsMsg = &msgstream.InsertMsg{
|
|
BaseMsg: msgstream.BaseMsg{EndTimestamp: test.MsgEndTs},
|
|
InsertRequest: internalpb.InsertRequest{
|
|
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Insert},
|
|
CollectionID: test.inMsgCollID,
|
|
SegmentID: test.inMsgSegID,
|
|
},
|
|
}
|
|
tsMessages := []msgstream.TsMsg{iMsg}
|
|
var msgStreamMsg Msg = flowgraph.GenerateMsgStreamMsg(tsMessages, 0, 0, nil, nil)
|
|
|
|
// Test
|
|
rt := ddn.Operate([]Msg{msgStreamMsg})
|
|
assert.Equal(t, test.expectedRtLen, len(rt[0].(*flowGraphMsg).insertMessages))
|
|
})
|
|
}
|
|
})
|
|
|
|
to.Run("Test DDNode Operate Delete Msg", func(te *testing.T) {
|
|
tests := []struct {
|
|
ddnCollID UniqueID
|
|
inMsgCollID UniqueID
|
|
|
|
MsgEndTs Timestamp
|
|
|
|
expectedRtLen int
|
|
description string
|
|
}{
|
|
{1, 1, 2000, 1, "normal"},
|
|
{1, 2, 4000, 0, "inMsgCollID(2) != ddnCollID"},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
te.Run(test.description, func(t *testing.T) {
|
|
factory := msgstream.NewPmsFactory()
|
|
deltaStream, err := factory.NewMsgStream(context.Background())
|
|
assert.Nil(t, err)
|
|
// Prepare ddNode states
|
|
ddn := ddNode{
|
|
collectionID: test.ddnCollID,
|
|
deltaMsgStream: deltaStream,
|
|
}
|
|
|
|
// Prepare delete messages
|
|
var dMsg msgstream.TsMsg = &msgstream.DeleteMsg{
|
|
BaseMsg: msgstream.BaseMsg{EndTimestamp: test.MsgEndTs},
|
|
DeleteRequest: internalpb.DeleteRequest{
|
|
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Delete},
|
|
CollectionID: test.inMsgCollID,
|
|
},
|
|
}
|
|
tsMessages := []msgstream.TsMsg{dMsg}
|
|
var msgStreamMsg Msg = flowgraph.GenerateMsgStreamMsg(tsMessages, 0, 0, nil, nil)
|
|
|
|
// Test
|
|
rt := ddn.Operate([]Msg{msgStreamMsg})
|
|
assert.Equal(t, test.expectedRtLen, len(rt[0].(*flowGraphMsg).deleteMessages))
|
|
})
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestFlowGraph_DDNode_filterMessages(te *testing.T) {
|
|
tests := []struct {
|
|
ddnFlushedSegments []UniqueID
|
|
ddnSegID2Ts map[UniqueID]Timestamp
|
|
|
|
inMsgSegID UniqueID
|
|
inMsgSegEntTs Timestamp
|
|
expectedOut bool
|
|
|
|
description string
|
|
}{
|
|
{[]UniqueID{1, 2, 3}, map[UniqueID]Timestamp{4: 1000, 5: 2000}, 1, 1500, true,
|
|
"Seg 1 in flushedSegs {1, 2, 3}"},
|
|
{[]UniqueID{1, 2, 3}, map[UniqueID]Timestamp{4: 1000, 5: 2000}, 2, 1500, true,
|
|
"Seg 2 in flushedSegs {1, 2, 3}"},
|
|
{[]UniqueID{1, 2, 3}, map[UniqueID]Timestamp{4: 1000, 5: 2000}, 3, 1500, true,
|
|
"Seg 3 in flushedSegs {1, 2, 3}"},
|
|
{[]UniqueID{1, 2, 3}, map[UniqueID]Timestamp{4: 1000, 5: 2000}, 4, 1500, false,
|
|
"Seg 4, inMsgSegEntTs(1500) > SegCheckPoint(1000)"},
|
|
{[]UniqueID{1, 2, 3}, map[UniqueID]Timestamp{4: 1000, 5: 2000}, 4, 500, true,
|
|
"Seg 4, inMsgSegEntTs(500) <= SegCheckPoint(1000)"},
|
|
{[]UniqueID{1, 2, 3}, map[UniqueID]Timestamp{4: 1000, 5: 2000}, 4, 1000, true,
|
|
"Seg 4 inMsgSegEntTs(1000) <= SegCheckPoint(1000)"},
|
|
{[]UniqueID{1, 2, 3}, map[UniqueID]Timestamp{4: 1000, 5: 2000}, 5, 1500, true,
|
|
"Seg 5 inMsgSegEntTs(1500) <= SegCheckPoint(2000)"},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
te.Run(test.description, func(t *testing.T) {
|
|
fs := []*datapb.SegmentInfo{}
|
|
for _, id := range test.ddnFlushedSegments {
|
|
s := &datapb.SegmentInfo{ID: id}
|
|
fs = append(fs, s)
|
|
}
|
|
factory := msgstream.NewPmsFactory()
|
|
deltaStream, err := factory.NewMsgStream(context.Background())
|
|
assert.Nil(t, err)
|
|
// Prepare ddNode states
|
|
ddn := ddNode{
|
|
flushedSegments: fs,
|
|
deltaMsgStream: deltaStream,
|
|
}
|
|
|
|
for k, v := range test.ddnSegID2Ts {
|
|
ddn.segID2SegInfo.Store(k, &datapb.SegmentInfo{DmlPosition: &internalpb.MsgPosition{Timestamp: v}})
|
|
}
|
|
|
|
// Prepare insert messages
|
|
var iMsg = &msgstream.InsertMsg{
|
|
BaseMsg: msgstream.BaseMsg{EndTimestamp: test.inMsgSegEntTs},
|
|
InsertRequest: internalpb.InsertRequest{
|
|
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Insert},
|
|
SegmentID: test.inMsgSegID,
|
|
},
|
|
}
|
|
|
|
// Test
|
|
rt := ddn.filterFlushedSegmentInsertMessages(iMsg)
|
|
assert.Equal(t, test.expectedOut, rt)
|
|
|
|
si, ok := ddn.segID2SegInfo.Load(iMsg.GetSegmentID())
|
|
if !rt {
|
|
assert.False(t, ok)
|
|
assert.Nil(t, si)
|
|
}
|
|
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFlowGraph_DDNode_isFlushed(te *testing.T) {
|
|
tests := []struct {
|
|
influshedSegment []UniqueID
|
|
inSeg UniqueID
|
|
|
|
expectedOut bool
|
|
|
|
description string
|
|
}{
|
|
{[]UniqueID{1, 2, 3}, 1, true,
|
|
"Input seg 1 in flushedSegs{1, 2, 3}"},
|
|
{[]UniqueID{1, 2, 3}, 2, true,
|
|
"Input seg 2 in flushedSegs{1, 2, 3}"},
|
|
{[]UniqueID{1, 2, 3}, 3, true,
|
|
"Input seg 3 in flushedSegs{1, 2, 3}"},
|
|
{[]UniqueID{1, 2, 3}, 4, false,
|
|
"Input seg 4 not in flushedSegs{1, 2, 3}"},
|
|
{[]UniqueID{}, 5, false,
|
|
"Input seg 5, no flushedSegs {}"},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
te.Run(test.description, func(t *testing.T) {
|
|
fs := []*datapb.SegmentInfo{}
|
|
for _, id := range test.influshedSegment {
|
|
s := &datapb.SegmentInfo{ID: id}
|
|
fs = append(fs, s)
|
|
}
|
|
factory := msgstream.NewPmsFactory()
|
|
deltaStream, err := factory.NewMsgStream(context.Background())
|
|
assert.Nil(t, err)
|
|
ddn := &ddNode{flushedSegments: fs, deltaMsgStream: deltaStream}
|
|
assert.Equal(t, test.expectedOut, ddn.isFlushed(test.inSeg))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFlowGraph_DDNode_isDropped(te *testing.T) {
|
|
genSegmentInfoByID := func(segmentID UniqueID) *datapb.SegmentInfo {
|
|
return &datapb.SegmentInfo{
|
|
ID: segmentID,
|
|
}
|
|
}
|
|
|
|
tests := []struct {
|
|
indroppedSegment []*datapb.SegmentInfo
|
|
inSeg UniqueID
|
|
|
|
expectedOut bool
|
|
|
|
description string
|
|
}{
|
|
{[]*datapb.SegmentInfo{genSegmentInfoByID(1), genSegmentInfoByID(2), genSegmentInfoByID(3)}, 1, true,
|
|
"Input seg 1 in droppedSegs{1, 2, 3}"},
|
|
{[]*datapb.SegmentInfo{genSegmentInfoByID(1), genSegmentInfoByID(2), genSegmentInfoByID(3)}, 2, true,
|
|
"Input seg 2 in droppedSegs{1, 2, 3}"},
|
|
{[]*datapb.SegmentInfo{genSegmentInfoByID(1), genSegmentInfoByID(2), genSegmentInfoByID(3)}, 3, true,
|
|
"Input seg 3 in droppedSegs{1, 2, 3}"},
|
|
{[]*datapb.SegmentInfo{genSegmentInfoByID(1), genSegmentInfoByID(2), genSegmentInfoByID(3)}, 4, false,
|
|
"Input seg 4 not in droppedSegs{1, 2, 3}"},
|
|
{[]*datapb.SegmentInfo{}, 5, false,
|
|
"Input seg 5, no droppedSegs {}"},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
te.Run(test.description, func(t *testing.T) {
|
|
factory := mockMsgStreamFactory{true, true}
|
|
deltaStream, err := factory.NewMsgStream(context.Background())
|
|
assert.Nil(t, err)
|
|
ddn := &ddNode{droppedSegments: test.indroppedSegment, deltaMsgStream: deltaStream}
|
|
assert.Equal(t, test.expectedOut, ddn.isDropped(test.inSeg))
|
|
})
|
|
}
|
|
}
|