tools/stress/src/org/apache/cassandra/stress/Operation.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.stress;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import com.google.common.util.concurrent.RateLimiter;

import org.apache.cassandra.stress.generate.*;
import org.apache.cassandra.stress.settings.SettingsLog;
import org.apache.cassandra.stress.settings.StressSettings;
import org.apache.cassandra.stress.util.JavaDriverClient;
import org.apache.cassandra.stress.util.ThriftClient;
import org.apache.cassandra.stress.util.Timer;
import org.apache.cassandra.thrift.InvalidRequestException;
import org.apache.cassandra.transport.SimpleClient;

public abstract class Operation
{
    public final StressSettings settings;
    public final Timer timer;
    protected final DataSpec spec;

    private final List<PartitionIterator> partitionCache = new ArrayList<>();
    protected List<PartitionIterator> partitions;

    public static final class DataSpec
    {
        public final PartitionGenerator partitionGenerator;
        final SeedManager seedManager;
        final Distribution partitionCount;
        final RatioDistribution useRatio;
        final Integer targetCount;

        public DataSpec(PartitionGenerator partitionGenerator, SeedManager seedManager, Distribution partitionCount, Integer targetCount)
        {
            this(partitionGenerator, seedManager, partitionCount, null, targetCount);
        }
        public DataSpec(PartitionGenerator partitionGenerator, SeedManager seedManager, Distribution partitionCount, RatioDistribution useRatio)
        {
            this(partitionGenerator, seedManager, partitionCount, useRatio, null);
        }
        private DataSpec(PartitionGenerator partitionGenerator, SeedManager seedManager, Distribution partitionCount, RatioDistribution useRatio, Integer targetCount)
        {
            this.partitionGenerator = partitionGenerator;
            this.seedManager = seedManager;
            this.partitionCount = partitionCount;
            this.useRatio = useRatio;
            this.targetCount = targetCount;
        }
    }

    public Operation(Timer timer, StressSettings settings, DataSpec spec)
    {
        this.timer = timer;
        this.settings = settings;
        this.spec = spec;
    }

    public static interface RunOp
    {
        public boolean run() throws Exception;
        public int partitionCount();
        public int rowCount();
    }

    boolean ready(WorkManager permits, RateLimiter rateLimiter)
    {
        int partitionCount = (int) spec.partitionCount.next();
        if (partitionCount <= 0)
            return false;
        partitionCount = permits.takePermits(partitionCount);
        if (partitionCount <= 0)
            return false;

        int i = 0;
        boolean success = true;
        for (; i < partitionCount && success ; i++)
        {
            if (i >= partitionCache.size())
                partitionCache.add(PartitionIterator.get(spec.partitionGenerator, spec.seedManager));

            success = false;
            while (!success)
            {
                Seed seed = spec.seedManager.next(this);
                if (seed == null)
                    break;

                success = reset(seed, partitionCache.get(i));
            }
        }
        partitionCount = i;

        if (rateLimiter != null)
            rateLimiter.acquire(partitionCount);

        partitions = partitionCache.subList(0, partitionCount);
        return !partitions.isEmpty();
    }

    protected boolean reset(Seed seed, PartitionIterator iterator)
    {
        if (spec.useRatio == null)
            return iterator.reset(seed, spec.targetCount, isWrite());
        else
            return iterator.reset(seed, spec.useRatio.next(), isWrite());
    }

    public boolean isWrite()
    {
        return false;
    }

    /**
     * Run operation
     * @param client Cassandra Thrift client connection
     * @throws IOException on any I/O error.
     */
    public abstract void run(ThriftClient client) throws IOException;

    public void run(SimpleClient client) throws IOException {
        throw new UnsupportedOperationException();
    }

    public void run(JavaDriverClient client) throws IOException {
        throw new UnsupportedOperationException();
    }

    public void timeWithRetry(RunOp run) throws IOException
    {
        timer.start();

        boolean success = false;
        String exceptionMessage = null;

        int tries = 0;
        for (; tries < settings.errors.tries; tries++)
        {
            try
            {
                success = run.run();
                break;
            }
            catch (Exception e)
            {
                switch (settings.log.level)
                {
                    case MINIMAL:
                        break;

                    case NORMAL:
                        System.err.println(e);
                        break;

                    case VERBOSE:
                        e.printStackTrace(System.err);
                        break;

                    default:
                        throw new AssertionError();
                }
                exceptionMessage = getExceptionMessage(e);
            }
        }

        timer.stop(run.partitionCount(), run.rowCount(), !success);

        if (!success)
        {
            error(String.format("Operation x%d on key(s) %s: %s%n",
                    tries,
                    key(),
                    (exceptionMessage == null)
                        ? "Data returned was not validated"
                        : "Error executing: " + exceptionMessage));
        }

    }

    private String key()
    {
        List<String> keys = new ArrayList<>();
        for (PartitionIterator partition : partitions)
            keys.add(partition.getKeyAsString());
        return keys.toString();
    }

    protected String getExceptionMessage(Exception e)
    {
        String className = e.getClass().getSimpleName();
        String message = (e instanceof InvalidRequestException) ? ((InvalidRequestException) e).getWhy() : e.getMessage();
        return (message == null) ? "(" + className + ")" : String.format("(%s): %s", className, message);
    }

    protected void error(String message) throws IOException
    {
        if (!settings.errors.ignore)
            throw new IOException(message);
        else if (settings.log.level.compareTo(SettingsLog.Level.MINIMAL) > 0)
            System.err.println(message);
    }

}