Repository: grpc/grpc-go
File: internal/transport/transport_test.go (Go)

Code before:
	dialer := func(_ context.Context, addr string) (net.Conn, error) {
		conn, err := net.Dial("tcp", addr)
		if err != nil {
			return nil, err
		}
		return &readHangingConn{Conn: conn, hangConn: hangConn, closed: isReaderHanging}, nil
	}
	copts := ConnectOptions{Dialer: dialer}
	copts.ChannelzParent = channelzSubChannel(t)
	// Create client transport with custom dialer
	ct, connErr := NewClientTransport(connectCtx, context.Background(), addr, copts, func(GoAwayReason) {})
	if connErr != nil {
		t.Fatalf("failed to create transport: %v", connErr)
	}

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	if _, err := ct.NewStream(ctx, &CallHdr{}); err != nil {
		t.Fatalf("Failed to open stream: %v", err)
	}

	transportClosed := make(chan struct{})
	go func() {
		ct.Close(errors.New("manually closed by client"))
		close(transportClosed)
	}()

	select {
	case <-transportClosed:
		t.Fatal("Transport closed before reader completed")
	case <-time.After(defaultTestShortTimeout):
	}
	close(hangConn)
	select {
	case <-transportClosed:
	case <-time.After(defaultTestTimeout):
		t.Fatal("Timeout when waiting for transport to close")
	}
}

// TestReadHeaderMultipleBuffers tests the stream when the gRPC headers are
// split across multiple buffers. It verifies that the reporting of the
// number of bytes read for flow control is correct.
func (s) TestReadHeaderMultipleBuffers(t *testing.T) {
	headerLen := 5
	recvBuffer := newRecvBuffer()
	recvBuffer.put(recvMsg{buffer: make(mem.SliceBuffer, 3)})
	recvBuffer.put(recvMsg{buffer: make(mem.SliceBuffer, headerLen-3)})
	bytesRead := 0
	s := Stream{
		requestRead: func(int) {},

Review comment (suggestion):
Add a comment at the start of this block? Something like:

```golang
// Closing the client transport will result in the underlying net.Conn being closed, 
// which will result in readHangingConn.Read() to hang. This will stall the exit of the reader goroutine,
// and will stall client transport's Close from returning.
```

Code after:
	}
	return n, err
}

// Tests the scenario where a client transport is closed and writing of the
// GOAWAY frame as part of the close does not complete because of a network
// hang. The test verifies that the client transport is closed without waiting
// for too long.
func (s) TestClientCloseReturnsEarlyWhenGoAwayWriteHangs(t *testing.T) {
	// Override timer for writing GOAWAY to 0 so that the connection write
	// always times out. It is equivalent of real network hang when conn
	// write for goaway doesn't finish in specified deadline
	origGoAwayLoopyTimeout := goAwayLoopyWriterTimeout
	goAwayLoopyWriterTimeout = time.Millisecond
	defer func() {
		goAwayLoopyWriterTimeout = origGoAwayLoopyTimeout
	}()

	// Create the server set up.
	connectCtx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	server := setUpServerOnly(t, 0, &ServerConfig{}, normal)
	defer server.stop()
	addr := resolver.Address{Addr: "localhost:" + server.port}
	isGreetingDone := &atomic.Bool{}
	hangConn := make(chan struct{})
	defer close(hangConn)
	dialer := func(_ context.Context, addr string) (net.Conn, error) {
		conn, err := net.Dial("tcp", addr)
		if err != nil {
			return nil, err
		}
		return &hangingConn{Conn: conn, hangConn: hangConn, startHanging: isGreetingDone}, nil
	}
	copts := ConnectOptions{Dialer: dialer}
	copts.ChannelzParent = channelzSubChannel(t)
	// Create client transport with custom dialer
	ct, connErr := NewClientTransport(connectCtx, context.Background(), addr, copts, func(GoAwayReason) {})
	if connErr != nil {
		t.Fatalf("failed to create transport: %v", connErr)
	}

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	if _, err := ct.NewStream(ctx, &CallHdr{}); err != nil {
		t.Fatalf("Failed to open stream: %v", err)
	}

	isGreetingDone.Store(true)
	ct.Close(errors.New("manually closed by client"))
}