Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,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);
}
}
Expand Down Expand Up @@ -105,6 +109,10 @@ public void end(@Nullable Throwable throwable) {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AdviceScope onEnter(
@Advice.This Connection connection, @Advice.Origin("#m") String methodName) {
if (JdbcSingletons.isWrapper(connection, Connection.class)) {
return null;
}

return AdviceScope.start(connection, methodName);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -68,5 +71,31 @@ public static Instrumenter<DataSource, DbInfo> dataSourceInstrumenter() {
return DATASOURCE_INSTRUMENTER;
}

private static final Cache<Class<?>, Boolean> wrapperClassCache = Cache.weak();
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does ClassValue work here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It doesn't work in the standard way because we need the actual object to decide whether it can be unwrapped. I think the most straightforward way we could use class value would be to use

  private static final ClassValue<IsWrapper> wrapperClassCache =
      new ClassValue<IsWrapper>() {
        @Override
        protected IsWrapper computeValue(Class<?> type) {
          return new IsWrapper();
        }
      };
  
  private static class IsWrapper {
    volatile boolean valueSet;
    volatile boolean value;
    
    boolean getValue(Wrapper object, Class<? extends Wrapper> clazz) {
      if (valueSet) {
        return value;
      }
      value = isWrapperInternal(object, clazz);
      valueSet = true;
      return true;
    }
  }

  /**
   * 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) {
    IsWrapper isWrapper = wrapperClassCache.get(clazz);
    return isWrapper.getValue(object, clazz);
  }

  private static boolean isWrapperInternal(Wrapper object, Class<? extends Wrapper> clazz) {
    try {
      // we are dealing with a wrapper when the object unwraps to a different instance
      if (object.isWrapperFor(clazz)) {
        Wrapper unwrapped = object.unwrap(clazz);
        if (object != unwrapped) {
          return true;
        }
      }
    } catch (SQLException | AbstractMethodError e) {
      // ignore
    }
    return false;
  }

or

  private static final int UNDEFINED = 0;
  private static final int WRAPPER = 1;
  private static final int NOT_WRAPPER = -1;
  private static final ClassValue<AtomicInteger> wrapperClassCache =
      new ClassValue<AtomicInteger>() {
        @Override
        protected AtomicInteger computeValue(Class<?> type) {
          return new AtomicInteger();
        }
      };

  /**
   * 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) {
    AtomicInteger isWrapper = wrapperClassCache.get(clazz);
    int value = isWrapper.get();
    if (value != UNDEFINED) {
      return value == WRAPPER;
    }
    boolean result = isWrapperInternal(object, clazz);
    isWrapper.set(result ? WRAPPER : NOT_WRAPPER);
    return result;
  }

  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)) {
        Wrapper unwrapped = object.unwrap(clazz);
        if (object != unwrapped) {
          return true;
        }
      }
    } catch (SQLException | AbstractMethodError e) {
      // ignore
    }
    return false;
  }

which is slightly more code than the weak cache, but not too bad really.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It doesn't work in the standard way because we need the actual object to decide whether it can be unwrapped

ah, makes sense

I like the IsWrapper version better than the AtomicInteger version, but I also like your existing implementation given that it's not really a standard usage of ClassValue anyways


/**
* 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() {}
}
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,9 @@ public static JdbcAdviceScope onEnter(@Advice.This PreparedStatement statement)
if (JdbcData.preparedStatement.get(statement) == null) {
return null;
}
if (JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
return null;
}

return JdbcAdviceScope.startPreparedStatement(CallDepth.forClass(Statement.class), statement);
}
Expand All @@ -127,6 +130,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);
}
}
Expand All @@ -141,6 +148,9 @@ public static void onExit(
if (!CAPTURE_QUERY_PARAMETERS) {
return;
}
if (JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
return;
}

String str = null;

Expand Down Expand Up @@ -173,6 +183,9 @@ public static void onExit(
if (!CAPTURE_QUERY_PARAMETERS) {
return;
}
if (JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
return;
}

String str = null;

Expand Down Expand Up @@ -205,6 +218,9 @@ public static void onExit(
if (!CAPTURE_QUERY_PARAMETERS) {
return;
}
if (JdbcSingletons.isWrapper(statement, PreparedStatement.class)) {
return;
}

String str = null;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,16 +59,24 @@ public void transform(TypeTransformer transformer) {
@SuppressWarnings("unused")
public static class StatementAdvice {

@Nullable
@Advice.OnMethodEnter(suppress = Throwable.class)
public static JdbcAdviceScope onEnter(
@Advice.Argument(0) String sql, @Advice.This Statement statement) {
if (JdbcSingletons.isWrapper(statement, Statement.class)) {
return null;
}

return JdbcAdviceScope.startStatement(CallDepth.forClass(Statement.class), sql, statement);
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void stopSpan(
@Advice.Thrown @Nullable Throwable throwable, @Advice.Enter JdbcAdviceScope adviceScope) {
adviceScope.end(throwable);
@Advice.Thrown @Nullable Throwable throwable,
@Advice.Enter @Nullable JdbcAdviceScope adviceScope) {
if (adviceScope != null) {
adviceScope.end(throwable);
}
}
}

Expand All @@ -80,6 +88,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);
}
}
Expand All @@ -96,15 +108,23 @@ public static void clearBatch(@Advice.This Statement statement) {
@SuppressWarnings("unused")
public static class ExecuteBatchAdvice {

@Nullable
@Advice.OnMethodEnter(suppress = Throwable.class)
public static JdbcAdviceScope onEnter(@Advice.This Statement statement) {
if (JdbcSingletons.isWrapper(statement, Statement.class)) {
return null;
}

return JdbcAdviceScope.startBatch(CallDepth.forClass(Statement.class), statement);
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void stopSpan(
@Advice.Thrown @Nullable Throwable throwable, @Advice.Enter JdbcAdviceScope adviceScope) {
adviceScope.end(throwable);
@Advice.Thrown @Nullable Throwable throwable,
@Advice.Enter @Nullable JdbcAdviceScope adviceScope) {
if (adviceScope != null) {
adviceScope.end(throwable);
}
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1330,7 +1330,7 @@ void testProxyStatement() throws Exception {
cleanup.deferCleanup(statement);
cleanup.deferCleanup(connection);

Statement proxyStatement = ProxyStatementFactory.proxyStatement(statement);
Statement proxyStatement = ProxyStatementFactory.proxyStatementWithCustomClassLoader(statement);
ResultSet resultSet =
testing.runWithSpan("parent", () -> proxyStatement.executeQuery("SELECT 3"));

Expand Down Expand Up @@ -1737,4 +1737,100 @@ static Stream<Arguments> transactionOperationsStream() throws SQLException {
Arguments.of(
"hsqldb", new JDBCDriver().connect(jdbcUrls.get("hsqldb"), null), "SA", "hsqldb:mem:"));
}

private static PreparedStatement wrapPreparedStatement(PreparedStatement statement) {
return ProxyStatementFactory.proxyPreparedStatement(
(proxy, method, args) -> {
if ("isWrapperFor".equals(method.getName())
&& args.length == 1
&& args[0] == PreparedStatement.class) {
return true;
}
if ("unwrap".equals(method.getName())
&& args.length == 1
&& args[0] == PreparedStatement.class) {
return statement;
}
return testing.runWithSpan("wrapper", () -> method.invoke(statement, args));
});
}

// test that tracing does not start from a wrapper
// regression test for
// https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/14733
@Test
void testPreparedStatementWrapper() throws SQLException {
Connection connection = new org.h2.Driver().connect(jdbcUrls.get("h2"), null);
Connection proxyConnection =
ProxyStatementFactory.proxy(
Connection.class,
(proxy, method, args) -> {
// we don't implement unwrapping here as that would also cause the executeQuery
// instrumentation to get skipped for the prepared statement wrapper
if ("prepareStatement".equals(method.getName())) {
return wrapPreparedStatement((PreparedStatement) method.invoke(connection, args));
}
return method.invoke(connection, args);
});
PreparedStatement statement = proxyConnection.prepareStatement("SELECT 3");
cleanup.deferCleanup(statement);
cleanup.deferCleanup(connection);

ResultSet resultSet = testing.runWithSpan("parent", () -> statement.executeQuery());

resultSet.next();
assertThat(resultSet.getInt(1)).isEqualTo(3);
testing.waitAndAssertTraces(
trace ->
trace.hasSpansSatisfyingExactly(
span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
span ->
span.hasName("wrapper").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0)),
span ->
span.hasName("SELECT " + dbNameLower)
.hasKind(SpanKind.CLIENT)
.hasParent(trace.getSpan(1))));
}

// test that tracing does not start from a wrapper
// regression test for
// https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/14733
@Test
void testStatementWrapper() throws SQLException {
Connection connection = new org.h2.Driver().connect(jdbcUrls.get("h2"), null);
Statement statement = connection.createStatement();
cleanup.deferCleanup(statement);
cleanup.deferCleanup(connection);

Statement proxyStatement =
ProxyStatementFactory.proxyStatement(
(proxy, method, args) -> {
if ("isWrapperFor".equals(method.getName())
&& args.length == 1
&& args[0] == Statement.class) {
return true;
}
if ("unwrap".equals(method.getName())
&& args.length == 1
&& args[0] == Statement.class) {
return statement;
}
return testing.runWithSpan("wrapper", () -> method.invoke(statement, args));
});
ResultSet resultSet =
testing.runWithSpan("parent", () -> proxyStatement.executeQuery("SELECT 3"));

resultSet.next();
assertThat(resultSet.getInt(1)).isEqualTo(3);
testing.waitAndAssertTraces(
trace ->
trace.hasSpansSatisfyingExactly(
span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
span ->
span.hasName("wrapper").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0)),
span ->
span.hasName("SELECT " + dbNameLower)
.hasKind(SpanKind.CLIENT)
.hasParent(trace.getSpan(1))));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -12,50 +12,61 @@

public final class ProxyStatementFactory {

public static Statement proxyStatement(Statement statement) throws Exception {
public static Statement proxyStatementWithCustomClassLoader(Statement statement)
throws Exception {
TestClassLoader classLoader = new TestClassLoader(ProxyStatementFactory.class.getClassLoader());
Class<?> testInterface = classLoader.loadClass(TestInterface.class.getName());
if (testInterface.getClassLoader() != classLoader) {
throw new IllegalStateException("wrong class loader");
}
InvocationHandler invocationHandler = (proxy, method, args) -> method.invoke(statement, args);
Statement proxyStatement =
(Statement)
Proxy.newProxyInstance(
classLoader, new Class<?>[] {Statement.class, testInterface}, invocationHandler);
// adding package private interface TestInterface to jdk proxy forces defining the proxy class
// in the same package as the package private interface
if (!proxyStatement
.getClass()
.getName()
.startsWith("io.opentelemetry.javaagent.instrumentation.jdbc.test")) {
throw new IllegalStateException("proxy statement is in wrong package");
}
return proxy(
Statement.class,
classLoader,
new Class<?>[] {Statement.class, testInterface},
invocationHandler);
}

return proxyStatement;
public static Statement proxyStatement(InvocationHandler invocationHandler) {
return proxy(Statement.class, invocationHandler);
}

public static PreparedStatement proxyPreparedStatement(PreparedStatement statement) {
InvocationHandler invocationHandler = (proxy, method, args) -> method.invoke(statement, args);
PreparedStatement proxyStatement =
(PreparedStatement)
Proxy.newProxyInstance(
ProxyStatementFactory.class.getClassLoader(),
new Class<?>[] {PreparedStatement.class, TestInterface.class},
invocationHandler);
return proxyPreparedStatement(invocationHandler);
}

public static PreparedStatement proxyPreparedStatement(InvocationHandler invocationHandler) {
return proxy(PreparedStatement.class, invocationHandler);
}

public static <T> T proxy(Class<T> clazz, InvocationHandler invocationHandler) {
return proxy(
clazz,
ProxyStatementFactory.class.getClassLoader(),
new Class<?>[] {clazz, TestInterface.class},
invocationHandler);
}

public static <T> T proxy(
Class<T> clazz,
ClassLoader classLoader,
Class<?>[] interfaces,
InvocationHandler invocationHandler) {
T proxy = clazz.cast(Proxy.newProxyInstance(classLoader, interfaces, invocationHandler));

// adding package private interface TestInterface to jdk proxy forces defining the proxy class
// in the same package as the package private interface
// by default we ignore jdk proxies, having the proxy in a different package ensures it gets
// instrumented
if (!proxyStatement
if (!proxy
.getClass()
.getName()
.startsWith("io.opentelemetry.javaagent.instrumentation.jdbc.test")) {
throw new IllegalStateException("proxy statement is in wrong package");
throw new IllegalStateException("proxy is in wrong package");
}

return proxyStatement;
return proxy;
}

private ProxyStatementFactory() {}
Expand Down