Test early gRPC headers sending

Author: alesj

integration-tests/grpc-interceptors/src/main/java/io/quarkus/grpc/examples/interceptors/EarlyHeaderClientInterceptor.java

@@ -0,0 +1,48 @@
+package io.quarkus.grpc.examples.interceptors;
+
+import java.util.concurrent.CompletableFuture;
+
+import io.grpc.CallOptions;
+import io.grpc.Channel;
+import io.grpc.ClientCall;
+import io.grpc.ClientInterceptor;
+import io.grpc.ForwardingClientCall;
+import io.grpc.ForwardingClientCallListener;
+import io.grpc.Metadata;
+import io.grpc.MethodDescriptor;
+import io.grpc.Status;
+
+public class EarlyHeaderClientInterceptor implements ClientInterceptor {
+
+    private static final Metadata.Key<String> HEADER = Metadata.Key.of("xx-acme-header", Metadata.ASCII_STRING_MARSHALLER);
+    private final CompletableFuture<String> headerFuture = new CompletableFuture<>();
+
+    public CompletableFuture<String> getHeaderFuture() {
+        return headerFuture;
+    }
+
+    @Override
+    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> methodDescriptor,
+            CallOptions callOptions, Channel next) {
+        return new ForwardingClientCall.SimpleForwardingClientCall<>(next.newCall(methodDescriptor, callOptions)) {
+            @Override
+            public void start(Listener<RespT> responseListener, Metadata headers) {
+                super.start(new ForwardingClientCallListener.SimpleForwardingClientCallListener<>(responseListener) {
+                    @Override
+                    public void onHeaders(Metadata headers) {
+                        headerFuture.complete(headers.get(HEADER));
+                        super.onHeaders(headers);
+                    }
+
+                    @Override
+                    public void onClose(Status status, Metadata trailers) {
+                        if (!headerFuture.isDone()) {
+                            headerFuture.completeExceptionally(new RuntimeException("Call closed before receiving headers"));
+                        }
+                        super.onClose(status, trailers);
+                    }
+                }, headers);
+            }
+        };
+    }
+}

integration-tests/grpc-interceptors/src/main/java/io/quarkus/grpc/examples/interceptors/EarlyHeaderServerInterceptor.java

@@ -0,0 +1,65 @@
+package io.quarkus.grpc.examples.interceptors;
+
+import jakarta.enterprise.context.ApplicationScoped;
+
+import io.grpc.ForwardingServerCall;
+import io.grpc.Metadata;
+import io.grpc.ServerCall;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
+import io.grpc.Status;
+import io.quarkus.grpc.GlobalInterceptor;
+
+@GlobalInterceptor
+@ApplicationScoped
+public class EarlyHeaderServerInterceptor implements ServerInterceptor {
+
+    private static final Metadata.Key<String> HEADER = Metadata.Key.of("xx-acme-header", Metadata.ASCII_STRING_MARSHALLER);
+
+    @Override
+    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers,
+            ServerCallHandler<ReqT, RespT> next) {
+        ServerCallDiscardingHeaders<ReqT, RespT> wrappedServerCall = new ServerCallDiscardingHeaders<>(call);
+        ServerCall.Listener<ReqT> serverCallListener = next.startCall(wrappedServerCall, headers);
+        if (wrappedServerCall.isClosed()) {
+            return new ServerCall.Listener<>() {
+            };
+        }
+
+        Metadata metadata = new Metadata();
+        metadata.put(HEADER, "whatever");
+        call.sendHeaders(metadata);
+
+        return serverCallListener;
+    }
+
+    private static class ServerCallDiscardingHeaders<ReqT, RespT>
+            extends ForwardingServerCall.SimpleForwardingServerCall<ReqT, RespT> {
+
+        private boolean closed;
+
+        protected ServerCallDiscardingHeaders(ServerCall<ReqT, RespT> delegate) {
+            super(delegate);
+        }
+
+        @Override
+        public void sendHeaders(Metadata headers) {
+            // headers have been sent already
+        }
+
+        @Override
+        public void sendMessage(RespT message) {
+            super.sendMessage(message);
+        }
+
+        @Override
+        public void close(Status status, Metadata trailers) {
+            closed = true;
+            super.close(status, trailers);
+        }
+
+        boolean isClosed() {
+            return closed;
+        }
+    }
+}

integration-tests/grpc-interceptors/src/test/java/io/quarkus/grpc/example/interceptors/HelloWorldServiceTestBase.java

@@ -17,6 +17,7 @@
 import examples.MutinyCopycatGrpc;
 import io.grpc.Channel;
 import io.grpc.stub.StreamObserver;
+import io.quarkus.grpc.examples.interceptors.EarlyHeaderClientInterceptor;
 import io.quarkus.grpc.examples.interceptors.HelloExceptionHandlerProvider;
 import io.quarkus.grpc.test.utils.GRPCTestUtils;
 import io.smallrye.mutiny.Multi;
@@ -53,6 +54,31 @@ private static void assertException(Throwable t) {
         Assertions.assertTrue(HelloExceptionHandlerProvider.invoked);
     }
 
+    @Test
+    void testEarlyHeaders() throws Exception {
+        EarlyHeaderClientInterceptor earlyHeaderClientInterceptor = new EarlyHeaderClientInterceptor();
+        GreeterGrpc.GreeterStub helloGrpcStub = GreeterGrpc.newStub(channel).withInterceptors(earlyHeaderClientInterceptor);
+        helloGrpcStub.sayHello(HelloRequest.newBuilder().setName("gRPC").build(), new StreamObserver<>() {
+            @Override
+            public void onNext(HelloReply helloReply) {
+                // ignore
+            }
+
+            @Override
+            public void onError(Throwable throwable) {
+                throwable.printStackTrace();
+            }
+
+            @Override
+            public void onCompleted() {
+                // ignore
+            }
+        });
+
+        String headerValue = earlyHeaderClientInterceptor.getHeaderFuture().get(5, TimeUnit.SECONDS);
+        Assertions.assertEquals(headerValue, "whatever");
+    }
+
     @Test
     public void testExceptionHandlerObserver() throws Exception {
         HelloExceptionHandlerProvider.invoked = false;