test/distributed/org/apache/cassandra/distributed/test/ClusterStorageUsageTest.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.cassandra.distributed.test;

import org.junit.Test;

import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.Keyspace;
import org.apache.cassandra.distributed.api.IInvokableInstance;
import org.apache.cassandra.distributed.api.IIsolatedExecutor.SerializableFunction;
import org.apache.cassandra.metrics.DefaultNameFactory;

import org.apache.cassandra.distributed.Cluster;
import org.apache.cassandra.distributed.api.ConsistencyLevel;

import static org.assertj.core.api.Assertions.assertThat;

/**
 * This class verifies the behavior of our global disk usage metrics across different cluster and replication
 * configurations. In addition, it verifies that they are properly exposed via the public metric registry.
 *
 * Disk usage metrics are characterized by how they handle compression and replication:
 *
 * "compressed" -> indicates raw disk usage
 * "uncompressed" -> indicates uncompressed file size (which is equivalent to "compressed" with no compression enabled)
 * "replicated" -> includes disk usage for data outside the node's primary range
 * "unreplicated" -> indicates disk usage scaled down by replication factor across the entire cluster
 */
public class ClusterStorageUsageTest extends TestBaseImpl
{
    private static final DefaultNameFactory FACTORY = new DefaultNameFactory("Storage");
    private static final int MUTATIONS = 1000;

    @Test
    public void testNoReplication() throws Throwable
    {
        // With a replication factor of 1 for our only user keyspace, system keyspaces using local replication, and
        // empty distributed system tables, replicated and unreplicated versions of our compressed and uncompressed
        // metrics should be equivalent.

        try (Cluster cluster = init(builder().withNodes(2).start(), 1))
        {
            populateUserKeyspace(cluster);
            verifyLoadMetricsWithoutReplication(cluster.get(1));
            verifyLoadMetricsWithoutReplication(cluster.get(2));
        }
    }

    private void verifyLoadMetricsWithoutReplication(IInvokableInstance node)
    {
        long compressedLoad = getLoad(node);
        long uncompressedLoad = getUncompressedLoad(node);
        assertThat(compressedLoad).isEqualTo(getUnreplicatedLoad(node));
        assertThat(uncompressedLoad).isEqualTo(getUnreplicatedUncompressedLoad(node));
        assertThat(uncompressedLoad).isGreaterThan(compressedLoad);
    }

    @Test
    public void testSimpleReplication() throws Throwable
    {
        // With a replication factor of 2 for our only user keyspace, disk space used by that keyspace should
        // be scaled down by a factor of 2, while contributions from system keyspaces are unaffected.

        try (Cluster cluster = init(builder().withNodes(3).start(), 2))
        {
            populateUserKeyspace(cluster);

            verifyLoadMetricsWithReplication(cluster.get(1));
            verifyLoadMetricsWithReplication(cluster.get(2));
            verifyLoadMetricsWithReplication(cluster.get(3));
        }
    }

    @Test
    public void testMultiDatacenterReplication() throws Throwable
    {
        // With a replication factor of 1 for our only user keyspace in two DCs, disk space used by that keyspace should
        // be scaled down by a factor of 2, while contributions from system keyspaces are unaffected.

        try (Cluster cluster = builder().withDC("DC1", 2).withDC("DC2", 2).start())
        {
            cluster.schemaChange("CREATE KEYSPACE " + KEYSPACE + " WITH replication = {'class': 'NetworkTopologyStrategy', 'DC1': 1, 'DC2': 1};");
            populateUserKeyspace(cluster);

            verifyLoadMetricsWithReplication(cluster.get(1));
            verifyLoadMetricsWithReplication(cluster.get(2));
            verifyLoadMetricsWithReplication(cluster.get(3));
            verifyLoadMetricsWithReplication(cluster.get(4));
        }
    }

    private void populateUserKeyspace(Cluster cluster)
    {
        cluster.schemaChange(withKeyspace("CREATE TABLE %s.tbl (pk int, v text, PRIMARY KEY (pk));"));

        for (int i = 0; i < MUTATIONS; i++) {
            cluster.coordinator(1).execute(withKeyspace("INSERT INTO %s.tbl (pk, v) VALUES (?,?)"), ConsistencyLevel.ALL, i, "compressable");
        }

        cluster.forEach((i) -> i.flush(KEYSPACE));
    }

    private void verifyLoadMetricsWithReplication(IInvokableInstance node)
    {
        long unreplicatedLoad = getUnreplicatedLoad(node);
        long expectedUnreplicatedLoad = computeUnreplicatedMetric(node, table -> table.metric.liveDiskSpaceUsed.getCount());
        assertThat(expectedUnreplicatedLoad).isEqualTo(unreplicatedLoad);
        assertThat(getLoad(node)).isGreaterThan(unreplicatedLoad);

        long unreplicatedUncompressedLoad = getUnreplicatedUncompressedLoad(node);
        long expectedUnreplicatedUncompressedLoad = computeUnreplicatedMetric(node, table -> table.metric.uncompressedLiveDiskSpaceUsed.getCount());
        assertThat(expectedUnreplicatedUncompressedLoad).isEqualTo(unreplicatedUncompressedLoad);
        assertThat(getUncompressedLoad(node)).isGreaterThan(unreplicatedUncompressedLoad);
    }

    private long getLoad(IInvokableInstance node)
    {
        return node.metrics().getCounter(FACTORY.createMetricName("Load").getMetricName());
    }

    private long getUncompressedLoad(IInvokableInstance node1)
    {
        return node1.metrics().getCounter(FACTORY.createMetricName("UncompressedLoad").getMetricName());
    }

    private long getUnreplicatedLoad(IInvokableInstance node)
    {
        return (Long) node.metrics().getGauge(FACTORY.createMetricName("UnreplicatedLoad").getMetricName());
    }

    private long getUnreplicatedUncompressedLoad(IInvokableInstance node)
    {
        return (Long) node.metrics().getGauge(FACTORY.createMetricName("UnreplicatedUncompressedLoad").getMetricName());
    }

    private long computeUnreplicatedMetric(IInvokableInstance node, SerializableFunction<ColumnFamilyStore, Long> metric)
    {
        return node.callOnInstance(() ->
                                   {
                                       long sum = 0;

                                       for (Keyspace keyspace : Keyspace.all())
                                           for (ColumnFamilyStore table : keyspace.getColumnFamilyStores())
                                               sum += metric.apply(table) / keyspace.getReplicationStrategy().getReplicationFactor().fullReplicas;

                                       return sum;
                                   });
    }
}