quic: ensure Endpoint.Close waits for connection loops to exit

In CL 822984, Endpoint.Close was modified to only return when its listen
loop has actually exited.

However, if Endpoint.packetConn is closed prior to Endpoint.Close being
called (as in the case of net/http), Endpoint.Close can still return
prematurely before all connection loops exit.

Fix this by making sure that we also wait for all connection loops to
exit.

For golang/go#70914

Change-Id: I7dbc985475e845c89e2fa9a7755718326a6a6964
Reviewed-on: https://go-review.googlesource.com/c/net/+/823524
LUCI-TryBot-Result: golang-scoped@luci-project-accounts.iam.gserviceaccount.com <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Damien Neil <dneil@google.com>

Author: nicholashusin

quic/endpoint.go

@@ -153,15 +153,21 @@ func (e *Endpoint) Close(ctx context.Context) error {
 	for _, c := range conns {
 		c.Abort(localTransportError{code: errNo})
 	}
+
 	select {
-	case <-e.closec:
 	case <-ctx.Done():
-		for _, c := range conns {
-			c.exit()
-		}
-		<-e.closec
+	case <-e.closec:
+	}
+	for _, c := range conns {
+		c.exit()
+	}
+	// We should only return once all conn loops and the listen loop exit.
+	// That is, there should no longer be any lingering goroutines.
+	for _, c := range conns {
+		<-c.donec
 	}
-	return ctx.Err() // nil if context hasn't expired
+	<-e.closec
+	return ctx.Err()
 }
 
 // Accept waits for and returns the next connection.

quic/endpoint_test.go

@@ -154,16 +154,27 @@ func TestEndpointClosePacketConn(t *testing.T) {
 }
 
 func TestEndpointCloseCanceledContext(t *testing.T) {
-	cli, _ := newLocalConnPair(t, &Config{}, &Config{})
-	// Closing an endpoint with a canceled context terminates connections
-	// immediately, without waiting for peer acknowledgement. However, it
-	// should still wait until the endpoint is actually closed before
-	// returning.
-	if err := cli.endpoint.Close(canceledContext()); !errors.Is(err, context.Canceled) {
-		t.Errorf("Endpoint.Close(canceledContext()) = %v, want %v", err, context.Canceled)
-	}
-	if err := cli.Wait(canceledContext()); !errors.Is(err, errConnClosed) {
-		t.Errorf("Conn.Wait(canceledContext()) = %v, want %v", err, errConnClosed)
+	for _, closePacketConn := range []bool{false, true} {
+		name := "packet conn open"
+		if closePacketConn {
+			name = "packet conn closed"
+		}
+		t.Run(name, func(t *testing.T) {
+			cli, _ := newLocalConnPair(t, &Config{}, &Config{})
+			if closePacketConn {
+				cli.endpoint.packetConn.Close()
+			}
+			// Closing an endpoint with a canceled context terminates connections
+			// immediately, without waiting for peer acknowledgement. However, it
+			// should still wait until the endpoint is actually closed before
+			// returning.
+			if err := cli.endpoint.Close(canceledContext()); !errors.Is(err, context.Canceled) {
+				t.Errorf("Endpoint.Close(canceledContext()) = %v, want %v", err, context.Canceled)
+			}
+			if err := cli.Wait(canceledContext()); !errors.Is(err, errConnClosed) {
+				t.Errorf("Conn.Wait(canceledContext()) = %v, want %v", err, errConnClosed)
+			}
+		})
 	}
 }