CNDB-15920: Add cassandra.sai.metrics.table_state.enabled flag to be able to enable/disable TableStateMetrics

Author: ekaterinadimitrova2

src/java/org/apache/cassandra/config/CassandraRelevantProperties.java

@@ -635,6 +635,12 @@ public enum CassandraRelevantProperties
      */
     SAI_INDEX_METRICS_ENABLED("cassandra.sai.metrics.index.enabled", "true"),
 
+    /**
+     * Whether to enable SAI table state metrics such as disk usage, queryable index count, and index build progress.
+     * These metrics include gauges for table-level SAI state information.
+     */
+    SAI_TABLE_STATE_METRICS_ENABLED("cassandra.sai.metrics.table_state.enabled", "true"),
+
     /**
      * If true, while creating or altering schema, NetworkTopologyStrategy won't check if the DC exists.
      * This is to remain compatible with older workflows that first change the replication before adding the nodes.

src/java/org/apache/cassandra/index/sai/StorageAttachedIndexGroup.java

@@ -74,6 +74,8 @@
 import org.apache.cassandra.utils.Throwables;
 import org.apache.lucene.index.CorruptIndexException;
 
+import static org.apache.cassandra.config.CassandraRelevantProperties.SAI_TABLE_STATE_METRICS_ENABLED;
+
 /**
  * Orchestrates building of storage-attached indices, and manages lifecycle of resources shared between them.
  */
@@ -85,7 +87,8 @@ public class StorageAttachedIndexGroup implements Index.Group, INotificationCons
     public static final Index.Group.Key GROUP_KEY = new Index.Group.Key(StorageAttachedIndexGroup.class);
 
     private final TableQueryMetrics queryMetrics;
-    private final TableStateMetrics stateMetrics;
+    @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
+    private final Optional<TableStateMetrics> stateMetrics;
     private final IndexGroupMetrics groupMetrics;
 
     private final Set<StorageAttachedIndex> indices = ConcurrentHashMap.newKeySet();
@@ -100,7 +103,9 @@ public class StorageAttachedIndexGroup implements Index.Group, INotificationCons
     {
         this.baseCfs = baseCfs;
         this.queryMetrics = new TableQueryMetrics(baseCfs.metadata());
-        this.stateMetrics = new TableStateMetrics(baseCfs.metadata(), this);
+        this.stateMetrics = SAI_TABLE_STATE_METRICS_ENABLED.getBoolean()
+                            ? Optional.of(new TableStateMetrics(baseCfs.metadata(), this))
+                            : Optional.empty();
         this.groupMetrics = new IndexGroupMetrics(baseCfs.metadata(), this);
         this.contextManager = new SSTableContextManager(baseCfs.getTracker());
         this.version = Version.current(baseCfs.keyspace.getName());
@@ -164,7 +169,7 @@ public void invalidate()
         // in case of removing last index from group,  sstable contexts should already been removed by StorageAttachedIndexGroup#removeIndex
         queryMetrics.release();
         groupMetrics.release();
-        stateMetrics.release();
+        stateMetrics.ifPresent(TableStateMetrics::release);
         baseCfs.getTracker().unsubscribe(this);
     }
 
@@ -176,7 +181,7 @@ public void unload()
         contextManager.clear();
         queryMetrics.release();
         groupMetrics.release();
-        stateMetrics.release();
+        stateMetrics.ifPresent(TableStateMetrics::release);
     }
 
     @Override
@@ -473,6 +478,12 @@ public TableQueryMetrics queryMetrics()
         return queryMetrics;
     }
 
+    // Needed by CNDB
+    public Optional<TableStateMetrics> stateMetrics()
+    {
+        return stateMetrics;
+    }
+
     public ColumnFamilyStore table()
     {
         return baseCfs;

src/java/org/apache/cassandra/index/sai/metrics/TableStateMetrics.java

@@ -28,7 +28,8 @@ public class TableStateMetrics extends AbstractMetrics
 {
     public static final String TABLE_STATE_METRIC_TYPE = "TableStateMetrics";
 
-    private final Gauge diskUsageBytes;
+    // Visible for CNDB
+    public final Gauge diskUsageBytes;
     private final Gauge diskUsagePercentageOfBaseTable;
     private final Gauge totalIndexCount;
     private final Gauge totalIndexBuildsInProgress;

test/unit/org/apache/cassandra/index/sai/metrics/StateMetricsTest.java

@@ -25,6 +25,7 @@
 
 import com.datastax.driver.core.ResultSet;
 
+import static org.apache.cassandra.config.CassandraRelevantProperties.SAI_TABLE_STATE_METRICS_ENABLED;
 import static org.apache.cassandra.index.sai.metrics.TableStateMetrics.TABLE_STATE_METRIC_TYPE;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.Assert.assertEquals;
@@ -49,9 +50,9 @@ public void testMetricRelease()
         createTable(String.format(CREATE_TABLE_TEMPLATE, keyspace, table));
         createIndex(String.format(CREATE_INDEX_TEMPLATE, index, keyspace, table, "v1"));
 
-        execute("INSERT INTO " + keyspace + "." + table + " (id1, v1, v2) VALUES ('0', 0, '0')");
+        execute("INSERT INTO " + keyspace + '.' + table + " (id1, v1, v2) VALUES ('0', 0, '0')");
 
-        ResultSet rows = executeNet("SELECT id1 FROM " + keyspace + "." + table + " WHERE v1 = 0");
+        ResultSet rows = executeNet("SELECT id1 FROM " + keyspace + '.' + table + " WHERE v1 = 0");
         assertEquals(1, rows.all().size());
         assertEquals(1L, getTableStateMetrics(keyspace, table, "TotalIndexCount"));
 
@@ -75,14 +76,14 @@ public void testMetricCreation()
         createIndex(String.format(CREATE_INDEX_TEMPLATE, index+"_v1", keyspace, table, "v1"));
         createIndex(String.format(CREATE_INDEX_TEMPLATE, index+"_v2", keyspace, table, "v2"));
 
-        execute("INSERT INTO " + keyspace + "." + table + " (id1, v1, v2) VALUES ('0', 0, '0')");
-        execute("INSERT INTO " + keyspace + "." + table + " (id1, v1, v2) VALUES ('1', 1, '1')");
-        execute("INSERT INTO " + keyspace + "." + table + " (id1, v1, v2) VALUES ('2', 2, '2')");
-        execute("INSERT INTO " + keyspace + "." + table + " (id1, v1, v2) VALUES ('3', 3, '3')");
+        execute("INSERT INTO " + keyspace + '.' + table + " (id1, v1, v2) VALUES ('0', 0, '0')");
+        execute("INSERT INTO " + keyspace + '.' + table + " (id1, v1, v2) VALUES ('1', 1, '1')");
+        execute("INSERT INTO " + keyspace + '.' + table + " (id1, v1, v2) VALUES ('2', 2, '2')");
+        execute("INSERT INTO " + keyspace + '.' + table + " (id1, v1, v2) VALUES ('3', 3, '3')");
 
         flush(keyspace, table);
 
-        ResultSet rows = executeNet("SELECT id1, v1, v2 FROM " + keyspace + "." + table + " WHERE v1 >= 0");
+        ResultSet rows = executeNet("SELECT id1, v1, v2 FROM " + keyspace + '.' + table + " WHERE v1 >= 0");
 
         int actualRows = rows.all().size();
         assertEquals(4, actualRows);
@@ -98,4 +99,55 @@ private int getTableStateMetrics(String keyspace, String table, String metricsNa
     {
         return (int) getMetricValue(objectNameNoIndex(metricsName, keyspace, table, TABLE_STATE_METRIC_TYPE));
     }
+
+    @Test
+    public void testTableStateMetricsEnabledAndDisabled()
+    {
+        testTableStateMetrics(true);
+        testTableStateMetrics(false);
+    }
+
+    private void testTableStateMetrics(boolean metricsEnabled)
+    {
+        SAI_TABLE_STATE_METRICS_ENABLED.setBoolean(metricsEnabled);
+
+        try
+        {
+            String table = "test_table_state_metrics_" + (metricsEnabled ? "enabled" : "disabled");
+            String index = "test_index_" + (metricsEnabled ? "enabled" : "disabled");
+
+            String keyspace = createKeyspace(CREATE_KEYSPACE_TEMPLATE);
+            createTable(String.format(CREATE_TABLE_TEMPLATE, keyspace, table));
+            createIndex(String.format(CREATE_INDEX_TEMPLATE, index, keyspace, table, "v1"));
+
+            // Test all TableStateMetrics Gauge metrics
+            assertTableStateMetricExistsIfEnabled(metricsEnabled, "TotalIndexCount", keyspace, table);
+            assertTableStateMetricExistsIfEnabled(metricsEnabled, "TotalQueryableIndexCount", keyspace, table);
+            assertTableStateMetricExistsIfEnabled(metricsEnabled, "TotalIndexBuildsInProgress", keyspace, table);
+            assertTableStateMetricExistsIfEnabled(metricsEnabled, "DiskUsedBytes", keyspace, table);
+            assertTableStateMetricExistsIfEnabled(metricsEnabled, "DiskPercentageOfBaseTable", keyspace, table);
+
+            // Test indexing operations to ensure null stateMetrics is handled gracefully
+            execute("INSERT INTO " + keyspace + '.' + table + " (id1, v1, v2) VALUES ('0', 0, '0')");
+            execute("INSERT INTO " + keyspace + '.' + table + " (id1, v1, v2) VALUES ('1', 1, '1')");
+            execute("INSERT INTO " + keyspace + '.' + table + " (id1, v1, v2) VALUES ('2', 2, '2')");
+
+            // Verify metrics still behave correctly after operations
+            assertTableStateMetricExistsIfEnabled(metricsEnabled, "TotalIndexCount", keyspace, table);
+            assertTableStateMetricExistsIfEnabled(metricsEnabled, "DiskUsedBytes", keyspace, table);
+        }
+        finally
+        {
+            // Reset property to default
+            SAI_TABLE_STATE_METRICS_ENABLED.setBoolean(true);
+        }
+    }
+
+    void assertTableStateMetricExistsIfEnabled(boolean shouldExist, String metricName, String keyspace, String table)
+    {
+        if (shouldExist)
+            assertMetricExists(objectNameNoIndex(metricName, keyspace, table, TABLE_STATE_METRIC_TYPE));
+        else
+            assertMetricDoesNotExist(objectNameNoIndex(metricName, keyspace, table, TABLE_STATE_METRIC_TYPE));
+    }
 }