net/http: fix TestTransportMaxConnsPerHost flakes

The testcase created a race between the close of the current connection
and the client grabbing a connection for the next request. The client
may receive the current connection which may be closed during its use.
We can have the trasnport close all idle connections thereby forcing the
client to receive a new connection.

Closing idle connections did not handle cleaning up host connection
counts for http/2. We will now decrement the host connection count for
http/2 connections.

Fixes #31784

Change-Id: Iefc0d0d7ed9fa3acd8b4f42004f1579fc1de63fd
Reviewed-on: https://go-review.googlesource.com/c/go/+/174950
Reviewed-by: Brad Fitzpatrick <[email protected]>
Run-TryBot: Brad Fitzpatrick <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>

Author: fraenkel

src/net/http/transport.go

@@ -1416,7 +1416,7 @@ func (t *Transport) dialConn(ctx context.Context, cm connectMethod) (*persistCon
 
 	if s := pconn.tlsState; s != nil && s.NegotiatedProtocolIsMutual && s.NegotiatedProtocol != "" {
 		if next, ok := t.TLSNextProto[s.NegotiatedProtocol]; ok {
-			return &persistConn{cacheKey: pconn.cacheKey, alt: next(cm.targetAddr, pconn.conn.(*tls.Conn))}, nil
+			return &persistConn{t: t, cacheKey: pconn.cacheKey, alt: next(cm.targetAddr, pconn.conn.(*tls.Conn))}, nil
 		}
 	}
 
@@ -2344,13 +2344,8 @@ func (pc *persistConn) closeLocked(err error) {
 	if pc.closed == nil {
 		pc.closed = err
 		if pc.alt != nil {
-			// Do nothing; can only get here via getConn's
-			// handlePendingDial's putOrCloseIdleConn when
-			// it turns out the abandoned connection in
-			// flight ended up negotiating an alternate
-			// protocol. We don't use the connection
-			// freelist for http2. That's done by the
-			// alternate protocol's RoundTripper.
+			// Clean up any host connection counting.
+			pc.t.decHostConnCount(pc.cacheKey)
 		} else {
 			if err != errCallerOwnsConn {
 				pc.conn.Close()

src/net/http/transport_test.go

@@ -22,7 +22,6 @@ import (
 	"fmt"
 	"go/token"
 	"internal/nettrace"
-	"internal/testenv"
 	"io"
 	"io/ioutil"
 	"log"
@@ -592,7 +591,7 @@ func TestTransportMaxConnsPerHostIncludeDialInProgress(t *testing.T) {
 
 func TestTransportMaxConnsPerHost(t *testing.T) {
 	defer afterTest(t)
-	testenv.SkipFlaky(t, 31784)
+
 	h := HandlerFunc(func(w ResponseWriter, r *Request) {
 		_, err := w.Write([]byte("foo"))
 		if err != nil {
@@ -666,6 +665,7 @@ func TestTransportMaxConnsPerHost(t *testing.T) {
 		}
 
 		(<-connCh).Close()
+		tr.CloseIdleConnections()
 
 		doReq()
 		expected++