net/http: add onClose hook to fake net listener

Avoids a race condition: If we set an onClose hook on a conn
created by a listener, then setting the hook can race with
the connection closing.

Change-Id: Ibadead3abbe4335d41f1e2cf84f4696fe98166b3
Reviewed-on: https://go-review.googlesource.com/c/go/+/658655
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
Auto-Submit: Damien Neil <dneil@google.com>

Author: neild

src/net/http/netconn_test.go

@@ -35,7 +35,8 @@ type fakeNetListener struct {
 	addr         netip.AddrPort
 	locPort      uint16
 
-	onDial func() // called when making a new connection
+	onDial  func()             // called when making a new connection
+	onClose func(*fakeNetConn) // called when closing a connection
 
 	trackConns bool // set this to record all created conns
 	conns      []*fakeNetConn
@@ -65,6 +66,8 @@ func (li *fakeNetListener) connect() *fakeNetConn {
 	locAddr := netip.AddrPortFrom(netip.AddrFrom4([4]byte{127, 0, 0, 1}), li.locPort)
 	li.locPort++
 	c0, c1 := fakeNetPipe(li.addr, locAddr)
+	c0.onClose = li.onClose
+	c1.onClose = li.onClose
 	li.queue = append(li.queue, c0)
 	if li.trackConns {
 		li.conns = append(li.conns, c0)
@@ -124,7 +127,7 @@ type fakeNetConn struct {
 	// peer is the other endpoint.
 	peer *fakeNetConn
 
-	onClose func() // called when closing
+	onClose func(*fakeNetConn) // called when closing
 }
 
 // Read reads data from the connection.
@@ -167,7 +170,7 @@ func (c *fakeNetConn) IsClosedByPeer() bool {
 // Close closes the connection.
 func (c *fakeNetConn) Close() error {
 	if c.onClose != nil {
-		c.onClose()
+		c.onClose(c)
 	}
 	// Local half of the conn is now closed.
 	c.loc.lock()

src/net/http/transport_test.go

@@ -4250,6 +4250,10 @@ func testTransportIdleConnRacesRequest(t testing.TB, mode testMode) {
 	cst.li.onDial = func() {
 		<-dialc
 	}
+	closec := make(chan struct{})
+	cst.li.onClose = func(*fakeNetConn) {
+		<-closec
+	}
 	ctx, cancel := context.WithCancel(context.Background())
 	req1c := make(chan error)
 	go func() {
@@ -4279,10 +4283,6 @@ func testTransportIdleConnRacesRequest(t testing.TB, mode testMode) {
 	//
 	// First: Wait for IdleConnTimeout. The net.Conn.Close blocks.
 	synctest.Wait()
-	closec := make(chan struct{})
-	cst.li.conns[0].peer.onClose = func() {
-		<-closec
-	}
 	time.Sleep(timeout)
 	synctest.Wait()
 	// Make a request, which will use a new connection (since the existing one is closing).