Fixed leaked goroutine

Author: aranjans

internal/transport/controlbuf.go

@@ -193,7 +193,7 @@ type goAway struct {
 	code         http2.ErrCode
 	debugData    []byte
 	headsUp      bool
-	closeConnErr error // if set, loopyWriter will exit, resulting in conn closure
+	closeConnErr error // if set, loopyWriter will exit with this error
 }
 
 func (*goAway) isTransportResponseFrame() bool { return false }

internal/transport/http2_client.go

@@ -407,6 +407,12 @@ func newHTTP2Client(connectCtx, ctx context.Context, addr resolver.Address, opts
 	// returning from this function.
 	readerErrCh := make(chan error, 1)
 	go t.reader(readerErrCh)
+	defer func() {
+		if err != nil {
+			close(t.writerDone)
+			t.Close(err)
+		}
+	}()
 
 	// Send connection preface to server.
 	n, err := t.conn.Write(clientPreface)
@@ -450,6 +456,9 @@ func newHTTP2Client(connectCtx, ctx context.Context, addr resolver.Address, opts
 	if err := t.framer.writer.Flush(); err != nil {
 		return nil, err
 	}
+	if err = <-readerErrCh; err != nil {
+		return nil, err
+	}
 	go func() {
 		t.loopy = newLoopyWriter(clientSide, t.framer, t.controlBuf, t.bdpEst, t.conn, t.logger)
 		t.loopy.ssGoAwayHandler = t.outgoingGoAwayHandler
@@ -460,17 +469,9 @@ func newHTTP2Client(connectCtx, ctx context.Context, addr resolver.Address, opts
 			// after draining any remaining incoming data.
 			t.conn.Close()
 		}
-		t.logger.Infof("Closing writerDone channel")
 		close(t.writerDone)
 	}()
-	defer func() {
-		if err == nil {
-			err = <-readerErrCh
-		}
-		if err != nil {
-			t.Close(err)
-		}
-	}()
+
 	return t, nil
 }
 
@@ -977,7 +978,7 @@ func (t *http2Client) closeStream(s *Stream, err error, rst bool, rstCode http2.
 
 // Close kicks off the shutdown process of the transport. This should be called
 // only once on a transport. Once it is called, the transport should not be
-// accessed any more.
+// accessed anymore.
 func (t *http2Client) Close(err error) {
 	t.mu.Lock()
 	// Make sure we only close once.
@@ -1006,7 +1007,9 @@ func (t *http2Client) Close(err error) {
 	// ever starts to take in an HTTP/2 error code the peer will be able to get more information about the reason
 	// behind the connection close.
 	t.controlBuf.put(&goAway{code: http2.ErrCodeNo, debugData: []byte(fmt.Sprintf("client shutdown with: %v", err)), closeConnErr: err})
-	<-t.writerDone
+	if t.writerDone != nil {
+		<-t.writerDone
+	}
 	t.cancel()
 	t.conn.Close()
 	channelz.RemoveEntry(t.channelz.ID)

internal/transport/transport_test.go

@@ -2670,7 +2670,7 @@ func (s) TestClientSendsAGoAwayFrame(t *testing.T) {
 		t.Fatalf("Error while listening: %v", err)
 	}
 	defer lis.Close()
-	// greetDone channel verifies that greet is done with the connection
+	// greetDone is used to notify when server is done greeting the client.
 	greetDone := make(chan struct{})
 	// successCh verifies that GOAWAY is received at server side
 	successCh := make(chan struct{})
@@ -2694,10 +2694,18 @@ func (s) TestClientSendsAGoAwayFrame(t *testing.T) {
 			t.Errorf("Error while writing settings %v", err)
 			return
 		}
-
-		sfr.ReadFrame()
-		sfr.ReadFrame()
-		sfr.ReadFrame()
+		fr, _ := sfr.ReadFrame()
+		if _, ok := fr.(*http2.SettingsFrame); !ok {
+			t.Errorf("Expected settings frame, got %v", fr)
+		}
+		fr, _ = sfr.ReadFrame()
+		if fr, ok := fr.(*http2.SettingsFrame); !ok && fr.IsAck() {
+			t.Errorf("Expected settings ACK frame, got %v", fr)
+		}
+		fr, _ = sfr.ReadFrame()
+		if fr, ok := fr.(*http2.HeadersFrame); !ok && fr.Flags.Has(http2.FlagHeadersEndStream) {
+			t.Errorf("Expected Headers frame with END_HEADERS frame, got %v", fr)
+		}
 		close(greetDone)
 
 		frame, err := sfr.ReadFrame()
@@ -2712,11 +2720,11 @@ func (s) TestClientSendsAGoAwayFrame(t *testing.T) {
 				t.Logf("Received goAway frame from client")
 				close(successCh)
 			} else {
+				t.Logf("Received unexpected goAway frame from client")
 				close(errorCh)
 			}
 			return
 		default:
-			// The client should have sent any frame other than GOAWAY
 			t.Logf("The server received a frame other than GOAWAY")
 			close(errorCh)
 			return
@@ -2731,7 +2739,7 @@ func (s) TestClientSendsAGoAwayFrame(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to open stream: %v", err)
 	}
-	// Wait until server receives the headers and settings frame as part of greet
+	// Wait until server receives the headers and settings frame as part of greet.
 	<-greetDone
 	ct.Close(errors.New("manually closed by client"))
 	t.Logf("Closed the client connection")

test/goaway_test.go

@@ -761,3 +761,66 @@ func (s) TestTwoGoAwayPingFrames(t *testing.T) {
 		t.Fatalf("Error waiting for graceful shutdown of the server: %v", err)
 	}
 }
+
+// TestClientSendsAGoAway tests the scenario where you get a go away ping
+// frames from the client during graceful shutdown.
+func (s) TestClientSendsAGoAway(t *testing.T) {
+	lis, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatalf("error listening: %v", err)
+	}
+	ctCh := testutils.NewChannel()
+	go func() {
+		conn, err := lis.Accept()
+		if err != nil {
+			t.Errorf("error in lis.Accept(): %v", err)
+		}
+		ct := newClientTester(t, conn)
+		ctCh.Send(ct)
+	}()
+	defer lis.Close()
+
+	cc, err := grpc.Dial(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("error dialing: %v", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	val, err := ctCh.Receive(ctx)
+	if err != nil {
+		t.Fatalf("timeout waiting for client transport (should be given after http2 creation)")
+	}
+	ct := val.(*clientTester)
+	goAwayReceived := make(chan struct{})
+	errCh := make(chan struct{})
+	go func() {
+		for {
+			f, err := ct.fr.ReadFrame()
+			if err != nil {
+				return
+			}
+			switch fr := f.(type) {
+			case *http2.GoAwayFrame:
+				fr = f.(*http2.GoAwayFrame)
+				if fr.ErrCode == http2.ErrCodeNo {
+					t.Logf("GoAway received from client")
+					close(goAwayReceived)
+				}
+			default:
+				t.Errorf("server tester received unexpected frame type %T", f)
+				close(errCh)
+			}
+		}
+	}()
+	cc.Close()
+	defer ct.conn.Close()
+	select {
+	case <-goAwayReceived:
+	case <-errCh:
+		t.Errorf("Error receiving the goAway: %v", err)
+	case <-ctx.Done():
+		t.Errorf("Context timed out")
+	}
+}