Merge pull request #91 from zilliztech/feature_grpcMsgSize

Support custom grpc msg size
zilliztech · Jul 15, 2024 · dd4ec13 · dd4ec13
2 parents bdf7617 + 80da0a2
commit dd4ec13
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 20 deletions.
diff --git a/core/config/config.go b/core/config/config.go
@@ -75,6 +75,9 @@ type Milvus2xConfig struct {
 	UserName string
 	Password string
 
+	GrpcMaxRecvMsgSize int
+	GrpcMaxSendMsgSize int
+
 	Version   string //internal param
 	hashCache atomic.Uint32
 }

diff --git a/core/config/resolve.go b/core/config/resolve.go
@@ -216,9 +216,11 @@ func resolveRemoteConfig(prefix string, v *viper.Viper) *RemoteConfig {
 
 func resolveTargetMilvus2xConfig(v *viper.Viper) *Milvus2xConfig {
 	return &Milvus2xConfig{
-		Endpoint: v.GetString("target.milvus2x.endpoint"),
-		UserName: v.GetString("target.milvus2x.username"),
-		Password: v.GetString("target.milvus2x.password"),
+		Endpoint:           v.GetString("target.milvus2x.endpoint"),
+		UserName:           v.GetString("target.milvus2x.username"),
+		Password:           v.GetString("target.milvus2x.password"),
+		GrpcMaxRecvMsgSize: v.GetInt("target.milvus2x.grpc.maxCallRecvMsgSize"),
+		GrpcMaxSendMsgSize: v.GetInt("target.milvus2x.grpc.maxCallSendMsgSize"),
 	}
 }
 
@@ -351,8 +353,10 @@ func resolveMilvus2xDumpWorkConfig(v *viper.Viper, workMode common.DumpMode) (*D
 
 func resolveSourceMilvus2xConfig(v *viper.Viper) *Milvus2xConfig {
 	return &Milvus2xConfig{
-		Endpoint: v.GetString("source.milvus2x.endpoint"),
-		UserName: v.GetString("source.milvus2x.username"),
-		Password: v.GetString("source.milvus2x.password"),
+		Endpoint:           v.GetString("source.milvus2x.endpoint"),
+		UserName:           v.GetString("source.milvus2x.username"),
+		Password:           v.GetString("source.milvus2x.password"),
+		GrpcMaxRecvMsgSize: v.GetInt("source.milvus2x.grpc.maxCallRecvMsgSize"),
+		GrpcMaxSendMsgSize: v.GetInt("source.milvus2x.grpc.maxCallSendMsgSize"),
 	}
 }
diff --git a/core/dbclient/milvus2x.go b/core/dbclient/milvus2x.go
@@ -10,6 +10,7 @@ import (
 	"github.com/zilliztech/milvus-migration/internal/log"
 	"github.com/zilliztech/milvus-migration/storage/milvus2x"
 	"go.uber.org/zap"
+	"google.golang.org/grpc"
 	"strconv"
 	"time"
 )
@@ -37,12 +38,29 @@ func NewMilvus2xClient(cfg *config.Milvus2xConfig) (*Milvus2x, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
-	if cfg.UserName == "" {
-		log.Info("[Milvus2x] find username is empty, will use NewDefaultGrpcClient() to new client")
-		milvus, err = client.NewDefaultGrpcClient(ctx, cfg.Endpoint)
+	if cfg.GrpcMaxRecvMsgSize <= 0 {
+		if cfg.UserName == "" {
+			log.Info("[Milvus2x] find username is empty, will use NewDefaultGrpcClient() to new client")
+			milvus, err = client.NewDefaultGrpcClient(ctx, cfg.Endpoint)
+		} else {
+			log.Info("[Milvus2x] find username not empty, will use NewDefaultGrpcClientWithURI() to new client")
+			milvus, err = client.NewDefaultGrpcClientWithURI(ctx, cfg.Endpoint, cfg.UserName, cfg.Password)
+		}
 	} else {
-		log.Info("[Milvus2x] find username not empty, will use NewDefaultGrpcClientWithURI() to new client")
-		milvus, err = client.NewDefaultGrpcClientWithURI(ctx, cfg.Endpoint, cfg.UserName, cfg.Password)
+		config := client.Config{
+			Address: cfg.Endpoint,
+			DialOptions: []grpc.DialOption{
+				grpc.WithDefaultCallOptions(
+					grpc.MaxCallRecvMsgSize(cfg.GrpcMaxRecvMsgSize),
+					grpc.MaxCallSendMsgSize(cfg.GrpcMaxSendMsgSize),
+				),
+			},
+		}
+		if cfg.UserName != "" {
+			config.Username = cfg.UserName
+			config.Password = cfg.Password
+		}
+		milvus, err = client.NewClient(ctx, config)
 	}
 	if err != nil {
 		log.Error("[Milvus2x] new milvus client error", zap.Error(err))
@@ -51,7 +69,9 @@ func NewMilvus2xClient(cfg *config.Milvus2xConfig) (*Milvus2x, error) {
 
 	log.Info("[Milvus2x] begin to test connect",
 		zap.String("endpoint", cfg.Endpoint),
-		zap.String("username", cfg.UserName))
+		zap.String("username", cfg.UserName),
+		zap.Int("GrpcMaxCallRecvMsgSize", cfg.GrpcMaxRecvMsgSize),
+		zap.Int("GrpcMaxCallSendMsgSize", cfg.GrpcMaxSendMsgSize))
 	_, err = milvus.HasCollection(ctx, "test")
 	if err != nil {
 		return nil, err

diff --git a/storage/milvus2x/milvus2_3_ver.go b/storage/milvus2x/milvus2_3_ver.go
@@ -9,6 +9,7 @@ import (
 	"github.com/zilliztech/milvus-migration/core/type/milvus2xtype"
 	"github.com/zilliztech/milvus-migration/internal/log"
 	"go.uber.org/zap"
+	"google.golang.org/grpc"
 	"io"
 	"strconv"
 	"time"
@@ -115,14 +116,32 @@ func _createMilvus23VerClient(cfg *config.Milvus2xConfig) (*Milvus23VerClient, e
 
 	var milvus client.Client
 	var err error
-	ctx := context.Background()
-
-	if cfg.UserName == "" {
-		log.Info("[Milvus23x] find username is empty, will use NewDefaultGrpcClient() to new client")
-		milvus, err = client.NewDefaultGrpcClient(ctx, cfg.Endpoint)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	if cfg.GrpcMaxRecvMsgSize <= 0 {
+		if cfg.UserName == "" {
+			log.Info("[Milvus23x] find username is empty, will use NewDefaultGrpcClient() to new client")
+			milvus, err = client.NewDefaultGrpcClient(ctx, cfg.Endpoint)
+		} else {
+			log.Info("[Milvus23x] find username not empty, will use NewDefaultGrpcClientWithURI() to new client")
+			milvus, err = client.NewDefaultGrpcClientWithURI(ctx, cfg.Endpoint, cfg.UserName, cfg.Password)
+		}
 	} else {
-		log.Info("[Milvus23x] find username not empty, will use NewDefaultGrpcClientWithURI() to new client")
-		milvus, err = client.NewDefaultGrpcClientWithURI(ctx, cfg.Endpoint, cfg.UserName, cfg.Password)
+		config := client.Config{
+			Address: cfg.Endpoint,
+			DialOptions: []grpc.DialOption{
+				grpc.WithDefaultCallOptions(
+					grpc.MaxCallRecvMsgSize(cfg.GrpcMaxRecvMsgSize),
+					grpc.MaxCallSendMsgSize(cfg.GrpcMaxSendMsgSize),
+				),
+			},
+		}
+		if cfg.UserName != "" {
+			config.Username = cfg.UserName
+			config.Password = cfg.Password
+		}
+		milvus, err = client.NewClient(ctx, config)
 	}
 	if err != nil {
 		log.Error("[Milvus23x] new milvus client error", zap.Error(err))
@@ -131,7 +150,10 @@ func _createMilvus23VerClient(cfg *config.Milvus2xConfig) (*Milvus23VerClient, e
 
 	log.Info("[Milvus23x] begin to test connect",
 		zap.String("endpoint", cfg.Endpoint),
-		zap.String("username", cfg.UserName))
+		zap.String("username", cfg.UserName),
+		zap.Int("GrpcMaxCallRecvMsgSize", cfg.GrpcMaxRecvMsgSize),
+		zap.Int("GrpcMaxCallSendMsgSize", cfg.GrpcMaxSendMsgSize))
+
 	_, err = milvus.HasCollection(ctx, "test")
 	if err != nil {
 		return nil, err