From 3d7b37741e4a94f62825e4a47d1221e6563078c6 Mon Sep 17 00:00:00 2001
From: Pawel Krolikowski <pawel@uber.com>
Date: Fri, 25 Jun 2021 16:07:55 -0700
Subject: [PATCH] Print response headers on error responses too

This would be super helpfull during debugging.

It is slightly unclear whether this should be printed on stdout or
stderr .. so I'm opting for the easiest option and printing it as is.
---
 CHANGELOG.md        |  2 +-
 integration_test.go | 20 +++++++++++++++++---
 main.go             | 29 ++++++++++++++++++-----------
 transport/grpc.go   | 13 +++++++++----
 4 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ffc48e97..97d53732 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,7 +2,7 @@ Changelog
 =========
 
 # Unreleased (0.21.0)
-* Nothing yet
+* New: Print response headers on error responses too.
 
 # 0.20.0 (2021-05-18)
 * Add `stream-delay-close-send` option which delays client send stream closure.
diff --git a/integration_test.go b/integration_test.go
index 2ce1dba8..dc26a34f 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -54,9 +54,9 @@ import (
 	yhttp "go.uber.org/yarpc/transport/http"
 	ytchan "go.uber.org/yarpc/transport/tchannel"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/reflection"
 	rpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
-	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
 
@@ -1190,6 +1190,13 @@ func TestGRPCReflectionSource(t *testing.T) {
 					Peers:       []string{addr.String()},
 				},
 			},
+			wantRes: `{
+  "headers": {
+    "content-type": "application/grpc"
+  }
+}
+
+`,
 			wantErr: "Failed while making call: code:unknown message:negative input\n",
 		},
 		{
@@ -1205,6 +1212,13 @@ func TestGRPCReflectionSource(t *testing.T) {
 					Peers:       []string{addr.String()},
 				},
 			},
+			wantRes: `{
+  "headers": {
+    "content-type": "application/grpc"
+  }
+}
+
+`,
 			wantErr: `Failed while making call: code:invalid-argument message:invalid username
 {
   "details": [
@@ -1222,8 +1236,8 @@ func TestGRPCReflectionSource(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.desc, func(t *testing.T) {
 			gotOut, gotErr := runTestWithOpts(tt.opts)
-			assert.Equal(t, gotErr, tt.wantErr)
-			assert.Equal(t, gotOut, tt.wantRes)
+			assert.Equal(t, tt.wantErr, gotErr, "bad error")
+			assert.Equal(t, tt.wantRes, gotOut, "bad response")
 		})
 	}
 }
diff --git a/main.go b/main.go
index 58a0584e..124a47f8 100644
--- a/main.go
+++ b/main.go
@@ -472,6 +472,23 @@ func makeContextWithTrace(ctx context.Context, t transport.Transport, request *t
 
 func makeInitialRequest(out output, transport transport.Transport, serializer encoding.Serializer, req *transport.Request) {
 	response, err := makeRequestWithTracePriority(transport, req, 1)
+
+	var bs []byte
+	outSerialized := map[string]interface{}{}
+	if response != nil && len(response.Headers) > 0 {
+		outSerialized["headers"] = response.Headers
+	}
+	defer func() {
+		if len(outSerialized) == 0 {
+			return
+		}
+		bs, err = json.MarshalIndent(outSerialized, "", "  ")
+		if err != nil {
+			out.Fatalf("Failed to convert map to JSON: %v\nMap: %+v\n", err, outSerialized)
+		}
+		out.Printf("%s\n\n", bs)
+	}()
+
 	if err != nil {
 		buffer := bytes.NewBufferString(err.Error())
 
@@ -502,20 +519,10 @@ func makeInitialRequest(out output, transport transport.Transport, serializer en
 	}
 
 	// Print the initial output body.
-	outSerialized := map[string]interface{}{
-		"body": responseMap,
-	}
-	if len(response.Headers) > 0 {
-		outSerialized["headers"] = response.Headers
-	}
+	outSerialized["body"] = responseMap
 	for k, v := range response.TransportFields {
 		outSerialized[k] = v
 	}
-	bs, err := json.MarshalIndent(outSerialized, "", "  ")
-	if err != nil {
-		out.Fatalf("Failed to convert map to JSON: %v\nMap: %+v\n", err, responseMap)
-	}
-	out.Printf("%s\n\n", bs)
 }
 
 // isYabTemplate is currently very conservative, it requires a file that exists
diff --git a/transport/grpc.go b/transport/grpc.go
index 18572cb9..b28647c8 100644
--- a/transport/grpc.go
+++ b/transport/grpc.go
@@ -128,11 +128,16 @@ func (t *grpcTransport) Call(ctx context.Context, request *Request) (*Response,
 
 	ctx, cancel := requestContextWithTimeout(ctx, request)
 	defer cancel()
+
+	var errs error
+
 	transportResponse, err := t.Outbound.Call(ctx, t.requestToYARPCRequest(request))
-	if err != nil {
-		return nil, err
-	}
-	return yarpcResponseToResponse(transportResponse)
+	errs = multierr.Append(errs, err)
+
+	r, err := yarpcResponseToResponse(transportResponse)
+	errs = multierr.Append(errs, err)
+
+	return r, errs
 }
 
 func (t *grpcTransport) CallStream(ctx context.Context, request *StreamRequest) (*transport.ClientStream, error) {