JDBC: use ContextStore when running as agent (#3987)

* JDBC: use ContextStore when running as agent

* address review comments

* use singular
This commit is contained in:
Lauri Tulmin 2021-08-28 00:14:40 +03:00 committed by GitHub
parent 1b78c64e3c
commit b61113fef6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 190 additions and 29 deletions

View File

@ -0,0 +1,42 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/
package io.opentelemetry.javaagent.instrumentation.jdbc;
import io.opentelemetry.instrumentation.api.caching.Cache;
import io.opentelemetry.instrumentation.jdbc.internal.DbInfo;
import io.opentelemetry.instrumentation.jdbc.internal.JdbcData;
import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
import java.sql.Connection;
import java.sql.PreparedStatement;
/**
* Provides means to associate extra info with JDBC {@link Connection} and {@link PreparedStatement}
* using {@link ContextStore} which is more efficient than using a weak map.
*/
public class AgentCacheFactory implements JdbcData.CacheFactory {
private final Cache<Connection, DbInfo> connectionCache;
private final Cache<PreparedStatement, String> preparedStatementCache;
public AgentCacheFactory() {
ContextStore<Connection, DbInfo> connectionContextStore =
InstrumentationContext.get(Connection.class, DbInfo.class);
connectionCache = connectionContextStore.asCache();
ContextStore<PreparedStatement, String> preparedStatementContextStore =
InstrumentationContext.get(PreparedStatement.class, String.class);
preparedStatementCache = preparedStatementContextStore.asCache();
}
@Override
public Cache<Connection, DbInfo> connectionInfoCache() {
return connectionCache;
}
@Override
public Cache<PreparedStatement, String> preparedStatementCache() {
return preparedStatementCache;
}
}

View File

@ -0,0 +1,36 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/
package io.opentelemetry.javaagent.instrumentation.jdbc;
import static net.bytebuddy.matcher.ElementMatchers.none;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
public class AgentCacheInstrumentation implements TypeInstrumentation {
@Override
public ElementMatcher<TypeDescription> typeMatcher() {
return none();
}
@Override
public void transform(TypeTransformer transformer) {
transformer.applyAdviceToMethod(
none(), AgentCacheInstrumentation.class.getName() + "$InitAdvice");
}
@SuppressWarnings("unused")
public static class InitAdvice {
public void init() {
// the sole purpose of this advice is to ensure that AgentDataStoreFactory is recognized
// as helper class and injected into class loader
AgentCacheFactory.class.getName();
}
}
}

View File

@ -12,7 +12,7 @@ import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.returns;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import io.opentelemetry.instrumentation.jdbc.internal.JdbcMaps;
import io.opentelemetry.instrumentation.jdbc.internal.JdbcData;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import java.sql.PreparedStatement;
@ -48,7 +48,7 @@ public class ConnectionInstrumentation implements TypeInstrumentation {
@Advice.OnMethodExit(suppress = Throwable.class)
public static void addDbInfo(
@Advice.Argument(0) String sql, @Advice.Return PreparedStatement statement) {
JdbcMaps.preparedStatements.put(statement, sql);
JdbcData.preparedStatement.put(statement, sql);
}
}
}

View File

@ -14,7 +14,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import io.opentelemetry.instrumentation.jdbc.internal.DbInfo;
import io.opentelemetry.instrumentation.jdbc.internal.JdbcConnectionUrlParser;
import io.opentelemetry.instrumentation.jdbc.internal.JdbcMaps;
import io.opentelemetry.instrumentation.jdbc.internal.JdbcData;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import java.sql.Connection;
@ -58,7 +58,7 @@ public class DriverInstrumentation implements TypeInstrumentation {
return;
}
DbInfo dbInfo = JdbcConnectionUrlParser.parse(url, props);
JdbcMaps.connectionInfo.put(connection, dbInfo);
JdbcData.connectionInfo.put(connection, dbInfo);
}
}
}

View File

@ -24,6 +24,7 @@ public class JdbcInstrumentationModule extends InstrumentationModule {
new ConnectionInstrumentation(),
new DriverInstrumentation(),
new PreparedStatementInstrumentation(),
new StatementInstrumentation());
new StatementInstrumentation(),
new AgentCacheInstrumentation());
}
}

View File

@ -13,7 +13,9 @@ import io.opentelemetry.api.trace.SpanKind
import io.opentelemetry.instrumentation.jdbc.TestConnection
import io.opentelemetry.instrumentation.jdbc.TestDriver
import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
import io.opentelemetry.javaagent.instrumentation.jdbc.AgentCacheFactory
import io.opentelemetry.semconv.trace.attributes.SemanticAttributes
import java.lang.reflect.Field
import java.sql.CallableStatement
import java.sql.Connection
import java.sql.DatabaseMetaData
@ -800,6 +802,17 @@ class JdbcInstrumentationTest extends AgentInstrumentationSpecification {
"getMetaData() uses PreparedStatement, test PreparedStatement" | true | { con, query -> con.prepareStatement(query).executeQuery() }
}
def "should use agent data store"() {
setup:
Class<?> clazz = Class.forName("io.opentelemetry.javaagent.shaded.instrumentation.jdbc.internal.JdbcData")
Field field = clazz.getDeclaredField("cacheFactory")
field.setAccessible(true)
def dataStoreFactory = field.get(null)
expect:
dataStoreFactory.getClass() == AgentCacheFactory
}
class DbCallingConnection extends TestConnection {
final boolean usePreparedStatement

View File

@ -19,7 +19,7 @@ public abstract class DbRequest {
@Nullable
public static DbRequest create(PreparedStatement statement) {
return create(statement, JdbcMaps.preparedStatements.get(statement));
return create(statement, JdbcData.preparedStatement.get(statement));
}
@Nullable

View File

@ -0,0 +1,59 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/
package io.opentelemetry.instrumentation.jdbc.internal;
import io.opentelemetry.instrumentation.api.caching.Cache;
import java.sql.Connection;
import java.sql.PreparedStatement;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/** Holds info associated with JDBC connections and prepared statements. */
public final class JdbcData {
private static final Logger logger = LoggerFactory.getLogger(JdbcData.class);
private static final CacheFactory cacheFactory = getCacheFactory();
public static Cache<Connection, DbInfo> connectionInfo = cacheFactory.connectionInfoCache();
public static Cache<PreparedStatement, String> preparedStatement =
cacheFactory.preparedStatementCache();
private JdbcData() {}
private static CacheFactory getCacheFactory() {
try {
// this class is provided by jdbc javaagent instrumentation
Class<?> clazz =
Class.forName("io.opentelemetry.javaagent.instrumentation.jdbc.AgentCacheFactory");
return (CacheFactory) clazz.getConstructor().newInstance();
} catch (ClassNotFoundException ignored) {
// ignored, this is expected when running as library instrumentation
} catch (Exception exception) {
logger.error("Failed to instantiate AgentCacheFactory", exception);
}
return new DefaultCacheFactory();
}
public interface CacheFactory {
Cache<Connection, DbInfo> connectionInfoCache();
Cache<PreparedStatement, String> preparedStatementCache();
}
private static class DefaultCacheFactory implements CacheFactory {
@Override
public Cache<Connection, DbInfo> connectionInfoCache() {
return Cache.newBuilder().setWeakKeys().build();
}
@Override
public Cache<PreparedStatement, String> preparedStatementCache() {
return Cache.newBuilder().setWeakKeys().build();
}
}
}

View File

@ -1,22 +0,0 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/
package io.opentelemetry.instrumentation.jdbc.internal;
import io.opentelemetry.instrumentation.api.caching.Cache;
import java.sql.Connection;
import java.sql.PreparedStatement;
/**
* JDBC instrumentation shares a global map of connection info.
*
* <p>Should be injected into the bootstrap classpath.
*/
public class JdbcMaps {
public static final Cache<Connection, DbInfo> connectionInfo =
Cache.newBuilder().setWeakKeys().build();
public static final Cache<PreparedStatement, String> preparedStatements =
Cache.newBuilder().setWeakKeys().build();
}

View File

@ -65,7 +65,7 @@ public final class JdbcUtils {
}
public static DbInfo extractDbInfo(Connection connection) {
return JdbcMaps.connectionInfo.computeIfAbsent(connection, JdbcUtils::computeDbInfo);
return JdbcData.connectionInfo.computeIfAbsent(connection, JdbcUtils::computeDbInfo);
}
public static DbInfo computeDbInfo(Connection connection) {

View File

@ -5,6 +5,9 @@
package io.opentelemetry.javaagent.instrumentation.api;
import io.opentelemetry.instrumentation.api.caching.Cache;
import java.util.function.Function;
/**
* Interface to represent context storage for instrumentations.
*
@ -62,4 +65,33 @@ public interface ContextStore<K, C> {
* @return old instance if it was present, or new instance
*/
C putIfAbsent(K key, Factory<C> contextFactory);
/**
* Adapt this context store instance to {@link Cache} interface.
*
* @return {@link Cache} backed by this context store instance
*/
default Cache<K, C> asCache() {
return new Cache<K, C>() {
@Override
public C computeIfAbsent(K key, Function<? super K, ? extends C> mappingFunction) {
return ContextStore.this.putIfAbsent(key, () -> mappingFunction.apply(key));
}
@Override
public C get(K key) {
return ContextStore.this.get(key);
}
@Override
public void put(K key, C value) {
ContextStore.this.put(key, value);
}
@Override
public void remove(K key) {
throw new UnsupportedOperationException("remove not supported");
}
};
}
}