enhance: remove the collectionBasicInfo class in the proxy metacache (#37874)

/kind improvement - issue: #37928 Signed-off-by: SimFG <bang.fu@zilliz.com>
2024-11-29 18:38:44 +08:00 · 2024-11-22 16:10:33 +08:00 · 2024-11-22 16:10:33 +08:00 · 7c5a8012cf
commit 7c5a8012cf
parent b48f531e63
5 changed files with 24 additions and 81 deletions
--- a/internal/proxy/meta_cache.go
+++ b/internal/proxy/meta_cache.go
@ -58,7 +58,7 @@ type Cache interface {
 	// GetCollectionName get collection's name and database by id
 	GetCollectionName(ctx context.Context, database string, collectionID int64) (string, error)
 	// GetCollectionInfo get collection's information by name or collection id, such as schema, and etc.
-	GetCollectionInfo(ctx context.Context, database, collectionName string, collectionID int64) (*collectionBasicInfo, error)
+	GetCollectionInfo(ctx context.Context, database, collectionName string, collectionID int64) (*collectionInfo, error)
 	// GetPartitionID get partition's identifier of specific collection.
 	GetPartitionID(ctx context.Context, database, collectionName string, partitionName string) (typeutil.UniqueID, error)
 	// GetPartitions get all partitions' id of specific collection.
@ -92,13 +92,6 @@ type Cache interface {
 	// AllocID is only using on requests that need to skip timestamp allocation, don't overuse it.
 	AllocID(ctx context.Context) (int64, error)
 }
-type collectionBasicInfo struct {
-	collID                typeutil.UniqueID
-	createdTimestamp      uint64
-	createdUtcTimestamp   uint64
-	consistencyLevel      commonpb.ConsistencyLevel
-	partitionKeyIsolation bool
-}

 type collectionInfo struct {
 	collID                typeutil.UniqueID
@ -272,20 +265,6 @@ type partitionInfo struct {
 	createdUtcTimestamp uint64
 }

-// getBasicInfo get a basic info by deep copy.
-func (info *collectionInfo) getBasicInfo() *collectionBasicInfo {
-	// Do a deep copy for all fields.
-	basicInfo := &collectionBasicInfo{
-		collID:                info.collID,
-		createdTimestamp:      info.createdTimestamp,
-		createdUtcTimestamp:   info.createdUtcTimestamp,
-		consistencyLevel:      info.consistencyLevel,
-		partitionKeyIsolation: info.partitionKeyIsolation,
-	}
-
-	return basicInfo
-}
-
 func (info *collectionInfo) isCollectionCached() bool {
 	return info != nil && info.collID != UniqueID(0) && info.schema != nil
 }
@ -556,7 +535,7 @@ func (m *MetaCache) GetCollectionName(ctx context.Context, database string, coll
 	return collInfo.schema.Name, nil
 }

-func (m *MetaCache) GetCollectionInfo(ctx context.Context, database string, collectionName string, collectionID int64) (*collectionBasicInfo, error) {
+func (m *MetaCache) GetCollectionInfo(ctx context.Context, database string, collectionName string, collectionID int64) (*collectionInfo, error) {
 	collInfo, ok := m.getCollection(database, collectionName, 0)

 	method := "GetCollectionInfo"
@ -571,11 +550,11 @@ func (m *MetaCache) GetCollectionInfo(ctx context.Context, database string, coll
 			return nil, err
 		}
 		metrics.ProxyUpdateCacheLatency.WithLabelValues(fmt.Sprint(paramtable.GetNodeID()), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
-		return collInfo.getBasicInfo(), nil
+		return collInfo, nil
 	}

 	metrics.ProxyCacheStatsCounter.WithLabelValues(fmt.Sprint(paramtable.GetNodeID()), method, metrics.CacheHitLabel).Inc()
-	return collInfo.getBasicInfo(), nil
+	return collInfo, nil
 }

 // GetCollectionInfo returns the collection information related to provided collection name
--- a/internal/proxy/mock_cache.go
+++ b/internal/proxy/mock_cache.go
@ -173,23 +173,23 @@ func (_c *MockCache_GetCollectionID_Call) RunAndReturn(run func(context.Context,
 }

 // GetCollectionInfo provides a mock function with given fields: ctx, database, collectionName, collectionID
-func (_m *MockCache) GetCollectionInfo(ctx context.Context, database string, collectionName string, collectionID int64) (*collectionBasicInfo, error) {
+func (_m *MockCache) GetCollectionInfo(ctx context.Context, database string, collectionName string, collectionID int64) (*collectionInfo, error) {
 	ret := _m.Called(ctx, database, collectionName, collectionID)

 	if len(ret) == 0 {
 		panic("no return value specified for GetCollectionInfo")
 	}

-	var r0 *collectionBasicInfo
+	var r0 *collectionInfo
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) (*collectionBasicInfo, error)); ok {
+	if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) (*collectionInfo, error)); ok {
 		return rf(ctx, database, collectionName, collectionID)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) *collectionBasicInfo); ok {
+	if rf, ok := ret.Get(0).(func(context.Context, string, string, int64) *collectionInfo); ok {
 		r0 = rf(ctx, database, collectionName, collectionID)
 	} else {
 		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*collectionBasicInfo)
+			r0 = ret.Get(0).(*collectionInfo)
 		}
 	}

@ -223,12 +223,12 @@ func (_c *MockCache_GetCollectionInfo_Call) Run(run func(ctx context.Context, da
 	return _c
 }

-func (_c *MockCache_GetCollectionInfo_Call) Return(_a0 *collectionBasicInfo, _a1 error) *MockCache_GetCollectionInfo_Call {
+func (_c *MockCache_GetCollectionInfo_Call) Return(_a0 *collectionInfo, _a1 error) *MockCache_GetCollectionInfo_Call {
 	_c.Call.Return(_a0, _a1)
 	return _c
 }

-func (_c *MockCache_GetCollectionInfo_Call) RunAndReturn(run func(context.Context, string, string, int64) (*collectionBasicInfo, error)) *MockCache_GetCollectionInfo_Call {
+func (_c *MockCache_GetCollectionInfo_Call) RunAndReturn(run func(context.Context, string, string, int64) (*collectionInfo, error)) *MockCache_GetCollectionInfo_Call {
 	_c.Call.Return(run)
 	return _c
 }
@ -1225,42 +1225,6 @@ func (_c *MockCache_RemoveDatabase_Call) RunAndReturn(run func(context.Context,
 	return _c
 }

-// RemovePartition provides a mock function with given fields: ctx, database, collectionName, partitionName
-func (_m *MockCache) RemovePartition(ctx context.Context, database string, collectionName string, partitionName string) {
-	_m.Called(ctx, database, collectionName, partitionName)
-}
-
-// MockCache_RemovePartition_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemovePartition'
-type MockCache_RemovePartition_Call struct {
-	*mock.Call
-}
-
-// RemovePartition is a helper method to define mock.On call
-//   - ctx context.Context
-//   - database string
-//   - collectionName string
-//   - partitionName string
-func (_e *MockCache_Expecter) RemovePartition(ctx interface{}, database interface{}, collectionName interface{}, partitionName interface{}) *MockCache_RemovePartition_Call {
-	return &MockCache_RemovePartition_Call{Call: _e.mock.On("RemovePartition", ctx, database, collectionName, partitionName)}
-}
-
-func (_c *MockCache_RemovePartition_Call) Run(run func(ctx context.Context, database string, collectionName string, partitionName string)) *MockCache_RemovePartition_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(string), args[2].(string), args[3].(string))
-	})
-	return _c
-}
-
-func (_c *MockCache_RemovePartition_Call) Return() *MockCache_RemovePartition_Call {
-	_c.Call.Return()
-	return _c
-}
-
-func (_c *MockCache_RemovePartition_Call) RunAndReturn(run func(context.Context, string, string, string)) *MockCache_RemovePartition_Call {
-	_c.Call.Return(run)
-	return _c
-}
-
 // UpdateCredential provides a mock function with given fields: credInfo
 func (_m *MockCache) UpdateCredential(credInfo *internalpb.CredentialInfo) {
 	_m.Called(credInfo)
--- a/internal/proxy/task_query_test.go
+++ b/internal/proxy/task_query_test.go
@ -1168,7 +1168,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
 		}
 		mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, nil)
 		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionBasicInfo{
+			&collectionInfo{
 				collID:           collID,
 				consistencyLevel: commonpb.ConsistencyLevel_Eventually,
 			}, nil).Once()
@ -1177,7 +1177,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
 		assert.True(t, skip)

 		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionBasicInfo{
+			&collectionInfo{
 				collID:           collID,
 				consistencyLevel: commonpb.ConsistencyLevel_Bounded,
 			}, nil).Once()
@ -1185,7 +1185,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
 		assert.True(t, skip)

 		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionBasicInfo{
+			&collectionInfo{
 				collID:           collID,
 				consistencyLevel: commonpb.ConsistencyLevel_Strong,
 			}, nil).Once()
@ -1195,7 +1195,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {

 	t.Run("request consistency level", func(t *testing.T) {
 		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionBasicInfo{
+			&collectionInfo{
 				collID:           collID,
 				consistencyLevel: commonpb.ConsistencyLevel_Eventually,
 			}, nil).Times(3)
@ -1267,7 +1267,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
 		mockMetaCache.ExpectedCalls = nil
 		mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, fmt.Errorf("mock error"))
 		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionBasicInfo{
+			&collectionInfo{
 				collID:           collID,
 				consistencyLevel: commonpb.ConsistencyLevel_Eventually,
 			}, nil)
--- a/internal/proxy/task_scheduler_test.go
+++ b/internal/proxy/task_scheduler_test.go
@ -621,7 +621,7 @@ func TestTaskScheduler_SkipAllocTimestamp(t *testing.T) {

 	mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, nil)
 	mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-		&collectionBasicInfo{
+		&collectionInfo{
 			collID:           collID,
 			consistencyLevel: commonpb.ConsistencyLevel_Eventually,
 		}, nil)
--- a/internal/proxy/task_search_test.go
+++ b/internal/proxy/task_search_test.go
@ -2655,7 +2655,7 @@ func TestSearchTask_Requery(t *testing.T) {
 	cache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collectionID, nil).Maybe()
 	cache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(schema, nil).Maybe()
 	cache.EXPECT().GetPartitions(mock.Anything, mock.Anything, mock.Anything).Return(map[string]int64{"_default": UniqueID(1)}, nil).Maybe()
-	cache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&collectionBasicInfo{}, nil).Maybe()
+	cache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&collectionInfo{}, nil).Maybe()
 	cache.EXPECT().GetShards(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(map[string][]nodeInfo{}, nil).Maybe()
 	cache.EXPECT().DeprecateShardCache(mock.Anything, mock.Anything).Return().Maybe()
 	globalMetaCache = cache
@ -2974,7 +2974,7 @@ func TestSearchTask_CanSkipAllocTimestamp(t *testing.T) {
 		}
 		mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, nil)
 		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionBasicInfo{
+			&collectionInfo{
 				collID:           collID,
 				consistencyLevel: commonpb.ConsistencyLevel_Eventually,
 			}, nil).Once()
@ -2983,7 +2983,7 @@ func TestSearchTask_CanSkipAllocTimestamp(t *testing.T) {
 		assert.True(t, skip)

 		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionBasicInfo{
+			&collectionInfo{
 				collID:           collID,
 				consistencyLevel: commonpb.ConsistencyLevel_Bounded,
 			}, nil).Once()
@ -2991,7 +2991,7 @@ func TestSearchTask_CanSkipAllocTimestamp(t *testing.T) {
 		assert.True(t, skip)

 		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionBasicInfo{
+			&collectionInfo{
 				collID:           collID,
 				consistencyLevel: commonpb.ConsistencyLevel_Strong,
 			}, nil).Once()
@ -3001,7 +3001,7 @@ func TestSearchTask_CanSkipAllocTimestamp(t *testing.T) {

 	t.Run("request consistency level", func(t *testing.T) {
 		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionBasicInfo{
+			&collectionInfo{
 				collID:           collID,
 				consistencyLevel: commonpb.ConsistencyLevel_Eventually,
 			}, nil).Times(3)
@ -3073,7 +3073,7 @@ func TestSearchTask_CanSkipAllocTimestamp(t *testing.T) {
 		mockMetaCache.ExpectedCalls = nil
 		mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, fmt.Errorf("mock error"))
 		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionBasicInfo{
+			&collectionInfo{
 				collID:           collID,
 				consistencyLevel: commonpb.ConsistencyLevel_Eventually,
 			}, nil)
@ -3128,7 +3128,7 @@ func (s *MaterializedViewTestSuite) SetupTest() {
 	s.mockMetaCache = NewMockCache(s.T())
 	s.mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(s.colID, nil)
 	s.mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-		&collectionBasicInfo{
+		&collectionInfo{
 			collID:                s.colID,
 			partitionKeyIsolation: true,
 		}, nil)