fix: WithConnPool should not be override

joway · Jul 30, 2021 · ae17c53 · ae17c53
1 parent c2f50ff
commit ae17c53
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 7 deletions.
diff --git a/client/option.go b/client/option.go
@@ -33,7 +33,6 @@ import (
 	"github.com/cloudwego/kitex/pkg/loadbalance"
 	"github.com/cloudwego/kitex/pkg/loadbalance/lbcache"
 	"github.com/cloudwego/kitex/pkg/remote"
-	connpool2 "github.com/cloudwego/kitex/pkg/remote/connpool"
 	"github.com/cloudwego/kitex/pkg/remote/trans/netpollmux"
 	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2"
 	"github.com/cloudwego/kitex/pkg/retry"
@@ -186,7 +185,7 @@ func WithShortConnection() Option {
 	return Option{F: func(o *client.Options, di *utils.Slice) {
 		di.Push("WithShortConnection")
 
-		o.RemoteOpt.ConnPool = connpool2.NewShortPool(o.Svr.ServiceName)
+		o.PoolCfg = new(connpool.IdleConfig)
 	}}
 }
 
@@ -196,7 +195,6 @@ func WithLongConnection(cfg connpool.IdleConfig) Option {
 		di.Push(fmt.Sprintf("WithLongConnection(%+v)", cfg))
 
 		o.PoolCfg = connpool.CheckPoolConfig(cfg)
-		o.RemoteOpt.ConnPool = connpool2.NewLongPool(o.Svr.ServiceName, *o.PoolCfg)
 	}}
 }
 

diff --git a/client/option_advanced.go b/client/option_advanced.go
@@ -26,6 +26,7 @@ import (
 	"github.com/cloudwego/kitex/pkg/acl"
 	"github.com/cloudwego/kitex/pkg/diagnosis"
 	"github.com/cloudwego/kitex/pkg/generic"
+	"github.com/cloudwego/kitex/pkg/klog"
 	"github.com/cloudwego/kitex/pkg/proxy"
 	"github.com/cloudwego/kitex/pkg/remote"
 	"github.com/cloudwego/kitex/pkg/remote/trans/netpoll"
@@ -133,12 +134,19 @@ func WithDialer(d remote.Dialer) Option {
 	}}
 }
 
-// WithConnPool to set connection pool
+// WithConnPool sets the connection pool.
+// Note that this option can only be specified once. If more
+// than one pool is specified by this option, only the first
+// one will be used.
 func WithConnPool(pool remote.ConnPool) Option {
 	return Option{F: func(o *client.Options, di *utils.Slice) {
 		di.Push(fmt.Sprintf("WithConnPool((%T)", pool))
 
-		o.RemoteOpt.ConnPool = pool
+		if o.RemoteOpt.ConnPool == nil {
+			o.RemoteOpt.ConnPool = pool
+		} else {
+			klog.Warnf("The connection pool has been initialized. The call to WithConnPool will not take effect.")
+		}
 	}}
 }
 

diff --git a/internal/client/option.go b/internal/client/option.go
@@ -145,8 +145,16 @@ func NewOptions(opts []Option) *Options {
 
 func (o *Options) initConnectionPool() {
 	if o.RemoteOpt.ConnPool == nil {
-		o.RemoteOpt.ConnPool = connpool.NewLongPool(o.Svr.ServiceName,
-			connpool2.IdleConfig{MaxIdlePerAddress: 10, MaxIdleGlobal: 100, MaxIdleTimeout: time.Minute})
+		if o.PoolCfg != nil {
+			var zero connpool2.IdleConfig
+			if *o.PoolCfg == zero {
+				o.RemoteOpt.ConnPool = connpool.NewShortPool(o.Svr.ServiceName)
+			} else {
+				o.RemoteOpt.ConnPool = connpool.NewLongPool(o.Svr.ServiceName, *o.PoolCfg)
+			}
+		} else {
+			o.RemoteOpt.ConnPool = connpool.NewLongPool(o.Svr.ServiceName, connpool2.IdleConfig{MaxIdlePerAddress: 10, MaxIdleGlobal: 100, MaxIdleTimeout: time.Minute})
+		}
 	}
 	pool := o.RemoteOpt.ConnPool
 	o.CloseCallbacks = append(o.CloseCallbacks, pool.Close)

diff --git a/pkg/connpool/config.go b/pkg/connpool/config.go
@@ -38,5 +38,13 @@ func CheckPoolConfig(config IdleConfig) *IdleConfig {
 	} else if config.MaxIdleTimeout < 3*time.Second {
 		config.MaxIdleTimeout = minMaxIdleTimeout
 	}
+
+	if config.MaxIdlePerAddress <= 0 {
+		config.MaxIdlePerAddress = 1
+	}
+
+	if config.MaxIdleGlobal <= 0 {
+		config.MaxIdleGlobal = 1
+	}
 	return &config
 }