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,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)

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)