From 0c62d6937c39aa5f08af41774cc7301479e7d56f Mon Sep 17 00:00:00 2001
From: Andy Bradshaw <abradshaw@palantir.com>
Date: Wed, 23 Oct 2024 14:54:01 -0400
Subject: [PATCH] Add mutation cache for schema pull requests

---
 .../cassandra/schema/LegacySchemaTables.java  | 26 ++++++++++++++-----
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/java/org/apache/cassandra/schema/LegacySchemaTables.java b/src/java/org/apache/cassandra/schema/LegacySchemaTables.java
index 13a7544c42..878675d8d7 100644
--- a/src/java/org/apache/cassandra/schema/LegacySchemaTables.java
+++ b/src/java/org/apache/cassandra/schema/LegacySchemaTables.java
@@ -21,11 +21,15 @@
 import java.nio.ByteBuffer;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.time.Duration;
 import java.util.*;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 import com.google.common.base.Function;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.MapDifference;
 import com.google.common.collect.Maps;
@@ -74,6 +78,10 @@ public class LegacySchemaTables
 
     public static final List<String> ALL = Arrays.asList(KEYSPACES, COLUMNFAMILIES, COLUMNS, TRIGGERS, USERTYPES, FUNCTIONS, AGGREGATES);
 
+    private static final int MUTATION_CACHE_MAX_SIZE = 5;
+    private static final Duration MUTATION_CACHE_EXPIRY = Duration.ofMinutes(5);
+    private static final LoadingCache<UUID, Collection<Mutation>> mutations = CacheBuilder.newBuilder().maximumSize(MUTATION_CACHE_MAX_SIZE).expireAfterAccess(MUTATION_CACHE_EXPIRY).build(new UUIDMutationCacheLoader());
+
     private static final CFMetaData Keyspaces =
         compile(KEYSPACES,
                 "keyspace definitions",
@@ -309,12 +317,7 @@ private static List<Row> getSchemaPartitionsForTable(String schemaTableName)
 
     public static Collection<Mutation> convertSchemaToMutations()
     {
-        Map<DecoratedKey, Mutation> mutationMap = new HashMap<>();
-
-        for (String table : ALL)
-            convertSchemaToMutations(mutationMap, table);
-
-        return mutationMap.values();
+        return mutations.getUnchecked(Schema.instance.getVersion());
     }
 
     private static void convertSchemaToMutations(Map<DecoratedKey, Mutation> mutationMap, String schemaTableName)
@@ -1506,4 +1509,15 @@ public static ByteBuffer functionSignatureWithNameAndTypes(AbstractFunction fun)
         return list.decompose(strList);
     }
 
+    private static class UUIDMutationCacheLoader extends CacheLoader<UUID, Collection<Mutation>>
+    {
+        @Override
+        public Collection<Mutation> load(UUID key)
+        {
+            Map<DecoratedKey, Mutation> mutationMap = new HashMap<>();
+            for (String table : ALL)
+                convertSchemaToMutations(mutationMap, table);
+            return mutationMap.values();
+        }
+    }
 }