Add ObservationApplier abstraction

micrometer-core/build.gradle

@@ -135,6 +135,9 @@ dependencies {
     optionalApi 'org.jetbrains.kotlin:kotlin-stdlib-jdk8'
     optionalApi 'org.jetbrains.kotlinx:kotlinx-coroutines-core'
 
+    // reactor
+    optionalApi 'io.projectreactor:reactor-core'
+
     testImplementation 'io.projectreactor:reactor-test'
     testImplementation project(":micrometer-observation-test")
     java11TestImplementation project(":micrometer-observation-test")

micrometer-core/src/main/kotlin/io/micrometer/core/aop/kotlin/ObservationToSuspendingMethodApplier.kt

@@ -0,0 +1,60 @@
+package io.micrometer.core.aop.kotlin
+
+import io.micrometer.core.instrument.kotlin.asContextElement
+import io.micrometer.observation.Observation
+import io.micrometer.observation.aop.applier.ObservationApplier
+import org.aspectj.lang.ProceedingJoinPoint
+import reactor.core.publisher.Mono
+import java.lang.reflect.Method
+import java.util.Arrays
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.CoroutineContext
+
+object ObservationToSuspendingMethodApplier : ObservationApplier {
+    /**
+     * Using such a simplistic implementation of [kotlin.coroutines.Continuation] interface is only enough,
+     * because of how AopUtils#invokeJoinpointUsingReflection() in Spring Framework is implemented.
+     * This means that this whole implementation is only usable in conjunction with Spring Framework.
+     * @see <a href="https://github.com/spring-projects/spring-framework/blob/main/spring-aop/src/main/java/org/springframework/aop/support/AopUtils.java">AopUtils.java</a>
+     * @see <a href="https://github.com/spring-projects/spring-framework/blob/main/spring-core/src/main/java/org/springframework/core/CoroutinesUtils.java">CoroutinesUtils.java</a>
+     */
+    private class DelegatingContinuation<T>(
+        override val context: CoroutineContext,
+        private val underlying: Continuation<T>
+    ) : Continuation<T> {
+        override fun resumeWith(result: Result<T>) {
+            underlying.resumeWith(result)
+        }
+    }
+
+    override fun isApplicable(pjp: ProceedingJoinPoint, method: Method): Boolean {
+        return method.parameterTypes.isNotEmpty()
+            && Continuation::class.java.isAssignableFrom(method.parameterTypes.last())
+    }
+
+    override fun applyAndProceed(pjp: ProceedingJoinPoint, method: Method, observation: Observation): Any {
+        observation.start()
+        return try {
+            val continuation = pjp.args.last() as Continuation<*>
+            val coroutineContext = continuation.context
+
+            val newCoroutineContext = coroutineContext + observation.openScope().use {
+                observation.observationRegistry.asContextElement()
+            }
+
+            val newArgs = Arrays.copyOf(pjp.args, pjp.args.size)
+            newArgs[newArgs.size - 1] = DelegatingContinuation(newCoroutineContext, continuation)
+
+            @Suppress("ReactiveStreamsUnusedPublisher")
+            (pjp.proceed(newArgs) as Mono<*>).doOnError { error ->
+                observation.error(error)
+            }.doOnTerminate {
+                observation.stop()
+            }
+        } catch (error: Throwable) {
+            observation.error(error)
+            observation.stop()
+            throw error
+        }
+    }
+}

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

@@ -21,12 +21,15 @@
 import io.micrometer.observation.ObservationRegistry;
 import io.micrometer.observation.annotation.Observed;
 import io.micrometer.observation.ObservationConvention;
+import io.micrometer.observation.aop.applier.ObservationApplier;
+import io.micrometer.observation.aop.applier.ObservationApplierRegistry;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.reflect.MethodSignature;
 
 import java.lang.reflect.Method;
+import java.util.Optional;
 import java.util.concurrent.CompletionStage;
 import java.util.function.Predicate;
 
@@ -84,6 +87,8 @@ public class ObservedAspect {
 
     private final Predicate<ProceedingJoinPoint> shouldSkip;
 
+    private final ObservationApplierRegistry observationApplierRegistry;
+
     public ObservedAspect(ObservationRegistry registry) {
         this(registry, null, DONT_SKIP_ANYTHING);
     }
@@ -100,9 +105,17 @@ public ObservedAspect(ObservationRegistry registry, Predicate<ProceedingJoinPoin
     public ObservedAspect(ObservationRegistry registry,
             @Nullable ObservationConvention<ObservedAspectContext> observationConvention,
             Predicate<ProceedingJoinPoint> shouldSkip) {
+        this(registry, observationConvention, shouldSkip, ObservationApplierRegistry.getInstance());
+    }
+
+    public ObservedAspect(ObservationRegistry registry,
+            @Nullable ObservationConvention<ObservedAspectContext> observationConvention,
+            Predicate<ProceedingJoinPoint> shouldSkip,
+            ObservationApplierRegistry observationApplierRegistry) {
         this.registry = registry;
         this.observationConvention = observationConvention;
         this.shouldSkip = shouldSkip;
+        this.observationApplierRegistry = observationApplierRegistry;
     }
 
     @Around("@within(io.micrometer.observation.annotation.Observed) and not @annotation(io.micrometer.observation.annotation.Observed)")
@@ -132,22 +145,10 @@ public Object observeMethod(ProceedingJoinPoint pjp) throws Throwable {
     private Object observe(ProceedingJoinPoint pjp, Method method, Observed observed) throws Throwable {
         Observation observation = ObservedAspectObservationDocumentation.of(pjp, observed, this.registry,
                 this.observationConvention);
-        if (CompletionStage.class.isAssignableFrom(method.getReturnType())) {
-            observation.start();
-            Observation.Scope scope = observation.openScope();
-            try {
-                return ((CompletionStage<?>) pjp.proceed())
-                    .whenComplete((result, error) -> stopObservation(observation, scope, error));
-            }
-            catch (Throwable error) {
-                stopObservation(observation, scope, error);
-                throw error;
-            }
-            finally {
-                scope.close();
-            }
-        }
-        else {
+        Optional<ObservationApplier> observationApplier = observationApplierRegistry.findApplicable(pjp, method);
+        if (observationApplier.isPresent()) {
+            return observationApplier.get().applyAndProceed(pjp, method, observation);
+        } else {
             return observation.observeChecked(() -> pjp.proceed());
         }
     }
@@ -171,14 +172,6 @@ private Method getMethod(ProceedingJoinPoint pjp) throws NoSuchMethodException {
         return method;
     }
 
-    private void stopObservation(Observation observation, Observation.Scope scope, @Nullable Throwable error) {
-        if (error != null) {
-            observation.error(error);
-        }
-        scope.close();
-        observation.stop();
-    }
-
     public static class ObservedAspectContext extends Observation.Context {
 
         private final ProceedingJoinPoint proceedingJoinPoint;

micrometer-observation/src/main/java/io/micrometer/observation/aop/applier/ObservationApplier.java

@@ -0,0 +1,18 @@
+package io.micrometer.observation.aop.applier;
+
+import io.micrometer.common.lang.NonNull;
+import io.micrometer.observation.Observation;
+import org.aspectj.lang.ProceedingJoinPoint;
+
+import java.lang.reflect.Method;
+
+public interface ObservationApplier {
+
+    boolean isApplicable(@NonNull ProceedingJoinPoint pjp, @NonNull Method method);
+
+    Object applyAndProceed(
+        @NonNull ProceedingJoinPoint pjp,
+        @NonNull Method method,
+        @NonNull Observation observation
+    ) throws Throwable;
+}

micrometer-observation/src/main/java/io/micrometer/observation/aop/applier/ObservationApplierRegistry.java

@@ -0,0 +1,42 @@
+package io.micrometer.observation.aop.applier;
+
+import org.aspectj.lang.ProceedingJoinPoint;
+
+import java.lang.reflect.Method;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+public class ObservationApplierRegistry {
+
+    private static ObservationApplierRegistry instance;
+
+    private final List<ObservationApplier> observationAppliers;
+
+    public ObservationApplierRegistry(List<ObservationApplier> observationAppliers) {
+        this.observationAppliers = new CopyOnWriteArrayList<>(observationAppliers);
+    }
+
+    public Optional<ObservationApplier> findApplicable(ProceedingJoinPoint pjp, Method method) {
+        for (ObservationApplier observationApplier : observationAppliers) {
+            if (observationApplier.isApplicable(pjp, method)) {
+                return Optional.of(observationApplier);
+            }
+        }
+        return Optional.empty();
+    }
+
+    public void register(ObservationApplier observationApplier) {
+        this.observationAppliers.add(observationApplier);
+    }
+
+    public static ObservationApplierRegistry getInstance() {
+        if (instance == null) {
+            instance = new ObservationApplierRegistry(
+                Collections.singletonList(new ObservationToCompletionStageApplier())
+            );
+        }
+        return instance;
+    }
+}

micrometer-observation/src/main/java/io/micrometer/observation/aop/applier/ObservationToCompletionStageApplier.java

@@ -0,0 +1,43 @@
+package io.micrometer.observation.aop.applier;
+
+import io.micrometer.common.lang.NonNull;
+import io.micrometer.common.lang.Nullable;
+import io.micrometer.observation.Observation;
+import org.aspectj.lang.ProceedingJoinPoint;
+
+import java.lang.reflect.Method;
+import java.util.concurrent.CompletionStage;
+
+public class ObservationToCompletionStageApplier implements ObservationApplier {
+    @Override
+    public boolean isApplicable(@NonNull ProceedingJoinPoint pjp, @NonNull Method method) {
+        return CompletionStage.class.isAssignableFrom(method.getReturnType());
+    }
+
+    @Override
+    public Object applyAndProceed(
+        @NonNull ProceedingJoinPoint pjp,
+        @NonNull Method method,
+        @NonNull Observation observation
+    ) throws Throwable {
+        observation.start();
+        Observation.Scope scope = observation.openScope();
+        try {
+            return ((CompletionStage<?>) pjp.proceed())
+                .whenComplete((result, error) -> stopObservation(observation, scope, error));
+        } catch (Throwable error) {
+            stopObservation(observation, scope, error);
+            throw error;
+        } finally {
+            scope.close();
+        }
+    }
+
+    private void stopObservation(Observation observation, Observation.Scope scope, @Nullable Throwable error) {
+        if (error != null) {
+            observation.error(error);
+        }
+        scope.close();
+        observation.stop();
+    }
+}