Fix status trailers in exception handler

Author: dsyer

samples/grpc-server/src/main/java/org/springframework/grpc/sample/GrpcServerApplication.java

@@ -5,7 +5,9 @@
 import org.springframework.context.annotation.Bean;
 import org.springframework.grpc.server.exception.GrpcExceptionHandler;
 
+import io.grpc.Metadata;
 import io.grpc.Status;
+import io.grpc.StatusException;
 
 @SpringBootApplication
 public class GrpcServerApplication {
@@ -18,7 +20,11 @@ public static void main(String[] args) {
 	public GrpcExceptionHandler globalInterceptor() {
 		return exception -> {
 			if (exception instanceof IllegalArgumentException) {
-				return Status.INVALID_ARGUMENT.withDescription(exception.getMessage()).asException();
+				Metadata metadata = new Metadata();
+				metadata.put(Metadata.Key.of("error-code", Metadata.ASCII_STRING_MARSHALLER), "INVALID_ARGUMENT");
+				StatusException result = Status.INVALID_ARGUMENT.withDescription(exception.getMessage())
+					.asException(metadata);
+				return result;
 			}
 			return null;
 		};

samples/grpc-server/src/test/java/org/springframework/grpc/sample/GrpcServerIntegrationTests.java

@@ -44,6 +44,7 @@
 
 import io.grpc.ForwardingServerCallListener;
 import io.grpc.ManagedChannel;
+import io.grpc.Metadata;
 import io.grpc.ServerCall.Listener;
 import io.grpc.ServerInterceptor;
 import io.grpc.Status.Code;
@@ -77,18 +78,19 @@ class ServerWithException {
 		void specificErrorResponse(@Autowired GrpcChannelFactory channels) {
 			SimpleGrpc.SimpleBlockingStub client = SimpleGrpc.newBlockingStub(channels.createChannel("0.0.0.0:0"));
 			assertThat(assertThrows(StatusRuntimeException.class,
-					() -> client.sayHello(HelloRequest.newBuilder().setName("error").build()))
+					() -> client.sayHello(HelloRequest.newBuilder().setName("internal").build()))
 				.getStatus()
-				.getCode()).isEqualTo(Code.INVALID_ARGUMENT);
+				.getCode()).isEqualTo(Code.UNKNOWN);
 		}
 
 		@Test
 		void defaultErrorResponseIsUnknown(@Autowired GrpcChannelFactory channels) {
 			SimpleGrpc.SimpleBlockingStub client = SimpleGrpc.newBlockingStub(channels.createChannel("0.0.0.0:0"));
-			assertThat(assertThrows(StatusRuntimeException.class,
-					() -> client.sayHello(HelloRequest.newBuilder().setName("internal").build()))
-				.getStatus()
-				.getCode()).isEqualTo(Code.UNKNOWN);
+			StatusRuntimeException status = assertThrows(StatusRuntimeException.class,
+					() -> client.sayHello(HelloRequest.newBuilder().setName("error").build()));
+			assertThat(status.getStatus().getCode()).isEqualTo(Code.INVALID_ARGUMENT);
+			assertThat(status.getTrailers().get(Metadata.Key.of("error-code", Metadata.ASCII_STRING_MARSHALLER)))
+				.isNotNull();
 		}
 
 	}

spring-grpc-core/src/main/java/org/springframework/grpc/server/exception/GrpcExceptionHandledServerCall.java

@@ -36,6 +36,7 @@ public void close(Status status, Metadata trailers) {
 		if (status.getCode() == Status.Code.UNKNOWN && status.getCause() != null) {
 			final Throwable cause = status.getCause();
 			final StatusException statusException = this.exceptionHandler.handleException(cause);
+			trailers.merge(statusException.getTrailers());
 			super.close(statusException.getStatus(), trailers);
 		}
 		else {

spring-grpc-core/src/main/java/org/springframework/grpc/server/exception/GrpcExceptionHandlerInterceptor.java

@@ -147,7 +147,7 @@ private void handle(Throwable t) {
 			catch (Throwable e) {
 			}
 			try {
-				this.call.close(status.getStatus(), headers(t));
+				this.call.close(status.getStatus(), headers(status));
 			}
 			catch (Throwable e) {
 				throw new IllegalStateException("Failed to close the call", e);