From 187aff68c51dd14f9aa23ba49eea178bfde7d528 Mon Sep 17 00:00:00 2001 From: jaime Date: Tue, 10 Jan 2023 15:31:46 +0800 Subject: [PATCH] [Cherry-Pick]Remove redundant LoadPrefix requests for Catalog ListCollections (#21551) (#21594) Signed-off-by: jaime --- internal/metastore/kv/rootcoord/kv_catalog.go | 25 ++- .../metastore/kv/rootcoord/kv_catalog_test.go | 177 ++++++++++++++---- 2 files changed, 159 insertions(+), 43 deletions(-) diff --git a/internal/metastore/kv/rootcoord/kv_catalog.go b/internal/metastore/kv/rootcoord/kv_catalog.go index 9d0ca50e5..ce1791faf 100644 --- a/internal/metastore/kv/rootcoord/kv_catalog.go +++ b/internal/metastore/kv/rootcoord/kv_catalog.go @@ -274,11 +274,8 @@ func (kc *Catalog) listFieldsAfter210(ctx context.Context, collectionID typeutil return fields, nil } -func (kc *Catalog) GetCollectionByID(ctx context.Context, collectionID typeutil.UniqueID, ts typeutil.Timestamp) (*model.Collection, error) { - collMeta, err := kc.loadCollection(ctx, collectionID, ts) - if err != nil { - return nil, err - } +func (kc *Catalog) appendPartitionAndFieldsInfo(ctx context.Context, collMeta *pb.CollectionInfo, + ts typeutil.Timestamp) (*model.Collection, error) { collection := model.UnmarshalCollectionModel(collMeta) @@ -286,13 +283,13 @@ func (kc *Catalog) GetCollectionByID(ctx context.Context, collectionID typeutil. return collection, nil } - partitions, err := kc.listPartitionsAfter210(ctx, collectionID, ts) + partitions, err := kc.listPartitionsAfter210(ctx, collection.CollectionID, ts) if err != nil { return nil, err } collection.Partitions = partitions - fields, err := kc.listFieldsAfter210(ctx, collectionID, ts) + fields, err := kc.listFieldsAfter210(ctx, collection.CollectionID, ts) if err != nil { return nil, err } @@ -301,6 +298,16 @@ func (kc *Catalog) GetCollectionByID(ctx context.Context, collectionID typeutil. return collection, nil } +func (kc *Catalog) GetCollectionByID(ctx context.Context, collectionID typeutil.UniqueID, + ts typeutil.Timestamp) (*model.Collection, error) { + collMeta, err := kc.loadCollection(ctx, collectionID, ts) + if err != nil { + return nil, err + } + + return kc.appendPartitionAndFieldsInfo(ctx, collMeta, ts) +} + func (kc *Catalog) CollectionExists(ctx context.Context, collectionID typeutil.UniqueID, ts typeutil.Timestamp) bool { _, err := kc.GetCollectionByID(ctx, collectionID, ts) return err == nil @@ -501,7 +508,7 @@ func (kc *Catalog) ListCollections(ctx context.Context, ts typeutil.Timestamp) ( zap.String("prefix", CollectionMetaPrefix), zap.Uint64("timestamp", ts), zap.Error(err)) - return nil, nil + return nil, err } colls := make(map[string]*model.Collection) @@ -512,7 +519,7 @@ func (kc *Catalog) ListCollections(ctx context.Context, ts typeutil.Timestamp) ( log.Warn("unmarshal collection info failed", zap.Error(err)) continue } - collection, err := kc.GetCollectionByID(ctx, collMeta.GetID(), ts) + collection, err := kc.appendPartitionAndFieldsInfo(ctx, &collMeta, ts) if err != nil { return nil, err } diff --git a/internal/metastore/kv/rootcoord/kv_catalog_test.go b/internal/metastore/kv/rootcoord/kv_catalog_test.go index 7da12324e..91f77414e 100644 --- a/internal/metastore/kv/rootcoord/kv_catalog_test.go +++ b/internal/metastore/kv/rootcoord/kv_catalog_test.go @@ -31,20 +31,6 @@ import ( "github.com/milvus-io/milvus/internal/util/typeutil" ) -type MockedTxnKV struct { - kv.TxnKV - loadWithPrefixFn func(key string) ([]string, []string, error) -} - -func (mc *MockedTxnKV) LoadWithPrefix(key string) ([]string, []string, error) { - return mc.loadWithPrefixFn(key) -} - -type MockedSnapShotKV struct { - mock.Mock - kv.SnapShotKV -} - var ( indexName = "idx" IndexID = 1 @@ -58,29 +44,152 @@ var ( } ) -func getStrIndexPb(t *testing.T) string { - idxPB := model.MarshalIndexModel(&index) - msg, err := proto.Marshal(idxPB) - assert.Nil(t, err) - return string(msg) -} +func TestCatalog_ListCollections(t *testing.T) { + ctx := context.Background() -func getStrSegIdxPb(idx model.Index, newSegIdx model.SegmentIndex) (string, error) { - segIdxInfo := &pb.SegmentIndexInfo{ - CollectionID: idx.CollectionID, - PartitionID: newSegIdx.PartitionID, - SegmentID: newSegIdx.SegmentID, - BuildID: newSegIdx.BuildID, - //EnableIndex: newSegIdx.EnableIndex, - CreateTime: newSegIdx.CreateTime, - FieldID: idx.FieldID, - IndexID: idx.IndexID, + coll1 := &pb.CollectionInfo{ + ID: 1, + PartitionIDs: []int64{100}, + PartitionNames: []string{"0"}, + PartitionCreatedTimestamps: []uint64{1}, + Schema: &schemapb.CollectionSchema{ + Name: "c1", + Fields: []*schemapb.FieldSchema{ + { + FieldID: 1, + Name: "f1", + }, + }, + }, } - msg, err := proto.Marshal(segIdxInfo) - if err != nil { - return "", err + + coll2 := &pb.CollectionInfo{ + ID: 2, + Schema: &schemapb.CollectionSchema{ + Name: "c1", + Fields: []*schemapb.FieldSchema{ + {}, + }, + }, } - return string(msg), nil + + targetErr := errors.New("fail") + + t.Run("load collection with prefix fail", func(t *testing.T) { + kv := mocks.NewSnapShotKV(t) + ts := uint64(1) + kv.On("LoadWithPrefix", CollectionMetaPrefix, ts). + Return(nil, nil, targetErr) + + kc := Catalog{Snapshot: kv} + ret, err := kc.ListCollections(ctx, ts) + assert.ErrorIs(t, err, targetErr) + assert.Nil(t, ret) + }) + + t.Run("list partition fail", func(t *testing.T) { + kv := mocks.NewSnapShotKV(t) + ts := uint64(1) + + bColl, err := proto.Marshal(coll2) + assert.NoError(t, err) + kv.On("LoadWithPrefix", CollectionMetaPrefix, ts). + Return(nil, []string{string(bColl)}, nil) + kv.On("LoadWithPrefix", mock.MatchedBy( + func(prefix string) bool { + return strings.HasPrefix(prefix, PartitionMetaPrefix) + }), ts). + Return(nil, nil, targetErr) + kc := Catalog{Snapshot: kv} + + ret, err := kc.ListCollections(ctx, ts) + assert.ErrorIs(t, err, targetErr) + assert.Nil(t, ret) + }) + + t.Run("list fields fail", func(t *testing.T) { + kv := mocks.NewSnapShotKV(t) + ts := uint64(1) + + bColl, err := proto.Marshal(coll2) + assert.NoError(t, err) + kv.On("LoadWithPrefix", CollectionMetaPrefix, ts). + Return(nil, []string{string(bColl)}, nil) + + partitionMeta := &pb.PartitionInfo{} + pm, err := proto.Marshal(partitionMeta) + assert.NoError(t, err) + + kv.On("LoadWithPrefix", mock.MatchedBy( + func(prefix string) bool { + return strings.HasPrefix(prefix, PartitionMetaPrefix) + }), ts). + Return(nil, []string{string(pm)}, nil) + + kv.On("LoadWithPrefix", mock.MatchedBy( + func(prefix string) bool { + return strings.HasPrefix(prefix, FieldMetaPrefix) + }), ts). + Return(nil, nil, targetErr) + kc := Catalog{Snapshot: kv} + + ret, err := kc.ListCollections(ctx, ts) + assert.ErrorIs(t, err, targetErr) + assert.Nil(t, ret) + }) + + t.Run("list collection ok for 210 version", func(t *testing.T) { + kv := mocks.NewSnapShotKV(t) + ts := uint64(1) + + bColl, err := proto.Marshal(coll1) + assert.NoError(t, err) + kv.On("LoadWithPrefix", CollectionMetaPrefix, ts). + Return(nil, []string{string(bColl)}, nil) + kc := Catalog{Snapshot: kv} + + ret, err := kc.ListCollections(ctx, ts) + assert.NoError(t, err) + assert.Equal(t, 1, len(ret)) + modCol := ret["c1"] + assert.Equal(t, coll1.ID, modCol.CollectionID) + }) + + t.Run("list collection ok for the newest version", func(t *testing.T) { + kv := mocks.NewSnapShotKV(t) + ts := uint64(1) + + bColl, err := proto.Marshal(coll2) + assert.NoError(t, err) + kv.On("LoadWithPrefix", CollectionMetaPrefix, ts). + Return(nil, []string{string(bColl)}, nil) + + partitionMeta := &pb.PartitionInfo{} + pm, err := proto.Marshal(partitionMeta) + assert.NoError(t, err) + + kv.On("LoadWithPrefix", mock.MatchedBy( + func(prefix string) bool { + return strings.HasPrefix(prefix, PartitionMetaPrefix) + }), ts). + Return(nil, []string{string(pm)}, nil) + + fieldMeta := &schemapb.FieldSchema{} + fm, err := proto.Marshal(fieldMeta) + assert.NoError(t, err) + + kv.On("LoadWithPrefix", mock.MatchedBy( + func(prefix string) bool { + return strings.HasPrefix(prefix, FieldMetaPrefix) + }), ts). + Return(nil, []string{string(fm)}, nil) + kc := Catalog{Snapshot: kv} + + ret, err := kc.ListCollections(ctx, ts) + assert.NoError(t, err) + assert.NotNil(t, ret) + assert.Equal(t, 1, len(ret)) + }) } func TestCatalog_loadCollection(t *testing.T) { -- GitLab