CASSANDRA-19986 Prepared statements cache computation improvements

patch by Berenguer Blasi; reviewed by Maxwell Guo, Stefan Miklosovic for CASSANDRA-19986

src/java/org/apache/cassandra/cql3/QueryHandler.java

@@ -58,6 +58,9 @@ public interface QueryHandler
 
     public static class Prepared
     {
+        // CASSANDRA-19986 Precomputed size might be available
+        public volatile int pstmntSize = -1;
+
         public final CQLStatement statement;
 
         public final MD5Digest resultMetadataId;

src/java/org/apache/cassandra/cql3/QueryProcessor.java

@@ -93,7 +93,7 @@ public class QueryProcessor implements QueryHandler
         preparedStatements = Caffeine.newBuilder()
                              .executor(MoreExecutors.directExecutor())
                              .maximumWeight(capacityToBytes(DatabaseDescriptor.getPreparedStatementsCacheSizeMB()))
-                             .weigher(QueryProcessor::measure)
+                             .weigher(QueryProcessor::getSizeOfPreparedStatementForCache)
                              .removalListener((key, prepared, cause) -> {
                                  MD5Digest md5Digest = (MD5Digest) key;
                                  if (cause.wasEvicted())
@@ -411,10 +411,14 @@ public class QueryProcessor implements QueryHandler
         CQLStatement statement = raw.prepare(clientState);
         statement.validate(clientState);
 
+        Prepared res;
         if (isInternal)
-            return new Prepared(statement, "", fullyQualified, keyspace);
+            res = new Prepared(statement, "", fullyQualified, keyspace);
         else
-            return new Prepared(statement, query, fullyQualified, keyspace);
+            res = new Prepared(statement, query, fullyQualified, keyspace);
+        res.pstmntSize = measurePstmnt(res);
+
+        return res;
     }
 
     public static UntypedResultSet executeInternal(String query, Object... values)
@@ -674,14 +678,14 @@ public class QueryProcessor implements QueryHandler
     {
         // Concatenate the current keyspace so we don't mix prepared statements between keyspace (#5352).
         // (if the keyspace is null, queryString has to have a fully-qualified keyspace so it's fine.
-        long statementSize = ObjectSizes.measureDeep(prepared.statement);
+        MD5Digest statementId = computeId(queryString, keyspace);
         // don't execute the statement if it's bigger than the allowed threshold
-        if (statementSize > capacityToBytes(DatabaseDescriptor.getPreparedStatementsCacheSizeMB()))
+        if (getSizeOfPreparedStatementForCache(statementId, prepared) > capacityToBytes(DatabaseDescriptor.getPreparedStatementsCacheSizeMB()))
             throw new InvalidRequestException(String.format("Prepared statement of size %d bytes is larger than allowed maximum of %d MB: %s...",
-                                                            statementSize,
+                                                            prepared.pstmntSize,
                                                             DatabaseDescriptor.getPreparedStatementsCacheSizeMB(),
                                                             queryString.substring(0, 200)));
-        MD5Digest statementId = computeId(queryString, keyspace);
+
         Prepared previous = preparedStatements.get(statementId, (ignored_) -> prepared);
         if (previous == prepared)
             SystemKeyspace.writePreparedStatement(keyspace, statementId, queryString);
@@ -797,9 +801,17 @@ public class QueryProcessor implements QueryHandler
         }
     }
 
-    private static int measure(Object key, Prepared value)
+    private static int measurePstmnt(Prepared value)
     {
-        return Ints.checkedCast(ObjectSizes.measureDeep(key) + ObjectSizes.measureDeep(value));
+        return Ints.checkedCast(ObjectSizes.measureDeep(value));
+    }
+
+    private static int getSizeOfPreparedStatementForCache(MD5Digest key, Prepared value)
+    {
+        if (value.pstmntSize < 0)
+            throw new IllegalStateException("Precomputed prepared statement size not available");
+
+        return Ints.checkedCast(key.size() + value.pstmntSize);
     }
 
     /**

src/java/org/apache/cassandra/utils/MD5Digest.java

@@ -113,4 +113,10 @@ public class MD5Digest
     {
         return localMD5Digest.get();
     }
+
+    public int size()
+    {
+        return bytes.length + //bytes
+               4; // int hashCode
+    }
 }

test/unit/org/apache/cassandra/cql3/PstmtPersistenceTest.java

@@ -21,7 +21,9 @@ import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
+import org.apache.cassandra.Util;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
@@ -153,7 +155,7 @@ public class PstmtPersistenceTest extends CQLTester
                     prepareStatement("INSERT INTO %s (key, val) VALUES (?, ?) USING TIMESTAMP " + cnt2, clientState);
 
                 // each new prepared statement should have caused an eviction
-                assertEquals("eviction count didn't increase by the expected number", numberOfEvictedStatements(), 10);
+                Util.spinAssertEquals("eviction count didn't increase by the expected number", 10L, this::numberOfEvictedStatements, 5, TimeUnit.SECONDS);
                 assertEquals("Number of statements in table and in cache don't match", numberOfStatementsInMemory(), numberOfStatementsOnDisk());
 
                 return;