milvus/internal/indexcoord/index_coord_mock.go
cai.zhang b3f6b67977
Alloc timestamp instead of using start position timestamp for segment index (#19820)
Signed-off-by: cai.zhang <cai.zhang@zilliz.com>

Signed-off-by: cai.zhang <cai.zhang@zilliz.com>
2022-10-18 17:05:26 +08:00

652 lines
22 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"
"math/rand"
"sync"
"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/internal/kv"
"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/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) (*milvuspb.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 commonpb.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 commonpb.StateCode) {
m.CallUpdateStateCode(stateCode)
}
// GetComponentStates gets the component states of the mocked IndexCoord.
func (m *Mock) GetComponentStates(ctx context.Context) (*milvuspb.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 (m *Mock) CheckHealth(ctx context.Context, req *milvuspb.CheckHealthRequest) (*milvuspb.CheckHealthResponse, error) {
return &milvuspb.CheckHealthResponse{
IsHealthy: true,
}, nil
}
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) (*milvuspb.ComponentStates, error) {
return &milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
NodeID: 1,
Role: typeutil.IndexCoordRole,
StateCode: commonpb.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) (*milvuspb.ComponentStates, error)
CallAllocID func(ctx context.Context, req *rootcoordpb.AllocIDRequest) (*rootcoordpb.AllocIDResponse, error)
CallAllocTimestamp func(ctx context.Context, req *rootcoordpb.AllocTimestampRequest) (*rootcoordpb.AllocTimestampResponse, error)
}
func (rcm *RootCoordMock) Init() error {
return rcm.CallInit()
}
func (rcm *RootCoordMock) Start() error {
return rcm.CallStart()
}
func (rcm *RootCoordMock) GetComponentStates(ctx context.Context) (*milvuspb.ComponentStates, error) {
return rcm.CallGetComponentStates(ctx)
}
func (rcm *RootCoordMock) AllocTimestamp(ctx context.Context, req *rootcoordpb.AllocTimestampRequest) (*rootcoordpb.AllocTimestampResponse, error) {
return rcm.CallAllocTimestamp(ctx, req)
}
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) (*milvuspb.ComponentStates, error) {
return &milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
NodeID: 1,
Role: typeutil.IndexCoordRole,
StateCode: commonpb.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
},
CallAllocTimestamp: func(ctx context.Context, req *rootcoordpb.AllocTimestampRequest) (*rootcoordpb.AllocTimestampResponse, error) {
return &rootcoordpb.AllocTimestampResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Timestamp: 1,
Count: req.Count,
}, nil
},
}
}
type DataCoordMock struct {
types.DataCoord
FuncLock sync.RWMutex
CallInit func() error
CallStart func() error
CallGetComponentStates func(ctx context.Context) (*milvuspb.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) (*milvuspb.ComponentStates, error) {
return dcm.CallGetComponentStates(ctx)
}
func (dcm *DataCoordMock) SetFunc(f func()) {
dcm.FuncLock.Lock()
defer dcm.FuncLock.Unlock()
f()
}
func (dcm *DataCoordMock) GetSegmentInfo(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) {
dcm.FuncLock.RLock()
defer dcm.FuncLock.RUnlock()
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) (*milvuspb.ComponentStates, error) {
return &milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
NodeID: 1,
Role: typeutil.IndexCoordRole,
StateCode: commonpb.StateCode_Healthy,
},
SubcomponentStates: nil,
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
}, nil
},
CallGetSegmentInfo: func(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) {
segInfos := make([]*datapb.SegmentInfo, 0)
for _, segID := range req.SegmentIDs {
segInfos = append(segInfos, &datapb.SegmentInfo{
ID: segID,
CollectionID: 100,
PartitionID: 200,
InsertChannel: "",
NumOfRows: 1026,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 0,
LastExpireTime: 0,
StartPosition: nil,
DmlPosition: nil,
Binlogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{
LogPath: "file1",
},
{
LogPath: "file2",
},
},
},
},
Statslogs: nil,
Deltalogs: nil,
CreatedByCompaction: false,
CompactionFrom: nil,
DroppedAt: 0,
})
}
return &datapb.GetSegmentInfoResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Infos: segInfos,
}, 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
},
}
}
type mockETCDKV struct {
kv.MetaKv
save func(string, string) error
load 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)
removeWithPrefix func(key string) error
}
func NewMockEtcdKV() *mockETCDKV {
return &mockETCDKV{
save: func(s string, s2 string) error {
return nil
},
load: func(s string) (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
},
removeWithPrefix: func(key string) error {
return nil
},
}
}
func NewMockEtcdKVWithReal(real kv.MetaKv) *mockETCDKV {
return &mockETCDKV{
save: func(s string, s2 string) error {
return real.Save(s, s2)
},
load: func(s string) (string, error) {
return real.Load(s)
},
remove: func(s string) error {
return real.Remove(s)
},
multiSave: func(m map[string]string) error {
return real.MultiSave(m)
},
loadWithRevisionAndVersions: func(s string) ([]string, []string, []int64, int64, error) {
return real.LoadWithRevisionAndVersions(s)
},
compareVersionAndSwap: func(key string, version int64, target string, opts ...clientv3.OpOption) (bool, error) {
return real.CompareVersionAndSwap(key, version, target, opts...)
},
loadWithPrefix: func(key string) ([]string, []string, error) {
return real.LoadWithPrefix(key)
},
loadWithPrefix2: func(key string) ([]string, []string, []int64, error) {
return real.LoadWithPrefix2(key)
},
loadWithRevision: func(key string) ([]string, []string, int64, error) {
return real.LoadWithRevision(key)
},
removeWithPrefix: func(key string) error {
return real.RemoveWithPrefix(key)
},
}
}
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)
}
func (mk *mockETCDKV) RemoveWithPrefix(key string) error {
return mk.removeWithPrefix(key)
}
func (mk *mockETCDKV) Load(key string) (string, error) {
return mk.load(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(ctx context.Context, prefix string) error {
return cmm.removeWithPrefix(prefix)
}
func (cmm *chunkManagerMock) ListWithPrefix(ctx context.Context, prefix string, recursive bool) ([]string, []time.Time, error) {
return cmm.listWithPrefix(prefix, recursive)
}
func (cmm *chunkManagerMock) Remove(ctx context.Context, key string) error {
return cmm.remove(key)
}