fix: datasource instrumentation to support custom connection (#14602)

Co-authored-by: Lauri Tulmin <[email protected]>

Author: wl2027

instrumentation/jdbc/javaagent/build.gradle.kts

@@ -31,6 +31,7 @@ dependencies {
   testLibrary("org.apache.tomcat:tomcat-juli:7.0.19")
   testLibrary("com.zaxxer:HikariCP:2.4.0")
   testLibrary("com.mchange:c3p0:0.9.5")
+  testLibrary("com.alibaba:druid:1.2.20")
 
   // some classes in earlier versions of derby were split out into derbytools in later versions
   latestDepTestLibrary("org.apache.derby:derbytools:latest.release")

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/datasource/DataSourceInstrumentation.java

@@ -13,6 +13,7 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.jdbc.internal.JdbcUtils;
+import io.opentelemetry.javaagent.bootstrap.CallDepth;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.bootstrap.jdbc.DbInfo;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
@@ -33,7 +34,7 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   @Override
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
-        named("getConnection").and(returns(named("java.sql.Connection"))),
+        named("getConnection").and(returns(implementsInterface(named("java.sql.Connection")))),
         DataSourceInstrumentation.class.getName() + "$GetConnectionAdvice");
   }
 
@@ -44,7 +45,13 @@ public static class GetConnectionAdvice {
     public static void start(
         @Advice.This DataSource ds,
         @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
+        @Advice.Local("otelScope") Scope scope,
+        @Advice.Local("otelCallDepth") CallDepth callDepth) {
+      callDepth = CallDepth.forClass(DataSource.class);
+      if (callDepth.getAndIncrement() > 0) {
+        return;
+      }
+
       Context parentContext = Java8BytecodeBridge.currentContext();
       if (!Java8BytecodeBridge.spanFromContext(parentContext).getSpanContext().isValid()) {
         // this instrumentation is already very noisy, and calls to getConnection outside of an
@@ -64,7 +71,12 @@ public static void stopSpan(
         @Advice.Return Connection connection,
         @Advice.Thrown Throwable throwable,
         @Advice.Local("otelContext") Context context,
-        @Advice.Local("otelScope") Scope scope) {
+        @Advice.Local("otelScope") Scope scope,
+        @Advice.Local("otelCallDepth") CallDepth callDepth) {
+      if (callDepth.decrementAndGet() > 0) {
+        return;
+      }
+
       if (scope == null) {
         return;
       }

instrumentation/jdbc/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jdbc/test/DruicDataSourceTest.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.jdbc.test;
+
+import static io.opentelemetry.api.trace.SpanKind.INTERNAL;
+import static io.opentelemetry.instrumentation.api.internal.SemconvStability.emitStableDatabaseSemconv;
+import static io.opentelemetry.instrumentation.testing.junit.db.SemconvStabilityUtil.maybeStable;
+import static io.opentelemetry.instrumentation.testing.junit.db.SemconvStabilityUtil.maybeStableDbSystemName;
+import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
+import static io.opentelemetry.semconv.incubating.CodeIncubatingAttributes.CODE_FUNCTION;
+import static io.opentelemetry.semconv.incubating.CodeIncubatingAttributes.CODE_NAMESPACE;
+import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_CONNECTION_STRING;
+import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_NAME;
+import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_SYSTEM;
+import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_USER;
+
+import com.alibaba.druid.pool.DruidDataSource;
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import java.sql.Connection;
+import java.sql.SQLException;
+import javax.sql.DataSource;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+@SuppressWarnings("deprecation") // using deprecated semconv
+class DruicDataSourceTest {
+
+  @RegisterExtension
+  static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
+
+  private DataSource dataSource;
+
+  @BeforeEach
+  void setUp() {
+    DruidDataSource druidDataSource = new DruidDataSource();
+    druidDataSource.setUrl("jdbc:h2:mem:test");
+    druidDataSource.setDriverClassName("org.h2.Driver");
+    druidDataSource.setUsername("sa");
+    druidDataSource.setPassword("");
+    druidDataSource.setMaxActive(1);
+    this.dataSource = druidDataSource;
+  }
+
+  @AfterEach
+  void tearDown() {
+    if (dataSource instanceof DruidDataSource) {
+      ((DruidDataSource) dataSource).close();
+    }
+  }
+
+  @Test
+  void testGetConnection() throws SQLException {
+    // In DruidDataSource we instrument both DruidPooledConnection getConnection() and the bridge
+    // method Connection getConnection(). Here we call Connection getConnection() that delegates
+    // to DruidPooledConnection getConnection(), and verify that only one span is created.
+    testing.runWithSpan(
+        "parent",
+        () -> {
+          try (Connection connection = dataSource.getConnection()) {
+            return null;
+          }
+        });
+
+    testing.waitAndAssertTraces(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("parent").hasKind(INTERNAL).hasNoParent(),
+                span ->
+                    span.hasName("DruidDataSource.getConnection")
+                        .hasKind(INTERNAL)
+                        .hasParent(trace.getSpan(0))
+                        .hasAttributesSatisfyingExactly(
+                            equalTo(CODE_NAMESPACE, "com.alibaba.druid.pool.DruidDataSource"),
+                            equalTo(CODE_FUNCTION, "getConnection"),
+                            equalTo(
+                                DB_CONNECTION_STRING,
+                                emitStableDatabaseSemconv() ? null : "h2:mem:"),
+                            equalTo(maybeStable(DB_NAME), "test"),
+                            equalTo(maybeStable(DB_SYSTEM), maybeStableDbSystemName("h2")),
+                            equalTo(DB_USER, emitStableDatabaseSemconv() ? null : "sa"))));
+  }
+}

instrumentation/jdbc/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jdbc/test/JdbcInstrumentationTest.java

@@ -1041,7 +1041,6 @@ void testGetConnection(
       throws SQLException {
     // Tomcat's pool doesn't work because the getConnection method is
     // implemented in a parent class that doesn't implement DataSource
-    boolean recursive = datasource instanceof EmbeddedDataSource;
 
     if (init != null) {
       init.accept(datasource);
@@ -1073,14 +1072,6 @@ void testGetConnection(
                               .hasKind(SpanKind.INTERNAL)
                               .hasParent(trace.getSpan(0))
                               .hasAttributesSatisfyingExactly(attributesAssertions)));
-          if (recursive) {
-            assertions.add(
-                span ->
-                    span.hasName(datasource.getClass().getSimpleName() + ".getConnection")
-                        .hasKind(SpanKind.INTERNAL)
-                        .hasParent(trace.getSpan(1))
-                        .hasAttributesSatisfyingExactly(attributesAssertions));
-          }
           trace.hasSpansSatisfyingExactly(assertions);
         });
   }