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

Code before:
}

func setUp(t *testing.T, port int, ht hType) (*server, *http2Client, func()) {
	copts := ConnectOptions{
		BufferPool: mem.DefaultBufferPool(),
	}
	return setUpWithOptions(t, port, &ServerConfig{BufferPool: mem.DefaultBufferPool()}, ht, copts)
}

func setUpWithOptions(t *testing.T, port int, sc *ServerConfig, ht hType, copts ConnectOptions) (*server, *http2Client, func()) {
	server := setUpServerOnly(t, port, sc, ht)
	addr := resolver.Address{Addr: "localhost:" + server.port}
	copts.ChannelzParent = channelzSubChannel(t)

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	t.Cleanup(cancel)
	connectCtx, cCancel := context.WithTimeout(context.Background(), 2*time.Second)
	ct, connErr := NewHTTP2Client(connectCtx, ctx, addr, copts, func(GoAwayReason) {})
	if connErr != nil {
		cCancel() // Do not cancel in success path.
		t.Fatalf("failed to create transport: %v", connErr)
	}
	return server, ct.(*http2Client), cCancel
}

type controllablePingServer struct {
	mu             sync.Mutex
	pingAck        bool
	loopDone       chan struct{}
	pingReceived   chan struct{}
	pingRegistered sync.Once
}

func (s *controllablePingServer) setPingAck(ack bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.pingAck = ack
}

func (s *controllablePingServer) serve(t *testing.T, conn net.Conn) {
	defer close(s.loopDone)
	// Read frame to consume the client preface.
	if _, err := io.ReadFull(conn, make([]byte, len(clientPreface))); err != nil {
		t.Errorf("Error while reading client preface: %v", err)
		return
	}
	// Read ping frames and ack checks.
	framer := http2.NewFramer(conn, conn)
	for {
		f, err := framer.ReadFrame()
		if err != nil {

Review comment (refactor):
If we remove the functionality to wait for the first ping being received, we can have a single field which is an atomic boolean for `pingAck`.  This should simplify things a bit.

Code after:
}

func setUp(t *testing.T, port int, ht hType) (*server, *http2Client, func()) {
	copts := ConnectOptions{
		BufferPool: mem.DefaultBufferPool(),
	}
	return setUpWithOptions(t, port, &ServerConfig{BufferPool: mem.DefaultBufferPool()}, ht, copts)
}

func setUpWithOptions(t *testing.T, port int, sc *ServerConfig, ht hType, copts ConnectOptions) (*server, *http2Client, func()) {
	server := setUpServerOnly(t, port, sc, ht)
	addr := resolver.Address{Addr: "localhost:" + server.port}
	copts.ChannelzParent = channelzSubChannel(t)

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	t.Cleanup(cancel)
	connectCtx, cCancel := context.WithTimeout(context.Background(), 2*time.Second)
	ct, connErr := NewHTTP2Client(connectCtx, ctx, addr, copts, func(GoAwayReason) {})
	if connErr != nil {
		cCancel() // Do not cancel in success path.
		t.Fatalf("failed to create transport: %v", connErr)
	}
	return server, ct.(*http2Client), cCancel
}

type controllablePingServer struct {
	pingAck atomic.Bool
}

func (s *controllablePingServer) setPingAck(ack bool) {
	s.pingAck.Store(ack)
}

func (s *controllablePingServer) serve(t *testing.T, conn net.Conn) {
	// Read frame to consume the client preface.
	if _, err := io.ReadFull(conn, make([]byte, len(clientPreface))); err != nil {
		t.Errorf("Error while reading client preface: %v", err)
		return
	}
	// Read ping frames and ack checks.
	framer := http2.NewFramer(conn, conn)
	for {
		f, err := framer.ReadFrame()
		if err != nil {
			return
		}
		if !s.pingAck.Load() {
			return
		}
		pf, ok := f.(*http2.PingFrame)
		if !ok {