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

Code before:
	ct.Close(errors.New("manually closed by client"))
}

type readHangingConn struct {
	net.Conn
	hangConn chan struct{}
	// variable needed to only make read hang when conn is closed
	closed *atomic.Bool
}

func (hc *readHangingConn) Read(b []byte) (n int, err error) {
	n, err = hc.Conn.Read(b)
	if hc.closed.Load() {
		<-hc.hangConn // hang the read till we want
	}
	return n, err
}

func (hc *readHangingConn) Close() error {
	hc.closed.Store(true)
	err := hc.Conn.Close()
	return err
}

// Tests that client does not close until the reader goroutine exits and closes
// once reader goroutine returns.
func (s) TestClientCloseReturnsAfterReaderCompletes(t *testing.T) {
	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}
	isReaderHanging := &atomic.Bool{}
	hangConn := make(chan struct{})
	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()

Review comment (suggestion):
```suggestion
// Tests that closing a client transport does not return until the reader goroutine exits.
```

Code after:
		close(transportClosed)
	}()

	// Wait for a short duration and ensure that the client transport's Close()
	// does not return.
	select {
	case <-transportClosed:
		t.Fatal("Transport closed before reader completed")
	case <-time.After(defaultTestShortTimeout):
	}

	// Closing the channel will unblock the reader goroutine and will ensure
	// that the client transport's Close() returns.
	close(readHangConn)
	select {
	case <-transportClosed:
	case <-time.After(defaultTestTimeout):
		t.Fatal("Timeout when waiting for transport to close")
	}
}

// hangingConn is a net.Conn wrapper for testing, simulating hanging connections
// after a GOAWAY frame is sent, of which Write operations pause until explicitly
// signaled or a timeout occurs.
type hangingConn struct {
	net.Conn
	hangConn     chan struct{}
	startHanging *atomic.Bool
}

func (hc *hangingConn) Write(b []byte) (n int, err error) {
	n, err = hc.Conn.Write(b)
	if hc.startHanging.Load() {
		<-hc.hangConn
	}
	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
	}()