From 1c1f93a51a91f5a2c7c4c7ad440a8fa7bf2ee99a Mon Sep 17 00:00:00 2001 From: weil <13622993145@163.com> Date: Tue, 16 Sep 2025 17:58:56 +0800 Subject: [PATCH] fix: datasource instrumentation to support custom connection (#14602) Co-authored-by: Lauri Tulmin --- .../jdbc/javaagent/build.gradle.kts | 1 + .../datasource/DataSourceInstrumentation.java | 18 +++- .../jdbc/test/DruicDataSourceTest.java | 88 +++++++++++++++++++ .../jdbc/test/JdbcInstrumentationTest.java | 9 -- 4 files changed, 104 insertions(+), 12 deletions(-) create mode 100644 instrumentation/jdbc/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jdbc/test/DruicDataSourceTest.java diff --git a/instrumentation/jdbc/javaagent/build.gradle.kts b/instrumentation/jdbc/javaagent/build.gradle.kts index 56308493f1..a03356066f 100644 --- a/instrumentation/jdbc/javaagent/build.gradle.kts +++ b/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") diff --git a/instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/datasource/DataSourceInstrumentation.java b/instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/datasource/DataSourceInstrumentation.java index 009c01d032..708e886a55 100644 --- a/instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/datasource/DataSourceInstrumentation.java +++ b/instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/datasource/DataSourceInstrumentation.java @@ -13,6 +13,7 @@ import static net.bytebuddy.matcher.ElementMatchers.returns; 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 class DataSourceInstrumentation implements TypeInstrumentation { @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 class DataSourceInstrumentation implements TypeInstrumentation { 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 class DataSourceInstrumentation implements TypeInstrumentation { @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; } diff --git a/instrumentation/jdbc/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jdbc/test/DruicDataSourceTest.java b/instrumentation/jdbc/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jdbc/test/DruicDataSourceTest.java new file mode 100644 index 0000000000..de506db95f --- /dev/null +++ b/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")))); + } +} diff --git a/instrumentation/jdbc/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jdbc/test/JdbcInstrumentationTest.java b/instrumentation/jdbc/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jdbc/test/JdbcInstrumentationTest.java index 6c0b89a0f9..9c887b5606 100644 --- a/instrumentation/jdbc/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jdbc/test/JdbcInstrumentationTest.java +++ b/instrumentation/jdbc/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jdbc/test/JdbcInstrumentationTest.java @@ -1041,7 +1041,6 @@ class JdbcInstrumentationTest { 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 @@ class JdbcInstrumentationTest { .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); }); }