feat: support short connection for grpc unary (cloudwego#473)

internal/client/option.go

@@ -158,13 +158,18 @@ func NewOptions(opts []Option) *Options {
 }
 
 func (o *Options) initRemoteOpt() {
+	var zero connpool2.IdleConfig
+
 	if o.Configs.TransportProtocol()&transport.GRPC == transport.GRPC {
+		if o.PoolCfg != nil && *o.PoolCfg == zero {
+			// grpc unary short connection
+			o.GRPCConnectOpts.ShortConn = true
+		}
 		o.RemoteOpt.ConnPool = nphttp2.NewConnPool(o.Svr.ServiceName, o.GRPCConnPoolSize, *o.GRPCConnectOpts)
 		o.RemoteOpt.CliHandlerFactory = nphttp2.NewCliTransHandlerFactory()
 	}
 	if o.RemoteOpt.ConnPool == nil {
 		if o.PoolCfg != nil {
-			var zero connpool2.IdleConfig
 			if *o.PoolCfg == zero {
 				o.RemoteOpt.ConnPool = connpool.NewShortPool(o.Svr.ServiceName)
 			} else {

pkg/remote/trans/nphttp2/conn_pool.go

@@ -123,6 +123,10 @@ func (p *connPool) newTransport(ctx context.Context, dialer remote.Dialer, netwo
 
 // Get pick or generate a net.Conn and return
 func (p *connPool) Get(ctx context.Context, network, address string, opt remote.ConnOption) (net.Conn, error) {
+	if p.connOpts.ShortConn {
+		return p.createShortConn(ctx, network, address, opt)
+	}
+
 	var (
 		trans *transports
 		conn  *clientConn
@@ -170,9 +174,28 @@ func (p *connPool) Get(ctx context.Context, network, address string, opt remote.
 
 // Put implements the ConnPool interface.
 func (p *connPool) Put(conn net.Conn) error {
+	if p.connOpts.ShortConn {
+		return p.release(conn)
+	}
+	return nil
+}
+
+func (p *connPool) release(conn net.Conn) error {
+	clientConn := conn.(*clientConn)
+	clientConn.tr.GracefulClose()
 	return nil
 }
 
+func (p *connPool) createShortConn(ctx context.Context, network, address string, opt remote.ConnOption) (net.Conn, error) {
+	// Notice: newTransport means new a connection, the timeout of connection cannot be set,
+	// so using context.Background() but not the ctx passed in as the parameter.
+	tr, err := p.newTransport(context.Background(), opt.Dialer, network, address, opt.ConnectTimeout, p.connOpts)
+	if err != nil {
+		return nil, err
+	}
+	return newClientConn(ctx, tr, address)
+}
+
 // Discard implements the ConnPool interface.
 func (p *connPool) Discard(conn net.Conn) error {
 	return nil

pkg/remote/trans/nphttp2/grpc/transport.go

@@ -566,6 +566,8 @@ type ConnectOptions struct {
 	ReadBufferSize uint32
 	// MaxHeaderListSize sets the max (uncompressed) size of header list that is prepared to be received.
 	MaxHeaderListSize *uint32
+	// ShortConn indicates whether the connection will be reused from grpc conn pool
+	ShortConn bool
 }
 
 // NewServerTransport creates a ServerTransport with conn or non-nil error