milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Improve ListSegments performance for datacoord catalog (#21477) 

Signed-off-by: jaime <yun.zhang@zilliz.com>
pull/21564/head
jaime 2023-01-06 14:33:36 +08:00 committed by GitHub
parent 79535931ab
commit 73d8509600
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
16 changed files with 1922 additions and 623 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
internal/datacoord/channel_store_test.go
View File

@ -17,68 +17,17 @@
package datacoord
import (
"errors"
"fmt"
"testing"
"github.com/milvus-io/milvus/internal/kv"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus/internal/kv"
"github.com/milvus-io/milvus/internal/kv/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
)
type mockTxnKv struct{}
func (m *mockTxnKv) Load(key string) (string, error) {
panic("not implemented") // TODO: Implement
}
func (m *mockTxnKv) MultiLoad(keys []string) ([]string, error) {
panic("not implemented") // TODO: Implement
}
func (m *mockTxnKv) LoadWithPrefix(key string) ([]string, []string, error) {
panic("not implemented") // TODO: Implement
}
func (m *mockTxnKv) Save(key string, value string) error {
panic("not implemented") // TODO: Implement
}
func (m *mockTxnKv) MultiSave(kvs map[string]string) error {
panic("not implemented") // TODO: Implement
}
func (m *mockTxnKv) Remove(key string) error {
panic("not implemented") // TODO: Implement
}
func (m *mockTxnKv) MultiRemove(keys []string) error {
panic("not implemented") // TODO: Implement
}
func (m *mockTxnKv) RemoveWithPrefix(key string) error {
panic("not implemented") // TODO: Implement
}
func (m *mockTxnKv) Close() {
panic("not implemented") // TODO: Implement
}
func (m *mockTxnKv) MultiSaveAndRemove(saves map[string]string, removals []string) error {
if len(saves)+len(removals) > 128 {
return errors.New("too many operations")
}
return nil
}
func (m *mockTxnKv) MultiRemoveWithPrefix(keys []string) error {
panic("not implemented") // TODO: Implement
}
func (m *mockTxnKv) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string) error {
panic("not implemented") // TODO: Implement
}
func genNodeChannelInfos(id int64, num int) *NodeChannelInfo {
channels := make([]*channel, 0, num)
for i := 0; i < num; i++ {
@ -118,6 +67,16 @@ func genChannelOperations(from, to int64, num int) ChannelOpSet {
}
func TestChannelStore_Update(t *testing.T) {
txnKv := mocks.NewTxnKV(t)
txnKv.On("MultiSaveAndRemove",
mock.Anything,
mock.Anything,
).Run(func(args mock.Arguments) {
saves := args.Get(0).(map[string]string)
removals := args.Get(1).([]string)
assert.False(t, len(saves)+len(removals) > 128, "too many operations")
}).Return(nil)
type fields struct {
store kv.TxnKV
channelsInfo map[int64]*NodeChannelInfo
@ -134,7 +93,7 @@ func TestChannelStore_Update(t *testing.T) {
{
"test more than 128 operations",
fields{
&mockTxnKv{},
txnKv,
map[int64]*NodeChannelInfo{
1: genNodeChannelInfos(1, 500),
2: {NodeID: 2},

29
internal/datacoord/compaction_test.go
View File

@ -22,21 +22,20 @@ import (
"testing"
"time"
"github.com/milvus-io/milvus/internal/util/metautil"
"github.com/milvus-io/milvus/internal/util/paramtable"
"github.com/milvus-io/milvus/internal/util/typeutil"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
memkv "github.com/milvus-io/milvus/internal/kv/mem"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/util/tsoutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/util/metautil"
"github.com/milvus-io/milvus/internal/util/paramtable"
"github.com/milvus-io/milvus/internal/util/tsoutil"
"github.com/milvus-io/milvus/internal/util/typeutil"
)
func Test_compactionPlanHandler_execCompactionPlan(t *testing.T) {
@ -299,7 +298,7 @@ func TestCompactionPlanHandler_handleMergeCompactionResult(t *testing.T) {
plans := map[int64]*compactionTask{1: task}
errMeta := &meta{
catalog: &datacoord.Catalog{Txn: &saveFailKV{TxnKV: memkv.NewMemoryKV()}},
catalog: &datacoord.Catalog{MetaKv: &saveFailKV{MetaKv: NewMetaMemoryKV()}},
segments: &SegmentsInfo{
map[int64]*SegmentInfo{
seg1.ID: {SegmentInfo: seg1},
@ -309,7 +308,7 @@ func TestCompactionPlanHandler_handleMergeCompactionResult(t *testing.T) {
}
meta := &meta{
catalog: &datacoord.Catalog{Txn: memkv.NewMemoryKV()},
catalog: &datacoord.Catalog{MetaKv: NewMetaMemoryKV()},
segments: &SegmentsInfo{
map[int64]*SegmentInfo{
seg1.ID: {SegmentInfo: seg1},
@ -445,7 +444,7 @@ func TestCompactionPlanHandler_completeCompaction(t *testing.T) {
plans := map[int64]*compactionTask{1: task}
meta := &meta{
catalog: &datacoord.Catalog{Txn: memkv.NewMemoryKV()},
catalog: &datacoord.Catalog{MetaKv: NewMetaMemoryKV()},
segments: &SegmentsInfo{
map[int64]*SegmentInfo{
seg1.ID: {SegmentInfo: seg1},
@ -537,7 +536,7 @@ func TestCompactionPlanHandler_completeCompaction(t *testing.T) {
plans := map[int64]*compactionTask{1: task}
meta := &meta{
catalog: &datacoord.Catalog{Txn: memkv.NewMemoryKV()},
catalog: &datacoord.Catalog{MetaKv: NewMetaMemoryKV()},
segments: &SegmentsInfo{
map[int64]*SegmentInfo{
seg1.ID: {SegmentInfo: seg1},

93
internal/datacoord/garbage_collector_test.go
View File

@ -27,25 +27,24 @@ import (
"testing"
"time"
"github.com/milvus-io/milvus-proto/go-api/schemapb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/util/funcutil"
"github.com/minio/minio-go/v7"
"github.com/minio/minio-go/v7/pkg/credentials"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/schemapb"
kvmocks "github.com/milvus-io/milvus/internal/kv/mocks"
"github.com/milvus-io/milvus/internal/metastore"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
catalogmocks "github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/util/funcutil"
)
func Test_garbageCollector_basic(t *testing.T) {
@ -327,7 +326,7 @@ func cleanupOSS(cli *minio.Client, bucket, root string) {
cli.RemoveBucket(context.TODO(), bucket)
}
func createMetaForRecycleUnusedIndexes(catalog *datacoord.Catalog) *meta {
func createMetaForRecycleUnusedIndexes(catalog metastore.DataCoordCatalog) *meta {
var (
ctx = context.Background()
collID = UniqueID(100)
@ -394,21 +393,33 @@ func createMetaForRecycleUnusedIndexes(catalog *datacoord.Catalog) *meta {
func TestGarbageCollector_recycleUnusedIndexes(t *testing.T) {
t.Run("success", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("DropIndex",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
gc := &garbageCollector{
meta: createMetaForRecycleUnusedIndexes(&datacoord.Catalog{Txn: &mockEtcdKv{}}),
meta: createMetaForRecycleUnusedIndexes(catalog),
}
gc.recycleUnusedIndexes()
})
t.Run("fail", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("DropIndex",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
gc := &garbageCollector{
meta: createMetaForRecycleUnusedIndexes(&datacoord.Catalog{Txn: &removeFailKV{}}),
meta: createMetaForRecycleUnusedIndexes(catalog),
}
gc.recycleUnusedIndexes()
})
}
func createMetaForRecycleUnusedSegIndexes(catalog *datacoord.Catalog) *meta {
func createMetaForRecycleUnusedSegIndexes(catalog metastore.DataCoordCatalog) *meta {
var (
ctx = context.Background()
collID = UniqueID(100)
@ -521,15 +532,31 @@ func createMetaForRecycleUnusedSegIndexes(catalog *datacoord.Catalog) *meta {
func TestGarbageCollector_recycleUnusedSegIndexes(t *testing.T) {
t.Run("success", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("DropSegmentIndex",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(nil)
gc := &garbageCollector{
meta: createMetaForRecycleUnusedSegIndexes(&datacoord.Catalog{Txn: &mockEtcdKv{}}),
meta: createMetaForRecycleUnusedSegIndexes(catalog),
}
gc.recycleUnusedSegIndexes()
})
t.Run("fail", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("DropSegmentIndex",
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
gc := &garbageCollector{
meta: createMetaForRecycleUnusedSegIndexes(&datacoord.Catalog{Txn: &removeFailKV{}}),
meta: createMetaForRecycleUnusedSegIndexes(catalog),
}
gc.recycleUnusedSegIndexes()
})
@ -676,7 +703,7 @@ func TestGarbageCollector_recycleUnusedIndexFiles(t *testing.T) {
cm.EXPECT().RemoveWithPrefix(mock.Anything, mock.Anything).Return(nil)
cm.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
gc := &garbageCollector{
meta: createMetaTableForRecycleUnusedIndexFiles(&datacoord.Catalog{Txn: &mockEtcdKv{}}),
meta: createMetaTableForRecycleUnusedIndexFiles(&datacoord.Catalog{MetaKv: kvmocks.NewMetaKv(t)}),
option: GcOption{
cli: cm,
},
@ -689,7 +716,7 @@ func TestGarbageCollector_recycleUnusedIndexFiles(t *testing.T) {
cm.EXPECT().RootPath().Return("root")
cm.EXPECT().ListWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, errors.New("error"))
gc := &garbageCollector{
meta: createMetaTableForRecycleUnusedIndexFiles(&datacoord.Catalog{Txn: &mockEtcdKv{}}),
meta: createMetaTableForRecycleUnusedIndexFiles(&datacoord.Catalog{MetaKv: kvmocks.NewMetaKv(t)}),
option: GcOption{
cli: cm,
},
@ -704,7 +731,7 @@ func TestGarbageCollector_recycleUnusedIndexFiles(t *testing.T) {
cm.EXPECT().ListWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return([]string{"a/b/c/", "a/b/600/", "a/b/601/", "a/b/602/"}, nil, nil)
cm.EXPECT().RemoveWithPrefix(mock.Anything, mock.Anything).Return(nil)
gc := &garbageCollector{
meta: createMetaTableForRecycleUnusedIndexFiles(&datacoord.Catalog{Txn: &mockEtcdKv{}}),
meta: createMetaTableForRecycleUnusedIndexFiles(&datacoord.Catalog{MetaKv: kvmocks.NewMetaKv(t)}),
option: GcOption{
cli: cm,
},
@ -719,7 +746,7 @@ func TestGarbageCollector_recycleUnusedIndexFiles(t *testing.T) {
cm.EXPECT().ListWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return([]string{"a/b/c/", "a/b/600/", "a/b/601/", "a/b/602/"}, nil, nil)
cm.EXPECT().RemoveWithPrefix(mock.Anything, mock.Anything).Return(errors.New("error"))
gc := &garbageCollector{
meta: createMetaTableForRecycleUnusedIndexFiles(&datacoord.Catalog{Txn: &mockEtcdKv{}}),
meta: createMetaTableForRecycleUnusedIndexFiles(&datacoord.Catalog{MetaKv: kvmocks.NewMetaKv(t)}),
option: GcOption{
cli: cm,
},
@ -729,8 +756,22 @@ func TestGarbageCollector_recycleUnusedIndexFiles(t *testing.T) {
}
func TestGarbageCollector_clearETCD(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("CreateSegmentIndex",
mock.Anything,
mock.Anything,
).Return(nil)
catalog.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
catalog.On("DropSegment",
mock.Anything,
mock.Anything,
).Return(nil)
m := &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: catalog,
segments: &SegmentsInfo{
map[UniqueID]*SegmentInfo{
segID: {

52
internal/datacoord/index_builder_test.go
View File

@ -22,17 +22,19 @@ import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus/internal/indexnode"
"github.com/milvus-io/milvus/internal/metastore"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
catalogmocks "github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/util/paramtable"
"github.com/stretchr/testify/assert"
)
var (
@ -601,7 +603,17 @@ func TestIndexBuilder(t *testing.T) {
Params.InitOnce()
ctx := context.Background()
mt := createMetaTable(&datacoord.Catalog{Txn: &mockEtcdKv{}})
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("CreateSegmentIndex",
mock.Anything,
mock.Anything,
).Return(nil)
catalog.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
mt := createMetaTable(catalog)
nodeManager := &IndexNodeManager{
ctx: ctx,
nodeClients: map[UniqueID]types.IndexNode{
@ -664,6 +676,17 @@ func TestIndexBuilder(t *testing.T) {
func TestIndexBuilder_Error(t *testing.T) {
Params.Init()
sc := catalogmocks.NewDataCoordCatalog(t)
sc.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
chunkManager := &mocks.ChunkManager{}
chunkManager.EXPECT().RootPath().Return("root")
ib := &indexBuilder{
@ -671,8 +694,7 @@ func TestIndexBuilder_Error(t *testing.T) {
tasks: map[int64]indexTaskState{
buildID: indexTaskInit,
},
meta: createMetaTable(&datacoord.Catalog{
Txn: &saveFailKV{}}),
meta: createMetaTable(ec),
chunkManager: chunkManager,
}
@ -728,7 +750,7 @@ func TestIndexBuilder_Error(t *testing.T) {
t.Run("assign task error", func(t *testing.T) {
paramtable.Get().Save(Params.CommonCfg.StorageType.Key, "local")
ib.tasks[buildID] = indexTaskInit
ib.meta.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
ib.meta.catalog = sc
ib.nodeManager = &IndexNodeManager{
ctx: context.Background(),
nodeClients: map[UniqueID]types.IndexNode{
@ -756,7 +778,7 @@ func TestIndexBuilder_Error(t *testing.T) {
})
t.Run("assign task fail", func(t *testing.T) {
paramtable.Get().Save(Params.CommonCfg.StorageType.Key, "local")
ib.meta.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
ib.meta.catalog = sc
ib.nodeManager = &IndexNodeManager{
ctx: context.Background(),
nodeClients: map[UniqueID]types.IndexNode{
@ -789,7 +811,7 @@ func TestIndexBuilder_Error(t *testing.T) {
t.Run("drop job error", func(t *testing.T) {
ib.meta.buildID2SegmentIndex[buildID].NodeID = nodeID
ib.meta.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
ib.meta.catalog = sc
ib.nodeManager = &IndexNodeManager{
ctx: context.Background(),
nodeClients: map[UniqueID]types.IndexNode{
@ -819,7 +841,7 @@ func TestIndexBuilder_Error(t *testing.T) {
t.Run("drop job fail", func(t *testing.T) {
ib.meta.buildID2SegmentIndex[buildID].NodeID = nodeID
ib.meta.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
ib.meta.catalog = sc
ib.nodeManager = &IndexNodeManager{
ctx: context.Background(),
nodeClients: map[UniqueID]types.IndexNode{
@ -850,7 +872,7 @@ func TestIndexBuilder_Error(t *testing.T) {
t.Run("get state error", func(t *testing.T) {
ib.meta.buildID2SegmentIndex[buildID].NodeID = nodeID
ib.meta.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
ib.meta.catalog = sc
ib.nodeManager = &IndexNodeManager{
ctx: context.Background(),
nodeClients: map[UniqueID]types.IndexNode{
@ -872,7 +894,7 @@ func TestIndexBuilder_Error(t *testing.T) {
t.Run("get state fail", func(t *testing.T) {
ib.meta.buildID2SegmentIndex[buildID].NodeID = nodeID
ib.meta.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
ib.meta.catalog = sc
ib.nodeManager = &IndexNodeManager{
ctx: context.Background(),
nodeClients: map[UniqueID]types.IndexNode{
@ -899,7 +921,7 @@ func TestIndexBuilder_Error(t *testing.T) {
t.Run("finish task fail", func(t *testing.T) {
ib.meta.buildID2SegmentIndex[buildID].NodeID = nodeID
ib.meta.catalog = &datacoord.Catalog{Txn: &saveFailKV{}}
ib.meta.catalog = ec
ib.nodeManager = &IndexNodeManager{
ctx: context.Background(),
nodeClients: map[UniqueID]types.IndexNode{
@ -935,7 +957,7 @@ func TestIndexBuilder_Error(t *testing.T) {
t.Run("task still in progress", func(t *testing.T) {
ib.meta.buildID2SegmentIndex[buildID].NodeID = nodeID
ib.meta.catalog = &datacoord.Catalog{Txn: &saveFailKV{}}
ib.meta.catalog = ec
ib.nodeManager = &IndexNodeManager{
ctx: context.Background(),
nodeClients: map[UniqueID]types.IndexNode{
@ -971,7 +993,7 @@ func TestIndexBuilder_Error(t *testing.T) {
t.Run("indexNode has no task", func(t *testing.T) {
ib.meta.buildID2SegmentIndex[buildID].NodeID = nodeID
ib.meta.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
ib.meta.catalog = sc
ib.nodeManager = &IndexNodeManager{
ctx: context.Background(),
nodeClients: map[UniqueID]types.IndexNode{
@ -999,7 +1021,7 @@ func TestIndexBuilder_Error(t *testing.T) {
t.Run("node not exist", func(t *testing.T) {
ib.meta.buildID2SegmentIndex[buildID].NodeID = nodeID
ib.meta.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
ib.meta.catalog = sc
ib.nodeManager = &IndexNodeManager{
ctx: context.Background(),
nodeClients: map[UniqueID]types.IndexNode{},

114
internal/datacoord/index_meta_test.go
View File

@ -19,14 +19,18 @@ package datacoord
import (
"context"
"errors"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus/internal/kv/mocks"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
catalogmocks "github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
@ -52,12 +56,17 @@ func TestMeta_CanCreateIndex(t *testing.T) {
},
}
)
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("CreateIndex",
mock.Anything,
mock.Anything,
).Return(nil)
m := &meta{
RWMutex: sync.RWMutex{},
ctx: context.Background(),
catalog: &datacoord.Catalog{
Txn: &mockEtcdKv{},
},
RWMutex: sync.RWMutex{},
ctx: context.Background(),
catalog: catalog,
collections: nil,
segments: nil,
channelCPs: nil,
@ -169,11 +178,9 @@ func TestMeta_HasSameReq(t *testing.T) {
}
)
m := &meta{
RWMutex: sync.RWMutex{},
ctx: context.Background(),
catalog: &datacoord.Catalog{
Txn: &mockEtcdKv{},
},
RWMutex: sync.RWMutex{},
ctx: context.Background(),
catalog: catalogmocks.NewDataCoordCatalog(t),
collections: nil,
segments: nil,
channelCPs: nil,
@ -257,10 +264,16 @@ func TestMeta_CreateIndex(t *testing.T) {
}
t.Run("success", func(t *testing.T) {
sc := catalogmocks.NewDataCoordCatalog(t)
sc.On("CreateIndex",
mock.Anything,
mock.Anything,
).Return(nil)
m := &meta{
RWMutex: sync.RWMutex{},
ctx: context.Background(),
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: sc,
indexes: make(map[UniqueID]map[UniqueID]*model.Index),
buildID2SegmentIndex: make(map[UniqueID]*model.SegmentIndex),
}
@ -270,10 +283,16 @@ func TestMeta_CreateIndex(t *testing.T) {
})
t.Run("save fail", func(t *testing.T) {
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("CreateIndex",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
m := &meta{
RWMutex: sync.RWMutex{},
ctx: context.Background(),
catalog: &datacoord.Catalog{Txn: &saveFailKV{}},
catalog: ec,
indexes: make(map[UniqueID]map[UniqueID]*model.Index),
buildID2SegmentIndex: make(map[UniqueID]*model.SegmentIndex),
}
@ -284,10 +303,22 @@ func TestMeta_CreateIndex(t *testing.T) {
}
func TestMeta_AddSegmentIndex(t *testing.T) {
sc := catalogmocks.NewDataCoordCatalog(t)
sc.On("CreateSegmentIndex",
mock.Anything,
mock.Anything,
).Return(nil)
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("CreateSegmentIndex",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
m := &meta{
RWMutex: sync.RWMutex{},
ctx: context.Background(),
catalog: &datacoord.Catalog{Txn: &saveFailKV{}},
catalog: ec,
indexes: make(map[UniqueID]map[UniqueID]*model.Index),
buildID2SegmentIndex: make(map[UniqueID]*model.SegmentIndex),
segments: &SegmentsInfo{
@ -330,7 +361,7 @@ func TestMeta_AddSegmentIndex(t *testing.T) {
})
t.Run("success", func(t *testing.T) {
m.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
m.catalog = sc
err := m.AddSegmentIndex(segmentIndex)
assert.NoError(t, err)
})
@ -359,7 +390,7 @@ func TestMeta_GetIndexIDByName(t *testing.T) {
m := &meta{
RWMutex: sync.RWMutex{},
ctx: context.Background(),
catalog: &datacoord.Catalog{Txn: &saveFailKV{}},
catalog: &datacoord.Catalog{MetaKv: &saveFailKV{}},
indexes: make(map[UniqueID]map[UniqueID]*model.Index),
buildID2SegmentIndex: make(map[UniqueID]*model.SegmentIndex),
}
@ -417,7 +448,7 @@ func TestMeta_GetSegmentIndexState(t *testing.T) {
m := &meta{
RWMutex: sync.RWMutex{},
ctx: context.Background(),
catalog: &datacoord.Catalog{Txn: &saveFailKV{}},
catalog: &datacoord.Catalog{MetaKv: &saveFailKV{}},
indexes: map[UniqueID]map[UniqueID]*model.Index{},
buildID2SegmentIndex: make(map[UniqueID]*model.SegmentIndex),
segments: &SegmentsInfo{
@ -628,8 +659,19 @@ func TestMeta_GetSegmentIndexStateOnField(t *testing.T) {
}
func TestMeta_MarkIndexAsDeleted(t *testing.T) {
sc := catalogmocks.NewDataCoordCatalog(t)
sc.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
m := &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: sc,
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
@ -663,13 +705,13 @@ func TestMeta_MarkIndexAsDeleted(t *testing.T) {
}
t.Run("fail", func(t *testing.T) {
m.catalog = &datacoord.Catalog{Txn: &saveFailKV{}}
m.catalog = ec
err := m.MarkIndexAsDeleted(collID, []UniqueID{indexID, indexID + 1, indexID + 2})
assert.Error(t, err)
})
t.Run("success", func(t *testing.T) {
m.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
m.catalog = sc
err := m.MarkIndexAsDeleted(collID, []UniqueID{indexID, indexID + 1, indexID + 2})
assert.NoError(t, err)
@ -682,7 +724,7 @@ func TestMeta_MarkIndexAsDeleted(t *testing.T) {
}
func TestMeta_GetSegmentIndexes(t *testing.T) {
m := createMetaTable(&datacoord.Catalog{Txn: &mockEtcdKv{}})
m := createMetaTable(&datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)})
t.Run("success", func(t *testing.T) {
segIndexes := m.GetSegmentIndexes(segID)
@ -960,9 +1002,15 @@ func TestMeta_IsIndexExist(t *testing.T) {
})
}
func updateSegmentIndexMeta() *meta {
func updateSegmentIndexMeta(t *testing.T) *meta {
sc := catalogmocks.NewDataCoordCatalog(t)
sc.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
return &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: sc,
segments: &SegmentsInfo{
segments: map[UniqueID]*SegmentInfo{
segID: {
@ -1036,7 +1084,12 @@ func updateSegmentIndexMeta() *meta {
}
func TestMeta_UpdateVersion(t *testing.T) {
m := updateSegmentIndexMeta()
m := updateSegmentIndexMeta(t)
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
t.Run("success", func(t *testing.T) {
err := m.UpdateVersion(buildID, nodeID)
@ -1044,7 +1097,7 @@ func TestMeta_UpdateVersion(t *testing.T) {
})
t.Run("fail", func(t *testing.T) {
m.catalog = &datacoord.Catalog{Txn: &saveFailKV{}}
m.catalog = ec
err := m.UpdateVersion(buildID, nodeID)
assert.Error(t, err)
})
@ -1056,7 +1109,7 @@ func TestMeta_UpdateVersion(t *testing.T) {
}
func TestMeta_FinishTask(t *testing.T) {
m := updateSegmentIndexMeta()
m := updateSegmentIndexMeta(t)
t.Run("success", func(t *testing.T) {
err := m.FinishTask(&indexpb.IndexTaskInfo{
@ -1071,7 +1124,7 @@ func TestMeta_FinishTask(t *testing.T) {
t.Run("fail", func(t *testing.T) {
m.catalog = &datacoord.Catalog{
Txn: &saveFailKV{},
MetaKv: &saveFailKV{},
}
err := m.FinishTask(&indexpb.IndexTaskInfo{
BuildID: buildID,
@ -1096,7 +1149,12 @@ func TestMeta_FinishTask(t *testing.T) {
}
func TestMeta_BuildIndex(t *testing.T) {
m := updateSegmentIndexMeta()
m := updateSegmentIndexMeta(t)
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
t.Run("success", func(t *testing.T) {
err := m.BuildIndex(buildID)
@ -1104,7 +1162,7 @@ func TestMeta_BuildIndex(t *testing.T) {
})
t.Run("fail", func(t *testing.T) {
m.catalog = &datacoord.Catalog{Txn: &saveFailKV{}}
m.catalog = ec
err := m.BuildIndex(buildID)
assert.Error(t, err)
})

50
internal/datacoord/index_service_test.go
View File

@ -18,13 +18,17 @@ package datacoord
import (
"context"
"errors"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus/internal/kv/mocks"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
catalogmocks "github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/storage"
@ -60,14 +64,21 @@ func TestServer_CreateIndex(t *testing.T) {
}
ctx = context.Background()
)
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("CreateIndex",
mock.Anything,
mock.Anything,
).Return(nil)
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: catalog,
indexes: map[UniqueID]map[UniqueID]*model.Index{},
},
allocator: newMockAllocator(),
notifyIndexChan: make(chan UniqueID, 1),
}
s.stateCode.Store(commonpb.StateCode_Healthy)
t.Run("success", func(t *testing.T) {
resp, err := s.CreateIndex(ctx, req)
@ -116,7 +127,7 @@ func TestServer_CreateIndex(t *testing.T) {
t.Run("save index fail", func(t *testing.T) {
s.meta.indexes = map[UniqueID]map[UniqueID]*model.Index{}
s.meta.catalog = &datacoord.Catalog{Txn: &saveFailKV{}}
s.meta.catalog = &datacoord.Catalog{MetaKv: &saveFailKV{}}
req.IndexParams = []*commonpb.KeyValuePair{
{
Key: "index_type",
@ -158,7 +169,7 @@ func TestServer_GetIndexState(t *testing.T) {
)
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
},
allocator: newMockAllocator(),
notifyIndexChan: make(chan UniqueID, 1),
@ -179,7 +190,7 @@ func TestServer_GetIndexState(t *testing.T) {
})
s.meta = &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
@ -276,7 +287,7 @@ func TestServer_GetSegmentIndexState(t *testing.T) {
)
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
indexes: map[UniqueID]map[UniqueID]*model.Index{},
segments: &SegmentsInfo{map[UniqueID]*SegmentInfo{}},
},
@ -413,7 +424,7 @@ func TestServer_GetIndexBuildProgress(t *testing.T) {
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
indexes: map[UniqueID]map[UniqueID]*model.Index{},
segments: &SegmentsInfo{map[UniqueID]*SegmentInfo{}},
},
@ -561,9 +572,15 @@ func TestServer_DescribeIndex(t *testing.T) {
}
)
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: catalog,
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
//finished
@ -817,9 +834,15 @@ func TestServer_DropIndex(t *testing.T) {
}
)
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: catalog,
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
//finished
@ -923,14 +946,19 @@ func TestServer_DropIndex(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Healthy)
t.Run("drop fail", func(t *testing.T) {
s.meta.catalog = &datacoord.Catalog{Txn: &saveFailKV{}}
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
s.meta.catalog = catalog
resp, err := s.DropIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
})
t.Run("drop one index", func(t *testing.T) {
s.meta.catalog = &datacoord.Catalog{Txn: &mockEtcdKv{}}
s.meta.catalog = catalog
resp, err := s.DropIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
@ -1009,7 +1037,7 @@ func TestServer_GetIndexInfos(t *testing.T) {
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{Txn: &mockEtcdKv{}},
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
//finished

6
internal/datacoord/meta.go
View File

@ -32,10 +32,8 @@ import (
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/schemapb"
"github.com/milvus-io/milvus/internal/common"
"github.com/milvus-io/milvus/internal/kv"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/metastore"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/metrics"
"github.com/milvus-io/milvus/internal/proto/datapb"
@ -80,10 +78,10 @@ type collectionInfo struct {
}
// NewMeta creates meta from provided `kv.TxnKV`
func newMeta(ctx context.Context, kv kv.TxnKV, chunkManagerRootPath string, chunkManager storage.ChunkManager) (*meta, error) {
func newMeta(ctx context.Context, catalog metastore.DataCoordCatalog, chunkManager storage.ChunkManager) (*meta, error) {
mt := &meta{
ctx: ctx,
catalog: &datacoord.Catalog{Txn: kv, ChunkManagerRootPath: chunkManagerRootPath},
catalog: catalog,
collections: make(map[UniqueID]*collectionInfo),
segments: NewSegmentsInfo(),
channelCPs: make(map[string]*internalpb.MsgPosition),

434
internal/datacoord/meta_test.go
View File

@ -18,312 +18,128 @@ package datacoord
import (
"context"
"errors"
"fmt"
"strings"
"testing"
"github.com/golang/protobuf/proto"
"github.com/samber/lo"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/golang/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus/internal/common"
"github.com/milvus-io/milvus/internal/kv"
memkv "github.com/milvus-io/milvus/internal/kv/mem"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
"github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/util"
)
type mockEtcdKv struct {
kv.TxnKV
}
func (mek *mockEtcdKv) LoadWithPrefix(key string) ([]string, []string, error) {
var val []byte
switch {
case strings.Contains(key, datacoord.SegmentPrefix):
segInfo := &datapb.SegmentInfo{ID: 1, Binlogs: []*datapb.FieldBinlog{getFieldBinlogPaths(1, "log1")}}
val, _ = proto.Marshal(segInfo)
case strings.Contains(key, datacoord.SegmentBinlogPathPrefix):
segInfo := getFieldBinlogPaths(1, "binlog1")
val, _ = proto.Marshal(segInfo)
case strings.Contains(key, datacoord.SegmentDeltalogPathPrefix):
segInfo := getFieldBinlogPaths(1, "deltalog1")
val, _ = proto.Marshal(segInfo)
case strings.Contains(key, datacoord.SegmentStatslogPathPrefix):
segInfo := getFieldBinlogPaths(1, "statslog1")
val, _ = proto.Marshal(segInfo)
case strings.Contains(key, datacoord.ChannelCheckpointPrefix):
channelCP := &internalpb.MsgPosition{
Timestamp: 1000,
}
val, _ = proto.Marshal(channelCP)
case strings.Contains(key, util.FieldIndexPrefix):
index := &datapb.FieldIndex{
IndexInfo: &datapb.IndexInfo{
CollectionID: 0,
FieldID: 1,
IndexName: "_default_idx_101",
IndexID: 0,
TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}},
IndexParams: []*commonpb.KeyValuePair{{Key: "index_type", Value: "HNSW"}},
IsAutoIndex: false,
UserIndexParams: nil,
},
Deleted: false,
CreateTime: 0,
}
val, _ = proto.Marshal(index)
case strings.Contains(key, util.SegmentIndexPrefix):
segIndex := &datapb.SegmentIndex{
CollectionID: 0,
PartitionID: 0,
SegmentID: 1,
NumRows: 1025,
IndexID: 0,
BuildID: 0,
NodeID: 0,
IndexVersion: 0,
State: commonpb.IndexState_Unissued,
FailReason: "",
IndexFileKeys: nil,
Deleted: false,
CreateTime: 0,
SerializeSize: 0,
WriteHandoff: false,
}
val, _ = proto.Marshal(segIndex)
default:
return nil, nil, fmt.Errorf("invalid key")
}
return nil, []string{string(val)}, nil
}
func (mek *mockEtcdKv) Save(key, value string) error {
return nil
}
func (mek *mockEtcdKv) MultiSave(kvs map[string]string) error {
return nil
}
func (mek *mockEtcdKv) Remove(key string) error {
return nil
}
func (mek *mockEtcdKv) MultiRemove(keys []string) error {
return nil
}
type mockKvLoadSegmentError struct {
kv.TxnKV
}
func (mek *mockKvLoadSegmentError) LoadWithPrefix(key string) ([]string, []string, error) {
if strings.Contains(key, datacoord.SegmentPrefix) {
return nil, nil, fmt.Errorf("segment LoadWithPrefix error")
}
return nil, nil, nil
}
type mockKvLoadBinlogError struct {
kv.TxnKV
}
func (mek *mockKvLoadBinlogError) LoadWithPrefix(key string) ([]string, []string, error) {
var val []byte
switch {
case strings.Contains(key, datacoord.SegmentPrefix):
segInfo := &datapb.SegmentInfo{ID: 1, Deltalogs: []*datapb.FieldBinlog{getFieldBinlogPaths(1, "delta_log_1")}}
val, _ = proto.Marshal(segInfo)
case strings.Contains(key, datacoord.SegmentBinlogPathPrefix):
return nil, nil, fmt.Errorf("LoadWithPrefix for binlogs error")
}
return nil, []string{string(val)}, nil
}
type mockKvLoadDeltaBinlogError struct {
kv.TxnKV
}
func (mek *mockKvLoadDeltaBinlogError) LoadWithPrefix(key string) ([]string, []string, error) {
var val []byte
switch {
case strings.Contains(key, datacoord.SegmentPrefix):
segInfo := &datapb.SegmentInfo{ID: 1, Binlogs: []*datapb.FieldBinlog{getFieldBinlogPaths(1, "binlog_1")}}
val, _ = proto.Marshal(segInfo)
case strings.Contains(key, datacoord.SegmentDeltalogPathPrefix):
return nil, nil, fmt.Errorf("LoadWithPrefix for deltalog error")
}
return nil, []string{string(val)}, nil
}
type mockKvLoadStatsBinlogError struct {
kv.TxnKV
}
func (mek *mockKvLoadStatsBinlogError) LoadWithPrefix(key string) ([]string, []string, error) {
var val []byte
switch {
case strings.Contains(key, datacoord.SegmentPrefix+"/"):
segInfo := &datapb.SegmentInfo{ID: 1, Binlogs: []*datapb.FieldBinlog{getFieldBinlogPaths(1, "binlog_1")}}
val, _ = proto.Marshal(segInfo)
case strings.Contains(key, datacoord.SegmentStatslogPathPrefix):
return nil, nil, fmt.Errorf("LoadWithPrefix for statslog error")
}
return nil, []string{string(val)}, nil
}
type mockKvIllegalSegment struct {
kv.TxnKV
}
func (mek *mockKvIllegalSegment) LoadWithPrefix(key string) ([]string, []string, error) {
var val []byte
switch {
case strings.Contains(key, datacoord.SegmentPrefix):
val = []byte{'i', 'l', 'l', 'e', 'g', 'a', 'l'}
}
return nil, []string{string(val)}, nil
}
type mockKvIllegalBinlog struct {
kv.TxnKV
}
func (mek *mockKvIllegalBinlog) LoadWithPrefix(key string) ([]string, []string, error) {
var val []byte
switch {
case strings.Contains(key, datacoord.SegmentBinlogPathPrefix):
val = []byte{'i', 'l', 'l', 'e', 'g', 'a', 'l'}
}
return nil, []string{string(val)}, nil
}
type mockKvIllegalDeltalog struct {
kv.TxnKV
}
func (mek *mockKvIllegalDeltalog) LoadWithPrefix(key string) ([]string, []string, error) {
var val []byte
switch {
case strings.Contains(key, datacoord.SegmentDeltalogPathPrefix):
val = []byte{'i', 'l', 'l', 'e', 'g', 'a', 'l'}
}
return nil, []string{string(val)}, nil
}
type mockKvIllegalStatslog struct {
kv.TxnKV
}
func (mek *mockKvIllegalStatslog) LoadWithPrefix(key string) ([]string, []string, error) {
var val []byte
switch {
case strings.Contains(key, datacoord.SegmentStatslogPathPrefix):
val = []byte{'i', 'l', 'l', 'e', 'g', 'a', 'l'}
}
return nil, []string{string(val)}, nil
}
type mockLoadIndexError struct {
kv.TxnKV
}
func (mek *mockLoadIndexError) LoadWithPrefix(key string) ([]string, []string, error) {
switch {
case strings.Contains(key, util.FieldIndexPrefix):
return nil, nil, fmt.Errorf("LoadWithPrefix for index error")
}
return nil, nil, nil
}
type mockLoadSegmentIndexError struct {
kv.TxnKV
}
func (mek *mockLoadSegmentIndexError) LoadWithPrefix(key string) ([]string, []string, error) {
switch {
case strings.Contains(key, util.SegmentIndexPrefix):
return nil, nil, fmt.Errorf("LoadWithPrefix for segment index error")
}
return nil, nil, nil
}
func TestMetaReloadFromKV(t *testing.T) {
t.Run("Test ReloadFromKV success", func(t *testing.T) {
fkv := &mockEtcdKv{}
_, err := newMeta(context.TODO(), fkv, "", nil)
assert.Nil(t, err)
})
t.Run("ListSegments fail", func(t *testing.T) {
catalog := mocks.NewDataCoordCatalog(t)
catalog.On("ListSegments",
mock.Anything,
).Return(nil, errors.New("error"))
// load segment error
t.Run("Test ReloadFromKV load segment fails", func(t *testing.T) {
fkv := &mockKvLoadSegmentError{}
_, err := newMeta(context.TODO(), fkv, "", nil)
assert.NotNil(t, err)
})
// illegal segment info
t.Run("Test ReloadFromKV unmarshal segment fails", func(t *testing.T) {
fkv := &mockKvIllegalSegment{}
_, err := newMeta(context.TODO(), fkv, "", nil)
assert.NotNil(t, err)
})
// load binlog/deltalog/statslog error
t.Run("Test ReloadFromKV load binlog fails", func(t *testing.T) {
fkv := &mockKvLoadBinlogError{}
_, err := newMeta(context.TODO(), fkv, "", nil)
assert.NotNil(t, err)
})
t.Run("Test ReloadFromKV load deltalog fails", func(t *testing.T) {
fkv := &mockKvLoadDeltaBinlogError{}
_, err := newMeta(context.TODO(), fkv, "", nil)
assert.NotNil(t, err)
})
t.Run("Test ReloadFromKV load statslog fails", func(t *testing.T) {
fkv := &mockKvLoadStatsBinlogError{}
_, err := newMeta(context.TODO(), fkv, "", nil)
assert.NotNil(t, err)
})
// illegal binlog/deltalog/statslog info
t.Run("Test ReloadFromKV unmarshal binlog fails", func(t *testing.T) {
fkv := &mockKvIllegalBinlog{}
_, err := newMeta(context.TODO(), fkv, "", nil)
assert.NotNil(t, err)
})
t.Run("Test ReloadFromKV unmarshal deltalog fails", func(t *testing.T) {
fkv := &mockKvIllegalDeltalog{}
_, err := newMeta(context.TODO(), fkv, "", nil)
assert.NotNil(t, err)
})
t.Run("Test ReloadFromKV unmarshal statslog fails", func(t *testing.T) {
fkv := &mockKvIllegalStatslog{}
_, err := newMeta(context.TODO(), fkv, "", nil)
assert.NotNil(t, err)
})
t.Run("Test ReloadFromKV load index fails", func(t *testing.T) {
fkv := &mockLoadIndexError{}
_, err := newMeta(context.TODO(), fkv, "", nil)
_, err := newMeta(context.TODO(), catalog, nil)
assert.Error(t, err)
})
t.Run("Test ReloadFromKV load segment index fails", func(t *testing.T) {
fkv := &mockLoadSegmentIndexError{}
_, err := newMeta(context.TODO(), fkv, "", nil)
t.Run("ListChannelCheckpoint fail", func(t *testing.T) {
catalog := mocks.NewDataCoordCatalog(t)
catalog.On("ListSegments",
mock.Anything,
).Return([]*datapb.SegmentInfo{}, nil)
catalog.On("ListChannelCheckpoint",
mock.Anything,
).Return(nil, errors.New("error"))
_, err := newMeta(context.TODO(), catalog, nil)
assert.Error(t, err)
})
t.Run("ListIndexes fail", func(t *testing.T) {
catalog := mocks.NewDataCoordCatalog(t)
catalog.On("ListSegments",
mock.Anything,
).Return([]*datapb.SegmentInfo{}, nil)
catalog.On("ListChannelCheckpoint",
mock.Anything,
).Return(map[string]*internalpb.MsgPosition{}, nil)
catalog.On("ListIndexes",
mock.Anything,
).Return(nil, errors.New("error"))
_, err := newMeta(context.TODO(), catalog, nil)
assert.Error(t, err)
})
t.Run("ListSegmentIndexes fails", func(t *testing.T) {
catalog := mocks.NewDataCoordCatalog(t)
catalog.On("ListSegments",
mock.Anything,
).Return([]*datapb.SegmentInfo{}, nil)
catalog.On("ListChannelCheckpoint",
mock.Anything,
).Return(map[string]*internalpb.MsgPosition{}, nil)
catalog.On("ListIndexes",
mock.Anything,
).Return([]*model.Index{}, nil)
catalog.On("ListSegmentIndexes",
mock.Anything,
).Return(nil, errors.New("error"))
_, err := newMeta(context.TODO(), catalog, nil)
assert.Error(t, err)
})
t.Run("ok", func(t *testing.T) {
catalog := mocks.NewDataCoordCatalog(t)
catalog.On("ListSegments",
mock.Anything,
).Return([]*datapb.SegmentInfo{
{
ID: 1,
CollectionID: 1,
PartitionID: 1,
State: commonpb.SegmentState_Flushed,
},
}, nil)
catalog.On("ListChannelCheckpoint",
mock.Anything,
).Return(map[string]*internalpb.MsgPosition{
"ch": {
ChannelName: "cn",
MsgID: []byte{},
Timestamp: 1000,
},
}, nil)
catalog.On("ListIndexes",
mock.Anything,
).Return([]*model.Index{
{
CollectionID: 1,
IndexID: 1,
IndexName: "dix",
CreateTime: 1,
},
}, nil)
catalog.On("ListSegmentIndexes",
mock.Anything,
).Return([]*model.SegmentIndex{
{
SegmentID: 1,
IndexID: 1,
},
}, nil)
_, err := newMeta(context.TODO(), catalog, nil)
assert.NoError(t, err)
})
}
func TestMeta_Basic(t *testing.T) {
@ -454,16 +270,18 @@ func TestMeta_Basic(t *testing.T) {
t.Run("Test segment with kv fails", func(t *testing.T) {
// inject error for `Save`
memoryKV := memkv.NewMemoryKV()
fkv := &saveFailKV{TxnKV: memoryKV}
meta, err := newMeta(context.TODO(), fkv, "", nil)
memoryKV := NewMetaMemoryKV()
fkv := &saveFailKV{MetaKv: memoryKV}
catalog := datacoord.NewCatalog(fkv, "", "")
meta, err := newMeta(context.TODO(), catalog, nil)
assert.Nil(t, err)
err = meta.AddSegment(NewSegmentInfo(&datapb.SegmentInfo{}))
assert.NotNil(t, err)
fkv2 := &removeFailKV{TxnKV: memoryKV}
meta, err = newMeta(context.TODO(), fkv2, "", nil)
fkv2 := &removeFailKV{MetaKv: memoryKV}
catalog = datacoord.NewCatalog(fkv2, "", "")
meta, err = newMeta(context.TODO(), catalog, nil)
assert.Nil(t, err)
// nil, since no segment yet
err = meta.DropSegment(0)
@ -475,7 +293,8 @@ func TestMeta_Basic(t *testing.T) {
err = meta.DropSegment(0)
assert.NotNil(t, err)
meta, err = newMeta(context.TODO(), fkv, "", nil)
catalog = datacoord.NewCatalog(fkv, "", "")
meta, err = newMeta(context.TODO(), catalog, nil)
assert.Nil(t, err)
})
@ -609,7 +428,7 @@ func TestGetUnFlushedSegments(t *testing.T) {
func TestUpdateFlushSegmentsInfo(t *testing.T) {
t.Run("normal", func(t *testing.T) {
meta, err := newMeta(context.TODO(), memkv.NewMemoryKV(), "", nil)
meta, err := newMemoryMeta()
assert.Nil(t, err)
segment1 := &SegmentInfo{SegmentInfo: &datapb.SegmentInfo{ID: 1, State: commonpb.SegmentState_Growing, Binlogs: []*datapb.FieldBinlog{getFieldBinlogPaths(1, getInsertLogPath("binlog0", 1))},
@ -647,7 +466,7 @@ func TestUpdateFlushSegmentsInfo(t *testing.T) {
})
t.Run("update non-existed segment", func(t *testing.T) {
meta, err := newMeta(context.TODO(), memkv.NewMemoryKV(), "", nil)
meta, err := newMemoryMeta()
assert.Nil(t, err)
err = meta.UpdateFlushSegmentsInfo(1, false, false, false, nil, nil, nil, nil, nil)
@ -655,7 +474,7 @@ func TestUpdateFlushSegmentsInfo(t *testing.T) {
})
t.Run("update checkpoints and start position of non existed segment", func(t *testing.T) {
meta, err := newMeta(context.TODO(), memkv.NewMemoryKV(), "", nil)
meta, err := newMemoryMeta()
assert.Nil(t, err)
segment1 := &SegmentInfo{SegmentInfo: &datapb.SegmentInfo{ID: 1, State: commonpb.SegmentState_Growing}}
@ -670,9 +489,10 @@ func TestUpdateFlushSegmentsInfo(t *testing.T) {
})
t.Run("test save etcd failed", func(t *testing.T) {
kv := memkv.NewMemoryKV()
kv := NewMetaMemoryKV()
failedKv := &saveFailKV{kv}
meta, err := newMeta(context.TODO(), failedKv, "", nil)
catalog := datacoord.NewCatalog(failedKv, "", "")
meta, err := newMeta(context.TODO(), catalog, nil)
assert.Nil(t, err)
segmentInfo := &SegmentInfo{
@ -728,7 +548,7 @@ func TestMeta_alterMetaStore(t *testing.T) {
}
m := &meta{
catalog: &datacoord.Catalog{Txn: memkv.NewMemoryKV()},
catalog: &datacoord.Catalog{MetaKv: NewMetaMemoryKV()},
segments: &SegmentsInfo{map[int64]*SegmentInfo{
1: {SegmentInfo: &datapb.SegmentInfo{
ID: 1,
@ -777,7 +597,7 @@ func TestMeta_PrepareCompleteCompactionMutation(t *testing.T) {
}
m := &meta{
catalog: &datacoord.Catalog{Txn: memkv.NewMemoryKV()},
catalog: &datacoord.Catalog{MetaKv: NewMetaMemoryKV()},
segments: prepareSegments,
}
@ -838,7 +658,7 @@ func TestMeta_PrepareCompleteCompactionMutation(t *testing.T) {
func Test_meta_SetSegmentCompacting(t *testing.T) {
type fields struct {
client kv.TxnKV
client kv.MetaKv
segments *SegmentsInfo
}
type args struct {
@ -853,7 +673,7 @@ func Test_meta_SetSegmentCompacting(t *testing.T) {
{
"test set segment compacting",
fields{
memkv.NewMemoryKV(),
NewMetaMemoryKV(),
&SegmentsInfo{
map[int64]*SegmentInfo{
1: {
@ -875,7 +695,7 @@ func Test_meta_SetSegmentCompacting(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
m := &meta{
catalog: &datacoord.Catalog{Txn: tt.fields.client},
catalog: &datacoord.Catalog{MetaKv: tt.fields.client},
segments: tt.fields.segments,
}
m.SetSegmentCompacting(tt.args.segmentID, tt.args.compacting)
@ -887,7 +707,7 @@ func Test_meta_SetSegmentCompacting(t *testing.T) {
func Test_meta_SetSegmentImporting(t *testing.T) {
type fields struct {
client kv.TxnKV
client kv.MetaKv
segments *SegmentsInfo
}
type args struct {
@ -902,7 +722,7 @@ func Test_meta_SetSegmentImporting(t *testing.T) {
{
"test set segment importing",
fields{
memkv.NewMemoryKV(),
NewMetaMemoryKV(),
&SegmentsInfo{
map[int64]*SegmentInfo{
1: {
@ -924,7 +744,7 @@ func Test_meta_SetSegmentImporting(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
m := &meta{
catalog: &datacoord.Catalog{Txn: tt.fields.client},
catalog: &datacoord.Catalog{MetaKv: tt.fields.client},
segments: tt.fields.segments,
}
m.SetSegmentCompacting(tt.args.segmentID, tt.args.importing)
@ -1086,7 +906,7 @@ func TestChannelCP(t *testing.T) {
}
t.Run("UpdateChannelCheckpoint", func(t *testing.T) {
meta, err := newMeta(context.TODO(), memkv.NewMemoryKV(), "", nil)
meta, err := newMemoryMeta()
assert.NoError(t, err)
// nil position
@ -1098,7 +918,7 @@ func TestChannelCP(t *testing.T) {
})
t.Run("GetChannelCheckpoint", func(t *testing.T) {
meta, err := newMeta(context.TODO(), memkv.NewMemoryKV(), "", nil)
meta, err := newMemoryMeta()
assert.NoError(t, err)
position := meta.GetChannelCheckpoint(mockVChannel)
@ -1113,7 +933,7 @@ func TestChannelCP(t *testing.T) {
})
t.Run("DropChannelCheckpoint", func(t *testing.T) {
meta, err := newMeta(context.TODO(), memkv.NewMemoryKV(), "", nil)
meta, err := newMemoryMeta()
assert.NoError(t, err)
err = meta.DropChannelCheckpoint(mockVChannel)

85
internal/datacoord/mock_test.go
View File

@ -22,11 +22,14 @@ import (
"sync/atomic"
"time"
clientv3 "go.etcd.io/etcd/client/v3"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/schemapb"
"github.com/milvus-io/milvus/internal/kv"
memkv "github.com/milvus-io/milvus/internal/kv/mem"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/proto/proxypb"
@ -36,9 +39,83 @@ import (
"github.com/milvus-io/milvus/internal/util/typeutil"
)
type metaMemoryKV struct {
memkv.MemoryKV
}
func NewMetaMemoryKV() *metaMemoryKV {
return &metaMemoryKV{MemoryKV: *memkv.NewMemoryKV()}
}
func (mm *metaMemoryKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error {
keys, values, err := mm.MemoryKV.LoadWithPrefix(prefix)
if err != nil {
return err
}
for i, k := range keys {
if err := fn([]byte(k), []byte(values[i])); err != nil {
return err
}
}
return nil
}
func (mm *metaMemoryKV) GetPath(key string) string {
panic("implement me")
}
func (mm *metaMemoryKV) LoadWithPrefix2(key string) ([]string, []string, []int64, error) {
panic("implement me")
}
func (mm *metaMemoryKV) LoadWithRevisionAndVersions(key string) ([]string, []string, []int64, int64, error) {
panic("implement me")
}
func (mm *metaMemoryKV) LoadWithRevision(key string) ([]string, []string, int64, error) {
panic("implement me")
}
func (mm *metaMemoryKV) Watch(key string) clientv3.WatchChan {
panic("implement me")
}
func (mm *metaMemoryKV) WatchWithPrefix(key string) clientv3.WatchChan {
panic("implement me")
}
func (mm *metaMemoryKV) WatchWithRevision(key string, revision int64) clientv3.WatchChan {
panic("implement me")
}
func (mm *metaMemoryKV) SaveWithLease(key, value string, id clientv3.LeaseID) error {
panic("implement me")
}
func (mm *metaMemoryKV) SaveWithIgnoreLease(key, value string) error {
panic("implement me")
}
func (mm *metaMemoryKV) Grant(ttl int64) (id clientv3.LeaseID, err error) {
panic("implement me")
}
func (mm *metaMemoryKV) KeepAlive(id clientv3.LeaseID) (<-chan *clientv3.LeaseKeepAliveResponse, error) {
panic("implement me")
}
func (mm *metaMemoryKV) CompareValueAndSwap(key, value, target string, opts ...clientv3.OpOption) (bool, error) {
panic("implement me")
}
func (mm *metaMemoryKV) CompareVersionAndSwap(key string, version int64, target string, opts ...clientv3.OpOption) (bool, error) {
panic("implement me")
}
func newMemoryMeta() (*meta, error) {
memoryKV := memkv.NewMemoryKV()
return newMeta(context.TODO(), memoryKV, "", nil)
catalog := datacoord.NewCatalog(NewMetaMemoryKV(), "", "")
return newMeta(context.TODO(), catalog, nil)
}
var _ allocator = (*MockAllocator)(nil)
@ -91,7 +168,7 @@ func (a *FailsAllocator) allocID(_ context.Context) (UniqueID, error) {
}
// a mock kv that always fail when do `Save`
type saveFailKV struct{ kv.TxnKV }
type saveFailKV struct{ kv.MetaKv }
// Save override behavior
func (kv *saveFailKV) Save(key, value string) error {
@ -103,7 +180,7 @@ func (kv *saveFailKV) MultiSave(kvs map[string]string) error {
}
// a mock kv that always fail when do `Remove`
type removeFailKV struct{ kv.TxnKV }
type removeFailKV struct{ kv.MetaKv }
// Remove override behavior, inject error
func (kv *removeFailKV) MultiRemove(key []string) error {

21
internal/datacoord/segment_manager_test.go
View File

@ -27,7 +27,6 @@ import (
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/schemapb"
memkv "github.com/milvus-io/milvus/internal/kv/mem"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/util/metautil"
@ -577,10 +576,10 @@ func TestTryToSealSegment(t *testing.T) {
t.Run("seal with segment policy with kv fails", func(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
memoryKV := memkv.NewMemoryKV()
fkv := &saveFailKV{TxnKV: memoryKV}
meta, err := newMeta(context.TODO(), memoryKV, "", nil)
memoryKV := NewMetaMemoryKV()
fkv := &saveFailKV{MetaKv: memoryKV}
catalog := datacoord.NewCatalog(memoryKV, "", "")
meta, err := newMeta(context.TODO(), catalog, nil)
assert.Nil(t, err)
schema := newTestSchema()
@ -592,7 +591,7 @@ func TestTryToSealSegment(t *testing.T) {
assert.Nil(t, err)
assert.EqualValues(t, 1, len(allocations))
segmentManager.meta.catalog = &datacoord.Catalog{Txn: fkv}
segmentManager.meta.catalog = &datacoord.Catalog{MetaKv: fkv}
ts, err := segmentManager.allocator.allocTimestamp(context.Background())
assert.Nil(t, err)
@ -603,10 +602,10 @@ func TestTryToSealSegment(t *testing.T) {
t.Run("seal with channel policy with kv fails", func(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
memoryKV := memkv.NewMemoryKV()
fkv := &saveFailKV{TxnKV: memoryKV}
meta, err := newMeta(context.TODO(), memoryKV, "", nil)
memoryKV := NewMetaMemoryKV()
fkv := &saveFailKV{MetaKv: memoryKV}
catalog := datacoord.NewCatalog(memoryKV, "", "")
meta, err := newMeta(context.TODO(), catalog, nil)
assert.Nil(t, err)
schema := newTestSchema()
@ -618,7 +617,7 @@ func TestTryToSealSegment(t *testing.T) {
assert.Nil(t, err)
assert.EqualValues(t, 1, len(allocations))
segmentManager.meta.catalog = &datacoord.Catalog{Txn: fkv}
segmentManager.meta.catalog = &datacoord.Catalog{MetaKv: fkv}
ts, err := segmentManager.allocator.allocTimestamp(context.Background())
assert.Nil(t, err)

8
internal/datacoord/server.go
View File

@ -37,6 +37,7 @@ import (
rootcoordclient "github.com/milvus-io/milvus/internal/distributed/rootcoord/client"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
"github.com/milvus-io/milvus/internal/metrics"
"github.com/milvus-io/milvus/internal/mq/msgstream"
"github.com/milvus-io/milvus/internal/mq/msgstream/mqwrapper"
@ -268,7 +269,7 @@ func (s *Server) Init() error {
return err
}
if err = s.initMeta(storageCli.RootPath(), storageCli); err != nil {
if err = s.initMeta(storageCli); err != nil {
return err
}
@ -454,13 +455,14 @@ func (s *Server) initSegmentManager() {
}
}
func (s *Server) initMeta(chunkManagerRootPath string, chunkManager storage.ChunkManager) error {
func (s *Server) initMeta(chunkManager storage.ChunkManager) error {
etcdKV := etcdkv.NewEtcdKV(s.etcdCli, Params.EtcdCfg.MetaRootPath.GetValue())
s.kvClient = etcdKV
reloadEtcdFn := func() error {
var err error
s.meta, err = newMeta(s.ctx, s.kvClient, chunkManagerRootPath, chunkManager)
catalog := datacoord.NewCatalog(etcdKV, chunkManager.RootPath(), Params.EtcdCfg.MetaRootPath.GetValue())
s.meta, err = newMeta(s.ctx, catalog, chunkManager)
if err != nil {
return err
}

5
internal/datacoord/server_test.go
View File

@ -3234,6 +3234,8 @@ func TestDataCoord_Import(t *testing.T) {
t.Run("no datanode available", func(t *testing.T) {
svr := newTestServer(t, nil)
Params.BaseTable.Save("minio.address", "minio:9000")
defer Params.BaseTable.Reset("minio.address")
resp, err := svr.Import(svr.ctx, &datapb.ImportTaskRequest{
ImportTask: &datapb.ImportTask{
CollectionId: 100,
@ -3245,9 +3247,6 @@ func TestDataCoord_Import(t *testing.T) {
closeTestServer(t, svr)
})
// just passed for ci, if test locally, need to replace it with localhost:9000
Params.BaseTable.Save("minio.address", "minio:9000")
t.Run("with closed server", func(t *testing.T) {
svr := newTestServer(t, nil)
closeTestServer(t, svr)

1
internal/metastore/catalog.go
View File

@ -98,6 +98,7 @@ func (t AlterType) String() string {
return ""
}
//go:generate mockery --name=DataCoordCatalog --with-expecter
type DataCoordCatalog interface {
ListSegments(ctx context.Context) ([]*datapb.SegmentInfo, error)
AddSegment(ctx context.Context, segment *datapb.SegmentInfo) error

248
internal/metastore/kv/datacoord/kv_catalog.go
View File

@ -19,10 +19,15 @@ package datacoord
import (
"context"
"fmt"
"path"
"strconv"
"strings"
"github.com/golang/protobuf/proto"
"go.uber.org/zap"
"golang.org/x/exp/maps"
"golang.org/x/sync/errgroup"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus/internal/kv"
"github.com/milvus-io/milvus/internal/log"
@ -35,76 +40,201 @@ import (
"github.com/milvus-io/milvus/internal/util/metautil"
"github.com/milvus-io/milvus/internal/util/segmentutil"
"github.com/milvus-io/milvus/internal/util/typeutil"
"go.uber.org/zap"
"golang.org/x/exp/maps"
)
var maxEtcdTxnNum = 128
var paginationSize = 2000
type Catalog struct {
Txn kv.TxnKV
MetaKv kv.MetaKv
ChunkManagerRootPath string
metaRootpath string
}
func NewCatalog(MetaKv kv.MetaKv, chunkManagerRootPath string, metaRootpath string) *Catalog {
return &Catalog{MetaKv: MetaKv, ChunkManagerRootPath: chunkManagerRootPath, metaRootpath: metaRootpath}
}
func (kc *Catalog) ListSegments(ctx context.Context) ([]*datapb.SegmentInfo, error) {
_, values, err := kc.Txn.LoadWithPrefix(SegmentPrefix)
group, _ := errgroup.WithContext(ctx)
segments := make([]*datapb.SegmentInfo, 0)
insertLogs := make(map[typeutil.UniqueID][]*datapb.FieldBinlog, 1)
deltaLogs := make(map[typeutil.UniqueID][]*datapb.FieldBinlog, 1)
statsLogs := make(map[typeutil.UniqueID][]*datapb.FieldBinlog, 1)
executeFn := func(binlogType storage.BinlogType, result map[typeutil.UniqueID][]*datapb.FieldBinlog) {
group.Go(func() error {
ret, err := kc.listBinlogs(binlogType)
if err != nil {
return err
}
maps.Copy(result, ret)
return nil
})
}
//execute list segment meta
executeFn(storage.InsertBinlog, insertLogs)
executeFn(storage.DeleteBinlog, deltaLogs)
executeFn(storage.StatsBinlog, statsLogs)
group.Go(func() error {
ret, err := kc.listSegments()
if err != nil {
return err
}
segments = append(segments, ret...)
return nil
})
err := group.Wait()
if err != nil {
return nil, err
}
// get segment info
var segments []*datapb.SegmentInfo
for _, value := range values {
segmentInfo := &datapb.SegmentInfo{}
err = proto.Unmarshal([]byte(value), segmentInfo)
if err != nil {
log.Error("unmarshal segment info error", zap.Int64("segmentID", segmentInfo.ID), zap.Int64("collID", segmentInfo.CollectionID), zap.Error(err))
return nil, err
kc.applyBinlogInfo(segments, insertLogs, deltaLogs, statsLogs)
return segments, nil
}
func (kc *Catalog) listSegments() ([]*datapb.SegmentInfo, error) {
segments := make([]*datapb.SegmentInfo, 0)
applyFn := func(key []byte, value []byte) error {
// due to SegmentStatslogPathPrefix has the same prefix with SegmentPrefix, so skip it.
if strings.Contains(string(key), SegmentStatslogPathPrefix) {
return nil
}
segmentInfo := &datapb.SegmentInfo{}
err := proto.Unmarshal(value, segmentInfo)
if err != nil {
return err
}
segments = append(segments, segmentInfo)
return nil
}
for _, segmentInfo := range segments {
if len(segmentInfo.Binlogs) == 0 {
binlogs, err := kc.unmarshalBinlog(storage.InsertBinlog, segmentInfo.CollectionID, segmentInfo.PartitionID, segmentInfo.ID)
if err != nil {
return nil, err
}
segmentInfo.Binlogs = binlogs
}
if len(segmentInfo.Deltalogs) == 0 {
deltalogs, err := kc.unmarshalBinlog(storage.DeleteBinlog, segmentInfo.CollectionID, segmentInfo.PartitionID, segmentInfo.ID)
if err != nil {
return nil, err
}
segmentInfo.Deltalogs = deltalogs
}
if len(segmentInfo.Statslogs) == 0 {
statslogs, err := kc.unmarshalBinlog(storage.StatsBinlog, segmentInfo.CollectionID, segmentInfo.PartitionID, segmentInfo.ID)
if err != nil {
return nil, err
}
segmentInfo.Statslogs = statslogs
}
err := kc.MetaKv.WalkWithPrefix(SegmentPrefix+"/", paginationSize, applyFn)
if err != nil {
return nil, err
}
return segments, nil
}
func (kc *Catalog) parseBinlogKey(key string, prefixIdx int) (int64, int64, int64, error) {
remainedKey := key[prefixIdx:]
keyWordGroup := strings.Split(remainedKey, "/")
if len(keyWordGroup) < 3 {
return 0, 0, 0, fmt.Errorf("parse key: %s faild, trimed key:%s", key, remainedKey)
}
collectionID, err := strconv.ParseInt(keyWordGroup[0], 10, 64)
if err != nil {
return 0, 0, 0, fmt.Errorf("parse key: %s faild, trimed key:%s, %w", key, remainedKey, err)
}
partitionID, err := strconv.ParseInt(keyWordGroup[1], 10, 64)
if err != nil {
return 0, 0, 0, fmt.Errorf("parse key: %s faild, trimed key:%s, %w", key, remainedKey, err)
}
segmentID, err := strconv.ParseInt(keyWordGroup[2], 10, 64)
if err != nil {
return 0, 0, 0, fmt.Errorf("parse key: %s faild, trimed key:%s, %w", key, remainedKey, err)
}
return collectionID, partitionID, segmentID, nil
}
func (kc *Catalog) listBinlogs(binlogType storage.BinlogType) (map[typeutil.UniqueID][]*datapb.FieldBinlog, error) {
ret := make(map[typeutil.UniqueID][]*datapb.FieldBinlog)
var err error
var logPathPrefix string
switch binlogType {
case storage.InsertBinlog:
logPathPrefix = SegmentBinlogPathPrefix
case storage.DeleteBinlog:
logPathPrefix = SegmentDeltalogPathPrefix
case storage.StatsBinlog:
logPathPrefix = SegmentStatslogPathPrefix
default:
err = fmt.Errorf("invalid binlog type: %d", binlogType)
}
if err != nil {
return nil, err
}
var prefixIdx int
if len(kc.metaRootpath) == 0 {
prefixIdx = len(logPathPrefix) + 1
} else {
prefixIdx = len(kc.metaRootpath) + 1 + len(logPathPrefix) + 1
}
applyFn := func(key []byte, value []byte) error {
fieldBinlog := &datapb.FieldBinlog{}
err := proto.Unmarshal(value, fieldBinlog)
if err != nil {
return fmt.Errorf("failed to unmarshal datapb.FieldBinlog: %d, err:%w", fieldBinlog.FieldID, err)
}
collectionID, partitionID, segmentID, err := kc.parseBinlogKey(string(key), prefixIdx)
if err != nil {
return fmt.Errorf("prefix:%s, %w", path.Join(kc.metaRootpath, logPathPrefix), err)
}
switch binlogType {
case storage.InsertBinlog:
fillLogPathByLogID(kc.ChunkManagerRootPath, storage.InsertBinlog, collectionID, partitionID, segmentID, fieldBinlog)
case storage.DeleteBinlog:
fillLogPathByLogID(kc.ChunkManagerRootPath, storage.DeleteBinlog, collectionID, partitionID, segmentID, fieldBinlog)
case storage.StatsBinlog:
fillLogPathByLogID(kc.ChunkManagerRootPath, storage.StatsBinlog, collectionID, partitionID, segmentID, fieldBinlog)
}
ret[segmentID] = append(ret[segmentID], fieldBinlog)
return nil
}
err = kc.MetaKv.WalkWithPrefix(logPathPrefix, paginationSize, applyFn)
if err != nil {
return nil, err
}
return ret, nil
}
func (kc *Catalog) applyBinlogInfo(segments []*datapb.SegmentInfo, insertLogs, deltaLogs,
statsLogs map[typeutil.UniqueID][]*datapb.FieldBinlog) {
for _, segmentInfo := range segments {
if len(segmentInfo.Binlogs) == 0 {
segmentInfo.Binlogs = insertLogs[segmentInfo.ID]
}
if len(segmentInfo.Deltalogs) == 0 {
segmentInfo.Deltalogs = deltaLogs[segmentInfo.ID]
}
if len(segmentInfo.Statslogs) == 0 {
segmentInfo.Statslogs = statsLogs[segmentInfo.ID]
}
}
}
func (kc *Catalog) AddSegment(ctx context.Context, segment *datapb.SegmentInfo) error {
kvs, err := buildSegmentAndBinlogsKvs(segment)
if err != nil {
return err
}
return kc.Txn.MultiSave(kvs)
return kc.MetaKv.MultiSave(kvs)
}
// LoadFromSegmentPath loads segment info from persistent storage by given segment path.
// # TESTING ONLY #
func (kc *Catalog) LoadFromSegmentPath(colID, partID, segID typeutil.UniqueID) (*datapb.SegmentInfo, error) {
v, err := kc.Txn.Load(buildSegmentPath(colID, partID, segID))
v, err := kc.MetaKv.Load(buildSegmentPath(colID, partID, segID))
if err != nil {
log.Error("(testing only) failed to load segment info by segment path")
return nil, err
@ -138,7 +268,7 @@ func (kc *Catalog) AlterSegments(ctx context.Context, newSegments []*datapb.Segm
kvsPiece := make(map[string]string)
currSize := 0
saveFn := func(partialKvs map[string]string) error {
return kc.Txn.MultiSave(partialKvs)
return kc.MetaKv.MultiSave(partialKvs)
}
for _, kvs := range kvsBySeg {
if currSize+len(kvs) >= maxEtcdTxnNum {
@ -173,7 +303,7 @@ func (kc *Catalog) AlterSegment(ctx context.Context, newSegment *datapb.SegmentI
}
maps.Copy(kvs, segmentKvs)
return kc.Txn.MultiSave(kvs)
return kc.MetaKv.MultiSave(kvs)
}
func (kc *Catalog) hasBinlogPrefix(segment *datapb.SegmentInfo) (bool, error) {
@ -241,7 +371,7 @@ func (kc *Catalog) AlterSegmentsAndAddNewSegment(ctx context.Context, segments [
maps.Copy(kvs, segmentKvs)
}
}
return kc.Txn.MultiSave(kvs)
return kc.MetaKv.MultiSave(kvs)
}
// RevertAlterSegmentsAndAddNewSegment reverts the metastore operation of AlterSegmentsAndAddNewSegment
@ -267,7 +397,7 @@ func (kc *Catalog) RevertAlterSegmentsAndAddNewSegment(ctx context.Context, oldS
removals = append(removals, binlogKeys...)
}
err := kc.Txn.MultiSaveAndRemove(kvs, removals)
err := kc.MetaKv.MultiSaveAndRemove(kvs, removals)
if err != nil {
log.Warn("batch save and remove segments failed", zap.Error(err))
return err
@ -295,7 +425,7 @@ func (kc *Catalog) SaveDroppedSegmentsInBatch(ctx context.Context, segments []*d
}
saveFn := func(partialKvs map[string]string) error {
return kc.Txn.MultiSave(partialKvs)
return kc.MetaKv.MultiSave(partialKvs)
}
if err := etcd.SaveByBatch(kvs, saveFn); err != nil {
return err
@ -309,12 +439,12 @@ func (kc *Catalog) DropSegment(ctx context.Context, segment *datapb.SegmentInfo)
keys := []string{segKey}
binlogKeys := buildBinlogKeys(segment)
keys = append(keys, binlogKeys...)
return kc.Txn.MultiRemove(keys)
return kc.MetaKv.MultiRemove(keys)
}
func (kc *Catalog) MarkChannelDeleted(ctx context.Context, channel string) error {
key := buildChannelRemovePath(channel)
err := kc.Txn.Save(key, RemoveFlagTomestone)
err := kc.MetaKv.Save(key, RemoveFlagTomestone)
if err != nil {
log.Error("Failed to mark channel dropped", zap.String("channel", channel), zap.Error(err))
return err
@ -325,7 +455,7 @@ func (kc *Catalog) MarkChannelDeleted(ctx context.Context, channel string) error
func (kc *Catalog) IsChannelDropped(ctx context.Context, channel string) bool {
key := buildChannelRemovePath(channel)
v, err := kc.Txn.Load(key)
v, err := kc.MetaKv.Load(key)
if err != nil || v != RemoveFlagTomestone {
return false
}
@ -335,11 +465,11 @@ func (kc *Catalog) IsChannelDropped(ctx context.Context, channel string) bool {
// DropChannel removes channel remove flag after whole procedure is finished
func (kc *Catalog) DropChannel(ctx context.Context, channel string) error {
key := buildChannelRemovePath(channel)
return kc.Txn.Remove(key)
return kc.MetaKv.Remove(key)
}
func (kc *Catalog) ListChannelCheckpoint(ctx context.Context) (map[string]*internalpb.MsgPosition, error) {
keys, values, err := kc.Txn.LoadWithPrefix(ChannelCheckpointPrefix)
keys, values, err := kc.MetaKv.LoadWithPrefix(ChannelCheckpointPrefix)
if err != nil {
return nil, err
}
@ -367,12 +497,12 @@ func (kc *Catalog) SaveChannelCheckpoint(ctx context.Context, vChannel string, p
if err != nil {
return err
}
return kc.Txn.Save(k, string(v))
return kc.MetaKv.Save(k, string(v))
}
func (kc *Catalog) DropChannelCheckpoint(ctx context.Context, vChannel string) error {
k := buildChannelCPKey(vChannel)
return kc.Txn.Remove(k)
return kc.MetaKv.Remove(k)
}
func (kc *Catalog) getBinlogsWithPrefix(binlogType storage.BinlogType, collectionID, partitionID,
@ -388,7 +518,7 @@ func (kc *Catalog) getBinlogsWithPrefix(binlogType storage.BinlogType, collectio
default:
return nil, nil, fmt.Errorf("invalid binlog type: %d", binlogType)
}
keys, values, err := kc.Txn.LoadWithPrefix(binlogPrefix)
keys, values, err := kc.MetaKv.LoadWithPrefix(binlogPrefix)
if err != nil {
return nil, nil, err
}
@ -425,7 +555,7 @@ func (kc *Catalog) CreateIndex(ctx context.Context, index *model.Index) error {
return err
}
err = kc.Txn.Save(key, string(value))
err = kc.MetaKv.Save(key, string(value))
if err != nil {
return err
}
@ -433,7 +563,7 @@ func (kc *Catalog) CreateIndex(ctx context.Context, index *model.Index) error {
}
func (kc *Catalog) ListIndexes(ctx context.Context) ([]*model.Index, error) {
_, values, err := kc.Txn.LoadWithPrefix(util.FieldIndexPrefix)
_, values, err := kc.MetaKv.LoadWithPrefix(util.FieldIndexPrefix)
if err != nil {
log.Error("list index meta fail", zap.String("prefix", util.FieldIndexPrefix), zap.Error(err))
return nil, err
@ -470,13 +600,13 @@ func (kc *Catalog) AlterIndexes(ctx context.Context, indexes []*model.Index) err
kvs[key] = string(value)
}
return kc.Txn.MultiSave(kvs)
return kc.MetaKv.MultiSave(kvs)
}
func (kc *Catalog) DropIndex(ctx context.Context, collID typeutil.UniqueID, dropIdxID typeutil.UniqueID) error {
key := BuildIndexKey(collID, dropIdxID)
err := kc.Txn.Remove(key)
err := kc.MetaKv.Remove(key)
if err != nil {
log.Error("drop collection index meta fail", zap.Int64("collectionID", collID),
zap.Int64("indexID", dropIdxID), zap.Error(err))
@ -493,7 +623,7 @@ func (kc *Catalog) CreateSegmentIndex(ctx context.Context, segIdx *model.Segment
if err != nil {
return err
}
err = kc.Txn.Save(key, string(value))
err = kc.MetaKv.Save(key, string(value))
if err != nil {
log.Error("failed to save segment index meta in etcd", zap.Int64("buildID", segIdx.BuildID),
zap.Int64("segmentID", segIdx.SegmentID), zap.Error(err))
@ -503,7 +633,7 @@ func (kc *Catalog) CreateSegmentIndex(ctx context.Context, segIdx *model.Segment
}
func (kc *Catalog) ListSegmentIndexes(ctx context.Context) ([]*model.SegmentIndex, error) {
_, values, err := kc.Txn.LoadWithPrefix(util.SegmentIndexPrefix)
_, values, err := kc.MetaKv.LoadWithPrefix(util.SegmentIndexPrefix)
if err != nil {
log.Error("list segment index meta fail", zap.String("prefix", util.SegmentIndexPrefix), zap.Error(err))
return nil, err
@ -538,13 +668,13 @@ func (kc *Catalog) AlterSegmentIndexes(ctx context.Context, segIdxes []*model.Se
}
kvs[key] = string(value)
}
return kc.Txn.MultiSave(kvs)
return kc.MetaKv.MultiSave(kvs)
}
func (kc *Catalog) DropSegmentIndex(ctx context.Context, collID, partID, segID, buildID typeutil.UniqueID) error {
key := BuildSegmentIndexKey(collID, partID, segID, buildID)
err := kc.Txn.Remove(key)
err := kc.MetaKv.Remove(key)
if err != nil {
log.Error("drop segment index meta fail", zap.Int64("buildID", buildID), zap.Error(err))
return err

332
internal/metastore/kv/datacoord/kv_catalog_test.go
View File

@ -19,8 +19,13 @@ package datacoord
import (
"context"
"errors"
"fmt"
"log"
"math/rand"
"os"
"strings"
"testing"
"time"
"github.com/golang/protobuf/proto"
"github.com/stretchr/testify/assert"
@ -29,20 +34,24 @@ import (
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus/internal/kv"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/kv/mocks"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/util/etcd"
"github.com/milvus-io/milvus/internal/util/metautil"
"github.com/milvus-io/milvus/internal/util/paramtable"
)
type MockedTxnKV struct {
kv.TxnKV
kv.MetaKv
multiSave func(kvs map[string]string) error
save func(key, value string) error
loadWithPrefix func(key string) ([]string, []string, error)
load func(key string) (string, error)
multiRemove func(keys []string) error
walkWithPrefix func(prefix string, paginationSize int, fn func([]byte, []byte) error) error
remove func(key string) error
}
@ -50,17 +59,18 @@ var (
logID = int64(99)
collectionID = int64(2)
partitionID = int64(1)
segmentID = int64(11)
segmentID2 = int64(1)
segmentID = int64(1)
segmentID2 = int64(11)
fieldID = int64(1)
rootPath = "a"
binlogPath = metautil.BuildInsertLogPath("a", collectionID, partitionID, segmentID, fieldID, logID)
deltalogPath = metautil.BuildDeltaLogPath("a", collectionID, partitionID, segmentID, logID)
statslogPath = metautil.BuildStatsLogPath("a", collectionID, partitionID, segmentID, fieldID, logID)
binlogPath = metautil.BuildInsertLogPath(rootPath, collectionID, partitionID, segmentID, fieldID, logID)
deltalogPath = metautil.BuildDeltaLogPath(rootPath, collectionID, partitionID, segmentID, logID)
statslogPath = metautil.BuildStatsLogPath(rootPath, collectionID, partitionID, segmentID, fieldID, logID)
binlogPath2 = metautil.BuildInsertLogPath("a", collectionID, partitionID, segmentID2, fieldID, logID)
deltalogPath2 = metautil.BuildDeltaLogPath("a", collectionID, partitionID, segmentID2, logID)
statslogPath2 = metautil.BuildStatsLogPath("a", collectionID, partitionID, segmentID2, fieldID, logID)
binlogPath2 = metautil.BuildInsertLogPath(rootPath, collectionID, partitionID, segmentID2, fieldID, logID)
deltalogPath2 = metautil.BuildDeltaLogPath(rootPath, collectionID, partitionID, segmentID2, logID)
statslogPath2 = metautil.BuildStatsLogPath(rootPath, collectionID, partitionID, segmentID2, fieldID, logID)
k1 = buildFieldBinlogPath(collectionID, partitionID, segmentID, fieldID)
k2 = buildFieldDeltalogPath(collectionID, partitionID, segmentID, fieldID)
@ -194,6 +204,10 @@ func (mc *MockedTxnKV) Load(key string) (string, error) {
return mc.load(key)
}
func (mc *MockedTxnKV) WalkWithPrefix(prefix string, paginationSize int, fn func([]byte, []byte) error) error {
return mc.walkWithPrefix(prefix, paginationSize, fn)
}
func (mc *MockedTxnKV) Remove(key string) error {
return mc.remove(key)
}
@ -201,11 +215,11 @@ func (mc *MockedTxnKV) Remove(key string) error {
func Test_ListSegments(t *testing.T) {
t.Run("load failed", func(t *testing.T) {
txn := &MockedTxnKV{}
txn.loadWithPrefix = func(key string) ([]string, []string, error) {
return nil, nil, errors.New("error")
txn.walkWithPrefix = func(prefix string, paginationSize int, fn func([]byte, []byte) error) error {
return errors.New("error")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
ret, err := catalog.ListSegments(context.TODO())
assert.Nil(t, ret)
assert.Error(t, err)
@ -242,11 +256,14 @@ func Test_ListSegments(t *testing.T) {
segBytes, err := proto.Marshal(segment1)
assert.NoError(t, err)
txn.loadWithPrefix = func(key string) ([]string, []string, error) {
return []string{k5}, []string{string(segBytes)}, nil
txn.walkWithPrefix = func(prefix string, paginationSize int, fn func([]byte, []byte) error) error {
if strings.HasPrefix(k5, prefix) {
return fn([]byte(k5), segBytes)
}
return nil
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
ret, err := catalog.ListSegments(context.TODO())
assert.NotNil(t, ret)
assert.NoError(t, err)
@ -262,28 +279,27 @@ func Test_ListSegments(t *testing.T) {
return nil
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.AddSegment(context.TODO(), segment1)
assert.Nil(t, err)
txn.loadWithPrefix = func(key string) ([]string, []string, error) {
if strings.HasPrefix(k5, key) {
return []string{k5}, []string{savedKvs[k5]}, nil
txn.walkWithPrefix = func(prefix string, paginationSize int, fn func(k []byte, v []byte) error) error {
if strings.HasPrefix(k5, prefix) {
return fn([]byte(k5), []byte(savedKvs[k5]))
}
if strings.HasPrefix(k1, key) {
return []string{k1}, []string{savedKvs[k1]}, nil
if strings.HasPrefix(k1, prefix) {
return fn([]byte(k1), []byte(savedKvs[k1]))
}
if strings.HasPrefix(k2, key) {
return []string{k2}, []string{savedKvs[k2]}, nil
if strings.HasPrefix(k2, prefix) {
return fn([]byte(k2), []byte(savedKvs[k2]))
}
if strings.HasPrefix(k3, prefix) {
return fn([]byte(k3), []byte(savedKvs[k3]))
}
if strings.HasPrefix(k3, key) {
return []string{k3}, []string{savedKvs[k3]}, nil
}
return nil, nil, errors.New("should not reach here")
return errors.New("should not reach here")
}
ret, err := catalog.ListSegments(context.TODO())
@ -301,7 +317,7 @@ func Test_AddSegments(t *testing.T) {
return errors.New("error")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
assert.Panics(t, func() {
catalog.AddSegment(context.TODO(), invalidSegment)
})
@ -313,7 +329,7 @@ func Test_AddSegments(t *testing.T) {
return errors.New("error")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.AddSegment(context.TODO(), segment1)
assert.Error(t, err)
})
@ -332,7 +348,7 @@ func Test_AddSegments(t *testing.T) {
return "", errors.New("key not found")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.AddSegment(context.TODO(), segment1)
assert.Nil(t, err)
adjustedSeg, err := catalog.LoadFromSegmentPath(segment1.CollectionID, segment1.PartitionID, segment1.ID)
@ -355,7 +371,7 @@ func Test_AlterSegments(t *testing.T) {
return errors.New("error")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
assert.Panics(t, func() {
catalog.AlterSegments(context.TODO(), []*datapb.SegmentInfo{invalidSegment})
})
@ -367,7 +383,7 @@ func Test_AlterSegments(t *testing.T) {
return errors.New("error")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.AlterSegments(context.TODO(), []*datapb.SegmentInfo{segment1})
assert.Error(t, err)
@ -381,7 +397,7 @@ func Test_AlterSegments(t *testing.T) {
return nil
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.AlterSegments(context.TODO(), []*datapb.SegmentInfo{})
assert.Nil(t, err)
@ -414,7 +430,7 @@ func Test_AlterSegments(t *testing.T) {
return "", errors.New("key not found")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.AlterSegments(context.TODO(), []*datapb.SegmentInfo{})
assert.Nil(t, err)
@ -462,7 +478,7 @@ func Test_AlterSegmentsAndAddNewSegment(t *testing.T) {
return errors.New("error")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.AlterSegmentsAndAddNewSegment(context.TODO(), []*datapb.SegmentInfo{}, segment1)
assert.Error(t, err)
})
@ -473,7 +489,7 @@ func Test_AlterSegmentsAndAddNewSegment(t *testing.T) {
return nil, nil, errors.New("error")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.AlterSegmentsAndAddNewSegment(context.TODO(), []*datapb.SegmentInfo{droppedSegment}, nil)
assert.Error(t, err)
})
@ -495,7 +511,7 @@ func Test_AlterSegmentsAndAddNewSegment(t *testing.T) {
return "", errors.New("key not found")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.AlterSegmentsAndAddNewSegment(context.TODO(), []*datapb.SegmentInfo{droppedSegment}, segment1)
assert.NoError(t, err)
@ -524,7 +540,7 @@ func Test_DropSegment(t *testing.T) {
return errors.New("error")
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.DropSegment(context.TODO(), segment1)
assert.Error(t, err)
})
@ -539,7 +555,7 @@ func Test_DropSegment(t *testing.T) {
return nil
}
catalog := &Catalog{txn, "a"}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.DropSegment(context.TODO(), segment1)
assert.NoError(t, err)
@ -552,10 +568,10 @@ func Test_DropSegment(t *testing.T) {
}
func Test_SaveDroppedSegmentsInBatch_SaveError(t *testing.T) {
txn := &mocks.TxnKV{}
txn := mocks.NewMetaKv(t)
txn.EXPECT().MultiSave(mock.Anything).Return(errors.New("mock error"))
catalog := &Catalog{txn, ""}
catalog := NewCatalog(txn, rootPath, "")
segments := []*datapb.SegmentInfo{
{
ID: 1,
@ -571,7 +587,7 @@ func Test_SaveDroppedSegmentsInBatch_MultiSave(t *testing.T) {
count = 0
kvSize = 0
)
txn := &mocks.TxnKV{}
txn := mocks.NewMetaKv(t)
txn.EXPECT().
MultiSave(mock.Anything).
Run(func(kvs map[string]string) {
@ -580,7 +596,7 @@ func Test_SaveDroppedSegmentsInBatch_MultiSave(t *testing.T) {
}).
Return(nil)
catalog := &Catalog{txn, ""}
catalog := NewCatalog(txn, rootPath, "")
// no segments
{
@ -629,18 +645,18 @@ func Test_SaveDroppedSegmentsInBatch_MultiSave(t *testing.T) {
func TestCatalog_RevertAlterSegmentsAndAddNewSegment(t *testing.T) {
t.Run("save error", func(t *testing.T) {
txn := &mocks.TxnKV{}
txn := mocks.NewMetaKv(t)
txn.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything).Return(errors.New("mock error"))
catalog := &Catalog{txn, ""}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.RevertAlterSegmentsAndAddNewSegment(context.TODO(), []*datapb.SegmentInfo{segment1}, droppedSegment)
assert.Error(t, err)
})
t.Run("revert successfully", func(t *testing.T) {
txn := &mocks.TxnKV{}
txn := mocks.NewMetaKv(t)
txn.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything).Return(nil)
catalog := &Catalog{txn, ""}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.RevertAlterSegmentsAndAddNewSegment(context.TODO(), []*datapb.SegmentInfo{segment1}, droppedSegment)
assert.NoError(t, err)
})
@ -660,9 +676,9 @@ func TestChannelCP(t *testing.T) {
assert.NoError(t, err)
t.Run("ListChannelCheckpoint", func(t *testing.T) {
txn := &mocks.TxnKV{}
txn := mocks.NewMetaKv(t)
txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
catalog := &Catalog{txn, ""}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, pos)
assert.NoError(t, err)
@ -673,33 +689,33 @@ func TestChannelCP(t *testing.T) {
})
t.Run("ListChannelCheckpoint failed", func(t *testing.T) {
txn := &mocks.TxnKV{}
catalog := &Catalog{txn, ""}
txn := mocks.NewMetaKv(t)
catalog := NewCatalog(txn, rootPath, "")
txn.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, errors.New("mock error"))
_, err = catalog.ListChannelCheckpoint(context.TODO())
assert.Error(t, err)
})
t.Run("SaveChannelCheckpoint", func(t *testing.T) {
txn := &mocks.TxnKV{}
txn := mocks.NewMetaKv(t)
txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
catalog := &Catalog{txn, ""}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, pos)
assert.NoError(t, err)
})
t.Run("SaveChannelCheckpoint failed", func(t *testing.T) {
txn := &mocks.TxnKV{}
catalog := &Catalog{txn, ""}
txn := mocks.NewMetaKv(t)
catalog := NewCatalog(txn, rootPath, "")
txn.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("mock error"))
err = catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, &internalpb.MsgPosition{})
assert.Error(t, err)
})
t.Run("DropChannelCheckpoint", func(t *testing.T) {
txn := &mocks.TxnKV{}
txn := mocks.NewMetaKv(t)
txn.EXPECT().Save(mock.Anything, mock.Anything).Return(nil)
catalog := &Catalog{txn, ""}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.SaveChannelCheckpoint(context.TODO(), mockVChannel, pos)
assert.NoError(t, err)
@ -713,8 +729,8 @@ func TestChannelCP(t *testing.T) {
})
t.Run("DropChannelCheckpoint failed", func(t *testing.T) {
txn := &mocks.TxnKV{}
catalog := &Catalog{txn, ""}
txn := mocks.NewMetaKv(t)
catalog := NewCatalog(txn, rootPath, "")
txn.EXPECT().Remove(mock.Anything).Return(errors.New("mock error"))
err = catalog.DropChannelCheckpoint(context.TODO(), mockVChannel)
assert.Error(t, err)
@ -722,16 +738,60 @@ func TestChannelCP(t *testing.T) {
}
func Test_MarkChannelDeleted_SaveError(t *testing.T) {
txn := &mocks.TxnKV{}
txn := mocks.NewMetaKv(t)
txn.EXPECT().
Save(mock.Anything, mock.Anything).
Return(errors.New("mock error"))
catalog := &Catalog{txn, ""}
catalog := NewCatalog(txn, rootPath, "")
err := catalog.MarkChannelDeleted(context.TODO(), "test_channel_1")
assert.Error(t, err)
}
func Test_parseBinlogKey(t *testing.T) {
catalog := NewCatalog(nil, "", "")
t.Run("parse collection id fail", func(t *testing.T) {
ret1, ret2, ret3, err := catalog.parseBinlogKey("root/err/1/1/1", 5)
assert.Error(t, err)
assert.Equal(t, int64(0), ret1)
assert.Equal(t, int64(0), ret2)
assert.Equal(t, int64(0), ret3)
})
t.Run("parse partition id fail", func(t *testing.T) {
ret1, ret2, ret3, err := catalog.parseBinlogKey("root/1/err/1/1", 5)
assert.Error(t, err)
assert.Equal(t, int64(0), ret1)
assert.Equal(t, int64(0), ret2)
assert.Equal(t, int64(0), ret3)
})
t.Run("parse segment id fail", func(t *testing.T) {
ret1, ret2, ret3, err := catalog.parseBinlogKey("root/1/1/err/1", 5)
assert.Error(t, err)
assert.Equal(t, int64(0), ret1)
assert.Equal(t, int64(0), ret2)
assert.Equal(t, int64(0), ret3)
})
t.Run("miss field", func(t *testing.T) {
ret1, ret2, ret3, err := catalog.parseBinlogKey("root/1/1/", 5)
assert.Error(t, err)
assert.Equal(t, int64(0), ret1)
assert.Equal(t, int64(0), ret2)
assert.Equal(t, int64(0), ret3)
})
t.Run("test ok", func(t *testing.T) {
ret1, ret2, ret3, err := catalog.parseBinlogKey("root/1/1/1/1", 5)
assert.NoError(t, err)
assert.Equal(t, int64(1), ret1)
assert.Equal(t, int64(1), ret2)
assert.Equal(t, int64(1), ret3)
})
}
func verifyBinlogs(t *testing.T, binlogBytes []byte) {
binlogs := &datapb.FieldBinlog{}
err := proto.Unmarshal([]byte(binlogBytes), binlogs)
@ -806,7 +866,7 @@ func TestCatalog_CreateIndex(t *testing.T) {
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.CreateIndex(context.Background(), &model.Index{})
@ -821,7 +881,7 @@ func TestCatalog_CreateIndex(t *testing.T) {
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.CreateIndex(context.Background(), &model.Index{})
@ -851,7 +911,7 @@ func TestCatalog_ListIndexes(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
indexes, err := catalog.ListIndexes(context.Background())
assert.NoError(t, err)
@ -865,7 +925,7 @@ func TestCatalog_ListIndexes(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
_, err := catalog.ListIndexes(context.Background())
assert.Error(t, err)
@ -878,7 +938,7 @@ func TestCatalog_ListIndexes(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
_, err := catalog.ListIndexes(context.Background())
assert.Error(t, err)
@ -903,7 +963,7 @@ func TestCatalog_AlterIndex(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.AlterIndex(context.Background(), i)
@ -929,7 +989,7 @@ func TestCatalog_AlterIndexes(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.AlterIndexes(context.Background(), []*model.Index{i})
@ -944,7 +1004,7 @@ func TestCatalog_DropIndex(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.DropIndex(context.Background(), 0, 0)
@ -958,7 +1018,7 @@ func TestCatalog_DropIndex(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.DropIndex(context.Background(), 0, 0)
@ -991,7 +1051,7 @@ func TestCatalog_CreateSegmentIndex(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.CreateSegmentIndex(context.Background(), segIdx)
@ -1005,7 +1065,7 @@ func TestCatalog_CreateSegmentIndex(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.CreateSegmentIndex(context.Background(), segIdx)
@ -1040,7 +1100,7 @@ func TestCatalog_ListSegmentIndexes(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
segIdxes, err := catalog.ListSegmentIndexes(context.Background())
@ -1055,7 +1115,7 @@ func TestCatalog_ListSegmentIndexes(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
_, err := catalog.ListSegmentIndexes(context.Background())
@ -1069,7 +1129,7 @@ func TestCatalog_ListSegmentIndexes(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
_, err := catalog.ListSegmentIndexes(context.Background())
@ -1102,7 +1162,7 @@ func TestCatalog_AlterSegmentIndex(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.AlterSegmentIndex(context.Background(), segIdx)
@ -1135,7 +1195,7 @@ func TestCatalog_AlterSegmentIndexes(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.AlterSegmentIndexes(context.Background(), []*model.SegmentIndex{segIdx})
@ -1151,7 +1211,7 @@ func TestCatalog_DropSegmentIndex(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.DropSegmentIndex(context.Background(), 0, 0, 0, 0)
@ -1165,10 +1225,122 @@ func TestCatalog_DropSegmentIndex(t *testing.T) {
},
}
catalog := &Catalog{
Txn: txn,
MetaKv: txn,
}
err := catalog.DropSegmentIndex(context.Background(), 0, 0, 0, 0)
assert.Error(t, err)
})
}
func BenchmarkCatalog_List1000Segments(b *testing.B) {
Params.Init()
etcdCli, err := etcd.GetEtcdClient(
Params.EtcdCfg.UseEmbedEtcd.GetAsBool(),
Params.EtcdCfg.EtcdUseSSL.GetAsBool(),
Params.EtcdCfg.Endpoints.GetAsStrings(),
Params.EtcdCfg.EtcdTLSCert.GetValue(),
Params.EtcdCfg.EtcdTLSKey.GetValue(),
Params.EtcdCfg.EtcdTLSCACert.GetValue(),
Params.EtcdCfg.EtcdTLSMinVersion.GetValue())
if err != nil {
log.Fatal(err)
}
defer etcdCli.Close()
randVal := rand.Int()
dataRootPath := fmt.Sprintf("/test/data/list-segment-%d", randVal)
etcdkv := etcdkv.NewEtcdKV(etcdCli, rootPath)
defer etcdkv.Close()
ctx := context.TODO()
catalog := NewCatalog(etcdkv, dataRootPath, rootPath)
generateSegments(ctx, catalog, 10, rootPath)
b.ResetTimer()
for i := 0; i < b.N; i++ {
segments, err := catalog.ListSegments(ctx)
assert.NoError(b, err)
for _, s := range segments {
assert.NotNil(b, s)
assert.NotNil(b, s.Binlogs)
assert.NotNil(b, s.Statslogs)
assert.NotNil(b, s.Deltalogs)
}
}
}
func generateSegments(ctx context.Context, catalog *Catalog, n int, rootPath string) {
rand.Seed(time.Now().UnixNano())
var collectionID int64
for i := 0; i < n; i++ {
if collectionID%25 == 0 {
collectionID = rand.Int63()
}
v := rand.Int63()
segment := addSegment(rootPath, collectionID, v, v, v)
err := catalog.AddSegment(ctx, segment)
if err != nil {
log.Fatal(err)
}
}
}
func addSegment(rootPath string, collectionID, partitionID, segmentID, fieldID int64) *datapb.SegmentInfo {
binlogs = []*datapb.FieldBinlog{
{
FieldID: fieldID,
Binlogs: []*datapb.Binlog{
{
EntriesNum: 10000,
LogPath: metautil.BuildInsertLogPath(rootPath, collectionID, partitionID, segmentID, fieldID, int64(rand.Int())),
},
},
},
}
deltalogs = []*datapb.FieldBinlog{
{
FieldID: fieldID,
Binlogs: []*datapb.Binlog{
{
EntriesNum: 5,
LogPath: metautil.BuildDeltaLogPath(rootPath, collectionID, partitionID, segmentID, int64(rand.Int())),
}},
},
}
statslogs = []*datapb.FieldBinlog{
{
FieldID: 1,
Binlogs: []*datapb.Binlog{
{
EntriesNum: 5,
LogPath: metautil.BuildStatsLogPath(rootPath, collectionID, partitionID, segmentID, fieldID, int64(rand.Int())),
},
},
},
}
return &datapb.SegmentInfo{
ID: segmentID,
CollectionID: collectionID,
PartitionID: partitionID,
NumOfRows: 10000,
State: commonpb.SegmentState_Flushed,
Binlogs: binlogs,
Deltalogs: deltalogs,
Statslogs: statslogs,
}
}
var Params = paramtable.Get()
func TestMain(m *testing.M) {
Params.Init()
code := m.Run()
os.Exit(code)
}

994
internal/metastore/mocks/DataCoordCatalog.go Normal file
View File

@ -0,0 +1,994 @@
// Code generated by mockery v2.16.0. DO NOT EDIT.
package mocks
import (
context "context"
datapb "github.com/milvus-io/milvus/internal/proto/datapb"
internalpb "github.com/milvus-io/milvus/internal/proto/internalpb"
mock "github.com/stretchr/testify/mock"
model "github.com/milvus-io/milvus/internal/metastore/model"
)
// DataCoordCatalog is an autogenerated mock type for the DataCoordCatalog type
type DataCoordCatalog struct {
mock.Mock
}
type DataCoordCatalog_Expecter struct {
mock *mock.Mock
}
func (_m *DataCoordCatalog) EXPECT() *DataCoordCatalog_Expecter {
return &DataCoordCatalog_Expecter{mock: &_m.Mock}
}
// AddSegment provides a mock function with given fields: ctx, segment
func (_m *DataCoordCatalog) AddSegment(ctx context.Context, segment *datapb.SegmentInfo) error {
ret := _m.Called(ctx, segment)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *datapb.SegmentInfo) error); ok {
r0 = rf(ctx, segment)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_AddSegment_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AddSegment'
type DataCoordCatalog_AddSegment_Call struct {
*mock.Call
}
// AddSegment is a helper method to define mock.On call
// - ctx context.Context
// - segment *datapb.SegmentInfo
func (_e *DataCoordCatalog_Expecter) AddSegment(ctx interface{}, segment interface{}) *DataCoordCatalog_AddSegment_Call {
return &DataCoordCatalog_AddSegment_Call{Call: _e.mock.On("AddSegment", ctx, segment)}
}
func (_c *DataCoordCatalog_AddSegment_Call) Run(run func(ctx context.Context, segment *datapb.SegmentInfo)) *DataCoordCatalog_AddSegment_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*datapb.SegmentInfo))
})
return _c
}
func (_c *DataCoordCatalog_AddSegment_Call) Return(_a0 error) *DataCoordCatalog_AddSegment_Call {
_c.Call.Return(_a0)
return _c
}
// AlterIndex provides a mock function with given fields: ctx, newIndex
func (_m *DataCoordCatalog) AlterIndex(ctx context.Context, newIndex *model.Index) error {
ret := _m.Called(ctx, newIndex)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *model.Index) error); ok {
r0 = rf(ctx, newIndex)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_AlterIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AlterIndex'
type DataCoordCatalog_AlterIndex_Call struct {
*mock.Call
}
// AlterIndex is a helper method to define mock.On call
// - ctx context.Context
// - newIndex *model.Index
func (_e *DataCoordCatalog_Expecter) AlterIndex(ctx interface{}, newIndex interface{}) *DataCoordCatalog_AlterIndex_Call {
return &DataCoordCatalog_AlterIndex_Call{Call: _e.mock.On("AlterIndex", ctx, newIndex)}
}
func (_c *DataCoordCatalog_AlterIndex_Call) Run(run func(ctx context.Context, newIndex *model.Index)) *DataCoordCatalog_AlterIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*model.Index))
})
return _c
}
func (_c *DataCoordCatalog_AlterIndex_Call) Return(_a0 error) *DataCoordCatalog_AlterIndex_Call {
_c.Call.Return(_a0)
return _c
}
// AlterIndexes provides a mock function with given fields: ctx, newIndexes
func (_m *DataCoordCatalog) AlterIndexes(ctx context.Context, newIndexes []*model.Index) error {
ret := _m.Called(ctx, newIndexes)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, []*model.Index) error); ok {
r0 = rf(ctx, newIndexes)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_AlterIndexes_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AlterIndexes'
type DataCoordCatalog_AlterIndexes_Call struct {
*mock.Call
}
// AlterIndexes is a helper method to define mock.On call
// - ctx context.Context
// - newIndexes []*model.Index
func (_e *DataCoordCatalog_Expecter) AlterIndexes(ctx interface{}, newIndexes interface{}) *DataCoordCatalog_AlterIndexes_Call {
return &DataCoordCatalog_AlterIndexes_Call{Call: _e.mock.On("AlterIndexes", ctx, newIndexes)}
}
func (_c *DataCoordCatalog_AlterIndexes_Call) Run(run func(ctx context.Context, newIndexes []*model.Index)) *DataCoordCatalog_AlterIndexes_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([]*model.Index))
})
return _c
}
func (_c *DataCoordCatalog_AlterIndexes_Call) Return(_a0 error) *DataCoordCatalog_AlterIndexes_Call {
_c.Call.Return(_a0)
return _c
}
// AlterSegment provides a mock function with given fields: ctx, newSegment, oldSegment
func (_m *DataCoordCatalog) AlterSegment(ctx context.Context, newSegment *datapb.SegmentInfo, oldSegment *datapb.SegmentInfo) error {
ret := _m.Called(ctx, newSegment, oldSegment)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *datapb.SegmentInfo, *datapb.SegmentInfo) error); ok {
r0 = rf(ctx, newSegment, oldSegment)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_AlterSegment_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AlterSegment'
type DataCoordCatalog_AlterSegment_Call struct {
*mock.Call
}
// AlterSegment is a helper method to define mock.On call
// - ctx context.Context
// - newSegment *datapb.SegmentInfo
// - oldSegment *datapb.SegmentInfo
func (_e *DataCoordCatalog_Expecter) AlterSegment(ctx interface{}, newSegment interface{}, oldSegment interface{}) *DataCoordCatalog_AlterSegment_Call {
return &DataCoordCatalog_AlterSegment_Call{Call: _e.mock.On("AlterSegment", ctx, newSegment, oldSegment)}
}
func (_c *DataCoordCatalog_AlterSegment_Call) Run(run func(ctx context.Context, newSegment *datapb.SegmentInfo, oldSegment *datapb.SegmentInfo)) *DataCoordCatalog_AlterSegment_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*datapb.SegmentInfo), args[2].(*datapb.SegmentInfo))
})
return _c
}
func (_c *DataCoordCatalog_AlterSegment_Call) Return(_a0 error) *DataCoordCatalog_AlterSegment_Call {
_c.Call.Return(_a0)
return _c
}
// AlterSegmentIndex provides a mock function with given fields: ctx, newSegIndex
func (_m *DataCoordCatalog) AlterSegmentIndex(ctx context.Context, newSegIndex *model.SegmentIndex) error {
ret := _m.Called(ctx, newSegIndex)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *model.SegmentIndex) error); ok {
r0 = rf(ctx, newSegIndex)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_AlterSegmentIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AlterSegmentIndex'
type DataCoordCatalog_AlterSegmentIndex_Call struct {
*mock.Call
}
// AlterSegmentIndex is a helper method to define mock.On call
// - ctx context.Context
// - newSegIndex *model.SegmentIndex
func (_e *DataCoordCatalog_Expecter) AlterSegmentIndex(ctx interface{}, newSegIndex interface{}) *DataCoordCatalog_AlterSegmentIndex_Call {
return &DataCoordCatalog_AlterSegmentIndex_Call{Call: _e.mock.On("AlterSegmentIndex", ctx, newSegIndex)}
}
func (_c *DataCoordCatalog_AlterSegmentIndex_Call) Run(run func(ctx context.Context, newSegIndex *model.SegmentIndex)) *DataCoordCatalog_AlterSegmentIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*model.SegmentIndex))
})
return _c
}
func (_c *DataCoordCatalog_AlterSegmentIndex_Call) Return(_a0 error) *DataCoordCatalog_AlterSegmentIndex_Call {
_c.Call.Return(_a0)
return _c
}
// AlterSegmentIndexes provides a mock function with given fields: ctx, newSegIdxes
func (_m *DataCoordCatalog) AlterSegmentIndexes(ctx context.Context, newSegIdxes []*model.SegmentIndex) error {
ret := _m.Called(ctx, newSegIdxes)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, []*model.SegmentIndex) error); ok {
r0 = rf(ctx, newSegIdxes)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_AlterSegmentIndexes_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AlterSegmentIndexes'
type DataCoordCatalog_AlterSegmentIndexes_Call struct {
*mock.Call
}
// AlterSegmentIndexes is a helper method to define mock.On call
// - ctx context.Context
// - newSegIdxes []*model.SegmentIndex
func (_e *DataCoordCatalog_Expecter) AlterSegmentIndexes(ctx interface{}, newSegIdxes interface{}) *DataCoordCatalog_AlterSegmentIndexes_Call {
return &DataCoordCatalog_AlterSegmentIndexes_Call{Call: _e.mock.On("AlterSegmentIndexes", ctx, newSegIdxes)}
}
func (_c *DataCoordCatalog_AlterSegmentIndexes_Call) Run(run func(ctx context.Context, newSegIdxes []*model.SegmentIndex)) *DataCoordCatalog_AlterSegmentIndexes_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([]*model.SegmentIndex))
})
return _c
}
func (_c *DataCoordCatalog_AlterSegmentIndexes_Call) Return(_a0 error) *DataCoordCatalog_AlterSegmentIndexes_Call {
_c.Call.Return(_a0)
return _c
}
// AlterSegments provides a mock function with given fields: ctx, newSegments
func (_m *DataCoordCatalog) AlterSegments(ctx context.Context, newSegments []*datapb.SegmentInfo) error {
ret := _m.Called(ctx, newSegments)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, []*datapb.SegmentInfo) error); ok {
r0 = rf(ctx, newSegments)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_AlterSegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AlterSegments'
type DataCoordCatalog_AlterSegments_Call struct {
*mock.Call
}
// AlterSegments is a helper method to define mock.On call
// - ctx context.Context
// - newSegments []*datapb.SegmentInfo
func (_e *DataCoordCatalog_Expecter) AlterSegments(ctx interface{}, newSegments interface{}) *DataCoordCatalog_AlterSegments_Call {
return &DataCoordCatalog_AlterSegments_Call{Call: _e.mock.On("AlterSegments", ctx, newSegments)}
}
func (_c *DataCoordCatalog_AlterSegments_Call) Run(run func(ctx context.Context, newSegments []*datapb.SegmentInfo)) *DataCoordCatalog_AlterSegments_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([]*datapb.SegmentInfo))
})
return _c
}
func (_c *DataCoordCatalog_AlterSegments_Call) Return(_a0 error) *DataCoordCatalog_AlterSegments_Call {
_c.Call.Return(_a0)
return _c
}
// AlterSegmentsAndAddNewSegment provides a mock function with given fields: ctx, segments, newSegment
func (_m *DataCoordCatalog) AlterSegmentsAndAddNewSegment(ctx context.Context, segments []*datapb.SegmentInfo, newSegment *datapb.SegmentInfo) error {
ret := _m.Called(ctx, segments, newSegment)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, []*datapb.SegmentInfo, *datapb.SegmentInfo) error); ok {
r0 = rf(ctx, segments, newSegment)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_AlterSegmentsAndAddNewSegment_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AlterSegmentsAndAddNewSegment'
type DataCoordCatalog_AlterSegmentsAndAddNewSegment_Call struct {
*mock.Call
}
// AlterSegmentsAndAddNewSegment is a helper method to define mock.On call
// - ctx context.Context
// - segments []*datapb.SegmentInfo
// - newSegment *datapb.SegmentInfo
func (_e *DataCoordCatalog_Expecter) AlterSegmentsAndAddNewSegment(ctx interface{}, segments interface{}, newSegment interface{}) *DataCoordCatalog_AlterSegmentsAndAddNewSegment_Call {
return &DataCoordCatalog_AlterSegmentsAndAddNewSegment_Call{Call: _e.mock.On("AlterSegmentsAndAddNewSegment", ctx, segments, newSegment)}
}
func (_c *DataCoordCatalog_AlterSegmentsAndAddNewSegment_Call) Run(run func(ctx context.Context, segments []*datapb.SegmentInfo, newSegment *datapb.SegmentInfo)) *DataCoordCatalog_AlterSegmentsAndAddNewSegment_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([]*datapb.SegmentInfo), args[2].(*datapb.SegmentInfo))
})
return _c
}
func (_c *DataCoordCatalog_AlterSegmentsAndAddNewSegment_Call) Return(_a0 error) *DataCoordCatalog_AlterSegmentsAndAddNewSegment_Call {
_c.Call.Return(_a0)
return _c
}
// CreateIndex provides a mock function with given fields: ctx, index
func (_m *DataCoordCatalog) CreateIndex(ctx context.Context, index *model.Index) error {
ret := _m.Called(ctx, index)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *model.Index) error); ok {
r0 = rf(ctx, index)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_CreateIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateIndex'
type DataCoordCatalog_CreateIndex_Call struct {
*mock.Call
}
// CreateIndex is a helper method to define mock.On call
// - ctx context.Context
// - index *model.Index
func (_e *DataCoordCatalog_Expecter) CreateIndex(ctx interface{}, index interface{}) *DataCoordCatalog_CreateIndex_Call {
return &DataCoordCatalog_CreateIndex_Call{Call: _e.mock.On("CreateIndex", ctx, index)}
}
func (_c *DataCoordCatalog_CreateIndex_Call) Run(run func(ctx context.Context, index *model.Index)) *DataCoordCatalog_CreateIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*model.Index))
})
return _c
}
func (_c *DataCoordCatalog_CreateIndex_Call) Return(_a0 error) *DataCoordCatalog_CreateIndex_Call {
_c.Call.Return(_a0)
return _c
}
// CreateSegmentIndex provides a mock function with given fields: ctx, segIdx
func (_m *DataCoordCatalog) CreateSegmentIndex(ctx context.Context, segIdx *model.SegmentIndex) error {
ret := _m.Called(ctx, segIdx)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *model.SegmentIndex) error); ok {
r0 = rf(ctx, segIdx)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_CreateSegmentIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateSegmentIndex'
type DataCoordCatalog_CreateSegmentIndex_Call struct {
*mock.Call
}
// CreateSegmentIndex is a helper method to define mock.On call
// - ctx context.Context
// - segIdx *model.SegmentIndex
func (_e *DataCoordCatalog_Expecter) CreateSegmentIndex(ctx interface{}, segIdx interface{}) *DataCoordCatalog_CreateSegmentIndex_Call {
return &DataCoordCatalog_CreateSegmentIndex_Call{Call: _e.mock.On("CreateSegmentIndex", ctx, segIdx)}
}
func (_c *DataCoordCatalog_CreateSegmentIndex_Call) Run(run func(ctx context.Context, segIdx *model.SegmentIndex)) *DataCoordCatalog_CreateSegmentIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*model.SegmentIndex))
})
return _c
}
func (_c *DataCoordCatalog_CreateSegmentIndex_Call) Return(_a0 error) *DataCoordCatalog_CreateSegmentIndex_Call {
_c.Call.Return(_a0)
return _c
}
// DropChannel provides a mock function with given fields: ctx, channel
func (_m *DataCoordCatalog) DropChannel(ctx context.Context, channel string) error {
ret := _m.Called(ctx, channel)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
r0 = rf(ctx, channel)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_DropChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropChannel'
type DataCoordCatalog_DropChannel_Call struct {
*mock.Call
}
// DropChannel is a helper method to define mock.On call
// - ctx context.Context
// - channel string
func (_e *DataCoordCatalog_Expecter) DropChannel(ctx interface{}, channel interface{}) *DataCoordCatalog_DropChannel_Call {
return &DataCoordCatalog_DropChannel_Call{Call: _e.mock.On("DropChannel", ctx, channel)}
}
func (_c *DataCoordCatalog_DropChannel_Call) Run(run func(ctx context.Context, channel string)) *DataCoordCatalog_DropChannel_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *DataCoordCatalog_DropChannel_Call) Return(_a0 error) *DataCoordCatalog_DropChannel_Call {
_c.Call.Return(_a0)
return _c
}
// DropChannelCheckpoint provides a mock function with given fields: ctx, vChannel
func (_m *DataCoordCatalog) DropChannelCheckpoint(ctx context.Context, vChannel string) error {
ret := _m.Called(ctx, vChannel)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
r0 = rf(ctx, vChannel)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_DropChannelCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropChannelCheckpoint'
type DataCoordCatalog_DropChannelCheckpoint_Call struct {
*mock.Call
}
// DropChannelCheckpoint is a helper method to define mock.On call
// - ctx context.Context
// - vChannel string
func (_e *DataCoordCatalog_Expecter) DropChannelCheckpoint(ctx interface{}, vChannel interface{}) *DataCoordCatalog_DropChannelCheckpoint_Call {
return &DataCoordCatalog_DropChannelCheckpoint_Call{Call: _e.mock.On("DropChannelCheckpoint", ctx, vChannel)}
}
func (_c *DataCoordCatalog_DropChannelCheckpoint_Call) Run(run func(ctx context.Context, vChannel string)) *DataCoordCatalog_DropChannelCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *DataCoordCatalog_DropChannelCheckpoint_Call) Return(_a0 error) *DataCoordCatalog_DropChannelCheckpoint_Call {
_c.Call.Return(_a0)
return _c
}
// DropIndex provides a mock function with given fields: ctx, collID, dropIdxID
func (_m *DataCoordCatalog) DropIndex(ctx context.Context, collID int64, dropIdxID int64) error {
ret := _m.Called(ctx, collID, dropIdxID)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, int64, int64) error); ok {
r0 = rf(ctx, collID, dropIdxID)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_DropIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropIndex'
type DataCoordCatalog_DropIndex_Call struct {
*mock.Call
}
// DropIndex is a helper method to define mock.On call
// - ctx context.Context
// - collID int64
// - dropIdxID int64
func (_e *DataCoordCatalog_Expecter) DropIndex(ctx interface{}, collID interface{}, dropIdxID interface{}) *DataCoordCatalog_DropIndex_Call {
return &DataCoordCatalog_DropIndex_Call{Call: _e.mock.On("DropIndex", ctx, collID, dropIdxID)}
}
func (_c *DataCoordCatalog_DropIndex_Call) Run(run func(ctx context.Context, collID int64, dropIdxID int64)) *DataCoordCatalog_DropIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(int64))
})
return _c
}
func (_c *DataCoordCatalog_DropIndex_Call) Return(_a0 error) *DataCoordCatalog_DropIndex_Call {
_c.Call.Return(_a0)
return _c
}
// DropSegment provides a mock function with given fields: ctx, segment
func (_m *DataCoordCatalog) DropSegment(ctx context.Context, segment *datapb.SegmentInfo) error {
ret := _m.Called(ctx, segment)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *datapb.SegmentInfo) error); ok {
r0 = rf(ctx, segment)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_DropSegment_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropSegment'
type DataCoordCatalog_DropSegment_Call struct {
*mock.Call
}
// DropSegment is a helper method to define mock.On call
// - ctx context.Context
// - segment *datapb.SegmentInfo
func (_e *DataCoordCatalog_Expecter) DropSegment(ctx interface{}, segment interface{}) *DataCoordCatalog_DropSegment_Call {
return &DataCoordCatalog_DropSegment_Call{Call: _e.mock.On("DropSegment", ctx, segment)}
}
func (_c *DataCoordCatalog_DropSegment_Call) Run(run func(ctx context.Context, segment *datapb.SegmentInfo)) *DataCoordCatalog_DropSegment_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*datapb.SegmentInfo))
})
return _c
}
func (_c *DataCoordCatalog_DropSegment_Call) Return(_a0 error) *DataCoordCatalog_DropSegment_Call {
_c.Call.Return(_a0)
return _c
}
// DropSegmentIndex provides a mock function with given fields: ctx, collID, partID, segID, buildID
func (_m *DataCoordCatalog) DropSegmentIndex(ctx context.Context, collID int64, partID int64, segID int64, buildID int64) error {
ret := _m.Called(ctx, collID, partID, segID, buildID)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, int64, int64, int64, int64) error); ok {
r0 = rf(ctx, collID, partID, segID, buildID)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_DropSegmentIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DropSegmentIndex'
type DataCoordCatalog_DropSegmentIndex_Call struct {
*mock.Call
}
// DropSegmentIndex is a helper method to define mock.On call
// - ctx context.Context
// - collID int64
// - partID int64
// - segID int64
// - buildID int64
func (_e *DataCoordCatalog_Expecter) DropSegmentIndex(ctx interface{}, collID interface{}, partID interface{}, segID interface{}, buildID interface{}) *DataCoordCatalog_DropSegmentIndex_Call {
return &DataCoordCatalog_DropSegmentIndex_Call{Call: _e.mock.On("DropSegmentIndex", ctx, collID, partID, segID, buildID)}
}
func (_c *DataCoordCatalog_DropSegmentIndex_Call) Run(run func(ctx context.Context, collID int64, partID int64, segID int64, buildID int64)) *DataCoordCatalog_DropSegmentIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64), args[2].(int64), args[3].(int64), args[4].(int64))
})
return _c
}
func (_c *DataCoordCatalog_DropSegmentIndex_Call) Return(_a0 error) *DataCoordCatalog_DropSegmentIndex_Call {
_c.Call.Return(_a0)
return _c
}
// IsChannelDropped provides a mock function with given fields: ctx, channel
func (_m *DataCoordCatalog) IsChannelDropped(ctx context.Context, channel string) bool {
ret := _m.Called(ctx, channel)
var r0 bool
if rf, ok := ret.Get(0).(func(context.Context, string) bool); ok {
r0 = rf(ctx, channel)
} else {
r0 = ret.Get(0).(bool)
}
return r0
}
// DataCoordCatalog_IsChannelDropped_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsChannelDropped'
type DataCoordCatalog_IsChannelDropped_Call struct {
*mock.Call
}
// IsChannelDropped is a helper method to define mock.On call
// - ctx context.Context
// - channel string
func (_e *DataCoordCatalog_Expecter) IsChannelDropped(ctx interface{}, channel interface{}) *DataCoordCatalog_IsChannelDropped_Call {
return &DataCoordCatalog_IsChannelDropped_Call{Call: _e.mock.On("IsChannelDropped", ctx, channel)}
}
func (_c *DataCoordCatalog_IsChannelDropped_Call) Run(run func(ctx context.Context, channel string)) *DataCoordCatalog_IsChannelDropped_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *DataCoordCatalog_IsChannelDropped_Call) Return(_a0 bool) *DataCoordCatalog_IsChannelDropped_Call {
_c.Call.Return(_a0)
return _c
}
// ListChannelCheckpoint provides a mock function with given fields: ctx
func (_m *DataCoordCatalog) ListChannelCheckpoint(ctx context.Context) (map[string]*internalpb.MsgPosition, error) {
ret := _m.Called(ctx)
var r0 map[string]*internalpb.MsgPosition
if rf, ok := ret.Get(0).(func(context.Context) map[string]*internalpb.MsgPosition); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(map[string]*internalpb.MsgPosition)
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// DataCoordCatalog_ListChannelCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListChannelCheckpoint'
type DataCoordCatalog_ListChannelCheckpoint_Call struct {
*mock.Call
}
// ListChannelCheckpoint is a helper method to define mock.On call
// - ctx context.Context
func (_e *DataCoordCatalog_Expecter) ListChannelCheckpoint(ctx interface{}) *DataCoordCatalog_ListChannelCheckpoint_Call {
return &DataCoordCatalog_ListChannelCheckpoint_Call{Call: _e.mock.On("ListChannelCheckpoint", ctx)}
}
func (_c *DataCoordCatalog_ListChannelCheckpoint_Call) Run(run func(ctx context.Context)) *DataCoordCatalog_ListChannelCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *DataCoordCatalog_ListChannelCheckpoint_Call) Return(_a0 map[string]*internalpb.MsgPosition, _a1 error) *DataCoordCatalog_ListChannelCheckpoint_Call {
_c.Call.Return(_a0, _a1)
return _c
}
// ListIndexes provides a mock function with given fields: ctx
func (_m *DataCoordCatalog) ListIndexes(ctx context.Context) ([]*model.Index, error) {
ret := _m.Called(ctx)
var r0 []*model.Index
if rf, ok := ret.Get(0).(func(context.Context) []*model.Index); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*model.Index)
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// DataCoordCatalog_ListIndexes_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListIndexes'
type DataCoordCatalog_ListIndexes_Call struct {
*mock.Call
}
// ListIndexes is a helper method to define mock.On call
// - ctx context.Context
func (_e *DataCoordCatalog_Expecter) ListIndexes(ctx interface{}) *DataCoordCatalog_ListIndexes_Call {
return &DataCoordCatalog_ListIndexes_Call{Call: _e.mock.On("ListIndexes", ctx)}
}
func (_c *DataCoordCatalog_ListIndexes_Call) Run(run func(ctx context.Context)) *DataCoordCatalog_ListIndexes_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *DataCoordCatalog_ListIndexes_Call) Return(_a0 []*model.Index, _a1 error) *DataCoordCatalog_ListIndexes_Call {
_c.Call.Return(_a0, _a1)
return _c
}
// ListSegmentIndexes provides a mock function with given fields: ctx
func (_m *DataCoordCatalog) ListSegmentIndexes(ctx context.Context) ([]*model.SegmentIndex, error) {
ret := _m.Called(ctx)
var r0 []*model.SegmentIndex
if rf, ok := ret.Get(0).(func(context.Context) []*model.SegmentIndex); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*model.SegmentIndex)
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// DataCoordCatalog_ListSegmentIndexes_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListSegmentIndexes'
type DataCoordCatalog_ListSegmentIndexes_Call struct {
*mock.Call
}
// ListSegmentIndexes is a helper method to define mock.On call
// - ctx context.Context
func (_e *DataCoordCatalog_Expecter) ListSegmentIndexes(ctx interface{}) *DataCoordCatalog_ListSegmentIndexes_Call {
return &DataCoordCatalog_ListSegmentIndexes_Call{Call: _e.mock.On("ListSegmentIndexes", ctx)}
}
func (_c *DataCoordCatalog_ListSegmentIndexes_Call) Run(run func(ctx context.Context)) *DataCoordCatalog_ListSegmentIndexes_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *DataCoordCatalog_ListSegmentIndexes_Call) Return(_a0 []*model.SegmentIndex, _a1 error) *DataCoordCatalog_ListSegmentIndexes_Call {
_c.Call.Return(_a0, _a1)
return _c
}
// ListSegments provides a mock function with given fields: ctx
func (_m *DataCoordCatalog) ListSegments(ctx context.Context) ([]*datapb.SegmentInfo, error) {
ret := _m.Called(ctx)
var r0 []*datapb.SegmentInfo
if rf, ok := ret.Get(0).(func(context.Context) []*datapb.SegmentInfo); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*datapb.SegmentInfo)
}
}
var r1 error
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// DataCoordCatalog_ListSegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListSegments'
type DataCoordCatalog_ListSegments_Call struct {
*mock.Call
}
// ListSegments is a helper method to define mock.On call
// - ctx context.Context
func (_e *DataCoordCatalog_Expecter) ListSegments(ctx interface{}) *DataCoordCatalog_ListSegments_Call {
return &DataCoordCatalog_ListSegments_Call{Call: _e.mock.On("ListSegments", ctx)}
}
func (_c *DataCoordCatalog_ListSegments_Call) Run(run func(ctx context.Context)) *DataCoordCatalog_ListSegments_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *DataCoordCatalog_ListSegments_Call) Return(_a0 []*datapb.SegmentInfo, _a1 error) *DataCoordCatalog_ListSegments_Call {
_c.Call.Return(_a0, _a1)
return _c
}
// MarkChannelDeleted provides a mock function with given fields: ctx, channel
func (_m *DataCoordCatalog) MarkChannelDeleted(ctx context.Context, channel string) error {
ret := _m.Called(ctx, channel)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
r0 = rf(ctx, channel)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_MarkChannelDeleted_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'MarkChannelDeleted'
type DataCoordCatalog_MarkChannelDeleted_Call struct {
*mock.Call
}
// MarkChannelDeleted is a helper method to define mock.On call
// - ctx context.Context
// - channel string
func (_e *DataCoordCatalog_Expecter) MarkChannelDeleted(ctx interface{}, channel interface{}) *DataCoordCatalog_MarkChannelDeleted_Call {
return &DataCoordCatalog_MarkChannelDeleted_Call{Call: _e.mock.On("MarkChannelDeleted", ctx, channel)}
}
func (_c *DataCoordCatalog_MarkChannelDeleted_Call) Run(run func(ctx context.Context, channel string)) *DataCoordCatalog_MarkChannelDeleted_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *DataCoordCatalog_MarkChannelDeleted_Call) Return(_a0 error) *DataCoordCatalog_MarkChannelDeleted_Call {
_c.Call.Return(_a0)
return _c
}
// RevertAlterSegmentsAndAddNewSegment provides a mock function with given fields: ctx, segments, removalSegment
func (_m *DataCoordCatalog) RevertAlterSegmentsAndAddNewSegment(ctx context.Context, segments []*datapb.SegmentInfo, removalSegment *datapb.SegmentInfo) error {
ret := _m.Called(ctx, segments, removalSegment)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, []*datapb.SegmentInfo, *datapb.SegmentInfo) error); ok {
r0 = rf(ctx, segments, removalSegment)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_RevertAlterSegmentsAndAddNewSegment_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RevertAlterSegmentsAndAddNewSegment'
type DataCoordCatalog_RevertAlterSegmentsAndAddNewSegment_Call struct {
*mock.Call
}
// RevertAlterSegmentsAndAddNewSegment is a helper method to define mock.On call
// - ctx context.Context
// - segments []*datapb.SegmentInfo
// - removalSegment *datapb.SegmentInfo
func (_e *DataCoordCatalog_Expecter) RevertAlterSegmentsAndAddNewSegment(ctx interface{}, segments interface{}, removalSegment interface{}) *DataCoordCatalog_RevertAlterSegmentsAndAddNewSegment_Call {
return &DataCoordCatalog_RevertAlterSegmentsAndAddNewSegment_Call{Call: _e.mock.On("RevertAlterSegmentsAndAddNewSegment", ctx, segments, removalSegment)}
}
func (_c *DataCoordCatalog_RevertAlterSegmentsAndAddNewSegment_Call) Run(run func(ctx context.Context, segments []*datapb.SegmentInfo, removalSegment *datapb.SegmentInfo)) *DataCoordCatalog_RevertAlterSegmentsAndAddNewSegment_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([]*datapb.SegmentInfo), args[2].(*datapb.SegmentInfo))
})
return _c
}
func (_c *DataCoordCatalog_RevertAlterSegmentsAndAddNewSegment_Call) Return(_a0 error) *DataCoordCatalog_RevertAlterSegmentsAndAddNewSegment_Call {
_c.Call.Return(_a0)
return _c
}
// SaveChannelCheckpoint provides a mock function with given fields: ctx, vChannel, pos
func (_m *DataCoordCatalog) SaveChannelCheckpoint(ctx context.Context, vChannel string, pos *internalpb.MsgPosition) error {
ret := _m.Called(ctx, vChannel, pos)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string, *internalpb.MsgPosition) error); ok {
r0 = rf(ctx, vChannel, pos)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_SaveChannelCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SaveChannelCheckpoint'
type DataCoordCatalog_SaveChannelCheckpoint_Call struct {
*mock.Call
}
// SaveChannelCheckpoint is a helper method to define mock.On call
// - ctx context.Context
// - vChannel string
// - pos *internalpb.MsgPosition
func (_e *DataCoordCatalog_Expecter) SaveChannelCheckpoint(ctx interface{}, vChannel interface{}, pos interface{}) *DataCoordCatalog_SaveChannelCheckpoint_Call {
return &DataCoordCatalog_SaveChannelCheckpoint_Call{Call: _e.mock.On("SaveChannelCheckpoint", ctx, vChannel, pos)}
}
func (_c *DataCoordCatalog_SaveChannelCheckpoint_Call) Run(run func(ctx context.Context, vChannel string, pos *internalpb.MsgPosition)) *DataCoordCatalog_SaveChannelCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].(*internalpb.MsgPosition))
})
return _c
}
func (_c *DataCoordCatalog_SaveChannelCheckpoint_Call) Return(_a0 error) *DataCoordCatalog_SaveChannelCheckpoint_Call {
_c.Call.Return(_a0)
return _c
}
// SaveDroppedSegmentsInBatch provides a mock function with given fields: ctx, segments
func (_m *DataCoordCatalog) SaveDroppedSegmentsInBatch(ctx context.Context, segments []*datapb.SegmentInfo) error {
ret := _m.Called(ctx, segments)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, []*datapb.SegmentInfo) error); ok {
r0 = rf(ctx, segments)
} else {
r0 = ret.Error(0)
}
return r0
}
// DataCoordCatalog_SaveDroppedSegmentsInBatch_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SaveDroppedSegmentsInBatch'
type DataCoordCatalog_SaveDroppedSegmentsInBatch_Call struct {
*mock.Call
}
// SaveDroppedSegmentsInBatch is a helper method to define mock.On call
// - ctx context.Context
// - segments []*datapb.SegmentInfo
func (_e *DataCoordCatalog_Expecter) SaveDroppedSegmentsInBatch(ctx interface{}, segments interface{}) *DataCoordCatalog_SaveDroppedSegmentsInBatch_Call {
return &DataCoordCatalog_SaveDroppedSegmentsInBatch_Call{Call: _e.mock.On("SaveDroppedSegmentsInBatch", ctx, segments)}
}
func (_c *DataCoordCatalog_SaveDroppedSegmentsInBatch_Call) Run(run func(ctx context.Context, segments []*datapb.SegmentInfo)) *DataCoordCatalog_SaveDroppedSegmentsInBatch_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([]*datapb.SegmentInfo))
})
return _c
}
func (_c *DataCoordCatalog_SaveDroppedSegmentsInBatch_Call) Return(_a0 error) *DataCoordCatalog_SaveDroppedSegmentsInBatch_Call {
_c.Call.Return(_a0)
return _c
}
type mockConstructorTestingTNewDataCoordCatalog interface {
mock.TestingT
Cleanup(func())
}
// NewDataCoordCatalog creates a new instance of DataCoordCatalog. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
func NewDataCoordCatalog(t mockConstructorTestingTNewDataCoordCatalog) *DataCoordCatalog {
mock := &DataCoordCatalog{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}