optimize: fix gRPC conn pool leak

Previously a closed/broken transport stayed in the nphttp2 client pool
indefinitely, because only a later put() could overwrite the slot. A
new OnClose callback now evicts the transport via transports.remove()
as soon as it closes, and all cliTransports accesses are serialized by
a new sync.RWMutex.

To avoid a lock-order inversion between transports.mu (which doPut
holds while calling IsActive -> http2Client.mu) and the onClose ->
remove path (which takes transports.mu while http2Client.Close
historically held t.mu), http2Client.Close now releases t.mu before
invoking onClose. The invariant transports.mu -> http2Client.mu is
documented on doPut, remove, close, and http2Client.Close.

GracefulClose is now invoked outside transports.mu in both put() and
close() to avoid re-entering the mutex through onClose -> remove on
transports with zero active streams. Dump() snapshots (addr, transport)
refs under RLock and calls each transport's Dump() after releasing the
lock.

The pool-side callback signatures now take (ctx, trans, err) instead
of (). A new grpc.ClientConfig + NewClientTransportWithConfig expose
the new shape; the deprecated NewClientTransport is preserved as a
thin adapter.

Note: the deprecated NewClientTransport path has a visible timing
change - onClose is now invoked after state is set to closing and
after http2Client.mu is released, whereas it previously ran before the
state change while the mutex was held. Callers relying on the old
ordering must migrate to NewClientTransportWithConfig.

Author: DMwangnima

pkg/remote/trans/nphttp2/conn_pool.go

@@ -23,7 +23,6 @@ import (
 	"runtime"
 	"runtime/debug"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"golang.org/x/sync/singleflight"
@@ -60,7 +59,7 @@ func NewConnPool(remoteService string, size uint32, connOpts grpc.ConnectOptions
 	}
 }
 
-// MuxPool manages a pool of long connections.
+// connPool manages a pool of gRPC long connections.
 type connPool struct {
 	size          uint32
 	sfg           singleflight.Group
@@ -69,76 +68,8 @@ type connPool struct {
 	connOpts      grpc.ConnectOptions
 }
 
-type transports struct {
-	index         uint32
-	size          uint32
-	cliTransports []grpc.ClientTransport
-}
-
-// get connection from the pool, load balance with round-robin.
-func (t *transports) get() grpc.ClientTransport {
-	idx := atomic.AddUint32(&t.index, 1)
-	return t.cliTransports[idx%t.size]
-}
-
-// put find the first empty position to put the connection to the pool.
-func (t *transports) put(trans grpc.ClientTransport) {
-	for i := 0; i < int(t.size); i++ {
-		cliTransport := t.cliTransports[i]
-		if cliTransport == nil {
-			t.cliTransports[i] = trans
-			return
-		}
-		if !cliTransport.(grpc.IsActive).IsActive() {
-			t.cliTransports[i].GracefulClose()
-			t.cliTransports[i] = trans
-			return
-		}
-	}
-}
-
-// close all connections of the pool.
-func (t *transports) close() {
-	for i := range t.cliTransports {
-		if c := t.cliTransports[i]; c != nil {
-			c.GracefulClose()
-		}
-	}
-}
-
 var _ remote.LongConnPool = (*connPool)(nil)
 
-func (p *connPool) newTransport(ctx context.Context, dialer remote.Dialer, network, address string,
-	connectTimeout time.Duration, opts grpc.ConnectOptions,
-) (grpc.ClientTransport, error) {
-	conn, err := dialer.DialTimeout(network, address, connectTimeout)
-	if err != nil {
-		return nil, err
-	}
-	if opts.TLSConfig != nil {
-		tlsConn, err := newTLSConn(conn, opts.TLSConfig)
-		if err != nil {
-			return nil, err
-		}
-		conn = tlsConn
-	}
-	return grpc.NewClientTransport(
-		ctx,
-		conn,
-		opts,
-		p.remoteService,
-		func(grpc.GoAwayReason) {
-			// remove connection from the pool.
-			// we do not need to close this grpc transport manually
-			// since grpc client is responsible for doing this.
-			p.conns.Delete(address)
-		},
-		func() {
-			// do nothing
-		},
-	)
-}
-
 // 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 {
@@ -147,47 +78,47 @@ func (p *connPool) Get(ctx context.Context, network, address string, opt remote.
 
 	var (
 		trans *transports
+		tr    grpc.ClientTransport
+		idx   uint32
 		conn  *clientConn
 		err   error
 	)
 
 	v, ok := p.conns.Load(address)
 	if ok {
 		trans = v.(*transports)
-		if tr := trans.get(); tr != nil {
-			if tr.(grpc.IsActive).IsActive() {
-				// Actually new a stream, reuse the connection (grpc.ClientTransport)
-				conn, err = newClientConn(ctx, tr, address)
-				if err == nil {
-					return conn, nil
-				}
-				klog.CtxDebugf(ctx, "KITEX: New grpc stream failed, network=%s, address=%s, error=%s", network, address, err.Error())
+		tr, idx = trans.getActiveTransport()
+		if tr != nil {
+			// Actually new a stream, reuse the connection (grpc.ClientTransport)
+			conn, err = newClientConn(ctx, tr, address)
+			if err == nil {
+				return conn, nil
 			}
+			klog.CtxDebugf(ctx, "KITEX: New grpc stream failed, network=%s, address=%s, error=%s", network, address, err.Error())
 		}
 	}
-	tr, err, _ := p.sfg.Do(address, func() (i interface{}, e 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
+	rawTr, dErr, _ := p.sfg.Do(address, func() (i interface{}, e error) {
+		// avoid creating duplicate transports
+		if existTrans, ok := p.conns.Load(address); ok {
+			trans = existTrans.(*transports)
+		} else if trans == nil {
+			trans = newTransports(p.size)
 		}
-		if trans == nil {
-			trans = &transports{
-				size:          p.size,
-				cliTransports: make([]grpc.ClientTransport, p.size),
-			}
+
+		res, cErr := trans.createTransport(idx, p.remoteService, opt.Dialer, network, address, opt.ConnectTimeout, p.connOpts)
+		if cErr != nil {
+			return nil, cErr
 		}
-		trans.put(tr) // the tr (connection) maybe not in the pool, but can be recycled by keepalive.
-		p.conns.Store(address, trans)
-		return tr, nil
+
+		p.conns.LoadOrStore(address, trans)
+		return res, nil
 	})
-	if err != nil {
-		klog.CtxErrorf(ctx, "KITEX: New grpc client connection failed, network=%s, address=%s, error=%s", network, address, err.Error())
-		return nil, err
+	if dErr != nil {
+		klog.CtxErrorf(ctx, "KITEX: New grpc client connection failed, network=%s, address=%s, error=%s", network, address, dErr.Error())
+		return nil, dErr
 	}
 	klog.CtxDebugf(ctx, "KITEX: New grpc client connection succeed, network=%s, address=%s", network, address)
-	return newClientConn(ctx, tr.(grpc.ClientTransport), address)
+	return newClientConn(ctx, rawTr.(grpc.ClientTransport), address)
 }
 
 // Put implements the ConnPool interface.
@@ -205,9 +136,7 @@ func (p *connPool) release(conn net.Conn) error {
 }
 
 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)
+	tr, err := newTransport(p.remoteService, opt.Dialer, network, address, opt.ConnectTimeout, p.connOpts, nil, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -224,8 +153,7 @@ func (p *connPool) Discard(conn net.Conn) error {
 
 // Clean implements the LongConnPool interface.
 func (p *connPool) Clean(network, address string) {
-	if v, ok := p.conns.Load(address); ok {
-		p.conns.Delete(address)
+	if v, ok := p.conns.LoadAndDelete(address); ok {
 		v.(*transports).close()
 	}
 }
@@ -240,6 +168,11 @@ func (p *connPool) Close() error {
 	return nil
 }
 
+type dumpEntry struct {
+	addr string
+	tr   grpc.ClientTransport
+}
+
 // Dump dumps the connection pool with the details of the underlying transport.
 func (p *connPool) Dump() interface{} {
 	defer func() {
@@ -252,31 +185,61 @@ func (p *connPool) Dump() interface{} {
 	// If mesh egress is not enabled, toAddr should be the address of the callee service.
 	// Otherwise, toAddr will be the same, so you should check the remoteAddress in each stream, which is read from the header.
 	poolDump := make(map[string]interface{}, p.size)
+	var cliTransDumps []dumpEntry
 	p.conns.Range(func(k, v interface{}) bool {
 		addr := k.(string)
-		ts := v.(*transports)
-		for _, t := range ts.cliTransports {
-			if t == nil {
-				continue
-			}
-			dumper, ok := t.(interface{ Dump() interface{} })
-			if !ok {
-				continue
-			}
-			var curr []interface{}
-			if poolDump[addr] == nil {
-				curr = make([]interface{}, 0)
-			} else {
-				curr = poolDump[addr].([]interface{})
-			}
-			curr = append(curr, dumper.Dump())
-			poolDump[addr] = curr
+		for _, tr := range v.(*transports).loadAll() {
+			cliTransDumps = append(cliTransDumps, dumpEntry{addr: addr, tr: tr})
 		}
 		return true
 	})
+
+	for _, cliTransDump := range cliTransDumps {
+		dumper, ok := cliTransDump.tr.(interface{ Dump() interface{} })
+		if !ok {
+			continue
+		}
+		var curr []interface{}
+		if poolDump[cliTransDump.addr] == nil {
+			curr = make([]interface{}, 0)
+		} else {
+			curr = poolDump[cliTransDump.addr].([]interface{})
+		}
+		curr = append(curr, dumper.Dump())
+		poolDump[cliTransDump.addr] = curr
+	}
 	return poolDump
 }
 
+// newTransport created a gRPC connection
+func newTransport(remoteService string,
+	dialer remote.Dialer, network, address string, connectTimeout time.Duration, opts grpc.ConnectOptions,
+	onGoAway func(context.Context, grpc.ClientTransport, grpc.GoAwayReason),
+	onClose func(context.Context, grpc.ClientTransport, error),
+) (grpc.ClientTransport, error) {
+	conn, err := dialer.DialTimeout(network, address, connectTimeout)
+	if err != nil {
+		return nil, err
+	}
+	if opts.TLSConfig != nil {
+		tlsConn, err := newTLSConn(conn, opts.TLSConfig)
+		if err != nil {
+			return nil, err
+		}
+		conn = tlsConn
+	}
+	return grpc.NewClientTransportWithConfig(
+		context.Background(), // gRPC connection does not need to be bound to a specific ctx
+		conn,
+		opts,
+		grpc.ClientConfig{
+			RemoteService: remoteService,
+			OnGoAway:      onGoAway,
+			OnClose:       onClose,
+		},
+	)
+}
+
 // newTLSConn constructs a client-side TLS connection and performs handshake.
 func newTLSConn(conn net.Conn, tlsCfg *tls.Config) (net.Conn, error) {
 	tlsConn := tls.Client(conn, tlsCfg)
@@ -285,3 +248,12 @@ func newTLSConn(conn net.Conn, tlsCfg *tls.Config) (net.Conn, error) {
 	}
 	return tlsConn, nil
 }
+
+func checkActive(trans grpc.ClientTransport) bool {
+	if trans == nil {
+		return false
+	}
+	// grpc.ClientTransport is implemented by *http2Client in pkg/remote/trans/nphttp2/grpc
+	// it has implements grpc.IsActive so we do not assert type
+	return trans.(grpc.IsActive).IsActive()
+}