better proxy unwrap

open-telemetry · Nov 13, 2024 · a7f57cf · a7f57cf
1 parent 54a9f00
commit a7f57cf
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 30 deletions.
diff --git a/...emetry/javaagent/instrumentation/spring/webmvc/v3_1/DispatcherServletInstrumentation.java b/...emetry/javaagent/instrumentation/spring/webmvc/v3_1/DispatcherServletInstrumentation.java
@@ -15,9 +15,9 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.bootstrap.IndyProxy;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import java.lang.reflect.Field;
 import java.util.List;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -67,20 +67,8 @@ public static void afterRefresh(
         return;
       }
       Object bean = springCtx.getBean("otelAutoDispatcherFilter");
-      OpenTelemetryHandlerMappingFilter filter;
-      if (bean instanceof OpenTelemetryHandlerMappingFilter) {
-        // TODO: remove this branch once advices are not inlined anymore as it's no longer relevant
-        // inline advice: no proxy class is used
-        filter = (OpenTelemetryHandlerMappingFilter) bean;
-      } else {
-        // non-inlined advice: proxy class is used
-        try {
-          Field delegate = bean.getClass().getField("delegate");
-          filter = (OpenTelemetryHandlerMappingFilter) delegate.get(bean);
-        } catch (NoSuchFieldException | IllegalAccessException e) {
-          throw new IllegalStateException(e);
-        }
-      }
+      OpenTelemetryHandlerMappingFilter filter =
+          IndyProxy.unwrapIfNeeded(bean, OpenTelemetryHandlerMappingFilter.class);
       filter.setHandlerMappings(handlerMappings);
     }
   }

diff --git a/...emetry/javaagent/instrumentation/spring/webmvc/v6_0/DispatcherServletInstrumentation.java b/...emetry/javaagent/instrumentation/spring/webmvc/v6_0/DispatcherServletInstrumentation.java
@@ -15,9 +15,9 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.bootstrap.IndyProxy;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import java.lang.reflect.Field;
 import java.util.List;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -68,19 +68,8 @@ public static void afterRefresh(
       }
 
       Object bean = springCtx.getBean("otelAutoDispatcherFilter");
-      OpenTelemetryHandlerMappingFilter filter;
-      if (bean instanceof OpenTelemetryHandlerMappingFilter) {
-        // inline advice: no proxy class is used
-        filter = (OpenTelemetryHandlerMappingFilter) bean;
-      } else {
-        // non-inlined advice: proxy class is used
-        try {
-          Field delegate = bean.getClass().getField("delegate");
-          filter = (OpenTelemetryHandlerMappingFilter) delegate.get(bean);
-        } catch (NoSuchFieldException | IllegalAccessException e) {
-          throw new IllegalStateException(e);
-        }
-      }
+      OpenTelemetryHandlerMappingFilter filter =
+          IndyProxy.unwrapIfNeeded(bean, OpenTelemetryHandlerMappingFilter.class);
       filter.setHandlerMappings(handlerMappings);
     }
   }

diff --git a/javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/IndyProxy.java b/javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/IndyProxy.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.bootstrap;
+
+import java.lang.reflect.Field;
+
+/** Interface added to proxy classes injected for indy instrumentation */
+@SuppressWarnings("InterfaceWithOnlyStatics") // TODO: remove this once we have a method
+public interface IndyProxy {
+
+  /**
+   * Unwraps the proxied object
+   *
+   * @return unwrapped object
+   */
+  default Object unwrap() {
+    try {
+      // current implementation based on introspection + public delegate field
+      Field delegate = this.getClass().getField("delegate");
+      return delegate.get(this);
+    } catch (NoSuchFieldException | IllegalAccessException e) {
+      throw new IllegalStateException(e);
+    }
+  }
+
+  /**
+   * Transparently unwraps an object that might have been proxied for indy instrumentation. When
+   * unwrapping is not needed, for example for inlined advices, the original object is returned
+   * effectively making this equivalent to a type cast.
+   *
+   * @param o object that might need unwrapping
+   * @param type expected unwrapped object type
+   * @param <T> type of the proxied object
+   * @return unwrapped object
+   * @throws IllegalArgumentException when object can't be cast or unwrapped to the desired type
+   */
+  static <T> T unwrapIfNeeded(Object o, Class<T> type) {
+    if (type.isAssignableFrom(o.getClass())) {
+      return type.cast(o);
+    }
+    if (o instanceof IndyProxy) {
+      Object unwrapped = ((IndyProxy) o).unwrap();
+      if (type.isAssignableFrom(unwrapped.getClass())) {
+        return type.cast(unwrapped);
+      }
+    }
+    throw new IllegalArgumentException("unexpected object unwrap");
+  }
+}
diff --git a/...c/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyProxyFactory.java b/...c/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyProxyFactory.java
@@ -5,12 +5,15 @@
 
 package io.opentelemetry.javaagent.tooling.instrumentation.indy;
 
+import io.opentelemetry.javaagent.bootstrap.IndyProxy;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
+import java.util.ArrayList;
 import java.util.List;
 import net.bytebuddy.ByteBuddy;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.method.ParameterDescription;
+import net.bytebuddy.description.type.TypeDefinition;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.dynamic.scaffold.subclass.ConstructorStrategy;
@@ -87,10 +90,12 @@ public IndyProxyFactory(Method bootstrapMethod, BootstrapArgsProvider bootstrapA
   public DynamicType.Unloaded<?> generateProxy(
       TypeDescription classToProxy, String proxyClassName) {
     TypeDescription.Generic superClass = classToProxy.getSuperClass();
+    List<TypeDefinition> interfaces = new ArrayList<>(classToProxy.getInterfaces());
+    interfaces.add(TypeDescription.ForLoadedType.of(IndyProxy.class));
     DynamicType.Builder<?> builder =
         new ByteBuddy()
             .subclass(superClass, ConstructorStrategy.Default.NO_CONSTRUCTORS)
-            .implement(classToProxy.getInterfaces())
+            .implement(interfaces)
             .name(proxyClassName)
             .annotateType(classToProxy.getDeclaredAnnotations())
             // field must be public to enable resolving the proxy target using introspection

diff --git a/...st/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyProxyFactoryTest.java b/...st/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyProxyFactoryTest.java
@@ -7,6 +7,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import io.opentelemetry.javaagent.bootstrap.IndyProxy;
 import io.opentelemetry.javaagent.tooling.instrumentation.indy.dummies.DummyAnnotation;
 import java.lang.invoke.CallSite;
 import java.lang.invoke.ConstantCallSite;
@@ -319,6 +320,28 @@ void verifyNonPublicMembersIgnored() throws Exception {
         .anySatisfy(method -> assertThat(method.getName()).isEqualTo("publicStaticMethod"));
   }
 
+  @Test
+  void verifyUnwrap() throws Exception {
+    Class<?> proxy = generateProxy(WrappingTest.class);
+    Object proxyInstance = proxy.getConstructor().newInstance();
+    assertThat(proxyInstance)
+        .describedAs("proxies should implement IndyProxy")
+        .isNotInstanceOf(WrappingTest.class)
+        .isInstanceOf(IndyProxy.class);
+
+    WrappingTest unwrapped = IndyProxy.unwrapIfNeeded(proxyInstance, WrappingTest.class);
+    assertThat(unwrapped).isInstanceOf(WrappingTest.class).isNotInstanceOf(IndyProxy.class);
+
+    assertThat(IndyProxy.unwrapIfNeeded(unwrapped, WrappingTest.class))
+        .describedAs("unwrap an already unwrapped is a no-op")
+        .isSameAs(unwrapped);
+  }
+
+  public static class WrappingTest {
+
+    public WrappingTest() {}
+  }
+
   private static Class<?> generateProxy(Class<?> clazz) {
     DynamicType.Unloaded<?> unloaded =
         proxyFactory.generateProxy(