diff --git a/CHANGES.txt b/CHANGES.txt index edf091da22..b83b88abe2 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -69,6 +69,7 @@ * cqlsh: add custom prompt support (CASSANDRA-5539) * Reuse prepared statements in hot auth queries (CASSANDRA-5594) * cqlsh: add vertical output option (see EXPAND) (CASSANDRA-5597) + * Add a rate limit option to stress (CASSANDRA-5004) Merged from 1.1: * Remove buggy thrift max message length option (CASSANDRA-5529) * Fix NPE in Pig's widerow mode (CASSANDRA-5488) diff --git a/tools/stress/src/org/apache/cassandra/stress/Session.java b/tools/stress/src/org/apache/cassandra/stress/Session.java index 7fca986241..f19d61e4a9 100644 --- a/tools/stress/src/org/apache/cassandra/stress/Session.java +++ b/tools/stress/src/org/apache/cassandra/stress/Session.java @@ -23,10 +23,8 @@ import java.net.UnknownHostException; import java.nio.ByteBuffer; import java.util.*; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; import com.yammer.metrics.Metrics; -import com.yammer.metrics.core.Histogram; import org.apache.cassandra.cli.transport.FramedTransportFactory; import org.apache.cassandra.config.CFMetaData; import org.apache.cassandra.config.EncryptionOptions; @@ -117,6 +115,7 @@ public class Session implements Serializable availableOptions.addOption("alg", SSL_ALGORITHM, true, "SSL: algorithm (default: SunX509)"); availableOptions.addOption("st", SSL_STORE_TYPE, true, "SSL: type of store"); availableOptions.addOption("ciphers", SSL_CIPHER_SUITES, true, "SSL: comma-separated list of encryption suites to use"); + availableOptions.addOption("th", "throttle", true, "Throttle the total number of operations per second to a maximum amount."); } private int numKeys = 1000 * 1000; @@ -143,6 +142,7 @@ public class Session implements Serializable private boolean trace = false; private boolean captureStatistics = true; public boolean use_native_protocol = false; + private double maxOpsPerSecond = Double.MAX_VALUE; private final String outFileName; @@ -285,6 +285,9 @@ public class Session implements Serializable if (cmd.hasOption("g")) keysPerCall = Integer.parseInt(cmd.getOptionValue("g")); + if (cmd.hasOption("th")) + maxOpsPerSecond = Double.parseDouble(cmd.getOptionValue("th")); + if (cmd.hasOption("e")) consistencyLevel = ConsistencyLevel.valueOf(cmd.getOptionValue("e").toUpperCase()); @@ -515,6 +518,11 @@ public class Session implements Serializable return threads; } + public double getMaxOpsPerSecond() + { + return maxOpsPerSecond; + } + public float getSkipKeys() { return skipKeys; diff --git a/tools/stress/src/org/apache/cassandra/stress/StressAction.java b/tools/stress/src/org/apache/cassandra/stress/StressAction.java index 13ac61148e..ea84647a77 100644 --- a/tools/stress/src/org/apache/cassandra/stress/StressAction.java +++ b/tools/stress/src/org/apache/cassandra/stress/StressAction.java @@ -23,6 +23,7 @@ import java.util.concurrent.SynchronousQueue; import java.util.concurrent.TimeUnit; import com.google.common.util.concurrent.Uninterruptibles; +import com.google.common.util.concurrent.RateLimiter; import com.yammer.metrics.stats.Snapshot; import org.apache.cassandra.stress.operations.*; import org.apache.cassandra.stress.util.CassandraClient; @@ -69,13 +70,14 @@ public class StressAction extends Thread int itemsPerThread = client.getKeysPerThread(); int modulo = client.getNumKeys() % threadCount; + RateLimiter rateLimiter = RateLimiter.create(client.getMaxOpsPerSecond()); // creating required type of the threads for the test for (int i = 0; i < threadCount; i++) { if (i == threadCount - 1) itemsPerThread += modulo; // last one is going to handle N + modulo items - consumers[i] = new Consumer(itemsPerThread); + consumers[i] = new Consumer(itemsPerThread, rateLimiter); } Producer producer = new Producer(); @@ -217,12 +219,14 @@ public class StressAction extends Thread private class Consumer extends Thread { private final int items; + private final RateLimiter rateLimiter; private volatile boolean stop = false; private volatile int returnCode = StressAction.SUCCESS; - public Consumer(int toConsume) + public Consumer(int toConsume, RateLimiter rateLimiter) { items = toConsume; + this.rateLimiter = rateLimiter; } public void run() @@ -238,6 +242,7 @@ public class StressAction extends Thread try { + rateLimiter.acquire(); operations.take().run(connection); // running job } catch (Exception e) @@ -266,6 +271,7 @@ public class StressAction extends Thread try { + rateLimiter.acquire(); operations.take().run(connection); // running job } catch (Exception e)