From 5f7175d59da8b8d20e1081619eadad8b60867e86 Mon Sep 17 00:00:00 2001
From: Jacek Lewandowski <lewandowski.jacek@gmail.com>
Date: Mon, 6 Mar 2023 13:10:28 +0100
Subject: [PATCH] Small fixes around Schema management

- Denormalize distributedAndLocalKeyspaces to remove perf degradation
- Set updated schema version once applied

patch by Jacek Lewandowski, reviewed by Stefan Miklosovic for CASSANDRA-18291
---
 CHANGES.txt                                   |  1 +
 .../org/apache/cassandra/schema/Schema.java   | 20 ++++++++-----------
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/CHANGES.txt b/CHANGES.txt
index 7c66dda36116..ae1a1a46e9fa 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -1,4 +1,5 @@
 4.1.1
+ * Fix too early schema version change in sysem local table (CASSANDRA-18291)
  * Fix copying of JAR of a trigger to temporary file (CASSANDRA-18264)
  * Fix possible NoSuchFileException when removing a snapshot (CASSANDRA-18211)
  * PaxosPrepare may add instances to the Electorate that are not in gossip (CASSANDRA-18194)
diff --git a/src/java/org/apache/cassandra/schema/Schema.java b/src/java/org/apache/cassandra/schema/Schema.java
index dee24c18cf5b..0dba167210d2 100644
--- a/src/java/org/apache/cassandra/schema/Schema.java
+++ b/src/java/org/apache/cassandra/schema/Schema.java
@@ -78,6 +78,7 @@ public class Schema implements SchemaProvider
     public static final Schema instance = new Schema();
 
     private volatile Keyspaces distributedKeyspaces = Keyspaces.none();
+    private volatile Keyspaces distributedAndLocalKeyspaces;
 
     private final Keyspaces localKeyspaces;
 
@@ -109,6 +110,7 @@ private Schema()
         this.localKeyspaces = (CassandraRelevantProperties.FORCE_LOAD_LOCAL_KEYSPACES.getBoolean() || isDaemonInitialized() || isToolInitialized())
                               ? Keyspaces.of(SchemaKeyspace.metadata(), SystemKeyspace.metadata())
                               : Keyspaces.none();
+        this.distributedAndLocalKeyspaces = this.localKeyspaces;
 
         this.localKeyspaces.forEach(this::loadNew);
         this.updateHandler = SchemaUpdateHandlerFactoryProvider.instance.get().getSchemaUpdateHandler(online, this::mergeAndUpdateVersion);
@@ -119,6 +121,7 @@ public Schema(boolean online, Keyspaces localKeyspaces, SchemaUpdateHandler upda
     {
         this.online = online;
         this.localKeyspaces = localKeyspaces;
+        this.distributedAndLocalKeyspaces = this.localKeyspaces;
         this.updateHandler = updateHandler;
     }
 
@@ -160,6 +163,7 @@ private synchronized void load(KeyspaceMetadata ksm)
             reload(previous, ksm);
 
         distributedKeyspaces = distributedKeyspaces.withAddedOrUpdated(ksm);
+        distributedAndLocalKeyspaces = distributedAndLocalKeyspaces.withAddedOrUpdated(ksm);
     }
 
     private synchronized void loadNew(KeyspaceMetadata ksm)
@@ -242,18 +246,9 @@ private Keyspace maybeRemoveKeyspaceInstance(String keyspaceName, Consumer<Keysp
         }
     }
 
-    /**
-     * @deprecated use {@link #distributedAndLocalKeyspaces()}
-     */
-    @Deprecated
-    public Keyspaces snapshot()
-    {
-        return distributedAndLocalKeyspaces();
-    }
-
     public Keyspaces distributedAndLocalKeyspaces()
     {
-        return Keyspaces.builder().add(localKeyspaces).add(distributedKeyspaces).build();
+        return distributedAndLocalKeyspaces;
     }
 
     public Keyspaces distributedKeyspaces()
@@ -282,6 +277,7 @@ public int largestGcgs()
     private synchronized void unload(KeyspaceMetadata ksm)
     {
         distributedKeyspaces = distributedKeyspaces.without(ksm.name);
+        distributedAndLocalKeyspaces = distributedAndLocalKeyspaces.without(ksm.name);
 
         this.tableMetadataRefCache = tableMetadataRefCache.withRemovedRefs(ksm);
 
@@ -591,12 +587,12 @@ public void reloadSchemaAndAnnounceVersion()
     @VisibleForTesting
     public synchronized void mergeAndUpdateVersion(SchemaTransformationResult result, boolean dropData)
     {
-        if (online)
-            SystemKeyspace.updateSchemaVersion(result.after.getVersion());
         result = localDiff(result);
         schemaChangeNotifier.notifyPreChanges(result);
         merge(result.diff, dropData);
         updateVersion(result.after.getVersion());
+        if (online)
+            SystemKeyspace.updateSchemaVersion(result.after.getVersion());
     }
 
     public SchemaTransformationResult transform(SchemaTransformation transformation)