fix: Make MetaCache field of GrpcClient

Related to milvus-io#809 Signed-off-by: Congqi Xia <[email protected]>
congqixia · Aug 21, 2024 · 67db196 · 67db196
1 parent 0c339b6
commit 67db196
Show file tree

Hide file tree

Showing 14 changed files with 118 additions and 106 deletions.
diff --git a/client/client_mock_test.go b/client/client_mock_test.go
@@ -29,7 +29,7 @@ type MockSuiteBase struct {
 	svr  *grpc.Server
 	mock *mocks.MilvusServiceServer
 
-	client Client
+	client *GrpcClient
 }
 
 func (s *MockSuiteBase) SetupSuite() {
@@ -71,7 +71,9 @@ func (s *MockSuiteBase) SetupTest() {
 	s.Require().NoError(err)
 	s.setupConnect()
 
-	s.client = c
+	grpcClient, ok := c.(*GrpcClient)
+	s.Require().True(ok)
+	s.client = grpcClient
 }
 
 func (s *MockSuiteBase) TearDownTest() {
@@ -80,8 +82,9 @@ func (s *MockSuiteBase) TearDownTest() {
 }
 
 func (s *MockSuiteBase) resetMock() {
-	MetaCache.reset()
+	// MetaCache.reset()
 	if s.mock != nil {
+		s.client.cache.reset()
 		s.mock.Calls = nil
 		s.mock.ExpectedCalls = nil
 		s.setupConnect()

diff --git a/client/collection.go b/client/collection.go
@@ -292,7 +292,7 @@ func (c *GrpcClient) DescribeCollection(ctx context.Context, collName string) (*
 		Schema:           collection.Schema,
 		ConsistencyLevel: collection.ConsistencyLevel,
 	}
-	MetaCache.setCollectionInfo(collName, &colInfo)
+	c.cache.setCollectionInfo(collName, &colInfo)
 	return collection, nil
 }
 
@@ -317,7 +317,7 @@ func (c *GrpcClient) DropCollection(ctx context.Context, collName string, opts .
 	}
 	err = handleRespStatus(resp)
 	if err == nil {
-		MetaCache.setCollectionInfo(collName, nil)
+		c.cache.setCollectionInfo(collName, nil)
 	}
 	return err
 }

diff --git a/client/collection_test.go b/client/collection_test.go
@@ -276,10 +276,8 @@ func (s *CollectionSuite) TestCreateCollection() {
 			WithField(entity.NewField().WithName("dynamic").WithDataType(entity.FieldTypeJSON).WithIsDynamic(true))
 		for _, tc := range cases {
 			s.Run(tc.tag, func() {
-				grpcClient, ok := c.(*GrpcClient)
-				s.Require().True(ok)
-				grpcClient.config.addFlags(tc.flag)
-				defer grpcClient.config.resetFlags(tc.flag)
+				c.config.addFlags(tc.flag)
+				defer c.config.resetFlags(tc.flag)
 
 				err := c.CreateCollection(ctx, sch, 1)
 				s.Error(err)

diff --git a/client/data.go b/client/data.go
@@ -43,18 +43,11 @@ func (c *GrpcClient) HybridSearch(ctx context.Context, collName string, partitio
 		return nil, ErrClientNotReady
 	}
 
-	var schema *entity.Schema
-	collInfo, ok := MetaCache.getCollectionInfo(collName)
-	if !ok {
-		coll, err := c.DescribeCollection(ctx, collName)
-		if err != nil {
-			return nil, err
-		}
-		schema = coll.Schema
-		collInfo, _ = MetaCache.getCollectionInfo(collName)
-	} else {
-		schema = collInfo.Schema
+	collInfo, err := c.getCollectionInfo(ctx, collName)
+	if err != nil {
+		return nil, err
 	}
+	schema := collInfo.Schema
 
 	sReqs := make([]*milvuspb.SearchRequest, 0, len(subRequests))
 	nq := 0
@@ -103,19 +96,13 @@ func (c *GrpcClient) Search(ctx context.Context, collName string, partitions []s
 	if c.Service == nil {
 		return []SearchResult{}, ErrClientNotReady
 	}
-	var schema *entity.Schema
-	collInfo, ok := MetaCache.getCollectionInfo(collName)
-	if !ok {
-		coll, err := c.DescribeCollection(ctx, collName)
-		if err != nil {
-			return nil, err
-		}
-		schema = coll.Schema
-	} else {
-		schema = collInfo.Schema
+	collInfo, err := c.getCollectionInfo(ctx, collName)
+	if err != nil {
+		return nil, err
 	}
+	schema := collInfo.Schema
 
-	option, err := makeSearchQueryOption(collName, opts...)
+	option, err := c.makeSearchQueryOption(collName, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -320,19 +307,13 @@ func (c *GrpcClient) Query(ctx context.Context, collectionName string, partition
 		return nil, ErrClientNotReady
 	}
 
-	var sch *entity.Schema
-	collInfo, ok := MetaCache.getCollectionInfo(collectionName)
-	if !ok {
-		coll, err := c.DescribeCollection(ctx, collectionName)
-		if err != nil {
-			return nil, err
-		}
-		sch = coll.Schema
-	} else {
-		sch = collInfo.Schema
+	collInfo, err := c.getCollectionInfo(ctx, collectionName)
+	if err != nil {
+		return nil, err
 	}
+	sch := collInfo.Schema
 
-	option, err := makeSearchQueryOption(collectionName, opts...)
+	option, err := c.makeSearchQueryOption(collectionName, opts...)
 	if err != nil {
 		return nil, err
 	}

diff --git a/client/database.go b/client/database.go
@@ -36,6 +36,7 @@ func (c *GrpcClient) UsingDatabase(ctx context.Context, dbName string) error {
 	if err != nil {
 		return err
 	}
+	c.cache.reset()
 
 	return nil
 }

diff --git a/client/grpc_client.go b/client/grpc_client.go
@@ -7,6 +7,8 @@ import (
 	"strconv"
 	"time"
 
+	"golang.org/x/sync/singleflight"
+
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-sdk-go/v2/common"
@@ -24,6 +26,9 @@ type GrpcClient struct {
 	Conn    *grpc.ClientConn             // grpc connection instance
 	Service milvuspb.MilvusServiceClient // Service client stub
 
+	cache *metaCache
+	sf    singleflight.Group
+
 	config *Config // No thread safety
 }
 
@@ -32,6 +37,8 @@ func (c *GrpcClient) connect(ctx context.Context, addr string, opts ...grpc.Dial
 	if addr == "" {
 		return fmt.Errorf("address is empty")
 	}
+
+	c.cache = newMetaCache()
 	conn, err := grpc.DialContext(ctx, addr, opts...)
 	if err != nil {
 		return err
@@ -92,6 +99,32 @@ func (c *GrpcClient) connectInternal(ctx context.Context) error {
 	return nil
 }
 
+func (c *GrpcClient) getCollectionInfo(ctx context.Context, collectionName string) (*collInfo, error) {
+	// try cache first
+	info, ok := c.cache.getCollectionInfo(collectionName)
+	if ok {
+		return info, nil
+	}
+
+	// use singleflight to fetch collection info
+	v, err, _ := c.sf.Do(collectionName, func() (interface{}, error) {
+		coll, err := c.DescribeCollection(ctx, collectionName)
+		if err != nil {
+			return nil, err
+		}
+		return &collInfo{
+			ID:               coll.ID,
+			Name:             coll.Name,
+			Schema:           coll.Schema,
+			ConsistencyLevel: coll.ConsistencyLevel,
+		}, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return v.(*collInfo), nil
+}
+
 // Close close the connection
 func (c *GrpcClient) Close() error {
 	if c.Conn != nil {

diff --git a/client/insert.go b/client/insert.go
@@ -35,17 +35,11 @@ func (c *GrpcClient) Insert(ctx context.Context, collName string, partitionName
 	if c.Service == nil {
 		return nil, ErrClientNotReady
 	}
-	var schema *entity.Schema
-	collInfo, ok := MetaCache.getCollectionInfo(collName)
-	if !ok {
-		coll, err := c.DescribeCollection(ctx, collName)
-		if err != nil {
-			return nil, err
-		}
-		schema = coll.Schema
-	} else {
-		schema = collInfo.Schema
+	collInfo, err := c.getCollectionInfo(ctx, collName)
+	if err != nil {
+		return nil, err
 	}
+	schema := collInfo.Schema
 
 	// convert columns to field data
 	fieldsData, rowSize, err := c.processInsertColumns(schema, columns...)
@@ -70,7 +64,7 @@ func (c *GrpcClient) Insert(ctx context.Context, collName string, partitionName
 	if err := handleRespStatus(resp.GetStatus()); err != nil {
 		return nil, err
 	}
-	MetaCache.setSessionTs(collName, resp.Timestamp)
+	c.cache.setSessionTs(collName, resp.Timestamp)
 	// 3. parse id column
 	return entity.IDColumns(schema, resp.GetIDs(), 0, -1)
 }
@@ -332,7 +326,7 @@ func (c *GrpcClient) DeleteByPks(ctx context.Context, collName string, partition
 	if err != nil {
 		return err
 	}
-	MetaCache.setSessionTs(collName, resp.Timestamp)
+	c.cache.setSessionTs(collName, resp.Timestamp)
 	return nil
 }
 
@@ -370,7 +364,7 @@ func (c *GrpcClient) Delete(ctx context.Context, collName string, partitionName
 	if err != nil {
 		return err
 	}
-	MetaCache.setSessionTs(collName, resp.Timestamp)
+	c.cache.setSessionTs(collName, resp.Timestamp)
 	return nil
 }
 
@@ -382,17 +376,11 @@ func (c *GrpcClient) Upsert(ctx context.Context, collName string, partitionName
 	if c.Service == nil {
 		return nil, ErrClientNotReady
 	}
-	var schema *entity.Schema
-	collInfo, ok := MetaCache.getCollectionInfo(collName)
-	if !ok {
-		coll, err := c.DescribeCollection(ctx, collName)
-		if err != nil {
-			return nil, err
-		}
-		schema = coll.Schema
-	} else {
-		schema = collInfo.Schema
+	collInfo, err := c.getCollectionInfo(ctx, collName)
+	if err != nil {
+		return nil, err
 	}
+	schema := collInfo.Schema
 
 	fieldsData, rowSize, err := c.processInsertColumns(schema, columns...)
 	if err != nil {
@@ -416,7 +404,7 @@ func (c *GrpcClient) Upsert(ctx context.Context, collName string, partitionName
 	if err := handleRespStatus(resp.GetStatus()); err != nil {
 		return nil, err
 	}
-	MetaCache.setSessionTs(collName, resp.Timestamp)
+	c.cache.setSessionTs(collName, resp.Timestamp)
 	// 3. parse id column
 	return entity.IDColumns(schema, resp.GetIDs(), 0, -1)
 }

diff --git a/client/iterator.go b/client/iterator.go
@@ -12,17 +12,11 @@ import (
 
 func (c *GrpcClient) SearchIterator(ctx context.Context, opt *SearchIteratorOption) (*SearchIterator, error) {
 	collectionName := opt.collectionName
-	var sch *entity.Schema
-	collInfo, ok := MetaCache.getCollectionInfo(collectionName)
-	if !ok {
-		coll, err := c.DescribeCollection(ctx, collectionName)
-		if err != nil {
-			return nil, err
-		}
-		sch = coll.Schema
-	} else {
-		sch = collInfo.Schema
+	collInfo, err := c.getCollectionInfo(ctx, collectionName)
+	if err != nil {
+		return nil, err
 	}
+	sch := collInfo.Schema
 
 	var vectorField *entity.Field
 	for _, field := range sch.Fields {
@@ -52,7 +46,7 @@ func (c *GrpcClient) SearchIterator(ctx context.Context, opt *SearchIteratorOpti
 		expr:      opt.expr,
 	}
 
-	err := itr.init(ctx)
+	err = itr.init(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -202,17 +196,11 @@ func (itr *SearchIterator) Next(ctx context.Context) (*SearchResult, error) {
 
 func (c *GrpcClient) QueryIterator(ctx context.Context, opt *QueryIteratorOption) (*QueryIterator, error) {
 	collectionName := opt.collectionName
-	var sch *entity.Schema
-	collInfo, ok := MetaCache.getCollectionInfo(collectionName)
-	if !ok {
-		coll, err := c.DescribeCollection(ctx, collectionName)
-		if err != nil {
-			return nil, err
-		}
-		sch = coll.Schema
-	} else {
-		sch = collInfo.Schema
+	collInfo, err := c.getCollectionInfo(ctx, collectionName)
+	if err != nil {
+		return nil, err
 	}
+	sch := collInfo.Schema
 
 	itr := &QueryIterator{
 		client: c,
@@ -227,7 +215,7 @@ func (c *GrpcClient) QueryIterator(ctx context.Context, opt *QueryIteratorOption
 		expr:      opt.expr,
 	}
 
-	err := itr.init(ctx)
+	err = itr.init(ctx)
 	if err != nil {
 		return nil, err
 	}

diff --git a/client/meta_cache.go b/client/meta_cache.go
@@ -20,9 +20,11 @@ type collInfo struct {
 	ConsistencyLevel entity.ConsistencyLevel
 }
 
-var MetaCache = metaCache{
-	sessionTsMap: make(map[string]uint64),
-	collInfoMap:  make(map[string]collInfo),
+func newMetaCache() *metaCache {
+	return &metaCache{
+		sessionTsMap: make(map[string]uint64),
+		collInfoMap:  make(map[string]collInfo),
+	}
 }
 
 // timestampMap collects the last-write-timestamp of every collection, which is required by session consistency level.