test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-11-29 18:38:44 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

Add server id validation interceptor (#26395)

Browse Source 
Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
This commit is contained in:
yihao.dai 2023-08-17 20:20:20 +08:00 committed by GitHub
parent 2539a19885
commit 63b86b32a6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
49 changed files with 1397 additions and 355 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
internal/datacoord/cluster_test.go
View File

@ -367,7 +367,7 @@ func (suite *ClusterSuite) TestUnregister() {
ctx, cancel := context.WithCancel(context.TODO())
defer cancel()
var mockSessionCreator = func(ctx context.Context, addr string) (types.DataNode, error) {
var mockSessionCreator = func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return newMockDataNodeClient(1, nil)
}
sessionManager := NewSessionManager(withSessionCreator(mockSessionCreator))
@ -414,7 +414,7 @@ func TestWatchIfNeeded(t *testing.T) {
ctx, cancel := context.WithCancel(context.TODO())
defer cancel()
var mockSessionCreator = func(ctx context.Context, addr string) (types.DataNode, error) {
var mockSessionCreator = func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return newMockDataNodeClient(1, nil)
}
sessionManager := NewSessionManager(withSessionCreator(mockSessionCreator))
@ -629,7 +629,7 @@ func TestCluster_ReCollectSegmentStats(t *testing.T) {
t.Run("recollect succeed", func(t *testing.T) {
ctx, cancel := context.WithCancel(context.TODO())
defer cancel()
var mockSessionCreator = func(ctx context.Context, addr string) (types.DataNode, error) {
var mockSessionCreator = func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return newMockDataNodeClient(1, nil)
}
sessionManager := NewSessionManager(withSessionCreator(mockSessionCreator))

2
internal/datacoord/indexnode_manager.go
View File

@ -86,7 +86,7 @@ func (nm *IndexNodeManager) AddNode(nodeID UniqueID, address string) error {
err error
)
nodeClient, err = nm.indexNodeCreator(context.TODO(), address)
nodeClient, err = nm.indexNodeCreator(context.TODO(), address, nodeID)
if err != nil {
log.Error("create IndexNode client fail", zap.Error(err))
return err

6
internal/datacoord/metrics_info_test.go
View File

@ -68,7 +68,7 @@ func TestGetDataNodeMetrics(t *testing.T) {
_, err = svr.getDataNodeMetrics(ctx, req, NewSession(&NodeInfo{}, nil))
assert.Error(t, err)
creator := func(ctx context.Context, addr string) (types.DataNode, error) {
creator := func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return newMockDataNodeClient(100, nil)
}
@ -80,8 +80,8 @@ func TestGetDataNodeMetrics(t *testing.T) {
assert.Equal(t, metricsinfo.ConstructComponentName(typeutil.DataNodeRole, 100), info.BaseComponentInfos.Name)
getMockFailedClientCreator := func(mockFunc func() (*milvuspb.GetMetricsResponse, error)) dataNodeCreatorFunc {
return func(ctx context.Context, addr string) (types.DataNode, error) {
cli, err := creator(ctx, addr)
return func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
cli, err := creator(ctx, addr, nodeID)
assert.NoError(t, err)
return &mockMetricDataNodeClient{DataNode: cli, mock: mockFunc}, nil
}

16
internal/datacoord/server.go
View File

@ -79,9 +79,9 @@ type (
Timestamp = typeutil.Timestamp
)
type dataNodeCreatorFunc func(ctx context.Context, addr string) (types.DataNode, error)
type dataNodeCreatorFunc func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error)
type indexNodeCreatorFunc func(ctx context.Context, addr string) (types.IndexNode, error)
type indexNodeCreatorFunc func(ctx context.Context, addr string, nodeID int64) (types.IndexNode, error)
type rootCoordCreatorFunc func(ctx context.Context, metaRootPath string, etcdClient *clientv3.Client) (types.RootCoord, error)
@ -220,12 +220,12 @@ func CreateServer(ctx context.Context, factory dependency.Factory, opts ...Optio
return s
}
func defaultDataNodeCreatorFunc(ctx context.Context, addr string) (types.DataNode, error) {
return datanodeclient.NewClient(ctx, addr)
func defaultDataNodeCreatorFunc(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return datanodeclient.NewClient(ctx, addr, nodeID)
}
func defaultIndexNodeCreatorFunc(ctx context.Context, addr string) (types.IndexNode, error) {
return indexnodeclient.NewClient(context.TODO(), addr, Params.DataCoordCfg.WithCredential.GetAsBool())
func defaultIndexNodeCreatorFunc(ctx context.Context, addr string, nodeID int64) (types.IndexNode, error) {
return indexnodeclient.NewClient(ctx, addr, nodeID, Params.DataCoordCfg.WithCredential.GetAsBool())
}
func defaultRootCoordCreatorFunc(ctx context.Context, metaRootPath string, client *clientv3.Client) (types.RootCoord, error) {
@ -427,11 +427,11 @@ func (s *Server) SetRootCoord(rootCoord types.RootCoord) {
s.rootCoordClient = rootCoord
}
func (s *Server) SetDataNodeCreator(f func(context.Context, string) (types.DataNode, error)) {
func (s *Server) SetDataNodeCreator(f func(context.Context, string, int64) (types.DataNode, error)) {
s.dataNodeCreator = f
}
func (s *Server) SetIndexNodeCreator(f func(context.Context, string) (types.IndexNode, error)) {
func (s *Server) SetIndexNodeCreator(f func(context.Context, string, int64) (types.IndexNode, error)) {
s.indexNodeCreator = f
}

18
internal/datacoord/server_test.go
View File

@ -3221,7 +3221,7 @@ func TestOptions(t *testing.T) {
t.Run("WithDataNodeCreator", func(t *testing.T) {
var target int64
var val = rand.Int63()
opt := WithDataNodeCreator(func(context.Context, string) (types.DataNode, error) {
opt := WithDataNodeCreator(func(context.Context, string, int64) (types.DataNode, error) {
target = val
return nil, nil
})
@ -3230,7 +3230,7 @@ func TestOptions(t *testing.T) {
factory := dependency.NewDefaultFactory(true)
svr := CreateServer(context.TODO(), factory, opt)
dn, err := svr.dataNodeCreator(context.Background(), "")
dn, err := svr.dataNodeCreator(context.Background(), "", 1)
assert.Nil(t, dn)
assert.NoError(t, err)
assert.Equal(t, target, val)
@ -3945,7 +3945,7 @@ func newTestServer(t *testing.T, receiveCh chan any, opts ...Option) *Server {
svr := CreateServer(context.TODO(), factory)
svr.SetEtcdClient(etcdCli)
svr.dataNodeCreator = func(ctx context.Context, addr string) (types.DataNode, error) {
svr.dataNodeCreator = func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return newMockDataNodeClient(0, receiveCh)
}
svr.rootCoordClientCreator = func(ctx context.Context, metaRootPath string, etcdCli *clientv3.Client) (types.RootCoord, error) {
@ -3997,7 +3997,7 @@ func newTestServerWithMeta(t *testing.T, receiveCh chan any, meta *meta, opts ..
svr := CreateServer(context.TODO(), factory, opts...)
svr.SetEtcdClient(etcdCli)
svr.dataNodeCreator = func(ctx context.Context, addr string) (types.DataNode, error) {
svr.dataNodeCreator = func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return newMockDataNodeClient(0, receiveCh)
}
svr.rootCoordClientCreator = func(ctx context.Context, metaRootPath string, etcdCli *clientv3.Client) (types.RootCoord, error) {
@ -4052,7 +4052,7 @@ func newTestServer2(t *testing.T, receiveCh chan any, opts ...Option) *Server {
svr := CreateServer(context.TODO(), factory, opts...)
svr.SetEtcdClient(etcdCli)
svr.dataNodeCreator = func(ctx context.Context, addr string) (types.DataNode, error) {
svr.dataNodeCreator = func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return newMockDataNodeClient(0, receiveCh)
}
svr.rootCoordClientCreator = func(ctx context.Context, metaRootPath string, etcdCli *clientv3.Client) (types.RootCoord, error) {
@ -4097,7 +4097,7 @@ func Test_CheckHealth(t *testing.T) {
data map[int64]*Session
}{data: map[int64]*Session{1: {
client: healthClient,
clientCreator: func(ctx context.Context, addr string) (types.DataNode, error) {
clientCreator: func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return healthClient, nil
},
}}}
@ -4122,7 +4122,7 @@ func Test_CheckHealth(t *testing.T) {
data map[int64]*Session
}{data: map[int64]*Session{1: {
client: unhealthClient,
clientCreator: func(ctx context.Context, addr string) (types.DataNode, error) {
clientCreator: func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return unhealthClient, nil
},
}}}
@ -4244,10 +4244,10 @@ func testDataCoordBase(t *testing.T, opts ...Option) *Server {
svr := CreateServer(ctx, factory, opts...)
svr.SetEtcdClient(etcdCli)
svr.SetDataNodeCreator(func(ctx context.Context, addr string) (types.DataNode, error) {
svr.SetDataNodeCreator(func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return newMockDataNodeClient(0, nil)
})
svr.SetIndexNodeCreator(func(ctx context.Context, addr string) (types.IndexNode, error) {
svr.SetIndexNodeCreator(func(ctx context.Context, addr string, nodeID int64) (types.IndexNode, error) {
return indexnode.NewMockIndexNodeComponent(ctx)
})
svr.SetRootCoord(newMockRootCoordService())

2
internal/datacoord/session.go
View File

@ -73,7 +73,7 @@ func (n *Session) GetOrCreateClient(ctx context.Context) (types.DataNode, error)
}
func (n *Session) initClient(ctx context.Context) (err error) {
if n.client, err = n.clientCreator(ctx, n.info.Address); err != nil {
if n.client, err = n.clientCreator(ctx, n.info.Address, n.info.NodeID); err != nil {
return
}
if err = n.client.Init(); err != nil {

4
internal/datacoord/session_manager.go
View File

@ -59,8 +59,8 @@ func withSessionCreator(creator dataNodeCreatorFunc) SessionOpt {
}
func defaultSessionCreator() dataNodeCreatorFunc {
return func(ctx context.Context, addr string) (types.DataNode, error) {
return grpcdatanodeclient.NewClient(ctx, addr)
return func(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
return grpcdatanodeclient.NewClient(ctx, addr, nodeID)
}
}

2
internal/distributed/datacoord/service.go
View File

@ -154,11 +154,13 @@ func (s *Server) startGrpcLoop(grpcPort int) {
otelgrpc.UnaryServerInterceptor(opts...),
logutil.UnaryTraceLoggerInterceptor,
interceptor.ClusterValidationUnaryServerInterceptor(),
interceptor.ServerIDValidationUnaryServerInterceptor(),
)),
grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
otelgrpc.StreamServerInterceptor(opts...),
logutil.StreamTraceLoggerInterceptor,
interceptor.ClusterValidationStreamServerInterceptor(),
interceptor.ServerIDValidationStreamServerInterceptor(),
)))
indexpb.RegisterIndexCoordServer(s.grpcServer, s)
datapb.RegisterDataCoordServer(s.grpcServer, s)

4
internal/distributed/datacoord/service_test.go
View File

@ -107,10 +107,10 @@ func (m *MockDataCoord) SetEtcdClient(etcdClient *clientv3.Client) {
func (m *MockDataCoord) SetRootCoord(rootCoord types.RootCoord) {
}
func (m *MockDataCoord) SetDataNodeCreator(func(context.Context, string) (types.DataNode, error)) {
func (m *MockDataCoord) SetDataNodeCreator(func(context.Context, string, int64) (types.DataNode, error)) {
}
func (m *MockDataCoord) SetIndexNodeCreator(func(context.Context, string) (types.IndexNode, error)) {
func (m *MockDataCoord) SetIndexNodeCreator(func(context.Context, string, int64) (types.IndexNode, error)) {
}
func (m *MockDataCoord) GetComponentStates(ctx context.Context) (*milvuspb.ComponentStates, error) {

3
internal/distributed/datanode/client/client.go
View File

@ -42,7 +42,7 @@ type Client struct {
}
// NewClient creates a client for DataNode.
func NewClient(ctx context.Context, addr string) (*Client, error) {
func NewClient(ctx context.Context, addr string, nodeID int64) (*Client, error) {
if addr == "" {
return nil, fmt.Errorf("address is empty")
}
@ -66,6 +66,7 @@ func NewClient(ctx context.Context, addr string) (*Client, error) {
client.grpcClient.SetRole(typeutil.DataNodeRole)
client.grpcClient.SetGetAddrFunc(client.getAddr)
client.grpcClient.SetNewGrpcClientFunc(client.newGrpcClient)
client.grpcClient.SetNodeID(nodeID)
return client, nil
}

4
internal/distributed/datanode/client/client_test.go
View File

@ -33,11 +33,11 @@ import (
func Test_NewClient(t *testing.T) {
proxy.Params.Init()
ctx := context.Background()
client, err := NewClient(ctx, "")
client, err := NewClient(ctx, "", 1)
assert.Nil(t, client)
assert.Error(t, err)
client, err = NewClient(ctx, "test")
client, err = NewClient(ctx, "test", 2)
assert.NoError(t, err)
assert.NotNil(t, client)

2
internal/distributed/datanode/service.go
View File

@ -138,11 +138,13 @@ func (s *Server) startGrpcLoop(grpcPort int) {
otelgrpc.UnaryServerInterceptor(opts...),
logutil.UnaryTraceLoggerInterceptor,
interceptor.ClusterValidationUnaryServerInterceptor(),
interceptor.ServerIDValidationUnaryServerInterceptor(),
)),
grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
otelgrpc.StreamServerInterceptor(opts...),
logutil.StreamTraceLoggerInterceptor,
interceptor.ClusterValidationStreamServerInterceptor(),
interceptor.ServerIDValidationStreamServerInterceptor(),
)))
datapb.RegisterDataNodeServer(s.grpcServer, s)

3
internal/distributed/indexnode/client/client.go
View File

@ -43,7 +43,7 @@ type Client struct {
}
// NewClient creates a new IndexNode client.
func NewClient(ctx context.Context, addr string, encryption bool) (*Client, error) {
func NewClient(ctx context.Context, addr string, nodeID int64, encryption bool) (*Client, error) {
if addr == "" {
return nil, fmt.Errorf("address is empty")
}
@ -67,6 +67,7 @@ func NewClient(ctx context.Context, addr string, encryption bool) (*Client, erro
client.grpcClient.SetRole(typeutil.IndexNodeRole)
client.grpcClient.SetGetAddrFunc(client.getAddr)
client.grpcClient.SetNewGrpcClientFunc(client.newGrpcClient)
client.grpcClient.SetNodeID(nodeID)
if encryption {
client.grpcClient.EnableEncryption()
}

6
internal/distributed/indexnode/client/client_test.go
View File

@ -40,11 +40,11 @@ import (
func Test_NewClient(t *testing.T) {
paramtable.Init()
ctx := context.Background()
client, err := NewClient(ctx, "", false)
client, err := NewClient(ctx, "", 1, false)
assert.Nil(t, client)
assert.Error(t, err)
client, err = NewClient(ctx, "test", false)
client, err = NewClient(ctx, "test", 2, false)
assert.NoError(t, err)
assert.NotNil(t, client)
@ -148,7 +148,7 @@ func TestIndexNodeClient(t *testing.T) {
err = ins.Run()
assert.NoError(t, err)
inc, err := NewClient(ctx, "localhost:21121", false)
inc, err := NewClient(ctx, "localhost:21121", paramtable.GetNodeID(), false)
assert.NoError(t, err)
assert.NotNil(t, inc)

2
internal/distributed/indexnode/service.go
View File

@ -110,11 +110,13 @@ func (s *Server) startGrpcLoop(grpcPort int) {
otelgrpc.UnaryServerInterceptor(opts...),
logutil.UnaryTraceLoggerInterceptor,
interceptor.ClusterValidationUnaryServerInterceptor(),
interceptor.ServerIDValidationUnaryServerInterceptor(),
)),
grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
otelgrpc.StreamServerInterceptor(opts...),
logutil.StreamTraceLoggerInterceptor,
interceptor.ClusterValidationStreamServerInterceptor(),
interceptor.ServerIDValidationStreamServerInterceptor(),
)))
indexpb.RegisterIndexNodeServer(s.grpcServer, s)
go funcutil.CheckGrpcReady(ctx, s.grpcErrChan)

3
internal/distributed/proxy/client/client.go
View File

@ -42,7 +42,7 @@ type Client struct {
}
// NewClient creates a new client instance
func NewClient(ctx context.Context, addr string) (*Client, error) {
func NewClient(ctx context.Context, addr string, nodeID int64) (*Client, error) {
if addr == "" {
return nil, fmt.Errorf("address is empty")
}
@ -66,6 +66,7 @@ func NewClient(ctx context.Context, addr string) (*Client, error) {
client.grpcClient.SetRole(typeutil.ProxyRole)
client.grpcClient.SetGetAddrFunc(client.getAddr)
client.grpcClient.SetNewGrpcClientFunc(client.newGrpcClient)
client.grpcClient.SetNodeID(nodeID)
return client, nil
}

4
internal/distributed/proxy/client/client_test.go
View File

@ -34,11 +34,11 @@ func Test_NewClient(t *testing.T) {
proxy.Params.Init()
ctx := context.Background()
client, err := NewClient(ctx, "")
client, err := NewClient(ctx, "", 1)
assert.Nil(t, client)
assert.Error(t, err)
client, err = NewClient(ctx, "test")
client, err = NewClient(ctx, "test", 2)
assert.NoError(t, err)
assert.NotNil(t, client)

7
internal/distributed/proxy/service.go
View File

@ -314,9 +314,12 @@ func (s *Server) startInternalGrpc(grpcPort int, errChan chan error) {
otelgrpc.UnaryServerInterceptor(opts...),
logutil.UnaryTraceLoggerInterceptor,
interceptor.ClusterValidationUnaryServerInterceptor(),
interceptor.ServerIDValidationUnaryServerInterceptor(),
)),
grpc.StreamInterceptor(interceptor.ClusterValidationStreamServerInterceptor()),
)
grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
interceptor.ClusterValidationStreamServerInterceptor(),
interceptor.ServerIDValidationStreamServerInterceptor(),
)))
proxypb.RegisterProxyServer(s.grpcInternalServer, s)
grpc_health_v1.RegisterHealthServer(s.grpcInternalServer, s)
errChan <- nil

2
internal/distributed/proxy/service_test.go
View File

@ -740,7 +740,7 @@ func (m *MockProxy) SetQueryCoordClient(queryCoord types.QueryCoord) {
}
func (m *MockProxy) SetQueryNodeCreator(func(ctx context.Context, addr string) (types.QueryNode, error)) {
func (m *MockProxy) SetQueryNodeCreator(func(ctx context.Context, addr string, nodeID int64) (types.QueryNode, error)) {
}

2
internal/distributed/querycoord/service.go
View File

@ -228,11 +228,13 @@ func (s *Server) startGrpcLoop(grpcPort int) {
otelgrpc.UnaryServerInterceptor(opts...),
logutil.UnaryTraceLoggerInterceptor,
interceptor.ClusterValidationUnaryServerInterceptor(),
interceptor.ServerIDValidationUnaryServerInterceptor(),
)),
grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
otelgrpc.StreamServerInterceptor(opts...),
logutil.StreamTraceLoggerInterceptor,
interceptor.ClusterValidationStreamServerInterceptor(),
interceptor.ServerIDValidationStreamServerInterceptor(),
)))
querypb.RegisterQueryCoordServer(s.grpcServer, s)

3
internal/distributed/querynode/client/client.go
View File

@ -41,7 +41,7 @@ type Client struct {
}
// NewClient creates a new QueryNode client.
func NewClient(ctx context.Context, addr string) (*Client, error) {
func NewClient(ctx context.Context, addr string, nodeID int64) (*Client, error) {
if addr == "" {
return nil, fmt.Errorf("addr is empty")
}
@ -65,6 +65,7 @@ func NewClient(ctx context.Context, addr string) (*Client, error) {
client.grpcClient.SetRole(typeutil.QueryNodeRole)
client.grpcClient.SetGetAddrFunc(client.getAddr)
client.grpcClient.SetNewGrpcClientFunc(client.newGrpcClient)
client.grpcClient.SetNodeID(nodeID)
return client, nil
}

4
internal/distributed/querynode/client/client_test.go
View File

@ -33,11 +33,11 @@ func Test_NewClient(t *testing.T) {
paramtable.Init()
ctx := context.Background()
client, err := NewClient(ctx, "")
client, err := NewClient(ctx, "", 1)
assert.Nil(t, client)
assert.Error(t, err)
client, err = NewClient(ctx, "test")
client, err = NewClient(ctx, "test", 2)
assert.NoError(t, err)
assert.NotNil(t, client)

2
internal/distributed/querynode/service.go
View File

@ -183,11 +183,13 @@ func (s *Server) startGrpcLoop(grpcPort int) {
otelgrpc.UnaryServerInterceptor(opts...),
logutil.UnaryTraceLoggerInterceptor,
interceptor.ClusterValidationUnaryServerInterceptor(),
interceptor.ServerIDValidationUnaryServerInterceptor(),
)),
grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
otelgrpc.StreamServerInterceptor(opts...),
logutil.StreamTraceLoggerInterceptor,
interceptor.ClusterValidationStreamServerInterceptor(),
interceptor.ServerIDValidationStreamServerInterceptor(),
)))
querypb.RegisterQueryNodeServer(s.grpcServer, s)

2
internal/distributed/rootcoord/service.go
View File

@ -255,11 +255,13 @@ func (s *Server) startGrpcLoop(port int) {
otelgrpc.UnaryServerInterceptor(opts...),
logutil.UnaryTraceLoggerInterceptor,
interceptor.ClusterValidationUnaryServerInterceptor(),
interceptor.ServerIDValidationUnaryServerInterceptor(),
)),
grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
otelgrpc.StreamServerInterceptor(opts...),
logutil.StreamTraceLoggerInterceptor,
interceptor.ClusterValidationStreamServerInterceptor(),
interceptor.ServerIDValidationStreamServerInterceptor(),
)))
rootcoordpb.RegisterRootCoordServer(s.grpcServer, s)

2
internal/distributed/rootcoord/service_test.go
View File

@ -84,7 +84,7 @@ func (m *mockCore) SetQueryCoord(types.QueryCoord) error {
return nil
}
func (m *mockCore) SetProxyCreator(func(ctx context.Context, addr string) (types.Proxy, error)) {
func (m *mockCore) SetProxyCreator(func(ctx context.Context, addr string, nodeID int64) (types.Proxy, error)) {
}
func (m *mockCore) Register() error {

194
internal/mocks/mock_datacoord.go
View File

@ -66,8 +66,8 @@ type MockDataCoord_AssignSegmentID_Call struct {
}
// AssignSegmentID is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.AssignSegmentIDRequest
// - ctx context.Context
// - req *datapb.AssignSegmentIDRequest
func (_e *MockDataCoord_Expecter) AssignSegmentID(ctx interface{}, req interface{}) *MockDataCoord_AssignSegmentID_Call {
return &MockDataCoord_AssignSegmentID_Call{Call: _e.mock.On("AssignSegmentID", ctx, req)}
}
@ -121,8 +121,8 @@ type MockDataCoord_BroadcastAlteredCollection_Call struct {
}
// BroadcastAlteredCollection is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.AlterCollectionRequest
// - ctx context.Context
// - req *datapb.AlterCollectionRequest
func (_e *MockDataCoord_Expecter) BroadcastAlteredCollection(ctx interface{}, req interface{}) *MockDataCoord_BroadcastAlteredCollection_Call {
return &MockDataCoord_BroadcastAlteredCollection_Call{Call: _e.mock.On("BroadcastAlteredCollection", ctx, req)}
}
@ -176,8 +176,8 @@ type MockDataCoord_CheckHealth_Call struct {
}
// CheckHealth is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.CheckHealthRequest
// - ctx context.Context
// - req *milvuspb.CheckHealthRequest
func (_e *MockDataCoord_Expecter) CheckHealth(ctx interface{}, req interface{}) *MockDataCoord_CheckHealth_Call {
return &MockDataCoord_CheckHealth_Call{Call: _e.mock.On("CheckHealth", ctx, req)}
}
@ -231,8 +231,8 @@ type MockDataCoord_CreateIndex_Call struct {
}
// CreateIndex is a helper method to define mock.On call
// - ctx context.Context
// - req *indexpb.CreateIndexRequest
// - ctx context.Context
// - req *indexpb.CreateIndexRequest
func (_e *MockDataCoord_Expecter) CreateIndex(ctx interface{}, req interface{}) *MockDataCoord_CreateIndex_Call {
return &MockDataCoord_CreateIndex_Call{Call: _e.mock.On("CreateIndex", ctx, req)}
}
@ -286,8 +286,8 @@ type MockDataCoord_DescribeIndex_Call struct {
}
// DescribeIndex is a helper method to define mock.On call
// - ctx context.Context
// - req *indexpb.DescribeIndexRequest
// - ctx context.Context
// - req *indexpb.DescribeIndexRequest
func (_e *MockDataCoord_Expecter) DescribeIndex(ctx interface{}, req interface{}) *MockDataCoord_DescribeIndex_Call {
return &MockDataCoord_DescribeIndex_Call{Call: _e.mock.On("DescribeIndex", ctx, req)}
}
@ -341,8 +341,8 @@ type MockDataCoord_DropIndex_Call struct {
}
// DropIndex is a helper method to define mock.On call
// - ctx context.Context
// - req *indexpb.DropIndexRequest
// - ctx context.Context
// - req *indexpb.DropIndexRequest
func (_e *MockDataCoord_Expecter) DropIndex(ctx interface{}, req interface{}) *MockDataCoord_DropIndex_Call {
return &MockDataCoord_DropIndex_Call{Call: _e.mock.On("DropIndex", ctx, req)}
}
@ -396,8 +396,8 @@ type MockDataCoord_DropVirtualChannel_Call struct {
}
// DropVirtualChannel is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.DropVirtualChannelRequest
// - ctx context.Context
// - req *datapb.DropVirtualChannelRequest
func (_e *MockDataCoord_Expecter) DropVirtualChannel(ctx interface{}, req interface{}) *MockDataCoord_DropVirtualChannel_Call {
return &MockDataCoord_DropVirtualChannel_Call{Call: _e.mock.On("DropVirtualChannel", ctx, req)}
}
@ -451,8 +451,8 @@ type MockDataCoord_Flush_Call struct {
}
// Flush is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.FlushRequest
// - ctx context.Context
// - req *datapb.FlushRequest
func (_e *MockDataCoord_Expecter) Flush(ctx interface{}, req interface{}) *MockDataCoord_Flush_Call {
return &MockDataCoord_Flush_Call{Call: _e.mock.On("Flush", ctx, req)}
}
@ -506,8 +506,8 @@ type MockDataCoord_GcConfirm_Call struct {
}
// GcConfirm is a helper method to define mock.On call
// - ctx context.Context
// - request *datapb.GcConfirmRequest
// - ctx context.Context
// - request *datapb.GcConfirmRequest
func (_e *MockDataCoord_Expecter) GcConfirm(ctx interface{}, request interface{}) *MockDataCoord_GcConfirm_Call {
return &MockDataCoord_GcConfirm_Call{Call: _e.mock.On("GcConfirm", ctx, request)}
}
@ -561,8 +561,8 @@ type MockDataCoord_GetCollectionStatistics_Call struct {
}
// GetCollectionStatistics is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.GetCollectionStatisticsRequest
// - ctx context.Context
// - req *datapb.GetCollectionStatisticsRequest
func (_e *MockDataCoord_Expecter) GetCollectionStatistics(ctx interface{}, req interface{}) *MockDataCoord_GetCollectionStatistics_Call {
return &MockDataCoord_GetCollectionStatistics_Call{Call: _e.mock.On("GetCollectionStatistics", ctx, req)}
}
@ -616,8 +616,8 @@ type MockDataCoord_GetCompactionState_Call struct {
}
// GetCompactionState is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.GetCompactionStateRequest
// - ctx context.Context
// - req *milvuspb.GetCompactionStateRequest
func (_e *MockDataCoord_Expecter) GetCompactionState(ctx interface{}, req interface{}) *MockDataCoord_GetCompactionState_Call {
return &MockDataCoord_GetCompactionState_Call{Call: _e.mock.On("GetCompactionState", ctx, req)}
}
@ -671,8 +671,8 @@ type MockDataCoord_GetCompactionStateWithPlans_Call struct {
}
// GetCompactionStateWithPlans is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.GetCompactionPlansRequest
// - ctx context.Context
// - req *milvuspb.GetCompactionPlansRequest
func (_e *MockDataCoord_Expecter) GetCompactionStateWithPlans(ctx interface{}, req interface{}) *MockDataCoord_GetCompactionStateWithPlans_Call {
return &MockDataCoord_GetCompactionStateWithPlans_Call{Call: _e.mock.On("GetCompactionStateWithPlans", ctx, req)}
}
@ -726,7 +726,7 @@ type MockDataCoord_GetComponentStates_Call struct {
}
// GetComponentStates is a helper method to define mock.On call
// - ctx context.Context
// - ctx context.Context
func (_e *MockDataCoord_Expecter) GetComponentStates(ctx interface{}) *MockDataCoord_GetComponentStates_Call {
return &MockDataCoord_GetComponentStates_Call{Call: _e.mock.On("GetComponentStates", ctx)}
}
@ -780,8 +780,8 @@ type MockDataCoord_GetFlushAllState_Call struct {
}
// GetFlushAllState is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.GetFlushAllStateRequest
// - ctx context.Context
// - req *milvuspb.GetFlushAllStateRequest
func (_e *MockDataCoord_Expecter) GetFlushAllState(ctx interface{}, req interface{}) *MockDataCoord_GetFlushAllState_Call {
return &MockDataCoord_GetFlushAllState_Call{Call: _e.mock.On("GetFlushAllState", ctx, req)}
}
@ -835,8 +835,8 @@ type MockDataCoord_GetFlushState_Call struct {
}
// GetFlushState is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.GetFlushStateRequest
// - ctx context.Context
// - req *milvuspb.GetFlushStateRequest
func (_e *MockDataCoord_Expecter) GetFlushState(ctx interface{}, req interface{}) *MockDataCoord_GetFlushState_Call {
return &MockDataCoord_GetFlushState_Call{Call: _e.mock.On("GetFlushState", ctx, req)}
}
@ -890,8 +890,8 @@ type MockDataCoord_GetFlushedSegments_Call struct {
}
// GetFlushedSegments is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.GetFlushedSegmentsRequest
// - ctx context.Context
// - req *datapb.GetFlushedSegmentsRequest
func (_e *MockDataCoord_Expecter) GetFlushedSegments(ctx interface{}, req interface{}) *MockDataCoord_GetFlushedSegments_Call {
return &MockDataCoord_GetFlushedSegments_Call{Call: _e.mock.On("GetFlushedSegments", ctx, req)}
}
@ -945,8 +945,8 @@ type MockDataCoord_GetIndexBuildProgress_Call struct {
}
// GetIndexBuildProgress is a helper method to define mock.On call
// - ctx context.Context
// - req *indexpb.GetIndexBuildProgressRequest
// - ctx context.Context
// - req *indexpb.GetIndexBuildProgressRequest
func (_e *MockDataCoord_Expecter) GetIndexBuildProgress(ctx interface{}, req interface{}) *MockDataCoord_GetIndexBuildProgress_Call {
return &MockDataCoord_GetIndexBuildProgress_Call{Call: _e.mock.On("GetIndexBuildProgress", ctx, req)}
}
@ -1000,8 +1000,8 @@ type MockDataCoord_GetIndexInfos_Call struct {
}
// GetIndexInfos is a helper method to define mock.On call
// - ctx context.Context
// - req *indexpb.GetIndexInfoRequest
// - ctx context.Context
// - req *indexpb.GetIndexInfoRequest
func (_e *MockDataCoord_Expecter) GetIndexInfos(ctx interface{}, req interface{}) *MockDataCoord_GetIndexInfos_Call {
return &MockDataCoord_GetIndexInfos_Call{Call: _e.mock.On("GetIndexInfos", ctx, req)}
}
@ -1055,8 +1055,8 @@ type MockDataCoord_GetIndexState_Call struct {
}
// GetIndexState is a helper method to define mock.On call
// - ctx context.Context
// - req *indexpb.GetIndexStateRequest
// - ctx context.Context
// - req *indexpb.GetIndexStateRequest
func (_e *MockDataCoord_Expecter) GetIndexState(ctx interface{}, req interface{}) *MockDataCoord_GetIndexState_Call {
return &MockDataCoord_GetIndexState_Call{Call: _e.mock.On("GetIndexState", ctx, req)}
}
@ -1110,8 +1110,8 @@ type MockDataCoord_GetIndexStatistics_Call struct {
}
// GetIndexStatistics is a helper method to define mock.On call
// - ctx context.Context
// - req *indexpb.GetIndexStatisticsRequest
// - ctx context.Context
// - req *indexpb.GetIndexStatisticsRequest
func (_e *MockDataCoord_Expecter) GetIndexStatistics(ctx interface{}, req interface{}) *MockDataCoord_GetIndexStatistics_Call {
return &MockDataCoord_GetIndexStatistics_Call{Call: _e.mock.On("GetIndexStatistics", ctx, req)}
}
@ -1165,8 +1165,8 @@ type MockDataCoord_GetInsertBinlogPaths_Call struct {
}
// GetInsertBinlogPaths is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.GetInsertBinlogPathsRequest
// - ctx context.Context
// - req *datapb.GetInsertBinlogPathsRequest
func (_e *MockDataCoord_Expecter) GetInsertBinlogPaths(ctx interface{}, req interface{}) *MockDataCoord_GetInsertBinlogPaths_Call {
return &MockDataCoord_GetInsertBinlogPaths_Call{Call: _e.mock.On("GetInsertBinlogPaths", ctx, req)}
}
@ -1220,8 +1220,8 @@ type MockDataCoord_GetMetrics_Call struct {
}
// GetMetrics is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.GetMetricsRequest
// - ctx context.Context
// - req *milvuspb.GetMetricsRequest
func (_e *MockDataCoord_Expecter) GetMetrics(ctx interface{}, req interface{}) *MockDataCoord_GetMetrics_Call {
return &MockDataCoord_GetMetrics_Call{Call: _e.mock.On("GetMetrics", ctx, req)}
}
@ -1275,8 +1275,8 @@ type MockDataCoord_GetPartitionStatistics_Call struct {
}
// GetPartitionStatistics is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.GetPartitionStatisticsRequest
// - ctx context.Context
// - req *datapb.GetPartitionStatisticsRequest
func (_e *MockDataCoord_Expecter) GetPartitionStatistics(ctx interface{}, req interface{}) *MockDataCoord_GetPartitionStatistics_Call {
return &MockDataCoord_GetPartitionStatistics_Call{Call: _e.mock.On("GetPartitionStatistics", ctx, req)}
}
@ -1330,8 +1330,8 @@ type MockDataCoord_GetRecoveryInfo_Call struct {
}
// GetRecoveryInfo is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.GetRecoveryInfoRequest
// - ctx context.Context
// - req *datapb.GetRecoveryInfoRequest
func (_e *MockDataCoord_Expecter) GetRecoveryInfo(ctx interface{}, req interface{}) *MockDataCoord_GetRecoveryInfo_Call {
return &MockDataCoord_GetRecoveryInfo_Call{Call: _e.mock.On("GetRecoveryInfo", ctx, req)}
}
@ -1385,8 +1385,8 @@ type MockDataCoord_GetRecoveryInfoV2_Call struct {
}
// GetRecoveryInfoV2 is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.GetRecoveryInfoRequestV2
// - ctx context.Context
// - req *datapb.GetRecoveryInfoRequestV2
func (_e *MockDataCoord_Expecter) GetRecoveryInfoV2(ctx interface{}, req interface{}) *MockDataCoord_GetRecoveryInfoV2_Call {
return &MockDataCoord_GetRecoveryInfoV2_Call{Call: _e.mock.On("GetRecoveryInfoV2", ctx, req)}
}
@ -1440,8 +1440,8 @@ type MockDataCoord_GetSegmentIndexState_Call struct {
}
// GetSegmentIndexState is a helper method to define mock.On call
// - ctx context.Context
// - req *indexpb.GetSegmentIndexStateRequest
// - ctx context.Context
// - req *indexpb.GetSegmentIndexStateRequest
func (_e *MockDataCoord_Expecter) GetSegmentIndexState(ctx interface{}, req interface{}) *MockDataCoord_GetSegmentIndexState_Call {
return &MockDataCoord_GetSegmentIndexState_Call{Call: _e.mock.On("GetSegmentIndexState", ctx, req)}
}
@ -1495,8 +1495,8 @@ type MockDataCoord_GetSegmentInfo_Call struct {
}
// GetSegmentInfo is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.GetSegmentInfoRequest
// - ctx context.Context
// - req *datapb.GetSegmentInfoRequest
func (_e *MockDataCoord_Expecter) GetSegmentInfo(ctx interface{}, req interface{}) *MockDataCoord_GetSegmentInfo_Call {
return &MockDataCoord_GetSegmentInfo_Call{Call: _e.mock.On("GetSegmentInfo", ctx, req)}
}
@ -1550,7 +1550,7 @@ type MockDataCoord_GetSegmentInfoChannel_Call struct {
}
// GetSegmentInfoChannel is a helper method to define mock.On call
// - ctx context.Context
// - ctx context.Context
func (_e *MockDataCoord_Expecter) GetSegmentInfoChannel(ctx interface{}) *MockDataCoord_GetSegmentInfoChannel_Call {
return &MockDataCoord_GetSegmentInfoChannel_Call{Call: _e.mock.On("GetSegmentInfoChannel", ctx)}
}
@ -1604,8 +1604,8 @@ type MockDataCoord_GetSegmentStates_Call struct {
}
// GetSegmentStates is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.GetSegmentStatesRequest
// - ctx context.Context
// - req *datapb.GetSegmentStatesRequest
func (_e *MockDataCoord_Expecter) GetSegmentStates(ctx interface{}, req interface{}) *MockDataCoord_GetSegmentStates_Call {
return &MockDataCoord_GetSegmentStates_Call{Call: _e.mock.On("GetSegmentStates", ctx, req)}
}
@ -1659,8 +1659,8 @@ type MockDataCoord_GetSegmentsByStates_Call struct {
}
// GetSegmentsByStates is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.GetSegmentsByStatesRequest
// - ctx context.Context
// - req *datapb.GetSegmentsByStatesRequest
func (_e *MockDataCoord_Expecter) GetSegmentsByStates(ctx interface{}, req interface{}) *MockDataCoord_GetSegmentsByStates_Call {
return &MockDataCoord_GetSegmentsByStates_Call{Call: _e.mock.On("GetSegmentsByStates", ctx, req)}
}
@ -1714,7 +1714,7 @@ type MockDataCoord_GetStatisticsChannel_Call struct {
}
// GetStatisticsChannel is a helper method to define mock.On call
// - ctx context.Context
// - ctx context.Context
func (_e *MockDataCoord_Expecter) GetStatisticsChannel(ctx interface{}) *MockDataCoord_GetStatisticsChannel_Call {
return &MockDataCoord_GetStatisticsChannel_Call{Call: _e.mock.On("GetStatisticsChannel", ctx)}
}
@ -1768,7 +1768,7 @@ type MockDataCoord_GetTimeTickChannel_Call struct {
}
// GetTimeTickChannel is a helper method to define mock.On call
// - ctx context.Context
// - ctx context.Context
func (_e *MockDataCoord_Expecter) GetTimeTickChannel(ctx interface{}) *MockDataCoord_GetTimeTickChannel_Call {
return &MockDataCoord_GetTimeTickChannel_Call{Call: _e.mock.On("GetTimeTickChannel", ctx)}
}
@ -1822,8 +1822,8 @@ type MockDataCoord_Import_Call struct {
}
// Import is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.ImportTaskRequest
// - ctx context.Context
// - req *datapb.ImportTaskRequest
func (_e *MockDataCoord_Expecter) Import(ctx interface{}, req interface{}) *MockDataCoord_Import_Call {
return &MockDataCoord_Import_Call{Call: _e.mock.On("Import", ctx, req)}
}
@ -1918,8 +1918,8 @@ type MockDataCoord_ManualCompaction_Call struct {
}
// ManualCompaction is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.ManualCompactionRequest
// - ctx context.Context
// - req *milvuspb.ManualCompactionRequest
func (_e *MockDataCoord_Expecter) ManualCompaction(ctx interface{}, req interface{}) *MockDataCoord_ManualCompaction_Call {
return &MockDataCoord_ManualCompaction_Call{Call: _e.mock.On("ManualCompaction", ctx, req)}
}
@ -1973,8 +1973,8 @@ type MockDataCoord_MarkSegmentsDropped_Call struct {
}
// MarkSegmentsDropped is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.MarkSegmentsDroppedRequest
// - ctx context.Context
// - req *datapb.MarkSegmentsDroppedRequest
func (_e *MockDataCoord_Expecter) MarkSegmentsDropped(ctx interface{}, req interface{}) *MockDataCoord_MarkSegmentsDropped_Call {
return &MockDataCoord_MarkSegmentsDropped_Call{Call: _e.mock.On("MarkSegmentsDropped", ctx, req)}
}
@ -2069,8 +2069,8 @@ type MockDataCoord_ReportDataNodeTtMsgs_Call struct {
}
// ReportDataNodeTtMsgs is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.ReportDataNodeTtMsgsRequest
// - ctx context.Context
// - req *datapb.ReportDataNodeTtMsgsRequest
func (_e *MockDataCoord_Expecter) ReportDataNodeTtMsgs(ctx interface{}, req interface{}) *MockDataCoord_ReportDataNodeTtMsgs_Call {
return &MockDataCoord_ReportDataNodeTtMsgs_Call{Call: _e.mock.On("ReportDataNodeTtMsgs", ctx, req)}
}
@ -2124,8 +2124,8 @@ type MockDataCoord_SaveBinlogPaths_Call struct {
}
// SaveBinlogPaths is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.SaveBinlogPathsRequest
// - ctx context.Context
// - req *datapb.SaveBinlogPathsRequest
func (_e *MockDataCoord_Expecter) SaveBinlogPaths(ctx interface{}, req interface{}) *MockDataCoord_SaveBinlogPaths_Call {
return &MockDataCoord_SaveBinlogPaths_Call{Call: _e.mock.On("SaveBinlogPaths", ctx, req)}
}
@ -2179,8 +2179,8 @@ type MockDataCoord_SaveImportSegment_Call struct {
}
// SaveImportSegment is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.SaveImportSegmentRequest
// - ctx context.Context
// - req *datapb.SaveImportSegmentRequest
func (_e *MockDataCoord_Expecter) SaveImportSegment(ctx interface{}, req interface{}) *MockDataCoord_SaveImportSegment_Call {
return &MockDataCoord_SaveImportSegment_Call{Call: _e.mock.On("SaveImportSegment", ctx, req)}
}
@ -2213,7 +2213,7 @@ type MockDataCoord_SetAddress_Call struct {
}
// SetAddress is a helper method to define mock.On call
// - address string
// - address string
func (_e *MockDataCoord_Expecter) SetAddress(address interface{}) *MockDataCoord_SetAddress_Call {
return &MockDataCoord_SetAddress_Call{Call: _e.mock.On("SetAddress", address)}
}
@ -2236,7 +2236,7 @@ func (_c *MockDataCoord_SetAddress_Call) RunAndReturn(run func(string)) *MockDat
}
// SetDataNodeCreator provides a mock function with given fields: _a0
func (_m *MockDataCoord) SetDataNodeCreator(_a0 func(context.Context, string) (types.DataNode, error)) {
func (_m *MockDataCoord) SetDataNodeCreator(_a0 func(context.Context, string, int64) (types.DataNode, error)) {
_m.Called(_a0)
}
@ -2246,14 +2246,14 @@ type MockDataCoord_SetDataNodeCreator_Call struct {
}
// SetDataNodeCreator is a helper method to define mock.On call
// - _a0 func(context.Context , string)(types.DataNode , error)
// - _a0 func(context.Context , string , int64)(types.DataNode , error)
func (_e *MockDataCoord_Expecter) SetDataNodeCreator(_a0 interface{}) *MockDataCoord_SetDataNodeCreator_Call {
return &MockDataCoord_SetDataNodeCreator_Call{Call: _e.mock.On("SetDataNodeCreator", _a0)}
}
func (_c *MockDataCoord_SetDataNodeCreator_Call) Run(run func(_a0 func(context.Context, string) (types.DataNode, error))) *MockDataCoord_SetDataNodeCreator_Call {
func (_c *MockDataCoord_SetDataNodeCreator_Call) Run(run func(_a0 func(context.Context, string, int64) (types.DataNode, error))) *MockDataCoord_SetDataNodeCreator_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(func(context.Context, string) (types.DataNode, error)))
run(args[0].(func(context.Context, string, int64) (types.DataNode, error)))
})
return _c
}
@ -2263,7 +2263,7 @@ func (_c *MockDataCoord_SetDataNodeCreator_Call) Return() *MockDataCoord_SetData
return _c
}
func (_c *MockDataCoord_SetDataNodeCreator_Call) RunAndReturn(run func(func(context.Context, string) (types.DataNode, error))) *MockDataCoord_SetDataNodeCreator_Call {
func (_c *MockDataCoord_SetDataNodeCreator_Call) RunAndReturn(run func(func(context.Context, string, int64) (types.DataNode, error))) *MockDataCoord_SetDataNodeCreator_Call {
_c.Call.Return(run)
return _c
}
@ -2279,7 +2279,7 @@ type MockDataCoord_SetEtcdClient_Call struct {
}
// SetEtcdClient is a helper method to define mock.On call
// - etcdClient *clientv3.Client
// - etcdClient *clientv3.Client
func (_e *MockDataCoord_Expecter) SetEtcdClient(etcdClient interface{}) *MockDataCoord_SetEtcdClient_Call {
return &MockDataCoord_SetEtcdClient_Call{Call: _e.mock.On("SetEtcdClient", etcdClient)}
}
@ -2302,7 +2302,7 @@ func (_c *MockDataCoord_SetEtcdClient_Call) RunAndReturn(run func(*clientv3.Clie
}
// SetIndexNodeCreator provides a mock function with given fields: _a0
func (_m *MockDataCoord) SetIndexNodeCreator(_a0 func(context.Context, string) (types.IndexNode, error)) {
func (_m *MockDataCoord) SetIndexNodeCreator(_a0 func(context.Context, string, int64) (types.IndexNode, error)) {
_m.Called(_a0)
}
@ -2312,14 +2312,14 @@ type MockDataCoord_SetIndexNodeCreator_Call struct {
}
// SetIndexNodeCreator is a helper method to define mock.On call
// - _a0 func(context.Context , string)(types.IndexNode , error)
// - _a0 func(context.Context , string , int64)(types.IndexNode , error)
func (_e *MockDataCoord_Expecter) SetIndexNodeCreator(_a0 interface{}) *MockDataCoord_SetIndexNodeCreator_Call {
return &MockDataCoord_SetIndexNodeCreator_Call{Call: _e.mock.On("SetIndexNodeCreator", _a0)}
}
func (_c *MockDataCoord_SetIndexNodeCreator_Call) Run(run func(_a0 func(context.Context, string) (types.IndexNode, error))) *MockDataCoord_SetIndexNodeCreator_Call {
func (_c *MockDataCoord_SetIndexNodeCreator_Call) Run(run func(_a0 func(context.Context, string, int64) (types.IndexNode, error))) *MockDataCoord_SetIndexNodeCreator_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(func(context.Context, string) (types.IndexNode, error)))
run(args[0].(func(context.Context, string, int64) (types.IndexNode, error)))
})
return _c
}
@ -2329,7 +2329,7 @@ func (_c *MockDataCoord_SetIndexNodeCreator_Call) Return() *MockDataCoord_SetInd
return _c
}
func (_c *MockDataCoord_SetIndexNodeCreator_Call) RunAndReturn(run func(func(context.Context, string) (types.IndexNode, error))) *MockDataCoord_SetIndexNodeCreator_Call {
func (_c *MockDataCoord_SetIndexNodeCreator_Call) RunAndReturn(run func(func(context.Context, string, int64) (types.IndexNode, error))) *MockDataCoord_SetIndexNodeCreator_Call {
_c.Call.Return(run)
return _c
}
@ -2345,7 +2345,7 @@ type MockDataCoord_SetRootCoord_Call struct {
}
// SetRootCoord is a helper method to define mock.On call
// - rootCoord types.RootCoord
// - rootCoord types.RootCoord
func (_e *MockDataCoord_Expecter) SetRootCoord(rootCoord interface{}) *MockDataCoord_SetRootCoord_Call {
return &MockDataCoord_SetRootCoord_Call{Call: _e.mock.On("SetRootCoord", rootCoord)}
}
@ -2399,8 +2399,8 @@ type MockDataCoord_SetSegmentState_Call struct {
}
// SetSegmentState is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.SetSegmentStateRequest
// - ctx context.Context
// - req *datapb.SetSegmentStateRequest
func (_e *MockDataCoord_Expecter) SetSegmentState(ctx interface{}, req interface{}) *MockDataCoord_SetSegmentState_Call {
return &MockDataCoord_SetSegmentState_Call{Call: _e.mock.On("SetSegmentState", ctx, req)}
}
@ -2454,8 +2454,8 @@ type MockDataCoord_ShowConfigurations_Call struct {
}
// ShowConfigurations is a helper method to define mock.On call
// - ctx context.Context
// - req *internalpb.ShowConfigurationsRequest
// - ctx context.Context
// - req *internalpb.ShowConfigurationsRequest
func (_e *MockDataCoord_Expecter) ShowConfigurations(ctx interface{}, req interface{}) *MockDataCoord_ShowConfigurations_Call {
return &MockDataCoord_ShowConfigurations_Call{Call: _e.mock.On("ShowConfigurations", ctx, req)}
}
@ -2591,8 +2591,8 @@ type MockDataCoord_UnsetIsImportingState_Call struct {
}
// UnsetIsImportingState is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.UnsetIsImportingStateRequest
// - ctx context.Context
// - req *datapb.UnsetIsImportingStateRequest
func (_e *MockDataCoord_Expecter) UnsetIsImportingState(ctx interface{}, req interface{}) *MockDataCoord_UnsetIsImportingState_Call {
return &MockDataCoord_UnsetIsImportingState_Call{Call: _e.mock.On("UnsetIsImportingState", ctx, req)}
}
@ -2646,8 +2646,8 @@ type MockDataCoord_UpdateChannelCheckpoint_Call struct {
}
// UpdateChannelCheckpoint is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.UpdateChannelCheckpointRequest
// - ctx context.Context
// - req *datapb.UpdateChannelCheckpointRequest
func (_e *MockDataCoord_Expecter) UpdateChannelCheckpoint(ctx interface{}, req interface{}) *MockDataCoord_UpdateChannelCheckpoint_Call {
return &MockDataCoord_UpdateChannelCheckpoint_Call{Call: _e.mock.On("UpdateChannelCheckpoint", ctx, req)}
}
@ -2701,8 +2701,8 @@ type MockDataCoord_UpdateSegmentStatistics_Call struct {
}
// UpdateSegmentStatistics is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.UpdateSegmentStatisticsRequest
// - ctx context.Context
// - req *datapb.UpdateSegmentStatisticsRequest
func (_e *MockDataCoord_Expecter) UpdateSegmentStatistics(ctx interface{}, req interface{}) *MockDataCoord_UpdateSegmentStatistics_Call {
return &MockDataCoord_UpdateSegmentStatistics_Call{Call: _e.mock.On("UpdateSegmentStatistics", ctx, req)}
}
@ -2756,8 +2756,8 @@ type MockDataCoord_WatchChannels_Call struct {
}
// WatchChannels is a helper method to define mock.On call
// - ctx context.Context
// - req *datapb.WatchChannelsRequest
// - ctx context.Context
// - req *datapb.WatchChannelsRequest
func (_e *MockDataCoord_Expecter) WatchChannels(ctx interface{}, req interface{}) *MockDataCoord_WatchChannels_Call {
return &MockDataCoord_WatchChannels_Call{Call: _e.mock.On("WatchChannels", ctx, req)}
}

927
internal/mocks/mock_proxy.go
View File

File diff suppressed because it is too large Load Diff

110
internal/mocks/mock_querycoord.go
View File

@ -64,8 +64,8 @@ type MockQueryCoord_CheckHealth_Call struct {
}
// CheckHealth is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.CheckHealthRequest
// - ctx context.Context
// - req *milvuspb.CheckHealthRequest
func (_e *MockQueryCoord_Expecter) CheckHealth(ctx interface{}, req interface{}) *MockQueryCoord_CheckHealth_Call {
return &MockQueryCoord_CheckHealth_Call{Call: _e.mock.On("CheckHealth", ctx, req)}
}
@ -119,8 +119,8 @@ type MockQueryCoord_CreateResourceGroup_Call struct {
}
// CreateResourceGroup is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.CreateResourceGroupRequest
// - ctx context.Context
// - req *milvuspb.CreateResourceGroupRequest
func (_e *MockQueryCoord_Expecter) CreateResourceGroup(ctx interface{}, req interface{}) *MockQueryCoord_CreateResourceGroup_Call {
return &MockQueryCoord_CreateResourceGroup_Call{Call: _e.mock.On("CreateResourceGroup", ctx, req)}
}
@ -174,8 +174,8 @@ type MockQueryCoord_DescribeResourceGroup_Call struct {
}
// DescribeResourceGroup is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.DescribeResourceGroupRequest
// - ctx context.Context
// - req *querypb.DescribeResourceGroupRequest
func (_e *MockQueryCoord_Expecter) DescribeResourceGroup(ctx interface{}, req interface{}) *MockQueryCoord_DescribeResourceGroup_Call {
return &MockQueryCoord_DescribeResourceGroup_Call{Call: _e.mock.On("DescribeResourceGroup", ctx, req)}
}
@ -229,8 +229,8 @@ type MockQueryCoord_DropResourceGroup_Call struct {
}
// DropResourceGroup is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.DropResourceGroupRequest
// - ctx context.Context
// - req *milvuspb.DropResourceGroupRequest
func (_e *MockQueryCoord_Expecter) DropResourceGroup(ctx interface{}, req interface{}) *MockQueryCoord_DropResourceGroup_Call {
return &MockQueryCoord_DropResourceGroup_Call{Call: _e.mock.On("DropResourceGroup", ctx, req)}
}
@ -284,7 +284,7 @@ type MockQueryCoord_GetComponentStates_Call struct {
}
// GetComponentStates is a helper method to define mock.On call
// - ctx context.Context
// - ctx context.Context
func (_e *MockQueryCoord_Expecter) GetComponentStates(ctx interface{}) *MockQueryCoord_GetComponentStates_Call {
return &MockQueryCoord_GetComponentStates_Call{Call: _e.mock.On("GetComponentStates", ctx)}
}
@ -338,8 +338,8 @@ type MockQueryCoord_GetMetrics_Call struct {
}
// GetMetrics is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.GetMetricsRequest
// - ctx context.Context
// - req *milvuspb.GetMetricsRequest
func (_e *MockQueryCoord_Expecter) GetMetrics(ctx interface{}, req interface{}) *MockQueryCoord_GetMetrics_Call {
return &MockQueryCoord_GetMetrics_Call{Call: _e.mock.On("GetMetrics", ctx, req)}
}
@ -393,8 +393,8 @@ type MockQueryCoord_GetPartitionStates_Call struct {
}
// GetPartitionStates is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.GetPartitionStatesRequest
// - ctx context.Context
// - req *querypb.GetPartitionStatesRequest
func (_e *MockQueryCoord_Expecter) GetPartitionStates(ctx interface{}, req interface{}) *MockQueryCoord_GetPartitionStates_Call {
return &MockQueryCoord_GetPartitionStates_Call{Call: _e.mock.On("GetPartitionStates", ctx, req)}
}
@ -448,8 +448,8 @@ type MockQueryCoord_GetReplicas_Call struct {
}
// GetReplicas is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.GetReplicasRequest
// - ctx context.Context
// - req *milvuspb.GetReplicasRequest
func (_e *MockQueryCoord_Expecter) GetReplicas(ctx interface{}, req interface{}) *MockQueryCoord_GetReplicas_Call {
return &MockQueryCoord_GetReplicas_Call{Call: _e.mock.On("GetReplicas", ctx, req)}
}
@ -503,8 +503,8 @@ type MockQueryCoord_GetSegmentInfo_Call struct {
}
// GetSegmentInfo is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.GetSegmentInfoRequest
// - ctx context.Context
// - req *querypb.GetSegmentInfoRequest
func (_e *MockQueryCoord_Expecter) GetSegmentInfo(ctx interface{}, req interface{}) *MockQueryCoord_GetSegmentInfo_Call {
return &MockQueryCoord_GetSegmentInfo_Call{Call: _e.mock.On("GetSegmentInfo", ctx, req)}
}
@ -558,8 +558,8 @@ type MockQueryCoord_GetShardLeaders_Call struct {
}
// GetShardLeaders is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.GetShardLeadersRequest
// - ctx context.Context
// - req *querypb.GetShardLeadersRequest
func (_e *MockQueryCoord_Expecter) GetShardLeaders(ctx interface{}, req interface{}) *MockQueryCoord_GetShardLeaders_Call {
return &MockQueryCoord_GetShardLeaders_Call{Call: _e.mock.On("GetShardLeaders", ctx, req)}
}
@ -613,7 +613,7 @@ type MockQueryCoord_GetStatisticsChannel_Call struct {
}
// GetStatisticsChannel is a helper method to define mock.On call
// - ctx context.Context
// - ctx context.Context
func (_e *MockQueryCoord_Expecter) GetStatisticsChannel(ctx interface{}) *MockQueryCoord_GetStatisticsChannel_Call {
return &MockQueryCoord_GetStatisticsChannel_Call{Call: _e.mock.On("GetStatisticsChannel", ctx)}
}
@ -667,7 +667,7 @@ type MockQueryCoord_GetTimeTickChannel_Call struct {
}
// GetTimeTickChannel is a helper method to define mock.On call
// - ctx context.Context
// - ctx context.Context
func (_e *MockQueryCoord_Expecter) GetTimeTickChannel(ctx interface{}) *MockQueryCoord_GetTimeTickChannel_Call {
return &MockQueryCoord_GetTimeTickChannel_Call{Call: _e.mock.On("GetTimeTickChannel", ctx)}
}
@ -762,8 +762,8 @@ type MockQueryCoord_ListResourceGroups_Call struct {
}
// ListResourceGroups is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.ListResourceGroupsRequest
// - ctx context.Context
// - req *milvuspb.ListResourceGroupsRequest
func (_e *MockQueryCoord_Expecter) ListResourceGroups(ctx interface{}, req interface{}) *MockQueryCoord_ListResourceGroups_Call {
return &MockQueryCoord_ListResourceGroups_Call{Call: _e.mock.On("ListResourceGroups", ctx, req)}
}
@ -817,8 +817,8 @@ type MockQueryCoord_LoadBalance_Call struct {
}
// LoadBalance is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.LoadBalanceRequest
// - ctx context.Context
// - req *querypb.LoadBalanceRequest
func (_e *MockQueryCoord_Expecter) LoadBalance(ctx interface{}, req interface{}) *MockQueryCoord_LoadBalance_Call {
return &MockQueryCoord_LoadBalance_Call{Call: _e.mock.On("LoadBalance", ctx, req)}
}
@ -872,8 +872,8 @@ type MockQueryCoord_LoadCollection_Call struct {
}
// LoadCollection is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.LoadCollectionRequest
// - ctx context.Context
// - req *querypb.LoadCollectionRequest
func (_e *MockQueryCoord_Expecter) LoadCollection(ctx interface{}, req interface{}) *MockQueryCoord_LoadCollection_Call {
return &MockQueryCoord_LoadCollection_Call{Call: _e.mock.On("LoadCollection", ctx, req)}
}
@ -927,8 +927,8 @@ type MockQueryCoord_LoadPartitions_Call struct {
}
// LoadPartitions is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.LoadPartitionsRequest
// - ctx context.Context
// - req *querypb.LoadPartitionsRequest
func (_e *MockQueryCoord_Expecter) LoadPartitions(ctx interface{}, req interface{}) *MockQueryCoord_LoadPartitions_Call {
return &MockQueryCoord_LoadPartitions_Call{Call: _e.mock.On("LoadPartitions", ctx, req)}
}
@ -1023,8 +1023,8 @@ type MockQueryCoord_ReleaseCollection_Call struct {
}
// ReleaseCollection is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.ReleaseCollectionRequest
// - ctx context.Context
// - req *querypb.ReleaseCollectionRequest
func (_e *MockQueryCoord_Expecter) ReleaseCollection(ctx interface{}, req interface{}) *MockQueryCoord_ReleaseCollection_Call {
return &MockQueryCoord_ReleaseCollection_Call{Call: _e.mock.On("ReleaseCollection", ctx, req)}
}
@ -1078,8 +1078,8 @@ type MockQueryCoord_ReleasePartitions_Call struct {
}
// ReleasePartitions is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.ReleasePartitionsRequest
// - ctx context.Context
// - req *querypb.ReleasePartitionsRequest
func (_e *MockQueryCoord_Expecter) ReleasePartitions(ctx interface{}, req interface{}) *MockQueryCoord_ReleasePartitions_Call {
return &MockQueryCoord_ReleasePartitions_Call{Call: _e.mock.On("ReleasePartitions", ctx, req)}
}
@ -1112,7 +1112,7 @@ type MockQueryCoord_SetAddress_Call struct {
}
// SetAddress is a helper method to define mock.On call
// - address string
// - address string
func (_e *MockQueryCoord_Expecter) SetAddress(address interface{}) *MockQueryCoord_SetAddress_Call {
return &MockQueryCoord_SetAddress_Call{Call: _e.mock.On("SetAddress", address)}
}
@ -1154,7 +1154,7 @@ type MockQueryCoord_SetDataCoord_Call struct {
}
// SetDataCoord is a helper method to define mock.On call
// - dataCoord types.DataCoord
// - dataCoord types.DataCoord
func (_e *MockQueryCoord_Expecter) SetDataCoord(dataCoord interface{}) *MockQueryCoord_SetDataCoord_Call {
return &MockQueryCoord_SetDataCoord_Call{Call: _e.mock.On("SetDataCoord", dataCoord)}
}
@ -1187,7 +1187,7 @@ type MockQueryCoord_SetEtcdClient_Call struct {
}
// SetEtcdClient is a helper method to define mock.On call
// - etcdClient *clientv3.Client
// - etcdClient *clientv3.Client
func (_e *MockQueryCoord_Expecter) SetEtcdClient(etcdClient interface{}) *MockQueryCoord_SetEtcdClient_Call {
return &MockQueryCoord_SetEtcdClient_Call{Call: _e.mock.On("SetEtcdClient", etcdClient)}
}
@ -1210,7 +1210,7 @@ func (_c *MockQueryCoord_SetEtcdClient_Call) RunAndReturn(run func(*clientv3.Cli
}
// SetQueryNodeCreator provides a mock function with given fields: _a0
func (_m *MockQueryCoord) SetQueryNodeCreator(_a0 func(context.Context, string) (types.QueryNode, error)) {
func (_m *MockQueryCoord) SetQueryNodeCreator(_a0 func(context.Context, string, int64) (types.QueryNode, error)) {
_m.Called(_a0)
}
@ -1220,14 +1220,14 @@ type MockQueryCoord_SetQueryNodeCreator_Call struct {
}
// SetQueryNodeCreator is a helper method to define mock.On call
// - _a0 func(context.Context , string)(types.QueryNode , error)
// - _a0 func(context.Context , string , int64)(types.QueryNode , error)
func (_e *MockQueryCoord_Expecter) SetQueryNodeCreator(_a0 interface{}) *MockQueryCoord_SetQueryNodeCreator_Call {
return &MockQueryCoord_SetQueryNodeCreator_Call{Call: _e.mock.On("SetQueryNodeCreator", _a0)}
}
func (_c *MockQueryCoord_SetQueryNodeCreator_Call) Run(run func(_a0 func(context.Context, string) (types.QueryNode, error))) *MockQueryCoord_SetQueryNodeCreator_Call {
func (_c *MockQueryCoord_SetQueryNodeCreator_Call) Run(run func(_a0 func(context.Context, string, int64) (types.QueryNode, error))) *MockQueryCoord_SetQueryNodeCreator_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(func(context.Context, string) (types.QueryNode, error)))
run(args[0].(func(context.Context, string, int64) (types.QueryNode, error)))
})
return _c
}
@ -1237,7 +1237,7 @@ func (_c *MockQueryCoord_SetQueryNodeCreator_Call) Return() *MockQueryCoord_SetQ
return _c
}
func (_c *MockQueryCoord_SetQueryNodeCreator_Call) RunAndReturn(run func(func(context.Context, string) (types.QueryNode, error))) *MockQueryCoord_SetQueryNodeCreator_Call {
func (_c *MockQueryCoord_SetQueryNodeCreator_Call) RunAndReturn(run func(func(context.Context, string, int64) (types.QueryNode, error))) *MockQueryCoord_SetQueryNodeCreator_Call {
_c.Call.Return(run)
return _c
}
@ -1262,7 +1262,7 @@ type MockQueryCoord_SetRootCoord_Call struct {
}
// SetRootCoord is a helper method to define mock.On call
// - rootCoord types.RootCoord
// - rootCoord types.RootCoord
func (_e *MockQueryCoord_Expecter) SetRootCoord(rootCoord interface{}) *MockQueryCoord_SetRootCoord_Call {
return &MockQueryCoord_SetRootCoord_Call{Call: _e.mock.On("SetRootCoord", rootCoord)}
}
@ -1316,8 +1316,8 @@ type MockQueryCoord_ShowCollections_Call struct {
}
// ShowCollections is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.ShowCollectionsRequest
// - ctx context.Context
// - req *querypb.ShowCollectionsRequest
func (_e *MockQueryCoord_Expecter) ShowCollections(ctx interface{}, req interface{}) *MockQueryCoord_ShowCollections_Call {
return &MockQueryCoord_ShowCollections_Call{Call: _e.mock.On("ShowCollections", ctx, req)}
}
@ -1371,8 +1371,8 @@ type MockQueryCoord_ShowConfigurations_Call struct {
}
// ShowConfigurations is a helper method to define mock.On call
// - ctx context.Context
// - req *internalpb.ShowConfigurationsRequest
// - ctx context.Context
// - req *internalpb.ShowConfigurationsRequest
func (_e *MockQueryCoord_Expecter) ShowConfigurations(ctx interface{}, req interface{}) *MockQueryCoord_ShowConfigurations_Call {
return &MockQueryCoord_ShowConfigurations_Call{Call: _e.mock.On("ShowConfigurations", ctx, req)}
}
@ -1426,8 +1426,8 @@ type MockQueryCoord_ShowPartitions_Call struct {
}
// ShowPartitions is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.ShowPartitionsRequest
// - ctx context.Context
// - req *querypb.ShowPartitionsRequest
func (_e *MockQueryCoord_Expecter) ShowPartitions(ctx interface{}, req interface{}) *MockQueryCoord_ShowPartitions_Call {
return &MockQueryCoord_ShowPartitions_Call{Call: _e.mock.On("ShowPartitions", ctx, req)}
}
@ -1563,8 +1563,8 @@ type MockQueryCoord_SyncNewCreatedPartition_Call struct {
}
// SyncNewCreatedPartition is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.SyncNewCreatedPartitionRequest
// - ctx context.Context
// - req *querypb.SyncNewCreatedPartitionRequest
func (_e *MockQueryCoord_Expecter) SyncNewCreatedPartition(ctx interface{}, req interface{}) *MockQueryCoord_SyncNewCreatedPartition_Call {
return &MockQueryCoord_SyncNewCreatedPartition_Call{Call: _e.mock.On("SyncNewCreatedPartition", ctx, req)}
}
@ -1618,8 +1618,8 @@ type MockQueryCoord_TransferNode_Call struct {
}
// TransferNode is a helper method to define mock.On call
// - ctx context.Context
// - req *milvuspb.TransferNodeRequest
// - ctx context.Context
// - req *milvuspb.TransferNodeRequest
func (_e *MockQueryCoord_Expecter) TransferNode(ctx interface{}, req interface{}) *MockQueryCoord_TransferNode_Call {
return &MockQueryCoord_TransferNode_Call{Call: _e.mock.On("TransferNode", ctx, req)}
}
@ -1673,8 +1673,8 @@ type MockQueryCoord_TransferReplica_Call struct {
}
// TransferReplica is a helper method to define mock.On call
// - ctx context.Context
// - req *querypb.TransferReplicaRequest
// - ctx context.Context
// - req *querypb.TransferReplicaRequest
func (_e *MockQueryCoord_Expecter) TransferReplica(ctx interface{}, req interface{}) *MockQueryCoord_TransferReplica_Call {
return &MockQueryCoord_TransferReplica_Call{Call: _e.mock.On("TransferReplica", ctx, req)}
}
@ -1707,7 +1707,7 @@ type MockQueryCoord_UpdateStateCode_Call struct {
}
// UpdateStateCode is a helper method to define mock.On call
// - stateCode commonpb.StateCode
// - stateCode commonpb.StateCode
func (_e *MockQueryCoord_Expecter) UpdateStateCode(stateCode interface{}) *MockQueryCoord_UpdateStateCode_Call {
return &MockQueryCoord_UpdateStateCode_Call{Call: _e.mock.On("UpdateStateCode", stateCode)}
}

2
internal/proxy/proxy.go
View File

@ -484,7 +484,7 @@ func (node *Proxy) SetQueryCoordClient(cli types.QueryCoord) {
node.queryCoord = cli
}
func (node *Proxy) SetQueryNodeCreator(f func(ctx context.Context, addr string) (types.QueryNode, error)) {
func (node *Proxy) SetQueryNodeCreator(f func(ctx context.Context, addr string, nodeID int64) (types.QueryNode, error)) {
node.shardMgr.SetClientCreatorFunc(f)
}

2
internal/proxy/proxy_rpc_test.go
View File

@ -52,7 +52,7 @@ func TestProxyRpcLimit(t *testing.T) {
go testServer.startGrpc(ctx, &wg, &p)
assert.NoError(t, testServer.waitForGrpcReady())
defer testServer.grpcServer.Stop()
client, err := grpcproxyclient.NewClient(ctx, "localhost:"+p.Port.GetValue())
client, err := grpcproxyclient.NewClient(ctx, "localhost:"+p.Port.GetValue(), 1)
assert.NoError(t, err)
proxy.stateCode.Store(commonpb.StateCode_Healthy)

8
internal/proxy/shard_client.go
View File

@ -11,7 +11,7 @@ import (
"github.com/milvus-io/milvus/internal/types"
)
type queryNodeCreatorFunc func(ctx context.Context, addr string) (types.QueryNode, error)
type queryNodeCreatorFunc func(ctx context.Context, addr string, nodeID int64) (types.QueryNode, error)
type nodeInfo struct {
nodeID UniqueID
@ -114,8 +114,8 @@ func withShardClientCreator(creator queryNodeCreatorFunc) shardClientMgrOpt {
return func(s shardClientMgr) { s.SetClientCreatorFunc(creator) }
}
func defaultQueryNodeClientCreator(ctx context.Context, addr string) (types.QueryNode, error) {
return qnClient.NewClient(ctx, addr)
func defaultQueryNodeClientCreator(ctx context.Context, addr string, nodeID int64) (types.QueryNode, error) {
return qnClient.NewClient(ctx, addr, nodeID)
}
// NewShardClientMgr creates a new shardClientMgr
@ -178,7 +178,7 @@ func (c *shardClientMgrImpl) UpdateShardLeaders(oldLeaders map[string][]nodeInfo
if c.clientCreator == nil {
return fmt.Errorf("clientCreator function is nil")
}
shardClient, err := c.clientCreator(context.Background(), node.address)
shardClient, err := c.clientCreator(context.Background(), node.address, node.nodeID)
if err != nil {
return err
}

2
internal/proxy/shard_client_test.go
View File

@ -31,7 +31,7 @@ func TestShardClientMgr_UpdateShardLeaders_CreatorNil(t *testing.T) {
}
func TestShardClientMgr_UpdateShardLeaders_Empty(t *testing.T) {
mockCreator := func(ctx context.Context, addr string) (types.QueryNode, error) {
mockCreator := func(ctx context.Context, addr string, nodeID int64) (types.QueryNode, error) {
return &mock.QueryNodeClient{}, nil
}
mgr := newShardClientMgr(withShardClientCreator(mockCreator))

2
internal/querycoordv2/server.go
View File

@ -569,7 +569,7 @@ func (s *Server) SetDataCoord(dataCoord types.DataCoord) error {
return nil
}
func (s *Server) SetQueryNodeCreator(f func(ctx context.Context, addr string) (types.QueryNode, error)) {
func (s *Server) SetQueryNodeCreator(f func(ctx context.Context, addr string, nodeID int64) (types.QueryNode, error)) {
s.queryNodeCreator = f
}

12
internal/querycoordv2/session/cluster.go
View File

@ -72,10 +72,10 @@ type QueryCluster struct {
stopOnce sync.Once
}
type QueryNodeCreator func(ctx context.Context, addr string) (types.QueryNode, error)
type QueryNodeCreator func(ctx context.Context, addr string, nodeID int64) (types.QueryNode, error)
func DefaultQueryNodeCreator(ctx context.Context, addr string) (types.QueryNode, error) {
return grpcquerynodeclient.NewClient(ctx, addr)
func DefaultQueryNodeCreator(ctx context.Context, addr string, nodeID int64) (types.QueryNode, error) {
return grpcquerynodeclient.NewClient(ctx, addr, nodeID)
}
func NewCluster(nodeManager *NodeManager, queryNodeCreator QueryNodeCreator) *QueryCluster {
@ -303,8 +303,8 @@ func (c *clients) getOrCreate(ctx context.Context, node *NodeInfo) (types.QueryN
return c.create(node)
}
func createNewClient(ctx context.Context, addr string, queryNodeCreator QueryNodeCreator) (types.QueryNode, error) {
newCli, err := queryNodeCreator(ctx, addr)
func createNewClient(ctx context.Context, addr string, nodeID int64, queryNodeCreator QueryNodeCreator) (types.QueryNode, error) {
newCli, err := queryNodeCreator(ctx, addr, nodeID)
if err != nil {
return nil, err
}
@ -323,7 +323,7 @@ func (c *clients) create(node *NodeInfo) (types.QueryNode, error) {
if cli, ok := c.clients[node.ID()]; ok {
return cli, nil
}
cli, err := createNewClient(context.Background(), node.Addr(), c.queryNodeCreator)
cli, err := createNewClient(context.Background(), node.Addr(), node.ID(), c.queryNodeCreator)
if err != nil {
return nil, err
}

2
internal/querynodev2/server.go
View File

@ -315,7 +315,7 @@ func (node *QueryNode) Init() error {
}
}
client, err := grpcquerynodeclient.NewClient(node.ctx, addr)
client, err := grpcquerynodeclient.NewClient(node.ctx, addr, nodeID)
if err != nil {
return nil, err
}

8
internal/rootcoord/proxy_client_manager.go
View File

@ -36,10 +36,10 @@ import (
"github.com/milvus-io/milvus/pkg/util/metricsinfo"
)
type proxyCreator func(ctx context.Context, addr string) (types.Proxy, error)
type proxyCreator func(ctx context.Context, addr string, nodeID int64) (types.Proxy, error)
func DefaultProxyCreator(ctx context.Context, addr string) (types.Proxy, error) {
cli, err := grpcproxyclient.NewClient(ctx, addr)
func DefaultProxyCreator(ctx context.Context, addr string, nodeID int64) (types.Proxy, error) {
cli, err := grpcproxyclient.NewClient(ctx, addr, nodeID)
if err != nil {
return nil, err
}
@ -107,7 +107,7 @@ func (p *proxyClientManager) updateProxyNumMetric() {
}
func (p *proxyClientManager) connect(session *sessionutil.Session) {
pc, err := p.creator(context.Background(), session.Address)
pc, err := p.creator(context.Background(), session.Address, session.ServerID)
if err != nil {
log.Warn("failed to create proxy client", zap.String("address", session.Address), zap.Int64("serverID", session.ServerID), zap.Error(err))
return

4
internal/rootcoord/proxy_client_manager_test.go
View File

@ -116,7 +116,7 @@ func TestProxyClientManager_GetProxyClients(t *testing.T) {
defer cli.Close()
assert.NoError(t, err)
core.etcdCli = cli
core.proxyCreator = func(ctx context.Context, addr string) (types.Proxy, error) {
core.proxyCreator = func(ctx context.Context, addr string, nodeID int64) (types.Proxy, error) {
return nil, errors.New("failed")
}
@ -148,7 +148,7 @@ func TestProxyClientManager_AddProxyClient(t *testing.T) {
defer cli.Close()
core.etcdCli = cli
core.proxyCreator = func(ctx context.Context, addr string) (types.Proxy, error) {
core.proxyCreator = func(ctx context.Context, addr string, nodeID int64) (types.Proxy, error) {
return nil, errors.New("failed")
}

2
internal/rootcoord/root_coord.go
View File

@ -252,7 +252,7 @@ func (c *Core) tsLoop() {
}
}
func (c *Core) SetProxyCreator(f func(ctx context.Context, addr string) (types.Proxy, error)) {
func (c *Core) SetProxyCreator(f func(ctx context.Context, addr string, nodeID int64) (types.Proxy, error)) {
c.proxyCreator = f
}

10
internal/types/types.go
View File

@ -396,10 +396,10 @@ type DataCoordComponent interface {
SetRootCoord(rootCoord RootCoord)
// SetDataNodeCreator set DataNode client creator func for DataCoord
SetDataNodeCreator(func(context.Context, string) (DataNode, error))
SetDataNodeCreator(func(context.Context, string, int64) (DataNode, error))
//SetIndexNodeCreator set Index client creator func for DataCoord
SetIndexNodeCreator(func(context.Context, string) (IndexNode, error))
SetIndexNodeCreator(func(context.Context, string, int64) (IndexNode, error))
}
// IndexNode is the interface `indexnode` package implements
@ -837,7 +837,7 @@ type RootCoordComponent interface {
SetQueryCoord(queryCoord QueryCoord) error
// SetProxyCreator set Proxy client creator func for RootCoord
SetProxyCreator(func(ctx context.Context, addr string) (Proxy, error))
SetProxyCreator(func(ctx context.Context, addr string, nodeID int64) (Proxy, error))
// GetMetrics notifies RootCoordComponent to collect metrics for specified component
GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)
@ -917,7 +917,7 @@ type ProxyComponent interface {
SetQueryCoordClient(queryCoord QueryCoord)
// SetQueryNodeCreator set QueryNode client creator func for Proxy
SetQueryNodeCreator(func(ctx context.Context, addr string) (QueryNode, error))
SetQueryNodeCreator(func(ctx context.Context, addr string, nodeID int64) (QueryNode, error))
// GetRateLimiter returns the rateLimiter in Proxy
GetRateLimiter() (Limiter, error)
@ -1542,5 +1542,5 @@ type QueryCoordComponent interface {
SetRootCoord(rootCoord RootCoord) error
// SetQueryNodeCreator set QueryNode client creator func for QueryCoord
SetQueryNodeCreator(func(ctx context.Context, addr string) (QueryNode, error))
SetQueryNodeCreator(func(ctx context.Context, addr string, nodeID int64) (QueryNode, error))
}

28
internal/util/grpcclient/client.go
View File

@ -201,10 +201,12 @@ func (c *ClientBase[T]) connect(ctx context.Context) error {
grpc.WithUnaryInterceptor(grpc_middleware.ChainUnaryClient(
otelgrpc.UnaryClientInterceptor(opts...),
interceptor.ClusterInjectionUnaryClientInterceptor(),
interceptor.ServerIDInjectionUnaryClientInterceptor(c.GetNodeID()),
)),
grpc.WithStreamInterceptor(grpc_middleware.ChainStreamClient(
otelgrpc.StreamClientInterceptor(opts...),
interceptor.ClusterInjectionStreamClientInterceptor(),
interceptor.ServerIDInjectionStreamClientInterceptor(c.GetNodeID()),
)),
grpc.WithDefaultServiceConfig(retryPolicy),
grpc.WithKeepaliveParams(keepalive.ClientParameters{
@ -239,10 +241,12 @@ func (c *ClientBase[T]) connect(ctx context.Context) error {
grpc.WithUnaryInterceptor(grpc_middleware.ChainUnaryClient(
otelgrpc.UnaryClientInterceptor(opts...),
interceptor.ClusterInjectionUnaryClientInterceptor(),
interceptor.ServerIDInjectionUnaryClientInterceptor(c.GetNodeID()),
)),
grpc.WithStreamInterceptor(grpc_middleware.ChainStreamClient(
otelgrpc.StreamClientInterceptor(opts...),
interceptor.ClusterInjectionStreamClientInterceptor(),
interceptor.ServerIDInjectionStreamClientInterceptor(c.GetNodeID()),
)),
grpc.WithDefaultServiceConfig(retryPolicy),
grpc.WithKeepaliveParams(keepalive.ClientParameters{
@ -279,6 +283,7 @@ func (c *ClientBase[T]) connect(ctx context.Context) error {
}
func (c *ClientBase[T]) callOnce(ctx context.Context, caller func(client T) (any, error)) (any, error) {
log := log.Ctx(ctx).With(zap.String("role", c.GetRole()))
client, err := c.GetGrpcClient(ctx)
if err != nil {
return generic.Zero[T](), err
@ -295,21 +300,20 @@ func (c *ClientBase[T]) callOnce(ctx context.Context, caller func(client T) (any
return generic.Zero[T](), err
}
if IsCrossClusterRoutingErr(err) {
log.Ctx(ctx).Warn("CrossClusterRoutingErr, start to reset connection",
zap.String("role", c.GetRole()),
zap.Error(err),
)
log.Warn("CrossClusterRoutingErr, start to reset connection", zap.Error(err))
c.resetConnection(client)
return ret, merr.ErrServiceUnavailable // For concealing ErrCrossClusterRouting from the client
}
if IsServerIDMismatchErr(err) {
log.Warn("Server ID mismatch, start to reset connection", zap.Error(err))
c.resetConnection(client)
return ret, err
}
if !funcutil.IsGrpcErr(err) {
log.Ctx(ctx).Warn("ClientBase:isNotGrpcErr", zap.Error(err))
log.Warn("ClientBase:isNotGrpcErr", zap.Error(err))
return generic.Zero[T](), err
}
log.Ctx(ctx).Info("ClientBase grpc error, start to reset connection",
zap.String("role", c.GetRole()),
zap.Error(err),
)
log.Info("ClientBase grpc error, start to reset connection", zap.Error(err))
c.resetConnection(client)
return ret, err
}
@ -398,3 +402,9 @@ func IsCrossClusterRoutingErr(err error) bool {
// hence it is not viable to employ the `errors.Is` for assessment.
return strings.Contains(err.Error(), merr.ErrCrossClusterRouting.Error())
}
func IsServerIDMismatchErr(err error) bool {
// GRPC utilizes `status.Status` to encapsulate errors,
// hence it is not viable to employ the `errors.Is` for assessment.
return strings.Contains(err.Error(), merr.ErrServerIDMismatch.Error())
}

19
pkg/util/interceptor/cluster_interceptor_test.go
View File

@ -29,25 +29,6 @@ import (
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
type mockSS struct {
grpc.ServerStream
ctx context.Context
}
func newMockSS(ctx context.Context) grpc.ServerStream {
return &mockSS{
ctx: ctx,
}
}
func (m *mockSS) Context() context.Context {
return m.ctx
}
func init() {
paramtable.Get().Init()
}
func TestClusterInterceptor(t *testing.T) {
t.Run("test ClusterInjectionUnaryClientInterceptor", func(t *testing.T) {
method := "MockMethod"

44
pkg/util/interceptor/interceptor_test.go Normal file
View File

@ -0,0 +1,44 @@
// 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 interceptor
import (
"context"
"google.golang.org/grpc"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
type mockSS struct {
grpc.ServerStream
ctx context.Context
}
func newMockSS(ctx context.Context) grpc.ServerStream {
return &mockSS{
ctx: ctx,
}
}
func (m *mockSS) Context() context.Context {
return m.ctx
}
func init() {
paramtable.Get().Init()
}

95
pkg/util/interceptor/server_id_interceptor.go Normal file
View File

@ -0,0 +1,95 @@
// 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 interceptor
import (
"context"
"fmt"
"strconv"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
const ServerIDKey = "ServerID"
// ServerIDValidationUnaryServerInterceptor returns a new unary server interceptor that
// verifies whether the target server ID of request matches with the server's ID and rejects it accordingly.
func ServerIDValidationUnaryServerInterceptor() grpc.UnaryServerInterceptor {
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
return handler(ctx, req)
}
values := md.Get(ServerIDKey)
if len(values) == 0 {
return handler(ctx, req)
}
serverID, err := strconv.ParseInt(values[0], 10, 64)
if err != nil {
return handler(ctx, req)
}
if serverID != paramtable.GetNodeID() {
return nil, merr.WrapErrServerIDMismatch(serverID, paramtable.GetNodeID())
}
return handler(ctx, req)
}
}
// ServerIDValidationStreamServerInterceptor returns a new streaming server interceptor that
// verifies whether the target server ID of request matches with the server's ID and rejects it accordingly.
func ServerIDValidationStreamServerInterceptor() grpc.StreamServerInterceptor {
return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
md, ok := metadata.FromIncomingContext(ss.Context())
if !ok {
return handler(srv, ss)
}
values := md.Get(ServerIDKey)
if len(values) == 0 {
return handler(srv, ss)
}
serverID, err := strconv.ParseInt(values[0], 10, 64)
if err != nil {
return handler(srv, ss)
}
if serverID != paramtable.GetNodeID() {
return merr.WrapErrServerIDMismatch(serverID, paramtable.GetNodeID())
}
return handler(srv, ss)
}
}
// ServerIDInjectionUnaryClientInterceptor returns a new unary client interceptor that
// injects target server ID into the request.
func ServerIDInjectionUnaryClientInterceptor(targetServerID int64) grpc.UnaryClientInterceptor {
return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
ctx = metadata.AppendToOutgoingContext(ctx, ServerIDKey, fmt.Sprint(targetServerID))
return invoker(ctx, method, req, reply, cc, opts...)
}
}
// ServerIDInjectionStreamClientInterceptor returns a new streaming client interceptor that
// injects target server ID into the request.
func ServerIDInjectionStreamClientInterceptor(targetServerID int64) grpc.StreamClientInterceptor {
return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
ctx = metadata.AppendToOutgoingContext(ctx, ServerIDKey, fmt.Sprint(targetServerID))
return streamer(ctx, desc, cc, method, opts...)
}
}

144
pkg/util/interceptor/server_id_interceptor_test.go Normal file
View File

@ -0,0 +1,144 @@
// 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 interceptor
import (
"context"
"fmt"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
func TestServerIDInterceptor(t *testing.T) {
t.Run("test ServerIDInjectionUnaryClientInterceptor", func(t *testing.T) {
method := "MockMethod"
req := &milvuspb.InsertRequest{}
serverID := int64(1)
var incomingContext context.Context
invoker := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
incomingContext = ctx
return nil
}
interceptor := ServerIDInjectionUnaryClientInterceptor(serverID)
ctx := metadata.NewOutgoingContext(context.Background(), metadata.New(make(map[string]string)))
err := interceptor(ctx, method, req, nil, nil, invoker)
assert.NoError(t, err)
md, ok := metadata.FromOutgoingContext(incomingContext)
assert.True(t, ok)
assert.Equal(t, fmt.Sprint(serverID), md.Get(ServerIDKey)[0])
})
t.Run("test ServerIDInjectionStreamClientInterceptor", func(t *testing.T) {
method := "MockMethod"
serverID := int64(1)
var incomingContext context.Context
streamer := func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
incomingContext = ctx
return nil, nil
}
interceptor := ServerIDInjectionStreamClientInterceptor(serverID)
ctx := metadata.NewOutgoingContext(context.Background(), metadata.New(make(map[string]string)))
_, err := interceptor(ctx, nil, nil, method, streamer)
assert.NoError(t, err)
md, ok := metadata.FromOutgoingContext(incomingContext)
assert.True(t, ok)
assert.Equal(t, fmt.Sprint(serverID), md.Get(ServerIDKey)[0])
})
t.Run("test ServerIDValidationUnaryServerInterceptor", func(t *testing.T) {
method := "MockMethod"
req := &milvuspb.InsertRequest{}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return nil, nil
}
serverInfo := &grpc.UnaryServerInfo{FullMethod: method}
interceptor := ServerIDValidationUnaryServerInterceptor()
// no md in context
_, err := interceptor(context.Background(), req, serverInfo, handler)
assert.NoError(t, err)
// no ServerID in md
ctx := metadata.NewIncomingContext(context.Background(), metadata.New(make(map[string]string)))
_, err = interceptor(ctx, req, serverInfo, handler)
assert.NoError(t, err)
// with invalid ServerID
md := metadata.Pairs(ServerIDKey, "@$#$%")
ctx = metadata.NewIncomingContext(context.Background(), md)
_, err = interceptor(ctx, req, serverInfo, handler)
assert.NoError(t, err)
// with mismatch ServerID
md = metadata.Pairs(ServerIDKey, "1234")
ctx = metadata.NewIncomingContext(context.Background(), md)
_, err = interceptor(ctx, req, serverInfo, handler)
assert.ErrorIs(t, err, merr.ErrServerIDMismatch)
// with same ServerID
md = metadata.Pairs(ServerIDKey, fmt.Sprint(paramtable.GetNodeID()))
ctx = metadata.NewIncomingContext(context.Background(), md)
_, err = interceptor(ctx, req, serverInfo, handler)
assert.NoError(t, err)
})
t.Run("test ServerIDValidationUnaryServerInterceptor", func(t *testing.T) {
handler := func(srv interface{}, stream grpc.ServerStream) error {
return nil
}
interceptor := ServerIDValidationStreamServerInterceptor()
// no md in context
err := interceptor(nil, newMockSS(context.Background()), nil, handler)
assert.NoError(t, err)
// no ServerID in md
ctx := metadata.NewIncomingContext(context.Background(), metadata.New(make(map[string]string)))
err = interceptor(nil, newMockSS(ctx), nil, handler)
assert.NoError(t, err)
// with invalid ServerID
md := metadata.Pairs(ServerIDKey, "@$#$%")
ctx = metadata.NewIncomingContext(context.Background(), md)
err = interceptor(nil, newMockSS(ctx), nil, handler)
assert.NoError(t, err)
// with mismatch ServerID
md = metadata.Pairs(ServerIDKey, "1234")
ctx = metadata.NewIncomingContext(context.Background(), md)
err = interceptor(nil, newMockSS(ctx), nil, handler)
assert.ErrorIs(t, err, merr.ErrServerIDMismatch)
// with same ServerID
md = metadata.Pairs(ServerIDKey, fmt.Sprint(paramtable.GetNodeID()))
ctx = metadata.NewIncomingContext(context.Background(), md)
err = interceptor(nil, newMockSS(ctx), nil, handler)
assert.NoError(t, err)
})
}

1
pkg/util/merr/errors.go
View File

@ -55,6 +55,7 @@ var (
ErrServiceInternal = newMilvusError("service internal error", 5, false) // Never return this error out of Milvus
ErrCrossClusterRouting = newMilvusError("cross cluster routing", 6, false)
ErrServiceDiskLimitExceeded = newMilvusError("disk limit exceeded", 7, false)
ErrServerIDMismatch = newMilvusError("server ID mismatch", 8, false)
// Collection related
ErrCollectionNotFound = newMilvusError("collection not found", 100, false)

1
pkg/util/merr/errors_test.go
View File

@ -77,6 +77,7 @@ func (s *ErrSuite) TestWrap() {
s.ErrorIs(WrapErrServiceInternal("never throw out"), ErrServiceInternal)
s.ErrorIs(WrapErrCrossClusterRouting("ins-0", "ins-1"), ErrCrossClusterRouting)
s.ErrorIs(WrapErrServiceDiskLimitExceeded(110, 100, "DLE"), ErrServiceDiskLimitExceeded)
s.ErrorIs(WrapErrServerIDMismatch(0, 1, "SIM"), ErrServerIDMismatch)
// Collection related
s.ErrorIs(WrapErrCollectionNotFound("test_collection", "failed to get collection"), ErrCollectionNotFound)

8
pkg/util/merr/utils.go
View File

@ -196,6 +196,14 @@ func WrapErrServiceDiskLimitExceeded(predict, limit float32, msg ...string) erro
return err
}
func WrapErrServerIDMismatch(expectedID, actualID int64, msg ...string) error {
err := errors.Wrapf(ErrServerIDMismatch, "expected=%s, actual=%s", expectedID, actualID)
if len(msg) > 0 {
err = errors.Wrap(err, strings.Join(msg, "; "))
}
return err
}
func WrapErrDatabaseNotFound(database any, msg ...string) error {
err := wrapWithField(ErrDatabaseNotfound, "database", database)
if len(msg) > 0 {

8
tests/integration/crossclusterrouting/cross_cluster_routing_test.go
View File

@ -118,13 +118,13 @@ func (s *CrossClusterRoutingSuite) SetupTest() {
s.NoError(err)
s.queryCoordClient, err = grpcquerycoordclient.NewClient(s.ctx, metaRoot, s.client)
s.NoError(err)
s.proxyClient, err = grpcproxyclient.NewClient(s.ctx, paramtable.Get().ProxyGrpcClientCfg.GetInternalAddress())
s.proxyClient, err = grpcproxyclient.NewClient(s.ctx, paramtable.Get().ProxyGrpcClientCfg.GetInternalAddress(), 1)
s.NoError(err)
s.dataNodeClient, err = grpcdatanodeclient.NewClient(s.ctx, paramtable.Get().DataNodeGrpcClientCfg.GetAddress())
s.dataNodeClient, err = grpcdatanodeclient.NewClient(s.ctx, paramtable.Get().DataNodeGrpcClientCfg.GetAddress(), 1)
s.NoError(err)
s.queryNodeClient, err = grpcquerynodeclient.NewClient(s.ctx, paramtable.Get().QueryNodeGrpcClientCfg.GetAddress())
s.queryNodeClient, err = grpcquerynodeclient.NewClient(s.ctx, paramtable.Get().QueryNodeGrpcClientCfg.GetAddress(), 1)
s.NoError(err)
s.indexNodeClient, err = grpcindexnodeclient.NewClient(s.ctx, paramtable.Get().IndexNodeGrpcClientCfg.GetAddress(), false)
s.indexNodeClient, err = grpcindexnodeclient.NewClient(s.ctx, paramtable.Get().IndexNodeGrpcClientCfg.GetAddress(), 1, false)
s.NoError(err)
// setup servers

8
tests/integration/minicluster.go
View File

@ -1179,7 +1179,7 @@ func (cluster *MiniCluster) UpdateClusterSize(clusterConfig ClusterConfig) error
return nil
}
func (cluster *MiniCluster) GetProxy(ctx context.Context, addr string) (types.Proxy, error) {
func (cluster *MiniCluster) GetProxy(ctx context.Context, addr string, nodeID int64) (types.Proxy, error) {
cluster.mu.RLock()
defer cluster.mu.RUnlock()
if cluster.Proxy.GetAddress() == addr {
@ -1188,7 +1188,7 @@ func (cluster *MiniCluster) GetProxy(ctx context.Context, addr string) (types.Pr
return nil, nil
}
func (cluster *MiniCluster) GetQueryNode(ctx context.Context, addr string) (types.QueryNode, error) {
func (cluster *MiniCluster) GetQueryNode(ctx context.Context, addr string, nodeID int64) (types.QueryNode, error) {
cluster.mu.RLock()
defer cluster.mu.RUnlock()
for _, queryNode := range cluster.QueryNodes {
@ -1199,7 +1199,7 @@ func (cluster *MiniCluster) GetQueryNode(ctx context.Context, addr string) (type
return nil, errors.New("no related queryNode found")
}
func (cluster *MiniCluster) GetDataNode(ctx context.Context, addr string) (types.DataNode, error) {
func (cluster *MiniCluster) GetDataNode(ctx context.Context, addr string, nodeID int64) (types.DataNode, error) {
cluster.mu.RLock()
defer cluster.mu.RUnlock()
for _, dataNode := range cluster.DataNodes {
@ -1210,7 +1210,7 @@ func (cluster *MiniCluster) GetDataNode(ctx context.Context, addr string) (types
return nil, errors.New("no related dataNode found")
}
func (cluster *MiniCluster) GetIndexNode(ctx context.Context, addr string) (types.IndexNode, error) {
func (cluster *MiniCluster) GetIndexNode(ctx context.Context, addr string, nodeID int64) (types.IndexNode, error) {
cluster.mu.RLock()
defer cluster.mu.RUnlock()
for _, indexNode := range cluster.IndexNodes {