admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-06 17:18:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

enhance: Refine task meta with key lock (#40613)

Browse Source 
issue: #39101

2.5 pr: #40146 #40353

Signed-off-by: Cai Zhang <cai.zhang@zilliz.com>
This commit is contained in:
cai.zhang 2025-03-14 15:44:22 +08:00 committed by GitHub
parent bf4fc6a8c6
commit 6dbe5d475e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
21 changed files with 1936 additions and 1784 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

666
internal/datacoord/compaction_trigger_test.go
View File

@ -169,7 +169,7 @@ func Test_compactionTrigger_force_without_index(t *testing.T) {
}, },
}, },
indexMeta: &indexMeta{ indexMeta: &indexMeta{
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{}, segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
indexes: map[UniqueID]map[UniqueID]*model.Index{}, indexes: map[UniqueID]map[UniqueID]*model.Index{},
}, },
collections: map[int64]*collectionInfo{ collections: map[int64]*collectionInfo{
@ -308,6 +308,109 @@ func Test_compactionTrigger_force(t *testing.T) {
}, },
} }
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)
tests := []struct { tests := []struct {
name string name string
fields fields fields fields
@ -340,109 +443,7 @@ func Test_compactionTrigger_force(t *testing.T) {
}, },
}, },
}, },
indexMeta: &indexMeta{ indexMeta: im,
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{
1: {
indexID: {
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,
},
},
2: {
indexID: {
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,
},
},
3: {
indexID: {
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,
},
},
},
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,
},
},
},
},
collections: map[int64]*collectionInfo{ collections: map[int64]*collectionInfo{
2: { 2: {
ID: 2, ID: 2,
@ -1200,21 +1201,53 @@ func Test_compactionTrigger_PrioritizedCandi(t *testing.T) {
mock0Allocator := newMockAllocator(t) mock0Allocator := newMockAllocator(t)
genSegIndex := func(segID, indexID UniqueID, numRows int64) map[UniqueID]*model.SegmentIndex { genSegIndex := func(segID, indexID UniqueID, numRows int64) *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex] {
return map[UniqueID]*model.SegmentIndex{ segIdx := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
indexID: { segIdx.Insert(indexID, &model.SegmentIndex{
SegmentID: segID, SegmentID: segID,
CollectionID: 2, CollectionID: 2,
PartitionID: 1, PartitionID: 1,
NumRows: numRows, NumRows: numRows,
IndexID: indexID, IndexID: indexID,
BuildID: segID, BuildID: segID,
NodeID: 0, NodeID: 0,
IndexVersion: 1, IndexVersion: 1,
IndexState: commonpb.IndexState_Finished, IndexState: commonpb.IndexState_Finished,
}, })
} return segIdx
} }
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 { tests := []struct {
name string name string
fields fields fields fields
@ -1228,39 +1261,8 @@ func Test_compactionTrigger_PrioritizedCandi(t *testing.T) {
// 8 small segments // 8 small segments
channelCPs: newChannelCps(), channelCPs: newChannelCps(),
segments: mockSegmentsInfo(20, 20, 20, 20, 20, 20), segments: mockSegmentsInfo(20, 20, 20, 20, 20, 20),
indexMeta: &indexMeta{ indexMeta: im,
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{
1: genSegIndex(1, indexID, 20),
2: genSegIndex(2, indexID, 20),
3: genSegIndex(3, indexID, 20),
4: genSegIndex(4, indexID, 20),
5: genSegIndex(5, indexID, 20),
6: genSegIndex(6, indexID, 20),
},
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: map[int64]*collectionInfo{ collections: map[int64]*collectionInfo{
2: { 2: {
ID: 2, ID: 2,
@ -1341,21 +1343,53 @@ func Test_compactionTrigger_SmallCandi(t *testing.T) {
vecFieldID := int64(201) vecFieldID := int64(201)
mock0Allocator := newMockAllocator(t) mock0Allocator := newMockAllocator(t)
genSegIndex := func(segID, indexID UniqueID, numRows int64) map[UniqueID]*model.SegmentIndex { genSegIndex := func(segID, indexID UniqueID, numRows int64) *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex] {
return map[UniqueID]*model.SegmentIndex{ segIdx := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
indexID: { segIdx.Insert(indexID, &model.SegmentIndex{
SegmentID: segID, SegmentID: segID,
CollectionID: 2, CollectionID: 2,
PartitionID: 1, PartitionID: 1,
NumRows: numRows, NumRows: numRows,
IndexID: indexID, IndexID: indexID,
BuildID: segID, BuildID: segID,
NodeID: 0, NodeID: 0,
IndexVersion: 1, IndexVersion: 1,
IndexState: commonpb.IndexState_Finished, IndexState: commonpb.IndexState_Finished,
}, })
} return segIdx
} }
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 { tests := []struct {
name string name string
fields fields fields fields
@ -1370,40 +1404,8 @@ func Test_compactionTrigger_SmallCandi(t *testing.T) {
channelCPs: newChannelCps(), channelCPs: newChannelCps(),
// 7 segments with 200MB each, the compaction is expected to be triggered // 7 segments with 200MB each, the compaction is expected to be triggered
// as the first 5 being merged, and 1 plus being squeezed. // as the first 5 being merged, and 1 plus being squeezed.
segments: mockSegmentsInfo(200, 200, 200, 200, 200, 200, 200), segments: mockSegmentsInfo(200, 200, 200, 200, 200, 200, 200),
indexMeta: &indexMeta{ indexMeta: im,
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{
1: genSegIndex(1, indexID, 20),
2: genSegIndex(2, indexID, 20),
3: genSegIndex(3, indexID, 20),
4: genSegIndex(4, indexID, 20),
5: genSegIndex(5, indexID, 20),
6: genSegIndex(6, indexID, 20),
7: genSegIndex(7, indexID, 20),
},
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: map[int64]*collectionInfo{ collections: map[int64]*collectionInfo{
2: { 2: {
ID: 2, ID: 2,
@ -1484,21 +1486,52 @@ func Test_compactionTrigger_SqueezeNonPlannedSegs(t *testing.T) {
} }
vecFieldID := int64(201) vecFieldID := int64(201)
genSegIndex := func(segID, indexID UniqueID, numRows int64) map[UniqueID]*model.SegmentIndex { genSegIndex := func(segID, indexID UniqueID, numRows int64) *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex] {
return map[UniqueID]*model.SegmentIndex{ segIdx := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
indexID: { segIdx.Insert(indexID, &model.SegmentIndex{
SegmentID: segID, SegmentID: segID,
CollectionID: 2, CollectionID: 2,
PartitionID: 1, PartitionID: 1,
NumRows: numRows, NumRows: numRows,
IndexID: indexID, IndexID: indexID,
BuildID: segID, BuildID: segID,
NodeID: 0, NodeID: 0,
IndexVersion: 1, IndexVersion: 1,
IndexState: commonpb.IndexState_Finished, IndexState: commonpb.IndexState_Finished,
}, })
} return segIdx
} }
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))
mock0Allocator := newMockAllocator(t) mock0Allocator := newMockAllocator(t)
tests := []struct { tests := []struct {
name string name string
@ -1513,39 +1546,8 @@ func Test_compactionTrigger_SqueezeNonPlannedSegs(t *testing.T) {
&meta{ &meta{
channelCPs: newChannelCps(), channelCPs: newChannelCps(),
segments: mockSegmentsInfo(600, 600, 600, 600, 260, 260), segments: mockSegmentsInfo(600, 600, 600, 600, 260, 260),
indexMeta: &indexMeta{ indexMeta: im,
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{
1: genSegIndex(1, indexID, 20),
2: genSegIndex(2, indexID, 20),
3: genSegIndex(3, indexID, 20),
4: genSegIndex(4, indexID, 20),
5: genSegIndex(5, indexID, 20),
6: genSegIndex(6, indexID, 20),
},
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: map[int64]*collectionInfo{ collections: map[int64]*collectionInfo{
2: { 2: {
ID: 2, ID: 2,
@ -2191,20 +2193,20 @@ func (s *CompactionTriggerSuite) genSeg(segID, numRows int64) *datapb.SegmentInf
} }
} }
func (s *CompactionTriggerSuite) genSegIndex(segID, indexID UniqueID, numRows int64) map[UniqueID]*model.SegmentIndex { func (s *CompactionTriggerSuite) genSegIndex(segID, indexID UniqueID, numRows int64) *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex] {
return map[UniqueID]*model.SegmentIndex{ segIdx := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
indexID: { segIdx.Insert(indexID, &model.SegmentIndex{
SegmentID: segID, SegmentID: segID,
CollectionID: s.collectionID, CollectionID: s.collectionID,
PartitionID: s.partitionID, PartitionID: s.partitionID,
NumRows: numRows, NumRows: numRows,
IndexID: indexID, IndexID: indexID,
BuildID: segID, BuildID: segID,
NodeID: 0, NodeID: 0,
IndexVersion: 1, IndexVersion: 1,
IndexState: commonpb.IndexState_Finished, IndexState: commonpb.IndexState_Finished,
}, })
} return segIdx
} }
func (s *CompactionTriggerSuite) SetupTest() { func (s *CompactionTriggerSuite) SetupTest() {
@ -2241,6 +2243,37 @@ func (s *CompactionTriggerSuite) SetupTest() {
lastFlushTime: time.Now(), lastFlushTime: time.Now(),
} }
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{ s.meta = &meta{
channelCPs: newChannelCps(), channelCPs: newChannelCps(),
catalog: catalog, catalog: catalog,
@ -2276,38 +2309,7 @@ func (s *CompactionTriggerSuite) SetupTest() {
}, },
}, },
}, },
indexMeta: &indexMeta{ indexMeta: im,
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{
1: s.genSegIndex(1, indexID, 60),
2: s.genSegIndex(2, indexID, 60),
3: s.genSegIndex(3, indexID, 60),
4: s.genSegIndex(4, indexID, 60),
5: s.genSegIndex(5, indexID, 26),
6: s.genSegIndex(6, indexID, 26),
},
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,
},
},
},
},
collections: map[int64]*collectionInfo{ collections: map[int64]*collectionInfo{
s.collectionID: { s.collectionID: {
ID: s.collectionID, ID: s.collectionID,
@ -2731,6 +2733,45 @@ func Test_compactionTrigger_generatePlans(t *testing.T) {
compactTime *compactTime compactTime *compactTime
expectedSize int64 expectedSize int64
} }
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 { tests := []struct {
name string name string
fields fields fields fields
@ -2758,46 +2799,7 @@ func Test_compactionTrigger_generatePlans(t *testing.T) {
}, },
}, },
indexMeta: &indexMeta{ indexMeta: &indexMeta{
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ segmentIndexes: segIndexes,
1: {
indexID: {
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,
},
},
2: {
indexID: {
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,
},
},
},
indexes: map[UniqueID]map[UniqueID]*model.Index{ indexes: map[UniqueID]map[UniqueID]*model.Index{
2: { 2: {
indexID: { indexID: {

248
internal/datacoord/garbage_collector_test.go
View File

@ -54,6 +54,7 @@ import (
"github.com/milvus-io/milvus/pkg/v2/util/lock" "github.com/milvus-io/milvus/pkg/v2/util/lock"
"github.com/milvus-io/milvus/pkg/v2/util/merr" "github.com/milvus-io/milvus/pkg/v2/util/merr"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable" "github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
) )
func Test_garbageCollector_basic(t *testing.T) { func Test_garbageCollector_basic(t *testing.T) {
@ -476,6 +477,45 @@ func createMetaForRecycleUnusedSegIndexes(catalog metastore.DataCoordCatalog) *m
}, },
}, },
} }
segIndexes := typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]]()
segIdx0 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx0.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1026,
IndexID: indexID,
BuildID: buildID,
NodeID: 1,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 10,
IndexFileKeys: []string{"file1", "file2"},
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx1.Insert(indexID, &model.SegmentIndex{
SegmentID: segID + 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 1026,
IndexID: indexID,
BuildID: buildID + 1,
NodeID: 1,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 10,
IndexFileKeys: []string{"file1", "file2"},
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIndexes.Insert(segID, segIdx0)
segIndexes.Insert(segID+1, segIdx1)
meta := &meta{ meta := &meta{
RWMutex: lock.RWMutex{}, RWMutex: lock.RWMutex{},
ctx: ctx, ctx: ctx,
@ -483,47 +523,8 @@ func createMetaForRecycleUnusedSegIndexes(catalog metastore.DataCoordCatalog) *m
collections: nil, collections: nil,
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
indexMeta: &indexMeta{ indexMeta: &indexMeta{
catalog: catalog, catalog: catalog,
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ segmentIndexes: segIndexes,
segID: {
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1026,
IndexID: indexID,
BuildID: buildID,
NodeID: 1,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 10,
IndexFileKeys: []string{"file1", "file2"},
IndexSerializedSize: 0,
WriteHandoff: false,
},
},
segID + 1: {
indexID: {
SegmentID: segID + 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 1026,
IndexID: indexID,
BuildID: buildID + 1,
NodeID: 1,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 10,
IndexFileKeys: []string{"file1", "file2"},
IndexSerializedSize: 0,
WriteHandoff: false,
},
},
},
indexes: map[UniqueID]map[UniqueID]*model.Index{}, indexes: map[UniqueID]map[UniqueID]*model.Index{},
segmentBuildInfo: newSegmentIndexBuildInfo(), segmentBuildInfo: newSegmentIndexBuildInfo(),
}, },
@ -641,6 +642,45 @@ func createMetaTableForRecycleUnusedIndexFiles(catalog *datacoord.Catalog) *meta
}, },
}, },
} }
segIndexes := typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]]()
segIdx0 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx0.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1026,
IndexID: indexID,
BuildID: buildID,
NodeID: 1,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 10,
IndexFileKeys: []string{"file1", "file2"},
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx1.Insert(indexID, &model.SegmentIndex{
SegmentID: segID + 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 1026,
IndexID: indexID,
BuildID: buildID + 1,
NodeID: 1,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 10,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIndexes.Insert(segID, segIdx0)
segIndexes.Insert(segID+1, segIdx1)
meta := &meta{ meta := &meta{
RWMutex: lock.RWMutex{}, RWMutex: lock.RWMutex{},
ctx: ctx, ctx: ctx,
@ -648,47 +688,8 @@ func createMetaTableForRecycleUnusedIndexFiles(catalog *datacoord.Catalog) *meta
collections: nil, collections: nil,
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
indexMeta: &indexMeta{ indexMeta: &indexMeta{
catalog: catalog, catalog: catalog,
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ segmentIndexes: segIndexes,
segID: {
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1026,
IndexID: indexID,
BuildID: buildID,
NodeID: 1,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 10,
IndexFileKeys: []string{"file1", "file2"},
IndexSerializedSize: 0,
WriteHandoff: false,
},
},
segID + 1: {
indexID: {
SegmentID: segID + 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 1026,
IndexID: indexID,
BuildID: buildID + 1,
NodeID: 1,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 10,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
},
},
indexes: map[UniqueID]map[UniqueID]*model.Index{ indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: { collID: {
indexID: { indexID: {
@ -1043,52 +1044,53 @@ func TestGarbageCollector_clearETCD(t *testing.T) {
}, },
}, },
} }
segIndexes := typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]]()
segIdx0 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx0.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 5000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: []string{"file1", "file2"},
IndexSerializedSize: 1024,
WriteHandoff: false,
})
segIdx1 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx1.Insert(indexID, &model.SegmentIndex{
SegmentID: segID + 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 5000,
IndexID: indexID,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: []string{"file3", "file4"},
IndexSerializedSize: 1024,
WriteHandoff: false,
})
segIndexes.Insert(segID, segIdx0)
segIndexes.Insert(segID+1, segIdx1)
m := &meta{ m := &meta{
catalog: catalog, catalog: catalog,
channelCPs: channelCPs, channelCPs: channelCPs,
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
indexMeta: &indexMeta{ indexMeta: &indexMeta{
catalog: catalog, keyLock: lock.NewKeyLock[UniqueID](),
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ catalog: catalog,
segID: { segmentIndexes: segIndexes,
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 5000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: []string{"file1", "file2"},
IndexSerializedSize: 1024,
WriteHandoff: false,
},
},
segID + 1: {
indexID: {
SegmentID: segID + 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 5000,
IndexID: indexID,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: []string{"file3", "file4"},
IndexSerializedSize: 1024,
WriteHandoff: false,
},
},
},
segmentBuildInfo: newSegmentIndexBuildInfo(), segmentBuildInfo: newSegmentIndexBuildInfo(),
indexes: map[UniqueID]map[UniqueID]*model.Index{ indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: { collID: {

299
internal/datacoord/index_meta.go
View File

@ -44,6 +44,7 @@ import (
"github.com/milvus-io/milvus/pkg/v2/proto/workerpb" "github.com/milvus-io/milvus/pkg/v2/proto/workerpb"
"github.com/milvus-io/milvus/pkg/v2/util/funcutil" "github.com/milvus-io/milvus/pkg/v2/util/funcutil"
"github.com/milvus-io/milvus/pkg/v2/util/indexparams" "github.com/milvus-io/milvus/pkg/v2/util/indexparams"
"github.com/milvus-io/milvus/pkg/v2/util/lock"
"github.com/milvus-io/milvus/pkg/v2/util/metricsinfo" "github.com/milvus-io/milvus/pkg/v2/util/metricsinfo"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable" "github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/timerecord" "github.com/milvus-io/milvus/pkg/v2/util/timerecord"
@ -52,19 +53,21 @@ import (
) )
type indexMeta struct { type indexMeta struct {
sync.RWMutex
ctx context.Context ctx context.Context
catalog metastore.DataCoordCatalog catalog metastore.DataCoordCatalog
// collectionIndexes records which indexes are on the collection // collectionIndexes records which indexes are on the collection
// collID -> indexID -> index // collID -> indexID -> index
indexes map[UniqueID]map[UniqueID]*model.Index fieldIndexLock sync.RWMutex
indexes map[UniqueID]map[UniqueID]*model.Index
// buildID2Meta records building index meta information of the segment // buildID2Meta records building index meta information of the segment
segmentBuildInfo *segmentBuildInfo segmentBuildInfo *segmentBuildInfo
// buildID -> lock
keyLock *lock.KeyLock[UniqueID]
// segmentID -> indexID -> segmentIndex // segmentID -> indexID -> segmentIndex
segmentIndexes map[UniqueID]map[UniqueID]*model.SegmentIndex segmentIndexes *typeutil.ConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]]
} }
func newIndexTaskStats(s *model.SegmentIndex) *metricsinfo.IndexTaskStats { func newIndexTaskStats(s *model.SegmentIndex) *metricsinfo.IndexTaskStats {
@ -86,7 +89,7 @@ func newIndexTaskStats(s *model.SegmentIndex) *metricsinfo.IndexTaskStats {
type segmentBuildInfo struct { type segmentBuildInfo struct {
// buildID2Meta records the meta information of the segment // buildID2Meta records the meta information of the segment
// buildID -> segmentIndex // buildID -> segmentIndex
buildID2SegmentIndex map[UniqueID]*model.SegmentIndex buildID2SegmentIndex *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]
// taskStats records the task stats of the segment // taskStats records the task stats of the segment
taskStats *expirable.LRU[UniqueID, *metricsinfo.IndexTaskStats] taskStats *expirable.LRU[UniqueID, *metricsinfo.IndexTaskStats]
} }
@ -94,28 +97,28 @@ type segmentBuildInfo struct {
func newSegmentIndexBuildInfo() *segmentBuildInfo { func newSegmentIndexBuildInfo() *segmentBuildInfo {
return &segmentBuildInfo{ return &segmentBuildInfo{
// build ID -> segment index // build ID -> segment index
buildID2SegmentIndex: make(map[UniqueID]*model.SegmentIndex), buildID2SegmentIndex: typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex](),
// build ID -> task stats // build ID -> task stats
taskStats: expirable.NewLRU[UniqueID, *metricsinfo.IndexTaskStats](1024, nil, time.Minute*30), taskStats: expirable.NewLRU[UniqueID, *metricsinfo.IndexTaskStats](1024, nil, time.Minute*30),
} }
} }
func (m *segmentBuildInfo) Add(segIdx *model.SegmentIndex) { func (m *segmentBuildInfo) Add(segIdx *model.SegmentIndex) {
m.buildID2SegmentIndex[segIdx.BuildID] = segIdx m.buildID2SegmentIndex.Insert(segIdx.BuildID, segIdx)
m.taskStats.Add(segIdx.BuildID, newIndexTaskStats(segIdx)) m.taskStats.Add(segIdx.BuildID, newIndexTaskStats(segIdx))
} }
func (m *segmentBuildInfo) Get(key UniqueID) (*model.SegmentIndex, bool) { func (m *segmentBuildInfo) Get(key UniqueID) (*model.SegmentIndex, bool) {
value, exists := m.buildID2SegmentIndex[key] value, exists := m.buildID2SegmentIndex.Get(key)
return value, exists return value, exists
} }
func (m *segmentBuildInfo) Remove(key UniqueID) { func (m *segmentBuildInfo) Remove(key UniqueID) {
delete(m.buildID2SegmentIndex, key) m.buildID2SegmentIndex.Remove(key)
} }
func (m *segmentBuildInfo) List() map[UniqueID]*model.SegmentIndex { func (m *segmentBuildInfo) List() []*model.SegmentIndex {
return m.buildID2SegmentIndex return m.buildID2SegmentIndex.Values()
} }
func (m *segmentBuildInfo) GetTaskStats() []*metricsinfo.IndexTaskStats { func (m *segmentBuildInfo) GetTaskStats() []*metricsinfo.IndexTaskStats {
@ -128,8 +131,9 @@ func newIndexMeta(ctx context.Context, catalog metastore.DataCoordCatalog) (*ind
ctx: ctx, ctx: ctx,
catalog: catalog, catalog: catalog,
indexes: make(map[UniqueID]map[UniqueID]*model.Index), indexes: make(map[UniqueID]map[UniqueID]*model.Index),
keyLock: lock.NewKeyLock[UniqueID](),
segmentBuildInfo: newSegmentIndexBuildInfo(), segmentBuildInfo: newSegmentIndexBuildInfo(),
segmentIndexes: make(map[UniqueID]map[UniqueID]*model.SegmentIndex), segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
} }
err := mt.reloadFromKV() err := mt.reloadFromKV()
if err != nil { if err != nil {
@ -174,12 +178,14 @@ func (m *indexMeta) updateCollectionIndex(index *model.Index) {
} }
func (m *indexMeta) updateSegmentIndex(segIdx *model.SegmentIndex) { func (m *indexMeta) updateSegmentIndex(segIdx *model.SegmentIndex) {
indexes, ok := m.segmentIndexes[segIdx.SegmentID] indexes, ok := m.segmentIndexes.Get(segIdx.SegmentID)
if ok { if ok {
indexes[segIdx.IndexID] = segIdx indexes.Insert(segIdx.IndexID, segIdx)
m.segmentIndexes.Insert(segIdx.SegmentID, indexes)
} else { } else {
m.segmentIndexes[segIdx.SegmentID] = make(map[UniqueID]*model.SegmentIndex) indexes := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
m.segmentIndexes[segIdx.SegmentID][segIdx.IndexID] = segIdx indexes.Insert(segIdx.IndexID, segIdx)
m.segmentIndexes.Insert(segIdx.SegmentID, indexes)
} }
m.segmentBuildInfo.Add(segIdx) m.segmentBuildInfo.Add(segIdx)
} }
@ -206,36 +212,37 @@ func (m *indexMeta) updateSegIndexMeta(segIdx *model.SegmentIndex, updateFunc fu
} }
func (m *indexMeta) updateIndexTasksMetrics() { func (m *indexMeta) updateIndexTasksMetrics() {
m.RLock() taskMetrics := make(map[indexpb.JobState]int)
defer m.RUnlock() taskMetrics[indexpb.JobState_JobStateNone] = 0
taskMetrics[indexpb.JobState_JobStateInit] = 0
taskMetrics := make(map[UniqueID]map[commonpb.IndexState]int) taskMetrics[indexpb.JobState_JobStateInProgress] = 0
taskMetrics[indexpb.JobState_JobStateFinished] = 0
taskMetrics[indexpb.JobState_JobStateFailed] = 0
taskMetrics[indexpb.JobState_JobStateRetry] = 0
for _, segIdx := range m.segmentBuildInfo.List() { for _, segIdx := range m.segmentBuildInfo.List() {
if segIdx.IsDeleted || !m.isIndexExist(segIdx.CollectionID, segIdx.IndexID) { if segIdx.IsDeleted || !m.IsIndexExist(segIdx.CollectionID, segIdx.IndexID) {
continue continue
} }
if _, ok := taskMetrics[segIdx.CollectionID]; !ok {
taskMetrics[segIdx.CollectionID] = make(map[commonpb.IndexState]int) switch segIdx.IndexState {
taskMetrics[segIdx.CollectionID][commonpb.IndexState_Unissued] = 0 case commonpb.IndexState_IndexStateNone:
taskMetrics[segIdx.CollectionID][commonpb.IndexState_InProgress] = 0 taskMetrics[indexpb.JobState_JobStateNone]++
taskMetrics[segIdx.CollectionID][commonpb.IndexState_Finished] = 0 case commonpb.IndexState_Unissued:
taskMetrics[segIdx.CollectionID][commonpb.IndexState_Failed] = 0 taskMetrics[indexpb.JobState_JobStateInit]++
case commonpb.IndexState_InProgress:
taskMetrics[indexpb.JobState_JobStateInProgress]++
case commonpb.IndexState_Finished:
taskMetrics[indexpb.JobState_JobStateFinished]++
case commonpb.IndexState_Failed:
taskMetrics[indexpb.JobState_JobStateFailed]++
case commonpb.IndexState_Retry:
taskMetrics[indexpb.JobState_JobStateRetry]++
} }
taskMetrics[segIdx.CollectionID][segIdx.IndexState]++
} }
for collID, m := range taskMetrics {
for k, v := range m { jobType := indexpb.JobType_JobTypeIndexJob.String()
switch k { for k, v := range taskMetrics {
case commonpb.IndexState_Unissued: metrics.TaskNum.WithLabelValues(jobType, k.String()).Set(float64(v))
metrics.IndexTaskNum.WithLabelValues(strconv.FormatInt(collID, 10), metrics.UnissuedIndexTaskLabel).Set(float64(v))
case commonpb.IndexState_InProgress:
metrics.IndexTaskNum.WithLabelValues(strconv.FormatInt(collID, 10), metrics.InProgressIndexTaskLabel).Set(float64(v))
case commonpb.IndexState_Finished:
metrics.IndexTaskNum.WithLabelValues(strconv.FormatInt(collID, 10), metrics.FinishedIndexTaskLabel).Set(float64(v))
case commonpb.IndexState_Failed:
metrics.IndexTaskNum.WithLabelValues(strconv.FormatInt(collID, 10), metrics.FailedIndexTaskLabel).Set(float64(v))
}
}
} }
log.Ctx(m.ctx).Info("update index metric", zap.Int("collectionNum", len(taskMetrics))) log.Ctx(m.ctx).Info("update index metric", zap.Int("collectionNum", len(taskMetrics)))
} }
@ -345,9 +352,10 @@ func checkParams(fieldIndex *model.Index, req *indexpb.CreateIndexRequest) bool
return !notEq return !notEq
} }
// CanCreateIndex currently is used in Unittest
func (m *indexMeta) CanCreateIndex(req *indexpb.CreateIndexRequest, isJson bool) (UniqueID, error) { func (m *indexMeta) CanCreateIndex(req *indexpb.CreateIndexRequest, isJson bool) (UniqueID, error) {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
indexes, ok := m.indexes[req.CollectionID] indexes, ok := m.indexes[req.CollectionID]
if !ok { if !ok {
@ -395,8 +403,8 @@ func (m *indexMeta) CanCreateIndex(req *indexpb.CreateIndexRequest, isJson bool)
// HasSameReq determine whether there are same indexing tasks. // HasSameReq determine whether there are same indexing tasks.
func (m *indexMeta) HasSameReq(req *indexpb.CreateIndexRequest) (bool, UniqueID) { func (m *indexMeta) HasSameReq(req *indexpb.CreateIndexRequest) (bool, UniqueID) {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
for _, fieldIndex := range m.indexes[req.CollectionID] { for _, fieldIndex := range m.indexes[req.CollectionID] {
if fieldIndex.IsDeleted { if fieldIndex.IsDeleted {
@ -420,8 +428,8 @@ func (m *indexMeta) HasSameReq(req *indexpb.CreateIndexRequest) (bool, UniqueID)
func (m *indexMeta) CreateIndex(ctx context.Context, index *model.Index) error { func (m *indexMeta) CreateIndex(ctx context.Context, index *model.Index) error {
log.Ctx(ctx).Info("meta update: CreateIndex", zap.Int64("collectionID", index.CollectionID), log.Ctx(ctx).Info("meta update: CreateIndex", zap.Int64("collectionID", index.CollectionID),
zap.Int64("fieldID", index.FieldID), zap.Int64("indexID", index.IndexID), zap.String("indexName", index.IndexName)) zap.Int64("fieldID", index.FieldID), zap.Int64("indexID", index.IndexID), zap.String("indexName", index.IndexName))
m.Lock() m.fieldIndexLock.Lock()
defer m.Unlock() defer m.fieldIndexLock.Unlock()
if err := m.catalog.CreateIndex(ctx, index); err != nil { if err := m.catalog.CreateIndex(ctx, index); err != nil {
log.Ctx(ctx).Error("meta update: CreateIndex save meta fail", zap.Int64("collectionID", index.CollectionID), log.Ctx(ctx).Error("meta update: CreateIndex save meta fail", zap.Int64("collectionID", index.CollectionID),
@ -437,8 +445,8 @@ func (m *indexMeta) CreateIndex(ctx context.Context, index *model.Index) error {
} }
func (m *indexMeta) AlterIndex(ctx context.Context, indexes ...*model.Index) error { func (m *indexMeta) AlterIndex(ctx context.Context, indexes ...*model.Index) error {
m.Lock() m.fieldIndexLock.Lock()
defer m.Unlock() defer m.fieldIndexLock.Unlock()
err := m.catalog.AlterIndexes(ctx, indexes) err := m.catalog.AlterIndexes(ctx, indexes)
if err != nil { if err != nil {
@ -454,10 +462,11 @@ func (m *indexMeta) AlterIndex(ctx context.Context, indexes ...*model.Index) err
// AddSegmentIndex adds the index meta corresponding the indexBuildID to meta table. // AddSegmentIndex adds the index meta corresponding the indexBuildID to meta table.
func (m *indexMeta) AddSegmentIndex(ctx context.Context, segIndex *model.SegmentIndex) error { func (m *indexMeta) AddSegmentIndex(ctx context.Context, segIndex *model.SegmentIndex) error {
m.Lock()
defer m.Unlock()
buildID := segIndex.BuildID buildID := segIndex.BuildID
m.keyLock.Lock(buildID)
defer m.keyLock.Unlock(buildID)
log.Ctx(ctx).Info("meta update: adding segment index", zap.Int64("collectionID", segIndex.CollectionID), log.Ctx(ctx).Info("meta update: adding segment index", zap.Int64("collectionID", segIndex.CollectionID),
zap.Int64("segmentID", segIndex.SegmentID), zap.Int64("indexID", segIndex.IndexID), zap.Int64("segmentID", segIndex.SegmentID), zap.Int64("indexID", segIndex.IndexID),
zap.Int64("buildID", buildID)) zap.Int64("buildID", buildID))
@ -477,8 +486,9 @@ func (m *indexMeta) AddSegmentIndex(ctx context.Context, segIndex *model.Segment
} }
func (m *indexMeta) GetIndexIDByName(collID int64, indexName string) map[int64]uint64 { func (m *indexMeta) GetIndexIDByName(collID int64, indexName string) map[int64]uint64 {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
indexID2CreateTs := make(map[int64]uint64) indexID2CreateTs := make(map[int64]uint64)
fieldIndexes, ok := m.indexes[collID] fieldIndexes, ok := m.indexes[collID]
@ -495,21 +505,22 @@ func (m *indexMeta) GetIndexIDByName(collID int64, indexName string) map[int64]u
} }
func (m *indexMeta) GetSegmentIndexState(collID, segmentID UniqueID, indexID UniqueID) *indexpb.SegmentIndexState { func (m *indexMeta) GetSegmentIndexState(collID, segmentID UniqueID, indexID UniqueID) *indexpb.SegmentIndexState {
m.RLock()
defer m.RUnlock()
state := &indexpb.SegmentIndexState{ state := &indexpb.SegmentIndexState{
SegmentID: segmentID, SegmentID: segmentID,
State: commonpb.IndexState_IndexStateNone, State: commonpb.IndexState_IndexStateNone,
FailReason: "", FailReason: "",
} }
m.fieldIndexLock.RLock()
fieldIndexes, ok := m.indexes[collID] fieldIndexes, ok := m.indexes[collID]
if !ok { if !ok {
state.FailReason = fmt.Sprintf("collection not exist with ID: %d", collID) state.FailReason = fmt.Sprintf("collection not exist with ID: %d", collID)
m.fieldIndexLock.RUnlock()
return state return state
} }
m.fieldIndexLock.RUnlock()
indexes, ok := m.segmentIndexes[segmentID] indexes, ok := m.segmentIndexes.Get(segmentID)
if !ok { if !ok {
state.State = commonpb.IndexState_Unissued state.State = commonpb.IndexState_Unissued
state.FailReason = fmt.Sprintf("segment index not exist with ID: %d", segmentID) state.FailReason = fmt.Sprintf("segment index not exist with ID: %d", segmentID)
@ -517,7 +528,7 @@ func (m *indexMeta) GetSegmentIndexState(collID, segmentID UniqueID, indexID Uni
} }
if index, ok := fieldIndexes[indexID]; ok && !index.IsDeleted { if index, ok := fieldIndexes[indexID]; ok && !index.IsDeleted {
if segIdx, ok := indexes[indexID]; ok { if segIdx, ok := indexes.Get(indexID); ok {
state.IndexName = index.IndexName state.IndexName = index.IndexName
state.State = segIdx.IndexState state.State = segIdx.IndexState
state.FailReason = segIdx.FailReason state.FailReason = segIdx.FailReason
@ -532,24 +543,24 @@ func (m *indexMeta) GetSegmentIndexState(collID, segmentID UniqueID, indexID Uni
} }
func (m *indexMeta) GetIndexedSegments(collectionID int64, segmentIDs, fieldIDs []UniqueID) []int64 { func (m *indexMeta) GetIndexedSegments(collectionID int64, segmentIDs, fieldIDs []UniqueID) []int64 {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock()
fieldIndexes, ok := m.indexes[collectionID] fieldIndexes, ok := m.indexes[collectionID]
if !ok { if !ok {
m.fieldIndexLock.RUnlock()
return nil return nil
} }
m.fieldIndexLock.RUnlock()
fieldIDSet := typeutil.NewUniqueSet(fieldIDs...) fieldIDSet := typeutil.NewUniqueSet(fieldIDs...)
checkSegmentState := func(indexes map[int64]*model.SegmentIndex) bool { checkSegmentState := func(indexes *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]) bool {
indexedFields := 0 indexedFields := 0
for indexID, index := range fieldIndexes { for indexID, index := range fieldIndexes {
if !fieldIDSet.Contain(index.FieldID) || index.IsDeleted { if !fieldIDSet.Contain(index.FieldID) || index.IsDeleted {
continue continue
} }
if segIdx, ok := indexes[indexID]; ok && segIdx.IndexState == commonpb.IndexState_Finished { if segIdx, ok := indexes.Get(indexID); ok && segIdx.IndexState == commonpb.IndexState_Finished {
indexedFields += 1 indexedFields += 1
} }
} }
@ -559,7 +570,7 @@ func (m *indexMeta) GetIndexedSegments(collectionID int64, segmentIDs, fieldIDs
ret := make([]int64, 0) ret := make([]int64, 0)
for _, sid := range segmentIDs { for _, sid := range segmentIDs {
if indexes, ok := m.segmentIndexes[sid]; ok { if indexes, ok := m.segmentIndexes.Get(sid); ok {
if checkSegmentState(indexes) { if checkSegmentState(indexes) {
ret = append(ret, sid) ret = append(ret, sid)
} }
@ -571,8 +582,8 @@ func (m *indexMeta) GetIndexedSegments(collectionID int64, segmentIDs, fieldIDs
// GetIndexesForCollection gets all indexes info with the specified collection. // GetIndexesForCollection gets all indexes info with the specified collection.
func (m *indexMeta) GetIndexesForCollection(collID UniqueID, indexName string) []*model.Index { func (m *indexMeta) GetIndexesForCollection(collID UniqueID, indexName string) []*model.Index {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
indexInfos := make([]*model.Index, 0) indexInfos := make([]*model.Index, 0)
for _, index := range m.indexes[collID] { for _, index := range m.indexes[collID] {
@ -587,8 +598,8 @@ func (m *indexMeta) GetIndexesForCollection(collID UniqueID, indexName string) [
} }
func (m *indexMeta) GetFieldIndexes(collID, fieldID UniqueID, indexName string) []*model.Index { func (m *indexMeta) GetFieldIndexes(collID, fieldID UniqueID, indexName string) []*model.Index {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
indexInfos := make([]*model.Index, 0) indexInfos := make([]*model.Index, 0)
for _, index := range m.indexes[collID] { for _, index := range m.indexes[collID] {
@ -607,8 +618,8 @@ func (m *indexMeta) MarkIndexAsDeleted(ctx context.Context, collID UniqueID, ind
log.Ctx(ctx).Info("IndexCoord metaTable MarkIndexAsDeleted", zap.Int64("collectionID", collID), log.Ctx(ctx).Info("IndexCoord metaTable MarkIndexAsDeleted", zap.Int64("collectionID", collID),
zap.Int64s("indexIDs", indexIDs)) zap.Int64s("indexIDs", indexIDs))
m.Lock() m.fieldIndexLock.Lock()
defer m.Unlock() defer m.fieldIndexLock.Unlock()
fieldIndexes, ok := m.indexes[collID] fieldIndexes, ok := m.indexes[collID]
if !ok { if !ok {
@ -641,23 +652,23 @@ func (m *indexMeta) MarkIndexAsDeleted(ctx context.Context, collID UniqueID, ind
} }
func (m *indexMeta) IsUnIndexedSegment(collectionID UniqueID, segID UniqueID) bool { func (m *indexMeta) IsUnIndexedSegment(collectionID UniqueID, segID UniqueID) bool {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock()
fieldIndexes, ok := m.indexes[collectionID] fieldIndexes, ok := m.indexes[collectionID]
if !ok { if !ok {
m.fieldIndexLock.RUnlock()
return false return false
} }
m.fieldIndexLock.RUnlock()
// the segment should be unindexed status if the fieldIndexes is not nil // the segment should be unindexed status if the fieldIndexes is not nil
segIndexInfos, ok := m.segmentIndexes[segID] segIndexInfos, ok := m.segmentIndexes.Get(segID)
if !ok || len(segIndexInfos) == 0 { if !ok || segIndexInfos.Len() == 0 {
return true return true
} }
for _, index := range fieldIndexes { for _, index := range fieldIndexes {
if _, ok := segIndexInfos[index.IndexID]; !index.IsDeleted { if !index.IsDeleted {
if !ok { if _, ok := segIndexInfos.Get(index.IndexID); !ok {
// the segment should be unindexed status if the segment index is not found within field indexes // the segment should be unindexed status if the segment index is not found within field indexes
return true return true
} }
@ -668,8 +679,6 @@ func (m *indexMeta) IsUnIndexedSegment(collectionID UniqueID, segID UniqueID) bo
} }
func (m *indexMeta) GetSegmentsIndexes(collectionID UniqueID, segIDs []UniqueID) map[int64]map[UniqueID]*model.SegmentIndex { func (m *indexMeta) GetSegmentsIndexes(collectionID UniqueID, segIDs []UniqueID) map[int64]map[UniqueID]*model.SegmentIndex {
m.RLock()
defer m.RUnlock()
segmentsIndexes := make(map[int64]map[UniqueID]*model.SegmentIndex) segmentsIndexes := make(map[int64]map[UniqueID]*model.SegmentIndex)
for _, segmentID := range segIDs { for _, segmentID := range segIDs {
segmentsIndexes[segmentID] = m.getSegmentIndexes(collectionID, segmentID) segmentsIndexes[segmentID] = m.getSegmentIndexes(collectionID, segmentID)
@ -678,16 +687,14 @@ func (m *indexMeta) GetSegmentsIndexes(collectionID UniqueID, segIDs []UniqueID)
} }
func (m *indexMeta) GetSegmentIndexes(collectionID UniqueID, segID UniqueID) map[UniqueID]*model.SegmentIndex { func (m *indexMeta) GetSegmentIndexes(collectionID UniqueID, segID UniqueID) map[UniqueID]*model.SegmentIndex {
m.RLock()
defer m.RUnlock()
return m.getSegmentIndexes(collectionID, segID) return m.getSegmentIndexes(collectionID, segID)
} }
// Note: thread-unsafe, don't call it outside indexMeta // Note: thread-unsafe, don't call it outside indexMeta
func (m *indexMeta) getSegmentIndexes(collectionID UniqueID, segID UniqueID) map[UniqueID]*model.SegmentIndex { func (m *indexMeta) getSegmentIndexes(collectionID UniqueID, segID UniqueID) map[UniqueID]*model.SegmentIndex {
ret := make(map[UniqueID]*model.SegmentIndex, 0) ret := make(map[UniqueID]*model.SegmentIndex, 0)
segIndexInfos, ok := m.segmentIndexes[segID] segIndexInfos, ok := m.segmentIndexes.Get(segID)
if !ok || len(segIndexInfos) == 0 { if !ok || segIndexInfos.Len() == 0 {
return ret return ret
} }
@ -696,7 +703,7 @@ func (m *indexMeta) getSegmentIndexes(collectionID UniqueID, segID UniqueID) map
return ret return ret
} }
for _, segIdx := range segIndexInfos { for _, segIdx := range segIndexInfos.Values() {
if index, ok := fieldIndexes[segIdx.IndexID]; ok && !index.IsDeleted { if index, ok := fieldIndexes[segIdx.IndexID]; ok && !index.IsDeleted {
ret[segIdx.IndexID] = model.CloneSegmentIndex(segIdx) ret[segIdx.IndexID] = model.CloneSegmentIndex(segIdx)
} }
@ -705,8 +712,8 @@ func (m *indexMeta) getSegmentIndexes(collectionID UniqueID, segID UniqueID) map
} }
func (m *indexMeta) GetFieldIDByIndexID(collID, indexID UniqueID) UniqueID { func (m *indexMeta) GetFieldIDByIndexID(collID, indexID UniqueID) UniqueID {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
if fieldIndexes, ok := m.indexes[collID]; ok { if fieldIndexes, ok := m.indexes[collID]; ok {
if index, ok := fieldIndexes[indexID]; ok { if index, ok := fieldIndexes[indexID]; ok {
@ -717,8 +724,9 @@ func (m *indexMeta) GetFieldIDByIndexID(collID, indexID UniqueID) UniqueID {
} }
func (m *indexMeta) GetIndexNameByID(collID, indexID UniqueID) string { func (m *indexMeta) GetIndexNameByID(collID, indexID UniqueID) string {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
if fieldIndexes, ok := m.indexes[collID]; ok { if fieldIndexes, ok := m.indexes[collID]; ok {
if index, ok := fieldIndexes[indexID]; ok { if index, ok := fieldIndexes[indexID]; ok {
return index.IndexName return index.IndexName
@ -728,8 +736,8 @@ func (m *indexMeta) GetIndexNameByID(collID, indexID UniqueID) string {
} }
func (m *indexMeta) GetIndexParams(collID, indexID UniqueID) []*commonpb.KeyValuePair { func (m *indexMeta) GetIndexParams(collID, indexID UniqueID) []*commonpb.KeyValuePair {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
fieldIndexes, ok := m.indexes[collID] fieldIndexes, ok := m.indexes[collID]
if !ok { if !ok {
@ -749,8 +757,8 @@ func (m *indexMeta) GetIndexParams(collID, indexID UniqueID) []*commonpb.KeyValu
} }
func (m *indexMeta) GetTypeParams(collID, indexID UniqueID) []*commonpb.KeyValuePair { func (m *indexMeta) GetTypeParams(collID, indexID UniqueID) []*commonpb.KeyValuePair {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
fieldIndexes, ok := m.indexes[collID] fieldIndexes, ok := m.indexes[collID]
if !ok { if !ok {
@ -770,9 +778,6 @@ func (m *indexMeta) GetTypeParams(collID, indexID UniqueID) []*commonpb.KeyValue
} }
func (m *indexMeta) GetIndexJob(buildID UniqueID) (*model.SegmentIndex, bool) { func (m *indexMeta) GetIndexJob(buildID UniqueID) (*model.SegmentIndex, bool) {
m.RLock()
defer m.RUnlock()
segIdx, ok := m.segmentBuildInfo.Get(buildID) segIdx, ok := m.segmentBuildInfo.Get(buildID)
if ok { if ok {
return model.CloneSegmentIndex(segIdx), true return model.CloneSegmentIndex(segIdx), true
@ -782,13 +787,9 @@ func (m *indexMeta) GetIndexJob(buildID UniqueID) (*model.SegmentIndex, bool) {
} }
func (m *indexMeta) IsIndexExist(collID, indexID UniqueID) bool { func (m *indexMeta) IsIndexExist(collID, indexID UniqueID) bool {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
return m.isIndexExist(collID, indexID)
}
func (m *indexMeta) isIndexExist(collID, indexID UniqueID) bool {
fieldIndexes, ok := m.indexes[collID] fieldIndexes, ok := m.indexes[collID]
if !ok { if !ok {
return false return false
@ -802,8 +803,8 @@ func (m *indexMeta) isIndexExist(collID, indexID UniqueID) bool {
// UpdateVersion updates the version and nodeID of the index meta, whenever the task is built once, the version will be updated once. // UpdateVersion updates the version and nodeID of the index meta, whenever the task is built once, the version will be updated once.
func (m *indexMeta) UpdateVersion(buildID, nodeID UniqueID) error { func (m *indexMeta) UpdateVersion(buildID, nodeID UniqueID) error {
m.Lock() m.keyLock.Lock(buildID)
defer m.Unlock() defer m.keyLock.Unlock(buildID)
log.Ctx(m.ctx).Info("IndexCoord metaTable UpdateVersion receive", zap.Int64("buildID", buildID), zap.Int64("nodeID", nodeID)) log.Ctx(m.ctx).Info("IndexCoord metaTable UpdateVersion receive", zap.Int64("buildID", buildID), zap.Int64("nodeID", nodeID))
segIdx, ok := m.segmentBuildInfo.Get(buildID) segIdx, ok := m.segmentBuildInfo.Get(buildID)
@ -821,8 +822,8 @@ func (m *indexMeta) UpdateVersion(buildID, nodeID UniqueID) error {
} }
func (m *indexMeta) FinishTask(taskInfo *workerpb.IndexTaskInfo) error { func (m *indexMeta) FinishTask(taskInfo *workerpb.IndexTaskInfo) error {
m.Lock() m.keyLock.Lock(taskInfo.GetBuildID())
defer m.Unlock() defer m.keyLock.Unlock(taskInfo.GetBuildID())
segIdx, ok := m.segmentBuildInfo.Get(taskInfo.GetBuildID()) segIdx, ok := m.segmentBuildInfo.Get(taskInfo.GetBuildID())
if !ok { if !ok {
@ -854,8 +855,8 @@ func (m *indexMeta) FinishTask(taskInfo *workerpb.IndexTaskInfo) error {
} }
func (m *indexMeta) DeleteTask(buildID int64) error { func (m *indexMeta) DeleteTask(buildID int64) error {
m.Lock() m.keyLock.Lock(buildID)
defer m.Unlock() defer m.keyLock.Unlock(buildID)
segIdx, ok := m.segmentBuildInfo.Get(buildID) segIdx, ok := m.segmentBuildInfo.Get(buildID)
if !ok { if !ok {
@ -878,8 +879,8 @@ func (m *indexMeta) DeleteTask(buildID int64) error {
// BuildIndex set the index state to be InProgress. It means IndexNode is building the index. // BuildIndex set the index state to be InProgress. It means IndexNode is building the index.
func (m *indexMeta) BuildIndex(buildID UniqueID) error { func (m *indexMeta) BuildIndex(buildID UniqueID) error {
m.Lock() m.keyLock.Lock(buildID)
defer m.Unlock() defer m.keyLock.Unlock(buildID)
segIdx, ok := m.segmentBuildInfo.Get(buildID) segIdx, ok := m.segmentBuildInfo.Get(buildID)
if !ok { if !ok {
@ -905,21 +906,18 @@ func (m *indexMeta) BuildIndex(buildID UniqueID) error {
} }
func (m *indexMeta) GetAllSegIndexes() map[int64]*model.SegmentIndex { func (m *indexMeta) GetAllSegIndexes() map[int64]*model.SegmentIndex {
m.RLock()
defer m.RUnlock()
tasks := m.segmentBuildInfo.List() tasks := m.segmentBuildInfo.List()
segIndexes := make(map[int64]*model.SegmentIndex, len(tasks)) segIndexes := make(map[int64]*model.SegmentIndex, len(tasks))
for buildID, segIndex := range tasks { for _, segIndex := range tasks {
segIndexes[buildID] = segIndex segIndexes[segIndex.BuildID] = segIndex
} }
return segIndexes return segIndexes
} }
// SetStoredIndexFileSizeMetric returns the total index files size of all segment for each collection. // SetStoredIndexFileSizeMetric returns the total index files size of all segment for each collection.
func (m *indexMeta) SetStoredIndexFileSizeMetric(collections map[UniqueID]*collectionInfo) uint64 { func (m *indexMeta) SetStoredIndexFileSizeMetric(collections map[UniqueID]*collectionInfo) uint64 {
m.RLock() m.fieldIndexLock.Lock()
defer m.RUnlock() defer m.fieldIndexLock.Unlock()
var total uint64 var total uint64
metrics.DataCoordStoredIndexFilesSize.Reset() metrics.DataCoordStoredIndexFilesSize.Reset()
@ -935,30 +933,41 @@ func (m *indexMeta) SetStoredIndexFileSizeMetric(collections map[UniqueID]*colle
return total return total
} }
func (m *indexMeta) RemoveSegmentIndex(ctx context.Context, collID, partID, segID, indexID, buildID UniqueID) error { func (m *indexMeta) removeSegmentIndex(ctx context.Context, collID, partID, segID, indexID, buildID UniqueID) error {
m.Lock()
defer m.Unlock()
err := m.catalog.DropSegmentIndex(ctx, collID, partID, segID, buildID) err := m.catalog.DropSegmentIndex(ctx, collID, partID, segID, buildID)
if err != nil { if err != nil {
return err return err
} }
if _, ok := m.segmentIndexes[segID]; ok { segIndexes, ok := m.segmentIndexes.Get(segID)
delete(m.segmentIndexes[segID], indexID) if ok {
segIndexes.Remove(indexID)
m.segmentIndexes.Insert(segID, segIndexes)
} }
if segIndexes.Len() == 0 {
if len(m.segmentIndexes[segID]) == 0 { m.segmentIndexes.Remove(segID)
delete(m.segmentIndexes, segID)
} }
m.segmentBuildInfo.Remove(buildID) m.segmentBuildInfo.Remove(buildID)
return nil return nil
} }
func (m *indexMeta) RemoveSegmentIndex(ctx context.Context, collID, partID, segID, indexID, buildID UniqueID) error {
return m.removeSegmentIndex(ctx, collID, partID, segID, indexID, buildID)
}
func (m *indexMeta) RemoveSegmentIndexByID(ctx context.Context, buildID UniqueID) error {
segIdx, ok := m.segmentBuildInfo.Get(buildID)
if !ok {
return nil
}
return m.removeSegmentIndex(ctx, segIdx.CollectionID, segIdx.PartitionID, segIdx.SegmentID, segIdx.IndexID, buildID)
}
func (m *indexMeta) GetDeletedIndexes() []*model.Index { func (m *indexMeta) GetDeletedIndexes() []*model.Index {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
deletedIndexes := make([]*model.Index, 0) deletedIndexes := make([]*model.Index, 0)
for _, fieldIndexes := range m.indexes { for _, fieldIndexes := range m.indexes {
@ -972,8 +981,8 @@ func (m *indexMeta) GetDeletedIndexes() []*model.Index {
} }
func (m *indexMeta) RemoveIndex(ctx context.Context, collID, indexID UniqueID) error { func (m *indexMeta) RemoveIndex(ctx context.Context, collID, indexID UniqueID) error {
m.Lock() m.fieldIndexLock.Lock()
defer m.Unlock() defer m.fieldIndexLock.Unlock()
log.Ctx(ctx).Info("IndexCoord meta table remove index", zap.Int64("collectionID", collID), zap.Int64("indexID", indexID)) log.Ctx(ctx).Info("IndexCoord meta table remove index", zap.Int64("collectionID", collID), zap.Int64("indexID", indexID))
err := m.catalog.DropIndex(ctx, collID, indexID) err := m.catalog.DropIndex(ctx, collID, indexID)
if err != nil { if err != nil {
@ -995,9 +1004,6 @@ func (m *indexMeta) RemoveIndex(ctx context.Context, collID, indexID UniqueID) e
} }
func (m *indexMeta) CheckCleanSegmentIndex(buildID UniqueID) (bool, *model.SegmentIndex) { func (m *indexMeta) CheckCleanSegmentIndex(buildID UniqueID) (bool, *model.SegmentIndex) {
m.RLock()
defer m.RUnlock()
if segIndex, ok := m.segmentBuildInfo.Get(buildID); ok { if segIndex, ok := m.segmentBuildInfo.Get(buildID); ok {
if segIndex.IndexState == commonpb.IndexState_Finished { if segIndex.IndexState == commonpb.IndexState_Finished {
return true, model.CloneSegmentIndex(segIndex) return true, model.CloneSegmentIndex(segIndex)
@ -1008,9 +1014,6 @@ func (m *indexMeta) CheckCleanSegmentIndex(buildID UniqueID) (bool, *model.Segme
} }
func (m *indexMeta) getSegmentsIndexStates(collectionID UniqueID, segmentIDs []UniqueID) map[int64]map[int64]*indexpb.SegmentIndexState { func (m *indexMeta) getSegmentsIndexStates(collectionID UniqueID, segmentIDs []UniqueID) map[int64]map[int64]*indexpb.SegmentIndexState {
m.RLock()
defer m.RUnlock()
ret := make(map[int64]map[int64]*indexpb.SegmentIndexState, 0) ret := make(map[int64]map[int64]*indexpb.SegmentIndexState, 0)
fieldIndexes, ok := m.indexes[collectionID] fieldIndexes, ok := m.indexes[collectionID]
if !ok { if !ok {
@ -1019,12 +1022,12 @@ func (m *indexMeta) getSegmentsIndexStates(collectionID UniqueID, segmentIDs []U
for _, segID := range segmentIDs { for _, segID := range segmentIDs {
ret[segID] = make(map[int64]*indexpb.SegmentIndexState) ret[segID] = make(map[int64]*indexpb.SegmentIndexState)
segIndexInfos, ok := m.segmentIndexes[segID] segIndexInfos, ok := m.segmentIndexes.Get(segID)
if !ok || len(segIndexInfos) == 0 { if !ok || segIndexInfos.Len() == 0 {
continue continue
} }
for _, segIdx := range segIndexInfos { for _, segIdx := range segIndexInfos.Values() {
if index, ok := fieldIndexes[segIdx.IndexID]; ok && !index.IsDeleted { if index, ok := fieldIndexes[segIdx.IndexID]; ok && !index.IsDeleted {
ret[segID][segIdx.IndexID] = &indexpb.SegmentIndexState{ ret[segID][segIdx.IndexID] = &indexpb.SegmentIndexState{
SegmentID: segID, SegmentID: segID,
@ -1077,8 +1080,9 @@ func (m *indexMeta) AreAllDiskIndex(collectionID int64, schema *schemapb.Collect
} }
func (m *indexMeta) HasIndex(collectionID int64) bool { func (m *indexMeta) HasIndex(collectionID int64) bool {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
indexes, ok := m.indexes[collectionID] indexes, ok := m.indexes[collectionID]
if ok { if ok {
for _, index := range indexes { for _, index := range indexes {
@ -1100,8 +1104,8 @@ func (m *indexMeta) TaskStatsJSON() string {
} }
func (m *indexMeta) GetIndexJSON(collectionID int64) string { func (m *indexMeta) GetIndexJSON(collectionID int64) string {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock() defer m.fieldIndexLock.RUnlock()
var indexMetrics []*metricsinfo.Index var indexMetrics []*metricsinfo.Index
for collID, indexes := range m.indexes { for collID, indexes := range m.indexes {
@ -1131,24 +1135,25 @@ func (m *indexMeta) GetIndexJSON(collectionID int64) string {
} }
func (m *indexMeta) GetSegmentIndexedFields(collectionID UniqueID, segmentID UniqueID) (bool, []*metricsinfo.IndexedField) { func (m *indexMeta) GetSegmentIndexedFields(collectionID UniqueID, segmentID UniqueID) (bool, []*metricsinfo.IndexedField) {
m.RLock() m.fieldIndexLock.RLock()
defer m.RUnlock()
fieldIndexes, ok := m.indexes[collectionID] fieldIndexes, ok := m.indexes[collectionID]
if !ok { if !ok {
// the segment should be unindexed status if the collection has no indexes // the segment should be unindexed status if the collection has no indexes
m.fieldIndexLock.RUnlock()
return false, []*metricsinfo.IndexedField{} return false, []*metricsinfo.IndexedField{}
} }
m.fieldIndexLock.RUnlock()
// the segment should be unindexed status if the segment indexes is not found // the segment should be unindexed status if the segment indexes is not found
segIndexInfos, ok := m.segmentIndexes[segmentID] segIndexInfos, ok := m.segmentIndexes.Get(segmentID)
if !ok || len(segIndexInfos) == 0 { if !ok || segIndexInfos.Len() == 0 {
return false, []*metricsinfo.IndexedField{} return false, []*metricsinfo.IndexedField{}
} }
isIndexed := true isIndexed := true
var segmentIndexes []*metricsinfo.IndexedField var segmentIndexes []*metricsinfo.IndexedField
for _, index := range fieldIndexes { for _, index := range fieldIndexes {
if si, ok := segIndexInfos[index.IndexID]; !index.IsDeleted { if si, ok := segIndexInfos.Get(index.IndexID); !index.IsDeleted {
buildID := int64(-1) buildID := int64(-1)
if !ok { if !ok {
// the segment should be unindexed status if the segment index is not found within field indexes // the segment should be unindexed status if the segment index is not found within field indexes

240
internal/datacoord/index_meta_test.go
View File

@ -19,7 +19,6 @@ package datacoord
import ( import (
"context" "context"
"sync"
"testing" "testing"
"time" "time"
@ -37,7 +36,9 @@ import (
"github.com/milvus-io/milvus/pkg/v2/common" "github.com/milvus-io/milvus/pkg/v2/common"
"github.com/milvus-io/milvus/pkg/v2/proto/indexpb" "github.com/milvus-io/milvus/pkg/v2/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v2/proto/workerpb" "github.com/milvus-io/milvus/pkg/v2/proto/workerpb"
"github.com/milvus-io/milvus/pkg/v2/util/lock"
"github.com/milvus-io/milvus/pkg/v2/util/metricsinfo" "github.com/milvus-io/milvus/pkg/v2/util/metricsinfo"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
) )
func TestReloadFromKV(t *testing.T) { func TestReloadFromKV(t *testing.T) {
@ -442,12 +443,12 @@ func TestMeta_HasSameReq(t *testing.T) {
func newSegmentIndexMeta(catalog metastore.DataCoordCatalog) *indexMeta { func newSegmentIndexMeta(catalog metastore.DataCoordCatalog) *indexMeta {
return &indexMeta{ return &indexMeta{
RWMutex: sync.RWMutex{}, keyLock: lock.NewKeyLock[UniqueID](),
ctx: context.Background(), ctx: context.Background(),
catalog: catalog, catalog: catalog,
indexes: make(map[UniqueID]map[UniqueID]*model.Index), indexes: make(map[UniqueID]map[UniqueID]*model.Index),
segmentBuildInfo: newSegmentIndexBuildInfo(), segmentBuildInfo: newSegmentIndexBuildInfo(),
segmentIndexes: make(map[UniqueID]map[UniqueID]*model.SegmentIndex), segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
} }
} }
@ -516,9 +517,8 @@ func TestMeta_AddSegmentIndex(t *testing.T) {
).Return(errors.New("fail")) ).Return(errors.New("fail"))
m := newSegmentIndexMeta(ec) m := newSegmentIndexMeta(ec)
m.segmentIndexes = map[UniqueID]map[UniqueID]*model.SegmentIndex{ m.segmentIndexes = typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]]()
1: make(map[UniqueID]*model.SegmentIndex, 0), m.segmentIndexes.Insert(1, typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]())
}
segmentIndex := &model.SegmentIndex{ segmentIndex := &model.SegmentIndex{
SegmentID: 1, SegmentID: 1,
@ -630,9 +630,8 @@ func TestMeta_GetSegmentIndexState(t *testing.T) {
metakv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(nil, nil, nil).Maybe() metakv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(nil, nil, nil).Maybe()
m := newSegmentIndexMeta(&datacoord.Catalog{MetaKv: metakv}) m := newSegmentIndexMeta(&datacoord.Catalog{MetaKv: metakv})
m.segmentIndexes = map[UniqueID]map[UniqueID]*model.SegmentIndex{ m.segmentIndexes = typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]]()
segID: make(map[UniqueID]*model.SegmentIndex, 0), m.segmentIndexes.Insert(1, typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]())
}
t.Run("collection has no index", func(t *testing.T) { t.Run("collection has no index", func(t *testing.T) {
state := m.GetSegmentIndexState(collID, segID, indexID) state := m.GetSegmentIndexState(collID, segID, indexID)
@ -735,26 +734,25 @@ func TestMeta_GetIndexedSegment(t *testing.T) {
) )
m := newSegmentIndexMeta(nil) m := newSegmentIndexMeta(nil)
m.segmentIndexes = map[UniqueID]map[UniqueID]*model.SegmentIndex{ segIdxes := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segID: { segIdxes.Insert(indexID, &model.SegmentIndex{
indexID: { SegmentID: segID,
SegmentID: segID, CollectionID: collID,
CollectionID: collID, PartitionID: partID,
PartitionID: partID, NumRows: 1025,
NumRows: 1025, IndexID: indexID,
IndexID: indexID, BuildID: buildID,
BuildID: buildID, NodeID: nodeID,
NodeID: nodeID, IndexVersion: 1,
IndexVersion: 1, IndexState: commonpb.IndexState_Finished,
IndexState: commonpb.IndexState_Finished, FailReason: "",
FailReason: "", IsDeleted: false,
IsDeleted: false, CreatedUTCTime: 10,
CreatedUTCTime: 10, IndexFileKeys: nil,
IndexFileKeys: nil, IndexSerializedSize: 0,
IndexSerializedSize: 0, })
},
}, m.segmentIndexes.Insert(segID, segIdxes)
}
m.indexes = map[UniqueID]map[UniqueID]*model.Index{ m.indexes = map[UniqueID]map[UniqueID]*model.Index{
collID: { collID: {
indexID: { indexID: {
@ -891,27 +889,17 @@ func TestMeta_GetSegmentIndexes(t *testing.T) {
t.Run("no index exist- field index empty", func(t *testing.T) { t.Run("no index exist- field index empty", func(t *testing.T) {
m := newSegmentIndexMeta(nil) m := newSegmentIndexMeta(nil)
m.segmentIndexes = map[UniqueID]map[UniqueID]*model.SegmentIndex{ segIdxes := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
1: { segIdxes.Insert(indexID, &model.SegmentIndex{})
1: &model.SegmentIndex{}, m.segmentIndexes.Insert(segID, segIdxes)
},
}
segIndexes := m.GetSegmentIndexes(collID, 1) segIndexes := m.GetSegmentIndexes(collID, 1)
assert.Equal(t, 0, len(segIndexes)) assert.Equal(t, 0, len(segIndexes))
}) })
t.Run("index exists", func(t *testing.T) { t.Run("index exists", func(t *testing.T) {
m := &indexMeta{ m := &indexMeta{
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
segID: {
indexID: &model.SegmentIndex{
CollectionID: collID,
SegmentID: segID,
IndexID: indexID,
IndexState: commonpb.IndexState_Finished,
},
},
},
indexes: map[UniqueID]map[UniqueID]*model.Index{ indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: { collID: {
indexID: { indexID: {
@ -930,6 +918,14 @@ func TestMeta_GetSegmentIndexes(t *testing.T) {
}, },
}, },
} }
segIdxes := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdxes.Insert(indexID, &model.SegmentIndex{
CollectionID: collID,
SegmentID: segID,
IndexID: indexID,
IndexState: commonpb.IndexState_Finished,
})
m.segmentIndexes.Insert(segID, segIdxes)
segIndexes := m.GetSegmentIndexes(collID, segID) segIndexes := m.GetSegmentIndexes(collID, segID)
assert.Equal(t, 1, len(segIndexes)) assert.Equal(t, 1, len(segIndexes))
@ -1200,28 +1196,10 @@ func updateSegmentIndexMeta(t *testing.T) *indexMeta {
IndexSerializedSize: 0, IndexSerializedSize: 0,
}) })
return &indexMeta{ m := &indexMeta{
catalog: sc, catalog: sc,
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ keyLock: lock.NewKeyLock[UniqueID](),
segID: { segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1025,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 0,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
},
},
},
indexes: map[UniqueID]map[UniqueID]*model.Index{ indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: { collID: {
indexID: { indexID: {
@ -1241,6 +1219,25 @@ func updateSegmentIndexMeta(t *testing.T) *indexMeta {
}, },
segmentBuildInfo: indexBuildInfo, segmentBuildInfo: indexBuildInfo,
} }
segIdxes := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdxes.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1025,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 0,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
})
m.segmentIndexes.Insert(segID, segIdxes)
return m
} }
func TestMeta_UpdateVersion(t *testing.T) { func TestMeta_UpdateVersion(t *testing.T) {
@ -1346,10 +1343,11 @@ func TestUpdateSegmentIndexNotExists(t *testing.T) {
}) })
}) })
assert.Equal(t, 1, len(m.segmentIndexes)) assert.Equal(t, 1, m.segmentIndexes.Len())
segmentIdx := m.segmentIndexes[1] segmentIdx, ok := m.segmentIndexes.Get(1)
assert.Equal(t, 1, len(segmentIdx)) assert.True(t, ok)
_, ok := segmentIdx[2] assert.Equal(t, 1, segmentIdx.Len())
_, ok = segmentIdx.Get(2)
assert.True(t, ok) assert.True(t, ok)
} }
@ -1488,19 +1486,67 @@ func TestRemoveSegmentIndex(t *testing.T) {
Return(nil) Return(nil)
m := &indexMeta{ m := &indexMeta{
catalog: catalog, catalog: catalog,
segmentIndexes: map[int64]map[int64]*model.SegmentIndex{ segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
segID: {
indexID: &model.SegmentIndex{},
},
},
segmentBuildInfo: newSegmentIndexBuildInfo(), segmentBuildInfo: newSegmentIndexBuildInfo(),
} }
m.segmentIndexes.Insert(segID, typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]())
err := m.RemoveSegmentIndex(context.TODO(), collID, partID, segID, indexID, buildID) err := m.RemoveSegmentIndex(context.TODO(), collID, partID, segID, indexID, buildID)
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, len(m.segmentIndexes), 0) assert.Equal(t, 0, m.segmentIndexes.Len())
assert.Equal(t, len(m.segmentBuildInfo.List()), 0)
})
}
func TestRemoveSegmentIndexByID(t *testing.T) {
t.Run("drop segment index fail", func(t *testing.T) {
expectedErr := errors.New("error")
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.EXPECT().
DropSegmentIndex(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(expectedErr)
catalog.EXPECT().CreateSegmentIndex(mock.Anything, mock.Anything).Return(nil)
m := newSegmentIndexMeta(catalog)
err := m.AddSegmentIndex(context.TODO(), &model.SegmentIndex{
SegmentID: 3,
CollectionID: 1,
PartitionID: 2,
NumRows: 1024,
IndexID: 1,
BuildID: 4,
})
assert.NoError(t, err)
err = m.RemoveSegmentIndexByID(context.TODO(), 4)
assert.Error(t, err)
assert.EqualError(t, err, "error")
})
t.Run("remove segment index ok", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.EXPECT().
DropSegmentIndex(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(nil)
catalog.EXPECT().CreateSegmentIndex(mock.Anything, mock.Anything).Return(nil)
m := newSegmentIndexMeta(catalog)
err := m.AddSegmentIndex(context.TODO(), &model.SegmentIndex{
SegmentID: 3,
CollectionID: 1,
PartitionID: 2,
NumRows: 1024,
IndexID: 1,
BuildID: 4,
})
assert.NoError(t, err)
err = m.RemoveSegmentIndexByID(context.TODO(), 4)
assert.NoError(t, err)
assert.Equal(t, m.segmentIndexes.Len(), 0)
assert.Equal(t, len(m.segmentBuildInfo.List()), 0) assert.Equal(t, len(m.segmentBuildInfo.List()), 0)
}) })
} }
@ -1639,27 +1685,25 @@ func TestMeta_GetSegmentIndexStatus(t *testing.T) {
}, },
}, },
} }
m.segmentIndexes = map[UniqueID]map[UniqueID]*model.SegmentIndex{ m.segmentIndexes = typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]]()
segID: { segIdx := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
indexID: { segIdx.Insert(indexID, &model.SegmentIndex{
SegmentID: segID, SegmentID: segID,
CollectionID: collID, CollectionID: collID,
PartitionID: partID, PartitionID: partID,
NumRows: 10250, NumRows: 10250,
IndexID: indexID, IndexID: indexID,
BuildID: buildID, BuildID: buildID,
NodeID: 1, NodeID: 1,
IndexVersion: 0, IndexVersion: 0,
IndexState: commonpb.IndexState_Finished, IndexState: commonpb.IndexState_Finished,
FailReason: "", FailReason: "",
IsDeleted: false, IsDeleted: false,
CreatedUTCTime: 12, CreatedUTCTime: 12,
IndexFileKeys: nil, IndexFileKeys: nil,
IndexSerializedSize: 0, IndexSerializedSize: 0,
}, })
}, m.segmentIndexes.Insert(segID, segIdx)
segID + 1: {},
}
t.Run("index exists", func(t *testing.T) { t.Run("index exists", func(t *testing.T) {
isIndexed, segmentIndexes := m.GetSegmentIndexedFields(collID, segID) isIndexed, segmentIndexes := m.GetSegmentIndexedFields(collID, segID)

880
internal/datacoord/index_service_test.go
View File

@ -47,6 +47,7 @@ import (
"github.com/milvus-io/milvus/pkg/v2/proto/indexpb" "github.com/milvus-io/milvus/pkg/v2/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v2/proto/workerpb" "github.com/milvus-io/milvus/pkg/v2/proto/workerpb"
"github.com/milvus-io/milvus/pkg/v2/util/merr" "github.com/milvus-io/milvus/pkg/v2/util/merr"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
) )
func TestServerId(t *testing.T) { func TestServerId(t *testing.T) {
@ -398,159 +399,159 @@ func TestServer_AlterIndex(t *testing.T) {
}, },
}, },
}, },
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
segID: {
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 1: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 3: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 4: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 5: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
},
segID - 1: {
indexID: {
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
},
indexID + 1: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
},
indexID + 3: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
CreatedUTCTime: createTS,
},
indexID + 4: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
CreatedUTCTime: createTS,
},
indexID + 5: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
},
},
},
} }
segIdx1 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx1.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+1, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+3, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+4, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+5, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
indexMeta.segmentIndexes.Insert(segID, segIdx1)
segIdx2 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx2.Insert(indexID, &model.SegmentIndex{
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
})
segIdx2.Insert(indexID+1, &model.SegmentIndex{
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
})
segIdx2.Insert(indexID+3, &model.SegmentIndex{
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
CreatedUTCTime: createTS,
})
segIdx2.Insert(indexID+4, &model.SegmentIndex{
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
CreatedUTCTime: createTS,
})
segIdx2.Insert(indexID+5, &model.SegmentIndex{
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
})
indexMeta.segmentIndexes.Insert(segID-1, segIdx2)
mockHandler := NewNMockHandler(t) mockHandler := NewNMockHandler(t)
@ -804,7 +805,7 @@ func TestServer_GetIndexState(t *testing.T) {
}, },
}, },
}, },
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{}, segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
}, },
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
@ -861,31 +862,29 @@ func TestServer_GetIndexState(t *testing.T) {
}, },
}, },
}, },
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
segID: {
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 3000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_IndexStateNone,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
},
},
}, },
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
} }
segIdx := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 3000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_IndexStateNone,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 0,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
s.meta.indexMeta.segmentIndexes.Insert(segID, segIdx)
for id, segment := range segments { for id, segment := range segments {
s.meta.segments.SetSegment(id, segment) s.meta.segments.SetSegment(id, segment)
} }
@ -1404,158 +1403,7 @@ func TestServer_DescribeIndex(t *testing.T) {
}, },
}, },
}, },
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
segID: {
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 1: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 3: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 4: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 5: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
},
segID - 1: {
indexID: {
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
},
indexID + 1: {
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
},
indexID + 3: {
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
CreatedUTCTime: createTS,
},
indexID + 4: {
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
CreatedUTCTime: createTS,
},
indexID + 5: {
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
},
},
},
}, },
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
@ -1563,6 +1411,158 @@ func TestServer_DescribeIndex(t *testing.T) {
allocator: mock0Allocator, allocator: mock0Allocator,
notifyIndexChan: make(chan UniqueID, 1), notifyIndexChan: make(chan UniqueID, 1),
} }
segIdx1 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx1.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+1, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+3, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+4, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+5, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
s.meta.indexMeta.segmentIndexes.Insert(segID, segIdx1)
segIdx2 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx2.Insert(indexID, &model.SegmentIndex{
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
})
segIdx2.Insert(indexID+1, &model.SegmentIndex{
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
})
segIdx2.Insert(indexID+3, &model.SegmentIndex{
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
CreatedUTCTime: createTS,
})
segIdx2.Insert(indexID+4, &model.SegmentIndex{
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
CreatedUTCTime: createTS,
})
segIdx2.Insert(indexID+5, &model.SegmentIndex{
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreatedUTCTime: createTS,
})
s.meta.indexMeta.segmentIndexes.Insert(segID-1, segIdx2)
for id, segment := range segments { for id, segment := range segments {
s.meta.segments.SetSegment(id, segment) s.meta.segments.SetSegment(id, segment)
} }
@ -1720,7 +1720,7 @@ func TestServer_ListIndexes(t *testing.T) {
}, },
}, },
}, },
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{}, segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
}, },
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
@ -1909,95 +1909,7 @@ func TestServer_GetIndexStatistics(t *testing.T) {
}, },
}, },
}, },
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
segID: {
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 1: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 3: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 4: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
indexID + 5: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
},
},
}, },
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
@ -2005,6 +1917,93 @@ func TestServer_GetIndexStatistics(t *testing.T) {
allocator: mock0Allocator, allocator: mock0Allocator,
notifyIndexChan: make(chan UniqueID, 1), notifyIndexChan: make(chan UniqueID, 1),
} }
segIdx1 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx1.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+1, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+3, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+4, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
segIdx1.Insert(indexID+5, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
s.meta.indexMeta.segmentIndexes.Insert(segID, segIdx1)
for id, segment := range segments { for id, segment := range segments {
s.meta.segments.SetSegment(id, segment) s.meta.segments.SetSegment(id, segment)
} }
@ -2156,7 +2155,7 @@ func TestServer_DropIndex(t *testing.T) {
}, },
}, },
}, },
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{}, segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
}, },
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
@ -2301,27 +2300,7 @@ func TestServer_GetIndexInfos(t *testing.T) {
}, },
}, },
}, },
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
segID: {
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
},
},
},
}, },
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
@ -2330,6 +2309,25 @@ func TestServer_GetIndexInfos(t *testing.T) {
allocator: mock0Allocator, allocator: mock0Allocator,
notifyIndexChan: make(chan UniqueID, 1), notifyIndexChan: make(chan UniqueID, 1),
} }
segIdx1 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx1.Insert(indexID, &model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: createTS,
IndexFileKeys: nil,
IndexSerializedSize: 0,
WriteHandoff: false,
})
s.meta.indexMeta.segmentIndexes.Insert(segID, segIdx1)
s.meta.segments.SetSegment(segID, &SegmentInfo{ s.meta.segments.SetSegment(segID, &SegmentInfo{
SegmentInfo: &datapb.SegmentInfo{ SegmentInfo: &datapb.SegmentInfo{
ID: segID, ID: segID,
@ -2395,7 +2393,7 @@ func TestMeta_GetHasUnindexTaskSegments(t *testing.T) {
segments: NewSegmentsInfo(), segments: NewSegmentsInfo(),
indexMeta: &indexMeta{ indexMeta: &indexMeta{
segmentBuildInfo: newSegmentIndexBuildInfo(), segmentBuildInfo: newSegmentIndexBuildInfo(),
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{}, segmentIndexes: typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]](),
indexes: map[UniqueID]map[UniqueID]*model.Index{ indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: { collID: {
indexID: { indexID: {
@ -2438,7 +2436,7 @@ func TestMeta_GetHasUnindexTaskSegments(t *testing.T) {
assert.Equal(t, 1, len(segments)) assert.Equal(t, 1, len(segments))
assert.Equal(t, segID, segments[0].ID) assert.Equal(t, segID, segments[0].ID)
m.indexMeta.segmentIndexes[segID] = make(map[UniqueID]*model.SegmentIndex) m.indexMeta.segmentIndexes.Insert(segID, typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]())
m.indexMeta.updateSegmentIndex(&model.SegmentIndex{ m.indexMeta.updateSegmentIndex(&model.SegmentIndex{
CollectionID: collID, CollectionID: collID,
SegmentID: segID, SegmentID: segID,

12
internal/datacoord/job_manager_test.go
View File

@ -16,6 +16,8 @@ import (
"github.com/milvus-io/milvus/internal/metastore/mocks" "github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/pkg/v2/proto/datapb" "github.com/milvus-io/milvus/pkg/v2/proto/datapb"
"github.com/milvus-io/milvus/pkg/v2/proto/indexpb" "github.com/milvus-io/milvus/pkg/v2/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v2/util/lock"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
) )
type jobManagerSuite struct { type jobManagerSuite struct {
@ -89,9 +91,11 @@ func (s *jobManagerSuite) TestJobManager_triggerStatsTaskLoop() {
}, },
}, },
statsTaskMeta: &statsTaskMeta{ statsTaskMeta: &statsTaskMeta{
ctx: ctx, ctx: ctx,
catalog: catalog, catalog: catalog,
tasks: make(map[int64]*indexpb.StatsTask), keyLock: lock.NewKeyLock[UniqueID](),
tasks: typeutil.NewConcurrentMap[UniqueID, *indexpb.StatsTask](),
segmentID2Tasks: typeutil.NewConcurrentMap[string, *indexpb.StatsTask](),
}, },
} }
@ -103,7 +107,7 @@ func (s *jobManagerSuite) TestJobManager_triggerStatsTaskLoop() {
scheduler: &taskScheduler{ scheduler: &taskScheduler{
allocator: alloc, allocator: alloc,
pendingTasks: newFairQueuePolicy(), pendingTasks: newFairQueuePolicy(),
runningTasks: make(map[UniqueID]Task), runningTasks: typeutil.NewConcurrentMap[UniqueID, Task](),
meta: mt, meta: mt,
taskStats: expirable.NewLRU[UniqueID, Task](512, nil, time.Minute*5), taskStats: expirable.NewLRU[UniqueID, Task](512, nil, time.Minute*5),
}, },

40
internal/datacoord/metrics_info_test.go
View File

@ -649,6 +649,25 @@ func TestGetDistJSON(t *testing.T) {
} }
func TestServer_getSegmentsJSON(t *testing.T) { func TestServer_getSegmentsJSON(t *testing.T) {
segIndexes := typeutil.NewConcurrentMap[UniqueID, *typeutil.ConcurrentMap[UniqueID, *model.SegmentIndex]]()
segIdx0 := typeutil.NewConcurrentMap[UniqueID, *model.SegmentIndex]()
segIdx0.Insert(10, &model.SegmentIndex{
SegmentID: 1000,
CollectionID: 1,
PartitionID: 2,
NumRows: 10250,
IndexID: 10,
BuildID: 10000,
NodeID: 1,
IndexVersion: 0,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 12,
IndexFileKeys: nil,
IndexSerializedSize: 0,
})
segIndexes.Insert(1000, segIdx0)
s := &Server{ s := &Server{
meta: &meta{ meta: &meta{
segments: &SegmentsInfo{ segments: &SegmentsInfo{
@ -664,26 +683,7 @@ func TestServer_getSegmentsJSON(t *testing.T) {
}, },
}, },
indexMeta: &indexMeta{ indexMeta: &indexMeta{
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{ segmentIndexes: segIndexes,
1000: {
10: &model.SegmentIndex{
SegmentID: 1000,
CollectionID: 1,
PartitionID: 2,
NumRows: 10250,
IndexID: 10,
BuildID: 10000,
NodeID: 1,
IndexVersion: 0,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreatedUTCTime: 12,
IndexFileKeys: nil,
IndexSerializedSize: 0,
},
},
},
indexes: map[UniqueID]map[UniqueID]*model.Index{ indexes: map[UniqueID]map[UniqueID]*model.Index{
1: { 1: {
10: &model.Index{ 10: &model.Index{

209
internal/datacoord/stats_task_meta.go
View File

@ -19,36 +19,40 @@ package datacoord
import ( import (
"context" "context"
"fmt" "fmt"
"strconv"
"sync"
"go.uber.org/zap" "go.uber.org/zap"
"google.golang.org/protobuf/proto" "google.golang.org/protobuf/proto"
"strconv"
"github.com/milvus-io/milvus/internal/metastore" "github.com/milvus-io/milvus/internal/metastore"
"github.com/milvus-io/milvus/pkg/v2/log" "github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/metrics" "github.com/milvus-io/milvus/pkg/v2/metrics"
"github.com/milvus-io/milvus/pkg/v2/proto/indexpb" "github.com/milvus-io/milvus/pkg/v2/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v2/proto/workerpb" "github.com/milvus-io/milvus/pkg/v2/proto/workerpb"
"github.com/milvus-io/milvus/pkg/v2/util/lock"
"github.com/milvus-io/milvus/pkg/v2/util/merr" "github.com/milvus-io/milvus/pkg/v2/util/merr"
"github.com/milvus-io/milvus/pkg/v2/util/timerecord" "github.com/milvus-io/milvus/pkg/v2/util/timerecord"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
) )
type statsTaskMeta struct { type statsTaskMeta struct {
sync.RWMutex
ctx context.Context ctx context.Context
catalog metastore.DataCoordCatalog catalog metastore.DataCoordCatalog
// taskID -> analyzeStats keyLock *lock.KeyLock[UniqueID]
tasks map[int64]*indexpb.StatsTask // taskID -> statsTask
tasks *typeutil.ConcurrentMap[UniqueID, *indexpb.StatsTask]
// segmentID + SubJobType -> statsTask
segmentID2Tasks *typeutil.ConcurrentMap[string, *indexpb.StatsTask]
} }
func newStatsTaskMeta(ctx context.Context, catalog metastore.DataCoordCatalog) (*statsTaskMeta, error) { func newStatsTaskMeta(ctx context.Context, catalog metastore.DataCoordCatalog) (*statsTaskMeta, error) {
stm := &statsTaskMeta{ stm := &statsTaskMeta{
ctx: ctx, ctx: ctx,
catalog: catalog, catalog: catalog,
tasks: make(map[int64]*indexpb.StatsTask), keyLock: lock.NewKeyLock[UniqueID](),
tasks: typeutil.NewConcurrentMap[UniqueID, *indexpb.StatsTask](),
segmentID2Tasks: typeutil.NewConcurrentMap[string, *indexpb.StatsTask](),
} }
if err := stm.reloadFromKV(); err != nil { if err := stm.reloadFromKV(); err != nil {
return nil, err return nil, err
@ -56,6 +60,10 @@ func newStatsTaskMeta(ctx context.Context, catalog metastore.DataCoordCatalog) (
return stm, nil return stm, nil
} }
func createSecondaryIndexKey(segmentID UniqueID, subJobType string) string {
return strconv.FormatUint(uint64(segmentID), 10) + "-" + subJobType
}
func (stm *statsTaskMeta) reloadFromKV() error { func (stm *statsTaskMeta) reloadFromKV() error {
record := timerecord.NewTimeRecorder("statsTaskMeta-reloadFromKV") record := timerecord.NewTimeRecorder("statsTaskMeta-reloadFromKV")
// load stats task // load stats task
@ -65,7 +73,10 @@ func (stm *statsTaskMeta) reloadFromKV() error {
return err return err
} }
for _, t := range statsTasks { for _, t := range statsTasks {
stm.tasks[t.GetTaskID()] = t stm.tasks.Insert(t.GetTaskID(), t)
secondaryKey := createSecondaryIndexKey(t.GetSegmentID(), t.GetSubJobType().String())
stm.segmentID2Tasks.Insert(secondaryKey, t)
} }
log.Info("statsTaskMeta reloadFromKV done", zap.Duration("duration", record.ElapseSpan())) log.Info("statsTaskMeta reloadFromKV done", zap.Duration("duration", record.ElapseSpan()))
@ -73,58 +84,54 @@ func (stm *statsTaskMeta) reloadFromKV() error {
} }
func (stm *statsTaskMeta) updateMetrics() { func (stm *statsTaskMeta) updateMetrics() {
stm.RLock() taskMetrics := make(map[indexpb.JobState]int)
defer stm.RUnlock() taskMetrics[indexpb.JobState_JobStateNone] = 0
taskMetrics[indexpb.JobState_JobStateInit] = 0
taskMetrics := make(map[UniqueID]map[indexpb.JobState]int) taskMetrics[indexpb.JobState_JobStateInProgress] = 0
for _, t := range stm.tasks { taskMetrics[indexpb.JobState_JobStateFinished] = 0
if _, ok := taskMetrics[t.GetCollectionID()]; !ok { taskMetrics[indexpb.JobState_JobStateFailed] = 0
taskMetrics[t.GetCollectionID()] = make(map[indexpb.JobState]int) taskMetrics[indexpb.JobState_JobStateRetry] = 0
taskMetrics[t.GetCollectionID()][indexpb.JobState_JobStateNone] = 0 allTasks := stm.tasks.Values()
taskMetrics[t.GetCollectionID()][indexpb.JobState_JobStateInit] = 0 for _, t := range allTasks {
taskMetrics[t.GetCollectionID()][indexpb.JobState_JobStateInProgress] = 0 taskMetrics[t.GetState()]++
taskMetrics[t.GetCollectionID()][indexpb.JobState_JobStateFinished] = 0
taskMetrics[t.GetCollectionID()][indexpb.JobState_JobStateFailed] = 0
taskMetrics[t.GetCollectionID()][indexpb.JobState_JobStateRetry] = 0
}
taskMetrics[t.GetCollectionID()][t.GetState()]++
} }
jobType := indexpb.JobType_JobTypeStatsJob.String() jobType := indexpb.JobType_JobTypeStatsJob.String()
for collID, m := range taskMetrics { for k, v := range taskMetrics {
for k, v := range m { metrics.TaskNum.WithLabelValues(jobType, k.String()).Set(float64(v))
metrics.TaskNum.WithLabelValues(strconv.FormatInt(collID, 10), jobType, k.String()).Set(float64(v))
}
} }
} }
func (stm *statsTaskMeta) AddStatsTask(t *indexpb.StatsTask) error { func (stm *statsTaskMeta) AddStatsTask(t *indexpb.StatsTask) error {
stm.Lock() taskID := t.GetTaskID()
defer stm.Unlock()
for _, st := range stm.tasks { secondaryKey := createSecondaryIndexKey(t.GetSegmentID(), t.GetSubJobType().String())
if st.GetTaskID() == t.GetTaskID() || (st.GetSegmentID() == t.GetSegmentID() && st.GetSubJobType() == t.GetSubJobType() && st.GetState() != indexpb.JobState_JobStateFailed) { task, alreadyExist := stm.segmentID2Tasks.Get(secondaryKey)
msg := fmt.Sprintf("stats task already exist in meta of segment %d with subJobType: %s", if alreadyExist {
t.GetSegmentID(), t.GetSubJobType().String()) msg := fmt.Sprintf("stats task already exist in meta of segment %d with subJobType: %s",
log.RatedWarn(10, msg, zap.Int64("taskID", t.GetTaskID()), zap.Int64("exist taskID", st.GetTaskID())) t.GetSegmentID(), t.GetSubJobType().String())
return merr.WrapErrTaskDuplicate(indexpb.JobType_JobTypeStatsJob.String(), msg) log.RatedWarn(10, msg, zap.Int64("taskID", t.GetTaskID()), zap.Int64("exist taskID", task.GetTaskID()))
} return merr.WrapErrTaskDuplicate(indexpb.JobType_JobTypeStatsJob.String(), msg)
} }
stm.keyLock.Lock(taskID)
defer stm.keyLock.Unlock(taskID)
log.Info("add stats task", zap.Int64("taskID", t.GetTaskID()), zap.Int64("originSegmentID", t.GetSegmentID()), log.Info("add stats task", zap.Int64("taskID", t.GetTaskID()), zap.Int64("originSegmentID", t.GetSegmentID()),
zap.Int64("targetSegmentID", t.GetTargetSegmentID()), zap.String("subJobType", t.GetSubJobType().String())) zap.Int64("targetSegmentID", t.GetTargetSegmentID()), zap.String("subJobType", t.GetSubJobType().String()))
t.State = indexpb.JobState_JobStateInit t.State = indexpb.JobState_JobStateInit
if err := stm.catalog.SaveStatsTask(stm.ctx, t); err != nil { if err := stm.catalog.SaveStatsTask(stm.ctx, t); err != nil {
log.Warn("adding stats task failed", log.Warn("adding stats task failed",
zap.Int64("taskID", t.GetTaskID()), zap.Int64("taskID", taskID),
zap.Int64("segmentID", t.GetSegmentID()), zap.Int64("segmentID", t.GetSegmentID()),
zap.String("subJobType", t.GetSubJobType().String()), zap.String("subJobType", t.GetSubJobType().String()),
zap.Error(err)) zap.Error(err))
return err return err
} }
stm.tasks[t.GetTaskID()] = t stm.tasks.Insert(taskID, t)
stm.segmentID2Tasks.Insert(secondaryKey, t)
log.Info("add stats task success", zap.Int64("taskID", t.GetTaskID()), zap.Int64("originSegmentID", t.GetSegmentID()), log.Info("add stats task success", zap.Int64("taskID", t.GetTaskID()), zap.Int64("originSegmentID", t.GetSegmentID()),
zap.Int64("targetSegmentID", t.GetTargetSegmentID()), zap.String("subJobType", t.GetSubJobType().String())) zap.Int64("targetSegmentID", t.GetTargetSegmentID()), zap.String("subJobType", t.GetSubJobType().String()))
@ -132,12 +139,12 @@ func (stm *statsTaskMeta) AddStatsTask(t *indexpb.StatsTask) error {
} }
func (stm *statsTaskMeta) DropStatsTask(taskID int64) error { func (stm *statsTaskMeta) DropStatsTask(taskID int64) error {
stm.Lock() stm.keyLock.Lock(taskID)
defer stm.Unlock() defer stm.keyLock.Unlock(taskID)
log.Info("drop stats task by taskID", zap.Int64("taskID", taskID)) log.Info("drop stats task by taskID", zap.Int64("taskID", taskID))
t, ok := stm.tasks[taskID] t, ok := stm.tasks.Get(taskID)
if !ok { if !ok {
log.Info("remove stats task success, task already not exist", zap.Int64("taskID", taskID)) log.Info("remove stats task success, task already not exist", zap.Int64("taskID", taskID))
return nil return nil
@ -150,17 +157,19 @@ func (stm *statsTaskMeta) DropStatsTask(taskID int64) error {
return err return err
} }
delete(stm.tasks, taskID) stm.tasks.Remove(taskID)
secondaryKey := createSecondaryIndexKey(t.GetSegmentID(), t.GetSubJobType().String())
stm.segmentID2Tasks.Remove(secondaryKey)
log.Info("remove stats task success", zap.Int64("taskID", taskID), zap.Int64("segmentID", t.SegmentID)) log.Info("remove stats task success", zap.Int64("taskID", taskID))
return nil return nil
} }
func (stm *statsTaskMeta) UpdateVersion(taskID, nodeID int64) error { func (stm *statsTaskMeta) UpdateVersion(taskID, nodeID int64) error {
stm.Lock() stm.keyLock.Lock(taskID)
defer stm.Unlock() defer stm.keyLock.Unlock(taskID)
t, ok := stm.tasks[taskID] t, ok := stm.tasks.Get(taskID)
if !ok { if !ok {
return fmt.Errorf("task %d not found", taskID) return fmt.Errorf("task %d not found", taskID)
} }
@ -178,17 +187,19 @@ func (stm *statsTaskMeta) UpdateVersion(taskID, nodeID int64) error {
return err return err
} }
stm.tasks[t.TaskID] = cloneT stm.tasks.Insert(taskID, cloneT)
secondaryKey := createSecondaryIndexKey(t.GetSegmentID(), t.GetSubJobType().String())
stm.segmentID2Tasks.Insert(secondaryKey, cloneT)
log.Info("update stats task version success", zap.Int64("taskID", taskID), zap.Int64("nodeID", nodeID), log.Info("update stats task version success", zap.Int64("taskID", taskID), zap.Int64("nodeID", nodeID),
zap.Int64("newVersion", cloneT.GetVersion())) zap.Int64("newVersion", cloneT.GetVersion()))
return nil return nil
} }
func (stm *statsTaskMeta) UpdateBuildingTask(taskID int64) error { func (stm *statsTaskMeta) UpdateBuildingTask(taskID int64) error {
stm.Lock() stm.keyLock.Lock(taskID)
defer stm.Unlock() defer stm.keyLock.Unlock(taskID)
t, ok := stm.tasks[taskID] t, ok := stm.tasks.Get(taskID)
if !ok { if !ok {
return fmt.Errorf("task %d not found", taskID) return fmt.Errorf("task %d not found", taskID)
} }
@ -204,17 +215,19 @@ func (stm *statsTaskMeta) UpdateBuildingTask(taskID int64) error {
return err return err
} }
stm.tasks[t.TaskID] = cloneT stm.tasks.Insert(taskID, cloneT)
secondaryKey := createSecondaryIndexKey(t.GetSegmentID(), t.GetSubJobType().String())
stm.segmentID2Tasks.Insert(secondaryKey, cloneT)
log.Info("update building stats task success", zap.Int64("taskID", taskID)) log.Info("update building stats task success", zap.Int64("taskID", taskID))
return nil return nil
} }
func (stm *statsTaskMeta) FinishTask(taskID int64, result *workerpb.StatsResult) error { func (stm *statsTaskMeta) FinishTask(taskID int64, result *workerpb.StatsResult) error {
stm.Lock() stm.keyLock.Lock(taskID)
defer stm.Unlock() defer stm.keyLock.Unlock(taskID)
t, ok := stm.tasks[taskID] t, ok := stm.tasks.Get(taskID)
if !ok { if !ok {
return fmt.Errorf("task %d not found", taskID) return fmt.Errorf("task %d not found", taskID)
} }
@ -231,77 +244,70 @@ func (stm *statsTaskMeta) FinishTask(taskID int64, result *workerpb.StatsResult)
return err return err
} }
stm.tasks[t.TaskID] = cloneT stm.tasks.Insert(taskID, cloneT)
secondaryKey := createSecondaryIndexKey(t.GetSegmentID(), t.GetSubJobType().String())
stm.segmentID2Tasks.Insert(secondaryKey, cloneT)
log.Info("finish stats task meta success", zap.Int64("taskID", taskID), zap.Int64("segmentID", t.SegmentID), log.Info("finish stats task meta success", zap.Int64("taskID", taskID), zap.Int64("segmentID", t.SegmentID),
zap.String("state", result.GetState().String()), zap.String("failReason", t.GetFailReason())) zap.String("state", result.GetState().String()), zap.String("failReason", t.GetFailReason()))
return nil return nil
} }
func (stm *statsTaskMeta) GetStatsTaskState(taskID int64) indexpb.JobState { func (stm *statsTaskMeta) GetStatsTask(taskID int64) *indexpb.StatsTask {
stm.RLock() t, _ := stm.tasks.Get(taskID)
defer stm.RUnlock() return t
}
t, ok := stm.tasks[taskID] func (stm *statsTaskMeta) GetStatsTaskState(taskID int64) indexpb.JobState {
t, ok := stm.tasks.Get(taskID)
if !ok { if !ok {
return indexpb.JobState_JobStateNone return indexpb.JobState_JobStateNone
} }
return t.GetState() return t.GetState()
} }
func (stm *statsTaskMeta) GetStatsTaskStateBySegmentID(segmentID int64, jobType indexpb.StatsSubJob) indexpb.JobState { func (stm *statsTaskMeta) GetStatsTaskStateBySegmentID(segmentID int64, subJobType indexpb.StatsSubJob) indexpb.JobState {
stm.RLock() state := indexpb.JobState_JobStateNone
defer stm.RUnlock()
for _, t := range stm.tasks { secondaryKey := createSecondaryIndexKey(segmentID, subJobType.String())
if segmentID == t.GetSegmentID() && jobType == t.GetSubJobType() { t, exists := stm.segmentID2Tasks.Get(secondaryKey)
return t.GetState() if exists {
} state = t.GetState()
} }
return state
return indexpb.JobState_JobStateNone
} }
func (stm *statsTaskMeta) CanCleanedTasks() []int64 { func (stm *statsTaskMeta) CanCleanedTasks() []int64 {
stm.RLock()
defer stm.RUnlock()
needCleanedTaskIDs := make([]int64, 0) needCleanedTaskIDs := make([]int64, 0)
for taskID, t := range stm.tasks { stm.tasks.Range(func(key UniqueID, value *indexpb.StatsTask) bool {
if t.GetCanRecycle() && (t.GetState() == indexpb.JobState_JobStateFinished || if value.GetCanRecycle() && (value.GetState() == indexpb.JobState_JobStateFinished ||
t.GetState() == indexpb.JobState_JobStateFailed) { value.GetState() == indexpb.JobState_JobStateFailed) {
needCleanedTaskIDs = append(needCleanedTaskIDs, taskID) needCleanedTaskIDs = append(needCleanedTaskIDs, key)
} }
} return true
})
return needCleanedTaskIDs return needCleanedTaskIDs
} }
func (stm *statsTaskMeta) GetAllTasks() map[int64]*indexpb.StatsTask { func (stm *statsTaskMeta) GetAllTasks() map[int64]*indexpb.StatsTask {
tasks := make(map[int64]*indexpb.StatsTask) tasks := make(map[int64]*indexpb.StatsTask)
stm.RLock() allTasks := stm.tasks.Values()
defer stm.RUnlock() for _, v := range allTasks {
for k, v := range stm.tasks { tasks[v.GetTaskID()] = proto.Clone(v).(*indexpb.StatsTask)
tasks[k] = proto.Clone(v).(*indexpb.StatsTask)
} }
return tasks return tasks
} }
func (stm *statsTaskMeta) GetStatsTaskBySegmentID(segmentID int64, subJobType indexpb.StatsSubJob) *indexpb.StatsTask { func (stm *statsTaskMeta) GetStatsTaskBySegmentID(segmentID int64, subJobType indexpb.StatsSubJob) *indexpb.StatsTask {
stm.RLock() secondaryKey := createSecondaryIndexKey(segmentID, subJobType.String())
defer stm.RUnlock() t, exists := stm.segmentID2Tasks.Get(secondaryKey)
if exists {
log.Info("get stats task by segmentID", zap.Int64("segmentID", segmentID), log.Info("get stats task by segmentID success",
zap.String("subJobType", subJobType.String())) zap.Int64("taskID", t.GetTaskID()),
zap.Int64("segmentID", segmentID),
for taskID, t := range stm.tasks { zap.String("subJobType", subJobType.String()))
if t.GetSegmentID() == segmentID && t.GetSubJobType() == subJobType { return t
log.Info("get stats task by segmentID success",
zap.Int64("taskID", taskID),
zap.Int64("segmentID", segmentID),
zap.String("subJobType", subJobType.String()))
return t
}
} }
log.Info("get stats task by segmentID failed, task not exist", zap.Int64("segmentID", segmentID), log.Info("get stats task by segmentID failed, task not exist", zap.Int64("segmentID", segmentID),
@ -310,12 +316,9 @@ func (stm *statsTaskMeta) GetStatsTaskBySegmentID(segmentID int64, subJobType in
} }
func (stm *statsTaskMeta) MarkTaskCanRecycle(taskID int64) error { func (stm *statsTaskMeta) MarkTaskCanRecycle(taskID int64) error {
stm.Lock()
defer stm.Unlock()
log.Info("mark stats task can recycle", zap.Int64("taskID", taskID)) log.Info("mark stats task can recycle", zap.Int64("taskID", taskID))
t, ok := stm.tasks[taskID] t, ok := stm.tasks.Get(taskID)
if !ok { if !ok {
return fmt.Errorf("task %d not found", taskID) return fmt.Errorf("task %d not found", taskID)
} }
@ -331,7 +334,9 @@ func (stm *statsTaskMeta) MarkTaskCanRecycle(taskID int64) error {
return err return err
} }
stm.tasks[t.TaskID] = cloneT stm.tasks.Insert(taskID, cloneT)
secondaryKey := createSecondaryIndexKey(t.GetSegmentID(), t.GetSubJobType().String())
stm.segmentID2Tasks.Insert(secondaryKey, cloneT)
log.Info("mark stats task can recycle success", zap.Int64("taskID", taskID), log.Info("mark stats task can recycle success", zap.Int64("taskID", taskID),
zap.Int64("segmentID", t.SegmentID), zap.Int64("segmentID", t.SegmentID),

26
internal/datacoord/stats_task_meta_test.go
View File

@ -108,7 +108,7 @@ func (s *statsTaskMetaSuite) Test_Method() {
catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once() catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once()
s.Error(m.AddStatsTask(t)) s.Error(m.AddStatsTask(t))
_, ok := m.tasks[1] _, ok := m.tasks.Get(1)
s.False(ok) s.False(ok)
}) })
@ -116,13 +116,13 @@ func (s *statsTaskMetaSuite) Test_Method() {
catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(nil).Once() catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(nil).Once()
s.NoError(m.AddStatsTask(t)) s.NoError(m.AddStatsTask(t))
_, ok := m.tasks[1] _, ok := m.tasks.Get(1)
s.True(ok) s.True(ok)
}) })
s.Run("already exist", func() { s.Run("already exist", func() {
s.Error(m.AddStatsTask(t)) s.Error(m.AddStatsTask(t))
_, ok := m.tasks[1] _, ok := m.tasks.Get(1)
s.True(ok) s.True(ok)
}) })
}) })
@ -132,13 +132,13 @@ func (s *statsTaskMetaSuite) Test_Method() {
catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(nil).Once() catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(nil).Once()
s.NoError(m.UpdateVersion(1, 1180)) s.NoError(m.UpdateVersion(1, 1180))
task, ok := m.tasks[1] task, ok := m.tasks.Get(1)
s.True(ok) s.True(ok)
s.Equal(int64(1), task.GetVersion()) s.Equal(int64(1), task.GetVersion())
}) })
s.Run("task not exist", func() { s.Run("task not exist", func() {
_, ok := m.tasks[100] _, ok := m.tasks.Get(100)
s.False(ok) s.False(ok)
s.Error(m.UpdateVersion(100, 1180)) s.Error(m.UpdateVersion(100, 1180))
@ -148,7 +148,7 @@ func (s *statsTaskMetaSuite) Test_Method() {
catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once() catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once()
s.Error(m.UpdateVersion(1, 1180)) s.Error(m.UpdateVersion(1, 1180))
task, ok := m.tasks[1] task, ok := m.tasks.Get(1)
s.True(ok) s.True(ok)
// still 1 // still 1
s.Equal(int64(1), task.GetVersion()) s.Equal(int64(1), task.GetVersion())
@ -160,7 +160,7 @@ func (s *statsTaskMetaSuite) Test_Method() {
catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once() catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once()
s.Error(m.UpdateBuildingTask(1)) s.Error(m.UpdateBuildingTask(1))
task, ok := m.tasks[1] task, ok := m.tasks.Get(1)
s.True(ok) s.True(ok)
s.Equal(indexpb.JobState_JobStateInit, task.GetState()) s.Equal(indexpb.JobState_JobStateInit, task.GetState())
s.Equal(int64(1180), task.GetNodeID()) s.Equal(int64(1180), task.GetNodeID())
@ -170,14 +170,14 @@ func (s *statsTaskMetaSuite) Test_Method() {
catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(nil).Once() catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(nil).Once()
s.NoError(m.UpdateBuildingTask(1)) s.NoError(m.UpdateBuildingTask(1))
task, ok := m.tasks[1] task, ok := m.tasks.Get(1)
s.True(ok) s.True(ok)
s.Equal(indexpb.JobState_JobStateInProgress, task.GetState()) s.Equal(indexpb.JobState_JobStateInProgress, task.GetState())
s.Equal(int64(1180), task.GetNodeID()) s.Equal(int64(1180), task.GetNodeID())
}) })
s.Run("task not exist", func() { s.Run("task not exist", func() {
_, ok := m.tasks[100] _, ok := m.tasks.Get(100)
s.False(ok) s.False(ok)
s.Error(m.UpdateBuildingTask(100)) s.Error(m.UpdateBuildingTask(100))
@ -217,7 +217,7 @@ func (s *statsTaskMetaSuite) Test_Method() {
catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once() catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once()
s.Error(m.FinishTask(1, result)) s.Error(m.FinishTask(1, result))
task, ok := m.tasks[1] task, ok := m.tasks.Get(1)
s.True(ok) s.True(ok)
s.Equal(indexpb.JobState_JobStateInProgress, task.GetState()) s.Equal(indexpb.JobState_JobStateInProgress, task.GetState())
}) })
@ -226,7 +226,7 @@ func (s *statsTaskMetaSuite) Test_Method() {
catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(nil).Once() catalog.EXPECT().SaveStatsTask(mock.Anything, mock.Anything).Return(nil).Once()
s.NoError(m.FinishTask(1, result)) s.NoError(m.FinishTask(1, result))
task, ok := m.tasks[1] task, ok := m.tasks.Get(1)
s.True(ok) s.True(ok)
s.Equal(indexpb.JobState_JobStateFinished, task.GetState()) s.Equal(indexpb.JobState_JobStateFinished, task.GetState())
}) })
@ -268,7 +268,7 @@ func (s *statsTaskMetaSuite) Test_Method() {
catalog.EXPECT().DropStatsTask(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once() catalog.EXPECT().DropStatsTask(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once()
s.Error(m.DropStatsTask(1)) s.Error(m.DropStatsTask(1))
_, ok := m.tasks[1] _, ok := m.tasks.Get(1)
s.True(ok) s.True(ok)
}) })
@ -276,7 +276,7 @@ func (s *statsTaskMetaSuite) Test_Method() {
catalog.EXPECT().DropStatsTask(mock.Anything, mock.Anything).Return(nil).Once() catalog.EXPECT().DropStatsTask(mock.Anything, mock.Anything).Return(nil).Once()
s.NoError(m.DropStatsTask(1)) s.NoError(m.DropStatsTask(1))
_, ok := m.tasks[1] _, ok := m.tasks.Get(1)
s.False(ok) s.False(ok)
s.NoError(m.DropStatsTask(1000)) s.NoError(m.DropStatsTask(1000))

13
internal/datacoord/task_analyze.go
View File

@ -277,14 +277,17 @@ func (at *analyzeTask) QueryResult(ctx context.Context, client types.DataNodeCli
// infos length is always one. // infos length is always one.
for _, result := range resp.GetAnalyzeJobResults().GetResults() { for _, result := range resp.GetAnalyzeJobResults().GetResults() {
if result.GetTaskID() == at.GetTaskID() { if result.GetTaskID() == at.GetTaskID() {
log.Ctx(ctx).Info("query analysis task info successfully",
zap.Int64("taskID", at.GetTaskID()), zap.String("result state", result.GetState().String()),
zap.String("failReason", result.GetFailReason()))
if result.GetState() == indexpb.JobState_JobStateFinished || result.GetState() == indexpb.JobState_JobStateFailed || if result.GetState() == indexpb.JobState_JobStateFinished || result.GetState() == indexpb.JobState_JobStateFailed ||
result.GetState() == indexpb.JobState_JobStateRetry { result.GetState() == indexpb.JobState_JobStateRetry {
log.Ctx(ctx).Info("query analysis task info successfully",
zap.Int64("taskID", at.GetTaskID()), zap.String("result state", result.GetState().String()),
zap.String("failReason", result.GetFailReason()))
// state is retry or finished or failed // state is retry or finished or failed
at.setResult(result) at.setResult(result)
} else if result.GetState() == indexpb.JobState_JobStateNone { } else if result.GetState() == indexpb.JobState_JobStateNone {
log.Ctx(ctx).Info("query analysis task info successfully",
zap.Int64("taskID", at.GetTaskID()), zap.String("result state", result.GetState().String()),
zap.String("failReason", result.GetFailReason()))
at.SetState(indexpb.JobState_JobStateRetry, "analyze task state is none in info response") at.SetState(indexpb.JobState_JobStateRetry, "analyze task state is none in info response")
} }
// inProgress or unissued/init, keep InProgress state // inProgress or unissued/init, keep InProgress state
@ -320,3 +323,7 @@ func (at *analyzeTask) DropTaskOnWorker(ctx context.Context, client types.DataNo
func (at *analyzeTask) SetJobInfo(meta *meta) error { func (at *analyzeTask) SetJobInfo(meta *meta) error {
return meta.analyzeMeta.FinishTask(at.GetTaskID(), at.taskInfo) return meta.analyzeMeta.FinishTask(at.GetTaskID(), at.taskInfo)
} }
func (at *analyzeTask) DropTaskMeta(ctx context.Context, meta *meta) error {
return meta.analyzeMeta.DropAnalyzeTask(ctx, at.GetTaskID())
}

13
internal/datacoord/task_index.go
View File

@ -317,14 +317,17 @@ func (it *indexBuildTask) QueryResult(ctx context.Context, node types.DataNodeCl
// indexInfos length is always one. // indexInfos length is always one.
for _, info := range resp.GetIndexJobResults().GetResults() { for _, info := range resp.GetIndexJobResults().GetResults() {
if info.GetBuildID() == it.GetTaskID() { if info.GetBuildID() == it.GetTaskID() {
log.Ctx(ctx).Info("query task index info successfully",
zap.Int64("taskID", it.GetTaskID()), zap.String("result state", info.GetState().String()),
zap.String("failReason", info.GetFailReason()))
if info.GetState() == commonpb.IndexState_Finished || info.GetState() == commonpb.IndexState_Failed || if info.GetState() == commonpb.IndexState_Finished || info.GetState() == commonpb.IndexState_Failed ||
info.GetState() == commonpb.IndexState_Retry { info.GetState() == commonpb.IndexState_Retry {
log.Ctx(ctx).Info("query task index info successfully",
zap.Int64("taskID", it.GetTaskID()), zap.String("result state", info.GetState().String()),
zap.String("failReason", info.GetFailReason()))
// state is retry or finished or failed // state is retry or finished or failed
it.setResult(info) it.setResult(info)
} else if info.GetState() == commonpb.IndexState_IndexStateNone { } else if info.GetState() == commonpb.IndexState_IndexStateNone {
log.Ctx(ctx).Info("query task index info successfully",
zap.Int64("taskID", it.GetTaskID()), zap.String("result state", info.GetState().String()),
zap.String("failReason", info.GetFailReason()))
it.SetState(indexpb.JobState_JobStateRetry, "index state is none in info response") it.SetState(indexpb.JobState_JobStateRetry, "index state is none in info response")
} }
// inProgress or unissued, keep InProgress state // inProgress or unissued, keep InProgress state
@ -358,3 +361,7 @@ func (it *indexBuildTask) DropTaskOnWorker(ctx context.Context, client types.Dat
func (it *indexBuildTask) SetJobInfo(meta *meta) error { func (it *indexBuildTask) SetJobInfo(meta *meta) error {
return meta.indexMeta.FinishTask(it.taskInfo) return meta.indexMeta.FinishTask(it.taskInfo)
} }
func (it *indexBuildTask) DropTaskMeta(ctx context.Context, meta *meta) error {
return meta.indexMeta.RemoveSegmentIndexByID(ctx, it.taskID)
}

122
internal/datacoord/task_scheduler.go
View File

@ -33,6 +33,7 @@ import (
"github.com/milvus-io/milvus/pkg/v2/proto/indexpb" "github.com/milvus-io/milvus/pkg/v2/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v2/proto/workerpb" "github.com/milvus-io/milvus/pkg/v2/proto/workerpb"
"github.com/milvus-io/milvus/pkg/v2/util/lock" "github.com/milvus-io/milvus/pkg/v2/util/lock"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
) )
type taskScheduler struct { type taskScheduler struct {
@ -43,9 +44,8 @@ type taskScheduler struct {
scheduleDuration time.Duration scheduleDuration time.Duration
collectMetricsDuration time.Duration collectMetricsDuration time.Duration
pendingTasks schedulePolicy pendingTasks schedulePolicy
runningTasks map[UniqueID]Task runningTasks *typeutil.ConcurrentMap[UniqueID, Task]
runningQueueLock sync.RWMutex
taskLock *lock.KeyLock[int64] taskLock *lock.KeyLock[int64]
@ -80,7 +80,7 @@ func newTaskScheduler(
cancel: cancel, cancel: cancel,
meta: metaTable, meta: metaTable,
pendingTasks: newFairQueuePolicy(), pendingTasks: newFairQueuePolicy(),
runningTasks: make(map[UniqueID]Task), runningTasks: typeutil.NewConcurrentMap[UniqueID, Task](),
notifyChan: make(chan struct{}, 1), notifyChan: make(chan struct{}, 1),
taskLock: lock.NewKeyLock[int64](), taskLock: lock.NewKeyLock[int64](),
scheduleDuration: Params.DataCoordCfg.IndexTaskSchedulerInterval.GetAsDuration(time.Millisecond), scheduleDuration: Params.DataCoordCfg.IndexTaskSchedulerInterval.GetAsDuration(time.Millisecond),
@ -125,6 +125,10 @@ func (s *taskScheduler) reloadFromMeta() {
State: segIndex.IndexState, State: segIndex.IndexState,
FailReason: segIndex.FailReason, FailReason: segIndex.FailReason,
}, },
req: &workerpb.CreateJobRequest{
ClusterID: Params.CommonCfg.ClusterPrefix.GetValue(),
BuildID: segIndex.BuildID,
},
queueTime: time.Now(), queueTime: time.Now(),
startTime: time.Now(), startTime: time.Now(),
endTime: time.Now(), endTime: time.Now(),
@ -133,9 +137,7 @@ func (s *taskScheduler) reloadFromMeta() {
case commonpb.IndexState_IndexStateNone, commonpb.IndexState_Unissued: case commonpb.IndexState_IndexStateNone, commonpb.IndexState_Unissued:
s.pendingTasks.Push(task) s.pendingTasks.Push(task)
case commonpb.IndexState_InProgress, commonpb.IndexState_Retry: case commonpb.IndexState_InProgress, commonpb.IndexState_Retry:
s.runningQueueLock.Lock() s.runningTasks.Insert(segIndex.BuildID, task)
s.runningTasks[segIndex.BuildID] = task
s.runningQueueLock.Unlock()
} }
} }
} }
@ -150,6 +152,10 @@ func (s *taskScheduler) reloadFromMeta() {
State: t.State, State: t.State,
FailReason: t.FailReason, FailReason: t.FailReason,
}, },
req: &workerpb.AnalyzeRequest{
ClusterID: Params.CommonCfg.ClusterPrefix.GetValue(),
TaskID: taskID,
},
queueTime: time.Now(), queueTime: time.Now(),
startTime: time.Now(), startTime: time.Now(),
endTime: time.Now(), endTime: time.Now(),
@ -158,9 +164,7 @@ func (s *taskScheduler) reloadFromMeta() {
case indexpb.JobState_JobStateNone, indexpb.JobState_JobStateInit: case indexpb.JobState_JobStateNone, indexpb.JobState_JobStateInit:
s.pendingTasks.Push(task) s.pendingTasks.Push(task)
case indexpb.JobState_JobStateInProgress, indexpb.JobState_JobStateRetry: case indexpb.JobState_JobStateInProgress, indexpb.JobState_JobStateRetry:
s.runningQueueLock.Lock() s.runningTasks.Insert(taskID, task)
s.runningTasks[taskID] = task
s.runningQueueLock.Unlock()
} }
} }
@ -176,6 +180,10 @@ func (s *taskScheduler) reloadFromMeta() {
State: t.GetState(), State: t.GetState(),
FailReason: t.GetFailReason(), FailReason: t.GetFailReason(),
}, },
req: &workerpb.CreateStatsRequest{
ClusterID: Params.CommonCfg.ClusterPrefix.GetValue(),
TaskID: taskID,
},
queueTime: time.Now(), queueTime: time.Now(),
startTime: time.Now(), startTime: time.Now(),
endTime: time.Now(), endTime: time.Now(),
@ -208,9 +216,7 @@ func (s *taskScheduler) reloadFromMeta() {
task.taskInfo.FailReason = "segment is not exist or is l0 compacting" task.taskInfo.FailReason = "segment is not exist or is l0 compacting"
} }
} }
s.runningQueueLock.Lock() s.runningTasks.Insert(taskID, task)
s.runningTasks[taskID] = task
s.runningQueueLock.Unlock()
} }
} }
} }
@ -228,34 +234,14 @@ func (s *taskScheduler) exist(taskID UniqueID) bool {
if exist { if exist {
return true return true
} }
_, ok := s.runningTasks.Get(taskID)
s.runningQueueLock.RLock()
defer s.runningQueueLock.RUnlock()
_, ok := s.runningTasks[taskID]
return ok return ok
} }
func (s *taskScheduler) getRunningTask(taskID UniqueID) Task {
s.runningQueueLock.RLock()
defer s.runningQueueLock.RUnlock()
return s.runningTasks[taskID]
}
func (s *taskScheduler) removeRunningTask(taskID UniqueID) {
s.runningQueueLock.Lock()
defer s.runningQueueLock.Unlock()
delete(s.runningTasks, taskID)
}
func (s *taskScheduler) enqueue(task Task) { func (s *taskScheduler) enqueue(task Task) {
defer s.notify() defer s.notify()
taskID := task.GetTaskID() taskID := task.GetTaskID()
_, ok := s.runningTasks.Get(taskID)
s.runningQueueLock.RLock()
_, ok := s.runningTasks[taskID]
s.runningQueueLock.RUnlock()
if !ok { if !ok {
s.pendingTasks.Push(task) s.pendingTasks.Push(task)
task.SetQueueTime(time.Now()) task.SetQueueTime(time.Now())
@ -265,25 +251,21 @@ func (s *taskScheduler) enqueue(task Task) {
func (s *taskScheduler) AbortTask(taskID int64) { func (s *taskScheduler) AbortTask(taskID int64) {
log.Ctx(s.ctx).Info("task scheduler receive abort task request", zap.Int64("taskID", taskID)) log.Ctx(s.ctx).Info("task scheduler receive abort task request", zap.Int64("taskID", taskID))
s.taskLock.Lock(taskID)
defer s.taskLock.Unlock(taskID)
task := s.pendingTasks.Get(taskID) task := s.pendingTasks.Get(taskID)
if task != nil { if task != nil {
s.taskLock.Lock(taskID)
task.SetState(indexpb.JobState_JobStateFailed, "canceled") task.SetState(indexpb.JobState_JobStateFailed, "canceled")
s.taskLock.Unlock(taskID) s.runningTasks.Insert(taskID, task)
s.pendingTasks.Remove(taskID)
return
} }
s.runningQueueLock.Lock() if runningTask, ok := s.runningTasks.Get(taskID); ok {
if task != nil {
s.runningTasks[taskID] = task
}
if runningTask, ok := s.runningTasks[taskID]; ok {
s.taskLock.Lock(taskID)
runningTask.SetState(indexpb.JobState_JobStateFailed, "canceled") runningTask.SetState(indexpb.JobState_JobStateFailed, "canceled")
s.taskLock.Unlock(taskID) s.runningTasks.Insert(taskID, runningTask)
} }
s.runningQueueLock.Unlock()
s.pendingTasks.Remove(taskID)
} }
func (s *taskScheduler) schedule() { func (s *taskScheduler) schedule() {
@ -326,34 +308,29 @@ func (s *taskScheduler) checkProcessingTasksLoop() {
} }
func (s *taskScheduler) checkProcessingTasks() { func (s *taskScheduler) checkProcessingTasks() {
runningTaskIDs := make([]UniqueID, 0) if s.runningTasks.Len() <= 0 {
s.runningQueueLock.RLock() return
for taskID := range s.runningTasks {
runningTaskIDs = append(runningTaskIDs, taskID)
} }
s.runningQueueLock.RUnlock() log.Ctx(s.ctx).Info("check running tasks", zap.Int("runningTask num", s.runningTasks.Len()))
log.Ctx(s.ctx).Info("check running tasks", zap.Int("runningTask num", len(runningTaskIDs)))
allRunningTasks := s.runningTasks.Values()
var wg sync.WaitGroup var wg sync.WaitGroup
sem := make(chan struct{}, 100) sem := make(chan struct{}, 100)
for _, taskID := range runningTaskIDs { for _, task := range allRunningTasks {
wg.Add(1) wg.Add(1)
sem <- struct{}{} sem <- struct{}{}
taskID := taskID go func(task Task) {
go func(taskID int64) {
defer wg.Done() defer wg.Done()
defer func() { defer func() {
<-sem <-sem
}() }()
task := s.getRunningTask(taskID) s.taskLock.Lock(task.GetTaskID())
s.taskLock.Lock(taskID)
suc := s.checkProcessingTask(task) suc := s.checkProcessingTask(task)
s.taskLock.Unlock(taskID) s.taskLock.Unlock(task.GetTaskID())
if suc { if suc {
s.removeRunningTask(taskID) s.runningTasks.Remove(task.GetTaskID())
} }
}(taskID) }(task)
} }
wg.Wait() wg.Wait()
} }
@ -410,13 +387,13 @@ func (s *taskScheduler) run() {
switch task.GetState() { switch task.GetState() {
case indexpb.JobState_JobStateNone: case indexpb.JobState_JobStateNone:
return if !s.processNone(task) {
s.pendingTasks.Push(task)
}
case indexpb.JobState_JobStateInit: case indexpb.JobState_JobStateInit:
s.pendingTasks.Push(task) s.pendingTasks.Push(task)
default: default:
s.runningQueueLock.Lock() s.runningTasks.Insert(task.GetTaskID(), task)
s.runningTasks[task.GetTaskID()] = task
s.runningQueueLock.Unlock()
} }
}(task, nodeID) }(task, nodeID)
} }
@ -433,7 +410,7 @@ func (s *taskScheduler) process(task Task, nodeID int64) bool {
switch task.GetState() { switch task.GetState() {
case indexpb.JobState_JobStateNone: case indexpb.JobState_JobStateNone:
return true return s.processNone(task)
case indexpb.JobState_JobStateInit: case indexpb.JobState_JobStateInit:
return s.processInit(task, nodeID) return s.processInit(task, nodeID)
default: default:
@ -505,11 +482,10 @@ func (s *taskScheduler) collectTaskMetrics() {
collectPendingMetricsFunc(taskID) collectPendingMetricsFunc(taskID)
} }
s.runningQueueLock.RLock() allRunningTasks := s.runningTasks.Values()
for _, task := range s.runningTasks { for _, task := range allRunningTasks {
collectRunningMetricsFunc(task) collectRunningMetricsFunc(task)
} }
s.runningQueueLock.RUnlock()
for taskType, queueingTime := range maxTaskQueueingTime { for taskType, queueingTime := range maxTaskQueueingTime {
metrics.DataCoordTaskExecuteLatency. metrics.DataCoordTaskExecuteLatency.
@ -577,6 +553,14 @@ func (s *taskScheduler) processInit(task Task, nodeID int64) bool {
return true return true
} }
func (s *taskScheduler) processNone(task Task) bool {
if err := task.DropTaskMeta(s.ctx, s.meta); err != nil {
log.Ctx(s.ctx).Warn("set job info failed", zap.Error(err))
return false
}
return true
}
func (s *taskScheduler) processFinished(task Task) bool { func (s *taskScheduler) processFinished(task Task) bool {
if err := task.SetJobInfo(s.meta); err != nil { if err := task.SetJobInfo(s.meta); err != nil {
log.Ctx(s.ctx).Warn("update task info failed", zap.Error(err)) log.Ctx(s.ctx).Warn("update task info failed", zap.Error(err))

852
internal/datacoord/task_scheduler_test.go
View File

File diff suppressed because it is too large Load Diff

26
internal/datacoord/task_stats.go
View File

@ -160,9 +160,10 @@ func (st *statsTask) UpdateMetaBuildingState(meta *meta) error {
} }
func (st *statsTask) PreCheck(ctx context.Context, dependency *taskScheduler) bool { func (st *statsTask) PreCheck(ctx context.Context, dependency *taskScheduler) bool {
log := log.Ctx(ctx).With(zap.Int64("taskID", st.taskID), zap.Int64("segmentID", st.segmentID)) log := log.Ctx(ctx).With(zap.Int64("taskID", st.taskID), zap.Int64("segmentID", st.segmentID),
zap.Int64("targetSegmentID", st.targetSegmentID))
statsMeta := dependency.meta.statsTaskMeta.GetStatsTaskBySegmentID(st.segmentID, st.subJobType) statsMeta := dependency.meta.statsTaskMeta.GetStatsTask(st.taskID)
if statsMeta == nil { if statsMeta == nil {
log.Warn("stats task meta is null, skip it") log.Warn("stats task meta is null, skip it")
st.SetState(indexpb.JobState_JobStateNone, "stats task meta is null") st.SetState(indexpb.JobState_JobStateNone, "stats task meta is null")
@ -241,6 +242,9 @@ func (st *statsTask) PreCheck(ctx context.Context, dependency *taskScheduler) bo
BinlogMaxSize: Params.DataNodeCfg.BinLogMaxSize.GetAsUint64(), BinlogMaxSize: Params.DataNodeCfg.BinLogMaxSize.GetAsUint64(),
} }
log.Info("stats task pre check successfully", zap.String("subJobType", st.subJobType.String()),
zap.Int64("num rows", segment.GetNumOfRows()), zap.Int64("task version", st.req.GetTaskVersion()))
return true return true
} }
@ -297,13 +301,16 @@ func (st *statsTask) QueryResult(ctx context.Context, client types.DataNodeClien
for _, result := range resp.GetStatsJobResults().GetResults() { for _, result := range resp.GetStatsJobResults().GetResults() {
if result.GetTaskID() == st.GetTaskID() { if result.GetTaskID() == st.GetTaskID() {
log.Ctx(ctx).Info("query stats task result success", zap.Int64("taskID", st.GetTaskID()),
zap.Int64("segmentID", st.segmentID), zap.String("result state", result.GetState().String()),
zap.String("failReason", result.GetFailReason()))
if result.GetState() == indexpb.JobState_JobStateFinished || result.GetState() == indexpb.JobState_JobStateRetry || if result.GetState() == indexpb.JobState_JobStateFinished || result.GetState() == indexpb.JobState_JobStateRetry ||
result.GetState() == indexpb.JobState_JobStateFailed { result.GetState() == indexpb.JobState_JobStateFailed {
log.Ctx(ctx).Info("query stats task result success", zap.Int64("taskID", st.GetTaskID()),
zap.Int64("segmentID", st.segmentID), zap.String("result state", result.GetState().String()),
zap.String("failReason", result.GetFailReason()))
st.setResult(result) st.setResult(result)
} else if result.GetState() == indexpb.JobState_JobStateNone { } else if result.GetState() == indexpb.JobState_JobStateNone {
log.Ctx(ctx).Info("query stats task result success", zap.Int64("taskID", st.GetTaskID()),
zap.Int64("segmentID", st.segmentID), zap.String("result state", result.GetState().String()),
zap.String("failReason", result.GetFailReason()))
st.SetState(indexpb.JobState_JobStateRetry, "stats task state is none in info response") st.SetState(indexpb.JobState_JobStateRetry, "stats task state is none in info response")
} }
// inProgress or unissued/init, keep InProgress state // inProgress or unissued/init, keep InProgress state
@ -374,3 +381,12 @@ func (st *statsTask) SetJobInfo(meta *meta) error {
zap.String("subJobType", st.subJobType.String()), zap.String("state", st.taskInfo.GetState().String())) zap.String("subJobType", st.subJobType.String()), zap.String("state", st.taskInfo.GetState().String()))
return nil return nil
} }
func (st *statsTask) DropTaskMeta(ctx context.Context, meta *meta) error {
if err := meta.statsTaskMeta.DropStatsTask(st.taskID); err != nil {
log.Ctx(ctx).Warn("drop stats task failed", zap.Int64("taskID", st.taskID), zap.Error(err))
return err
}
log.Ctx(ctx).Info("drop stats task success", zap.Int64("taskID", st.taskID))
return nil
}

47
internal/datacoord/task_stats_test.go
View File

@ -19,7 +19,6 @@ package datacoord
import ( import (
"context" "context"
"fmt" "fmt"
"sync"
"testing" "testing"
"time" "time"
@ -35,6 +34,8 @@ import (
"github.com/milvus-io/milvus/pkg/v2/proto/datapb" "github.com/milvus-io/milvus/pkg/v2/proto/datapb"
"github.com/milvus-io/milvus/pkg/v2/proto/indexpb" "github.com/milvus-io/milvus/pkg/v2/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v2/proto/workerpb" "github.com/milvus-io/milvus/pkg/v2/proto/workerpb"
"github.com/milvus-io/milvus/pkg/v2/util/lock"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
) )
type statsTaskSuite struct { type statsTaskSuite struct {
@ -55,6 +56,24 @@ func (s *statsTaskSuite) SetupSuite() {
s.segID = 1179 s.segID = 1179
s.targetID = 1180 s.targetID = 1180
tasks := typeutil.NewConcurrentMap[UniqueID, *indexpb.StatsTask]()
statsTask := &indexpb.StatsTask{
CollectionID: 1,
PartitionID: 2,
SegmentID: s.segID,
InsertChannel: "ch1",
TaskID: s.taskID,
SubJobType: indexpb.StatsSubJob_Sort,
Version: 0,
NodeID: 0,
State: indexpb.JobState_JobStateInit,
FailReason: "",
}
tasks.Insert(s.taskID, statsTask)
secondaryIndex := typeutil.NewConcurrentMap[string, *indexpb.StatsTask]()
secondaryKey := createSecondaryIndexKey(statsTask.GetSegmentID(), statsTask.GetSubJobType().String())
secondaryIndex.Insert(secondaryKey, statsTask)
s.mt = &meta{ s.mt = &meta{
segments: &SegmentsInfo{ segments: &SegmentsInfo{
segments: map[int64]*SegmentInfo{ segments: map[int64]*SegmentInfo{
@ -109,23 +128,11 @@ func (s *statsTaskSuite) SetupSuite() {
}, },
statsTaskMeta: &statsTaskMeta{ statsTaskMeta: &statsTaskMeta{
RWMutex: sync.RWMutex{}, keyLock: lock.NewKeyLock[UniqueID](),
ctx: context.Background(), ctx: context.Background(),
catalog: nil, catalog: nil,
tasks: map[int64]*indexpb.StatsTask{ tasks: tasks,
s.taskID: { segmentID2Tasks: secondaryIndex,
CollectionID: 1,
PartitionID: 2,
SegmentID: s.segID,
InsertChannel: "ch1",
TaskID: s.taskID,
SubJobType: indexpb.StatsSubJob_Sort,
Version: 0,
NodeID: 0,
State: indexpb.JobState_JobStateInit,
FailReason: "",
},
},
}, },
} }
} }
@ -595,7 +602,9 @@ func (s *statsTaskSuite) TestTaskStats_PreCheck() {
s.NoError(st.SetJobInfo(s.mt)) s.NoError(st.SetJobInfo(s.mt))
s.NotNil(s.mt.GetHealthySegment(context.TODO(), s.targetID)) s.NotNil(s.mt.GetHealthySegment(context.TODO(), s.targetID))
s.Equal(indexpb.JobState_JobStateFinished, s.mt.statsTaskMeta.tasks[s.taskID].GetState()) t, ok := s.mt.statsTaskMeta.tasks.Get(s.taskID)
s.True(ok)
s.Equal(indexpb.JobState_JobStateFinished, t.GetState())
s.Equal(datapb.SegmentLevel_L2, s.mt.GetHealthySegment(context.TODO(), s.targetID).GetLevel()) s.Equal(datapb.SegmentLevel_L2, s.mt.GetHealthySegment(context.TODO(), s.targetID).GetLevel())
}) })
}) })

1
internal/datacoord/types.go
View File

@ -46,4 +46,5 @@ type Task interface {
SetEndTime(time.Time) SetEndTime(time.Time)
GetEndTime() time.Time GetEndTime() time.Time
GetTaskType() string GetTaskType() string
DropTaskMeta(ctx context.Context, meta *meta) error
} }

13
internal/datanode/index/scheduler.go
View File

@ -33,7 +33,7 @@ import (
// TaskQueue is a queue used to store tasks. // TaskQueue is a queue used to store tasks.
type TaskQueue interface { type TaskQueue interface {
utChan() <-chan int utChan() <-chan struct{}
utEmpty() bool utEmpty() bool
utFull() bool utFull() bool
addUnissuedTask(t Task) error addUnissuedTask(t Task) error
@ -54,12 +54,12 @@ type IndexTaskQueue struct {
// maxTaskNum should keep still // maxTaskNum should keep still
maxTaskNum int64 maxTaskNum int64
utBufChan chan int // to block scheduler utBufChan chan struct{} // to block scheduler
sched *TaskScheduler sched *TaskScheduler
} }
func (queue *IndexTaskQueue) utChan() <-chan int { func (queue *IndexTaskQueue) utChan() <-chan struct{} {
return queue.utBufChan return queue.utBufChan
} }
@ -79,7 +79,10 @@ func (queue *IndexTaskQueue) addUnissuedTask(t Task) error {
return errors.New("index task queue is full") return errors.New("index task queue is full")
} }
queue.unissuedTasks.PushBack(t) queue.unissuedTasks.PushBack(t)
queue.utBufChan <- 1 select {
case queue.utBufChan <- struct{}{}:
default:
}
return nil return nil
} }
@ -159,7 +162,7 @@ func NewIndexBuildTaskQueue(sched *TaskScheduler) *IndexTaskQueue {
activeTasks: make(map[string]Task), activeTasks: make(map[string]Task),
maxTaskNum: 1024, maxTaskNum: 1024,
utBufChan: make(chan int, 1024), utBufChan: make(chan struct{}, 1024),
sched: sched, sched: sched,
} }
} }

5
pkg/metrics/datacoord_metrics.go
View File

@ -309,6 +309,7 @@ var (
}, []string{statusLabelName}) }, []string{statusLabelName})
// IndexTaskNum records the number of index tasks of each type. // IndexTaskNum records the number of index tasks of each type.
// Deprecated: please ues TaskNum after v2.5.5.
IndexTaskNum = prometheus.NewGaugeVec( IndexTaskNum = prometheus.NewGaugeVec(
prometheus.GaugeOpts{ prometheus.GaugeOpts{
Namespace: milvusNamespace, Namespace: milvusNamespace,
@ -350,7 +351,7 @@ var (
Help: "latency of task execute operation", Help: "latency of task execute operation",
Buckets: longTaskBuckets, Buckets: longTaskBuckets,
}, []string{ }, []string{
taskTypeLabel, TaskTypeLabel,
statusLabelName, statusLabelName,
}) })
@ -361,7 +362,7 @@ var (
Subsystem: typeutil.DataCoordRole, Subsystem: typeutil.DataCoordRole,
Name: "task_count", Name: "task_count",
Help: "number of index tasks of each type", Help: "number of index tasks of each type",
}, []string{collectionIDLabelName, taskTypeLabel, taskStateLabel}) }, []string{TaskTypeLabel, TaskStateLabel})
) )
// RegisterDataCoord registers DataCoord metrics // RegisterDataCoord registers DataCoord metrics

4
pkg/metrics/metrics.go
View File

@ -136,8 +136,8 @@ const (
LoadedLabel = "loaded" LoadedLabel = "loaded"
NumEntitiesAllLabel = "all" NumEntitiesAllLabel = "all"
taskTypeLabel = "task_type" TaskTypeLabel = "task_type"
taskStateLabel = "task_state" TaskStateLabel = "task_state"
) )
var ( var (

2
pkg/metrics/proxy_metrics.go
View File

@ -435,7 +435,7 @@ var (
Subsystem: typeutil.ProxyRole, Subsystem: typeutil.ProxyRole,
Name: "queue_task_num", Name: "queue_task_num",
Help: "", Help: "",
}, []string{nodeIDLabelName, queueTypeLabelName, taskStateLabel}) }, []string{nodeIDLabelName, queueTypeLabelName, TaskStateLabel})
ProxyParseExpressionLatency = prometheus.NewHistogramVec( ProxyParseExpressionLatency = prometheus.NewHistogramVec(
prometheus.HistogramOpts{ prometheus.HistogramOpts{

2
pkg/metrics/querycoord_metrics.go
View File

@ -132,7 +132,7 @@ var (
Name: "task_latency", Name: "task_latency",
Help: "latency of all kind of task in query coord scheduler scheduler", Help: "latency of all kind of task in query coord scheduler scheduler",
Buckets: longTaskBuckets, Buckets: longTaskBuckets,
}, []string{collectionIDLabelName, taskTypeLabel, channelNameLabelName}) }, []string{collectionIDLabelName, TaskTypeLabel, channelNameLabelName})
QueryCoordResourceGroupInfo = prometheus.NewGaugeVec( QueryCoordResourceGroupInfo = prometheus.NewGaugeVec(
prometheus.GaugeOpts{ prometheus.GaugeOpts{