Support grpc connection pooling

See also #637 Signed-off-by: Congqi Xia <[email protected]>
milvus-io · Dec 8, 2023 · 3ef070c · 3ef070c
1 parent 43226a3
commit 3ef070c
Show file tree

Hide file tree

Showing 21 changed files with 470 additions and 241 deletions.
diff --git a/client/admin.go b/client/admin.go
@@ -20,10 +20,12 @@ import (
 
 // GetVersion returns milvus server version information.
 func (c *GrpcClient) GetVersion(ctx context.Context) (string, error) {
-	if c.Service == nil {
+	service := c.Service(ctx)
+	if service == nil {
 		return "", ErrClientNotReady
 	}
-	resp, err := c.Service.GetVersion(ctx, &milvuspb.GetVersionRequest{})
+	defer service.Close()
+	resp, err := service.GetVersion(ctx, &milvuspb.GetVersionRequest{})
 	if err != nil {
 		return "", err
 	}
@@ -32,10 +34,12 @@ func (c *GrpcClient) GetVersion(ctx context.Context) (string, error) {
 
 // CheckHealth returns milvus state
 func (c *GrpcClient) CheckHealth(ctx context.Context) (*entity.MilvusState, error) {
-	if c.Service == nil {
+	service := c.Service(ctx)
+	if service == nil {
 		return nil, ErrClientNotReady
 	}
-	resp, err := c.Service.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
+	defer service.Close()
+	resp, err := service.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
 	if err != nil {
 		return nil, err
 	}

diff --git a/client/alias.go b/client/alias.go
@@ -24,17 +24,19 @@ import (
 
 // CreateAlias creates an alias for collection
 func (c *GrpcClient) CreateAlias(ctx context.Context, collName string, alias string) error {
-	if c.Service == nil {
+	service := c.Service(ctx)
+	if service == nil {
 		return ErrClientNotReady
 	}
+	defer service.Close()
 
 	req := &milvuspb.CreateAliasRequest{
 		DbName:         "", // reserved
 		CollectionName: collName,
 		Alias:          alias,
 	}
 
-	resp, err := c.Service.CreateAlias(ctx, req)
+	resp, err := service.CreateAlias(ctx, req)
 	if err != nil {
 		return err
 	}
@@ -47,16 +49,18 @@ func (c *GrpcClient) CreateAlias(ctx context.Context, collName string, alias str
 
 // DropAlias drops the specified Alias
 func (c *GrpcClient) DropAlias(ctx context.Context, alias string) error {
-	if c.Service == nil {
+	service := c.Service(ctx)
+	if service == nil {
 		return ErrClientNotReady
 	}
+	defer service.Close()
 
 	req := &milvuspb.DropAliasRequest{
 		DbName: "", // reserved
 		Alias:  alias,
 	}
 
-	resp, err := c.Service.DropAlias(ctx, req)
+	resp, err := service.DropAlias(ctx, req)
 	if err != nil {
 		return err
 	}
@@ -69,17 +73,19 @@ func (c *GrpcClient) DropAlias(ctx context.Context, alias string) error {
 
 // AlterAlias changes collection alias to provided alias
 func (c *GrpcClient) AlterAlias(ctx context.Context, collName string, alias string) error {
-	if c.Service == nil {
+	service := c.Service(ctx)
+	if service == nil {
 		return ErrClientNotReady
 	}
+	defer service.Close()
 
 	req := &milvuspb.AlterAliasRequest{
 		DbName:         "", // reserved
 		CollectionName: collName,
 		Alias:          alias,
 	}
 
-	resp, err := c.Service.AlterAlias(ctx, req)
+	resp, err := service.AlterAlias(ctx, req)
 	if err != nil {
 		return err
 	}

diff --git a/client/authentication.go b/client/authentication.go
@@ -9,14 +9,16 @@ import (
 
 // CreateCredential create new user and password
 func (c *GrpcClient) CreateCredential(ctx context.Context, username string, password string) error {
-	if c.Service == nil {
+	service := c.Service(ctx)
+	if service == nil {
 		return ErrClientNotReady
 	}
+	defer service.Close()
 	req := &milvuspb.CreateCredentialRequest{
 		Username: username,
 		Password: crypto.Base64Encode(password),
 	}
-	resp, err := c.Service.CreateCredential(ctx, req)
+	resp, err := service.CreateCredential(ctx, req)
 	if err != nil {
 		return err
 	}
@@ -29,15 +31,17 @@ func (c *GrpcClient) CreateCredential(ctx context.Context, username string, pass
 
 // UpdateCredential update password for a user
 func (c *GrpcClient) UpdateCredential(ctx context.Context, username string, oldPassword string, newPassword string) error {
-	if c.Service == nil {
+	service := c.Service(ctx)
+	if service == nil {
 		return ErrClientNotReady
 	}
+	defer service.Close()
 	req := &milvuspb.UpdateCredentialRequest{
 		Username:    username,
 		OldPassword: crypto.Base64Encode(oldPassword),
 		NewPassword: crypto.Base64Encode(newPassword),
 	}
-	resp, err := c.Service.UpdateCredential(ctx, req)
+	resp, err := service.UpdateCredential(ctx, req)
 	if err != nil {
 		return err
 	}
@@ -50,13 +54,15 @@ func (c *GrpcClient) UpdateCredential(ctx context.Context, username string, oldP
 
 // DeleteCredential delete a user
 func (c *GrpcClient) DeleteCredential(ctx context.Context, username string) error {
-	if c.Service == nil {
+	service := c.Service(ctx)
+	if service == nil {
 		return ErrClientNotReady
 	}
+	defer service.Close()
 	req := &milvuspb.DeleteCredentialRequest{
 		Username: username,
 	}
-	resp, err := c.Service.DeleteCredential(ctx, req)
+	resp, err := service.DeleteCredential(ctx, req)
 	if err != nil {
 		return err
 	}
@@ -69,11 +75,13 @@ func (c *GrpcClient) DeleteCredential(ctx context.Context, username string) erro
 
 // ListCredUsers list all usernames
 func (c *GrpcClient) ListCredUsers(ctx context.Context) ([]string, error) {
-	if c.Service == nil {
+	service := c.Service(ctx)
+	if service == nil {
 		return nil, ErrClientNotReady
 	}
+	defer service.Close()
 	req := &milvuspb.ListCredUsersRequest{}
-	resp, err := c.Service.ListCredUsers(ctx, req)
+	resp, err := service.ListCredUsers(ctx, req)
 	if err != nil {
 		return nil, err
 	}

diff --git a/client/client.go b/client/client.go
@@ -18,6 +18,8 @@ import (
 
 	"google.golang.org/grpc"
 
+	grpcpool "github.com/processout/grpc-go-pool"
+
 	"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
 
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
@@ -238,11 +240,16 @@ func NewClient(ctx context.Context, config Config) (Client, error) {
 	// Parse grpc options
 	options := c.config.getDialOption()
 
-	// Connect the grpc server.
-	if err := c.connect(ctx, addr, options...); err != nil {
+	factory := func() (*grpc.ClientConn, error) {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
+		defer cancel()
+		return c.connect(ctx, addr, options...)
+	}
+	pool, err := grpcpool.New(factory, c.config.ConnPoolInit, c.config.ConnPoolMax, c.config.ConnPoolIdleTimeout)
+	if err != nil {
 		return nil, err
 	}
-
+	c.connPool = pool
 	return c, nil
 }
 

diff --git a/client/client_mock_test.go b/client/client_mock_test.go
@@ -66,6 +66,8 @@ func (s *MockSuiteBase) SetupTest() {
 			grpc.WithTransportCredentials(insecure.NewCredentials()),
 			grpc.WithContextDialer(s.mockDialer),
 		},
+		ConnPoolInit: 1,
+		ConnPoolMax:  1,
 	})
 	s.Require().NoError(err)
 	s.setupConnect()

diff --git a/client/client_test.go b/client/client_test.go
@@ -12,8 +12,10 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/examples/helloworld/helloworld"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/reflection"
@@ -150,6 +152,8 @@ func TestGrpcClientNil(t *testing.T) {
 			if m.Name == "Close" || m.Name == "Connect" || // skip connect & close
 				m.Name == "UsingDatabase" || // skip use database
 				m.Name == "Search" || // type alias MetricType treated as string
+				m.Name == "QueryByPks" ||
+				m.Name == "Service" ||
 				m.Name == "CalcDistance" ||
 				m.Name == "ManualCompaction" || // time.Duration hard to detect in reflect
 				m.Name == "Insert" || m.Name == "Upsert" { // complex methods with ...
@@ -221,7 +225,7 @@ func TestGrpcClientConnect(t *testing.T) {
 			Config{
 				Address: "bufnet",
 				DialOptions: []grpc.DialOption{
-					grpc.WithBlock(), grpc.WithInsecure(), grpc.WithContextDialer(bufDialer),
+					grpc.WithBlock(), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithContextDialer(bufDialer),
 				},
 			})
 		assert.Nil(t, err)
@@ -303,7 +307,10 @@ func TestGrpcClientRetryPolicy(t *testing.T) {
 	assert.Nil(t, err)
 	defer client.Close()
 
-	greeterClient := helloworld.NewGreeterClient(client.(*GrpcClient).Conn)
+	conn, err := client.(*GrpcClient).connPool.Get(context.TODO())
+	require.NoError(t, err)
+
+	greeterClient := helloworld.NewGreeterClient(conn)
 	ctx := context.Background()
 	name := fmt.Sprintf("hello world %d", time.Now().Second())
 	res, err := greeterClient.SayHello(ctx, &helloworld.HelloRequest{Name: name})