milvus/internal/indexcoord/index_coord_mock.go
congqixia c91bb0b015
Unify rootPath in configs and ChunkManager (#18808)
Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
2022-08-25 19:32:53 +08:00

565 lines
19 KiB
Go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package indexcoord
import (
"context"
"errors"
"math/rand"
"time"
clientv3 "go.etcd.io/etcd/client/v3"
"github.com/milvus-io/milvus/internal/kv"
"github.com/milvus-io/milvus/internal/proto/commonpb"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/proto/milvuspb"
"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/util/typeutil"
)
// Mock is an alternative to IndexCoord, it will return specific results based on specific parameters.
type Mock struct {
types.IndexCoord
CallInit func() error
CallStart func() error
CallStop func() error
CallGetComponentStates func(ctx context.Context) (*internalpb.ComponentStates, error)
CallGetStatisticsChannel func(ctx context.Context) (*milvuspb.StringResponse, error)
CallRegister func() error
CallSetEtcdClient func(etcdClient *clientv3.Client)
CallSetDataCoord func(dataCoord types.DataCoord) error
CallSetRootCoord func(rootCoord types.RootCoord) error
CallUpdateStateCode func(stateCode internalpb.StateCode)
CallCreateIndex func(ctx context.Context, req *indexpb.CreateIndexRequest) (*commonpb.Status, error)
CallGetIndexState func(ctx context.Context, req *indexpb.GetIndexStateRequest) (*indexpb.GetIndexStateResponse, error)
CallGetSegmentIndexState func(ctx context.Context, req *indexpb.GetSegmentIndexStateRequest) (*indexpb.GetSegmentIndexStateResponse, error)
CallGetIndexInfos func(ctx context.Context, req *indexpb.GetIndexInfoRequest) (*indexpb.GetIndexInfoResponse, error)
CallDescribeIndex func(ctx context.Context, req *indexpb.DescribeIndexRequest) (*indexpb.DescribeIndexResponse, error)
CallGetIndexBuildProgress func(ctx context.Context, req *indexpb.GetIndexBuildProgressRequest) (*indexpb.GetIndexBuildProgressResponse, error)
CallDropIndex func(ctx context.Context, req *indexpb.DropIndexRequest) (*commonpb.Status, error)
CallShowConfigurations func(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error)
CallGetMetrics func(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)
}
// Init initializes the Mock of IndexCoord. When param `Failure` is true, it will return an error.
func (m *Mock) Init() error {
return m.CallInit()
}
// Start starts the Mock of IndexCoord. When param `Failure` is true, it will return an error.
func (m *Mock) Start() error {
return m.CallStart()
}
// Stop stops the Mock of IndexCoord. When param `Failure` is true, it will return an error.
func (m *Mock) Stop() error {
return m.CallStop()
}
// Register registers an IndexCoord role in ETCD, if Param `Failure` is true, it will return an error.
func (m *Mock) Register() error {
return m.CallRegister()
}
func (m *Mock) SetEtcdClient(client *clientv3.Client) {
m.CallSetEtcdClient(client)
}
func (m *Mock) SetDataCoord(dataCoord types.DataCoord) error {
return m.CallSetDataCoord(dataCoord)
}
func (m *Mock) SetRootCoord(rootCoord types.RootCoord) error {
return m.CallSetRootCoord(rootCoord)
}
func (m *Mock) UpdateStateCode(stateCode internalpb.StateCode) {
m.CallUpdateStateCode(stateCode)
}
// GetComponentStates gets the component states of the mocked IndexCoord.
func (m *Mock) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
return m.CallGetComponentStates(ctx)
}
// GetStatisticsChannel gets the statistics channel of the mocked IndexCoord, if Param `Failure` is true, it will return an error.
func (m *Mock) GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
return m.CallGetStatisticsChannel(ctx)
}
func (m *Mock) CreateIndex(ctx context.Context, req *indexpb.CreateIndexRequest) (*commonpb.Status, error) {
return m.CallCreateIndex(ctx, req)
}
func (m *Mock) GetIndexState(ctx context.Context, req *indexpb.GetIndexStateRequest) (*indexpb.GetIndexStateResponse, error) {
return m.CallGetIndexState(ctx, req)
}
func (m *Mock) GetSegmentIndexState(ctx context.Context, req *indexpb.GetSegmentIndexStateRequest) (*indexpb.GetSegmentIndexStateResponse, error) {
return m.CallGetSegmentIndexState(ctx, req)
}
func (m *Mock) GetIndexInfos(ctx context.Context, req *indexpb.GetIndexInfoRequest) (*indexpb.GetIndexInfoResponse, error) {
return m.CallGetIndexInfos(ctx, req)
}
func (m *Mock) DescribeIndex(ctx context.Context, req *indexpb.DescribeIndexRequest) (*indexpb.DescribeIndexResponse, error) {
return m.CallDescribeIndex(ctx, req)
}
func (m *Mock) GetIndexBuildProgress(ctx context.Context, req *indexpb.GetIndexBuildProgressRequest) (*indexpb.GetIndexBuildProgressResponse, error) {
return m.CallGetIndexBuildProgress(ctx, req)
}
func (m *Mock) DropIndex(ctx context.Context, req *indexpb.DropIndexRequest) (*commonpb.Status, error) {
return m.CallDropIndex(ctx, req)
}
func (m *Mock) ShowConfigurations(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) {
return m.CallShowConfigurations(ctx, req)
}
func (m *Mock) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
return m.CallGetMetrics(ctx, req)
}
func NewIndexCoordMock() *Mock {
return &Mock{
CallInit: func() error {
return nil
},
CallStart: func() error {
return nil
},
CallRegister: func() error {
return nil
},
CallStop: func() error {
return nil
},
CallSetEtcdClient: func(etcdClient *clientv3.Client) {
},
CallSetDataCoord: func(dataCoord types.DataCoord) error {
return nil
},
CallSetRootCoord: func(rootCoord types.RootCoord) error {
return nil
},
CallGetComponentStates: func(ctx context.Context) (*internalpb.ComponentStates, error) {
return &internalpb.ComponentStates{
State: &internalpb.ComponentInfo{
NodeID: 1,
Role: typeutil.IndexCoordRole,
StateCode: internalpb.StateCode_Healthy,
},
SubcomponentStates: nil,
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
}, nil
},
CallGetStatisticsChannel: func(ctx context.Context) (*milvuspb.StringResponse, error) {
return &milvuspb.StringResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
}, nil
},
CallCreateIndex: func(ctx context.Context, req *indexpb.CreateIndexRequest) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
},
CallGetIndexState: func(ctx context.Context, req *indexpb.GetIndexStateRequest) (*indexpb.GetIndexStateResponse, error) {
return &indexpb.GetIndexStateResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
State: commonpb.IndexState_Finished,
}, nil
},
CallGetSegmentIndexState: func(ctx context.Context, req *indexpb.GetSegmentIndexStateRequest) (*indexpb.GetSegmentIndexStateResponse, error) {
segmentStates := make([]*indexpb.SegmentIndexState, 0)
for _, segID := range req.SegmentIDs {
segmentStates = append(segmentStates, &indexpb.SegmentIndexState{
SegmentID: segID,
State: commonpb.IndexState_Finished,
})
}
return &indexpb.GetSegmentIndexStateResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
States: segmentStates,
}, nil
},
CallGetIndexInfos: func(ctx context.Context, req *indexpb.GetIndexInfoRequest) (*indexpb.GetIndexInfoResponse, error) {
segmentInfos := make(map[int64]*indexpb.SegmentInfo)
filePaths := make([]*indexpb.IndexFilePathInfo, 0)
for _, segID := range req.SegmentIDs {
filePaths = append(filePaths, &indexpb.IndexFilePathInfo{
SegmentID: segID,
IndexName: "default",
IndexFilePaths: []string{"file1", "file2"},
})
segmentInfos[segID] = &indexpb.SegmentInfo{
CollectionID: req.CollectionID,
SegmentID: segID,
EnableIndex: true,
IndexInfos: filePaths,
}
}
return &indexpb.GetIndexInfoResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
SegmentInfo: segmentInfos,
}, nil
},
CallDescribeIndex: func(ctx context.Context, req *indexpb.DescribeIndexRequest) (*indexpb.DescribeIndexResponse, error) {
return &indexpb.DescribeIndexResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
IndexInfos: []*indexpb.IndexInfo{
{
CollectionID: 1,
FieldID: 0,
IndexName: "default",
IndexID: 0,
TypeParams: nil,
IndexParams: nil,
},
},
}, nil
},
CallGetIndexBuildProgress: func(ctx context.Context, req *indexpb.GetIndexBuildProgressRequest) (*indexpb.GetIndexBuildProgressResponse, error) {
return &indexpb.GetIndexBuildProgressResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
IndexedRows: 10240,
TotalRows: 10240,
}, nil
},
CallDropIndex: func(ctx context.Context, req *indexpb.DropIndexRequest) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
},
CallShowConfigurations: func(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) {
return &internalpb.ShowConfigurationsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
}, nil
},
CallGetMetrics: func(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
return &milvuspb.GetMetricsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
ComponentName: typeutil.IndexCoordRole,
}, nil
},
}
}
type RootCoordMock struct {
types.RootCoord
CallInit func() error
CallStart func() error
CallGetComponentStates func(ctx context.Context) (*internalpb.ComponentStates, error)
CallAllocID func(ctx context.Context, req *rootcoordpb.AllocIDRequest) (*rootcoordpb.AllocIDResponse, error)
}
func (rcm *RootCoordMock) Init() error {
return rcm.CallInit()
}
func (rcm *RootCoordMock) Start() error {
return rcm.CallStart()
}
func (rcm *RootCoordMock) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
return rcm.CallGetComponentStates(ctx)
}
func (rcm *RootCoordMock) AllocID(ctx context.Context, req *rootcoordpb.AllocIDRequest) (*rootcoordpb.AllocIDResponse, error) {
return rcm.CallAllocID(ctx, req)
}
func NewRootCoordMock() *RootCoordMock {
return &RootCoordMock{
CallInit: func() error {
return nil
},
CallStart: func() error {
return nil
},
CallGetComponentStates: func(ctx context.Context) (*internalpb.ComponentStates, error) {
return &internalpb.ComponentStates{
State: &internalpb.ComponentInfo{
NodeID: 1,
Role: typeutil.IndexCoordRole,
StateCode: internalpb.StateCode_Healthy,
},
SubcomponentStates: nil,
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
}, nil
},
CallAllocID: func(ctx context.Context, req *rootcoordpb.AllocIDRequest) (*rootcoordpb.AllocIDResponse, error) {
return &rootcoordpb.AllocIDResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
ID: rand.Int63(),
Count: req.Count,
}, nil
},
}
}
type DataCoordMock struct {
types.DataCoord
CallInit func() error
CallStart func() error
CallGetComponentStates func(ctx context.Context) (*internalpb.ComponentStates, error)
CallGetSegmentInfo func(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error)
CallGetFlushedSegment func(ctx context.Context, req *datapb.GetFlushedSegmentsRequest) (*datapb.GetFlushedSegmentsResponse, error)
CallAcquireSegmentLock func(ctx context.Context, req *datapb.AcquireSegmentLockRequest) (*commonpb.Status, error)
CallReleaseSegmentLock func(ctx context.Context, req *datapb.ReleaseSegmentLockRequest) (*commonpb.Status, error)
}
func (dcm *DataCoordMock) Init() error {
return dcm.CallInit()
}
func (dcm *DataCoordMock) Start() error {
return dcm.CallStart()
}
func (dcm *DataCoordMock) GetComponentStates(ctx context.Context) (*internalpb.ComponentStates, error) {
return dcm.CallGetComponentStates(ctx)
}
func (dcm *DataCoordMock) GetSegmentInfo(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) {
return dcm.CallGetSegmentInfo(ctx, req)
}
func (dcm *DataCoordMock) AcquireSegmentLock(ctx context.Context, req *datapb.AcquireSegmentLockRequest) (*commonpb.Status, error) {
return dcm.CallAcquireSegmentLock(ctx, req)
}
func (dcm *DataCoordMock) ReleaseSegmentLock(ctx context.Context, req *datapb.ReleaseSegmentLockRequest) (*commonpb.Status, error) {
return dcm.CallReleaseSegmentLock(ctx, req)
}
func (dcm *DataCoordMock) GetFlushedSegments(ctx context.Context, req *datapb.GetFlushedSegmentsRequest) (*datapb.GetFlushedSegmentsResponse, error) {
return dcm.CallGetFlushedSegment(ctx, req)
}
func NewDataCoordMock() *DataCoordMock {
return &DataCoordMock{
CallInit: func() error {
return nil
},
CallStart: func() error {
return nil
},
CallGetComponentStates: func(ctx context.Context) (*internalpb.ComponentStates, error) {
return &internalpb.ComponentStates{
State: &internalpb.ComponentInfo{
NodeID: 1,
Role: typeutil.IndexCoordRole,
StateCode: internalpb.StateCode_Healthy,
},
SubcomponentStates: nil,
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
}, nil
},
CallGetSegmentInfo: func(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) {
return &datapb.GetSegmentInfoResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
}, nil
},
CallGetFlushedSegment: func(ctx context.Context, req *datapb.GetFlushedSegmentsRequest) (*datapb.GetFlushedSegmentsResponse, error) {
return &datapb.GetFlushedSegmentsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
}, nil
},
CallAcquireSegmentLock: func(ctx context.Context, req *datapb.AcquireSegmentLockRequest) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
},
CallReleaseSegmentLock: func(ctx context.Context, req *datapb.ReleaseSegmentLockRequest) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
}, nil
},
}
}
// ChunkManagerMock is mock
// deprecated
type ChunkManagerMock struct {
storage.ChunkManager
Fail bool
Err bool
}
func (cmm *ChunkManagerMock) Exist(path string) (bool, error) {
if cmm.Err {
return false, errors.New("path not exist")
}
if cmm.Fail {
return false, nil
}
return true, nil
}
func (cmm *ChunkManagerMock) RemoveWithPrefix(prefix string) error {
if cmm.Err {
return errors.New("error occurred")
}
if cmm.Fail {
return nil
}
return nil
}
type mockETCDKV struct {
kv.MetaKv
save func(string, string) error
remove func(string) error
multiSave func(map[string]string) error
watchWithRevision func(string, int64) clientv3.WatchChan
loadWithRevisionAndVersions func(string) ([]string, []string, []int64, int64, error)
compareVersionAndSwap func(key string, version int64, target string, opts ...clientv3.OpOption) (bool, error)
loadWithPrefix2 func(key string) ([]string, []string, []int64, error)
loadWithPrefix func(key string) ([]string, []string, error)
loadWithRevision func(key string) ([]string, []string, int64, error)
}
func NewMockEtcdKV() *mockETCDKV {
return &mockETCDKV{
save: func(s string, s2 string) error {
return nil
},
remove: func(s string) error {
return nil
},
multiSave: func(m map[string]string) error {
return nil
},
loadWithRevisionAndVersions: func(s string) ([]string, []string, []int64, int64, error) {
return []string{}, []string{}, []int64{}, 0, nil
},
compareVersionAndSwap: func(key string, version int64, target string, opts ...clientv3.OpOption) (bool, error) {
return true, nil
},
loadWithPrefix2: func(key string) ([]string, []string, []int64, error) {
return []string{}, []string{}, []int64{}, nil
},
loadWithRevision: func(key string) ([]string, []string, int64, error) {
return []string{}, []string{}, 0, nil
},
}
}
func (mk *mockETCDKV) Save(key string, value string) error {
return mk.save(key, value)
}
func (mk *mockETCDKV) Remove(key string) error {
return mk.remove(key)
}
func (mk *mockETCDKV) MultiSave(kvs map[string]string) error {
return mk.multiSave(kvs)
}
func (mk *mockETCDKV) LoadWithRevisionAndVersions(prefix string) ([]string, []string, []int64, int64, error) {
return mk.loadWithRevisionAndVersions(prefix)
}
func (mk *mockETCDKV) CompareVersionAndSwap(key string, version int64, target string, opts ...clientv3.OpOption) (bool, error) {
return mk.compareVersionAndSwap(key, version, target, opts...)
}
func (mk *mockETCDKV) LoadWithPrefix(key string) ([]string, []string, error) {
return mk.loadWithPrefix(key)
}
func (mk *mockETCDKV) LoadWithPrefix2(key string) ([]string, []string, []int64, error) {
return mk.loadWithPrefix2(key)
}
func (mk *mockETCDKV) WatchWithRevision(key string, revision int64) clientv3.WatchChan {
return mk.watchWithRevision(key, revision)
}
func (mk *mockETCDKV) LoadWithRevision(key string) ([]string, []string, int64, error) {
return mk.loadWithRevision(key)
}
type chunkManagerMock struct {
storage.ChunkManager
removeWithPrefix func(string) error
listWithPrefix func(string, bool) ([]string, []time.Time, error)
remove func(string) error
}
func (cmm *chunkManagerMock) RootPath() string {
return ""
}
func (cmm *chunkManagerMock) RemoveWithPrefix(prefix string) error {
return cmm.removeWithPrefix(prefix)
}
func (cmm *chunkManagerMock) ListWithPrefix(prefix string, recursive bool) ([]string, []time.Time, error) {
return cmm.listWithPrefix(prefix, recursive)
}
func (cmm *chunkManagerMock) Remove(key string) error {
return cmm.remove(key)
}