mirror of https://github.com/milvus-io/milvus.git
536 lines
15 KiB
Go
536 lines
15 KiB
Go
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
|
//
|
|
// Licensed 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 querycoord
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"math/rand"
|
|
"net"
|
|
"path"
|
|
"strconv"
|
|
"sync"
|
|
|
|
"google.golang.org/grpc"
|
|
|
|
"github.com/milvus-io/milvus/internal/kv"
|
|
minioKV "github.com/milvus-io/milvus/internal/kv/minio"
|
|
"github.com/milvus-io/milvus/internal/log"
|
|
"github.com/milvus-io/milvus/internal/msgstream"
|
|
"github.com/milvus-io/milvus/internal/proto/commonpb"
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
|
"github.com/milvus-io/milvus/internal/proto/etcdpb"
|
|
"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/proxypb"
|
|
"github.com/milvus-io/milvus/internal/proto/querypb"
|
|
"github.com/milvus-io/milvus/internal/proto/schemapb"
|
|
qn "github.com/milvus-io/milvus/internal/querynode"
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
"github.com/milvus-io/milvus/internal/types"
|
|
"github.com/milvus-io/milvus/internal/util/funcutil"
|
|
"github.com/milvus-io/milvus/internal/util/retry"
|
|
)
|
|
|
|
const (
|
|
defaultCollectionID = UniqueID(2021)
|
|
defaultPartitionID = UniqueID(2021)
|
|
defaultSegmentID = UniqueID(2021)
|
|
defaultShardsNum = 2
|
|
)
|
|
|
|
func genCollectionSchema(collectionID UniqueID, isBinary bool) *schemapb.CollectionSchema {
|
|
var fieldVec schemapb.FieldSchema
|
|
if isBinary {
|
|
fieldVec = schemapb.FieldSchema{
|
|
FieldID: UniqueID(101),
|
|
Name: "vec",
|
|
IsPrimaryKey: false,
|
|
DataType: schemapb.DataType_BinaryVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "dim",
|
|
Value: "128",
|
|
},
|
|
},
|
|
IndexParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "metric_type",
|
|
Value: "JACCARD",
|
|
},
|
|
},
|
|
}
|
|
} else {
|
|
fieldVec = schemapb.FieldSchema{
|
|
FieldID: UniqueID(101),
|
|
Name: "vec",
|
|
IsPrimaryKey: false,
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "dim",
|
|
Value: "16",
|
|
},
|
|
},
|
|
IndexParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: "metric_type",
|
|
Value: "L2",
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
return &schemapb.CollectionSchema{
|
|
AutoID: true,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 0, Name: "row_id", IsPrimaryKey: false, Description: "row_id", DataType: schemapb.DataType_Int64},
|
|
{FieldID: 1, Name: "Ts", IsPrimaryKey: false, Description: "Ts", DataType: schemapb.DataType_Int64},
|
|
{FieldID: 100, Name: "field_age", IsPrimaryKey: false, Description: "int64", DataType: schemapb.DataType_Int64},
|
|
&fieldVec,
|
|
},
|
|
}
|
|
}
|
|
|
|
func genETCDCollectionMeta(collectionID UniqueID, isBinary bool) *etcdpb.CollectionMeta {
|
|
schema := genCollectionSchema(collectionID, isBinary)
|
|
collectionMeta := etcdpb.CollectionMeta{
|
|
ID: collectionID,
|
|
Schema: schema,
|
|
CreateTime: Timestamp(0),
|
|
PartitionIDs: []UniqueID{defaultPartitionID},
|
|
}
|
|
|
|
return &collectionMeta
|
|
}
|
|
|
|
func generateInsertBinLog(collectionID UniqueID, partitionID UniqueID, segmentID UniqueID, keyPrefix string, kv kv.BaseKV) (map[int64]string, error) {
|
|
const (
|
|
msgLength = 1000
|
|
DIM = 16
|
|
)
|
|
|
|
idData := make([]int64, 0)
|
|
for n := 0; n < msgLength; n++ {
|
|
idData = append(idData, int64(n))
|
|
}
|
|
|
|
var timestamps []int64
|
|
for n := 0; n < msgLength; n++ {
|
|
timestamps = append(timestamps, int64(n+1))
|
|
}
|
|
|
|
var fieldAgeData []int64
|
|
for n := 0; n < msgLength; n++ {
|
|
fieldAgeData = append(fieldAgeData, int64(n))
|
|
}
|
|
|
|
fieldVecData := make([]float32, 0)
|
|
for n := 0; n < msgLength; n++ {
|
|
for i := 0; i < DIM; i++ {
|
|
fieldVecData = append(fieldVecData, float32(n*i)*0.1)
|
|
}
|
|
}
|
|
|
|
insertData := &storage.InsertData{
|
|
Data: map[int64]storage.FieldData{
|
|
0: &storage.Int64FieldData{
|
|
NumRows: []int64{msgLength},
|
|
Data: idData,
|
|
},
|
|
1: &storage.Int64FieldData{
|
|
NumRows: []int64{msgLength},
|
|
Data: timestamps,
|
|
},
|
|
100: &storage.Int64FieldData{
|
|
NumRows: []int64{msgLength},
|
|
Data: fieldAgeData,
|
|
},
|
|
101: &storage.FloatVectorFieldData{
|
|
NumRows: []int64{msgLength},
|
|
Data: fieldVecData,
|
|
Dim: DIM,
|
|
},
|
|
},
|
|
}
|
|
|
|
// buffer data to binLogs
|
|
collMeta := genETCDCollectionMeta(collectionID, false)
|
|
inCodec := storage.NewInsertCodec(collMeta)
|
|
binLogs, _, err := inCodec.Serialize(partitionID, segmentID, insertData)
|
|
|
|
if err != nil {
|
|
log.Debug("insert data serialize error")
|
|
return nil, err
|
|
}
|
|
|
|
// binLogs -> minIO/S3
|
|
segIDStr := strconv.FormatInt(segmentID, 10)
|
|
keyPrefix = path.Join(keyPrefix, segIDStr)
|
|
|
|
fieldID2Paths := make(map[int64]string)
|
|
for _, blob := range binLogs {
|
|
uid := rand.Int63n(100000000)
|
|
path := path.Join(keyPrefix, blob.Key, strconv.FormatInt(uid, 10))
|
|
err = kv.Save(path, string(blob.Value[:]))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
fieldID, err := strconv.Atoi(blob.Key)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
fieldID2Paths[int64(fieldID)] = path
|
|
}
|
|
|
|
return fieldID2Paths, nil
|
|
}
|
|
|
|
type rootCoordMock struct {
|
|
types.RootCoord
|
|
CollectionIDs []UniqueID
|
|
Col2partition map[UniqueID][]UniqueID
|
|
sync.RWMutex
|
|
}
|
|
|
|
func newRootCoordMock() *rootCoordMock {
|
|
collectionIDs := make([]UniqueID, 0)
|
|
col2partition := make(map[UniqueID][]UniqueID)
|
|
|
|
return &rootCoordMock{
|
|
CollectionIDs: collectionIDs,
|
|
Col2partition: col2partition,
|
|
}
|
|
}
|
|
|
|
func (rc *rootCoordMock) createCollection(collectionID UniqueID) {
|
|
rc.Lock()
|
|
defer rc.Unlock()
|
|
|
|
if _, ok := rc.Col2partition[collectionID]; !ok {
|
|
rc.CollectionIDs = append(rc.CollectionIDs, collectionID)
|
|
rc.Col2partition[collectionID] = make([]UniqueID, 0)
|
|
}
|
|
}
|
|
|
|
func (rc *rootCoordMock) createPartition(collectionID UniqueID, partitionID UniqueID) error {
|
|
rc.Lock()
|
|
defer rc.Unlock()
|
|
|
|
if _, ok := rc.Col2partition[collectionID]; ok {
|
|
rc.Col2partition[collectionID] = append(rc.Col2partition[collectionID], partitionID)
|
|
}
|
|
|
|
return errors.New("collection not exist")
|
|
}
|
|
|
|
func (rc *rootCoordMock) ShowPartitions(ctx context.Context, in *milvuspb.ShowPartitionsRequest) (*milvuspb.ShowPartitionsResponse, error) {
|
|
collectionID := in.CollectionID
|
|
status := &commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_Success,
|
|
}
|
|
if partitionIDs, ok := rc.Col2partition[collectionID]; ok {
|
|
response := &milvuspb.ShowPartitionsResponse{
|
|
Status: status,
|
|
PartitionIDs: partitionIDs,
|
|
}
|
|
|
|
return response, nil
|
|
}
|
|
|
|
rc.createCollection(collectionID)
|
|
rc.createPartition(collectionID, defaultPartitionID)
|
|
|
|
return &milvuspb.ShowPartitionsResponse{
|
|
Status: status,
|
|
PartitionIDs: rc.Col2partition[collectionID],
|
|
}, nil
|
|
}
|
|
|
|
func (rc *rootCoordMock) ReleaseDQLMessageStream(ctx context.Context, in *proxypb.ReleaseDQLMessageStreamRequest) (*commonpb.Status, error) {
|
|
return &commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_Success,
|
|
}, nil
|
|
}
|
|
|
|
func (rc *rootCoordMock) DescribeSegment(ctx context.Context, req *milvuspb.DescribeSegmentRequest) (*milvuspb.DescribeSegmentResponse, error) {
|
|
return nil, errors.New("describeSegment fail")
|
|
}
|
|
|
|
type dataCoordMock struct {
|
|
types.DataCoord
|
|
minioKV kv.BaseKV
|
|
collections []UniqueID
|
|
col2DmChannels map[UniqueID][]*datapb.VchannelInfo
|
|
partitionID2Segment map[UniqueID]UniqueID
|
|
Segment2Binlog map[UniqueID][]*datapb.SegmentBinlogs
|
|
assignedSegmentID UniqueID
|
|
}
|
|
|
|
func newDataCoordMock(ctx context.Context) (*dataCoordMock, error) {
|
|
collectionIDs := make([]UniqueID, 0)
|
|
col2DmChannels := make(map[UniqueID][]*datapb.VchannelInfo)
|
|
partitionID2Segment := make(map[UniqueID]UniqueID)
|
|
segment2Binglog := make(map[UniqueID][]*datapb.SegmentBinlogs)
|
|
|
|
// create minio client
|
|
option := &minioKV.Option{
|
|
Address: Params.MinioEndPoint,
|
|
AccessKeyID: Params.MinioAccessKeyID,
|
|
SecretAccessKeyID: Params.MinioSecretAccessKey,
|
|
UseSSL: Params.MinioUseSSLStr,
|
|
BucketName: Params.MinioBucketName,
|
|
CreateBucket: true,
|
|
}
|
|
kv, err := minioKV.NewMinIOKV(ctx, option)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &dataCoordMock{
|
|
minioKV: kv,
|
|
collections: collectionIDs,
|
|
col2DmChannels: col2DmChannels,
|
|
partitionID2Segment: partitionID2Segment,
|
|
Segment2Binlog: segment2Binglog,
|
|
assignedSegmentID: defaultSegmentID,
|
|
}, nil
|
|
}
|
|
|
|
func (data *dataCoordMock) GetRecoveryInfo(ctx context.Context, req *datapb.GetRecoveryInfoRequest) (*datapb.GetRecoveryInfoResponse, error) {
|
|
collectionID := req.CollectionID
|
|
partitionID := req.PartitionID
|
|
|
|
if _, ok := data.col2DmChannels[collectionID]; !ok {
|
|
segmentID := data.assignedSegmentID
|
|
data.partitionID2Segment[partitionID] = segmentID
|
|
fieldID2Paths, err := generateInsertBinLog(collectionID, partitionID, segmentID, "queryCoorf-mockDataCoord", data.minioKV)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
fieldBinlogs := make([]*datapb.FieldBinlog, 0)
|
|
for fieldID, path := range fieldID2Paths {
|
|
fieldBinlog := &datapb.FieldBinlog{
|
|
FieldID: fieldID,
|
|
Binlogs: []string{path},
|
|
}
|
|
fieldBinlogs = append(fieldBinlogs, fieldBinlog)
|
|
}
|
|
data.Segment2Binlog[segmentID] = make([]*datapb.SegmentBinlogs, 0)
|
|
segmentBinlog := &datapb.SegmentBinlogs{
|
|
SegmentID: segmentID,
|
|
FieldBinlogs: fieldBinlogs,
|
|
}
|
|
data.Segment2Binlog[segmentID] = append(data.Segment2Binlog[segmentID], segmentBinlog)
|
|
|
|
channelInfos := make([]*datapb.VchannelInfo, 0)
|
|
data.collections = append(data.collections, collectionID)
|
|
collectionName := funcutil.RandomString(8)
|
|
for i := 0; i < defaultShardsNum; i++ {
|
|
vChannel := fmt.Sprintf("%s_%d_%d_v", collectionName, collectionID, i)
|
|
channelInfo := &datapb.VchannelInfo{
|
|
CollectionID: collectionID,
|
|
ChannelName: vChannel,
|
|
SeekPosition: &internalpb.MsgPosition{
|
|
ChannelName: vChannel,
|
|
},
|
|
FlushedSegments: []int64{segmentID},
|
|
}
|
|
channelInfos = append(channelInfos, channelInfo)
|
|
}
|
|
data.col2DmChannels[collectionID] = channelInfos
|
|
}
|
|
|
|
segmentID := data.partitionID2Segment[partitionID]
|
|
return &datapb.GetRecoveryInfoResponse{
|
|
Status: &commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_Success,
|
|
},
|
|
Channels: data.col2DmChannels[collectionID],
|
|
Binlogs: data.Segment2Binlog[segmentID],
|
|
}, nil
|
|
}
|
|
|
|
type indexCoordMock struct {
|
|
types.IndexCoord
|
|
}
|
|
|
|
func newIndexCoordMock() *indexCoordMock {
|
|
return &indexCoordMock{}
|
|
}
|
|
|
|
func (c *indexCoordMock) GetIndexFilePaths(ctx context.Context, req *indexpb.GetIndexFilePathsRequest) (*indexpb.GetIndexFilePathsResponse, error) {
|
|
return nil, errors.New("get index file path fail")
|
|
}
|
|
|
|
type queryNodeServerMock struct {
|
|
querypb.QueryNodeServer
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
queryNode *qn.QueryNode
|
|
grpcErrChan chan error
|
|
grpcServer *grpc.Server
|
|
|
|
addQueryChannels func() (*commonpb.Status, error)
|
|
watchDmChannels func() (*commonpb.Status, error)
|
|
loadSegment func() (*commonpb.Status, error)
|
|
releaseCollection func() (*commonpb.Status, error)
|
|
releasePartition func() (*commonpb.Status, error)
|
|
releaseSegment func() (*commonpb.Status, error)
|
|
}
|
|
|
|
func newQueryNodeServerMock(ctx context.Context) *queryNodeServerMock {
|
|
ctx1, cancel := context.WithCancel(ctx)
|
|
factory := msgstream.NewPmsFactory()
|
|
|
|
return &queryNodeServerMock{
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
queryNode: qn.NewQueryNode(ctx1, factory),
|
|
grpcErrChan: make(chan error),
|
|
|
|
addQueryChannels: returnSuccessResult,
|
|
watchDmChannels: returnSuccessResult,
|
|
loadSegment: returnSuccessResult,
|
|
releaseCollection: returnSuccessResult,
|
|
releasePartition: returnSuccessResult,
|
|
releaseSegment: returnSuccessResult,
|
|
}
|
|
}
|
|
|
|
func (qs *queryNodeServerMock) init() error {
|
|
qn.Params.Init()
|
|
qn.Params.MetaRootPath = Params.MetaRootPath
|
|
qn.Params.QueryNodeIP = funcutil.GetLocalIP()
|
|
grpcPort := Params.Port
|
|
|
|
go func() {
|
|
var lis net.Listener
|
|
var err error
|
|
|
|
err = retry.Do(qs.ctx, func() error {
|
|
addr := ":" + strconv.Itoa(grpcPort)
|
|
lis, err = net.Listen("tcp", addr)
|
|
if err == nil {
|
|
qn.Params.QueryNodePort = int64(lis.Addr().(*net.TCPAddr).Port)
|
|
} else {
|
|
// set port=0 to get next available port
|
|
grpcPort = 0
|
|
}
|
|
return err
|
|
}, retry.Attempts(10))
|
|
|
|
if err != nil {
|
|
log.Error(err.Error())
|
|
}
|
|
|
|
qs.grpcServer = grpc.NewServer()
|
|
querypb.RegisterQueryNodeServer(qs.grpcServer, qs)
|
|
if err = qs.grpcServer.Serve(lis); err != nil {
|
|
log.Error(err.Error())
|
|
}
|
|
}()
|
|
|
|
rootCoord := newRootCoordMock()
|
|
indexCoord := newIndexCoordMock()
|
|
qs.queryNode.SetRootCoord(rootCoord)
|
|
qs.queryNode.SetIndexCoord(indexCoord)
|
|
err := qs.queryNode.Init()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if err = qs.queryNode.Register(); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (qs *queryNodeServerMock) start() error {
|
|
return qs.queryNode.Start()
|
|
}
|
|
|
|
func (qs *queryNodeServerMock) stop() error {
|
|
qs.cancel()
|
|
if qs.grpcServer != nil {
|
|
qs.grpcServer.GracefulStop()
|
|
}
|
|
|
|
err := qs.queryNode.Stop()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (qs *queryNodeServerMock) run() error {
|
|
if err := qs.init(); err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := qs.start(); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (qs *queryNodeServerMock) AddQueryChannel(ctx context.Context, req *querypb.AddQueryChannelRequest) (*commonpb.Status, error) {
|
|
return qs.addQueryChannels()
|
|
}
|
|
|
|
func (qs *queryNodeServerMock) WatchDmChannels(ctx context.Context, req *querypb.WatchDmChannelsRequest) (*commonpb.Status, error) {
|
|
return qs.watchDmChannels()
|
|
}
|
|
|
|
func (qs *queryNodeServerMock) LoadSegments(ctx context.Context, req *querypb.LoadSegmentsRequest) (*commonpb.Status, error) {
|
|
return qs.loadSegment()
|
|
}
|
|
|
|
func (qs *queryNodeServerMock) ReleaseCollection(ctx context.Context, req *querypb.ReleaseCollectionRequest) (*commonpb.Status, error) {
|
|
return qs.releaseCollection()
|
|
}
|
|
|
|
func (qs *queryNodeServerMock) ReleasePartitions(ctx context.Context, req *querypb.ReleasePartitionsRequest) (*commonpb.Status, error) {
|
|
return qs.releasePartition()
|
|
}
|
|
|
|
func (qs *queryNodeServerMock) ReleaseSegments(ctx context.Context, req *querypb.ReleaseSegmentsRequest) (*commonpb.Status, error) {
|
|
return qs.releaseSegment()
|
|
}
|
|
|
|
func startQueryNodeServer(ctx context.Context) (*queryNodeServerMock, error) {
|
|
node := newQueryNodeServerMock(ctx)
|
|
err := node.run()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return node, nil
|
|
}
|
|
|
|
func returnSuccessResult() (*commonpb.Status, error) {
|
|
return &commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_Success,
|
|
}, nil
|
|
}
|
|
|
|
func returnFailedResult() (*commonpb.Status, error) {
|
|
return &commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_UnexpectedError,
|
|
}, errors.New("query node do task failed")
|
|
}
|