milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

enhance: Move datacoord broker into separate package (#28876) 

See also #28861

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
pull/28812/head^2
congqixia 2023-12-01 10:22:34 +08:00 committed by GitHub
parent 8036ee13fa
commit 038eebba4d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 301 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
internal/datacoord/coordinator_broker.go → internal/datacoord/broker/coordinator_broker.go
View File

@ -13,7 +13,8 @@
// 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 datacoord
package broker
import (
"context"
@ -39,19 +40,21 @@ type Broker interface {
HasCollection(ctx context.Context, collectionID int64) (bool, error)
}
type CoordinatorBroker struct {
type coordinatorBroker struct {
rootCoord types.RootCoordClient
}
func NewCoordinatorBroker(rootCoord types.RootCoordClient) *CoordinatorBroker {
return &CoordinatorBroker{
func NewCoordinatorBroker(rootCoord types.RootCoordClient) *coordinatorBroker {
return &coordinatorBroker{
rootCoord: rootCoord,
}
}
func (b *CoordinatorBroker) DescribeCollectionInternal(ctx context.Context, collectionID int64) (*milvuspb.DescribeCollectionResponse, error) {
func (b *coordinatorBroker) DescribeCollectionInternal(ctx context.Context, collectionID int64) (*milvuspb.DescribeCollectionResponse, error) {
ctx, cancel := context.WithTimeout(ctx, paramtable.Get().QueryCoordCfg.BrokerTimeout.GetAsDuration(time.Millisecond))
defer cancel()
log := log.Ctx(ctx).With(zap.Int64("collectionID", collectionID))
resp, err := b.rootCoord.DescribeCollectionInternal(ctx, &milvuspb.DescribeCollectionRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_DescribeCollection),
@ -60,19 +63,19 @@ func (b *CoordinatorBroker) DescribeCollectionInternal(ctx context.Context, coll
// please do not specify the collection name alone after database feature.
CollectionID: collectionID,
})
if err = VerifyResponse(resp, err); err != nil {
log.Warn("DescribeCollectionInternal failed",
zap.Int64("collectionID", collectionID),
zap.Error(err))
if err := merr.CheckRPCCall(resp, err); err != nil {
log.Warn("DescribeCollectionInternal failed", zap.Error(err))
return nil, err
}
return resp, nil
}
func (b *CoordinatorBroker) ShowPartitionsInternal(ctx context.Context, collectionID int64) ([]int64, error) {
func (b *coordinatorBroker) ShowPartitionsInternal(ctx context.Context, collectionID int64) ([]int64, error) {
ctx, cancel := context.WithTimeout(ctx, paramtable.Get().QueryCoordCfg.BrokerTimeout.GetAsDuration(time.Millisecond))
defer cancel()
log := log.Ctx(ctx).With(zap.Int64("collectionID", collectionID))
resp, err := b.rootCoord.ShowPartitionsInternal(ctx, &milvuspb.ShowPartitionsRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_ShowPartitions),
@ -81,19 +84,20 @@ func (b *CoordinatorBroker) ShowPartitionsInternal(ctx context.Context, collecti
// please do not specify the collection name alone after database feature.
CollectionID: collectionID,
})
if err = VerifyResponse(resp, err); err != nil {
if err := merr.CheckRPCCall(resp, err); err != nil {
log.Warn("ShowPartitionsInternal failed",
zap.Int64("collectionID", collectionID),
zap.Error(err))
return nil, err
}
return resp.PartitionIDs, nil
return resp.GetPartitionIDs(), nil
}
func (b *CoordinatorBroker) ShowCollections(ctx context.Context, dbName string) (*milvuspb.ShowCollectionsResponse, error) {
func (b *coordinatorBroker) ShowCollections(ctx context.Context, dbName string) (*milvuspb.ShowCollectionsResponse, error) {
ctx, cancel := context.WithTimeout(ctx, paramtable.Get().QueryCoordCfg.BrokerTimeout.GetAsDuration(time.Millisecond))
defer cancel()
log := log.Ctx(ctx).With(zap.String("dbName", dbName))
resp, err := b.rootCoord.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_ShowCollections),
@ -101,7 +105,7 @@ func (b *CoordinatorBroker) ShowCollections(ctx context.Context, dbName string)
DbName: dbName,
})
if err = VerifyResponse(resp, err); err != nil {
if err := merr.CheckRPCCall(resp, err); err != nil {
log.Warn("ShowCollections failed",
zap.String("dbName", dbName),
zap.Error(err))
@ -111,13 +115,14 @@ func (b *CoordinatorBroker) ShowCollections(ctx context.Context, dbName string)
return resp, nil
}
func (b *CoordinatorBroker) ListDatabases(ctx context.Context) (*milvuspb.ListDatabasesResponse, error) {
func (b *coordinatorBroker) ListDatabases(ctx context.Context) (*milvuspb.ListDatabasesResponse, error) {
ctx, cancel := context.WithTimeout(ctx, paramtable.Get().QueryCoordCfg.BrokerTimeout.GetAsDuration(time.Millisecond))
defer cancel()
log := log.Ctx(ctx)
resp, err := b.rootCoord.ListDatabases(ctx, &milvuspb.ListDatabasesRequest{
Base: commonpbutil.NewMsgBase(commonpbutil.WithMsgType(commonpb.MsgType_ListDatabases)),
})
if err = VerifyResponse(resp, err); err != nil {
if err := merr.CheckRPCCall(resp, err); err != nil {
log.Warn("failed to ListDatabases", zap.Error(err))
return nil, err
}
@ -125,7 +130,7 @@ func (b *CoordinatorBroker) ListDatabases(ctx context.Context) (*milvuspb.ListDa
}
// HasCollection communicates with RootCoord and check whether this collection exist from the user's perspective.
func (b *CoordinatorBroker) HasCollection(ctx context.Context, collectionID int64) (bool, error) {
func (b *coordinatorBroker) HasCollection(ctx context.Context, collectionID int64) (bool, error) {
ctx, cancel := context.WithTimeout(ctx, paramtable.Get().QueryCoordCfg.BrokerTimeout.GetAsDuration(time.Millisecond))
defer cancel()
resp, err := b.rootCoord.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{
@ -139,9 +144,6 @@ func (b *CoordinatorBroker) HasCollection(ctx context.Context, collectionID int6
if err != nil {
return false, err
}
if resp == nil {
return false, errNilResponse
}
err = merr.Error(resp.GetStatus())
if errors.Is(err, merr.ErrCollectionNotFound) {
return false, nil

273
internal/datacoord/broker/coordinator_broker_test.go Normal file
View File

@ -0,0 +1,273 @@
// 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 broker
import (
"context"
"fmt"
"math/rand"
"testing"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"google.golang.org/grpc"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
type BrokerSuite struct {
suite.Suite
rootCoordClient *mocks.MockRootCoordClient
broker *coordinatorBroker
}
func (s *BrokerSuite) SetupSuite() {
paramtable.Init()
}
func (s *BrokerSuite) SetupTest() {
s.rootCoordClient = mocks.NewMockRootCoordClient(s.T())
s.broker = NewCoordinatorBroker(s.rootCoordClient)
}
func (s *BrokerSuite) TearDownTest() {
if s.rootCoordClient != nil {
s.rootCoordClient.AssertExpectations(s.T())
}
s.rootCoordClient = nil
}
func (s *BrokerSuite) TestDescribeCollectionInternal() {
s.Run("return_success", func() {
s.SetupTest()
collID := int64(1000 + rand.Intn(500))
s.rootCoordClient.EXPECT().DescribeCollectionInternal(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.DescribeCollectionRequest, options ...grpc.CallOption) (*milvuspb.DescribeCollectionResponse, error) {
s.Equal(collID, req.GetCollectionID())
return &milvuspb.DescribeCollectionResponse{
Status: merr.Status(nil),
CollectionID: collID,
CollectionName: "test_collection",
}, nil
})
resp, err := s.broker.DescribeCollectionInternal(context.Background(), collID)
s.NoError(err)
s.Equal(collID, resp.GetCollectionID())
s.Equal("test_collection", resp.GetCollectionName())
s.TearDownTest()
})
s.Run("return_error", func() {
s.SetupTest()
collID := int64(1000 + rand.Intn(500))
s.rootCoordClient.EXPECT().DescribeCollectionInternal(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.DescribeCollectionRequest, options ...grpc.CallOption) (*milvuspb.DescribeCollectionResponse, error) {
s.Equal(collID, req.GetCollectionID())
return nil, errors.New("mocked")
})
_, err := s.broker.DescribeCollectionInternal(context.Background(), collID)
s.Error(err)
s.TearDownTest()
})
}
func (s *BrokerSuite) TestShowPartitionsInternal() {
s.Run("return_success", func() {
s.SetupTest()
collID := int64(1000 + rand.Intn(500))
s.rootCoordClient.EXPECT().ShowPartitionsInternal(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.ShowPartitionsRequest, options ...grpc.CallOption) (*milvuspb.ShowPartitionsResponse, error) {
s.Equal(collID, req.GetCollectionID())
return &milvuspb.ShowPartitionsResponse{
Status: merr.Status(nil),
PartitionIDs: []int64{1, 2, 3},
PartitionNames: []string{"_default_1", "_default_2", "_default_3"},
CreatedTimestamps: []uint64{0, 0, 0},
}, nil
})
resp, err := s.broker.ShowPartitionsInternal(context.Background(), collID)
s.NoError(err)
s.ElementsMatch([]int64{1, 2, 3}, resp)
s.TearDownTest()
})
s.Run("return_error", func() {
s.SetupTest()
collID := int64(1000 + rand.Intn(500))
s.rootCoordClient.EXPECT().ShowPartitionsInternal(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.ShowPartitionsRequest, options ...grpc.CallOption) (*milvuspb.ShowPartitionsResponse, error) {
s.Equal(collID, req.GetCollectionID())
return nil, errors.New("mocked")
})
_, err := s.broker.ShowPartitionsInternal(context.Background(), collID)
s.Error(err)
s.TearDownTest()
})
}
func (s *BrokerSuite) TestShowCollections() {
s.Run("return_success", func() {
s.SetupTest()
dbName := fmt.Sprintf("db_%d", rand.Intn(100))
s.rootCoordClient.EXPECT().ShowCollections(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.ShowCollectionsRequest, options ...grpc.CallOption) (*milvuspb.ShowCollectionsResponse, error) {
s.Equal(dbName, req.GetDbName())
return &milvuspb.ShowCollectionsResponse{
Status: merr.Status(nil),
CollectionIds: []int64{1, 2, 3},
CollectionNames: []string{"coll_1", "coll_2", "coll_3"},
CreatedTimestamps: []uint64{0, 0, 0},
}, nil
})
resp, err := s.broker.ShowCollections(context.Background(), dbName)
s.NoError(err)
s.ElementsMatch([]int64{1, 2, 3}, resp.GetCollectionIds())
s.ElementsMatch([]string{"coll_1", "coll_2", "coll_3"}, resp.GetCollectionNames())
s.TearDownTest()
})
s.Run("return_error", func() {
s.SetupTest()
dbName := fmt.Sprintf("db_%d", rand.Intn(100))
s.rootCoordClient.EXPECT().ShowCollections(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.ShowCollectionsRequest, options ...grpc.CallOption) (*milvuspb.ShowCollectionsResponse, error) {
s.Equal(dbName, req.GetDbName())
return nil, errors.New("mocked")
})
_, err := s.broker.ShowCollections(context.Background(), dbName)
s.Error(err)
s.TearDownTest()
})
}
func (s *BrokerSuite) TestListDatabases() {
s.Run("return_success", func() {
s.SetupTest()
s.rootCoordClient.EXPECT().ListDatabases(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.ListDatabasesRequest, options ...grpc.CallOption) (*milvuspb.ListDatabasesResponse, error) {
return &milvuspb.ListDatabasesResponse{
Status: merr.Status(nil),
DbNames: []string{"db_1", "db_2", "db_3"},
}, nil
})
resp, err := s.broker.ListDatabases(context.Background())
s.NoError(err)
s.ElementsMatch([]string{"db_1", "db_2", "db_3"}, resp.GetDbNames())
s.TearDownTest()
})
s.Run("return_error", func() {
s.SetupTest()
s.rootCoordClient.EXPECT().ListDatabases(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.ListDatabasesRequest, options ...grpc.CallOption) (*milvuspb.ListDatabasesResponse, error) {
return nil, errors.New("mocked")
})
_, err := s.broker.ListDatabases(context.Background())
s.Error(err)
s.TearDownTest()
})
}
func (s *BrokerSuite) TestHasCollection() {
s.Run("return_success", func() {
s.SetupTest()
collID := int64(1000 + rand.Intn(500))
s.rootCoordClient.EXPECT().DescribeCollection(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.DescribeCollectionRequest, options ...grpc.CallOption) (*milvuspb.DescribeCollectionResponse, error) {
s.Equal(collID, req.GetCollectionID())
return &milvuspb.DescribeCollectionResponse{
Status: merr.Status(nil),
CollectionID: collID,
CollectionName: "test_collection",
}, nil
})
result, err := s.broker.HasCollection(context.Background(), collID)
s.NoError(err)
s.True(result)
s.TearDownTest()
})
s.Run("return_collection_not_found", func() {
s.SetupTest()
collID := int64(1000 + rand.Intn(500))
s.rootCoordClient.EXPECT().DescribeCollection(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.DescribeCollectionRequest, options ...grpc.CallOption) (*milvuspb.DescribeCollectionResponse, error) {
s.Equal(collID, req.GetCollectionID())
return &milvuspb.DescribeCollectionResponse{
Status: merr.Status(merr.WrapErrCollectionNotFound(collID)),
}, nil
})
result, err := s.broker.HasCollection(context.Background(), collID)
s.NoError(err)
s.False(result)
s.TearDownTest()
})
s.Run("return_error", func() {
s.SetupTest()
collID := int64(1000 + rand.Intn(500))
s.rootCoordClient.EXPECT().DescribeCollection(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.DescribeCollectionRequest, options ...grpc.CallOption) (*milvuspb.DescribeCollectionResponse, error) {
s.Equal(collID, req.GetCollectionID())
return nil, errors.New("mocked")
})
_, err := s.broker.HasCollection(context.Background(), collID)
s.Error(err)
s.TearDownTest()
})
}
func TestBrokerSuite(t *testing.T) {
suite.Run(t, new(BrokerSuite))
}

5
internal/datacoord/server.go
View File

@ -34,6 +34,7 @@ import (
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/datacoord/broker"
datanodeclient "github.com/milvus-io/milvus/internal/distributed/datanode/client"
indexnodeclient "github.com/milvus-io/milvus/internal/distributed/indexnode/client"
rootcoordclient "github.com/milvus-io/milvus/internal/distributed/rootcoord/client"
@ -152,7 +153,7 @@ type Server struct {
indexEngineVersionManager IndexEngineVersionManager
// manage ways that data coord access other coord
broker Broker
broker broker.Broker
}
// ServerHelper datacoord server injection helper
@ -332,7 +333,7 @@ func (s *Server) initDataCoord() error {
}
log.Info("init rootcoord client done")
s.broker = NewCoordinatorBroker(s.rootCoordClient)
s.broker = broker.NewCoordinatorBroker(s.rootCoordClient)
storageCli, err := s.newChunkManagerFactory()
if err != nil {

5
internal/datacoord/server_test.go
View File

@ -44,6 +44,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/datacoord/broker"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
@ -3690,7 +3691,7 @@ func TestGetFlushAllState(t *testing.T) {
var err error
svr.meta = &meta{}
svr.rootCoordClient = mocks.NewMockRootCoordClient(t)
svr.broker = NewCoordinatorBroker(svr.rootCoordClient)
svr.broker = broker.NewCoordinatorBroker(svr.rootCoordClient)
if test.ListDatabaseFailed {
svr.rootCoordClient.(*mocks.MockRootCoordClient).EXPECT().ListDatabases(mock.Anything, mock.Anything).
Return(&milvuspb.ListDatabasesResponse{
@ -3777,7 +3778,7 @@ func TestGetFlushAllStateWithDB(t *testing.T) {
var err error
svr.meta = &meta{}
svr.rootCoordClient = mocks.NewMockRootCoordClient(t)
svr.broker = NewCoordinatorBroker(svr.rootCoordClient)
svr.broker = broker.NewCoordinatorBroker(svr.rootCoordClient)
if test.DbExist {
svr.rootCoordClient.(*mocks.MockRootCoordClient).EXPECT().ListDatabases(mock.Anything, mock.Anything).