Exclude jdbc wrappers from instrumentation

Author: laurit

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/ConnectionInstrumentation.java

@@ -61,6 +61,10 @@ public static class PrepareAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void addDbInfo(
         @Advice.Argument(0) String sql, @Advice.Return PreparedStatement statement) {
+      if (JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
+        return;
+      }
+
       JdbcData.preparedStatement.set(statement, sql);
     }
   }
@@ -74,6 +78,10 @@ public static void onEnter(
         @Advice.Origin("#m") String methodName,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
+      if (JdbcSingletons.isWrapper(connection, Connection.class)) {
+        return;
+      }
+
       Context parentContext = currentContext();
       DbRequest request =
           DbRequest.createTransaction(connection, methodName.toUpperCase(Locale.ROOT));

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/JdbcSingletons.java

@@ -11,12 +11,15 @@
 import io.opentelemetry.instrumentation.api.incubator.semconv.net.PeerServiceAttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.api.internal.cache.Cache;
 import io.opentelemetry.instrumentation.jdbc.internal.DbRequest;
 import io.opentelemetry.instrumentation.jdbc.internal.JdbcInstrumenterFactory;
 import io.opentelemetry.instrumentation.jdbc.internal.JdbcNetworkAttributesGetter;
 import io.opentelemetry.javaagent.bootstrap.internal.AgentCommonConfig;
 import io.opentelemetry.javaagent.bootstrap.internal.AgentInstrumentationConfig;
 import io.opentelemetry.javaagent.bootstrap.jdbc.DbInfo;
+import java.sql.SQLException;
+import java.sql.Wrapper;
 import java.util.Collections;
 import javax.sql.DataSource;
 
@@ -68,5 +71,31 @@ public static Instrumenter<DataSource, DbInfo> dataSourceInstrumenter() {
     return DATASOURCE_INSTRUMENTER;
   }
 
+  private static final Cache<Class<?>, Boolean> wrapperClassCache = Cache.weak();
+
+  /**
+   * Returns true if the given object is a wrapper and shouldn't be instrumented. We'll instrument
+   * the underlying object called by the wrapper instead.
+   */
+  public static <T extends Wrapper> boolean isWrapper(T object, Class<T> clazz) {
+    return wrapperClassCache.computeIfAbsent(
+        object.getClass(), key -> isWrapperInternal(object, clazz));
+  }
+
+  private static <T extends Wrapper> boolean isWrapperInternal(T object, Class<T> clazz) {
+    try {
+      // we are dealing with a wrapper when the object unwraps to a different instance
+      if (object.isWrapperFor(clazz)) {
+        T unwrapped = object.unwrap(clazz);
+        if (object != unwrapped) {
+          return true;
+        }
+      }
+    } catch (SQLException | AbstractMethodError e) {
+      // ignore
+    }
+    return false;
+  }
+
   private JdbcSingletons() {}
 }

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/PreparedStatementInstrumentation.java

@@ -118,6 +118,10 @@ public static void onEnter(
         return;
       }
 
+      if (JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
+        return;
+      }
+
       // Connection#getMetaData() may execute a Statement or PreparedStatement to retrieve DB info
       // this happens before the DB CLIENT span is started (and put in the current context), so this
       // instrumentation runs again and the shouldStartSpan() check always returns true - and so on
@@ -165,6 +169,10 @@ public static class AddBatchAdvice {
 
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void addBatch(@Advice.This PreparedStatement statement) {
+      if (JdbcSingletons.isWrapper(statement, Statement.class)) {
+        return;
+      }
+
       JdbcData.addPreparedStatementBatch(statement);
     }
   }
@@ -179,6 +187,9 @@ public static void onExit(
       if (!CAPTURE_QUERY_PARAMETERS) {
         return;
       }
+      if (JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
+        return;
+      }
 
       String str = null;
 
@@ -211,6 +222,9 @@ public static void onExit(
       if (!CAPTURE_QUERY_PARAMETERS) {
         return;
       }
+      if (JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
+        return;
+      }
 
       String str = null;
 
@@ -243,6 +257,9 @@ public static void onExit(
       if (!CAPTURE_QUERY_PARAMETERS) {
         return;
       }
+      if (JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
+        return;
+      }
 
       String str = null;
 

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/StatementInstrumentation.java

@@ -72,6 +72,10 @@ public static void onEnter(
         @Advice.Local("otelRequest") DbRequest request,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
+      if (JdbcSingletons.isWrapper(statement, Statement.class)) {
+        return;
+      }
+
       // Connection#getMetaData() may execute a Statement or PreparedStatement to retrieve DB info
       // this happens before the DB CLIENT span is started (and put in the current context), so this
       // instrumentation runs again and the shouldStartSpan() check always returns true - and so on
@@ -102,7 +106,7 @@ public static void stopSpan(
         @Advice.Local("otelRequest") DbRequest request,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
-      if (callDepth.decrementAndGet() > 0) {
+      if (callDepth == null || callDepth.decrementAndGet() > 0) {
         return;
       }
 
@@ -121,6 +125,10 @@ public static void addBatch(@Advice.This Statement statement, @Advice.Argument(0
       if (statement instanceof PreparedStatement) {
         return;
       }
+      if (JdbcSingletons.isWrapper(statement, Statement.class)) {
+        return;
+      }
+
       JdbcData.addStatementBatch(statement, sql);
     }
   }
@@ -144,6 +152,10 @@ public static void onEnter(
         @Advice.Local("otelRequest") DbRequest request,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
+      if (JdbcSingletons.isWrapper(statement, Statement.class)) {
+        return;
+      }
+
       // Connection#getMetaData() may execute a Statement or PreparedStatement to retrieve DB info
       // this happens before the DB CLIENT span is started (and put in the current context), so this
       // instrumentation runs again and the shouldStartSpan() check always returns true - and so on
@@ -190,7 +202,7 @@ public static void stopSpan(
         @Advice.Local("otelRequest") DbRequest request,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
-      if (callDepth.decrementAndGet() > 0) {
+      if (callDepth == null || callDepth.decrementAndGet() > 0) {
         return;
       }