Fix completion stage returns null in aspects (#5742)

Handling the case was missing when `null` was returned from an annotated method with return type CompletionStage, resulting in a NullPointerException.

Fixes #5741

Author: laststem

micrometer-core/src/main/java/io/micrometer/core/aop/CountedAspect.java

@@ -69,6 +69,7 @@
  *
  * @author Ali Dehghani
  * @author Jonatan Ivanov
+ * @author Jeonggi Kim
  * @since 1.2.0
  * @see Counted
  */
@@ -212,8 +213,17 @@ public Object interceptAndRecord(ProceedingJoinPoint pjp, Counted counted) throw
 
         if (stopWhenCompleted) {
             try {
-                return ((CompletionStage<?>) pjp.proceed())
-                    .whenComplete((result, throwable) -> recordCompletionResult(pjp, counted, throwable));
+                Object result = pjp.proceed();
+                if (result == null) {
+                    if (!counted.recordFailuresOnly()) {
+                        record(pjp, counted, DEFAULT_EXCEPTION_TAG_VALUE, RESULT_TAG_SUCCESS_VALUE);
+                    }
+                    return result;
+                }
+                else {
+                    CompletionStage<?> stage = ((CompletionStage<?>) result);
+                    return stage.whenComplete((res, throwable) -> recordCompletionResult(pjp, counted, throwable));
+                }
             }
             catch (Throwable e) {
                 record(pjp, counted, e.getClass().getSimpleName(), RESULT_TAG_FAILURE_VALUE);

micrometer-core/src/main/java/io/micrometer/core/aop/TimedAspect.java

@@ -77,6 +77,7 @@
  * @author Johnny Lim
  * @author Nejc Korasa
  * @author Jonatan Ivanov
+ * @author Jeonggi Kim
  * @since 1.0.0
  */
 @Aspect
@@ -230,8 +231,16 @@ private Object processWithTimer(ProceedingJoinPoint pjp, Timed timed, String met
 
         if (stopWhenCompleted) {
             try {
-                return ((CompletionStage<?>) pjp.proceed()).whenComplete(
-                        (result, throwable) -> record(pjp, timed, metricName, sample, getExceptionTag(throwable)));
+                Object result = pjp.proceed();
+                if (result == null) {
+                    record(pjp, timed, metricName, sample, DEFAULT_EXCEPTION_TAG_VALUE);
+                    return result;
+                }
+                else {
+                    CompletionStage<?> stage = ((CompletionStage<?>) result);
+                    return stage.whenComplete(
+                            (res, throwable) -> record(pjp, timed, metricName, sample, getExceptionTag(throwable)));
+                }
             }
             catch (Throwable e) {
                 record(pjp, timed, metricName, sample, e.getClass().getSimpleName());
@@ -297,8 +306,15 @@ private Object processWithLongTaskTimer(ProceedingJoinPoint pjp, Timed timed, St
 
         if (stopWhenCompleted) {
             try {
-                return ((CompletionStage<?>) pjp.proceed())
-                    .whenComplete((result, throwable) -> sample.ifPresent(this::stopTimer));
+                Object result = pjp.proceed();
+                if (result == null) {
+                    sample.ifPresent(this::stopTimer);
+                    return result;
+                }
+                else {
+                    CompletionStage<?> stage = ((CompletionStage<?>) result);
+                    return stage.whenComplete((res, throwable) -> sample.ifPresent(this::stopTimer));
+                }
             }
             catch (Throwable e) {
                 sample.ifPresent(this::stopTimer);

micrometer-core/src/test/java/io/micrometer/core/aop/CountedAspectTest.java

@@ -211,6 +211,29 @@ void pjpFunctionThrows() {
         assertThat(counter.getId().getDescription()).isNull();
     }
 
+    @Test
+    void countedWithSuccessfulMetricsWhenReturnsCompletionStageNull() {
+        CompletableFuture<?> completableFuture = asyncCountedService.successButNull();
+        assertThat(completableFuture).isNull();
+
+        assertThat(meterRegistry.get("metric.success")
+            .tag("method", "successButNull")
+            .tag("class", getClass().getName() + "$AsyncCountedService")
+            .tag("extra", "tag")
+            .tag("exception", "none")
+            .tag("result", "success")
+            .counter()
+            .count()).isOne();
+    }
+
+    @Test
+    void countedWithoutSuccessfulMetricsWhenReturnsCompletionStageNull() {
+        CompletableFuture<?> completableFuture = asyncCountedService.successButNullWithoutMetrics();
+        assertThat(completableFuture).isNull();
+
+        assertThatThrownBy(() -> meterRegistry.get("metric.none").counter()).isInstanceOf(MeterNotFoundException.class);
+    }
+
     static class CountedService {
 
         @Counted(value = "metric.none", recordFailuresOnly = true)
@@ -272,6 +295,16 @@ CompletableFuture<?> emptyMetricName(GuardedResult guardedResult) {
             return supplyAsync(guardedResult::get);
         }
 
+        @Counted(value = "metric.success", extraTags = { "extra", "tag" })
+        CompletableFuture<?> successButNull() {
+            return null;
+        }
+
+        @Counted(value = "metric.none", recordFailuresOnly = true)
+        CompletableFuture<?> successButNullWithoutMetrics() {
+            return null;
+        }
+
     }
 
     static class GuardedResult {

micrometer-core/src/test/java/io/micrometer/core/aop/TimedAspectTest.java

@@ -217,6 +217,29 @@ void timeMethodWhenCompletedExceptionally() {
             .count()).isEqualTo(1);
     }
 
+    @Test
+    void timeMethodWhenReturnCompletionStageNull() {
+        MeterRegistry registry = new SimpleMeterRegistry();
+
+        AspectJProxyFactory pf = new AspectJProxyFactory(new AsyncTimedService());
+        pf.addAspect(new TimedAspect(registry));
+
+        AsyncTimedService service = pf.getProxy();
+
+        CompletableFuture<?> completableFuture = service.callNull();
+        assertThat(completableFuture).isNull();
+
+        assertThat(registry.getMeters()).isNotEmpty();
+
+        assertThat(registry.get("callNull")
+            .tag("class", getClass().getName() + "$AsyncTimedService")
+            .tag("method", "callNull")
+            .tag("extra", "tag")
+            .tag("exception", "none")
+            .timer()
+            .count()).isEqualTo(1);
+    }
+
     @Test
     void timeMethodWithLongTaskTimerWhenCompleted() {
         MeterRegistry registry = new SimpleMeterRegistry();
@@ -277,6 +300,32 @@ void timeMethodWithLongTaskTimerWhenCompletedExceptionally() {
             .activeTasks()).isEqualTo(0);
     }
 
+    @Test
+    void timeMethodWithLongTaskTimerWhenReturnCompletionStageNull() {
+        MeterRegistry registry = new SimpleMeterRegistry();
+
+        AspectJProxyFactory pf = new AspectJProxyFactory(new AsyncTimedService());
+        pf.addAspect(new TimedAspect(registry));
+
+        AsyncTimedService service = pf.getProxy();
+
+        CompletableFuture<?> completableFuture = service.longCallNull();
+        assertThat(completableFuture).isNull();
+
+        assertThat(registry.get("longCallNull")
+            .tag("class", getClass().getName() + "$AsyncTimedService")
+            .tag("method", "longCallNull")
+            .tag("extra", "tag")
+            .longTaskTimers()).hasSize(1);
+
+        assertThat(registry.find("longCallNull")
+            .tag("class", getClass().getName() + "$AsyncTimedService")
+            .tag("method", "longCallNull")
+            .tag("extra", "tag")
+            .longTaskTimer()
+            .activeTasks()).isEqualTo(0);
+    }
+
     @Test
     void timeMethodFailureWhenCompletedExceptionally() {
         MeterRegistry failingRegistry = new FailingMeterRegistry();
@@ -641,11 +690,21 @@ CompletableFuture<?> call(GuardedResult guardedResult) {
             return supplyAsync(guardedResult::get);
         }
 
+        @Timed(value = "callNull", extraTags = { "extra", "tag" })
+        CompletableFuture<?> callNull() {
+            return null;
+        }
+
         @Timed(value = "longCall", extraTags = { "extra", "tag" }, longTask = true)
         CompletableFuture<?> longCall(GuardedResult guardedResult) {
             return supplyAsync(guardedResult::get);
         }
 
+        @Timed(value = "longCallNull", extraTags = { "extra", "tag" }, longTask = true)
+        CompletableFuture<?> longCallNull() {
+            return null;
+        }
+
     }
 
     static class GuardedResult {

micrometer-observation/src/main/java/io/micrometer/observation/aop/ObservedAspect.java

@@ -68,6 +68,7 @@
  * </pre>
  *
  * @author Jonatan Ivanov
+ * @author Jeonggi Kim
  * @since 1.10.0
  */
 @Aspect
@@ -135,8 +136,15 @@ private Object observe(ProceedingJoinPoint pjp, Method method, Observed observed
             observation.start();
             Observation.Scope scope = observation.openScope();
             try {
-                return ((CompletionStage<?>) pjp.proceed())
-                    .whenComplete((result, error) -> stopObservation(observation, scope, error));
+                Object result = pjp.proceed();
+                if (result == null) {
+                    stopObservation(observation, scope, null);
+                    return result;
+                }
+                else {
+                    CompletionStage<?> stage = (CompletionStage<?>) result;
+                    return stage.whenComplete((res, error) -> stopObservation(observation, scope, error));
+                }
             }
             catch (Throwable error) {
                 stopObservation(observation, scope, error);

micrometer-observation/src/test/java/io/micrometer/observation/aop/ObservedAspectTests.java

@@ -360,6 +360,29 @@ void skipPredicateShouldTakeEffectForClass() {
         TestObservationRegistryAssert.assertThat(registry).doesNotHaveAnyObservation();
     }
 
+    @Test
+    void annotatedAsyncClassCallWithNullShouldBeObserved() {
+        registry.observationConfig().observationHandler(new ObservationTextPublisher());
+
+        AspectJProxyFactory pf = new AspectJProxyFactory(new ObservedClassLevelAnnotatedService());
+        pf.addAspect(new ObservedAspect(registry));
+
+        ObservedClassLevelAnnotatedService service = pf.getProxy();
+        CompletableFuture<String> asyncResult = service.asyncNull();
+        assertThat(asyncResult).isNull();
+
+        TestObservationRegistryAssert.assertThat(registry)
+            .doesNotHaveAnyRemainingCurrentObservation()
+            .hasSingleObservationThat()
+            .hasNameEqualTo("test.class")
+            .hasContextualNameEqualTo("test.class#call")
+            .hasLowCardinalityKeyValue("abc", "123")
+            .hasLowCardinalityKeyValue("test", "42")
+            .hasLowCardinalityKeyValue("class", ObservedClassLevelAnnotatedService.class.getName())
+            .hasLowCardinalityKeyValue("method", "asyncNull")
+            .doesNotHaveError();
+    }
+
     static class ObservedService {
 
         @Observed(name = "test.call", contextualName = "test#call",
@@ -444,6 +467,10 @@ CompletableFuture<String> async(FakeAsyncTask fakeAsyncTask) {
                     contextSnapshot.wrapExecutor(Executors.newSingleThreadExecutor()));
         }
 
+        CompletableFuture<String> asyncNull() {
+            return null;
+        }
+
     }
 
     static class FakeAsyncTask implements Supplier<String> {