test/distributed/org/apache/cassandra/distributed/test/thresholds/AbstractClientSizeWarning.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.thresholds;

import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Random;
import java.util.function.Consumer;

import com.google.common.collect.ImmutableSet;
import org.junit.Assert;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;

import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.SimpleStatement;
import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.cql3.QueryProcessor;
import org.apache.cassandra.distributed.Cluster;
import org.apache.cassandra.distributed.api.ConsistencyLevel;
import org.apache.cassandra.distributed.api.Feature;
import org.apache.cassandra.distributed.api.ICluster;
import org.apache.cassandra.distributed.api.ICoordinator;
import org.apache.cassandra.distributed.api.IInvokableInstance;
import org.apache.cassandra.distributed.api.SimpleQueryResult;
import org.apache.cassandra.distributed.test.JavaDriverUtils;
import org.apache.cassandra.distributed.test.TestBaseImpl;
import org.apache.cassandra.exceptions.ReadSizeAbortException;
import org.apache.cassandra.exceptions.RequestFailureReason;
import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.QueryState;
import org.apache.cassandra.service.reads.thresholds.CoordinatorWarnings;
import org.assertj.core.api.Condition;

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

public abstract class AbstractClientSizeWarning extends TestBaseImpl
{
    private static final String CQL_PK_READ = "SELECT * FROM " + KEYSPACE + ".tbl WHERE pk=1";
    private static final String CQL_TABLE_SCAN = "SELECT * FROM " + KEYSPACE + ".tbl";

    private static final Random RANDOM = new Random(0);
    protected static ICluster<IInvokableInstance> CLUSTER;
    protected static com.datastax.driver.core.Cluster JAVA_DRIVER;
    protected static com.datastax.driver.core.Session JAVA_DRIVER_SESSION;

    @BeforeClass
    public static void setupClass() throws IOException
    {
        Cluster.Builder builder = Cluster.build(3);
        builder.withConfig(c -> c.with(Feature.NATIVE_PROTOCOL, Feature.GOSSIP));
        CLUSTER = builder.start();
        JAVA_DRIVER = JavaDriverUtils.create(CLUSTER);
        JAVA_DRIVER_SESSION = JAVA_DRIVER.connect();
    }

    protected abstract long totalWarnings();
    protected abstract long totalAborts();
    protected abstract void assertWarnings(List<String> warnings);
    protected abstract void assertAbortWarnings(List<String> warnings);
    protected boolean shouldFlush()
    {
        return false;
    }

    @Before
    public void setup()
    {
        CLUSTER.schemaChange("DROP KEYSPACE IF EXISTS " + KEYSPACE);
        init(CLUSTER);
        // disable key cache so RowIndexEntry is read each time
        CLUSTER.schemaChange("CREATE TABLE " + KEYSPACE + ".tbl (pk int, ck int, v blob, PRIMARY KEY (pk, ck)) WITH caching = { 'keys' : 'NONE'}");
    }

    @Test
    public void noWarningsSinglePartition()
    {
        noWarnings(CQL_PK_READ);
    }

    @Test
    public void noWarningsScan()
    {
        noWarnings(CQL_TABLE_SCAN);
    }

    public void noWarnings(String cql)
    {
        CLUSTER.coordinator(1).execute("INSERT INTO " + KEYSPACE + ".tbl (pk, ck, v) VALUES (1, 1, ?)", ConsistencyLevel.ALL, bytes(128));
        CLUSTER.coordinator(1).execute("INSERT INTO " + KEYSPACE + ".tbl (pk, ck, v) VALUES (1, 2, ?)", ConsistencyLevel.ALL, bytes(128));
        if (shouldFlush())
            CLUSTER.stream().forEach(i -> i.flush(KEYSPACE));

        Consumer<List<String>> test = warnings ->
                                      Assert.assertEquals(Collections.emptyList(), warnings);

        for (boolean b : Arrays.asList(true, false))
        {
            enable(b);
            checkpointHistogram();
            SimpleQueryResult result = CLUSTER.coordinator(1).executeWithResult(cql, ConsistencyLevel.ALL);
            test.accept(result.warnings());
            if (b)
            {
                assertHistogramUpdated();
            }
            else
            {
                assertHistogramNotUpdated();
            }
            test.accept(driverQueryAll(cql).getExecutionInfo().getWarnings());
            if (b)
            {
                assertHistogramUpdated();
            }
            else
            {
                assertHistogramNotUpdated();
            }
            assertWarnAborts(0, 0, 0);
        }
    }

    @Test
    public void warnThresholdSinglePartition()
    {
        warnThreshold(CQL_PK_READ, false);
    }

    @Test
    public void warnThresholdScan()
    {
        warnThreshold(CQL_TABLE_SCAN, false);
    }

    @Test
    public void warnThresholdSinglePartitionWithReadRepair()
    {
        warnThreshold(CQL_PK_READ, true);
    }

    @Test
    public void warnThresholdScanWithReadRepair()
    {
        warnThreshold(CQL_TABLE_SCAN, true);
    }

    protected int warnThresholdRowCount()
    {
        return 2;
    }

    public void warnThreshold(String cql, boolean triggerReadRepair)
    {
        for (int i = 0; i < warnThresholdRowCount(); i++)
        {
            if (triggerReadRepair)
            {
                int finalI = i;
                // cell timestamps will not match (even though the values match) which will trigger a read-repair
                CLUSTER.stream().forEach(node -> node.executeInternal("INSERT INTO " + KEYSPACE + ".tbl (pk, ck, v) VALUES (1, ?, ?)", finalI + 1, bytes(512)));
            }
            else
            {
                CLUSTER.coordinator(1).execute("INSERT INTO " + KEYSPACE + ".tbl (pk, ck, v) VALUES (1, ?, ?)", ConsistencyLevel.ALL, i + 1, bytes(512));
            }
        }

        if (shouldFlush())
            CLUSTER.stream().forEach(i -> i.flush(KEYSPACE));

        enable(true);
        checkpointHistogram();
        SimpleQueryResult result = CLUSTER.coordinator(1).executeWithResult(cql, ConsistencyLevel.ALL);
        assertWarnings(result.warnings());
        assertHistogramUpdated();
        assertWarnAborts(1, 0, 0);
        assertWarnings(driverQueryAll(cql).getExecutionInfo().getWarnings());
        assertHistogramUpdated();
        assertWarnAborts(2, 0, 0);

        enable(false);
        result = CLUSTER.coordinator(1).executeWithResult(cql, ConsistencyLevel.ALL);
        assertThat(result.warnings()).isEmpty();
        assertHistogramNotUpdated();
        assertThat(driverQueryAll(cql).getExecutionInfo().getWarnings()).isEmpty();
        assertHistogramNotUpdated();
        assertWarnAborts(2, 0, 0);
    }

    @Test
    public void failThresholdSinglePartitionTrackingEnabled() throws UnknownHostException
    {
        failThresholdEnabled(CQL_PK_READ);
    }

    @Test
    public void failThresholdSinglePartitionTrackingDisabled() throws UnknownHostException
    {
        failThresholdDisabled(CQL_PK_READ);
    }

    @Test
    public void failThresholdScanTrackingEnabled() throws UnknownHostException
    {
        failThresholdEnabled(CQL_TABLE_SCAN);
    }

    @Test
    public void failThresholdScanTrackingDisabled() throws UnknownHostException
    {
        failThresholdDisabled(CQL_TABLE_SCAN);
    }

    protected int failThresholdRowCount()
    {
        return 5;
    }

    public void failThresholdEnabled(String cql) throws UnknownHostException
    {
        ICoordinator node = CLUSTER.coordinator(1);
        for (int i = 0; i < failThresholdRowCount(); i++)
            node.execute("INSERT INTO " + KEYSPACE + ".tbl (pk, ck, v) VALUES (1, ?, ?)", ConsistencyLevel.ALL, i + 1, bytes(512));

        if (shouldFlush())
            CLUSTER.stream().forEach(i -> i.flush(KEYSPACE));

        enable(true);
        checkpointHistogram();
        List<String> warnings = CLUSTER.get(1).callsOnInstance(() -> {
            ClientWarn.instance.captureWarnings();
            CoordinatorWarnings.init();
            try
            {
                QueryProcessor.execute(cql, org.apache.cassandra.db.ConsistencyLevel.ALL, QueryState.forInternalCalls());
                Assert.fail("Expected query failure");
            }
            catch (ReadSizeAbortException e)
            {
                // expected, client transport returns an error message and includes client warnings
            }
            CoordinatorWarnings.done();
            CoordinatorWarnings.reset();
            return ClientWarn.instance.getWarnings();
        }).call();
        assertAbortWarnings(warnings);
        assertHistogramUpdated();
        assertWarnAborts(0, 1, 1);

        try
        {
            driverQueryAll(cql);
            Assert.fail("Query should have thrown ReadFailureException");
        }
        catch (com.datastax.driver.core.exceptions.ReadFailureException e)
        {
            // without changing the client can't produce a better message...
            // client does NOT include the message sent from the server in the exception; so the message doesn't work
            // well in this case
            assertThat(e.getMessage()).contains("responses were required but only 0 replica responded");
            ImmutableSet<InetAddress> expectedKeys = ImmutableSet.of(InetAddress.getByAddress(new byte[]{ 127, 0, 0, 1 }), InetAddress.getByAddress(new byte[]{ 127, 0, 0, 2 }), InetAddress.getByAddress(new byte[]{ 127, 0, 0, 3 }));
            assertThat(e.getFailuresMap())
            .hasSizeBetween(1, 3)
            // coordinator changes from run to run, so can't assert map as the key is dynamic... so assert the domain of keys and the single value expect
            .containsValue(RequestFailureReason.READ_SIZE.code)
            .hasKeySatisfying(new Condition<InetAddress>() {
                public boolean matches(InetAddress value)
                {
                    return expectedKeys.contains(value);
                }
            });
        }
        assertHistogramUpdated();
        assertWarnAborts(0, 2, 1);
    }

    public void failThresholdDisabled(String cql) throws UnknownHostException
    {
        ICoordinator node = CLUSTER.coordinator(1);
        for (int i = 0; i < failThresholdRowCount(); i++)
            node.execute("INSERT INTO " + KEYSPACE + ".tbl (pk, ck, v) VALUES (1, ?, ?)", ConsistencyLevel.ALL, i + 1, bytes(512));

        if (shouldFlush())
            CLUSTER.stream().forEach(i -> i.flush(KEYSPACE));

        // query should no longer fail
        enable(false);
        checkpointHistogram();
        SimpleQueryResult result = node.executeWithResult(cql, ConsistencyLevel.ALL);
        assertThat(result.warnings()).isEmpty();
        assertHistogramNotUpdated();
        assertThat(driverQueryAll(cql).getExecutionInfo().getWarnings()).isEmpty();
        assertHistogramNotUpdated();
        assertWarnAborts(0, 0, 0);
    }

    protected static void enable(boolean value)
    {
        CLUSTER.stream().forEach(i -> i.runOnInstance(() -> DatabaseDescriptor.setReadThresholdsEnabled(value)));
    }

    protected static ByteBuffer bytes(int size)
    {
        byte[] b = new byte[size];
        RANDOM.nextBytes(b);
        return ByteBuffer.wrap(b);
    }

    protected static ResultSet driverQueryAll(String cql)
    {
        return JAVA_DRIVER_SESSION.execute(new SimpleStatement(cql).setConsistencyLevel(com.datastax.driver.core.ConsistencyLevel.ALL));
    }

    protected abstract long[] getHistogram();
    private static long[] previous = new long[0];
    protected void assertHistogramUpdated()
    {
        long[] latestCount = getHistogram();
        try
        {
            // why notEquals?  timing can cause 1 replica to not process before the failure makes it to the test
            // for this reason it is possible 1 replica was not updated but the others were; by expecting everyone
            // to update the test will become flaky
            assertThat(latestCount).isNotEqualTo(previous);
        }
        finally
        {
            previous = latestCount;
        }
    }

    protected void assertHistogramNotUpdated()
    {
        long[] latestCount = getHistogram();
        try
        {
            assertThat(latestCount).isEqualTo(previous);
        }
        finally
        {
            previous = latestCount;
        }
    }

    private void checkpointHistogram()
    {
        previous = getHistogram();
    }

    private static long GLOBAL_READ_ABORTS = 0;
    protected void assertWarnAborts(int warns, int aborts, int globalAborts)
    {
        assertThat(totalWarnings()).as("warnings").isEqualTo(warns);
        assertThat(totalAborts()).as("aborts").isEqualTo(aborts);
        long expectedGlobalAborts = GLOBAL_READ_ABORTS + globalAborts;
        assertThat(totalReadAborts()).as("global aborts").isEqualTo(expectedGlobalAborts);
        GLOBAL_READ_ABORTS = expectedGlobalAborts;
    }

    protected static long totalReadAborts()
    {
        return CLUSTER.stream().mapToLong(i ->
                                          i.metrics().getCounter("org.apache.cassandra.metrics.ClientRequest.Aborts.Read-ALL")
                                          + i.metrics().getCounter("org.apache.cassandra.metrics.ClientRequest.Aborts.RangeSlice")
        ).sum();
    }
}