Skip to content

Commit de27d6f

Browse files
QuChen88QuChen88
committed
Address some review comments. Fix perf testing program.
1 parent c60c8b4 commit de27d6f

File tree

5 files changed

+140
-58
lines changed

5 files changed

+140
-58
lines changed

‎benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/PgCacheBenchmarks.java‎

Lines changed: 35 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package software.amazon.jdbc.benchmarks;
22

33
import org.openjdk.jmh.annotations.*;
4+
import org.openjdk.jmh.infra.Blackhole;
45
import java.sql.*;
56
import java.util.*;
67
import java.util.concurrent.TimeUnit;
@@ -89,36 +90,54 @@ public void warmUpDataSet() throws SQLException {
8990
}
9091
}
9192

92-
private void validateResultSet(ResultSet rs) throws SQLException {
93+
private void validateResultSet(ResultSet rs, Blackhole b) throws SQLException {
9394
while (rs.next()) {
94-
assert rs.getInt(1) >= 0;
95-
assert rs.getInt(2) >= 0;
96-
assert rs.getString(3) != null;
97-
assert rs.getString(4) != null;
98-
assert rs.getDouble(5) >= 0.0;
99-
assert rs.getDate(6) != null;
100-
assert rs.getTime(7) != null;
101-
assert rs.getTime(8) != null;
102-
assert rs.getTimestamp(9) != null;
103-
assert rs.getTimestamp(10) != null;
104-
assert !rs.wasNull();
95+
b.consume(rs.getInt(1));
96+
b.consume(rs.getInt(2));
97+
b.consume(rs.getString(3));
98+
b.consume(rs.getString(4));
99+
b.consume(rs.getDouble(5));
100+
b.consume(rs.getDate(6));
101+
b.consume(rs.getTime(7));
102+
b.consume(rs.getTime(8));
103+
b.consume(rs.getTimestamp(9));
104+
b.consume(rs.getTimestamp(10));
105+
b.consume(rs.wasNull());
105106
}
106107
}
107108

108109
@Benchmark
109-
public void runBenchmarkPrimaryKeyLookup() throws SQLException {
110+
public void runBenchmarkPrimaryKeyLookupNoCaching(Blackhole b) throws SQLException {
111+
try (Statement stmt = connection.createStatement();
112+
ResultSet rs = stmt.executeQuery("SELECT * FROM test where id = " + counter)) {
113+
validateResultSet(rs, b);
114+
}
115+
counter++;
116+
}
117+
118+
@Benchmark
119+
public void runBenchmarkNonIndexedLookupNoCaching(Blackhole b) throws SQLException {
120+
try (Statement stmt = connection.createStatement();
121+
ResultSet rs = stmt.executeQuery("SELECT * FROM test where int_col = " + counter*10)) {
122+
validateResultSet(rs, b);
123+
}
124+
counter++;
125+
}
126+
127+
@Benchmark
128+
public void runBenchmarkPrimaryKeyLookupWithCaching(Blackhole b) throws SQLException {
110129
try (Statement stmt = connection.createStatement();
111130
ResultSet rs = stmt.executeQuery("/*+ CACHE_PARAM(ttl=172800s) */ SELECT * FROM test where id = " + counter)) {
112-
validateResultSet(rs);
131+
validateResultSet(rs, b);
113132
}
114133
counter++;
115134
}
116135

117136
@Benchmark
118-
public void runBenchmarkNonIndexedLookup() throws SQLException {
137+
public void runBenchmarkNonIndexedLookupWithCaching(Blackhole b) throws SQLException {
119138
try (Statement stmt = connection.createStatement();
120139
ResultSet rs = stmt.executeQuery("/*+ CACHE_PARAM(ttl=172800s) */ SELECT * FROM test where int_col = " + counter*10)) {
121-
validateResultSet(rs);
140+
validateResultSet(rs, b);
122141
}
123142
counter++;
124143
}

‎wrapper/build.gradle.kts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,14 +44,14 @@ dependencies {
4444
optionalImplementation("com.mchange:c3p0:0.11.0")
4545
optionalImplementation("org.apache.httpcomponents:httpclient:4.5.14")
4646
optionalImplementation("com.fasterxml.jackson.core:jackson-databind:2.19.0")
47+
optionalImplementation("org.apache.commons:commons-pool2:2.11.1")
4748
optionalImplementation("org.jsoup:jsoup:1.21.1")
4849
optionalImplementation("com.amazonaws:aws-xray-recorder-sdk-core:2.18.2")
50+
optionalImplementation("io.lettuce:lettuce-core:6.6.0.RELEASE")
4951
optionalImplementation("io.opentelemetry:opentelemetry-api:1.52.0")
5052
optionalImplementation("io.opentelemetry:opentelemetry-sdk:1.52.0")
5153
optionalImplementation("io.opentelemetry:opentelemetry-sdk-metrics:1.52.0")
5254

53-
compileOnly("io.lettuce:lettuce-core:6.6.0.RELEASE")
54-
compileOnly("org.apache.commons:commons-pool2:2.11.1")
5555
compileOnly("org.checkerframework:checker-qual:3.49.5")
5656
compileOnly("com.mysql:mysql-connector-j:9.4.0")
5757
compileOnly("org.postgresql:postgresql:42.7.7")
@@ -110,6 +110,7 @@ dependencies {
110110
testImplementation("de.vandermeer:asciitable:0.3.2")
111111
testImplementation("org.hibernate:hibernate-core:5.6.15.Final") // the latest version compatible with Java 8
112112
testImplementation("jakarta.persistence:jakarta.persistence-api:2.2.3")
113+
testImplementation("com.fasterxml.jackson.datatype:jackson-datatype-jsr310:2.19.2")
113114
}
114115

115116
repositories {

‎wrapper/src/main/java/software/amazon/jdbc/plugin/cache/CacheConnection.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ private void createConnectionPool(boolean isRead) {
116116
String[] hostnameAndPort = serverAddr.split(":");
117117
RedisURI redisUriCluster = RedisURI.Builder.redis(hostnameAndPort[0])
118118
.withPort(Integer.parseInt(hostnameAndPort[1]))
119-
.withSsl(useSSL).withVerifyPeer(false).withLibraryName("aws-jdbc-lettuce").build();
119+
.withSsl(useSSL).withVerifyPeer(false).withLibraryName("aws-sql-jdbc-lettuce").build();
120120

121121
RedisClient client = RedisClient.create(resources, redisUriCluster);
122122
GenericObjectPool<StatefulRedisConnection<byte[], byte[]>> pool = new GenericObjectPool<>(

‎wrapper/src/main/java/software/amazon/jdbc/plugin/cache/DataRemoteCachePlugin.java‎

Lines changed: 35 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,11 @@
2626
import java.util.Properties;
2727
import java.util.Set;
2828
import java.util.logging.Logger;
29+
import software.amazon.jdbc.AwsWrapperProperty;
2930
import software.amazon.jdbc.JdbcCallable;
3031
import software.amazon.jdbc.JdbcMethod;
3132
import software.amazon.jdbc.PluginService;
33+
import software.amazon.jdbc.PropertyDefinition;
3234
import software.amazon.jdbc.plugin.AbstractConnectionPlugin;
3335
import software.amazon.jdbc.states.SessionStateService;
3436
import software.amazon.jdbc.util.Messages;
@@ -41,7 +43,6 @@
4143

4244
public class DataRemoteCachePlugin extends AbstractConnectionPlugin {
4345
private static final Logger LOGGER = Logger.getLogger(DataRemoteCachePlugin.class.getName());
44-
private static final int MAX_CACHEABLE_QUERY_SIZE = 16000;
4546
private static final String QUERY_HINT_START_PATTERN = "/*+";
4647
private static final String QUERY_HINT_END_PATTERN = "*/";
4748
private static final String CACHE_PARAM_PATTERN = "CACHE_PARAM(";
@@ -55,6 +56,7 @@ public class DataRemoteCachePlugin extends AbstractConnectionPlugin {
5556
JdbcMethod.CALLABLESTATEMENT_EXECUTE.methodName,
5657
JdbcMethod.CALLABLESTATEMENT_EXECUTEQUERY.methodName)));
5758

59+
private int maxCacheableQuerySize;
5860
private PluginService pluginService;
5961
private TelemetryFactory telemetryFactory;
6062
private TelemetryCounter cacheHitCounter;
@@ -63,6 +65,17 @@ public class DataRemoteCachePlugin extends AbstractConnectionPlugin {
6365
private TelemetryCounter malformedHintCounter;
6466
private TelemetryCounter cacheBypassCounter;
6567
private CacheConnection cacheConnection;
68+
private String dbUserName;
69+
70+
private static final AwsWrapperProperty CACHE_MAX_QUERY_SIZE =
71+
new AwsWrapperProperty(
72+
"cacheMaxQuerySize",
73+
"16384",
74+
"The max query size for remote caching");
75+
76+
static {
77+
PropertyDefinition.registerPluginProperties(DataRemoteCachePlugin.class);
78+
}
6679

6780
public DataRemoteCachePlugin(final PluginService pluginService, final Properties properties) {
6881
try {
@@ -78,7 +91,9 @@ public DataRemoteCachePlugin(final PluginService pluginService, final Properties
7891
this.totalQueryCounter = telemetryFactory.createCounter("JdbcCacheTotalQueryCount");
7992
this.malformedHintCounter = telemetryFactory.createCounter("JdbcCacheMalformedQueryHint");
8093
this.cacheBypassCounter = telemetryFactory.createCounter("JdbcCacheBypassCount");
94+
this.maxCacheableQuerySize = CACHE_MAX_QUERY_SIZE.getInteger(properties);
8195
this.cacheConnection = new CacheConnection(properties);
96+
this.dbUserName = PropertyDefinition.USER.getString(properties);
8297
}
8398

8499
// Used for unit testing purposes only
@@ -93,25 +108,36 @@ public Set<String> getSubscribedMethods() {
93108

94109
private String getCacheQueryKey(String query) {
95110
// Check some basic session states. The important ones for caching include (but not limited to):
96-
// schema name, username which can affect the query result from the DB in addition to the query string
111+
// schema name, username which can affect the query result from the DB in addition to the query string
97112
try {
98113
Connection currentConn = pluginService.getCurrentConnection();
99114
DatabaseMetaData metadata = currentConn.getMetaData();
100115
// Fetch and record the schema name if the session state doesn't currently have it
101116
SessionStateService sessionStateService = pluginService.getSessionStateService();
117+
String catalog = sessionStateService.getCatalog().orElse(null);
102118
String schema = sessionStateService.getSchema().orElse(null);
103-
if (schema == null) {
119+
if (catalog == null && schema == null) {
104120
// Fetch the current schema name and store it in sessionStateService
121+
catalog = currentConn.getCatalog();
105122
schema = currentConn.getSchema();
106-
sessionStateService.setSchema(schema);
123+
if (catalog != null) sessionStateService.setCatalog(catalog);
124+
if (schema != null) sessionStateService.setSchema(schema);
107125
}
108126

127+
if (dbUserName == null) {
128+
// For MySQL, metadata username is actually <UserName>@<ip>. We just need the part before '@'.
129+
dbUserName = metadata.getUserName();
130+
int nameIndexEnd = dbUserName.indexOf('@');
131+
if (nameIndexEnd > 0) {
132+
dbUserName = dbUserName.substring(0, nameIndexEnd);
133+
}
134+
}
109135
LOGGER.finest("DB driver protocol " + pluginService.getDriverProtocol()
110136
+ ", database product: " + metadata.getDatabaseProductName() + " " + metadata.getDatabaseProductVersion()
111-
+ ", schema: " + schema + ", user: " + metadata.getUserName()
137+
+ ", catalog: " + catalog + ", schema: " + schema + ", user: " + dbUserName
112138
+ ", driver: " + metadata.getDriverName() + " " + metadata.getDriverVersion());
113139
// The cache key contains the schema name, user name, and the query string
114-
String[] words = {schema, metadata.getUserName(), query};
140+
String[] words = {catalog, schema, dbUserName, query};
115141
return String.join("_", words);
116142
} catch (SQLException e) {
117143
LOGGER.warning("Error getting session state: " + e.getMessage());
@@ -239,11 +265,11 @@ public <T, E extends Exception> T execute(
239265
int endOfQueryHint = 0;
240266
Integer configuredQueryTtl = null;
241267
// Queries longer than 16KB is not cacheable
242-
if ((sql.length() < MAX_CACHEABLE_QUERY_SIZE) && sql.startsWith(QUERY_HINT_START_PATTERN)) {
268+
if ((sql.length() < maxCacheableQuerySize) && sql.startsWith(QUERY_HINT_START_PATTERN)) {
243269
endOfQueryHint = sql.indexOf(QUERY_HINT_END_PATTERN);
244270
if (endOfQueryHint > 0) {
245-
configuredQueryTtl = getTtlForQuery(sql.substring(2, endOfQueryHint).trim());
246-
mainQuery = sql.substring(endOfQueryHint + 2).trim();
271+
configuredQueryTtl = getTtlForQuery(sql.substring(QUERY_HINT_START_PATTERN.length(), endOfQueryHint).trim());
272+
mainQuery = sql.substring(endOfQueryHint + QUERY_HINT_END_PATTERN.length()).trim();
247273
}
248274
}
249275

0 commit comments

Comments
 (0)