client: fix ClientStream.Header() behavior (#6557)

binarylog/binarylog_end2end_test.go

@@ -31,10 +31,12 @@ import (
 	"github.com/golang/protobuf/proto"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/binarylog"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/grpclog"
 	iblog "google.golang.org/grpc/internal/binarylog"
 	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 
@@ -1059,3 +1061,39 @@ func (s) TestServerBinaryLogFullDuplexError(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+// TestCanceledStatus ensures a server that responds with a Canceled status has
+// its trailers logged appropriately and is not treated as a canceled RPC.
+func (s) TestCanceledStatus(t *testing.T) {
+	defer testSink.clear()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	const statusMsgWant = "server returned Canceled"
+	ss := &stubserver.StubServer{
+		UnaryCallF: func(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
+			grpc.SetTrailer(ctx, metadata.Pairs("key", "value"))
+			return nil, status.Error(codes.Canceled, statusMsgWant)
+		},
+	}
+	if err := ss.Start(nil); err != nil {
+		t.Fatalf("Error starting endpoint server: %v", err)
+	}
+	defer ss.Stop()
+
+	if _, err := ss.Client.UnaryCall(ctx, &testpb.SimpleRequest{}); status.Code(err) != codes.Canceled {
+		t.Fatalf("Received unexpected error from UnaryCall: %v; want Canceled", err)
+	}
+
+	got := testSink.logEntries(true)
+	last := got[len(got)-1]
+	if last.Type != binlogpb.GrpcLogEntry_EVENT_TYPE_SERVER_TRAILER ||
+		last.GetTrailer().GetStatusCode() != uint32(codes.Canceled) ||
+		last.GetTrailer().GetStatusMessage() != statusMsgWant ||
+		len(last.GetTrailer().GetMetadata().GetEntry()) != 1 ||
+		last.GetTrailer().GetMetadata().GetEntry()[0].GetKey() != "key" ||
+		string(last.GetTrailer().GetMetadata().GetEntry()[0].GetValue()) != "value" {
+		t.Fatalf("Got binary log: %+v; want last entry is server trailing with status Canceled", got)
+	}
+}

internal/transport/http2_client.go

@@ -1505,30 +1505,28 @@ func (t *http2Client) operateHeaders(frame *http2.MetaHeadersFrame) {
 		return
 	}
 
-	isHeader := false
-
-	// If headerChan hasn't been closed yet
-	if atomic.CompareAndSwapUint32(&s.headerChanClosed, 0, 1) {
-		s.headerValid = true
-		if !endStream {
-			// HEADERS frame block carries a Response-Headers.
-			isHeader = true
+	// For headers, set them in s.header and close headerChan.  For trailers or
+	// trailers-only, closeStream will set the trailers and close headerChan as
+	// needed.
+	if !endStream {
+		// If headerChan hasn't been closed yet (expected, given we checked it
+		// above, but something else could have potentially closed the whole
+		// stream).
+		if atomic.CompareAndSwapUint32(&s.headerChanClosed, 0, 1) {
+			s.headerValid = true
 			// These values can be set without any synchronization because
 			// stream goroutine will read it only after seeing a closed
 			// headerChan which we'll close after setting this.
 			s.recvCompress = recvCompress
 			if len(mdata) > 0 {
 				s.header = mdata
 			}
-		} else {
-			// HEADERS frame block carries a Trailers-Only.
-			s.noHeaders = true
+			close(s.headerChan)
 		}
-		close(s.headerChan)
 	}
 
 	for _, sh := range t.statsHandlers {
-		if isHeader {
+		if !endStream {
 			inHeader := &stats.InHeader{
 				Client:      true,
 				WireLength:  int(frame.Header().Length),
@@ -1554,9 +1552,10 @@ func (t *http2Client) operateHeaders(frame *http2.MetaHeadersFrame) {
 		statusGen = status.New(rawStatusCode, grpcMessage)
 	}
 
-	// if client received END_STREAM from server while stream was still active, send RST_STREAM
-	rst := s.getState() == streamActive
-	t.closeStream(s, io.EOF, rst, http2.ErrCodeNo, statusGen, mdata, true)
+	// If client received END_STREAM from server while stream was still active,
+	// send RST_STREAM.
+	rstStream := s.getState() == streamActive
+	t.closeStream(s, io.EOF, rstStream, http2.ErrCodeNo, statusGen, mdata, true)
 }
 
 // readServerPreface reads and handles the initial settings frame from the

internal/transport/transport.go

@@ -43,10 +43,6 @@ import (
 	"google.golang.org/grpc/tap"
 )
 
-// ErrNoHeaders is used as a signal that a trailers only response was received,
-// and is not a real error.
-var ErrNoHeaders = errors.New("stream has no headers")
-
 const logLevel = 2
 
 type bufferPool struct {
@@ -390,14 +386,10 @@ func (s *Stream) Header() (metadata.MD, error) {
 	}
 	s.waitOnHeader()
 
-	if !s.headerValid {
+	if !s.headerValid || s.noHeaders {
 		return nil, s.status.Err()
 	}
 
-	if s.noHeaders {
-		return nil, ErrNoHeaders
-	}
-
 	return s.header.Copy(), nil
 }
 

rpc_util.go

@@ -867,15 +867,18 @@ func Errorf(c codes.Code, format string, a ...any) error {
 	return status.Errorf(c, format, a...)
 }
 
+var errContextCanceled = status.Error(codes.Canceled, context.Canceled.Error())
+var errContextDeadline = status.Error(codes.DeadlineExceeded, context.DeadlineExceeded.Error())
+
 // toRPCErr converts an error into an error from the status package.
 func toRPCErr(err error) error {
 	switch err {
 	case nil, io.EOF:
 		return err
 	case context.DeadlineExceeded:
-		return status.Error(codes.DeadlineExceeded, err.Error())
+		return errContextDeadline
 	case context.Canceled:
-		return status.Error(codes.Canceled, err.Error())
+		return errContextCanceled
 	case io.ErrUnexpectedEOF:
 		return status.Error(codes.Internal, err.Error())
 	}

stream.go

@@ -789,23 +789,23 @@ func (cs *clientStream) withRetry(op func(a *csAttempt) error, onSuccess func())
 
 func (cs *clientStream) Header() (metadata.MD, error) {
 	var m metadata.MD
-	noHeader := false
 	err := cs.withRetry(func(a *csAttempt) error {
 		var err error
 		m, err = a.s.Header()
-		if err == transport.ErrNoHeaders {
-			noHeader = true
-			return nil
-		}
 		return toRPCErr(err)
 	}, cs.commitAttemptLocked)
 
+	if m == nil && err == nil {
+		// The stream ended with success.  Finish the clientStream.
+		err = io.EOF
+	}
+
 	if err != nil {
 		cs.finish(err)
 		return nil, err
 	}
 
-	if len(cs.binlogs) != 0 && !cs.serverHeaderBinlogged && !noHeader {
+	if len(cs.binlogs) != 0 && !cs.serverHeaderBinlogged && m != nil {
 		// Only log if binary log is on and header has not been logged, and
 		// there is actually headers to log.
 		logEntry := &binarylog.ServerHeader{
@@ -821,6 +821,7 @@ func (cs *clientStream) Header() (metadata.MD, error) {
 			binlog.Log(cs.ctx, logEntry)
 		}
 	}
+
 	return m, nil
 }
 
@@ -929,24 +930,6 @@ func (cs *clientStream) RecvMsg(m any) error {
 	if err != nil || !cs.desc.ServerStreams {
 		// err != nil or non-server-streaming indicates end of stream.
 		cs.finish(err)
-
-		if len(cs.binlogs) != 0 {
-			// finish will not log Trailer. Log Trailer here.
-			logEntry := &binarylog.ServerTrailer{
-				OnClientSide: true,
-				Trailer:      cs.Trailer(),
-				Err:          err,
-			}
-			if logEntry.Err == io.EOF {
-				logEntry.Err = nil
-			}
-			if peer, ok := peer.FromContext(cs.Context()); ok {
-				logEntry.PeerAddr = peer.Addr
-			}
-			for _, binlog := range cs.binlogs {
-				binlog.Log(cs.ctx, logEntry)
-			}
-		}
 	}
 	return err
 }
@@ -1002,18 +985,30 @@ func (cs *clientStream) finish(err error) {
 			}
 		}
 	}
+
 	cs.mu.Unlock()
-	// For binary logging. only log cancel in finish (could be caused by RPC ctx
-	// canceled or ClientConn closed). Trailer will be logged in RecvMsg.
-	//
-	// Only one of cancel or trailer needs to be logged. In the cases where
-	// users don't call RecvMsg, users must have already canceled the RPC.
-	if len(cs.binlogs) != 0 && status.Code(err) == codes.Canceled {
-		c := &binarylog.Cancel{
-			OnClientSide: true,
-		}
-		for _, binlog := range cs.binlogs {
-			binlog.Log(cs.ctx, c)
+	// Only one of cancel or trailer needs to be logged.
+	if len(cs.binlogs) != 0 {
+		switch err {
+		case errContextCanceled, errContextDeadline, ErrClientConnClosing:
+			c := &binarylog.Cancel{
+				OnClientSide: true,
+			}
+			for _, binlog := range cs.binlogs {
+				binlog.Log(cs.ctx, c)
+			}
+		default:
+			logEntry := &binarylog.ServerTrailer{
+				OnClientSide: true,
+				Trailer:      cs.Trailer(),
+				Err:          err,
+			}
+			if peer, ok := peer.FromContext(cs.Context()); ok {
+				logEntry.PeerAddr = peer.Addr
+			}
+			for _, binlog := range cs.binlogs {
+				binlog.Log(cs.ctx, logEntry)
+			}
 		}
 	}
 	if err == nil {

test/end2end_test.go

@@ -6328,12 +6328,11 @@ func (s) TestGlobalBinaryLoggingOptions(t *testing.T) {
 			return &testpb.SimpleResponse{}, nil
 		},
 		FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
-			for {
-				_, err := stream.Recv()
-				if err == io.EOF {
-					return nil
-				}
+			_, err := stream.Recv()
+			if err == io.EOF {
+				return nil
 			}
+			return status.Errorf(codes.Unknown, "expected client to call CloseSend")
 		},
 	}
 

test/retry_test.go

@@ -211,6 +211,11 @@ func (s) TestRetryStreaming(t *testing.T) {
 			return nil
 		}
 	}
+	sHdr := func() serverOp {
+		return func(stream testgrpc.TestService_FullDuplexCallServer) error {
+			return stream.SendHeader(metadata.Pairs("test_header", "test_value"))
+		}
+	}
 	sRes := func(b byte) serverOp {
 		return func(stream testgrpc.TestService_FullDuplexCallServer) error {
 			msg := res(b)
@@ -222,7 +227,7 @@ func (s) TestRetryStreaming(t *testing.T) {
 	}
 	sErr := func(c codes.Code) serverOp {
 		return func(stream testgrpc.TestService_FullDuplexCallServer) error {
-			return status.New(c, "").Err()
+			return status.New(c, "this is a test error").Err()
 		}
 	}
 	sCloseSend := func() serverOp {
@@ -270,7 +275,7 @@ func (s) TestRetryStreaming(t *testing.T) {
 	}
 	cErr := func(c codes.Code) clientOp {
 		return func(stream testgrpc.TestService_FullDuplexCallClient) error {
-			want := status.New(c, "").Err()
+			want := status.New(c, "this is a test error").Err()
 			if c == codes.OK {
 				want = io.EOF
 			}
@@ -309,6 +314,11 @@ func (s) TestRetryStreaming(t *testing.T) {
 	cHdr := func() clientOp {
 		return func(stream testgrpc.TestService_FullDuplexCallClient) error {
 			_, err := stream.Header()
+			if err == io.EOF {
+				// The stream ended successfully; convert to nil to avoid
+				// erroring the test case.
+				err = nil
+			}
 			return err
 		}
 	}
@@ -362,9 +372,13 @@ func (s) TestRetryStreaming(t *testing.T) {
 			sReq(1), sRes(3), sErr(codes.Unavailable),
 		},
 		clientOps: []clientOp{cReq(1), cRes(3), cErr(codes.Unavailable)},
+	}, {
+		desc:      "Retry via ClientStream.Header()",
+		serverOps: []serverOp{sReq(1), sErr(codes.Unavailable), sReq(1), sAttempts(1)},
+		clientOps: []clientOp{cReq(1), cHdr() /* this should cause a retry */, cErr(codes.OK)},
 	}, {
 		desc:      "No retry after header",
-		serverOps: []serverOp{sReq(1), sErr(codes.Unavailable)},
+		serverOps: []serverOp{sReq(1), sHdr(), sErr(codes.Unavailable)},
 		clientOps: []clientOp{cReq(1), cHdr(), cErr(codes.Unavailable)},
 	}, {
 		desc:      "No retry after context",