Add "cancelled" information to the GrpcServerObservationContext (micr…

…ometer-metrics#5463) Add gRPC cancellation information to the `GrpcServerObservationContext`. Also, improve gRPC unit test reliability by adding proper wait mechanisms. Closes micrometer-metricsgh-5301 Signed-off-by: Tadaya Tsuyukubo <[email protected]>
jonatan-ivanov · Sep 4, 2024 · a720efd · a720efd
1 parent 5cb8706
commit a720efd
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 1 deletion.
diff --git a/...src/main/java/io/micrometer/core/instrument/binder/grpc/GrpcServerObservationContext.java b/...src/main/java/io/micrometer/core/instrument/binder/grpc/GrpcServerObservationContext.java
@@ -50,6 +50,8 @@ public class GrpcServerObservationContext extends RequestReplyReceiverContext<Me
 
     private Metadata trailers;
 
+    private boolean cancelled;
+
     public GrpcServerObservationContext(Getter<Metadata> getter) {
         super(getter);
     }
@@ -140,4 +142,21 @@ public void setTrailers(Metadata trailers) {
         this.trailers = trailers;
     }
 
+    /**
+     * Indicate whether the request is cancelled or not.
+     * @return {@code true} if the request is cancelled
+     * @since 1.14
+     */
+    public boolean isCancelled() {
+        return this.cancelled;
+    }
+
+    /**
+     * Set {@code true} when the request is cancelled.
+     * @since 1.14
+     */
+    public void setCancelled(boolean cancelled) {
+        this.cancelled = cancelled;
+    }
+
 }
diff --git a/...re/src/main/java/io/micrometer/core/instrument/binder/grpc/ObservationGrpcServerCall.java b/...re/src/main/java/io/micrometer/core/instrument/binder/grpc/ObservationGrpcServerCall.java
@@ -67,6 +67,7 @@ public void close(Status status, Metadata trailers) {
         GrpcServerObservationContext context = (GrpcServerObservationContext) this.observation.getContext();
         context.setStatusCode(status.getCode());
         context.setTrailers(trailersToKeep);
+        context.setCancelled(isCancelled());
         super.close(status, trailers);
     }
 

diff --git a/...ter-core/src/test/java/io/micrometer/core/instrument/binder/grpc/GrpcObservationTest.java b/...ter-core/src/test/java/io/micrometer/core/instrument/binder/grpc/GrpcObservationTest.java
@@ -536,6 +536,10 @@ void observationShouldBeCapturedByInterceptor() {
             SimpleRequest request = SimpleRequest.newBuilder().setRequestMessage("Hello").build();
             stub.unaryRpc(request);
 
+            // await until server side processing finishes, otherwise context name might
+            // not be populated.
+            await().until(serverHandler::isContextStopped);
+
             assertThat(scopeAwareServerInterceptor.lastObservation).isNotNull().satisfies((observation -> {
                 assertThat(observation.getContext().getContextualName())
                     .isEqualTo("grpc.testing.SimpleService/UnaryRpc");
@@ -570,7 +574,11 @@ void cancel() {
             assertThat(future.isCancelled()).isTrue();
             TestObservationRegistryAssert.assertThat(observationRegistry)
                 .hasAnObservation(observationContextAssert -> observationContextAssert.hasNameEqualTo("grpc.client")
-                    .hasLowCardinalityKeyValue("grpc.status_code", "CANCELLED"));
+                    .hasLowCardinalityKeyValue("grpc.status_code", "CANCELLED"))
+                .hasAnObservation(observationContextAssert -> observationContextAssert.hasNameEqualTo("grpc.server")
+                    .satisfies(observation -> assertThat(observation).isInstanceOfSatisfying(
+                            GrpcServerObservationContext.class,
+                            context -> assertThat(context.isCancelled()).isTrue())));
             assertThat(serverHandler.getEvents()).contains(GrpcServerEvents.CANCELLED);
         }
 
@@ -744,6 +752,8 @@ static class ContextAndEventHoldingObservationHandler<T extends Observation.Cont
 
         private final Class<T> contextClass;
 
+        private boolean contextStopped;
+
         ContextAndEventHoldingObservationHandler(Class<T> contextClass) {
             this.contextClass = contextClass;
         }
@@ -767,10 +777,21 @@ T getContext() {
             return this.contextHolder.get();
         }
 
+        @Override
+        public void onStop(T context) {
+            if (context.equals(this.contextHolder.get())) {
+                this.contextStopped = true;
+            }
+        }
+
         List<Event> getEvents() {
             return this.events;
         }
 
+        public boolean isContextStopped() {
+            return this.contextStopped;
+        }
+
     }
 
     static class ClientHeaderInterceptor implements ClientInterceptor {