Merge pull request #2963 from newrelic/jdbc-generic-nf-replacement

jtduffy · web-flow · commit 91fe8ffb988d · 2026-07-02T07:45:31.000-04:00
Jdbc generic new fields replacement
diff --git a/.fleetControl/schemaGeneration/reference-newrelic.yml b/.fleetControl/schemaGeneration/reference-newrelic.yml
@@ -690,6 +690,11 @@ common: &default_settings
     # Default: 10
     gc_cpu_threshold: 10
 
+  # JDBC Statements are cached in a weak key map by default. Set enabled:false to use a vanilla ConcurrentHashMap instead.
+  # Users who set enabled:false MUST call Statement.close() on every JDBC Statement.
+  jdbc_statement_weak_key_caching:
+    enabled: true
+
   # Class transformer can be used to disable all agent instrumentation or specific instrumentation modules.
   # All instrumentation modules can be found here: https://github.com/newrelic/newrelic-java-agent/tree/main/instrumentation
   class_transformer:
diff --git a/.fleetControl/schemas/config.json b/.fleetControl/schemas/config.json
@@ -1197,6 +1197,17 @@
             "additionalProperties": true,
             "description": "Circuit breaker protection settings."
         },
+        "jdbc_statement_weak_key_caching": {
+            "type": "object",
+            "properties": {
+                "enabled": {
+                    "type": "boolean",
+                    "default": true
+                }
+            },
+            "additionalProperties": true,
+            "description": "JDBC Statements are cached in a weak key map by default. Set enabled:false to use a vanilla ConcurrentHashMap instead. Users who set enabled:false MUST call Statement.close() on every JDBC Statement."
+        },
         "class_transformer": {
             "type": "object",
             "properties": {
diff --git a/agent-bridge-datastore/build.gradle b/agent-bridge-datastore/build.gradle
@@ -40,6 +40,7 @@ dependencies {
     implementation(project(":agent-bridge"))
 
     testImplementation('org.awaitility:awaitility:4.2.0')
+    testImplementation("org.mockito:mockito-inline:4.11.0")
     testImplementation(fileTree(dir: 'src/test/resources/com/newrelic/agent/bridge', include: '*.jar'))
 }
 
diff --git a/agent-bridge-datastore/src/main/java/com/newrelic/agent/bridge/datastore/JdbcHelper.java b/agent-bridge-datastore/src/main/java/com/newrelic/agent/bridge/datastore/JdbcHelper.java
@@ -8,7 +8,6 @@
 package com.newrelic.agent.bridge.datastore;
 
 import com.newrelic.agent.bridge.AgentBridge;
-import com.newrelic.agent.bridge.NoOpTransaction;
 import com.newrelic.api.agent.Logger;
 import com.newrelic.api.agent.NewRelic;
 
@@ -41,10 +40,25 @@ protected Boolean initialValue() {
         }
     };
 
+    //statement cache settings
+    /**
+     * In extreme high-throughput scenarios, weak keyed eviction may throttle the CPU when it tries to perform maintenance on statement caches.
+     * Setting -Dnewrelic.config.jdbc_statement_weak_key_caching.enabled=false will convert the caches to ordinary ConcurrentHashMaps
+     * to alleviate this maintenance overhead.
+     * <p>
+     * Statements are only removed from the caches on Statement.close(), so this config MUST be enabled with caution. If the user does not properly close their
+     * statements, setting -Dnewrelic.config.jdbc_statement_weak_key_caching.enabled=false will cause memory issues. The user MUST explicitly call
+     * Statement.close() (not the implicit Statement.closeOnCompletion()) for cleanup to be guaranteed.
+     */
+    private static final String JDBC_STATEMENT_WEAK_KEY_CACHING_ENABLED = "jdbc_statement_weak_key_caching.enabled";
+    private static final boolean JDBC_STATEMENT_WEAK_KEY_CACHING_ENABLED_DEFAULT = Boolean.TRUE;
+    private static final int INITIAL_STATEMENT_CACHE_CAPACITY = 1024;
+    private static final Map<Statement, Object[]> statementToParams = initStatementCache("statementToParams");
+    private static final Map<Statement, String> statementToSql = initStatementCache("statementToSql");
+
     // This will contain every vendor type that we detected on the client system
     private static final Map<String, DatabaseVendor> typeToVendorLookup = new ConcurrentHashMap<>(10);
     private static final Map<Class<?>, DatabaseVendor> classToVendorLookup = AgentBridge.collectionFactory.createConcurrentWeakKeyedMap();
-    private static final Map<Statement, String> statementToSql = AgentBridge.collectionFactory.createConcurrentWeakKeyedMap();
     private static final Map<Connection, String> connectionToIdentifier = AgentBridge.collectionFactory.createConcurrentWeakKeyedMap();
     private static final Map<Connection, String> connectionToURL = AgentBridge.collectionFactory.createConcurrentWeakKeyedMap();
     public static final String UNKNOWN = "unknown";
@@ -58,6 +72,31 @@ protected Boolean initialValue() {
     private static volatile Boolean isSqlMetadataCommentsEnabled = null;
     private static volatile String cachedServiceGuid = null;
 
+    public static Object[] getParams(Statement statement) {
+        return statementToParams.get(statement);
+    }
+
+    public static void putParams(Statement statement, Object[] params) {
+        statementToParams.put(statement, params);
+    }
+
+    public static String getSql(Statement statement) {
+        AgentBridge.getAgent().getLogger().log(Level.FINEST, "Getting sql for statement: {0}", statement);
+        return statementToSql.get(statement);
+    }
+
+    public static void putSql(Statement statement, String sql) {
+        AgentBridge.getAgent().getLogger().log(Level.FINEST, "Storing sql for statement: {0}", statement);
+        statementToSql.put(statement, sql);
+    }
+
+    public static void removeStatement(Statement statement) {
+        if (statement != null) {
+            statementToParams.remove(statement);
+            statementToSql.remove(statement);
+        }
+    }
+
     public static void putVendor(Class<?> driverOrDatastoreClass, DatabaseVendor databaseVendor) {
         classToVendorLookup.put(driverOrDatastoreClass, databaseVendor);
         AgentBridge.getAgent().getLogger().log(Level.FINEST, "Storing class: {0}, vendor: {1}", driverOrDatastoreClass, databaseVendor);
@@ -68,7 +107,7 @@ public static void putVendor(Class<?> driverOrDatastoreClass, DatabaseVendor dat
 
     public static DatabaseVendor getVendor(Class<?> driverOrDatastoreClass, String url) {
         DatabaseVendor vendor = classToVendorLookup.get(driverOrDatastoreClass);
-        AgentBridge.getAgent().getLogger().log(Level.FINEST,"Getting class: {0}, url: {1}, vendor: {2}", driverOrDatastoreClass, url, vendor);
+        AgentBridge.getAgent().getLogger().log(Level.FINEST, "Getting class: {0}, url: {1}, vendor: {2}", driverOrDatastoreClass, url, vendor);
 
         if (vendor != null) {
             return vendor;
@@ -150,16 +189,6 @@ public static String getCachedDatabaseName(Connection connection) {
         return null;
     }
 
-    public static void putSql(Statement statement, String sql) {
-        AgentBridge.getAgent().getLogger().log(Level.FINEST, "Storing sql for statement: {0}", statement);
-        statementToSql.put(statement, sql);
-    }
-
-    public static String getSql(Statement statement) {
-        AgentBridge.getAgent().getLogger().log(Level.FINEST, "Getting sql for statement: {0}", statement);
-        return statementToSql.get(statement);
-    }
-
     public static Object[] growParameterArray(Object[] params, int missingIndex) {
         int length = Math.max(10, (int) (missingIndex * 1.2));
         Object[] newParams = new Object[length];
@@ -224,7 +253,7 @@ public static String getConnectionURL(Connection connection) {
      *
      * @param connectionString URL connection string to parse.
      * @return identifier parsed from connection string if vendor is part of supported in-memory JDBC drivers,
-     *         {@link JdbcHelper#UNKNOWN} otherwise.
+     * {@link JdbcHelper#UNKNOWN} otherwise.
      */
     public static String parseInMemoryIdentifier(String connectionString) {
         if (connectionString == null) {
@@ -257,8 +286,6 @@ public static String parseInMemoryIdentifier(String connectionString) {
         return UNKNOWN;
     }
 
-
-
     /**
      * Parse and cache identifier of in-memory database from Connection string url.
      *
@@ -325,7 +352,6 @@ public static void invalidateMetadataCommentConfig() {
      * the original SQL statement.
      *
      * @param sql the target SQL statement
-     *
      * @return a SQL statement which might have the metadata comment prepended to it
      */
     public static String addSqlMetadataCommentIfNeeded(String sql) {
@@ -412,6 +438,29 @@ private static Boolean isSqlMetadataCommentsEnabled() {
         return (isSqlMetadataCommentsEnabled != null) ? isSqlMetadataCommentsEnabled : Boolean.FALSE;
     }
 
+    /**
+     * Checks the config to see whether weak key caching is enabled.
+     * <p>
+     * If weak key caching is enabled (default) a Caffeine-backed Weak Keyed cache will be returned.
+     * If weak key caching is not enabled a vanilla Concurrent Hash Map will be returned.
+     */
+    static <V> Map<Statement, V> initStatementCache(String cacheName) {
+        boolean weakKeyCachingEnabled = NewRelic.getAgent()
+                .getConfig()
+                .getValue(JDBC_STATEMENT_WEAK_KEY_CACHING_ENABLED, JDBC_STATEMENT_WEAK_KEY_CACHING_ENABLED_DEFAULT);
+        if (weakKeyCachingEnabled) {
+            NewRelic.getAgent().getLogger().log(Level.INFO, "JDBC Statement weak key caching is enabled. Using default Weak Keyed Cache for {0}.", cacheName);
+            return AgentBridge.collectionFactory.createConcurrentWeakKeyedMap();
+        } else {
+            NewRelic.getAgent()
+                    .getLogger()
+                    .log(Level.INFO,
+                            "JDBC Statement weak key caching is disabled. Using a ConcurrentHashMap for {0}. All JDBC Statements MUST be closed when this setting is enabled.",
+                            cacheName);
+            return AgentBridge.collectionFactory.createVanillaJavaConcurrentHashMap(INITIAL_STATEMENT_CACHE_CAPACITY);
+        }
+    }
+
     /**
      * Retrieves the cached entity GUID value and initialize the value on first access.
      * If the entity GUID is not yet set (empty string), this method will retry on subsequent
diff --git a/agent-bridge-datastore/src/test/java/com/newrelic/agent/bridge/datastore/JdbcHelperTest.java b/agent-bridge-datastore/src/test/java/com/newrelic/agent/bridge/datastore/JdbcHelperTest.java
@@ -9,21 +9,27 @@
 
 import com.newrelic.agent.bridge.Agent;
 import com.newrelic.agent.bridge.AgentBridge;
-import org.junit.After;
 import org.junit.Assert;
-import org.junit.Before;
 import org.junit.Test;
+import com.newrelic.api.agent.NewRelic;
+import org.mockito.Answers;
+import org.mockito.MockedStatic;
 import org.mockito.Mockito;
 
 import java.sql.Connection;
 import java.sql.DatabaseMetaData;
 import java.sql.Driver;
 import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mockStatic;
 
 public class JdbcHelperTest {
 
@@ -313,4 +319,68 @@ public void testEntityGuidCaching_handlesNull() {
             JdbcHelper.resetEntityGuidCache();
         }
     }
+
+    @Test
+    public void testInitStatementCachesReturnsDefaultCacheType() {
+        Class<?> expectedDefaultCacheType = AgentBridge.collectionFactory.createConcurrentWeakKeyedMap().getClass();
+        Class<?> actualDefaultCacheType = JdbcHelper.initStatementCache("test_cache").getClass();
+        assertEquals("Default statement cache type should be a weak keyed cache", expectedDefaultCacheType, actualDefaultCacheType);
+    }
+
+    @Test
+    public void testInitStatementCacheWeakKeysDisabledReturnsVanillaCHM() {
+        try (MockedStatic<NewRelic> nRMockedStatic = mockStatic(NewRelic.class, Answers.RETURNS_DEEP_STUBS)) {
+            addDefaultConfigValToAgentMock(nRMockedStatic);
+            nRMockedStatic.when(() -> NewRelic.getAgent().getConfig().getValue(
+                    eq("jdbc_statement_weak_key_caching.enabled"), any(Boolean.class)))
+                    .thenReturn(false);
+
+            Map<Statement, String> cache = JdbcHelper.initStatementCache("test_cache_disabled");
+
+            Class<?> expectedCacheType = AgentBridge.collectionFactory.createVanillaJavaConcurrentHashMap(1024).getClass();
+            assertEquals("When weak key caching is disabled, should return a ConcurrentHashMap",
+                    expectedCacheType, cache.getClass());
+        }
+    }
+
+    @Test
+    public void testInitStatementCacheWeakKeysEnabledReturnsDefaultCacheType() {
+        try (MockedStatic<NewRelic> nRMockedStatic = mockStatic(NewRelic.class, Answers.RETURNS_DEEP_STUBS)) {
+            addDefaultConfigValToAgentMock(nRMockedStatic);
+            nRMockedStatic.when(() -> NewRelic.getAgent().getConfig().getValue(
+                    eq("jdbc_statement_weak_key_caching.enabled"), any(Boolean.class)))
+                    .thenReturn(true);
+
+            Map<Statement, String> cache = JdbcHelper.initStatementCache("test_cache_enabled");
+
+            Class<?> expectedCacheType = AgentBridge.collectionFactory.createConcurrentWeakKeyedMap().getClass();
+            assertEquals("When weak key caching is enabled, should return a weak keyed cache",
+                    expectedCacheType, cache.getClass());
+        }
+    }
+
+    @Test
+    public void testRemoveStatementClearsBothCaches() throws SQLException {
+        Statement st = Mockito.mock(Statement.class);
+        String fakeSql = "SELECT * FROM foo";
+        Object[] fakeParams = {"bar", 6};
+
+        //Add the st to both caches and make sure it's in there.
+        JdbcHelper.putSql(st, fakeSql);
+        JdbcHelper.putParams(st, fakeParams);
+        assertEquals(fakeSql, JdbcHelper.getSql(st));
+        assertEquals(fakeParams, JdbcHelper.getParams(st));
+
+        //Remove should remove from both caches.
+        JdbcHelper.removeStatement(st);
+        assertNull(JdbcHelper.getSql(st));
+        assertNull(JdbcHelper.getParams(st));
+    }
+
+    //This config setting is read by JdbcHelper when it is initialized, so it needs to be added to the static mock.
+    private void addDefaultConfigValToAgentMock(MockedStatic<NewRelic> nRMockedStatic){
+        nRMockedStatic.when(() -> NewRelic.getAgent().getConfig().getValue(
+                eq("jdbc_helper_cache_expire_time"), any()))
+                .thenReturn(7200);
+    }
 }
diff --git a/instrumentation/jdbc-generic/src/main/java/java/sql/Connection_Weaved.java b/instrumentation/jdbc-generic/src/main/java/java/sql/Connection_Weaved.java
@@ -18,66 +18,66 @@ public abstract class Connection_Weaved {
     public PreparedStatement_Weaved prepareStatement(String sql) throws SQLException {
         sql = JdbcHelper.addSqlMetadataCommentIfNeeded(sql);
         PreparedStatement_Weaved preparedStatement = Weaver.callOriginal();
-        preparedStatement.preparedSql = sql;
+        JdbcHelper.putSql((Statement) preparedStatement, sql);
         return preparedStatement;
     }
 
     public CallableStatement prepareCall(String sql) throws SQLException {
         sql = JdbcHelper.addSqlMetadataCommentIfNeeded(sql);
         CallableStatement callableStatement = Weaver.callOriginal();
-        ((PreparedStatement_Weaved) callableStatement).preparedSql = sql;
+        JdbcHelper.putSql(callableStatement, sql);
         return callableStatement;
     }
 
     public PreparedStatement_Weaved prepareStatement(String sql, int resultSetType, int resultSetConcurrency)
             throws SQLException {
         sql = JdbcHelper.addSqlMetadataCommentIfNeeded(sql);
         PreparedStatement_Weaved preparedStatement = Weaver.callOriginal();
-        preparedStatement.preparedSql = sql;
+        JdbcHelper.putSql((Statement) preparedStatement, sql);
         return preparedStatement;
     }
 
     public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
         sql = JdbcHelper.addSqlMetadataCommentIfNeeded(sql);
         CallableStatement callableStatement = Weaver.callOriginal();
-        ((PreparedStatement_Weaved) callableStatement).preparedSql = sql;
+        JdbcHelper.putSql(callableStatement, sql);
         return callableStatement;
     }
 
     public PreparedStatement_Weaved prepareStatement(String sql, int resultSetType, int resultSetConcurrency,
             int resultSetHoldability) throws SQLException {
         sql = JdbcHelper.addSqlMetadataCommentIfNeeded(sql);
         PreparedStatement_Weaved preparedStatement = Weaver.callOriginal();
-        preparedStatement.preparedSql = sql;
+        JdbcHelper.putSql((Statement) preparedStatement, sql);
         return preparedStatement;
     }
 
     public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency,
             int resultSetHoldability) throws SQLException {
         sql = JdbcHelper.addSqlMetadataCommentIfNeeded(sql);
         CallableStatement callableStatement = Weaver.callOriginal();
-        ((PreparedStatement_Weaved) callableStatement).preparedSql = sql;
+        JdbcHelper.putSql(callableStatement, sql);
         return callableStatement;
     }
 
     public PreparedStatement_Weaved prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
         sql = JdbcHelper.addSqlMetadataCommentIfNeeded(sql);
         PreparedStatement_Weaved preparedStatement = Weaver.callOriginal();
-        preparedStatement.preparedSql = sql;
+        JdbcHelper.putSql((Statement) preparedStatement, sql);
         return preparedStatement;
     }
 
     public PreparedStatement_Weaved prepareStatement(String sql, int[] columnIndexes) throws SQLException {
         sql = JdbcHelper.addSqlMetadataCommentIfNeeded(sql);
         PreparedStatement_Weaved preparedStatement = Weaver.callOriginal();
-        preparedStatement.preparedSql = sql;
+        JdbcHelper.putSql((Statement) preparedStatement, sql);
         return preparedStatement;
     }
 
     public PreparedStatement_Weaved prepareStatement(String sql, String[] columnNames) throws SQLException {
         sql = JdbcHelper.addSqlMetadataCommentIfNeeded(sql);
         PreparedStatement_Weaved preparedStatement = Weaver.callOriginal();
-        preparedStatement.preparedSql = sql;
+        JdbcHelper.putSql((Statement) preparedStatement, sql);
         return preparedStatement;
     }
 }
diff --git a/instrumentation/jdbc-generic/src/main/java/java/sql/PreparedStatement_Weaved.java b/instrumentation/jdbc-generic/src/main/java/java/sql/PreparedStatement_Weaved.java
diff --git a/instrumentation/jdbc-generic/src/main/java/java/sql/Statement_Weaved.java b/instrumentation/jdbc-generic/src/main/java/java/sql/Statement_Weaved.java

Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@ dependencies {`
`40`	`40`	`implementation(project(":agent-bridge"))`
`41`	`41`
`42`	`42`	`testImplementation('org.awaitility:awaitility:4.2.0')`
	`43`	`+ testImplementation("org.mockito:mockito-inline:4.11.0")`
`43`	`44`	`testImplementation(fileTree(dir: 'src/test/resources/com/newrelic/agent/bridge', include: '*.jar'))`
`44`	`45`	`}`
`45`	`46`