chore: use a cache for the statement parser

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractStatementParser.java

@@ -25,6 +25,10 @@
 import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheStats;
+import com.google.common.cache.Weigher;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
@@ -59,6 +63,13 @@ public abstract class AbstractStatementParser {
               Dialect.POSTGRESQL,
               PostgreSQLStatementParser.class);
 
+  @VisibleForTesting
+  static void resetParsers() {
+    synchronized (lock) {
+      INSTANCES.clear();
+    }
+  }
+
   /**
    * Get an instance of {@link AbstractStatementParser} for the specified dialect.
    *
@@ -171,7 +182,7 @@ private static ParsedStatement ddl(Statement statement, String sqlWithoutComment
     private static ParsedStatement query(
         Statement statement, String sqlWithoutComments, QueryOptions defaultQueryOptions) {
       return new ParsedStatement(
-          StatementType.QUERY, statement, sqlWithoutComments, defaultQueryOptions, false);
+          StatementType.QUERY, null, statement, sqlWithoutComments, defaultQueryOptions, false);
     }
 
     private static ParsedStatement update(
@@ -193,7 +204,7 @@ private ParsedStatement(
       this.type = StatementType.CLIENT_SIDE;
       this.clientSideStatement = clientSideStatement;
       this.statement = statement;
-      this.sqlWithoutComments = sqlWithoutComments;
+      this.sqlWithoutComments = Preconditions.checkNotNull(sqlWithoutComments);
       this.returningClause = false;
     }
 
@@ -202,28 +213,48 @@ private ParsedStatement(
         Statement statement,
         String sqlWithoutComments,
         boolean returningClause) {
-      this(type, statement, sqlWithoutComments, null, returningClause);
+      this(type, null, statement, sqlWithoutComments, null, returningClause);
     }
 
     private ParsedStatement(StatementType type, Statement statement, String sqlWithoutComments) {
-      this(type, statement, sqlWithoutComments, null, false);
+      this(type, null, statement, sqlWithoutComments, null, false);
     }
 
     private ParsedStatement(
         StatementType type,
+        ClientSideStatementImpl clientSideStatement,
         Statement statement,
         String sqlWithoutComments,
         QueryOptions defaultQueryOptions,
         boolean returningClause) {
       Preconditions.checkNotNull(type);
-      Preconditions.checkNotNull(statement);
       this.type = type;
-      this.clientSideStatement = null;
-      this.statement = mergeQueryOptions(statement, defaultQueryOptions);
-      this.sqlWithoutComments = sqlWithoutComments;
+      this.clientSideStatement = clientSideStatement;
+      this.statement = statement == null ? null : mergeQueryOptions(statement, defaultQueryOptions);
+      this.sqlWithoutComments = Preconditions.checkNotNull(sqlWithoutComments);
       this.returningClause = returningClause;
     }
 
+    private ParsedStatement copy(Statement statement, QueryOptions defaultQueryOptions) {
+      return new ParsedStatement(
+          this.type,
+          this.clientSideStatement,
+          statement,
+          this.sqlWithoutComments,
+          defaultQueryOptions,
+          this.returningClause);
+    }
+
+    private ParsedStatement forCache() {
+      return new ParsedStatement(
+          this.type,
+          this.clientSideStatement,
+          null,
+          this.sqlWithoutComments,
+          null,
+          this.returningClause);
+    }
+
     @Override
     public int hashCode() {
       return Objects.hash(
@@ -361,8 +392,58 @@ ClientSideStatement getClientSideStatement() {
   static final Set<String> dmlStatements = ImmutableSet.of("INSERT", "UPDATE", "DELETE");
   private final Set<ClientSideStatementImpl> statements;
 
+  /** The default maximum size of the statement cache in Mb. */
+  public static final int DEFAULT_MAX_STATEMENT_CACHE_SIZE_MB = 5;
+
+  private static int getMaxStatementCacheSize() {
+    String stringValue = System.getProperty("spanner.statement_cache_size_mb");
+    if (stringValue == null) {
+      return DEFAULT_MAX_STATEMENT_CACHE_SIZE_MB;
+    }
+    int value = 0;
+    try {
+      value = Integer.parseInt(stringValue);
+    } catch (NumberFormatException ignore) {
+    }
+    return Math.max(value, 0);
+  }
+
+  private static boolean isRecordStatementCacheStats() {
+    return "true"
+        .equalsIgnoreCase(System.getProperty("spanner.record_statement_cache_stats", "false"));
+  }
+
+  /**
+   * Cache for parsed statements. This prevents statements that are executed multiple times by the
+   * application to be parsed over and over again. The default maximum size is 5Mb.
+   */
+  private final Cache<String, ParsedStatement> statementCache;
+
   AbstractStatementParser(Set<ClientSideStatementImpl> statements) {
     this.statements = Collections.unmodifiableSet(statements);
+    int maxCacheSize = getMaxStatementCacheSize();
+    if (maxCacheSize > 0) {
+      CacheBuilder<String, ParsedStatement> cacheBuilder =
+          CacheBuilder.newBuilder()
+              // Set the max size to (approx) 5MB (by default).
+              .maximumWeight(maxCacheSize * 1024L * 1024L)
+              // We do length*2 because Java uses 2 bytes for each char.
+              .weigher(
+                  (Weigher<String, ParsedStatement>)
+                      (key, value) -> 2 * key.length() + 2 * value.sqlWithoutComments.length())
+              .concurrencyLevel(Runtime.getRuntime().availableProcessors());
+      if (isRecordStatementCacheStats()) {
+        cacheBuilder.recordStats();
+      }
+      this.statementCache = cacheBuilder.build();
+    } else {
+      this.statementCache = null;
+    }
+  }
+
+  @VisibleForTesting
+  CacheStats getStatementCacheStats() {
+    return statementCache == null ? null : statementCache.stats();
   }
 
   @VisibleForTesting
@@ -383,6 +464,20 @@ public ParsedStatement parse(Statement statement) {
   }
 
   ParsedStatement parse(Statement statement, QueryOptions defaultQueryOptions) {
+    if (statementCache == null) {
+      return internalParse(statement, defaultQueryOptions);
+    }
+
+    ParsedStatement parsedStatement = statementCache.getIfPresent(statement.getSql());
+    if (parsedStatement == null) {
+      parsedStatement = internalParse(statement, null);
+      statementCache.put(statement.getSql(), parsedStatement.forCache());
+      return parsedStatement;
+    }
+    return parsedStatement.copy(statement, defaultQueryOptions);
+  }
+
+  private ParsedStatement internalParse(Statement statement, QueryOptions defaultQueryOptions) {
     String sql = removeCommentsAndTrim(statement.getSql());
     ClientSideStatementImpl client = parseClientSideStatement(sql);
     if (client != null) {

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java

@@ -294,6 +294,7 @@ public void stressTest() throws Exception {
             () -> {
               while (!stopMaintenance.get()) {
                 runMaintenanceLoop(clock, pool, 1);
+                Uninterruptibles.sleepUninterruptibly(1L, TimeUnit.MILLISECONDS);
               }
             })
         .start();

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementParserTest.java

@@ -22,6 +22,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeTrue;
@@ -34,6 +35,7 @@
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.ClientSideStatementImpl.CompileException;
 import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
+import com.google.common.cache.CacheStats;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.truth.Truth;
@@ -43,9 +45,12 @@
 import java.util.List;
 import java.util.Scanner;
 import java.util.Set;
+import java.util.UUID;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import org.junit.AfterClass;
 import org.junit.Before;
+import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
@@ -81,6 +86,17 @@ public static Object[] data() {
 
   private AbstractStatementParser parser;
 
+  @BeforeClass
+  public static void enableStatementCacheStats() {
+    AbstractStatementParser.resetParsers();
+    System.setProperty("spanner.record_statement_cache_stats", "true");
+  }
+
+  @AfterClass
+  public static void disableStatementCacheStats() {
+    System.clearProperty("spanner.record_statement_cache_stats");
+  }
+
   @Before
   public void setupParser() {
     parser = AbstractStatementParser.getInstance(dialect);
@@ -1609,6 +1625,63 @@ public void testSkipMultiLineComment() {
         skipMultiLineComment("/* foo /* inner comment */ not in inner comment */ bar", 0));
   }
 
+  @Test
+  public void testStatementCache_NonParameterizedStatement() {
+    CacheStats statsBefore = parser.getStatementCacheStats();
+
+    String sql = "select foo from bar where id=" + UUID.randomUUID();
+    ParsedStatement parsedStatement1 = parser.parse(Statement.of(sql));
+    assertEquals(StatementType.QUERY, parsedStatement1.getType());
+
+    ParsedStatement parsedStatement2 = parser.parse(Statement.of(sql));
+    assertEquals(StatementType.QUERY, parsedStatement2.getType());
+
+    // Even though the parsed statements are cached, the returned instances are not the same.
+    // This makes sure that statements with the same SQL string and different parameter values
+    // can use the cache.
+    assertNotSame(parsedStatement1, parsedStatement2);
+
+    CacheStats statsAfter = parser.getStatementCacheStats();
+    CacheStats stats = statsAfter.minus(statsBefore);
+
+    // The first query had a cache miss. The second a cache hit.
+    assertEquals(1, stats.missCount());
+    assertEquals(1, stats.hitCount());
+  }
+
+  @Test
+  public void testStatementCache_ParameterizedStatement() {
+    CacheStats statsBefore = parser.getStatementCacheStats();
+
+    String sql =
+        "select "
+            + UUID.randomUUID()
+            + " from bar where id="
+            + (dialect == Dialect.POSTGRESQL ? "$1" : "@p1");
+    Statement statement1 = Statement.newBuilder(sql).bind("p1").to(1L).build();
+    Statement statement2 = Statement.newBuilder(sql).bind("p1").to(2L).build();
+
+    ParsedStatement parsedStatement1 = parser.parse(statement1);
+    assertEquals(StatementType.QUERY, parsedStatement1.getType());
+    assertEquals(parsedStatement1.getStatement(), statement1);
+
+    ParsedStatement parsedStatement2 = parser.parse(statement2);
+    assertEquals(StatementType.QUERY, parsedStatement2.getType());
+    assertEquals(parsedStatement2.getStatement(), statement2);
+
+    // Even though the parsed statements are cached, the returned instances are not the same.
+    // This makes sure that statements with the same SQL string and different parameter values
+    // can use the cache.
+    assertNotSame(parsedStatement1, parsedStatement2);
+
+    CacheStats statsAfter = parser.getStatementCacheStats();
+    CacheStats stats = statsAfter.minus(statsBefore);
+
+    // The first query had a cache miss. The second a cache hit.
+    assertEquals(1, stats.missCount());
+    assertEquals(1, stats.hitCount());
+  }
+
   private void assertUnclosedLiteral(String sql) {
     try {
       parser.convertPositionalParametersToNamedParameters('?', sql);