2022-10-11 11:39:22 +08:00
|
|
|
// 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.
|
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
package mocks
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
"net"
|
|
|
|
"sync"
|
|
|
|
"testing"
|
|
|
|
"time"
|
|
|
|
|
2023-06-25 17:20:43 +08:00
|
|
|
"github.com/stretchr/testify/mock"
|
|
|
|
clientv3 "go.etcd.io/etcd/client/v3"
|
|
|
|
"go.uber.org/zap"
|
|
|
|
"google.golang.org/grpc"
|
|
|
|
|
2023-06-09 01:28:37 +08:00
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
|
2023-04-06 19:14:32 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/proto/querypb"
|
2022-09-15 18:48:32 +08:00
|
|
|
. "github.com/milvus-io/milvus/internal/querycoordv2/params"
|
|
|
|
"github.com/milvus-io/milvus/internal/util/sessionutil"
|
2023-04-06 19:14:32 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/log"
|
2023-09-04 09:57:09 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/util/merr"
|
2023-04-06 19:14:32 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/util/typeutil"
|
2022-09-15 18:48:32 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
type MockQueryNode struct {
|
|
|
|
*MockQueryNodeServer
|
|
|
|
|
|
|
|
ID int64
|
|
|
|
addr string
|
|
|
|
ctx context.Context
|
|
|
|
cancel context.CancelFunc
|
|
|
|
session *sessionutil.Session
|
|
|
|
server *grpc.Server
|
|
|
|
|
|
|
|
rwmutex sync.RWMutex
|
|
|
|
channels map[int64][]string
|
|
|
|
channelVersion map[string]int64
|
|
|
|
segments map[int64]map[string][]int64
|
|
|
|
segmentVersion map[int64]int64
|
|
|
|
}
|
|
|
|
|
2022-11-04 14:25:38 +08:00
|
|
|
func NewMockQueryNode(t *testing.T, etcdCli *clientv3.Client, nodeID int64) *MockQueryNode {
|
2022-09-15 18:48:32 +08:00
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
|
node := &MockQueryNode{
|
|
|
|
MockQueryNodeServer: NewMockQueryNodeServer(t),
|
|
|
|
ctx: ctx,
|
|
|
|
cancel: cancel,
|
2023-10-27 07:36:12 +08:00
|
|
|
session: sessionutil.NewSessionWithEtcd(ctx, Params.EtcdCfg.MetaRootPath.GetValue(), etcdCli),
|
2022-09-15 18:48:32 +08:00
|
|
|
channels: make(map[int64][]string),
|
|
|
|
segments: make(map[int64]map[string][]int64),
|
2022-11-04 14:25:38 +08:00
|
|
|
ID: nodeID,
|
2022-09-15 18:48:32 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
return node
|
|
|
|
}
|
|
|
|
|
|
|
|
func (node *MockQueryNode) Start() error {
|
|
|
|
// Start gRPC server
|
|
|
|
lis, err := net.Listen("tcp", "localhost:0")
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
node.addr = lis.Addr().String()
|
|
|
|
node.server = grpc.NewServer()
|
|
|
|
querypb.RegisterQueryNodeServer(node.server, node)
|
|
|
|
go func() {
|
|
|
|
err = node.server.Serve(lis)
|
|
|
|
}()
|
|
|
|
|
2023-10-11 21:01:35 +08:00
|
|
|
successStatus := merr.Success()
|
2022-09-15 18:48:32 +08:00
|
|
|
node.EXPECT().GetDataDistribution(mock.Anything, mock.Anything).Return(&querypb.GetDataDistributionResponse{
|
|
|
|
Status: successStatus,
|
|
|
|
NodeID: node.ID,
|
|
|
|
Channels: node.getAllChannels(),
|
|
|
|
Segments: node.getAllSegments(),
|
|
|
|
}, nil).Maybe()
|
|
|
|
|
|
|
|
node.EXPECT().WatchDmChannels(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *querypb.WatchDmChannelsRequest) {
|
|
|
|
node.rwmutex.Lock()
|
|
|
|
defer node.rwmutex.Unlock()
|
|
|
|
|
|
|
|
node.channels[req.GetCollectionID()] = append(node.channels[req.GetCollectionID()],
|
|
|
|
req.GetInfos()[0].GetChannelName())
|
|
|
|
}).Return(successStatus, nil).Maybe()
|
|
|
|
node.EXPECT().LoadSegments(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *querypb.LoadSegmentsRequest) {
|
|
|
|
node.rwmutex.Lock()
|
|
|
|
defer node.rwmutex.Unlock()
|
|
|
|
|
|
|
|
shardSegments, ok := node.segments[req.GetCollectionID()]
|
|
|
|
if !ok {
|
|
|
|
shardSegments = make(map[string][]int64)
|
|
|
|
node.segments[req.GetCollectionID()] = shardSegments
|
|
|
|
}
|
|
|
|
segment := req.GetInfos()[0]
|
|
|
|
shardSegments[segment.GetInsertChannel()] = append(shardSegments[segment.GetInsertChannel()],
|
|
|
|
segment.GetSegmentID())
|
|
|
|
node.segmentVersion[segment.GetSegmentID()] = req.GetVersion()
|
|
|
|
}).Return(successStatus, nil).Maybe()
|
2023-03-07 09:51:50 +08:00
|
|
|
node.EXPECT().GetComponentStates(mock.Anything, mock.AnythingOfType("*milvuspb.GetComponentStatesRequest")).
|
|
|
|
Call.Return(func(context.Context, *milvuspb.GetComponentStatesRequest) *milvuspb.ComponentStates {
|
|
|
|
select {
|
|
|
|
case <-node.ctx.Done():
|
|
|
|
return nil
|
|
|
|
default:
|
2023-09-06 17:43:14 +08:00
|
|
|
return &milvuspb.ComponentStates{
|
|
|
|
Status: successStatus,
|
|
|
|
}
|
2023-03-07 09:51:50 +08:00
|
|
|
}
|
|
|
|
}, func(context.Context, *milvuspb.GetComponentStatesRequest) error {
|
|
|
|
select {
|
|
|
|
case <-node.ctx.Done():
|
|
|
|
return grpc.ErrServerStopped
|
|
|
|
default:
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
}).Maybe()
|
2022-09-15 18:48:32 +08:00
|
|
|
|
2022-11-04 14:25:38 +08:00
|
|
|
// Register
|
2022-09-20 15:54:49 +08:00
|
|
|
node.session.Init(typeutil.QueryNodeRole, node.addr, false, true)
|
2022-11-04 14:25:38 +08:00
|
|
|
node.session.ServerID = node.ID
|
2022-09-20 15:54:49 +08:00
|
|
|
node.session.Register()
|
|
|
|
log.Debug("mock QueryNode started",
|
|
|
|
zap.Int64("nodeID", node.ID),
|
|
|
|
zap.String("nodeAddr", node.addr))
|
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
2022-12-06 22:59:19 +08:00
|
|
|
func (node *MockQueryNode) Stopping() {
|
|
|
|
node.session.GoingStop()
|
|
|
|
}
|
|
|
|
|
2022-09-15 18:48:32 +08:00
|
|
|
func (node *MockQueryNode) Stop() {
|
|
|
|
node.cancel()
|
|
|
|
node.server.GracefulStop()
|
|
|
|
node.session.Revoke(time.Second)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (node *MockQueryNode) getAllChannels() []*querypb.ChannelVersionInfo {
|
|
|
|
node.rwmutex.RLock()
|
|
|
|
defer node.rwmutex.RUnlock()
|
|
|
|
|
|
|
|
ret := make([]*querypb.ChannelVersionInfo, 0)
|
|
|
|
for collection, channels := range node.channels {
|
|
|
|
for _, channel := range channels {
|
|
|
|
ret = append(ret, &querypb.ChannelVersionInfo{
|
|
|
|
Channel: channel,
|
|
|
|
Collection: collection,
|
|
|
|
Version: node.channelVersion[channel],
|
|
|
|
})
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return ret
|
|
|
|
}
|
|
|
|
|
|
|
|
func (node *MockQueryNode) getAllSegments() []*querypb.SegmentVersionInfo {
|
|
|
|
node.rwmutex.RLock()
|
|
|
|
defer node.rwmutex.RUnlock()
|
|
|
|
|
|
|
|
ret := make([]*querypb.SegmentVersionInfo, 0)
|
|
|
|
for collection, shardSegments := range node.segments {
|
|
|
|
for shard, segments := range shardSegments {
|
|
|
|
for _, segment := range segments {
|
|
|
|
ret = append(ret, &querypb.SegmentVersionInfo{
|
|
|
|
ID: segment,
|
|
|
|
Collection: collection,
|
|
|
|
Channel: shard,
|
|
|
|
Version: node.segmentVersion[segment],
|
|
|
|
})
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return ret
|
|
|
|
}
|