support custom grpcMsgSize param

zilliztech · Jul 15, 2024 · 80da0a2 · 80da0a2
1 parent 8460436
commit 80da0a2
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 56 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
@@ -38,36 +38,40 @@ 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)
-	//} 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)
-	//}
-	//if err != nil {
-	//	log.Error("[Milvus2x] new milvus client error", zap.Error(err))
-	//	return nil, err
-	//}
-
-	milvus, err = client.NewClient(ctx, client.Config{
-		Address:  cfg.Endpoint,
-		Username: cfg.UserName,
-		Password: cfg.Password,
-		DialOptions: []grpc.DialOption{
-			grpc.WithDefaultCallOptions(
-				grpc.MaxCallRecvMsgSize(67108864),
-				grpc.MaxCallSendMsgSize(268435456),
-			),
-		},
-	})
+	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 {
+		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))
 		return nil, err
 	}
 
 	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
@@ -116,38 +116,44 @@ 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)
-	//} 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)
-	//}
-	//if err != nil {
-	//	log.Error("[Milvus23x] new milvus client error", zap.Error(err))
-	//	return nil, err
-	//}
-
-	milvus, err = client.NewClient(ctx, client.Config{
-		Address:  cfg.Endpoint,
-		Username: cfg.UserName,
-		Password: cfg.Password,
-		DialOptions: []grpc.DialOption{
-			grpc.WithDefaultCallOptions(
-				grpc.MaxCallRecvMsgSize(67108864),
-				grpc.MaxCallSendMsgSize(268435456),
-			),
-		},
-	})
+	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 {
+		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))
 		return nil, err
 	}
 
 	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