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

Code before:
	}

	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"))
}

// deadlineTestConn is a net.Conn wrapper used to assert that deadlines are set
// during http2Client.Close().
type deadlineTestConn struct {
	net.Conn
	observedReadDeadline  atomic.Bool
	observedWriteDeadline atomic.Bool
}

func (c *deadlineTestConn) SetReadDeadline(t time.Time) error {
	c.observedReadDeadline.Store(true)
	return c.Conn.SetReadDeadline(t)
}

func (c *deadlineTestConn) SetWriteDeadline(t time.Time) error {
	c.observedWriteDeadline.Store(true)
	return c.Conn.SetWriteDeadline(t)
}

// Tests that connection read and write deadlines are set as expected during
// Close().
func (s) TestCloseSetsConnectionDeadlines(t *testing.T) {
	dialer := func(_ context.Context, addr string) (net.Conn, error) {
		conn, err := net.Dial("tcp", addr)
		if err != nil {
			return nil, err
		}
		return &deadlineTestConn{Conn: conn}, nil
	}
	co := ConnectOptions{
		Dialer: dialer,
	}
	server, client, cancel := setUpWithOptions(t, 0, &ServerConfig{}, normal, co)
	defer cancel()
	defer server.stop()
	client.Close(fmt.Errorf("closed manually by test"))
	dConn := client.conn.(*deadlineTestConn)
	if !dConn.observedReadDeadline.Load() {
		t.Errorf("Connection read deadline was never set")
	}
	if !dConn.observedWriteDeadline.Load() {
		t.Errorf("Connection write deadline was never set")

Review comment (bug):
I believe we need these atomics because `t.Close` is being called after unblocking the reader of the `readerDone` channel here:
https://github.com/grpc/grpc-go/blob/3c743c9d3f630d12ec805055c304cc30fa6ea89d/internal/transport/http2_client.go#L1637-L1642

Using an atomic will satisfy the race detector, but will make the test flake if the deadline assertion code executed before the deadline is set.

@eshitachandwani is there a reason to close `readerDone` before closing the transport? I tried moving `close(t.readerDone)` below and some tests started failing. I'm not sure if that's a real deadlock of an issue in the tests. 

Code after:
	}

	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"))
}

// deadlineTestConn is a net.Conn wrapper used to assert that deadlines are set
// during http2Client.Close().
type deadlineTestConn struct {
	net.Conn
	// We use atomic.Bool here since there may be more than one call to
	// http2Client.Close -- which sets these deadlines -- and not all of them
	// from the same goroutine as our test. In fact we only care about the first
	// such invocation, which *does* come from the main goroutine of our test,
	// but the race detector can't know that and complains (understandably)
	// about writes from those successive calls when these variables are not
	// atomic.Bool.
	//
	// For more detailed background, see
	// https://github.com/grpc/grpc-go/pull/8534#discussion_r2297717445 .
	observedReadDeadline  atomic.Bool
	observedWriteDeadline atomic.Bool
}

func (c *deadlineTestConn) SetReadDeadline(t time.Time) error {
	c.observedReadDeadline.Store(true)
	return c.Conn.SetReadDeadline(t)
}

func (c *deadlineTestConn) SetWriteDeadline(t time.Time) error {
	c.observedWriteDeadline.Store(true)
	return c.Conn.SetWriteDeadline(t)
}

// Tests that connection read and write deadlines are set as expected during
// Close().
func (s) TestCloseSetsConnectionDeadlines(t *testing.T) {
	dialer := func(_ context.Context, addr string) (net.Conn, error) {
		conn, err := net.Dial("tcp", addr)
		if err != nil {
			return nil, err
		}
		return &deadlineTestConn{Conn: conn}, nil
	}
	co := ConnectOptions{
		Dialer: dialer,
	}