Files
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

2346 lines
70 KiB
Go

package datacoord
import (
"context"
"testing"
"time"
"github.com/bytedance/mockey"
"github.com/samber/lo"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus/internal/datacoord/broker"
"github.com/milvus-io/milvus/internal/metastore/model"
mocks2 "github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/objectstorage"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v3/proto/rootcoordpb"
"github.com/milvus-io/milvus/pkg/v3/proto/workerpb"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
)
func TestGetQueryVChanPositionsRetrieveM2N(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
channel := "ch1"
svr.meta.AddCollection(&collectionInfo{
ID: 1,
Partitions: []int64{0},
Schema: schema,
StartPositions: []*commonpb.KeyDataPair{
{
Key: channel,
Data: []byte{8, 9, 10},
},
},
})
err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
CollectionID: 1,
FieldID: 2,
IndexID: 1,
})
require.NoError(t, err)
segArgs := []struct {
segID int64
state commonpb.SegmentState
level datapb.SegmentLevel
indexed bool
}{
{100, commonpb.SegmentState_Growing, datapb.SegmentLevel_L1, false},
{200, commonpb.SegmentState_Flushing, datapb.SegmentLevel_L1, false},
{300, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L1, false},
{400, commonpb.SegmentState_Dropped, datapb.SegmentLevel_L1, true},
{401, commonpb.SegmentState_Dropped, datapb.SegmentLevel_L1, true},
{402, commonpb.SegmentState_Dropped, datapb.SegmentLevel_L1, true},
{403, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L1, true},
{404, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L1, false},
// (400[indexed], 401[indexed], 402(indexed) -> 403(indexed), 404(no index))
{500, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L2, true},
{600, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L2, false},
{700, commonpb.SegmentState_Flushed, datapb.SegmentLevel_L0, false},
{800, commonpb.SegmentState_Dropped, datapb.SegmentLevel_L1, false},
}
compactFroms := []int64{400, 401, 402}
compactTos := []int64{403, 404}
for _, arg := range segArgs {
seg := NewSegmentInfo(&datapb.SegmentInfo{
ID: arg.segID,
CollectionID: 1,
InsertChannel: channel,
State: arg.state,
Level: arg.level,
DmlPosition: &msgpb.MsgPosition{
ChannelName: channel,
MsgID: []byte{1, 2, 3},
},
NumOfRows: 2048,
})
if lo.Contains(compactTos, arg.segID) {
seg.CompactionFrom = compactFroms
}
err := svr.meta.AddSegment(context.TODO(), seg)
require.NoError(t, err)
if arg.indexed {
err = svr.meta.indexMeta.AddSegmentIndex(context.TODO(), &model.SegmentIndex{
SegmentID: arg.segID,
BuildID: arg.segID,
IndexID: 1,
})
assert.NoError(t, err)
err = svr.meta.indexMeta.FinishTask(&workerpb.IndexTaskInfo{
BuildID: arg.segID,
State: commonpb.IndexState_Finished,
})
assert.NoError(t, err)
}
}
info := svr.handler.GetQueryVChanPositions(&channelMeta{Name: channel, CollectionID: 1}, -1)
totalSegs := len(info.GetLevelZeroSegmentIds()) +
len(info.GetUnflushedSegmentIds()) +
len(info.GetFlushedSegmentIds()) +
len(info.GetDroppedSegmentIds())
assert.EqualValues(t, 1, info.CollectionID)
assert.EqualValues(t, len(segArgs)-2, totalSegs)
assert.ElementsMatch(t, []int64{700}, info.GetLevelZeroSegmentIds())
assert.ElementsMatch(t, []int64{100}, info.GetUnflushedSegmentIds())
assert.ElementsMatch(t, []int64{200, 300, 400, 401, 402, 500, 600}, info.GetFlushedSegmentIds())
assert.ElementsMatch(t, []int64{800}, info.GetDroppedSegmentIds())
assert.Empty(t, info.GetUnflushedSegments())
assert.Empty(t, info.GetFlushedSegments())
assert.Empty(t, info.GetDroppedSegments())
assert.Empty(t, info.GetIndexedSegments())
assert.Empty(t, info.GetIndexedSegmentIds())
}
func TestGetQueryVChanPositions(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
svr.meta.AddCollection(&collectionInfo{
ID: 0,
Partitions: []int64{0, 1},
Schema: schema,
StartPositions: []*commonpb.KeyDataPair{
{
Key: "ch1",
Data: []byte{8, 9, 10},
},
},
})
svr.meta.AddCollection(&collectionInfo{
ID: 1,
Schema: schema,
StartPositions: []*commonpb.KeyDataPair{
{
Key: "ch0",
Data: []byte{8, 9, 10},
},
},
})
err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
CollectionID: 0,
FieldID: 2,
IndexID: 1,
})
assert.NoError(t, err)
s1 := &datapb.SegmentInfo{
ID: 1,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 0,
},
NumOfRows: 2048,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s1))
assert.NoError(t, err)
err = svr.meta.indexMeta.AddSegmentIndex(context.TODO(), &model.SegmentIndex{
SegmentID: 1,
BuildID: 1,
IndexID: 1,
})
assert.NoError(t, err)
err = svr.meta.indexMeta.FinishTask(&workerpb.IndexTaskInfo{
BuildID: 1,
State: commonpb.IndexState_Finished,
})
assert.NoError(t, err)
s2 := &datapb.SegmentInfo{
ID: 2,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Growing,
StartPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{8, 9, 10},
MsgGroup: "",
},
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s2))
assert.NoError(t, err)
s3 := &datapb.SegmentInfo{
ID: 3,
CollectionID: 0,
PartitionID: 1,
InsertChannel: "ch1",
State: commonpb.SegmentState_Growing,
StartPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{8, 9, 10},
MsgGroup: "",
},
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{11, 12, 13},
MsgGroup: "",
Timestamp: 2,
},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s3))
assert.NoError(t, err)
s4 := &datapb.SegmentInfo{
ID: 4,
CollectionID: 0,
PartitionID: common.AllPartitionsID,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
StartPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{8, 9, 10},
MsgGroup: "",
Timestamp: 1,
},
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{11, 12, 13},
MsgGroup: "",
Timestamp: 2,
},
Level: datapb.SegmentLevel_L0,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s4))
assert.NoError(t, err)
t.Run("get unexisted channel", func(t *testing.T) {
vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "chx1", CollectionID: 0})
assert.Empty(t, vchan.UnflushedSegmentIds)
assert.Empty(t, vchan.FlushedSegmentIds)
})
t.Run("empty collection", func(t *testing.T) {
infos := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch0_suffix", CollectionID: 1})
assert.EqualValues(t, 1, infos.CollectionID)
assert.EqualValues(t, 0, len(infos.FlushedSegmentIds))
assert.EqualValues(t, 0, len(infos.UnflushedSegmentIds))
assert.EqualValues(t, 0, len(infos.GetLevelZeroSegmentIds()))
})
t.Run("filter partition", func(t *testing.T) {
infos := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0}, 1)
assert.EqualValues(t, 0, infos.CollectionID)
assert.EqualValues(t, 1, len(infos.GetLevelZeroSegmentIds()))
assert.EqualValues(t, uint64(1), infos.GetDeleteCheckpoint().GetTimestamp())
})
t.Run("empty collection with passed positions", func(t *testing.T) {
vchannel := "ch_no_segment_1"
pchannel := funcutil.ToPhysicalChannel(vchannel)
infos := svr.handler.GetQueryVChanPositions(&channelMeta{
Name: vchannel,
CollectionID: 0,
StartPosition: &msgpb.MsgPosition{
ChannelName: pchannel,
MsgID: []byte{14, 15, 16},
},
})
assert.EqualValues(t, 0, infos.CollectionID)
assert.EqualValues(t, vchannel, infos.ChannelName)
assert.EqualValues(t, 0, len(infos.GetLevelZeroSegmentIds()))
})
}
func TestGetQueryVChanPositions_PartitionStats(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
collectionID := int64(0)
partitionID := int64(1)
vchannel := "test_vchannel"
version := int64(100)
svr.meta.AddCollection(&collectionInfo{
ID: collectionID,
Schema: schema,
})
svr.meta.partitionStatsMeta.partitionStatsInfos = map[string]map[int64]*partitionStatsInfo{
vchannel: {
partitionID: {
currentVersion: version,
infos: map[int64]*datapb.PartitionStatsInfo{
version: {Version: version},
},
},
partitionID + 1: {
currentVersion: version + 1,
infos: map[int64]*datapb.PartitionStatsInfo{
version + 1: {Version: version + 1},
},
},
},
}
vChannelInfo := svr.handler.GetQueryVChanPositions(&channelMeta{Name: vchannel, CollectionID: collectionID}, partitionID)
statsVersions := vChannelInfo.GetPartitionStatsVersions()
assert.Equal(t, 1, len(statsVersions))
assert.Equal(t, int64(100), statsVersions[partitionID])
vChannelInfo2 := svr.handler.GetQueryVChanPositions(&channelMeta{Name: vchannel, CollectionID: collectionID})
statsVersions2 := vChannelInfo2.GetPartitionStatsVersions()
assert.Equal(t, 2, len(statsVersions2))
}
func TestGetQueryVChanPositions_Retrieve_unIndexed(t *testing.T) {
t.Run("ab GC-ed, cde unIndexed", func(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
svr.meta.AddCollection(&collectionInfo{
ID: 0,
Schema: schema,
})
err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
CollectionID: 0,
FieldID: 2,
IndexID: 1,
})
assert.NoError(t, err)
c := &datapb.SegmentInfo{
ID: 1,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
CompactionFrom: []int64{99, 100}, // a, b which have been GC-ed
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(c))
assert.NoError(t, err)
d := &datapb.SegmentInfo{
ID: 2,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(d))
assert.NoError(t, err)
e := &datapb.SegmentInfo{
ID: 3,
CollectionID: 0,
PartitionID: 1,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
CompactionFrom: []int64{1, 2}, // c, d
NumOfRows: 2048,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(e))
assert.NoError(t, err)
// vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0})
// assert.EqualValues(t, 2, len(vchan.FlushedSegmentIds))
// assert.EqualValues(t, 0, len(vchan.UnflushedSegmentIds))
// assert.ElementsMatch(t, []int64{c.GetID(), d.GetID()}, vchan.FlushedSegmentIds) // expected c, d
})
t.Run("a GC-ed, bcde unIndexed", func(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
svr.meta.AddCollection(&collectionInfo{
ID: 0,
Schema: schema,
})
err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
CollectionID: 0,
FieldID: 2,
IndexID: 1,
})
assert.NoError(t, err)
a := &datapb.SegmentInfo{
ID: 99,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(a))
assert.NoError(t, err)
c := &datapb.SegmentInfo{
ID: 1,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
CompactionFrom: []int64{99, 100}, // a, b which have been GC-ed
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(c))
assert.NoError(t, err)
d := &datapb.SegmentInfo{
ID: 2,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(d))
assert.NoError(t, err)
e := &datapb.SegmentInfo{
ID: 3,
CollectionID: 0,
PartitionID: 1,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
CompactionFrom: []int64{1, 2}, // c, d
NumOfRows: 2048,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(e))
assert.NoError(t, err)
// vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0})
// assert.EqualValues(t, 2, len(vchan.FlushedSegmentIds))
// assert.EqualValues(t, 0, len(vchan.UnflushedSegmentIds))
// assert.ElementsMatch(t, []int64{c.GetID(), d.GetID()}, vchan.FlushedSegmentIds) // expected c, d
})
t.Run("ab GC-ed, c unIndexed, de indexed", func(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
svr.meta.AddCollection(&collectionInfo{
ID: 0,
Schema: schema,
})
err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
CollectionID: 0,
FieldID: 2,
IndexID: 1,
})
assert.NoError(t, err)
c := &datapb.SegmentInfo{
ID: 1,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
CompactionFrom: []int64{99, 100}, // a, b which have been GC-ed
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(c))
assert.NoError(t, err)
d := &datapb.SegmentInfo{
ID: 2,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(d))
assert.NoError(t, err)
err = svr.meta.indexMeta.AddSegmentIndex(context.TODO(), &model.SegmentIndex{
SegmentID: 2,
BuildID: 1,
IndexID: 1,
})
assert.NoError(t, err)
err = svr.meta.indexMeta.FinishTask(&workerpb.IndexTaskInfo{
BuildID: 1,
State: commonpb.IndexState_Finished,
})
assert.NoError(t, err)
e := &datapb.SegmentInfo{
ID: 3,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
CompactionFrom: []int64{1, 2}, // c, d
NumOfRows: 2048,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(e))
assert.NoError(t, err)
err = svr.meta.indexMeta.AddSegmentIndex(context.TODO(), &model.SegmentIndex{
SegmentID: 3,
BuildID: 2,
IndexID: 1,
})
assert.NoError(t, err)
err = svr.meta.indexMeta.FinishTask(&workerpb.IndexTaskInfo{
BuildID: 2,
State: commonpb.IndexState_Finished,
})
assert.NoError(t, err)
// vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0})
// assert.EqualValues(t, 1, len(vchan.FlushedSegmentIds))
// assert.EqualValues(t, 0, len(vchan.UnflushedSegmentIds))
// assert.ElementsMatch(t, []int64{e.GetID()}, vchan.FlushedSegmentIds) // expected e
})
t.Run("complex derivation", func(t *testing.T) {
// numbers indicate segmentID, letters indicate segment information
// i: indexed, u: unindexed, g: gced
// 1i, 2i, 3g 4i, 5i, 6i
// | | | | | |
// \ | / \ | /
// \ | / \ | /
// 7u, [8i,9i,10i] [11u, 12i]
// | | | | | |
// \ | / \ / |
// \ | / \ / |
// [13u] [14i, 15u] 12i [19u](unsorted)
// | | | | |
// \ / \ / |
// \ / \ / |
// [16u] [17u] [18u](unsorted) [20u](sorted) [21i](unsorted)
// all leaf nodes are [1,2,3,4,5,6,7], but because segment3 has been gced, the leaf node becomes [7,8,9,10,4,5,6]
// should be returned: flushed: [7, 8, 9, 10, 4, 5, 6, 20, 21], growing: [18]
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
svr.meta.AddCollection(&collectionInfo{
ID: 0,
Partitions: []int64{0},
Schema: schema,
})
err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
CollectionID: 0,
FieldID: 2,
IndexID: 1,
})
assert.NoError(t, err)
seg1 := &datapb.SegmentInfo{
ID: 1,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
assert.NoError(t, err)
seg2 := &datapb.SegmentInfo{
ID: 2,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
assert.NoError(t, err)
// seg3 was GCed
seg4 := &datapb.SegmentInfo{
ID: 4,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg4))
assert.NoError(t, err)
seg5 := &datapb.SegmentInfo{
ID: 5,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg5))
assert.NoError(t, err)
seg6 := &datapb.SegmentInfo{
ID: 6,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg6))
assert.NoError(t, err)
seg7 := &datapb.SegmentInfo{
ID: 7,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 2048,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg7))
assert.NoError(t, err)
seg8 := &datapb.SegmentInfo{
ID: 8,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{1, 2, 3},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg8))
assert.NoError(t, err)
seg9 := &datapb.SegmentInfo{
ID: 9,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{1, 2, 3},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg9))
assert.NoError(t, err)
seg10 := &datapb.SegmentInfo{
ID: 10,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{1, 2, 3},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg10))
assert.NoError(t, err)
seg11 := &datapb.SegmentInfo{
ID: 11,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 2048,
CompactionFrom: []int64{4, 5, 6},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg11))
assert.NoError(t, err)
seg12 := &datapb.SegmentInfo{
ID: 12,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{4, 5, 6},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg12))
assert.NoError(t, err)
seg13 := &datapb.SegmentInfo{
ID: 13,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 2047,
CompactionFrom: []int64{7, 8, 9},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg13))
assert.NoError(t, err)
seg14 := &datapb.SegmentInfo{
ID: 14,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{10, 11},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg14))
assert.NoError(t, err)
seg15 := &datapb.SegmentInfo{
ID: 15,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 2048,
CompactionFrom: []int64{10, 11},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg15))
assert.NoError(t, err)
seg16 := &datapb.SegmentInfo{
ID: 16,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 2048,
CompactionFrom: []int64{13, 14},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg16))
assert.NoError(t, err)
seg17 := &datapb.SegmentInfo{
ID: 17,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 2048,
CompactionFrom: []int64{12, 15},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg17))
assert.NoError(t, err)
seg18 := &datapb.SegmentInfo{
ID: 18,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Level: datapb.SegmentLevel_L1,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 2048,
CompactionFrom: []int64{},
IsInvisible: true,
IsSorted: false,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg18))
assert.NoError(t, err)
seg19 := &datapb.SegmentInfo{
ID: 19,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
Level: datapb.SegmentLevel_L1,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 2048,
CompactionFrom: []int64{},
IsInvisible: true,
IsSorted: false,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg19))
assert.NoError(t, err)
seg20 := &datapb.SegmentInfo{
ID: 20,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Level: datapb.SegmentLevel_L1,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 2048,
CompactionFrom: []int64{19},
CreatedByCompaction: true,
IsInvisible: false,
IsSorted: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg20))
assert.NoError(t, err)
seg21 := &datapb.SegmentInfo{
ID: 21,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Level: datapb.SegmentLevel_L1,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{},
IsInvisible: false,
IsSorted: false,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg21))
assert.NoError(t, err)
vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0})
assert.ElementsMatch(t, []int64{7, 8, 9, 10, 4, 5, 6, 20, 21}, vchan.FlushedSegmentIds)
assert.ElementsMatch(t, []int64{18}, vchan.UnflushedSegmentIds)
assert.ElementsMatch(t, []int64{1, 2, 19}, vchan.DroppedSegmentIds)
})
t.Run("compaction iterate", func(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
svr.meta.AddCollection(&collectionInfo{
ID: 0,
Partitions: []int64{0},
Schema: schema,
})
err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
CollectionID: 0,
FieldID: 2,
IndexID: 1,
})
assert.NoError(t, err)
seg1 := &datapb.SegmentInfo{
ID: 1,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
assert.NoError(t, err)
seg2 := &datapb.SegmentInfo{
ID: 2,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
assert.NoError(t, err)
seg3 := &datapb.SegmentInfo{
ID: 3,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{1, 2},
IsInvisible: true,
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg3))
assert.NoError(t, err)
seg4 := &datapb.SegmentInfo{
ID: 4,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{1, 2},
IsInvisible: true,
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg4))
assert.NoError(t, err)
seg5 := &datapb.SegmentInfo{
ID: 5,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{3},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg5))
assert.NoError(t, err)
seg6 := &datapb.SegmentInfo{
ID: 6,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{4},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg6))
assert.NoError(t, err)
seg7 := &datapb.SegmentInfo{
ID: 7,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{5, 6},
IsInvisible: true,
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg7))
assert.NoError(t, err)
seg8 := &datapb.SegmentInfo{
ID: 8,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{5, 6},
IsInvisible: true,
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg8))
assert.NoError(t, err)
seg9 := &datapb.SegmentInfo{
ID: 9,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{7},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg9))
assert.NoError(t, err)
seg10 := &datapb.SegmentInfo{
ID: 10,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{8},
CreatedByCompaction: true,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg10))
assert.NoError(t, err)
vchan := svr.handler.GetQueryVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0})
assert.ElementsMatch(t, []int64{5, 6}, vchan.FlushedSegmentIds)
assert.ElementsMatch(t, []int64{}, vchan.UnflushedSegmentIds)
assert.ElementsMatch(t, []int64{1, 2}, vchan.DroppedSegmentIds)
})
}
func TestGetCurrentSegmentsView(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
svr.meta.AddCollection(&collectionInfo{
ID: 0,
Partitions: []int64{0},
Schema: schema,
})
err := svr.meta.indexMeta.CreateIndex(context.TODO(), &model.Index{
CollectionID: 0,
FieldID: 2,
IndexID: 1,
})
assert.NoError(t, err)
seg1 := &datapb.SegmentInfo{
ID: 1,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Dropped,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg1))
assert.NoError(t, err)
seg2 := &datapb.SegmentInfo{
ID: 2,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{1},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg2))
assert.NoError(t, err)
seg3 := &datapb.SegmentInfo{
ID: 3,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg3))
assert.NoError(t, err)
seg4 := &datapb.SegmentInfo{
ID: 4,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg4))
assert.NoError(t, err)
seg5 := &datapb.SegmentInfo{
ID: 5,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{3, 4},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg5))
assert.NoError(t, err)
seg6 := &datapb.SegmentInfo{
ID: 6,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
CompactionFrom: []int64{3, 4},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg6))
assert.NoError(t, err)
seg7 := &datapb.SegmentInfo{
ID: 7,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Level: datapb.SegmentLevel_L0,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg7))
assert.NoError(t, err)
seg8 := &datapb.SegmentInfo{
ID: 8,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Growing,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg8))
assert.NoError(t, err)
seg9 := &datapb.SegmentInfo{
ID: 9,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Importing,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
MsgGroup: "",
Timestamp: 1,
},
NumOfRows: 100,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(seg9))
assert.NoError(t, err)
view := svr.handler.GetCurrentSegmentsView(context.Background(), &channelMeta{Name: "ch1", CollectionID: 0})
assert.ElementsMatch(t, []int64{2, 3, 4}, view.FlushedSegmentIDs)
assert.ElementsMatch(t, []int64{8}, view.GrowingSegmentIDs)
assert.ElementsMatch(t, []int64{1}, view.DroppedSegmentIDs)
assert.ElementsMatch(t, []int64{7}, view.L0SegmentIDs)
assert.ElementsMatch(t, []int64{9}, view.ImportingSegmentIDs)
}
func TestShouldDropChannel(t *testing.T) {
type myRootCoord struct {
mocks2.MockRootCoordClient
}
myRoot := &myRootCoord{}
myRoot.EXPECT().AllocTimestamp(mock.Anything, mock.Anything).Return(&rootcoordpb.AllocTimestampResponse{
Status: merr.Success(),
Timestamp: tsoutil.ComposeTSByTime(time.Now()),
Count: 1,
}, nil)
myRoot.EXPECT().AllocID(mock.Anything, mock.Anything).Return(&rootcoordpb.AllocIDResponse{
Status: merr.Success(),
ID: int64(tsoutil.ComposeTSByTime(time.Now())),
Count: 1,
}, nil)
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
svr.meta.AddCollection(&collectionInfo{
ID: 0,
Schema: schema,
StartPositions: []*commonpb.KeyDataPair{
{
Key: "ch1",
Data: []byte{8, 9, 10},
},
},
})
svr.meta.AddCollection(&collectionInfo{
ID: 1,
Schema: schema,
StartPositions: []*commonpb.KeyDataPair{
{
Key: "ch0",
Data: []byte{8, 9, 10},
},
},
})
t.Run("channel name not in kv ", func(t *testing.T) {
assert.False(t, svr.handler.CheckShouldDropChannel("ch99"))
})
t.Run("channel in remove flag", func(t *testing.T) {
err := svr.meta.catalog.MarkChannelDeleted(context.TODO(), "ch1")
require.NoError(t, err)
assert.True(t, svr.handler.CheckShouldDropChannel("ch1"))
})
}
func TestGetDataVChanPositions(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
schema := newTestSchema()
svr.meta.AddCollection(&collectionInfo{
ID: 0,
Schema: schema,
StartPositions: []*commonpb.KeyDataPair{
{
Key: "ch1",
Data: []byte{8, 9, 10},
},
},
})
svr.meta.AddCollection(&collectionInfo{
ID: 1,
Schema: schema,
StartPositions: []*commonpb.KeyDataPair{
{
Key: "ch0",
Data: []byte{8, 9, 10},
},
},
})
s1 := &datapb.SegmentInfo{
ID: 1,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
},
}
err := svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s1))
require.Nil(t, err)
s2 := &datapb.SegmentInfo{
ID: 2,
CollectionID: 0,
PartitionID: 0,
InsertChannel: "ch1",
State: commonpb.SegmentState_Growing,
StartPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{8, 9, 10},
},
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{1, 2, 3},
Timestamp: 1,
},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s2))
require.Nil(t, err)
s3 := &datapb.SegmentInfo{
ID: 3,
CollectionID: 0,
PartitionID: 1,
InsertChannel: "ch1",
State: commonpb.SegmentState_Growing,
StartPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{8, 9, 10},
},
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{11, 12, 13},
Timestamp: 2,
},
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(s3))
require.Nil(t, err)
l0Segment := &datapb.SegmentInfo{
ID: 4,
CollectionID: 0,
PartitionID: 1,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
StartPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{8, 9, 10},
},
DmlPosition: &msgpb.MsgPosition{
ChannelName: "ch1",
MsgID: []byte{11, 12, 13},
Timestamp: 2,
},
Level: datapb.SegmentLevel_L0,
}
err = svr.meta.AddSegment(context.TODO(), NewSegmentInfo(l0Segment))
require.Nil(t, err)
t.Run("get unexisted channel", func(t *testing.T) {
vchan := svr.handler.GetDataVChanPositions(&channelMeta{Name: "chx1", CollectionID: 0}, allPartitionID)
assert.Empty(t, vchan.UnflushedSegmentIds)
assert.Empty(t, vchan.FlushedSegmentIds)
})
t.Run("get existed channel", func(t *testing.T) {
vchan := svr.handler.GetDataVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0}, allPartitionID)
assert.EqualValues(t, 1, len(vchan.FlushedSegmentIds))
assert.EqualValues(t, 1, vchan.FlushedSegmentIds[0])
assert.EqualValues(t, 2, len(vchan.UnflushedSegmentIds))
assert.ElementsMatch(t, []int64{s2.ID, s3.ID}, vchan.UnflushedSegmentIds)
assert.EqualValues(t, 1, len(vchan.LevelZeroSegmentIds))
assert.ElementsMatch(t, []int64{l0Segment.ID}, vchan.LevelZeroSegmentIds)
})
t.Run("empty collection", func(t *testing.T) {
infos := svr.handler.GetDataVChanPositions(&channelMeta{Name: "ch0_suffix", CollectionID: 1}, allPartitionID)
assert.EqualValues(t, 1, infos.CollectionID)
assert.EqualValues(t, 0, len(infos.FlushedSegmentIds))
assert.EqualValues(t, 0, len(infos.UnflushedSegmentIds))
})
t.Run("filter partition", func(t *testing.T) {
infos := svr.handler.GetDataVChanPositions(&channelMeta{Name: "ch1", CollectionID: 0}, 1)
assert.EqualValues(t, 0, infos.CollectionID)
assert.EqualValues(t, 0, len(infos.FlushedSegmentIds))
assert.EqualValues(t, 1, len(infos.UnflushedSegmentIds))
})
t.Run("empty collection with passed positions", func(t *testing.T) {
vchannel := "ch_no_segment_1"
pchannel := funcutil.ToPhysicalChannel(vchannel)
infos := svr.handler.GetDataVChanPositions(&channelMeta{
Name: vchannel,
CollectionID: 0,
StartPosition: &msgpb.MsgPosition{
ChannelName: pchannel,
MsgID: []byte{14, 15, 16},
},
}, allPartitionID)
assert.EqualValues(t, 0, infos.CollectionID)
assert.EqualValues(t, vchannel, infos.ChannelName)
})
}
func TestGetSnapshotSeekPositions_ReturnsSortedPositionsAndMinTs(t *testing.T) {
handler := &ServerHandler{s: &Server{}}
mockChannels := mockey.Mock((*Server).getChannelsByCollectionID).To(func(s *Server, ctx context.Context, collectionID int64) ([]RWChannel, error) {
assert.Equal(t, UniqueID(100), collectionID)
return []RWChannel{
&channelMeta{Name: "ch-b", CollectionID: 100},
&channelMeta{Name: "ch-a", CollectionID: 100},
}, nil
}).Build()
defer mockChannels.UnPatch()
chAPosition := &msgpb.MsgPosition{
Timestamp: 100,
MsgID: []byte{1},
}
chBPosition := &msgpb.MsgPosition{
ChannelName: "ch-b",
Timestamp: 1000,
MsgID: []byte{2},
}
mockSeekPosition := mockey.Mock((*ServerHandler).GetChannelSeekPosition).To(func(h *ServerHandler, channel RWChannel, partitionIDs ...UniqueID) *msgpb.MsgPosition {
assert.Equal(t, []UniqueID{10, 20}, partitionIDs)
switch channel.GetName() {
case "ch-a":
return chAPosition
case "ch-b":
return chBPosition
default:
t.Fatalf("unexpected channel: %s", channel.GetName())
return nil
}
}).Build()
defer mockSeekPosition.UnPatch()
positions, minTs, err := handler.GetSnapshotSeekPositions(context.Background(), 100, 10, 20)
require.NoError(t, err)
assert.Equal(t, uint64(100), minTs)
require.Len(t, positions, 2)
assert.Equal(t, "ch-a", positions[0].GetChannelName())
assert.Equal(t, uint64(100), positions[0].GetTimestamp())
assert.Equal(t, []byte{1}, positions[0].GetMsgID())
assert.Equal(t, "ch-b", positions[1].GetChannelName())
assert.Equal(t, uint64(1000), positions[1].GetTimestamp())
assert.Equal(t, []byte{2}, positions[1].GetMsgID())
positions[0].Timestamp = 9999
positions, minTs, err = handler.GetSnapshotSeekPositions(context.Background(), 100, 10, 20)
require.NoError(t, err)
assert.Equal(t, uint64(100), minTs)
require.Len(t, positions, 2)
assert.Equal(t, "ch-a", positions[0].GetChannelName())
assert.Equal(t, uint64(100), positions[0].GetTimestamp())
}
func TestGetDeltaLogFromCompactTo(t *testing.T) {
t.Run("no compaction children", func(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
// Add a segment with no compaction children
seg := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1001,
CollectionID: 100,
PartitionID: 10,
State: commonpb.SegmentState_Flushed,
})
err := svr.meta.AddSegment(context.TODO(), seg)
require.NoError(t, err)
// Should return empty delta logs
deltaLogs, err := svr.handler.GetDeltaLogFromCompactTo(context.Background(), 1001)
assert.NoError(t, err)
assert.Empty(t, deltaLogs)
})
t.Run("segment not found", func(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
// Query non-existent segment
_, err := svr.handler.GetDeltaLogFromCompactTo(context.Background(), 9999)
assert.Error(t, err)
})
t.Run("single level compaction", func(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
// Create parent segment
parent := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1000,
CollectionID: 100,
PartitionID: 10,
State: commonpb.SegmentState_Dropped,
})
err := svr.meta.AddSegment(context.TODO(), parent)
require.NoError(t, err)
// Create child segment with delta logs
child := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1001,
CollectionID: 100,
PartitionID: 10,
State: commonpb.SegmentState_Flushed,
CompactionFrom: []int64{1000},
Deltalogs: []*datapb.FieldBinlog{
{
FieldID: 0,
Binlogs: []*datapb.Binlog{
{LogID: 1, LogSize: 100},
},
},
},
})
err = svr.meta.AddSegment(context.TODO(), child)
require.NoError(t, err)
// Should return child's delta logs
deltaLogs, err := svr.handler.GetDeltaLogFromCompactTo(context.Background(), 1000)
assert.NoError(t, err)
assert.Len(t, deltaLogs, 1)
assert.Equal(t, int64(0), deltaLogs[0].FieldID)
})
t.Run("multi-level compaction", func(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
// Create compaction chain: seg1 -> seg2 -> seg3
seg1 := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1001,
CollectionID: 100,
PartitionID: 10,
State: commonpb.SegmentState_Dropped,
})
err := svr.meta.AddSegment(context.TODO(), seg1)
require.NoError(t, err)
seg2 := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1002,
CollectionID: 100,
PartitionID: 10,
State: commonpb.SegmentState_Dropped,
CompactionFrom: []int64{1001},
Deltalogs: []*datapb.FieldBinlog{
{
FieldID: 0,
Binlogs: []*datapb.Binlog{
{LogID: 2, LogSize: 100},
},
},
},
})
err = svr.meta.AddSegment(context.TODO(), seg2)
require.NoError(t, err)
seg3 := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1003,
CollectionID: 100,
PartitionID: 10,
State: commonpb.SegmentState_Flushed,
CompactionFrom: []int64{1002},
Deltalogs: []*datapb.FieldBinlog{
{
FieldID: 0,
Binlogs: []*datapb.Binlog{
{LogID: 3, LogSize: 100},
},
},
},
})
err = svr.meta.AddSegment(context.TODO(), seg3)
require.NoError(t, err)
// Should return all delta logs from the compaction chain
deltaLogs, err := svr.handler.GetDeltaLogFromCompactTo(context.Background(), 1001)
assert.NoError(t, err)
assert.Len(t, deltaLogs, 2) // seg2 and seg3 delta logs
})
t.Run("multi-child compaction", func(t *testing.T) {
svr := newTestServer(t)
defer closeTestServer(t, svr)
// Create parent segment
parent := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1000,
CollectionID: 100,
PartitionID: 10,
State: commonpb.SegmentState_Dropped,
})
err := svr.meta.AddSegment(context.TODO(), parent)
require.NoError(t, err)
// Create multiple children from the same parent
child1 := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1001,
CollectionID: 100,
PartitionID: 10,
State: commonpb.SegmentState_Flushed,
CompactionFrom: []int64{1000},
Deltalogs: []*datapb.FieldBinlog{
{
FieldID: 0,
Binlogs: []*datapb.Binlog{
{LogID: 1, LogSize: 100},
},
},
},
})
err = svr.meta.AddSegment(context.TODO(), child1)
require.NoError(t, err)
child2 := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1002,
CollectionID: 100,
PartitionID: 10,
State: commonpb.SegmentState_Flushed,
CompactionFrom: []int64{1000},
Deltalogs: []*datapb.FieldBinlog{
{
FieldID: 0,
Binlogs: []*datapb.Binlog{
{LogID: 2, LogSize: 100},
},
},
},
})
err = svr.meta.AddSegment(context.TODO(), child2)
require.NoError(t, err)
// Should return delta logs from both children
deltaLogs, err := svr.handler.GetDeltaLogFromCompactTo(context.Background(), 1000)
assert.NoError(t, err)
assert.Len(t, deltaLogs, 2)
})
}
func TestGenSnapshot(t *testing.T) {
schema := newTestSchema()
// Create minimal components for testing
mockIndexMeta := &indexMeta{}
mockMeta := &meta{
indexMeta: mockIndexMeta,
}
// Create mock mixCoord
mockMixCoord := mocks2.NewMixCoord(t)
mockMixCoord.EXPECT().DescribeCollectionInternal(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
Schema: schema,
ShardsNum: 2,
NumPartitions: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
Properties: []*commonpb.KeyValuePair{},
CollectionID: 200,
VirtualChannelNames: []string{"dml_0_200v0", "dml_1_200v1"},
}, nil).Maybe()
mockMixCoord.EXPECT().ShowPartitionsInternal(mock.Anything, mock.Anything).Return(&milvuspb.ShowPartitionsResponse{
Status: merr.Success(),
PartitionIDs: []int64{0, 1},
PartitionNames: []string{"_default", "partition1"},
}, nil).Maybe()
// Create handler with mock broker
mockBroker := broker.NewCoordinatorBroker(mockMixCoord)
handler := &ServerHandler{
s: &Server{
broker: mockBroker,
meta: mockMeta,
},
}
// Setup mocks for other methods
mock1 := mockey.Mock((*ServerHandler).GetSnapshotSeekPositions).To(func(h *ServerHandler, ctx context.Context, collectionID UniqueID, partitionIDs ...UniqueID) ([]*msgpb.MsgPosition, uint64, error) {
return []*msgpb.MsgPosition{
{ChannelName: "dml_0_200v0", Timestamp: 12345, MsgID: []byte{1}},
}, uint64(12345), nil
}).Build()
defer mock1.UnPatch()
mock4 := mockey.Mock((*indexMeta).GetIndexesForCollection).To(func(im *indexMeta, collectionID UniqueID, fieldName string) []*model.Index {
return []*model.Index{
{
IndexID: 1,
CollectionID: 200,
FieldID: 2,
IndexName: "test_index",
CreateTime: 10000, // Earlier than snapshot timestamp
},
}
}).Build()
defer mock4.UnPatch()
mock5 := mockey.Mock((*meta).SelectSegments).To(func(m *meta, ctx context.Context, filters ...SegmentFilter) []*SegmentInfo {
seg := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1001,
CollectionID: 200,
PartitionID: 0,
InsertChannel: "dml_0_200v0",
State: commonpb.SegmentState_Flushed,
StartPosition: &msgpb.MsgPosition{Timestamp: 10000},
Binlogs: []*datapb.FieldBinlog{
{
FieldID: 1,
Binlogs: []*datapb.Binlog{
{LogID: 1, LogSize: 100},
},
},
},
})
return []*SegmentInfo{seg}
}).Build()
defer mock5.UnPatch()
mock6 := mockey.Mock((*meta).GetSegment).To(func(m *meta, ctx context.Context, segmentID int64) *SegmentInfo {
if segmentID == 1001 {
return NewSegmentInfo(&datapb.SegmentInfo{
ID: 1001,
CollectionID: 200,
PartitionID: 0,
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
FieldID: 1,
Binlogs: []*datapb.Binlog{
{LogID: 1, LogSize: 100},
},
},
},
Deltalogs: []*datapb.FieldBinlog{},
})
}
return nil
}).Build()
defer mock6.UnPatch()
mock7 := mockey.Mock((*indexMeta).getSegmentIndexes).To(func(im *indexMeta, collectionID, segmentID int64) map[int64]*model.SegmentIndex {
return map[int64]*model.SegmentIndex{}
}).Build()
defer mock7.UnPatch()
mock8 := mockey.Mock((*meta).GetCompactionTo).To(func(m *meta, segmentID int64) ([]*SegmentInfo, bool) {
return nil, false // No compaction children
}).Build()
defer mock8.UnPatch()
// Test GenSnapshot
snapshotData, err := handler.GenSnapshot(context.Background(), 200)
assert.NoError(t, err)
assert.NotNil(t, snapshotData)
assert.Equal(t, int64(200), snapshotData.SnapshotInfo.CollectionId)
assert.Equal(t, int64(12345), snapshotData.SnapshotInfo.CreateTs)
assert.NotNil(t, snapshotData.Collection)
assert.Equal(t, int64(2), snapshotData.Collection.NumShards)
assert.Equal(t, 1, len(snapshotData.Indexes))
assert.Equal(t, 1, len(snapshotData.Segments))
assert.Equal(t, int64(1001), snapshotData.Segments[0].SegmentId)
// Verify VirtualChannelNames is populated from DescribeCollectionInternal response
assert.Equal(t, []string{"dml_0_200v0", "dml_1_200v1"}, snapshotData.Collection.VirtualChannelNames)
}
func TestGenSnapshot_UsesPerChannelSeekPositions(t *testing.T) {
schema := newTestSchema()
segCh1 := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1001,
CollectionID: 200,
PartitionID: 0,
InsertChannel: "ch-1",
State: commonpb.SegmentState_Flushed,
StartPosition: &msgpb.MsgPosition{ChannelName: "ch-1", Timestamp: 500},
DmlPosition: &msgpb.MsgPosition{ChannelName: "ch-1", Timestamp: 600},
Binlogs: []*datapb.FieldBinlog{
{FieldID: 1, Binlogs: []*datapb.Binlog{{LogID: 1, LogSize: 100}}},
},
})
segCh2 := NewSegmentInfo(&datapb.SegmentInfo{
ID: 1002,
CollectionID: 200,
PartitionID: 0,
InsertChannel: "ch-2",
State: commonpb.SegmentState_Flushed,
StartPosition: &msgpb.MsgPosition{ChannelName: "ch-2", Timestamp: 500},
DmlPosition: &msgpb.MsgPosition{ChannelName: "ch-2", Timestamp: 600},
Binlogs: []*datapb.FieldBinlog{
{FieldID: 1, Binlogs: []*datapb.Binlog{{LogID: 2, LogSize: 100}}},
},
})
mockMeta := &meta{indexMeta: &indexMeta{}}
handler := &ServerHandler{
s: &Server{
broker: broker.NewCoordinatorBroker(newMockMixCoord()),
meta: mockMeta,
},
}
mockDescribe := mockey.Mock((*mockMixCoord).DescribeCollectionInternal).To(
func(m *mockMixCoord, ctx context.Context, req *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
Schema: schema,
ShardsNum: 2,
NumPartitions: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
CollectionID: 200,
VirtualChannelNames: []string{"ch-1", "ch-2"},
}, nil
}).Build()
defer mockDescribe.UnPatch()
mockShowPartitions := mockey.Mock((*mockMixCoord).ShowPartitionsInternal).To(
func(m *mockMixCoord, ctx context.Context, req *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
return &milvuspb.ShowPartitionsResponse{
Status: merr.Success(),
PartitionIDs: []int64{0},
PartitionNames: []string{"_default"},
}, nil
}).Build()
defer mockShowPartitions.UnPatch()
mockSeekPositions := mockey.Mock((*ServerHandler).GetSnapshotSeekPositions).To(
func(h *ServerHandler, ctx context.Context, collectionID UniqueID, partitionIDs ...UniqueID) ([]*msgpb.MsgPosition, uint64, error) {
return []*msgpb.MsgPosition{
{ChannelName: "ch-1", Timestamp: 100, MsgID: []byte{1}},
{ChannelName: "ch-2", Timestamp: 1000, MsgID: []byte{2}},
}, uint64(100), nil
}).Build()
defer mockSeekPositions.UnPatch()
mockIndexes := mockey.Mock((*indexMeta).GetIndexesForCollection).To(
func(im *indexMeta, collectionID UniqueID, fieldName string) []*model.Index {
return []*model.Index{}
}).Build()
defer mockIndexes.UnPatch()
mockSelectSegments := mockey.Mock((*meta).SelectSegments).To(
func(m *meta, ctx context.Context, filters ...SegmentFilter) []*SegmentInfo {
candidates := []*SegmentInfo{segCh1, segCh2}
var result []*SegmentInfo
for _, seg := range candidates {
pass := true
for _, filter := range filters {
if !filter.Match(seg) {
pass = false
break
}
}
if pass {
result = append(result, seg)
}
}
return result
}).Build()
defer mockSelectSegments.UnPatch()
mockCompactionTo := mockey.Mock((*meta).GetCompactionTo).To(func(m *meta, segmentID int64) ([]*SegmentInfo, bool) {
return nil, false
}).Build()
defer mockCompactionTo.UnPatch()
mockSegmentIndexes := mockey.Mock((*indexMeta).getSegmentIndexes).To(
func(im *indexMeta, collectionID, segmentID int64) map[int64]*model.SegmentIndex {
return map[int64]*model.SegmentIndex{}
}).Build()
defer mockSegmentIndexes.UnPatch()
snapshotData, err := handler.GenSnapshot(context.Background(), 200)
require.NoError(t, err)
require.NotNil(t, snapshotData)
require.NotNil(t, snapshotData.SnapshotInfo)
assert.Equal(t, int64(100), snapshotData.SnapshotInfo.CreateTs)
require.Len(t, snapshotData.SnapshotInfo.GetChannelSeekPositions(), 2)
assert.Equal(t, "ch-1", snapshotData.SnapshotInfo.GetChannelSeekPositions()[0].GetChannelName())
assert.Equal(t, uint64(100), snapshotData.SnapshotInfo.GetChannelSeekPositions()[0].GetTimestamp())
assert.Equal(t, "ch-2", snapshotData.SnapshotInfo.GetChannelSeekPositions()[1].GetChannelName())
assert.Equal(t, uint64(1000), snapshotData.SnapshotInfo.GetChannelSeekPositions()[1].GetTimestamp())
require.Len(t, snapshotData.Segments, 1)
assert.Equal(t, int64(1002), snapshotData.Segments[0].GetSegmentId())
assert.Equal(t, "ch-2", snapshotData.Segments[0].GetChannelName())
}
func TestGenSnapshot_RejectsSegmentWithoutChannelSeekPosition(t *testing.T) {
schema := newTestSchema()
var seg *SegmentInfo
mockMeta := &meta{indexMeta: &indexMeta{}}
handler := &ServerHandler{
s: &Server{
broker: broker.NewCoordinatorBroker(newMockMixCoord()),
meta: mockMeta,
},
}
mockDescribe := mockey.Mock((*mockMixCoord).DescribeCollectionInternal).To(
func(m *mockMixCoord, ctx context.Context, req *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
Schema: schema,
ShardsNum: 1,
NumPartitions: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
CollectionID: 200,
VirtualChannelNames: []string{"ch-1"},
}, nil
}).Build()
defer mockDescribe.UnPatch()
mockShowPartitions := mockey.Mock((*mockMixCoord).ShowPartitionsInternal).To(
func(m *mockMixCoord, ctx context.Context, req *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
return &milvuspb.ShowPartitionsResponse{
Status: merr.Success(),
PartitionIDs: []int64{0},
PartitionNames: []string{"_default"},
}, nil
}).Build()
defer mockShowPartitions.UnPatch()
mockSeekPositions := mockey.Mock((*ServerHandler).GetSnapshotSeekPositions).To(
func(h *ServerHandler, ctx context.Context, collectionID UniqueID, partitionIDs ...UniqueID) ([]*msgpb.MsgPosition, uint64, error) {
return []*msgpb.MsgPosition{
{ChannelName: "ch-1", Timestamp: 1000, MsgID: []byte{1}},
}, uint64(1000), nil
}).Build()
defer mockSeekPositions.UnPatch()
mockIndexes := mockey.Mock((*indexMeta).GetIndexesForCollection).To(
func(im *indexMeta, collectionID UniqueID, fieldName string) []*model.Index {
return []*model.Index{}
}).Build()
defer mockIndexes.UnPatch()
mockSelectSegments := mockey.Mock((*meta).SelectSegments).To(
func(m *meta, ctx context.Context, filters ...SegmentFilter) []*SegmentInfo {
for _, filter := range filters {
if !filter.Match(seg) {
return nil
}
}
return []*SegmentInfo{seg}
}).Build()
defer mockSelectSegments.UnPatch()
for _, test := range []struct {
name string
insertChannel string
}{
{name: "unknown channel", insertChannel: "ch-missing"},
{name: "empty channel", insertChannel: ""},
} {
t.Run(test.name, func(t *testing.T) {
seg = NewSegmentInfo(&datapb.SegmentInfo{
ID: 1001,
CollectionID: 200,
PartitionID: 0,
InsertChannel: test.insertChannel,
State: commonpb.SegmentState_Flushed,
StartPosition: &msgpb.MsgPosition{ChannelName: test.insertChannel, Timestamp: 500},
DmlPosition: &msgpb.MsgPosition{ChannelName: test.insertChannel, Timestamp: 600},
Binlogs: []*datapb.FieldBinlog{
{FieldID: 1, Binlogs: []*datapb.Binlog{{LogID: 1, LogSize: 100}}},
},
})
snapshotData, err := handler.GenSnapshot(context.Background(), 200)
require.Error(t, err)
assert.Nil(t, snapshotData)
assert.Contains(t, err.Error(), "missing snapshot channel seek position")
})
}
}
func TestUncompressIndexFilesPreservesIndexStorePathVersion(t *testing.T) {
handler := &ServerHandler{
s: &Server{
meta: &meta{
chunkManager: storage.NewLocalChunkManager(objectstorage.RootPath("files")),
indexMeta: &indexMeta{},
},
},
}
mockGetSegmentIndexes := mockey.Mock((*indexMeta).getSegmentIndexes).To(func(im *indexMeta, collectionID, segmentID int64) map[int64]*model.SegmentIndex {
return map[int64]*model.SegmentIndex{
10: {
CollectionID: collectionID,
PartitionID: 20,
SegmentID: segmentID,
IndexID: 10,
BuildID: 30,
IndexVersion: 40,
IndexState: commonpb.IndexState_Finished,
IndexFileKeys: []string{"index_data"},
IndexStorePathVersion: indexpb.IndexStorePathVersion_INDEX_STORE_PATH_VERSION_COLLECTION_ROOTED,
},
}
}).Build()
defer mockGetSegmentIndexes.UnPatch()
mockFieldID := mockey.Mock((*indexMeta).GetFieldIDByIndexID).Return(int64(100)).Build()
defer mockFieldID.UnPatch()
mockIndexName := mockey.Mock((*indexMeta).GetIndexNameByID).Return("vec_idx").Build()
defer mockIndexName.UnPatch()
mockIndexParams := mockey.Mock((*indexMeta).GetIndexParams).Return(nil).Build()
defer mockIndexParams.UnPatch()
mockTypeParams := mockey.Mock((*indexMeta).GetTypeParams).Return(nil).Build()
defer mockTypeParams.UnPatch()
indexFiles := uncompressIndexFiles(handler, 1, 2)
require.Len(t, indexFiles, 1)
assert.Equal(t, indexpb.IndexStorePathVersion_INDEX_STORE_PATH_VERSION_COLLECTION_ROOTED, indexFiles[0].GetIndexStorePathVersion())
assert.Equal(t, []string{"files/index_v1/1/20/2/30/40/index_data"}, indexFiles[0].GetIndexFilePaths())
}
func TestUncompressJsonStatsForSnapshotPaths(t *testing.T) {
handler := &ServerHandler{
s: &Server{
meta: &meta{
chunkManager: storage.NewLocalChunkManager(objectstorage.RootPath("files")),
},
},
}
t.Run("storage v2 path", func(t *testing.T) {
segInfo := &datapb.SegmentInfo{
CollectionID: 1,
PartitionID: 2,
ID: 3,
}
stats := &datapb.JsonKeyStats{
FieldID: 100,
BuildID: 10,
Version: 2,
JsonKeyStatsDataFormat: 3,
Files: []string{"meta.json", "shared_key_index/.managed.json_0"},
}
got := uncompressJSONStats(handler, segInfo, stats)
assert.Equal(t, []string{
"files/json_stats/3/10/2/1/2/3/100/meta.json",
"files/json_stats/3/10/2/1/2/3/100/shared_key_index/.managed.json_0",
}, got.GetFiles())
assert.Equal(t, []string{"meta.json", "shared_key_index/.managed.json_0"}, stats.GetFiles())
})
t.Run("storage v3 manifest path", func(t *testing.T) {
segInfo := &datapb.SegmentInfo{
CollectionID: 1,
PartitionID: 2,
ID: 3,
ManifestPath: `{"ver":7,"base_path":"files/insert_log/1/2/3"}`,
}
stats := &datapb.JsonKeyStats{
FieldID: 100,
Files: []string{"meta.json", "shredding_data/data.parquet"},
}
got := uncompressJSONStats(handler, segInfo, stats)
assert.Equal(t, []string{
"files/insert_log/1/2/3/_stats/json_stats.100/meta.json",
"files/insert_log/1/2/3/_stats/json_stats.100/shredding_data/data.parquet",
}, got.GetFiles())
assert.Equal(t, []string{"meta.json", "shredding_data/data.parquet"}, stats.GetFiles())
})
}
// TestGenSnapshot_CommitTimestamp verifies that import segments with a
// commit_timestamp are correctly included/excluded based on snapshotTs.
func TestGenSnapshot_CommitTimestamp(t *testing.T) {
schema := newTestSchema()
// importSeg has start_position.ts=1000 and commit_timestamp=5000.
// Its effective timestamp is 5000.
importSeg := NewSegmentInfo(&datapb.SegmentInfo{
ID: 999,
CollectionID: 200,
PartitionID: 0,
InsertChannel: "dml_0_200v0",
State: commonpb.SegmentState_Flushed,
CommitTimestamp: 5000,
StartPosition: &msgpb.MsgPosition{Timestamp: 1000},
Binlogs: []*datapb.FieldBinlog{
{FieldID: 1, Binlogs: []*datapb.Binlog{{LogID: 1, LogSize: 100}}},
},
})
setupHandlerMocks := func(t *testing.T, snapshotTs uint64) (*ServerHandler, []*mockey.Mocker) {
mockIndexMeta := &indexMeta{}
mockMeta := &meta{indexMeta: mockIndexMeta}
mockMixCoord := mocks2.NewMixCoord(t)
mockMixCoord.EXPECT().DescribeCollectionInternal(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
Schema: schema,
ShardsNum: 2,
NumPartitions: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
Properties: []*commonpb.KeyValuePair{},
CollectionID: 200,
VirtualChannelNames: []string{"dml_0_200v0"},
}, nil).Maybe()
mockMixCoord.EXPECT().ShowPartitionsInternal(mock.Anything, mock.Anything).Return(&milvuspb.ShowPartitionsResponse{
Status: merr.Success(),
PartitionIDs: []int64{0},
PartitionNames: []string{"_default"},
}, nil).Maybe()
mockBroker := broker.NewCoordinatorBroker(mockMixCoord)
handler := &ServerHandler{
s: &Server{broker: mockBroker, meta: mockMeta},
}
var mockers []*mockey.Mocker
m1 := mockey.Mock((*ServerHandler).GetSnapshotSeekPositions).To(func(h *ServerHandler, ctx context.Context, collectionID UniqueID, partitionIDs ...UniqueID) ([]*msgpb.MsgPosition, uint64, error) {
return []*msgpb.MsgPosition{
{ChannelName: "dml_0_200v0", Timestamp: snapshotTs, MsgID: []byte{1}},
}, snapshotTs, nil
}).Build()
mockers = append(mockers, m1)
m2 := mockey.Mock((*indexMeta).GetIndexesForCollection).To(func(im *indexMeta, collectionID UniqueID, fieldName string) []*model.Index {
return []*model.Index{}
}).Build()
mockers = append(mockers, m2)
m3 := mockey.Mock((*meta).SelectSegments).To(func(m *meta, ctx context.Context, filters ...SegmentFilter) []*SegmentInfo {
// Apply each filter to the candidate segment list.
candidates := []*SegmentInfo{importSeg}
var result []*SegmentInfo
for _, seg := range candidates {
pass := true
for _, f := range filters {
if !f.Match(seg) {
pass = false
break
}
}
if pass {
result = append(result, seg)
}
}
return result
}).Build()
mockers = append(mockers, m3)
m4 := mockey.Mock((*meta).GetCompactionTo).To(func(m *meta, segmentID int64) ([]*SegmentInfo, bool) {
return nil, false
}).Build()
mockers = append(mockers, m4)
m5 := mockey.Mock((*indexMeta).getSegmentIndexes).To(func(im *indexMeta, collectionID, segmentID int64) map[int64]*model.SegmentIndex {
return map[int64]*model.SegmentIndex{}
}).Build()
mockers = append(mockers, m5)
m6 := mockey.Mock((*meta).GetSegment).To(func(m *meta, ctx context.Context, segmentID int64) *SegmentInfo {
return nil
}).Build()
mockers = append(mockers, m6)
return handler, mockers
}
t.Run("import segment excluded when snapshotTs < commit_timestamp", func(t *testing.T) {
// snapshotTs=3000 < commit_ts=5000 → segment must NOT appear
handler, mockers := setupHandlerMocks(t, 3000)
for _, m := range mockers {
defer m.UnPatch()
}
snapshotData, err := handler.GenSnapshot(context.Background(), 200)
assert.NoError(t, err)
assert.NotNil(t, snapshotData)
assert.Equal(t, 0, len(snapshotData.Segments), "segment with commit_ts=5000 must not appear at snapshotTs=3000")
})
t.Run("import segment included when snapshotTs >= commit_timestamp", func(t *testing.T) {
// snapshotTs=6000 >= commit_ts=5000 → segment must appear
handler, mockers := setupHandlerMocks(t, 6000)
for _, m := range mockers {
defer m.UnPatch()
}
snapshotData, err := handler.GenSnapshot(context.Background(), 200)
assert.NoError(t, err)
assert.NotNil(t, snapshotData)
assert.Equal(t, 1, len(snapshotData.Segments), "segment with commit_ts=5000 must appear at snapshotTs=6000")
assert.Equal(t, int64(999), snapshotData.Segments[0].SegmentId)
})
}