closes #191

Author: Alexander Furer

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

@@ -1,20 +1,32 @@
 package org.lognet.springboot.grpc;
 
+import io.grpc.Attributes;
+import io.grpc.MethodDescriptor;
+import io.grpc.Status;
 import io.grpc.examples.GreeterGrpc;
+import io.grpc.examples.GreeterOuterClass;
 import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.core.instrument.Tag;
 import io.micrometer.core.instrument.Timer;
 import io.micrometer.core.instrument.simple.SimpleConfig;
 import org.awaitility.Awaitility;
 import org.junit.Before;
 import org.junit.runner.RunWith;
+import org.lognet.springboot.grpc.autoconfigure.metrics.RequestAwareGRpcMetricsTagsContributor;
 import org.lognet.springboot.grpc.context.LocalRunningGrpcPort;
 import org.lognet.springboot.grpc.demo.DemoApp;
+import org.mockito.Mockito;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.context.TestConfiguration;
+import org.springframework.boot.test.mock.mockito.SpyBean;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Import;
 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;
@@ -27,7 +39,43 @@
 @RunWith(SpringRunner.class)
 @SpringBootTest(classes = {DemoApp.class}, webEnvironment = NONE, properties = {"grpc.port=0"})
 @ActiveProfiles("measure")
+@Import(GrpcMeterTest.Config.class)
 public class GrpcMeterTest extends GrpcServerTestBase {
+    @TestConfiguration
+    static class Config{
+        @Bean
+        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) {
+                    return Collections.singletonList(Tag.of("hello",request.getName()));
+                }
+
+                @Override
+                public Iterable<Tag> getTags(Status status, MethodDescriptor<?, ?> methodDescriptor, Attributes attributes) {
+                    return Collections.singletonList(Tag.of("customTagName",status.getCode().name()));
+                }
+            };
+        }
+        @Bean
+        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) {
+                    return Collections.emptyList();
+                }
+
+                @Override
+                public Iterable<Tag> getTags(Status status, MethodDescriptor<?, ?> methodDescriptor, Attributes attributes) {
+                    return Collections.emptyList();
+                }
+            };
+        }
+    }
+
+    @SpyBean
+    private RequestAwareGRpcMetricsTagsContributor<GreeterOuterClass.Person> shouldNotBeInvoked;
+
     @Autowired
     private MeterRegistry registry;
 
@@ -65,5 +113,27 @@ protected void afterGreeting()  {
         final String methodTag = timer.getId().getTag("method");
         assertThat(methodTag,notNullValue());
         assertThat(methodTag,is(GreeterGrpc.getSayHelloMethod().getFullMethodName()));
+
+        final String resultTag = timer.getId().getTag("result");
+        assertThat(resultTag,notNullValue());
+        assertThat(resultTag,is(Status.OK.getCode().name()));
+
+        //from contributor
+
+        final String helloTag = timer.getId().getTag("hello");
+        assertThat(helloTag,notNullValue());
+        assertThat(helloTag,is(name));
+
+        final String customTag = timer.getId().getTag("customTagName");
+        assertThat(customTag,notNullValue());
+        assertThat(customTag,is(Status.OK.getCode().name()));
+
+        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());
+
+
+
+
+
     }
 }

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

@@ -49,6 +49,8 @@ public abstract class GrpcServerTestBase {
     @Autowired
     protected GRpcServerProperties gRpcServerProperties;
 
+    protected String name="John";
+
     @Before
     public   void setupChannels() throws IOException {
         if(gRpcServerProperties.isEnabled()) {
@@ -102,7 +104,7 @@ public void shutdownChannels() {
     public void simpleGreeting() throws ExecutionException, InterruptedException {
 
 
-        String name ="John";
+
         final GreeterGrpc.GreeterFutureStub greeterFutureStub = GreeterGrpc.newFutureStub(selectedChanel);
         final GreeterOuterClass.HelloRequest helloRequest =GreeterOuterClass.HelloRequest.newBuilder().setName(name).build();
         final String reply = beforeGreeting(greeterFutureStub).sayHello(helloRequest).get().getMessage();

grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/autoconfigure/metrics/GRpcMetricsAutoConfiguration.java

@@ -1,19 +1,23 @@
 package org.lognet.springboot.grpc.autoconfigure.metrics;
 
 import io.grpc.ForwardingServerCall;
+import io.grpc.ForwardingServerCallListener;
 import io.grpc.Grpc;
 import io.grpc.Metadata;
 import io.grpc.ServerCall;
 import io.grpc.ServerCallHandler;
 import io.grpc.ServerInterceptor;
 import io.grpc.Status;
 import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.core.instrument.Tag;
 import io.micrometer.core.instrument.Timer;
 import lombok.Setter;
 import lombok.experimental.Accessors;
+import lombok.extern.slf4j.Slf4j;
 import org.lognet.springboot.grpc.GRpcGlobalInterceptor;
 import org.lognet.springboot.grpc.GRpcService;
 import org.lognet.springboot.grpc.autoconfigure.GRpcServerProperties;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.actuate.autoconfigure.metrics.CompositeMeterRegistryAutoConfiguration;
 import org.springframework.boot.actuate.autoconfigure.metrics.MetricsAutoConfiguration;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
@@ -27,7 +31,13 @@
 import org.springframework.core.Ordered;
 
 import java.net.SocketAddress;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
 import java.util.Optional;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 @Configuration
 @AutoConfigureAfter({MetricsAutoConfiguration.class, CompositeMeterRegistryAutoConfiguration.class})
@@ -52,73 +62,136 @@ static class GrpcServiceCondition {
 
     }
 
-    static class MonitoringServerCall<ReqT, RespT> extends  ForwardingServerCall.SimpleForwardingServerCall<ReqT, RespT> {
+    static class MonitoringServerCall<ReqT, RespT> extends ForwardingServerCall.SimpleForwardingServerCall<ReqT, RespT> {
+
 
-        private final boolean addAddressTag;
         private MeterRegistry registry;
+
         final Timer.Sample start;
-        protected MonitoringServerCall(ServerCall<ReqT, RespT> delegate, MeterRegistry registry, boolean addAddressTag) {
+
+        private Collection<GRpcMetricsTagsContributor> tagsContributors;
+        private  List<Tag> additionalTags;
+
+
+
+        protected MonitoringServerCall(ServerCall<ReqT, RespT> delegate, MeterRegistry registry, Collection<GRpcMetricsTagsContributor> tagsContributors) {
             super(delegate);
             this.start = Timer.start(registry);
             this.registry = registry;
-            this.addAddressTag = addAddressTag;
-
-
+            this.tagsContributors = tagsContributors;
         }
+
         @Override
         public void close(Status status, Metadata trailers) {
-            final Timer.Builder timerBuilder = Timer.builder("grpc.server.calls")
-                    .tag("method", getMethodDescriptor().getFullMethodName())
-                    .tag("result", status.getCode().name());
-            if (addAddressTag){
-                Optional.ofNullable(delegate().getAttributes())
-                        .map(a->a.get(Grpc.TRANSPORT_ATTR_LOCAL_ADDR))
-                        .map(SocketAddress::toString)
-                        .ifPresent(a->timerBuilder.tag("address",a));
-            }
+
+            final Timer.Builder timerBuilder = Timer.builder("grpc.server.calls");
+            tagsContributors.forEach(c->
+                timerBuilder.tags(c.getTags(status,getMethodDescriptor(),getAttributes()))
+            );
+            Optional.ofNullable(additionalTags)
+                    .ifPresent(timerBuilder::tags);
+
             start.stop(timerBuilder.register(registry));
 
             super.close(status, trailers);
         }
+
+        public void addTags(List<Tag> tags) {
+            additionalTags  = tags;
+        }
     }
 
+    @Slf4j
     static class MonitoringServerInterceptor implements ServerInterceptor, Ordered {
 
-
         private MeterRegistry registry;
-        private boolean addAddressTag;
+
+
+
+        private Collection<GRpcMetricsTagsContributor> tagsContributors ;
+
+        @Autowired
+        public void setTagsContributors(Collection<GRpcMetricsTagsContributor> tagsContributors) {
+            this.tagsContributors = tagsContributors;
+        }
+
         @Setter
         @Accessors(fluent = true)
         private Integer order;
 
-
-        public MonitoringServerInterceptor(MeterRegistry registry,boolean addAddressTag) {
+        public MonitoringServerInterceptor(MeterRegistry registry) {
             this.registry = registry;
-            this.addAddressTag = addAddressTag;
         }
 
         @Override
         public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
-            return next.startCall(new MonitoringServerCall<>(call,registry,addAddressTag), headers);
 
+            final MonitoringServerCall<ReqT, RespT> monitoringServerCall = new MonitoringServerCall<>(call, registry,tagsContributors);
+            final ServerCall.Listener<ReqT> measuredCall = next.startCall(monitoringServerCall, headers);
+            if (call.getMethodDescriptor().getType().clientSendsOneMessage()) {
+                return new ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT>(measuredCall) {
+                    @Override
+                    public void onMessage(ReqT message) {
+
+                            final Stream<Tag> fd = tagsContributors
+                                    .stream()
+                                    .filter(RequestAwareGRpcMetricsTagsContributor.class::isInstance)
+                                    .map(RequestAwareGRpcMetricsTagsContributor.class::cast)
+                                    .filter(c -> c.accepts(message))
+                                    .flatMap(c -> {
+                                        try {
+                                            return StreamSupport.stream(c.getTags(message,monitoringServerCall.getMethodDescriptor(),monitoringServerCall.getAttributes()).spliterator(), false);
+                                        }catch (Throwable t){
+                                            log.error("Failed to  execute tag contributor",t);
+                                            return Stream.empty();
+                                        }
+                                    });
+
+                            monitoringServerCall.addTags(fd.collect(Collectors.toList()));
+
+                        super.onMessage(message);
+
+                    }
+                };
+            } else {
+                return measuredCall;
+            }
         }
 
-
-
         @Override
         public int getOrder() {
-            return Optional.ofNullable(order).orElse(HIGHEST_PRECEDENCE+20);
+            return Optional.ofNullable(order).orElse(HIGHEST_PRECEDENCE + 20);
         }
     }
 
     @Bean
     @GRpcGlobalInterceptor
-    public ServerInterceptor measure(MeterRegistry registry, GRpcServerProperties properties,GRpcMetricsProperties metricsProperties){
+    public ServerInterceptor measure(MeterRegistry registry, GRpcMetricsProperties metricsProperties) {
+
+        return new MonitoringServerInterceptor(registry)
+                .order(metricsProperties.getInterceptorOrder());
+    }
+
+    @Bean
+    public GRpcMetricsTagsContributor defaultTagsContributor(GRpcServerProperties properties) {
         final Boolean hasMultipleAddresses = Optional.ofNullable(properties.getNettyServer())
                 .map(GRpcServerProperties.NettyServerProperties::getAdditionalListenAddresses)
                 .map(l -> !l.isEmpty())
                 .orElse(false);
-        return  new MonitoringServerInterceptor(registry,hasMultipleAddresses)
-                .order(metricsProperties.getInterceptorOrder());
+
+        return (status, methodDescriptor, attributes) -> {
+
+            final ArrayList<Tag> tags = new ArrayList<>();
+            tags.add(Tag.of("result", status.getCode().name()));
+            tags.add(Tag.of("method", methodDescriptor.getFullMethodName()));
+            if (hasMultipleAddresses) {
+                Optional.ofNullable(attributes)
+                        .map(a -> a.get(Grpc.TRANSPORT_ATTR_LOCAL_ADDR))
+                        .map(SocketAddress::toString)
+                        .map(a -> Tag.of("address", a))
+                        .ifPresent(tags::add);
+            }
+            return tags;
+        };
     }
 }

grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/autoconfigure/metrics/GRpcMetricsTagsContributor.java

@@ -0,0 +1,11 @@
+package org.lognet.springboot.grpc.autoconfigure.metrics;
+
+import io.grpc.Attributes;
+import io.grpc.MethodDescriptor;
+import io.grpc.Status;
+import io.micrometer.core.instrument.Tag;
+@FunctionalInterface
+public interface GRpcMetricsTagsContributor  {
+
+    Iterable<Tag> getTags(Status status, MethodDescriptor<?,?> methodDescriptor, Attributes attributes);
+}

grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/autoconfigure/metrics/RequestAwareGRpcMetricsTagsContributor.java

@@ -0,0 +1,27 @@
+package org.lognet.springboot.grpc.autoconfigure.metrics;
+
+import io.grpc.Attributes;
+import io.grpc.MethodDescriptor;
+import io.grpc.Status;
+import io.micrometer.core.instrument.Tag;
+
+import java.util.Collections;
+
+public abstract class RequestAwareGRpcMetricsTagsContributor<ReqT> implements GRpcMetricsTagsContributor  {
+    private Class<ReqT> tClass;
+
+    protected RequestAwareGRpcMetricsTagsContributor(Class<ReqT> tClass) {
+        this.tClass = tClass;
+    }
+    public final boolean accepts(Object o){
+        return tClass.isInstance(o);
+    }
+
+    @Override
+    public Iterable<Tag> getTags(Status status, MethodDescriptor<?, ?> methodDescriptor, Attributes attributes) {
+        return Collections.emptyList();
+    }
+
+    public abstract Iterable<Tag> getTags(ReqT request,MethodDescriptor<?, ?> methodDescriptor, Attributes attributes);
+
+}