Make sure replication factor overrides are applied to system_auth / system_distributed / system_traces keyspaces

Author: Eduard Tudenhoefner

management-api-agent/src/main/java/com/datastax/mgmtapi/Agent.java

@@ -5,21 +5,14 @@
  */
 package com.datastax.mgmtapi;
 
-import com.datastax.mgmtapi.interceptors.AuthSchemaInterceptor;
+import com.datastax.mgmtapi.interceptors.SystemDistributedReplicationInterceptor;
 import com.datastax.mgmtapi.interceptors.CassandraDaemonInterceptor;
 import com.datastax.mgmtapi.interceptors.CassandraRoleManagerInterceptor;
 import com.datastax.mgmtapi.interceptors.QueryHandlerInterceptor;
 import net.bytebuddy.agent.builder.AgentBuilder;
-import net.bytebuddy.description.type.TypeDescription;
-import net.bytebuddy.dynamic.ClassFileLocator;
-import net.bytebuddy.dynamic.loading.ClassInjector;
 import org.apache.cassandra.gms.GossiperInterceptor;
 
-import java.io.File;
 import java.lang.instrument.Instrumentation;
-import java.nio.file.Files;
-import java.util.HashMap;
-import java.util.Map;
 
 import static net.bytebuddy.matcher.ElementMatchers.any;
 import static net.bytebuddy.matcher.ElementMatchers.isSynthetic;
@@ -45,8 +38,8 @@ public static void premain(String arg, Instrumentation inst) throws Exception {
                 //Auth Setup
                 .type(CassandraRoleManagerInterceptor.type())
                 .transform(CassandraRoleManagerInterceptor.transformer())
-                .type(AuthSchemaInterceptor.type())
-                .transform(AuthSchemaInterceptor.transformer())
+                .type(SystemDistributedReplicationInterceptor.type())
+                .transform(SystemDistributedReplicationInterceptor.transformer())
                 .installOn(inst);
     }
 }

management-api-agent/src/main/java/com/datastax/mgmtapi/interceptors/AuthSchemaInterceptor.java management-api-agent/src/main/java/com/datastax/mgmtapi/interceptors/SystemDistributedReplicationInterceptor.java

@@ -30,7 +30,9 @@
 import org.apache.cassandra.schema.KeyspaceParams;
 import org.apache.cassandra.schema.ReplicationParams;
 
-public class AuthSchemaInterceptor
+import static net.bytebuddy.matcher.ElementMatchers.nameEndsWith;
+
+public class SystemDistributedReplicationInterceptor
 {
     private static final String SYSTEM_DISTRIBUTED_NTS_DC_OVERRIDE_PROPERTY = "cassandra.system_distributed_replication_dc_names";
     private static final String SYSTEM_DISTRIBUTED_NTS_RF_OVERRIDE_PROPERTY = "cassandra.system_distributed_replication_per_dc";
@@ -53,22 +55,22 @@ public class AuthSchemaInterceptor
         }
         catch (Throwable t)
         {
-            LoggerFactory.getLogger(AuthSchemaInterceptor.class).error("Error parsing system distributed replication override properties", t);
+            LoggerFactory.getLogger(SystemDistributedReplicationInterceptor.class).error("Error parsing system distributed replication override properties", t);
         }
 
         if (rfOverride != null && !dcOverride.isEmpty())
         {
             //Validate reasonable defaults
             if (rfOverride <= 0 || rfOverride > 5)
             {
-                LoggerFactory.getLogger(AuthSchemaInterceptor.class).error("Invalid value for {}", SYSTEM_DISTRIBUTED_NTS_RF_OVERRIDE_PROPERTY);
+                LoggerFactory.getLogger(SystemDistributedReplicationInterceptor.class).error("Invalid value for {}", SYSTEM_DISTRIBUTED_NTS_RF_OVERRIDE_PROPERTY);
             }
             else
             {
                 for (String dc : dcOverride)
                     ntsOverride.put(dc, String.valueOf(rfOverride));
 
-                LoggerFactory.getLogger(AuthSchemaInterceptor.class).info("Using override for distributed system keyspaces: {}", ntsOverride);
+                LoggerFactory.getLogger(SystemDistributedReplicationInterceptor.class).info("Using override for distributed system keyspaces: {}", ntsOverride);
             }
         }
 
@@ -78,7 +80,9 @@ public class AuthSchemaInterceptor
 
     public static ElementMatcher<? super TypeDescription> type()
     {
-        return ElementMatchers.nameEndsWith(".AuthKeyspace");
+        return nameEndsWith(".AuthKeyspace")
+                .or(nameEndsWith(".TraceKeyspace"))
+                .or(nameEndsWith(".SystemDistributedKeyspace"));
     }
 
     public static AgentBuilder.Transformer transformer()
@@ -88,7 +92,7 @@ public static AgentBuilder.Transformer transformer()
             @Override
             public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule javaModule)
             {
-                return builder.method(ElementMatchers.named("metadata")).intercept(MethodDelegation.to(AuthSchemaInterceptor.class));
+                return builder.method(ElementMatchers.named("metadata")).intercept(MethodDelegation.to(SystemDistributedReplicationInterceptor.class));
             }
         };
     }