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

3617 lines
102 KiB
Go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package datacoord
import (
"context"
"reflect"
"sort"
satomic "sync/atomic"
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"go.uber.org/atomic"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/compaction"
"github.com/milvus-io/milvus/internal/datacoord/allocator"
"github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/util/lifetime"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
type spyCompactionInspector struct {
t *testing.T
spyChan chan *datapb.CompactionPlan
meta *meta
}
// getCompactionTasksNum implements CompactionInspector.
func (h *spyCompactionInspector) getCompactionTasksNum(filters ...compactionTaskFilter) int {
return 0
}
func (h *spyCompactionInspector) getCompactionTasksNumBySignalID(signalID int64) int {
return 0
}
func (h *spyCompactionInspector) getCompactionInfo(ctx context.Context, signalID int64) *compactionInfo {
return nil
}
var _ CompactionInspector = (*spyCompactionInspector)(nil)
func (h *spyCompactionInspector) removeTasksByChannel(channel string) {}
// enqueueCompaction start to execute plan and return immediately
func (h *spyCompactionInspector) enqueueCompaction(task *datapb.CompactionTask) error {
t := newMixCompactionTask(task, nil, h.meta, newMockVersionManager())
alloc := newMock0Allocator(h.t)
t.allocator = alloc
plan, err := t.BuildCompactionRequest()
h.spyChan <- plan
return err
}
func (h *spyCompactionInspector) checkAndSetSegmentStating(channel string, segmentID int64) bool {
return false
}
// isFull return true if the task pool is full
func (h *spyCompactionInspector) isFull() bool {
return false
}
func (h *spyCompactionInspector) start() {}
func (h *spyCompactionInspector) stop() {}
func newMockVersionManager() IndexEngineVersionManager {
return &versionManagerImpl{}
}
var _ CompactionInspector = (*spyCompactionInspector)(nil)
func Test_compactionTrigger_force_without_index(t *testing.T) {
catalog := mocks.NewDataCoordCatalog(t)
catalog.EXPECT().AlterSegments(mock.Anything, mock.Anything).Return(nil).Maybe()
collectionID := int64(11)
binlogs := []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 1},
},
},
}
deltaLogs := []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 1},
},
},
}
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: 101,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
}
segInfo := &datapb.SegmentInfo{
ID: 1,
CollectionID: collectionID,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: binlogs,
Deltalogs: deltaLogs,
IsSorted: true,
}
collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()
collections.Insert(collectionID, &collectionInfo{
ID: collectionID,
Schema: schema,
})
m := &meta{
catalog: catalog,
channelCPs: newChannelCps(),
segments: &SegmentsInfo{
segments: map[int64]*SegmentInfo{
1: {
SegmentInfo: segInfo,
},
},
secondaryIndexes: segmentInfoIndexes{
coll2Segments: map[UniqueID]map[UniqueID]*SegmentInfo{
collectionID: {
1: {
SegmentInfo: segInfo,
},
},
},
},
},
indexMeta: &indexMeta{
segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
indexes: map[UniqueID]map[UniqueID]*model.Index{},
},
collections: collections,
}
inspector := &spyCompactionInspector{t: t, spyChan: make(chan *datapb.CompactionPlan, 1), meta: m}
tr := &compactionTrigger{
meta: m,
handler: newMockHandlerWithMeta(m),
allocator: newMock0Allocator(t),
signals: make(chan *compactionSignal, 100),
manualSignals: make(chan *compactionSignal, 100),
inspector: inspector,
globalTrigger: nil,
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
tr.closeWaiter.Add(1)
go func() {
defer tr.closeWaiter.Done()
tr.work()
}()
defer tr.stop()
_, err := tr.TriggerCompaction(context.TODO(), NewCompactionSignal().WithCollectionID(collectionID).WithIsForce(true))
assert.NoError(t, err)
select {
case val := <-inspector.spyChan:
assert.Equal(t, 1, len(val.SegmentBinlogs))
return
case <-time.After(3 * time.Second):
assert.Fail(t, "failed to get plan")
return
}
}
func Test_compactionTrigger_force(t *testing.T) {
paramtable.Init()
type fields struct {
meta *meta
allocator allocator.Allocator
signals chan *compactionSignal
inspector CompactionInspector
globalTrigger *time.Ticker
}
catalog := mocks.NewDataCoordCatalog(t)
catalog.EXPECT().AlterSegments(mock.Anything, mock.Anything).Return(nil).Maybe()
vecFieldID := int64(201)
indexID := int64(1001)
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
}
mock0Allocator := newMock0Allocator(t)
seg1 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 1},
},
},
},
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 1},
},
},
},
IsSorted: true,
},
}
seg2 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 2,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 2},
},
},
},
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 2},
},
},
},
IsSorted: true,
},
}
seg3 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 3,
CollectionID: 1111,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
IsSorted: true,
},
}
collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()
collections.Insert(2, &collectionInfo{
ID: 2,
Schema: schema,
Properties: map[string]string{
common.CollectionTTLConfigKey: "0",
},
})
collections.Insert(1111, &collectionInfo{
ID: 1111,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
},
Properties: map[string]string{
common.CollectionTTLConfigKey: "error",
},
})
collections.Insert(1000, &collectionInfo{
ID: 1000,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
},
})
// error (has no vector field)
collections.Insert(2000, &collectionInfo{
ID: 2000,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_Int16,
},
},
},
})
// error (has no dim)
collections.Insert(3000, &collectionInfo{
ID: 3000,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{},
},
},
},
},
})
// error (dim parse fail)
collections.Insert(4000, &collectionInfo{
ID: 4000,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128error",
},
},
},
},
},
})
collections.Insert(10000, &collectionInfo{
ID: 10000,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
},
})
im := &indexMeta{
segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
indexes: map[UniqueID]map[UniqueID]*model.Index{
2: {
indexID: {
TenantID: "",
CollectionID: 2,
FieldID: vecFieldID,
IndexID: indexID,
IndexName: "_default_idx",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "HNSW",
},
},
IsAutoIndex: false,
UserIndexParams: nil,
},
},
1000: {
indexID: {
TenantID: "",
CollectionID: 1000,
FieldID: vecFieldID,
IndexID: indexID,
IndexName: "_default_idx",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "DISKANN",
},
},
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
}
segIdx1 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx1.Insert(indexID, &model.SegmentIndex{
SegmentID: 1,
CollectionID: 2,
PartitionID: 1,
NumRows: 100,
IndexID: indexID,
BuildID: 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx2 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx2.Insert(indexID, &model.SegmentIndex{
SegmentID: 2,
CollectionID: 2,
PartitionID: 1,
NumRows: 100,
IndexID: indexID,
BuildID: 2,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx3 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx3.Insert(indexID, &model.SegmentIndex{
SegmentID: 3,
CollectionID: 1111,
PartitionID: 1,
NumRows: 100,
IndexID: indexID,
BuildID: 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
im.segmentIndexes.Insert(1, segIdx1)
im.segmentIndexes.Insert(2, segIdx2)
im.segmentIndexes.Insert(3, segIdx3)
params, err := compaction.GenerateJSONParams(nil)
if err != nil {
panic(err)
}
preAllocateIDExpansionFactor := paramtable.Get().DataCoordCfg.CompactionPreAllocateIDExpansionFactor.GetAsInt64()
preAllocatedSegmentIDBegin := int64(100)
preAllocatedSegmentIDEnd := preAllocatedSegmentIDBegin + preAllocateIDExpansionFactor
preAllocatedLogIDEnd := preAllocatedSegmentIDBegin + 4*preAllocateIDExpansionFactor
tests := []struct {
name string
fields fields
collectionID UniqueID
wantErr bool
wantSegIDs []int64
wantPlans []*datapb.CompactionPlan
}{
{
"test force compaction",
fields{
&meta{
catalog: catalog,
channelCPs: newChannelCps(),
segments: &SegmentsInfo{
segments: map[int64]*SegmentInfo{
1: seg1,
2: seg2,
3: seg3,
},
secondaryIndexes: segmentInfoIndexes{
coll2Segments: map[UniqueID]map[UniqueID]*SegmentInfo{
2: {
seg1.GetID(): seg1,
seg2.GetID(): seg2,
},
1111: {
seg3.GetID(): seg3,
},
},
},
},
indexMeta: im,
collections: collections,
},
mock0Allocator,
nil,
&spyCompactionInspector{t: t, spyChan: make(chan *datapb.CompactionPlan, 1)},
nil,
},
2,
false,
[]int64{
1, 2,
},
[]*datapb.CompactionPlan{
{
PlanID: preAllocatedSegmentIDEnd,
SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{
{
SegmentID: 1,
FieldBinlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 1},
},
},
},
Field2StatslogPaths: nil,
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 1},
},
},
},
InsertChannel: "ch1",
CollectionID: 2,
PartitionID: 1,
IsSorted: true,
},
{
SegmentID: 2,
FieldBinlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 2},
},
},
},
Field2StatslogPaths: nil,
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 2},
},
},
},
InsertChannel: "ch1",
CollectionID: 2,
PartitionID: 1,
IsSorted: true,
},
},
// StartTime: 0,
BeginLogID: 100,
Type: datapb.CompactionType_MixCompaction,
Channel: "ch1",
TotalRows: 200,
Schema: schema,
PreAllocatedSegmentIDs: &datapb.IDRange{Begin: preAllocatedSegmentIDBegin, End: preAllocatedSegmentIDEnd},
PreAllocatedLogIDs: &datapb.IDRange{Begin: preAllocatedSegmentIDBegin, End: preAllocatedLogIDEnd},
MaxSize: 1073741824,
SlotUsage: paramtable.Get().DataCoordCfg.MixCompactionSlotUsage.GetAsInt64(),
JsonParams: params,
},
},
},
}
for _, tt := range tests {
tt.fields.inspector.(*spyCompactionInspector).meta = tt.fields.meta
t.Run(tt.name, func(t *testing.T) {
tr := &compactionTrigger{
meta: tt.fields.meta,
handler: newMockHandlerWithMeta(tt.fields.meta),
allocator: tt.fields.allocator,
signals: make(chan *compactionSignal, 100),
manualSignals: make(chan *compactionSignal, 100),
inspector: tt.fields.inspector,
globalTrigger: tt.fields.globalTrigger,
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
tr.closeWaiter.Add(1)
go func() {
defer tr.closeWaiter.Done()
tr.work()
}()
defer tr.stop()
_, err := tr.TriggerCompaction(context.TODO(), NewCompactionSignal().WithCollectionID(tt.collectionID).WithIsForce(true))
assert.Equal(t, tt.wantErr, err != nil)
spy := (tt.fields.inspector).(*spyCompactionInspector)
select {
case plan := <-spy.spyChan:
plan.StartTime = 0
sortPlanCompactionBinlogs(plan)
assert.True(t, proto.Equal(tt.wantPlans[0], plan))
return
case <-time.After(3 * time.Second):
assert.Fail(t, "timeout")
return
}
})
t.Run(tt.name+" with DiskANN index", func(t *testing.T) {
for _, segment := range tt.fields.meta.segments.GetSegments() {
// Collection 1000 means it has DiskANN index
delete(tt.fields.meta.segments.secondaryIndexes.coll2Segments[segment.GetCollectionID()], segment.GetID())
segment.CollectionID = 1000
_, ok := tt.fields.meta.segments.secondaryIndexes.coll2Segments[segment.GetCollectionID()]
if !ok {
tt.fields.meta.segments.secondaryIndexes.coll2Segments[segment.GetCollectionID()] = make(map[UniqueID]*SegmentInfo)
}
tt.fields.meta.segments.secondaryIndexes.coll2Segments[segment.GetCollectionID()][segment.GetID()] = segment
}
tr := &compactionTrigger{
meta: tt.fields.meta,
handler: newMockHandlerWithMeta(tt.fields.meta),
allocator: tt.fields.allocator,
signals: make(chan *compactionSignal, 100),
manualSignals: make(chan *compactionSignal, 100),
inspector: tt.fields.inspector,
globalTrigger: tt.fields.globalTrigger,
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
tr.closeWaiter.Add(1)
go func() {
defer tr.closeWaiter.Done()
tr.work()
}()
defer tr.stop()
tt.collectionID = 1000
_, err := tr.TriggerCompaction(context.TODO(), NewCompactionSignal().WithCollectionID(tt.collectionID).WithIsForce(true))
assert.Equal(t, tt.wantErr, err != nil)
// expect max row num = 2048*1024*1024/(128*4) = 4194304
// assert.EqualValues(t, 4194304, tt.fields.meta.segments.GetSegments()[0].MaxRowNum)
spy := (tt.fields.inspector).(*spyCompactionInspector)
select {
case plan := <-spy.spyChan:
assert.NotNil(t, plan)
return
case <-time.After(3 * time.Second):
assert.Fail(t, "timeout")
return
}
})
t.Run(tt.name+" with getCompact error", func(t *testing.T) {
for _, segment := range tt.fields.meta.segments.GetSegments() {
segment.CollectionID = 1111
}
tr := &compactionTrigger{
meta: tt.fields.meta,
handler: newMockHandlerWithMeta(tt.fields.meta),
allocator: tt.fields.allocator,
signals: make(chan *compactionSignal, 100),
manualSignals: make(chan *compactionSignal, 100),
inspector: tt.fields.inspector,
globalTrigger: tt.fields.globalTrigger,
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
tr.closeWaiter.Add(1)
go func() {
defer tr.closeWaiter.Done()
tr.work()
}()
defer tr.stop()
{
// test getCompactTime fail for handle global signal
signal := NewCompactionSignal().WithCollectionID(1111).WithIsForce(true)
tr.TriggerCompaction(context.TODO(), signal)
spy := (tt.fields.inspector).(*spyCompactionInspector)
hasPlan := true
select {
case <-spy.spyChan:
hasPlan = true
case <-time.After(2 * time.Second):
hasPlan = false
}
assert.Equal(t, false, hasPlan)
}
{
// test getCompactTime fail for handle signal
signal := NewCompactionSignal().WithCollectionID(1111).WithIsForce(true).WithSegmentIDs(3)
tr.TriggerCompaction(context.TODO(), signal)
spy := (tt.fields.inspector).(*spyCompactionInspector)
hasPlan := true
select {
case <-spy.spyChan:
hasPlan = true
case <-time.After(2 * time.Second):
hasPlan = false
}
assert.Equal(t, false, hasPlan)
}
})
}
}
// test force compaction with too many Segment
func Test_compactionTrigger_force_maxSegmentLimit(t *testing.T) {
type fields struct {
meta *meta
allocator allocator.Allocator
signals chan *compactionSignal
inspector CompactionInspector
globalTrigger *time.Ticker
}
type args struct {
collectionID int64
compactTime *compactTime
}
vecFieldID := int64(201)
segmentInfos := &SegmentsInfo{
segments: make(map[UniqueID]*SegmentInfo),
secondaryIndexes: segmentInfoIndexes{
coll2Segments: make(map[UniqueID]map[UniqueID]*SegmentInfo),
},
}
indexMeta := newSegmentIndexMeta(nil)
indexMeta.indexes = map[UniqueID]map[UniqueID]*model.Index{
2: {
indexID: {
TenantID: "",
CollectionID: 2,
FieldID: vecFieldID,
IndexID: indexID,
IndexName: "_default_idx",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "HNSW",
},
},
IsAutoIndex: false,
UserIndexParams: nil,
},
},
}
segmentInfos.secondaryIndexes.coll2Segments[2] = make(map[UniqueID]*SegmentInfo)
nSegments := 50
for i := UniqueID(0); i < UniqueID(nSegments); i++ {
info := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: i,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300000,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log1"},
},
},
},
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "deltalog1"},
},
},
},
IsSorted: true,
},
}
indexMeta.updateSegmentIndex(&model.SegmentIndex{
SegmentID: i,
CollectionID: 2,
PartitionID: 1,
NumRows: 100,
IndexID: indexID,
BuildID: i,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
})
segmentInfos.segments[i] = info
segmentInfos.secondaryIndexes.coll2Segments[2][i] = info
}
mock0Allocator := newMockAllocator(t)
collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()
collections.Insert(2, &collectionInfo{
ID: 2,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
},
})
tests := []struct {
name string
fields fields
args args
wantErr bool
wantPlans []*datapb.CompactionPlan
}{
{
"test many segments",
fields{
&meta{
segments: segmentInfos,
channelCPs: newChannelCps(),
collections: collections,
indexMeta: indexMeta,
},
mock0Allocator,
nil,
&spyCompactionInspector{t: t, spyChan: make(chan *datapb.CompactionPlan, 2)},
nil,
},
args{
2,
&compactTime{},
},
false,
[]*datapb.CompactionPlan{
{
PlanID: 2,
SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{
{
SegmentID: 1,
FieldBinlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log1"},
},
},
},
Field2StatslogPaths: nil,
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "deltalog1"},
},
},
},
IsSorted: true,
},
{
SegmentID: 2,
FieldBinlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log2"},
},
},
},
Field2StatslogPaths: nil,
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "deltalog2"},
},
},
},
IsSorted: true,
},
},
BeginLogID: 100,
PreAllocatedLogIDs: &datapb.IDRange{Begin: 200, End: 2000},
StartTime: 3,
Type: datapb.CompactionType_MixCompaction,
Channel: "ch1",
MaxSize: 1342177280,
},
},
},
}
for _, tt := range tests {
(tt.fields.inspector).(*spyCompactionInspector).meta = tt.fields.meta
t.Run(tt.name, func(t *testing.T) {
tr := &compactionTrigger{
meta: tt.fields.meta,
handler: newMockHandlerWithMeta(tt.fields.meta),
allocator: tt.fields.allocator,
signals: make(chan *compactionSignal, 100),
manualSignals: make(chan *compactionSignal, 100),
inspector: tt.fields.inspector,
globalTrigger: tt.fields.globalTrigger,
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
tr.closeWaiter.Add(1)
go func() {
defer tr.closeWaiter.Done()
tr.work()
}()
defer tr.stop()
_, err := tr.TriggerCompaction(context.TODO(), NewCompactionSignal().WithCollectionID(tt.args.collectionID).WithIsForce(true))
assert.Equal(t, tt.wantErr, err != nil)
spy := (tt.fields.inspector).(*spyCompactionInspector)
select {
case plan := <-spy.spyChan:
assert.NotEmpty(t, plan)
assert.Equal(t, len(plan.SegmentBinlogs), nSegments)
case <-time.After(2 * time.Second):
assert.Fail(t, "timeout")
}
})
}
}
func sortPlanCompactionBinlogs(plan *datapb.CompactionPlan) {
sort.Slice(plan.SegmentBinlogs, func(i, j int) bool {
return plan.SegmentBinlogs[i].SegmentID < plan.SegmentBinlogs[j].SegmentID
})
}
// Test no compaction selection
func Test_compactionTrigger_noplan(t *testing.T) {
type fields struct {
meta *meta
allocator allocator.Allocator
signals chan *compactionSignal
inspector CompactionInspector
globalTrigger *time.Ticker
}
type args struct {
collectionID int64
compactTime *compactTime
}
Params.Save(Params.DataCoordCfg.MinSegmentToMerge.Key, "4")
defer Params.Save(Params.DataCoordCfg.MinSegmentToMerge.Key, Params.DataCoordCfg.MinSegmentToMerge.DefaultValue)
vecFieldID := int64(201)
mock0Allocator := newMockAllocator(t)
im := newSegmentIndexMeta(nil)
im.indexes[2] = make(map[UniqueID]*model.Index)
collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()
collections.Insert(2, &collectionInfo{
ID: 2,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
},
})
tests := []struct {
name string
fields fields
args args
wantErr bool
wantPlans []*datapb.CompactionPlan
}{
{
"test no plan",
fields{
&meta{
indexMeta: im,
// 4 segment
channelCPs: newChannelCps(),
segments: &SegmentsInfo{
segments: map[int64]*SegmentInfo{
1: {
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 200,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log1", LogSize: 100, MemorySize: 100},
},
},
},
},
lastFlushTime: time.Now(),
},
2: {
SegmentInfo: &datapb.SegmentInfo{
ID: 2,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 200,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log2", LogSize: Params.DataCoordCfg.SegmentMaxSize.GetAsInt64()*1024*1024 - 1, MemorySize: Params.DataCoordCfg.SegmentMaxSize.GetAsInt64()*1024*1024 - 1},
},
},
},
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "deltalog2"},
},
},
},
},
lastFlushTime: time.Now(),
},
},
},
collections: collections,
},
mock0Allocator,
make(chan *compactionSignal, 1),
&spyCompactionInspector{t: t, spyChan: make(chan *datapb.CompactionPlan, 1)},
nil,
},
args{
2,
&compactTime{},
},
false,
nil,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tr := &compactionTrigger{
meta: tt.fields.meta,
handler: newMockHandlerWithMeta(tt.fields.meta),
allocator: tt.fields.allocator,
signals: make(chan *compactionSignal, 100),
inspector: tt.fields.inspector,
globalTrigger: tt.fields.globalTrigger,
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
tr.start()
defer tr.stop()
_, err := tr.TriggerCompaction(context.TODO(), NewCompactionSignal())
assert.Equal(t, tt.wantErr, err != nil)
spy := (tt.fields.inspector).(*spyCompactionInspector)
select {
case val := <-spy.spyChan:
assert.Fail(t, "we expect no compaction generated", val)
return
case <-time.After(3 * time.Second):
return
}
})
}
}
func mockSegment(segID, rows, deleteRows, sizeInMB int64) *datapb.SegmentInfo {
return &datapb.SegmentInfo{
ID: segID,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: sizeInMB,
MaxRowNum: 150,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: rows, LogPath: "log1", LogSize: 100 * 1024 * 1024, MemorySize: sizeInMB * 1024 * 1024},
},
},
},
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: deleteRows, LogPath: "deltalog1"},
},
},
},
IsSorted: true,
}
}
func mockSegmentsInfo(sizeInMB ...int64) *SegmentsInfo {
segments := make(map[int64]*SegmentInfo, len(sizeInMB))
collectionID := int64(2)
channel := "ch1"
coll2Segments := make(map[UniqueID]map[UniqueID]*SegmentInfo)
coll2Segments[collectionID] = make(map[UniqueID]*SegmentInfo)
channel2Segments := make(map[string]map[UniqueID]*SegmentInfo)
channel2Segments[channel] = make(map[UniqueID]*SegmentInfo)
for i, size := range sizeInMB {
segId := int64(i + 1)
info := &SegmentInfo{
SegmentInfo: mockSegment(segId, size, 1, size),
lastFlushTime: time.Now().Add(-100 * time.Minute),
}
segments[segId] = info
coll2Segments[collectionID][segId] = info
channel2Segments[channel][segId] = info
}
return &SegmentsInfo{
segments: segments,
secondaryIndexes: segmentInfoIndexes{
coll2Segments: coll2Segments,
channel2Segments: channel2Segments,
},
}
}
// Test compaction with prioritized candi
func Test_compactionTrigger_PrioritizedCandi(t *testing.T) {
type fields struct {
meta *meta
allocator allocator.Allocator
signals chan *compactionSignal
inspector CompactionInspector
globalTrigger *time.Ticker
}
vecFieldID := int64(201)
mock0Allocator := newMockAllocator(t)
genSegIndex := func(segID, indexID UniqueID, numRows int64) *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex] {
segIdx := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: 2,
PartitionID: 1,
NumRows: numRows,
IndexID: indexID,
BuildID: segID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
})
return segIdx
}
collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()
collections.Insert(2, &collectionInfo{
ID: 2,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
},
})
im := &indexMeta{
segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
indexes: map[UniqueID]map[UniqueID]*model.Index{
2: {
indexID: {
TenantID: "",
CollectionID: 2,
FieldID: vecFieldID,
IndexID: indexID,
IndexName: "_default_idx",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "HNSW",
},
},
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
}
im.segmentIndexes.Insert(1, genSegIndex(1, indexID, 20))
im.segmentIndexes.Insert(2, genSegIndex(2, indexID, 20))
im.segmentIndexes.Insert(3, genSegIndex(3, indexID, 20))
im.segmentIndexes.Insert(4, genSegIndex(4, indexID, 20))
im.segmentIndexes.Insert(5, genSegIndex(5, indexID, 20))
im.segmentIndexes.Insert(6, genSegIndex(6, indexID, 20))
tests := []struct {
name string
fields fields
wantErr bool
wantPlans []*datapb.CompactionPlan
}{
{
"test small segment",
fields{
&meta{
// 8 small segments
channelCPs: newChannelCps(),
segments: mockSegmentsInfo(20, 20, 20, 20, 20, 20),
indexMeta: im,
collections: collections,
},
mock0Allocator,
make(chan *compactionSignal, 1),
&spyCompactionInspector{t: t, spyChan: make(chan *datapb.CompactionPlan, 1)},
nil,
},
false,
nil,
},
}
for _, tt := range tests {
(tt.fields.inspector).(*spyCompactionInspector).meta = tt.fields.meta
t.Run(tt.name, func(t *testing.T) {
tt.fields.meta.channelCPs.checkpoints["ch1"] = &msgpb.MsgPosition{
Timestamp: tsoutil.ComposeTSByTime(time.Now()),
MsgID: []byte{1, 2, 3, 4},
}
tr := &compactionTrigger{
meta: tt.fields.meta,
handler: newMockHandlerWithMeta(tt.fields.meta),
allocator: tt.fields.allocator,
signals: make(chan *compactionSignal, 100),
inspector: tt.fields.inspector,
globalTrigger: tt.fields.globalTrigger,
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
tr.start()
defer tr.stop()
_, err := tr.TriggerCompaction(context.TODO(), NewCompactionSignal())
assert.Equal(t, tt.wantErr, err != nil)
spy := (tt.fields.inspector).(*spyCompactionInspector)
select {
case val := <-spy.spyChan:
// 6 segments in the final pick list
assert.Equal(t, 6, len(val.SegmentBinlogs))
return
case <-time.After(3 * time.Second):
assert.Fail(t, "failed to get plan")
return
}
})
}
}
// Test compaction with small candi
func Test_compactionTrigger_SmallCandi(t *testing.T) {
type fields struct {
meta *meta
allocator allocator.Allocator
signals chan *compactionSignal
inspector CompactionInspector
globalTrigger *time.Ticker
}
type args struct {
collectionID int64
compactTime *compactTime
}
vecFieldID := int64(201)
mock0Allocator := newMockAllocator(t)
genSegIndex := func(segID, indexID UniqueID, numRows int64) *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex] {
segIdx := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: 2,
PartitionID: 1,
NumRows: numRows,
IndexID: indexID,
BuildID: segID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
})
return segIdx
}
collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()
collections.Insert(2, &collectionInfo{
ID: 2,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
},
},
},
})
im := &indexMeta{
segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
indexes: map[UniqueID]map[UniqueID]*model.Index{
2: {
indexID: {
TenantID: "",
CollectionID: 2,
FieldID: vecFieldID,
IndexID: indexID,
IndexName: "_default_idx",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "HNSW",
},
},
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
}
im.segmentIndexes.Insert(1, genSegIndex(1, indexID, 20))
im.segmentIndexes.Insert(2, genSegIndex(2, indexID, 20))
im.segmentIndexes.Insert(3, genSegIndex(3, indexID, 20))
im.segmentIndexes.Insert(4, genSegIndex(4, indexID, 20))
im.segmentIndexes.Insert(5, genSegIndex(5, indexID, 20))
im.segmentIndexes.Insert(6, genSegIndex(6, indexID, 20))
im.segmentIndexes.Insert(7, genSegIndex(7, indexID, 20))
tests := []struct {
name string
fields fields
args args
wantErr bool
wantPlans []*datapb.CompactionPlan
}{
{
"test small segment",
fields{
&meta{
channelCPs: newChannelCps(),
// 7 segments with 200MB each, the compaction is expected to be triggered
// as the first 5 being merged, and 1 plus being squeezed.
segments: mockSegmentsInfo(200, 200, 200, 200, 200, 200, 200),
indexMeta: im,
collections: collections,
},
mock0Allocator,
make(chan *compactionSignal, 1),
&spyCompactionInspector{t: t, spyChan: make(chan *datapb.CompactionPlan, 1)},
nil,
},
args{
2,
&compactTime{},
},
false,
nil,
},
}
for _, tt := range tests {
(tt.fields.inspector).(*spyCompactionInspector).meta = tt.fields.meta
t.Run(tt.name, func(t *testing.T) {
tt.fields.meta.channelCPs.checkpoints["ch1"] = &msgpb.MsgPosition{
Timestamp: tsoutil.ComposeTSByTime(time.Now()),
MsgID: []byte{1, 2, 3, 4},
}
tr := &compactionTrigger{
meta: tt.fields.meta,
handler: newMockHandlerWithMeta(tt.fields.meta),
allocator: tt.fields.allocator,
signals: make(chan *compactionSignal, 100),
inspector: tt.fields.inspector,
globalTrigger: tt.fields.globalTrigger,
indexEngineVersionManager: newMockVersionManager(),
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
tr.start()
defer tr.stop()
_, err := tr.TriggerCompaction(context.TODO(), NewCompactionSignal())
assert.Equal(t, tt.wantErr, err != nil)
spy := (tt.fields.inspector).(*spyCompactionInspector)
select {
case val := <-spy.spyChan:
// 6 segments in the final pick list.
// 5 generated by the origin plan, 1 was added as additional segment.
assert.Equal(t, len(val.SegmentBinlogs), 6)
return
case <-time.After(3 * time.Second):
assert.Fail(t, "failed to get plan")
return
}
})
}
}
// Test segment compaction target size
func Test_compactionTrigger_noplan_random_size(t *testing.T) {
type fields struct {
meta *meta
allocator allocator.Allocator
signals chan *compactionSignal
inspector CompactionInspector
globalTrigger *time.Ticker
}
type args struct {
collectionID int64
compactTime *compactTime
}
segmentInfos := &SegmentsInfo{
segments: make(map[UniqueID]*SegmentInfo),
secondaryIndexes: segmentInfoIndexes{
coll2Segments: map[UniqueID]map[UniqueID]*SegmentInfo{2: {}},
channel2Segments: map[string]map[UniqueID]*SegmentInfo{"ch1": {}},
},
}
size := []int64{
510, 500, 480, 300, 250, 200, 128, 128, 128, 127,
40, 40, 40, 40, 40, 40, 40, 40, 40, 40,
20, 20, 20, 20, 20, 20, 20, 20, 20, 20,
10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
}
vecFieldID := int64(201)
indexMeta := newSegmentIndexMeta(nil)
indexMeta.indexes = map[UniqueID]map[UniqueID]*model.Index{
2: {
indexID: {
TenantID: "",
CollectionID: 2,
FieldID: vecFieldID,
IndexID: indexID,
IndexName: "_default_idx",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "HNSW",
},
},
IsAutoIndex: false,
UserIndexParams: nil,
},
},
}
for i := UniqueID(0); i < 50; i++ {
info := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: i,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: size[i],
MaxRowNum: 512,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log1", LogSize: size[i] * 2 * 1024 * 1024, MemorySize: size[i] * 2 * 1024 * 1024},
},
},
},
IsSorted: true,
},
lastFlushTime: time.Now(),
}
indexMeta.updateSegmentIndex(&model.SegmentIndex{
SegmentID: i,
CollectionID: 2,
PartitionID: 1,
NumRows: 100,
IndexID: indexID,
BuildID: i,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
})
segmentInfos.segments[i] = info
segmentInfos.secondaryIndexes.coll2Segments[2][i] = info
segmentInfos.secondaryIndexes.channel2Segments["ch1"][i] = info
}
mock0Allocator := newMockAllocator(t)
collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()
collections.Insert(2, &collectionInfo{
ID: 2,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
},
})
tests := []struct {
name string
fields fields
args args
wantErr bool
wantPlans []*datapb.CompactionPlan
}{
{
"test rand size segment",
fields{
&meta{
channelCPs: newChannelCps(),
segments: segmentInfos,
collections: collections,
indexMeta: indexMeta,
},
mock0Allocator,
make(chan *compactionSignal, 1),
&spyCompactionInspector{t: t, spyChan: make(chan *datapb.CompactionPlan, 10)},
nil,
},
args{
2,
&compactTime{},
},
false,
nil,
},
}
for _, tt := range tests {
(tt.fields.inspector).(*spyCompactionInspector).meta = tt.fields.meta
t.Run(tt.name, func(t *testing.T) {
tt.fields.meta.channelCPs.checkpoints["ch1"] = &msgpb.MsgPosition{
Timestamp: tsoutil.ComposeTSByTime(time.Now()),
MsgID: []byte{1, 2, 3, 4},
}
tr := &compactionTrigger{
meta: tt.fields.meta,
handler: newMockHandlerWithMeta(tt.fields.meta),
allocator: tt.fields.allocator,
signals: make(chan *compactionSignal, 100),
inspector: tt.fields.inspector,
globalTrigger: tt.fields.globalTrigger,
indexEngineVersionManager: newMockVersionManager(),
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
tr.start()
defer tr.stop()
_, err := tr.TriggerCompaction(context.TODO(), NewCompactionSignal())
assert.Equal(t, tt.wantErr, err != nil)
spy := (tt.fields.inspector).(*spyCompactionInspector)
// should be split into two plans
var plans []*datapb.CompactionPlan
ticker := time.NewTicker(3 * time.Second)
defer ticker.Stop()
WAIT:
for {
select {
case val := <-spy.spyChan:
plans = append(plans, val)
case <-ticker.C:
break WAIT
}
}
assert.Equal(t, 4, len(plans))
// plan 1: 250 + 20 * 10 + 3 * 20
// plan 2: 200 + 7 * 20 + 4 * 40
// plan 3: 128 + 6 * 40 + 127
// plan 4: 300 + 128 + 128 ( < 512 * 1.25)
// assert.Equal(t, 24, len(plans[0].GetInputSegments()))
// assert.Equal(t, 12, len(plans[1].GetInputSegments()))
// assert.Equal(t, 8, len(plans[2].GetInputSegments()))
// assert.Equal(t, 3, len(plans[3].GetInputSegments()))
})
}
}
// Test shouldDoSingleCompaction
func Test_compactionTrigger_shouldDoSingleCompaction(t *testing.T) {
indexMeta := newSegmentIndexMeta(nil)
mock0Allocator := newMockAllocator(t)
trigger := newCompactionTrigger(&meta{
indexMeta: indexMeta,
channelCPs: newChannelCps(),
}, &compactionInspector{}, mock0Allocator, newMockHandler(), newIndexEngineVersionManager())
// Test too many deltalogs.
var binlogs []*datapb.FieldBinlog
for i := UniqueID(0); i < 1000; i++ {
binlogs = append(binlogs, &datapb.FieldBinlog{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log1", LogSize: 100, MemorySize: 100},
},
})
}
info := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Deltalogs: binlogs,
},
}
couldDo := trigger.ShouldDoSingleCompaction(info, &compactTime{})
assert.True(t, couldDo)
// Test too many stats log
info = &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Statslogs: binlogs,
},
}
couldDo = trigger.ShouldDoSingleCompaction(info, &compactTime{})
assert.False(t, couldDo)
// Test expire triggered compaction
var binlogs2 []*datapb.FieldBinlog
for i := UniqueID(0); i < 100; i++ {
binlogs2 = append(binlogs2, &datapb.FieldBinlog{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log1", LogSize: 100000, TimestampFrom: 300, TimestampTo: 500, MemorySize: 100000},
},
})
}
for i := UniqueID(0); i < 100; i++ {
binlogs2 = append(binlogs2, &datapb.FieldBinlog{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log1", LogSize: 1000000, TimestampFrom: 300, TimestampTo: 1000, MemorySize: 1000000},
},
})
}
info2 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 600,
NumOfRows: 10000,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: binlogs2,
},
}
// expire time < Timestamp To
couldDo = trigger.ShouldDoSingleCompaction(info2, &compactTime{expireTime: 300})
assert.False(t, couldDo)
// didn't reach single compaction size 10 * 1024 * 1024
couldDo = trigger.ShouldDoSingleCompaction(info2, &compactTime{expireTime: 600})
assert.False(t, couldDo)
// expire time < Timestamp False
couldDo = trigger.ShouldDoSingleCompaction(info2, &compactTime{expireTime: 1200})
assert.True(t, couldDo)
// Test Delete triggered compaction
var binlogs3 []*datapb.FieldBinlog
for i := UniqueID(0); i < 100; i++ {
binlogs3 = append(binlogs2, &datapb.FieldBinlog{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log1", LogSize: 100000, TimestampFrom: 300, TimestampTo: 500, MemorySize: 100000},
},
})
}
info3 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 700,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: binlogs3,
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 200, LogPath: "deltalog1"},
},
},
},
},
}
// deltalog is large enough, should do compaction
couldDo = trigger.ShouldDoSingleCompaction(info3, &compactTime{})
assert.True(t, couldDo)
mockVersionManager := NewMockVersionManager(t)
mockVersionManager.On("GetCurrentIndexEngineVersion").Return(int32(2)).Maybe()
mockVersionManager.On("GetCurrentScalarIndexEngineVersion").Return(int32(2)).Maybe()
mockVersionManager.On("ResolveVecIndexVersion").Return(int32(5)).Maybe()
trigger.indexEngineVersionManager = mockVersionManager
info4 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 600,
NumOfRows: 10000,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: binlogs2,
},
}
info5 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 600,
NumOfRows: 10000,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: binlogs2,
},
}
info6 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 600,
NumOfRows: 10000,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: binlogs2,
},
}
indexMeta.updateSegmentIndex(&model.SegmentIndex{
SegmentID: 1,
IndexID: 101,
CurrentIndexVersion: 1,
IndexFileKeys: []string{"index1"},
})
indexMeta.indexes = map[UniqueID]map[UniqueID]*model.Index{
2: {
101: {
CollectionID: 2,
IndexID: 101,
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "HNSW"},
},
},
},
}
// expire time < Timestamp To, but index engine version is 2 which is larger than CurrentIndexVersion in segmentIndex
Params.Save(Params.DataCoordCfg.AutoUpgradeSegmentIndex.Key, "true")
defer Params.Save(Params.DataCoordCfg.AutoUpgradeSegmentIndex.Key, "false")
couldDo = trigger.ShouldDoSingleCompaction(info4, &compactTime{expireTime: 300})
assert.True(t, couldDo)
indexMeta.updateSegmentIndex(&model.SegmentIndex{
SegmentID: 1,
IndexID: 101,
CurrentIndexVersion: 2,
IndexFileKeys: []string{"index1"},
})
// expire time < Timestamp To, and index engine version is 2 which is equal CurrentIndexVersion in segmentIndex
couldDo = trigger.ShouldDoSingleCompaction(info5, &compactTime{expireTime: 300})
assert.False(t, couldDo)
Params.Save(Params.DataCoordCfg.ForceRebuildSegmentIndex.Key, "true")
defer Params.Save(Params.DataCoordCfg.ForceRebuildSegmentIndex.Key, "false")
Params.Save(Params.DataCoordCfg.TargetVecIndexVersion.Key, "5")
defer Params.Save(Params.DataCoordCfg.TargetVecIndexVersion.Key, "-1")
couldDo = trigger.ShouldDoSingleCompaction(info5, &compactTime{expireTime: 300})
assert.True(t, couldDo)
indexMeta.updateSegmentIndex(&model.SegmentIndex{
SegmentID: 1,
IndexID: 101,
CurrentIndexVersion: 1,
IndexFileKeys: nil,
})
// expire time < Timestamp To, and index engine version is 2 which is larger than CurrentIndexVersion in segmentIndex but indexFileKeys is nil
couldDo = trigger.ShouldDoSingleCompaction(info6, &compactTime{expireTime: 300})
assert.False(t, couldDo)
// Test import segment: old row timestamps should not trigger TTL compaction when commit_timestamp is recent
t.Run("import segment not TTL-triggered by old row timestamps", func(t *testing.T) {
// Row timestamps are very old (ts=1000), but commit_timestamp is in the future (now + 1 day).
// expireTime is 5000 which is > row timestamps but < commit_timestamp.
// Without the fix, rows appear expired. With the fix, commit_timestamp is used as effective ts.
now := time.Now()
commitTs := tsoutil.ComposeTSByTime(now.Add(24 * time.Hour)) // future: definitely not expired
expireTime := uint64(5000) // > old row ts (1000) but < commitTs
var importBinlogs []*datapb.FieldBinlog
for i := 0; i < 100; i++ {
importBinlogs = append(importBinlogs, &datapb.FieldBinlog{
Binlogs: []*datapb.Binlog{
// TimestampTo=1000: very old, looks expired compared to expireTime=5000
{EntriesNum: 5, LogPath: "log1", LogSize: 100000, TimestampFrom: 500, TimestampTo: 1000, MemorySize: 100000},
},
})
}
importSegment := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 2000,
NumOfRows: 500,
MaxRowNum: 1000,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: importBinlogs,
CommitTimestamp: commitTs,
},
}
// Without commit_timestamp fix, all rows look expired (TimestampTo=1000 < expireTime=5000),
// and 100% expired ratio would trigger compaction. With the fix, commit_timestamp is used
// as the effective ts, which is far in the future, so no TTL compaction is triggered.
couldDo := trigger.ShouldDoSingleCompaction(importSegment, &compactTime{expireTime: expireTime})
assert.False(t, couldDo, "import segment with recent commit_timestamp should not be TTL-compacted due to old row timestamps")
})
}
func Test_compactionTrigger_shouldDoSingleCompaction_CommitTimestamp(t *testing.T) {
indexMeta := newSegmentIndexMeta(nil)
mock0Allocator := newMockAllocator(t)
trigger := newCompactionTrigger(&meta{
indexMeta: indexMeta,
channelCPs: newChannelCps(),
}, &compactionInspector{}, mock0Allocator, newMockHandler(), newIndexEngineVersionManager())
// Import segment: binlog TimestampFrom=100, TimestampTo=500 (very old)
// commit_timestamp=2000 (current time)
// expireTime=1000 — between old binlog ts and commit_ts
// Without fix: would trigger (500 < 1000). With fix: must NOT trigger (max(500,2000)=2000 > 1000).
importSeg := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 99,
CollectionID: 1,
NumOfRows: 1000,
State: commonpb.SegmentState_Flushed,
CommitTimestamp: 2000,
Binlogs: []*datapb.FieldBinlog{
{Binlogs: []*datapb.Binlog{
{EntriesNum: 900, TimestampFrom: 100, TimestampTo: 500, MemorySize: 1024 * 1024},
}},
},
},
}
// expireTime=1000: old binlog ts (500) < 1000, but commit_ts (2000) > 1000 → NOT expired
shouldCompact := trigger.ShouldDoSingleCompaction(importSeg, &compactTime{expireTime: 1000})
assert.False(t, shouldCompact,
"import segment with commit_ts=2000 must NOT be TTL-compacted at expireTime=1000")
// At expireTime=3000 > commit_ts=2000 → should compact
shouldCompact = trigger.ShouldDoSingleCompaction(importSeg, &compactTime{expireTime: 3000})
assert.True(t, shouldCompact,
"import segment with commit_ts=2000 MUST be TTL-compacted at expireTime=3000")
// Normal segment (commit_ts=0): old behavior preserved
normalSeg := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 100,
NumOfRows: 1000,
State: commonpb.SegmentState_Flushed,
// CommitTimestamp = 0
Binlogs: []*datapb.FieldBinlog{
{Binlogs: []*datapb.Binlog{
{EntriesNum: 900, TimestampFrom: 100, TimestampTo: 500, MemorySize: 1024 * 1024},
}},
},
},
}
// Normal: 500 < 1000 → should compact (old behavior unchanged)
shouldCompact = trigger.ShouldDoSingleCompaction(normalSeg, &compactTime{expireTime: 1000})
assert.True(t, shouldCompact,
"normal segment with binlog ts=500 MUST be TTL-compacted at expireTime=1000 (unchanged behavior)")
}
func Test_compactionTrigger_ShouldStrictCompactExpiry(t *testing.T) {
trigger := &compactionTrigger{}
now := time.Now()
expireTS := tsoutil.ComposeTSByTime(now)
compact := &compactTime{
expireTime: expireTS,
}
segment := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 10,
CollectionID: 20,
PartitionID: 30,
InsertChannel: "test-channel",
},
}
t.Run("no tolerance, fromTs before expire time => should compact", func(t *testing.T) {
Params.Save(Params.DataCoordCfg.CompactionExpiryTolerance.Key, "0")
defer Params.Save(Params.DataCoordCfg.CompactionExpiryTolerance.Key, "-1") // reset
fromTs := tsoutil.ComposeTSByTime(now.Add(-time.Hour))
shouldCompact := trigger.ShouldCompactExpiry(fromTs, compact, segment)
assert.True(t, shouldCompact)
})
t.Run("negative tolerance, disable force expiry compaction => should not compact", func(t *testing.T) {
fromTs := tsoutil.ComposeTSByTime(now.Add(-time.Hour))
shouldCompact := trigger.ShouldCompactExpiry(fromTs, compact, segment)
assert.False(t, shouldCompact)
})
t.Run("with tolerance, fromTs within tolerance => should NOT compact", func(t *testing.T) {
Params.Save(Params.DataCoordCfg.CompactionExpiryTolerance.Key, "2")
defer Params.Save(Params.DataCoordCfg.CompactionExpiryTolerance.Key, "-1") // reset
fromTs := tsoutil.ComposeTSByTime(now.Add(-time.Hour)) // within 2h tolerance
shouldCompact := trigger.ShouldCompactExpiry(fromTs, compact, segment)
assert.False(t, shouldCompact)
})
t.Run("with tolerance, fromTs before expireTime - tolerance => should compact", func(t *testing.T) {
Params.Save(Params.DataCoordCfg.CompactionExpiryTolerance.Key, "2")
defer Params.Save(Params.DataCoordCfg.CompactionExpiryTolerance.Key, "-1") // reset
fromTs := tsoutil.ComposeTSByTime(now.Add(-3 * time.Hour)) // earlier than expireTime - 30m
shouldCompact := trigger.ShouldCompactExpiry(fromTs, compact, segment)
assert.True(t, shouldCompact)
})
}
func Test_compactionTrigger_new(t *testing.T) {
type args struct {
meta *meta
inspector CompactionInspector
allocator allocator.Allocator
}
tests := []struct {
name string
args args
}{
{
"test new trigger",
args{
&meta{},
&compactionInspector{},
allocator.NewMockAllocator(t),
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := newCompactionTrigger(tt.args.meta, tt.args.inspector, tt.args.allocator, newMockHandler(), newMockVersionManager())
assert.Equal(t, tt.args.meta, got.meta)
assert.Equal(t, tt.args.inspector, got.inspector)
assert.Equal(t, tt.args.allocator, got.allocator)
})
}
}
func TestCompactionTriggerKeepsMixedSchemaVersionSegments(t *testing.T) {
paramtable.Init()
collectionID := int64(1)
partitionID := int64(10)
channel := "ch-1"
schema := newTestSchema()
schema.Version = 5
mt := &meta{
segments: NewSegmentsInfo(),
indexMeta: newSegmentIndexMeta(nil),
collections: typeutil.NewConcurrentMap[UniqueID, *collectionInfo](),
}
mt.collections.Insert(collectionID, &collectionInfo{ID: collectionID, Schema: schema})
for _, item := range []struct {
id int64
schemaVersion int32
}{
{id: 101, schemaVersion: 3},
{id: 102, schemaVersion: 5},
{id: 103, schemaVersion: 4},
} {
mt.segments.SetSegment(item.id, &SegmentInfo{SegmentInfo: &datapb.SegmentInfo{
ID: item.id,
CollectionID: collectionID,
PartitionID: partitionID,
InsertChannel: channel,
State: commonpb.SegmentState_Flushed,
Level: datapb.SegmentLevel_L1,
IsSorted: true,
NumOfRows: 10,
MaxRowNum: 100,
SchemaVersion: item.schemaVersion,
Binlogs: []*datapb.FieldBinlog{{FieldID: 1, Binlogs: []*datapb.Binlog{{EntriesNum: 10, MemorySize: 1}}}},
}})
}
inspector := &spyCompactionInspector{t: t, spyChan: make(chan *datapb.CompactionPlan, 1), meta: mt}
trigger := newCompactionTrigger(mt, inspector, newMock0Allocator(t), newMockHandlerWithMeta(mt), newMockVersionManager())
err := trigger.handleSignal(&compactionSignal{
id: 1,
collectionID: collectionID,
partitionID: partitionID,
channel: channel,
isForce: true,
})
assert.NoError(t, err)
select {
case plan := <-inspector.spyChan:
var got []int64
for _, segment := range plan.GetSegmentBinlogs() {
got = append(got, segment.GetSegmentID())
}
assert.ElementsMatch(t, []int64{101, 102, 103}, got)
case <-time.After(time.Second):
assert.Fail(t, "expected compaction plan for mixed schema version segments")
}
}
func Test_compactionTrigger_getCompactTime(t *testing.T) {
coll := &collectionInfo{
ID: 1,
Schema: newTestSchema(),
Partitions: []UniqueID{1},
Properties: map[string]string{
common.CollectionTTLConfigKey: "10",
},
}
now := tsoutil.ComposeTSByTime(time.Now())
ct, err := getCompactTime(now, coll)
assert.NoError(t, err)
assert.NotNil(t, ct)
}
func Test_TirggerCompaction_WaitResult(t *testing.T) {
originValue := Params.DataCoordCfg.EnableAutoCompaction.GetValue()
Params.Save(Params.DataCoordCfg.EnableAutoCompaction.Key, "true")
defer func() {
Params.Save(Params.DataCoordCfg.EnableAutoCompaction.Key, originValue)
}()
m := &meta{
channelCPs: newChannelCps(),
segments: NewSegmentsInfo(), collections: typeutil.NewConcurrentMap[UniqueID, *collectionInfo](),
}
got := newCompactionTrigger(m, &compactionInspector{}, newMockAllocator(t),
&ServerHandler{
&Server{
meta: m,
},
}, newMockVersionManager())
got.signals = make(chan *compactionSignal, 1)
{
_, err := got.TriggerCompaction(context.TODO(), NewCompactionSignal().
WithCollectionID(1).
WithPartitionID(1).
WithSegmentIDs(1).
WithChannel("a").
WithWaitResult(false))
assert.NoError(t, err)
}
{
_, err := got.TriggerCompaction(context.TODO(), NewCompactionSignal().
WithCollectionID(2).
WithPartitionID(2).
WithSegmentIDs(2).
WithChannel("b").
WithWaitResult(false))
assert.Error(t, err)
}
var i satomic.Value
i.Store(0)
check := func() {
for {
select {
case signal := <-got.signals:
x := i.Load().(int)
i.Store(x + 1)
assert.EqualValues(t, 1, signal.collectionID)
default:
return
}
}
}
check()
assert.Equal(t, 1, i.Load().(int))
var j satomic.Value
j.Store(0)
go func() {
timeoutCtx, cancelFunc := context.WithTimeout(context.Background(), time.Second)
defer cancelFunc()
for {
select {
case signal := <-got.signals:
x := j.Load().(int)
j.Store(x + 1)
switch x {
case 0:
assert.EqualValues(t, 3, signal.collectionID)
case 1:
assert.EqualValues(t, 4, signal.collectionID)
}
signal.Notify(nil)
case <-timeoutCtx.Done():
return
}
}
}()
{
_, err := got.TriggerCompaction(context.TODO(), NewCompactionSignal().
WithCollectionID(3).
WithPartitionID(3).
WithSegmentIDs(3).
WithChannel("c").
WithWaitResult(false))
assert.NoError(t, err)
}
{
_, err := got.TriggerCompaction(context.TODO(), NewCompactionSignal().
WithCollectionID(4).
WithPartitionID(4).
WithSegmentIDs(4).
WithChannel("d").
WithWaitResult(true))
assert.NoError(t, err)
}
assert.Eventually(t, func() bool {
return j.Load().(int) == 2
}, 2*time.Second, 500*time.Millisecond)
}
type CompactionTriggerSuite struct {
suite.Suite
collectionID int64
partitionID int64
channel string
indexID int64
vecFieldID int64
meta *meta
tr *compactionTrigger
allocator *allocator.MockAllocator
handler *NMockHandler
inspector *MockCompactionInspector
versionManager *MockVersionManager
}
func (s *CompactionTriggerSuite) SetupSuite() {
paramtable.Init()
}
func (s *CompactionTriggerSuite) genSeg(segID, numRows int64) *datapb.SegmentInfo {
return &datapb.SegmentInfo{
ID: segID,
CollectionID: s.collectionID,
PartitionID: s.partitionID,
LastExpireTime: 100,
NumOfRows: numRows,
MaxRowNum: 110,
InsertChannel: s.channel,
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogPath: "log1", LogSize: 100, MemorySize: 100},
},
},
},
IsSorted: true,
}
}
func (s *CompactionTriggerSuite) genSegIndex(segID, indexID UniqueID, numRows int64) *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex] {
segIdx := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: s.collectionID,
PartitionID: s.partitionID,
NumRows: numRows,
IndexID: indexID,
BuildID: segID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
})
return segIdx
}
func (s *CompactionTriggerSuite) SetupTest() {
s.collectionID = 100
s.partitionID = 200
s.indexID = 300
s.vecFieldID = 400
s.channel = "dml_0_100v0"
catalog := mocks.NewDataCoordCatalog(s.T())
catalog.EXPECT().SaveChannelCheckpoint(mock.Anything, s.channel, mock.Anything).Return(nil)
seg1 := &SegmentInfo{
SegmentInfo: s.genSeg(1, 60),
lastFlushTime: time.Now().Add(-100 * time.Minute),
}
seg2 := &SegmentInfo{
SegmentInfo: s.genSeg(2, 60),
lastFlushTime: time.Now(),
}
seg3 := &SegmentInfo{
SegmentInfo: s.genSeg(3, 60),
lastFlushTime: time.Now(),
}
seg4 := &SegmentInfo{
SegmentInfo: s.genSeg(4, 60),
lastFlushTime: time.Now(),
}
seg5 := &SegmentInfo{
SegmentInfo: s.genSeg(5, 60),
lastFlushTime: time.Now(),
}
seg6 := &SegmentInfo{
SegmentInfo: s.genSeg(6, 60),
lastFlushTime: time.Now(),
}
collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()
collections.Insert(s.collectionID, &collectionInfo{
ID: s.collectionID,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: s.vecFieldID,
DataType: schemapb.DataType_FloatVector,
},
},
},
})
im := &indexMeta{
segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
indexes: map[UniqueID]map[UniqueID]*model.Index{
s.collectionID: {
s.indexID: {
TenantID: "",
CollectionID: s.collectionID,
FieldID: s.vecFieldID,
IndexID: s.indexID,
IndexName: "_default_idx",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "HNSW",
},
},
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
}
im.segmentIndexes.Insert(1, s.genSegIndex(1, indexID, 60))
im.segmentIndexes.Insert(2, s.genSegIndex(2, indexID, 60))
im.segmentIndexes.Insert(3, s.genSegIndex(3, indexID, 60))
im.segmentIndexes.Insert(4, s.genSegIndex(4, indexID, 60))
im.segmentIndexes.Insert(5, s.genSegIndex(5, indexID, 60))
im.segmentIndexes.Insert(6, s.genSegIndex(6, indexID, 60))
s.meta = &meta{
channelCPs: newChannelCps(),
catalog: catalog,
segments: &SegmentsInfo{
segments: map[int64]*SegmentInfo{
1: seg1,
2: seg2,
3: seg3,
4: seg4,
5: seg5,
6: seg6,
},
secondaryIndexes: segmentInfoIndexes{
coll2Segments: map[UniqueID]map[UniqueID]*SegmentInfo{
s.collectionID: {
1: seg1,
2: seg2,
3: seg3,
4: seg4,
5: seg5,
6: seg6,
},
},
channel2Segments: map[string]map[UniqueID]*SegmentInfo{
s.channel: {
1: seg1,
2: seg2,
3: seg3,
4: seg4,
5: seg5,
6: seg6,
},
},
},
},
indexMeta: im,
collections: collections,
}
s.meta.UpdateChannelCheckpoint(context.TODO(), s.channel, &msgpb.MsgPosition{
ChannelName: s.channel,
Timestamp: tsoutil.ComposeTSByTime(time.Now()),
MsgID: []byte{1, 2, 3, 4},
})
s.allocator = allocator.NewMockAllocator(s.T())
s.inspector = NewMockCompactionInspector(s.T())
s.handler = NewNMockHandler(s.T())
s.versionManager = NewMockVersionManager(s.T())
s.tr = newCompactionTrigger(
s.meta,
s.inspector,
s.allocator,
s.handler,
s.versionManager,
)
s.tr.testingOnly = true
}
func (s *CompactionTriggerSuite) TestHandleSignal() {
s.Run("getCompaction_failed", func() {
defer s.SetupTest()
tr := s.tr
s.inspector.EXPECT().isFull().Return(false)
// s.allocator.EXPECT().AllocTimestamp(mock.Anything).Return(10000, nil)
s.handler.EXPECT().GetCollection(mock.Anything, int64(100)).Return(nil, errors.New("mocked"))
tr.handleSignal(&compactionSignal{
collectionID: s.collectionID,
partitionID: s.partitionID,
channel: s.channel,
isForce: false,
})
// suite shall check inspector.enqueueCompaction never called
})
s.Run("collectionAutoCompactionConfigError", func() {
defer s.SetupTest()
tr := s.tr
s.inspector.EXPECT().isFull().Return(false)
// s.allocator.EXPECT().AllocTimestamp(mock.Anything).Return(10000, nil)
s.handler.EXPECT().GetCollection(mock.Anything, int64(100)).Return(&collectionInfo{
Properties: map[string]string{
common.CollectionAutoCompactionKey: "bad_value",
},
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: s.vecFieldID,
DataType: schemapb.DataType_FloatVector,
},
},
},
}, nil)
tr.handleSignal(&compactionSignal{
collectionID: s.collectionID,
partitionID: s.partitionID,
channel: s.channel,
isForce: false,
})
// suite shall check inspector.enqueueCompaction never called
})
s.Run("collectionAutoCompactionDisabled", func() {
defer s.SetupTest()
tr := s.tr
s.inspector.EXPECT().isFull().Return(false)
// s.allocator.EXPECT().AllocTimestamp(mock.Anything).Return(10000, nil)
s.handler.EXPECT().GetCollection(mock.Anything, int64(100)).Return(&collectionInfo{
Properties: map[string]string{
common.CollectionAutoCompactionKey: "false",
},
ID: s.collectionID,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: s.vecFieldID,
DataType: schemapb.DataType_FloatVector,
},
},
},
}, nil)
tr.handleSignal(&compactionSignal{
collectionID: s.collectionID,
partitionID: s.partitionID,
channel: s.channel,
isForce: false,
})
// suite shall check inspector.enqueueCompaction never called
})
s.Run("collectionAutoCompactionDisabled_force", func() {
defer s.SetupTest()
tr := s.tr
s.inspector.EXPECT().isFull().Return(false)
// s.allocator.EXPECT().AllocTimestamp(mock.Anything).Return(10000, nil)
// s.allocator.EXPECT().AllocID(mock.Anything).Return(20000, nil)
start := int64(20000)
s.allocator.EXPECT().AllocN(mock.Anything).RunAndReturn(func(i int64) (int64, int64, error) {
return start, start + i, nil
})
s.allocator.EXPECT().AllocTimestamp(mock.Anything).Return(10000, nil)
s.handler.EXPECT().GetCollection(mock.Anything, int64(100)).Return(&collectionInfo{
Properties: map[string]string{
common.CollectionAutoCompactionKey: "false",
},
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: s.vecFieldID,
DataType: schemapb.DataType_FloatVector,
},
},
},
}, nil)
s.inspector.EXPECT().enqueueCompaction(mock.Anything).Return(nil)
tr.handleSignal(&compactionSignal{
collectionID: s.collectionID,
partitionID: s.partitionID,
channel: s.channel,
isForce: true,
})
})
s.Run("force compaction preallocates segment IDs from input size", func() {
defer s.SetupTest()
pt := paramtable.Get()
pt.Save(pt.DataCoordCfg.IndexBasedCompaction.Key, "false")
defer pt.Reset(pt.DataCoordCfg.IndexBasedCompaction.Key)
pt.Save(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key, "1")
defer pt.Reset(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key)
pt.Save(pt.DataCoordCfg.SegmentMaxSize.Key, "100")
defer pt.Reset(pt.DataCoordCfg.SegmentMaxSize.Key)
const mb = 1024 * 1024
s.meta.segments.segments[1].Binlogs[0].Binlogs[0].MemorySize = 250 * mb
tr := s.tr
handler := NewNMockHandler(s.T())
handler.EXPECT().GetCollection(mock.Anything, s.collectionID).
Return(&collectionInfo{
ID: s.collectionID,
Schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: s.vecFieldID, DataType: schemapb.DataType_FloatVector},
},
},
}, nil).Once()
tr.handler = handler
const (
startID = int64(20000)
planID = int64(20003)
endID = int64(20004)
)
s.allocator.EXPECT().AllocN(int64(4)).Return(startID, endID, nil).Once()
s.inspector.EXPECT().enqueueCompaction(mock.MatchedBy(func(task *datapb.CompactionTask) bool {
s.EqualValues(planID, task.GetPlanID())
s.ElementsMatch([]int64{1}, task.GetInputSegments())
s.EqualValues(startID, task.GetPreAllocatedSegmentIDs().GetBegin())
s.EqualValues(planID, task.GetPreAllocatedSegmentIDs().GetEnd())
return true
})).Return(nil).Once()
err := tr.handleSignal(NewCompactionSignal().
WithCollectionID(s.collectionID).
WithPartitionID(s.partitionID).
WithChannel(s.channel).
WithSegmentIDs(1).
WithIsForce(true))
s.NoError(err)
})
}
func TestIsCollectionAutoCompactionEnabledExternal(t *testing.T) {
coll := &collectionInfo{
ID: 1,
Schema: &schemapb.CollectionSchema{
ExternalSource: "s3://external",
Fields: []*schemapb.FieldSchema{
{
FieldID: 1,
Name: "external_pk",
DataType: schemapb.DataType_Int64,
ExternalField: "pk_col",
},
},
},
}
assert.False(t, isCollectionAutoCompactionEnabled(coll))
}
func (s *CompactionTriggerSuite) TestHandleGlobalSignal() {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: common.StartOfUserFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
{
FieldID: common.StartOfUserFieldID + 1,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
}
s.Run("GetCollection_failed", func() {
defer s.SetupTest()
tr := s.tr
s.inspector.EXPECT().isFull().Return(false)
s.handler.EXPECT().GetCollection(mock.Anything, int64(100)).Return(nil, errors.New("mocked"))
err := tr.handleSignal(NewCompactionSignal().
WithCollectionID(s.collectionID).
WithPartitionID(s.partitionID).
WithChannel(s.channel))
s.Error(err)
})
s.Run("collectionAutoCompactionConfigError", func() {
defer s.SetupTest()
tr := s.tr
s.inspector.EXPECT().isFull().Return(false)
s.allocator.EXPECT().AllocTimestamp(mock.Anything).Return(10000, nil)
s.handler.EXPECT().GetCollection(mock.Anything, int64(100)).Return(&collectionInfo{
Schema: schema,
Properties: map[string]string{
common.CollectionAutoCompactionKey: "bad_value",
},
}, nil)
s.NotPanics(func() {
err := tr.handleSignal(NewCompactionSignal().
WithCollectionID(s.collectionID).
WithPartitionID(s.partitionID).
WithChannel(s.channel))
s.NoError(err)
}, "bad configuration shall not cause panicking")
})
s.Run("collectionAutoCompactionDisabled", func() {
defer s.SetupTest()
tr := s.tr
s.inspector.EXPECT().isFull().Return(false)
s.handler.EXPECT().GetCollection(mock.Anything, int64(100)).Return(&collectionInfo{
Schema: schema,
Properties: map[string]string{
common.CollectionAutoCompactionKey: "false",
},
}, nil)
err := tr.handleSignal(NewCompactionSignal().
WithCollectionID(s.collectionID).
WithPartitionID(s.partitionID).
WithChannel(s.channel))
s.NoError(err)
})
s.Run("collectionAutoCompactionDisabled_force", func() {
defer s.SetupTest()
tr := s.tr
s.inspector.EXPECT().isFull().Return(false)
start := int64(20000)
s.allocator.EXPECT().AllocN(mock.Anything).RunAndReturn(func(i int64) (int64, int64, error) {
return start, start + i, nil
}).Maybe()
s.allocator.EXPECT().AllocTimestamp(mock.Anything).Return(10000, nil)
s.handler.EXPECT().GetCollection(mock.Anything, int64(100)).Return(&collectionInfo{
Schema: schema,
Properties: map[string]string{
common.CollectionAutoCompactionKey: "false",
},
}, nil)
err := tr.handleSignal(NewCompactionSignal().
WithCollectionID(s.collectionID).
WithPartitionID(s.partitionID).
WithChannel(s.channel))
s.NoError(err)
})
}
func (s *CompactionTriggerSuite) TestSqueezeSmallSegments() {
expectedSize := int64(70000)
smallsegments := []*SegmentInfo{
{SegmentInfo: &datapb.SegmentInfo{ID: 3}, size: *atomic.NewInt64(69999)},
{SegmentInfo: &datapb.SegmentInfo{ID: 1}, size: *atomic.NewInt64(100)},
}
largeSegment := &SegmentInfo{SegmentInfo: &datapb.SegmentInfo{ID: 2}, size: *atomic.NewInt64(expectedSize)}
buckets := [][]*SegmentInfo{{largeSegment}}
s.Require().Equal(1, len(buckets))
s.Require().Equal(1, len(buckets[0]))
remaining := s.tr.squeezeSmallSegmentsToBuckets(smallsegments, buckets, expectedSize)
s.Equal(1, len(remaining))
s.EqualValues(3, remaining[0].ID)
s.Equal(1, len(buckets))
s.Equal(2, len(buckets[0]))
mlog.Info(context.TODO(), "buckets", mlog.Any("buckets", buckets))
}
func TestCompactionTriggerSuite(t *testing.T) {
suite.Run(t, new(CompactionTriggerSuite))
}
func Test_compactionTrigger_generatePlans(t *testing.T) {
catalog := mocks.NewDataCoordCatalog(t)
catalog.EXPECT().AlterSegments(mock.Anything, mock.Anything).Return(nil).Maybe()
vecFieldID := int64(201)
indexID := int64(1001)
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
}
mock0Allocator := newMock0Allocator(t)
seg1 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{
EntriesNum: 5,
LogID: 1,
TimestampFrom: 1000,
TimestampTo: 2000,
},
},
},
},
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 1, MemorySize: 1 * 1024 * 1024},
},
},
},
IsSorted: true,
},
}
seg2 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 2,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{
EntriesNum: 5,
LogID: 2,
MemorySize: 3 * 1024 * 1024 * 1024,
TimestampFrom: 2000,
TimestampTo: 3000,
},
},
},
},
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 2},
},
},
},
IsSorted: true,
},
}
type fields struct {
meta *meta
allocator allocator.Allocator
signals chan *compactionSignal
inspector CompactionInspector
globalTrigger *time.Ticker
}
type args struct {
segments []*SegmentInfo
signal *compactionSignal
compactTime *compactTime
expectedSize int64
}
collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()
collections.Insert(2, &collectionInfo{
ID: 2,
Schema: schema,
Properties: map[string]string{
common.CollectionTTLConfigKey: "0",
},
})
segIndexes := typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]]()
segIdx0 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx0.Insert(indexID, &model.SegmentIndex{
SegmentID: 1,
CollectionID: 2,
PartitionID: 1,
NumRows: 100,
IndexID: indexID,
BuildID: 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIndexes.Insert(1, segIdx0)
segIdx1 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx1.Insert(indexID, &model.SegmentIndex{
SegmentID: 2,
CollectionID: 2,
PartitionID: 1,
NumRows: 100,
IndexID: indexID,
BuildID: 2,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIndexes.Insert(2, segIdx1)
tests := []struct {
name string
fields fields
args args
want []*typeutil.Pair[int64, []int64]
}{
{
name: "force trigger on large segments",
fields: fields{
meta: &meta{
catalog: catalog,
channelCPs: newChannelCps(),
segments: &SegmentsInfo{
segments: map[int64]*SegmentInfo{
1: seg1,
2: seg2,
},
secondaryIndexes: segmentInfoIndexes{
coll2Segments: map[UniqueID]map[UniqueID]*SegmentInfo{
2: {
seg1.GetID(): seg1,
seg2.GetID(): seg2,
},
},
},
},
indexMeta: &indexMeta{
segmentIndexes: segIndexes,
indexes: map[UniqueID]map[UniqueID]*model.Index{
2: {
indexID: {
TenantID: "",
CollectionID: 2,
FieldID: vecFieldID,
IndexID: indexID,
IndexName: "_default_idx",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "HNSW",
},
},
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
},
collections: collections,
},
allocator: mock0Allocator,
signals: nil,
inspector: &spyCompactionInspector{t: t, spyChan: make(chan *datapb.CompactionPlan, 1)},
globalTrigger: nil,
},
args: args{
segments: []*SegmentInfo{seg1, seg2},
signal: &compactionSignal{collectionID: 2, partitionID: 1, channel: "ch1", isForce: true},
compactTime: nil,
expectedSize: 1024 * 1024 * 1024,
},
want: []*typeutil.Pair[int64, []int64]{
{A: 100, B: []int64{1}}, {A: 100, B: []int64{2}},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tr := &compactionTrigger{
meta: tt.fields.meta,
handler: newMockHandlerWithMeta(tt.fields.meta),
allocator: tt.fields.allocator,
signals: tt.fields.signals,
inspector: tt.fields.inspector,
globalTrigger: tt.fields.globalTrigger,
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
if got := tr.generatePlans(tt.args.segments, tt.args.signal, tt.args.compactTime, tt.args.expectedSize); !reflect.DeepEqual(got, tt.want) {
t.Errorf("compactionTrigger.generatePlans() = %+v, want %+v", got, tt.want)
}
})
}
}
func Test_compactionTrigger_generatePlansByTime(t *testing.T) {
catalog := mocks.NewDataCoordCatalog(t)
catalog.EXPECT().AlterSegments(mock.Anything, mock.Anything).Return(nil).Maybe()
vecFieldID := int64(201)
indexID := int64(1001)
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: vecFieldID,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
}
mock0Allocator := newMock0Allocator(t)
seg1 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{
EntriesNum: 5,
LogID: 1,
TimestampFrom: 1000,
TimestampTo: 2000,
},
},
},
},
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 1, MemorySize: 1 * 1024 * 1024},
},
},
},
IsSorted: true,
},
}
seg2 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 2,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{
EntriesNum: 5,
LogID: 2,
TimestampFrom: 1000,
TimestampTo: 2000,
},
},
},
},
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 2},
},
},
},
IsSorted: true,
},
}
seg3 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 3,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{
EntriesNum: 5,
LogID: 3,
TimestampFrom: 3000,
TimestampTo: 4000,
},
},
},
},
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{EntriesNum: 5, LogID: 3},
},
},
},
IsSorted: true,
},
}
type fields struct {
meta *meta
allocator allocator.Allocator
signals chan *compactionSignal
inspector CompactionInspector
globalTrigger *time.Ticker
}
type args struct {
segments []*SegmentInfo
signal *compactionSignal
compactTime *compactTime
expectedSize int64
}
collections := typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()
collections.Insert(2, &collectionInfo{
ID: 2,
Schema: schema,
Properties: map[string]string{
common.CollectionTTLConfigKey: "0",
},
})
segIndexes := typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]]()
segIdx0 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx0.Insert(indexID, &model.SegmentIndex{
SegmentID: 1,
CollectionID: 2,
PartitionID: 1,
NumRows: 100,
IndexID: indexID,
BuildID: 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIndexes.Insert(1, segIdx0)
segIdx1 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx1.Insert(indexID, &model.SegmentIndex{
SegmentID: 2,
CollectionID: 2,
PartitionID: 1,
NumRows: 100,
IndexID: indexID,
BuildID: 2,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIndexes.Insert(2, segIdx1)
segIdx2 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx2.Insert(indexID, &model.SegmentIndex{
SegmentID: 3,
CollectionID: 2,
PartitionID: 1,
NumRows: 100,
IndexID: indexID,
BuildID: 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIndexes.Insert(3, segIdx2)
tests := []struct {
name string
fields fields
args args
want []*typeutil.Pair[int64, []int64]
}{
{
name: "test time-based compaction",
fields: fields{
meta: &meta{
catalog: catalog,
channelCPs: newChannelCps(),
segments: &SegmentsInfo{
segments: map[int64]*SegmentInfo{
1: seg1,
2: seg2,
3: seg3,
},
secondaryIndexes: segmentInfoIndexes{
coll2Segments: map[UniqueID]map[UniqueID]*SegmentInfo{
2: {
seg1.GetID(): seg1,
seg2.GetID(): seg2,
seg3.GetID(): seg3,
},
},
},
},
indexMeta: &indexMeta{
segmentIndexes: segIndexes,
indexes: map[UniqueID]map[UniqueID]*model.Index{
2: {
indexID: {
TenantID: "",
CollectionID: 2,
FieldID: vecFieldID,
IndexID: indexID,
IndexName: "_default_idx",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "HNSW",
},
},
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
},
collections: collections,
},
allocator: mock0Allocator,
signals: nil,
inspector: &spyCompactionInspector{t: t, spyChan: make(chan *datapb.CompactionPlan, 1)},
globalTrigger: nil,
},
args: args{
segments: []*SegmentInfo{seg1, seg2, seg3},
signal: &compactionSignal{collectionID: 2, partitionID: 1, channel: "ch1", isForce: false},
compactTime: &compactTime{startTime: 1000, collectionTTL: time.Hour},
expectedSize: 1024 * 1024 * 1024,
},
want: []*typeutil.Pair[int64, []int64]{
{A: 300, B: []int64{1, 2, 3}},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tr := &compactionTrigger{
meta: tt.fields.meta,
handler: newMockHandlerWithMeta(tt.fields.meta),
allocator: tt.fields.allocator,
signals: tt.fields.signals,
inspector: tt.fields.inspector,
globalTrigger: tt.fields.globalTrigger,
closeCh: lifetime.NewSafeChan(),
testingOnly: true,
}
got := tr.generatePlans(tt.args.segments, tt.args.signal, tt.args.compactTime, tt.args.expectedSize)
for i, pair := range got {
t.Logf("got[%d]: totalRows=%d, segmentIDs=%v", i, pair.A, pair.B)
}
if !reflect.DeepEqual(got, tt.want) {
t.Errorf("compactionTrigger.generatePlans() = %+v, want %+v", got, tt.want)
}
})
}
}
func Test_compactionTrigger_ShouldCompactExpiryWithTTLField(t *testing.T) {
trigger := &compactionTrigger{}
ts := time.Now()
segment := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
ExpirQuantiles: []int64{ts.UnixMicro(), ts.Add(time.Hour).UnixMicro(), ts.Add(2 * time.Hour).UnixMicro(), ts.Add(3 * time.Hour).UnixMicro(), ts.Add(4 * time.Hour).UnixMicro()},
},
}
startTime := tsoutil.ComposeTSByTime(ts.Add(time.Minute))
ct := &compactTime{startTime: startTime, collectionTTL: 0}
shouldCompact := trigger.ShouldCompactExpiryWithTTLField(ct, segment)
assert.True(t, shouldCompact)
startTime = tsoutil.ComposeTSByTime(ts.Add(-1 * time.Hour))
ct = &compactTime{startTime: startTime, collectionTTL: 0}
shouldCompact = trigger.ShouldCompactExpiryWithTTLField(ct, segment)
assert.False(t, shouldCompact)
origin := Params.DataCoordCfg.SingleCompactionRatioThreshold.GetValue()
defer Params.Save(Params.DataCoordCfg.SingleCompactionRatioThreshold.Key, origin)
Params.Save(Params.DataCoordCfg.SingleCompactionRatioThreshold.Key, "0.1")
startTime = tsoutil.ComposeTSByTime(ts.Add(time.Minute))
ct = &compactTime{startTime: startTime, collectionTTL: 0}
shouldCompact = trigger.ShouldCompactExpiryWithTTLField(ct, segment)
assert.True(t, shouldCompact)
Params.Save(Params.DataCoordCfg.SingleCompactionRatioThreshold.Key, "5")
startTime = tsoutil.ComposeTSByTime(ts.Add(time.Minute))
ct = &compactTime{startTime: startTime, collectionTTL: 0}
shouldCompact = trigger.ShouldCompactExpiryWithTTLField(ct, segment)
assert.False(t, shouldCompact)
segment2 := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
PartitionID: 1,
LastExpireTime: 100,
NumOfRows: 100,
MaxRowNum: 300,
InsertChannel: "ch1",
State: commonpb.SegmentState_Flushed,
},
}
startTime = tsoutil.ComposeTSByTime(ts.Add(time.Minute))
ct = &compactTime{startTime: startTime, collectionTTL: 0}
shouldCompact = trigger.ShouldCompactExpiryWithTTLField(ct, segment2)
assert.False(t, shouldCompact)
}
func Test_compactionTrigger_ShouldCompactExpiryWithTTLField_CommitTimestamp(t *testing.T) {
trigger := &compactionTrigger{}
ts := time.Now()
// expirQuantiles contain very old timestamps (simulating imported data with stale row timestamps)
oldTs := ts.Add(-24 * time.Hour)
segment := &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{
ID: 1,
CollectionID: 2,
CommitTimestamp: tsoutil.ComposeTSByTime(ts),
ExpirQuantiles: []int64{
oldTs.UnixMicro(),
oldTs.Add(time.Minute).UnixMicro(),
oldTs.Add(2 * time.Minute).UnixMicro(),
oldTs.Add(3 * time.Minute).UnixMicro(),
oldTs.Add(4 * time.Minute).UnixMicro(),
},
},
}
// expirQuantiles are based on original row timestamps; ShouldCompactExpiryWithTTLField
// does not adjust for commit_timestamp, so stale quantiles still trigger compaction.
startTime := tsoutil.ComposeTSByTime(ts.Add(time.Minute))
ct := &compactTime{startTime: startTime, collectionTTL: 0}
shouldCompact := trigger.ShouldCompactExpiryWithTTLField(ct, segment)
assert.True(t, shouldCompact, "stale expirQuantiles should still trigger TTL compaction")
}
func newTestIndexMeta(collID, segID, indexID int64, indexType string, segIdx *model.SegmentIndex) *indexMeta {
im := newSegmentIndexMeta(nil)
im.indexes[collID] = map[UniqueID]*model.Index{
indexID: {
CollectionID: collID,
IndexID: indexID,
IndexName: "test_idx",
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: indexType},
},
},
}
segIdxMap := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdxMap.Insert(indexID, segIdx)
im.segmentIndexes.Insert(segID, segIdxMap)
return im
}
func Test_ShouldRebuildSegmentIndex_AutoUpgrade_ScalarUsesCorrectField(t *testing.T) {
paramtable.Init()
Params.Save("dataCoord.autoUpgradeSegmentIndex", "true")
collID, segID, indexID := int64(1), int64(100), int64(10)
t.Run("scalar auto-upgrade compares CurrentScalarIndexVersion", func(t *testing.T) {
// Scalar index with CurrentIndexVersion=5 (vector field) but CurrentScalarIndexVersion=1
// Engine scalar version is 2, so 1 < 2 should trigger rebuild
segIdx := &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
IndexID: indexID,
IndexFileKeys: []string{"file1"},
CurrentIndexVersion: 5, // vector version field - should NOT be used for scalar
CurrentScalarIndexVersion: 1, // scalar version field - should be used
}
im := newTestIndexMeta(collID, segID, indexID, "INVERTED", segIdx)
mockVM := NewMockVersionManager(t)
mockVM.On("GetCurrentScalarIndexEngineVersion").Return(int32(2)).Maybe()
mockVM.On("GetCurrentIndexEngineVersion").Return(int32(5)).Maybe()
trigger := &compactionTrigger{
meta: &meta{indexMeta: im, channelCPs: newChannelCps()},
indexEngineVersionManager: mockVM,
}
segment := &SegmentInfo{SegmentInfo: &datapb.SegmentInfo{ID: segID, CollectionID: collID}}
assert.True(t, trigger.ShouldRebuildSegmentIndex(segment))
})
t.Run("scalar auto-upgrade no rebuild when version matches", func(t *testing.T) {
segIdx := &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
IndexID: indexID,
IndexFileKeys: []string{"file1"},
CurrentIndexVersion: 1, // vector version - irrelevant
CurrentScalarIndexVersion: 2, // matches engine version
}
im := newTestIndexMeta(collID, segID, indexID, "INVERTED", segIdx)
mockVM := NewMockVersionManager(t)
mockVM.On("GetCurrentScalarIndexEngineVersion").Return(int32(2)).Maybe()
mockVM.On("GetCurrentIndexEngineVersion").Return(int32(5)).Maybe()
trigger := &compactionTrigger{
meta: &meta{indexMeta: im, channelCPs: newChannelCps()},
indexEngineVersionManager: mockVM,
}
segment := &SegmentInfo{SegmentInfo: &datapb.SegmentInfo{ID: segID, CollectionID: collID}}
assert.False(t, trigger.ShouldRebuildSegmentIndex(segment))
})
Params.Save("dataCoord.autoUpgradeSegmentIndex", "false")
}
func Test_ShouldRebuildSegmentIndex_ForceRebuild_ScalarUsesCorrectField(t *testing.T) {
paramtable.Init()
Params.Save("dataCoord.autoUpgradeSegmentIndex", "false")
collID, segID, indexID := int64(1), int64(100), int64(10)
t.Run("scalar force-rebuild compares CurrentScalarIndexVersion", func(t *testing.T) {
Params.Save("dataCoord.forceRebuildScalarSegmentIndex", "true")
Params.Save("dataCoord.targetScalarIndexVersion", "3")
defer func() {
Params.Save("dataCoord.forceRebuildScalarSegmentIndex", "false")
Params.Save("dataCoord.targetScalarIndexVersion", "-1")
}()
// CurrentScalarIndexVersion=2 != resolved target=3, should trigger
segIdx := &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
IndexID: indexID,
IndexFileKeys: []string{"file1"},
CurrentIndexVersion: 3, // vector version - should NOT be used
CurrentScalarIndexVersion: 2, // scalar version - should be used
}
im := newTestIndexMeta(collID, segID, indexID, "INVERTED", segIdx)
mockVM := NewMockVersionManager(t)
mockVM.On("ResolveScalarIndexVersion").Return(int32(3)).Maybe()
trigger := &compactionTrigger{
meta: &meta{indexMeta: im, channelCPs: newChannelCps()},
indexEngineVersionManager: mockVM,
}
segment := &SegmentInfo{SegmentInfo: &datapb.SegmentInfo{ID: segID, CollectionID: collID}}
assert.True(t, trigger.ShouldRebuildSegmentIndex(segment))
})
t.Run("scalar force-rebuild no trigger when version matches resolved target", func(t *testing.T) {
Params.Save("dataCoord.forceRebuildScalarSegmentIndex", "true")
Params.Save("dataCoord.targetScalarIndexVersion", "3")
defer func() {
Params.Save("dataCoord.forceRebuildScalarSegmentIndex", "false")
Params.Save("dataCoord.targetScalarIndexVersion", "-1")
}()
segIdx := &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
IndexID: indexID,
IndexFileKeys: []string{"file1"},
CurrentIndexVersion: 1, // irrelevant
CurrentScalarIndexVersion: 3, // matches resolved target
}
im := newTestIndexMeta(collID, segID, indexID, "INVERTED", segIdx)
mockVM := NewMockVersionManager(t)
mockVM.On("ResolveScalarIndexVersion").Return(int32(3)).Maybe()
trigger := &compactionTrigger{
meta: &meta{indexMeta: im, channelCPs: newChannelCps()},
indexEngineVersionManager: mockVM,
}
segment := &SegmentInfo{SegmentInfo: &datapb.SegmentInfo{ID: segID, CollectionID: collID}}
assert.False(t, trigger.ShouldRebuildSegmentIndex(segment))
})
}
func Test_ShouldRebuildSegmentIndex_ForceRebuild_TargetExceedsMax_Converges(t *testing.T) {
paramtable.Init()
Params.Save("dataCoord.autoUpgradeSegmentIndex", "false")
collID, segID, indexID := int64(1), int64(100), int64(10)
t.Run("vec force-rebuild with target>max converges after clamp", func(t *testing.T) {
Params.Save("dataCoord.forceRebuildSegmentIndex", "true")
Params.Save("dataCoord.targetVecIndexVersion", "30")
defer func() {
Params.Save("dataCoord.forceRebuildSegmentIndex", "false")
Params.Save("dataCoord.targetVecIndexVersion", "-1")
}()
// Index was already rebuilt to clamped version (20), should NOT trigger again
segIdx := &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
IndexID: indexID,
IndexFileKeys: []string{"file1"},
CurrentIndexVersion: 20, // matches resolved (clamped) target
}
im := newTestIndexMeta(collID, segID, indexID, "HNSW", segIdx)
mockVM := NewMockVersionManager(t)
// ResolveVecIndexVersion clamps target=30 to max=20
mockVM.On("ResolveVecIndexVersion").Return(int32(20)).Maybe()
trigger := &compactionTrigger{
meta: &meta{indexMeta: im, channelCPs: newChannelCps()},
indexEngineVersionManager: mockVM,
}
segment := &SegmentInfo{SegmentInfo: &datapb.SegmentInfo{ID: segID, CollectionID: collID}}
assert.False(t, trigger.ShouldRebuildSegmentIndex(segment))
})
t.Run("scalar force-rebuild with target>max converges after clamp", func(t *testing.T) {
Params.Save("dataCoord.forceRebuildScalarSegmentIndex", "true")
Params.Save("dataCoord.targetScalarIndexVersion", "10")
defer func() {
Params.Save("dataCoord.forceRebuildScalarSegmentIndex", "false")
Params.Save("dataCoord.targetScalarIndexVersion", "-1")
}()
// Index was already rebuilt to clamped version (5), should NOT trigger again
segIdx := &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
IndexID: indexID,
IndexFileKeys: []string{"file1"},
CurrentScalarIndexVersion: 5, // matches resolved (clamped) target
}
im := newTestIndexMeta(collID, segID, indexID, "INVERTED", segIdx)
mockVM := NewMockVersionManager(t)
// ResolveScalarIndexVersion clamps target=10 to max=5
mockVM.On("ResolveScalarIndexVersion").Return(int32(5)).Maybe()
trigger := &compactionTrigger{
meta: &meta{indexMeta: im, channelCPs: newChannelCps()},
indexEngineVersionManager: mockVM,
}
segment := &SegmentInfo{SegmentInfo: &datapb.SegmentInfo{ID: segID, CollectionID: collID}}
assert.False(t, trigger.ShouldRebuildSegmentIndex(segment))
})
}