support adding request-dependent tags to stream calls

Author: jGleitz

grpc-spring-boot-starter-demo/src/main/java/org/lognet/springboot/grpc/demo/GreeterService.java

@@ -7,7 +7,6 @@
 import lombok.extern.slf4j.Slf4j;
 import org.lognet.springboot.grpc.GRpcService;
 import org.lognet.springboot.grpc.security.GrpcSecurity;
-import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.security.access.annotation.Secured;
 import org.springframework.security.access.prepost.PostAuthorize;
 import org.springframework.security.access.prepost.PreAuthorize;
@@ -28,6 +27,27 @@ public void sayHello(GreeterOuterClass.HelloRequest request, StreamObserver<Gree
         log.info("Returning " +message);
     }
 
+    @Override public StreamObserver<GreeterOuterClass.HelloRequest> sayManyHellos(
+        StreamObserver<GreeterOuterClass.HelloReply> responseObserver
+    ) {
+        return new StreamObserver<GreeterOuterClass.HelloRequest>() {
+            @Override public void onNext(GreeterOuterClass.HelloRequest request) {
+                String message = "Hello " + request.getName();
+                final GreeterOuterClass.HelloReply.Builder replyBuilder = GreeterOuterClass.HelloReply.newBuilder().setMessage(message);
+                responseObserver.onNext(replyBuilder.build());
+                log.info("Returning " + message);
+            }
+
+            @Override
+            public void onError(Throwable t) {}
+
+            @Override
+            public void onCompleted() {
+                responseObserver.onCompleted();
+            }
+        };
+    }
+
     @Override
     @Secured("SCOPE_profile")
     public void sayAuthHello(Empty request, StreamObserver<GreeterOuterClass.HelloReply> responseObserver) {

grpc-spring-boot-starter-demo/src/main/proto/greeter.proto

@@ -8,6 +8,7 @@ option java_package = "io.grpc.examples";
 service Greeter {
     // Sends a greeting
     rpc SayHello ( HelloRequest) returns (  HelloReply) {}
+    rpc SayManyHellos (stream HelloRequest) returns (stream HelloReply) {}
     rpc SayAuthHello ( google.protobuf.Empty) returns (  HelloReply) {}
     rpc SayAuthOnlyHello ( google.protobuf.Empty) returns (  HelloReply) {}
     rpc SayPreAuthHello ( Person) returns (  Person) {}

grpc-spring-boot-starter-demo/src/test/java/org/lognet/springboot/grpc/GrpcMeterTest.java

@@ -1,17 +1,34 @@
 package org.lognet.springboot.grpc;
 
+import static io.grpc.MethodDescriptor.MethodType.BIDI_STREAMING;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.NONE;
+
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Objects;
+import java.util.concurrent.TimeUnit;
+
 import io.grpc.Attributes;
 import io.grpc.MethodDescriptor;
 import io.grpc.Status;
 import io.grpc.examples.GreeterGrpc;
 import io.grpc.examples.GreeterOuterClass;
+import io.grpc.stub.StreamObserver;
 import io.micrometer.core.instrument.MeterRegistry;
 import io.micrometer.core.instrument.Tag;
+import io.micrometer.core.instrument.Tags;
 import io.micrometer.core.instrument.Timer;
-import io.micrometer.core.instrument.simple.SimpleConfig;
+import io.micrometer.core.instrument.search.MeterNotFoundException;
 import io.micrometer.prometheus.PrometheusConfig;
 import org.awaitility.Awaitility;
 import org.junit.Before;
+import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.lognet.springboot.grpc.autoconfigure.metrics.RequestAwareGRpcMetricsTagsContributor;
 import org.lognet.springboot.grpc.context.LocalRunningGrpcPort;
@@ -26,17 +43,6 @@
 import org.springframework.test.context.ActiveProfiles;
 import org.springframework.test.context.junit4.SpringRunner;
 
-import java.time.Duration;
-import java.util.Collections;
-import java.util.concurrent.TimeUnit;
-
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.containsString;
-import static org.hamcrest.Matchers.greaterThan;
-import static org.hamcrest.Matchers.is;
-import static org.hamcrest.Matchers.notNullValue;
-import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.NONE;
-
 @RunWith(SpringRunner.class)
 @SpringBootTest(classes = {DemoApp.class}, webEnvironment = NONE, properties = {"grpc.port=0"})
 @ActiveProfiles("measure")
@@ -48,7 +54,7 @@ static class Config{
         public RequestAwareGRpcMetricsTagsContributor<GreeterOuterClass.HelloRequest> helloContributor(){
             return new RequestAwareGRpcMetricsTagsContributor<GreeterOuterClass.HelloRequest>(GreeterOuterClass.HelloRequest.class) {
                 @Override
-                public Iterable<Tag> getTags(GreeterOuterClass.HelloRequest request, MethodDescriptor<?, ?> methodDescriptor, Attributes attributes) {
+                public Iterable<Tag> addTags(GreeterOuterClass.HelloRequest request, MethodDescriptor<?, ?> methodDescriptor, Attributes attributes, Tags tags) {
                     return Collections.singletonList(Tag.of("hello",request.getName()));
                 }
 
@@ -62,7 +68,7 @@ public Iterable<Tag> getTags(Status status, MethodDescriptor<?, ?> methodDescrip
         public RequestAwareGRpcMetricsTagsContributor<GreeterOuterClass.Person> shouldNotBeInvoked(){
             return new RequestAwareGRpcMetricsTagsContributor<GreeterOuterClass.Person>(GreeterOuterClass.Person.class) {
                 @Override
-                public Iterable<Tag> getTags(GreeterOuterClass.Person request, MethodDescriptor<?, ?> methodDescriptor, Attributes attributes) {
+                public Iterable<Tag> addTags(GreeterOuterClass.Person request, MethodDescriptor<?, ?> methodDescriptor, Attributes attributes, Tags tags) {
                     return Collections.emptyList();
                 }
 
@@ -72,6 +78,26 @@ public Iterable<Tag> getTags(Status status, MethodDescriptor<?, ?> methodDescrip
                 }
             };
         }
+
+        @Bean
+        public RequestAwareGRpcMetricsTagsContributor<GreeterOuterClass.HelloRequest> multiHelloContributor() {
+            return new RequestAwareGRpcMetricsTagsContributor<GreeterOuterClass.HelloRequest>(GreeterOuterClass.HelloRequest.class, BIDI_STREAMING) {
+                @Override
+                public Tags addTags(GreeterOuterClass.HelloRequest request, MethodDescriptor<?, ?> methodDescriptor, Attributes attributes, Tags existingTags) {
+                    String existingTag = existingTags.stream()
+                        .filter(tag -> tag.getKey().equals("many-hellos"))
+                        .findAny()
+                        .map(Tag::getValue)
+                        .orElse("");
+                    return Tags.of("many-hellos", existingTag.isEmpty() ? request.getName() : existingTag + ", " + request.getName());
+                }
+
+                @Override
+                public Iterable<Tag> getTags(Status status, MethodDescriptor<?, ?> methodDescriptor, Attributes attributes) {
+                    return Collections.singletonList(Tag.of("endTag", status.getCode().name()));
+                }
+            };
+        }
     }
 
     @SpyBean
@@ -131,10 +157,41 @@ protected void afterGreeting()  {
 
         Mockito.verify(shouldNotBeInvoked,Mockito.times(1)).getTags(Mockito.any(Status.class),Mockito.any(),Mockito.any());
         Mockito.verify(shouldNotBeInvoked,Mockito.never()).getTags(Mockito.any(GreeterOuterClass.Person.class),Mockito.any(),Mockito.any());
+        Mockito.verify(shouldNotBeInvoked, Mockito.never())
+            .addTags(Mockito.any(GreeterOuterClass.Person.class), Mockito.any(), Mockito.any(), Mockito.any());
+    }
 
+    @Test
+    public void tagsForStream() {
+        final GreeterGrpc.GreeterStub greeterFutureStub = GreeterGrpc.newStub(selectedChanel);
+        io.grpc.stub.StreamObserver<GreeterOuterClass.HelloRequest> helloInput =
+            greeterFutureStub.sayManyHellos(new StreamObserver<GreeterOuterClass.HelloReply>() {
+                @Override
+                public void onNext(GreeterOuterClass.HelloReply value) {}
 
+                @Override
+                public void onError(Throwable t) {}
 
-
-
+                @Override
+                public void onCompleted() {}
+            });
+        Arrays.asList("a", "b", "c", "d").stream()
+            .map(name -> GreeterOuterClass.HelloRequest.newBuilder().setName(name).build())
+            .forEach(helloInput::onNext);
+        helloInput.onCompleted();
+
+        final Timer timer = Awaitility
+            .waitAtMost(Duration.ofMillis(registryConfig.step().toMillis() * 2))
+            .ignoreExceptionsInstanceOf(MeterNotFoundException.class)
+            .until(
+                () -> registry.get("grpc.server.calls")
+                    .tags("method", "Greeter/SayManyHellos")
+                    .timer(),
+                Objects::nonNull
+            );
+
+        assertThat(timer.totalTime(TimeUnit.MILLISECONDS), greaterThan(0d));
+        assertThat(timer.getId().getTag("many-hellos"), is("a, b, c, d"));
+        assertThat(timer.getId().getTag("endTag"), is("OK"));
     }
 }