Support TimeUUIDType column names in Stress.java tool

patch by Pavel Yaskevich; reviewed by Brandon Williams for CASSANDRA-3541

git-svn-id: https://svn.apache.org/repos/asf/cassandra/branches/cassandra-1.0@1208748 13f79535-47bb-0310-9956-ffa450edef68

CHANGES.txt

@@ -2,6 +2,8 @@
  * add command to stop compactions (CASSANDRA-1740)
  * filter out unavailable cipher suites when using encryption (CASSANDRA-3178)
  * (HADOOP) add old-style api support for CFIF and CFRR (CASSANDRA-2799)
+ * Support TimeUUIDType column names in Stress.java tool (CASSANDRA-3541)
+
 
 1.0.5
  * revert CASSANDRA-3407 (see CASSANDRA-3540)

tools/stress/src/org/apache/cassandra/stress/Session.java

@@ -27,9 +27,7 @@ import java.util.concurrent.atomic.AtomicLong;
 
 import org.apache.cassandra.config.CFMetaData;
 import org.apache.cassandra.config.ConfigurationException;
-import org.apache.cassandra.db.compaction.AbstractCompactionStrategy;
+import org.apache.cassandra.db.marshal.*;
-import org.apache.cassandra.db.marshal.AbstractType;
-import org.apache.cassandra.db.marshal.TypeParser;
 import org.apache.commons.cli.*;
 
 import org.apache.cassandra.db.ColumnFamilyType;
@@ -50,6 +48,8 @@ public class Session implements Serializable
     public static final String DEFAULT_COMPARATOR = "AsciiType";
     public static final String DEFAULT_VALIDATOR  = "BytesType";
 
+    private static InetAddress localInetAddress;
+
     public final AtomicInteger operations;
     public final AtomicInteger keys;
     public final AtomicLong    latency;
@@ -89,6 +89,7 @@ public class Session implements Serializable
         availableOptions.addOption("I",  "compression",          true,   "Specify the compression to use for sstable, default:no compression");
         availableOptions.addOption("Q",  "query-names",          true,   "Comma-separated list of column names to retrieve from each row.");
         availableOptions.addOption("Z",  "compaction-strategy",  true,   "CompactionStrategy to use.");
+        availableOptions.addOption("U",  "comparator",           true,   "Column Comparator to use. Currently supported types are: TimeUUIDType, AsciiType, UTF8Type.");
     }
 
     private int numKeys          = 1000 * 1000;
@@ -131,6 +132,8 @@ public class Session implements Serializable
     protected float sigma;
 
     public final InetAddress sendToDaemon;
+    public final String comparator;
+    public final boolean timeUUIDComparator;
 
     public Session(String[] arguments) throws IllegalArgumentException
     {
@@ -321,6 +324,35 @@ public class Session implements Serializable
                     System.exit(1);
                 }
             }
+
+            if (cmd.hasOption("U"))
+            {
+                AbstractType parsed = null;
+
+                try
+                {
+                    parsed = TypeParser.parse(cmd.getOptionValue("U"));
+                }
+                catch (ConfigurationException e)
+                {
+                    System.err.println(e.getMessage());
+                    System.exit(1);
+                }
+
+                comparator = cmd.getOptionValue("U");
+                timeUUIDComparator = parsed instanceof TimeUUIDType;
+
+                if (!(parsed instanceof TimeUUIDType || parsed instanceof AsciiType || parsed instanceof UTF8Type))
+                {
+                    System.err.println("Currently supported types are: TimeUUIDType, AsciiType, UTF8Type.");
+                    System.exit(1);
+                }
+            }
+            else
+            {
+                comparator = null;
+                timeUUIDComparator = false;
+            }
         }
         catch (ParseException e)
         {
@@ -464,6 +496,7 @@ public class Session implements Serializable
     public void createKeySpaces()
     {
         KsDef keyspace = new KsDef();
+        String defaultComparator = comparator == null ? DEFAULT_COMPARATOR : comparator;
 
         // column family for standard columns
         CfDef standardCfDef = new CfDef("Keyspace1", "Standard1");
@@ -471,7 +504,7 @@ public class Session implements Serializable
         if (compression != null)
             compressionOptions.put("sstable_compression", compression);
 
-        standardCfDef.setComparator_type(DEFAULT_COMPARATOR)
+        standardCfDef.setComparator_type(defaultComparator)
                      .setDefault_validation_class(DEFAULT_VALIDATOR)
                      .setCompression_options(compressionOptions);
 
@@ -485,7 +518,7 @@ public class Session implements Serializable
         // column family with super columns
         CfDef superCfDef = new CfDef("Keyspace1", "Super1").setColumn_type("Super");
         superCfDef.setComparator_type(DEFAULT_COMPARATOR)
-                  .setSubcomparator_type(DEFAULT_COMPARATOR)
+                  .setSubcomparator_type(defaultComparator)
                   .setDefault_validation_class(DEFAULT_VALIDATOR)
                   .setCompression_options(compressionOptions);
 
@@ -574,4 +607,20 @@ public class Session implements Serializable
         return client;
     }
 
+    public static InetAddress getLocalAddress()
+    {
+        if (localInetAddress == null)
+        {
+            try
+            {
+                localInetAddress = InetAddress.getLocalHost();
+            }
+            catch (UnknownHostException e)
+            {
+                throw new RuntimeException(e);
+            }
+        }
+
+        return localInetAddress;
+    }
 }

tools/stress/src/org/apache/cassandra/stress/operations/Inserter.java

@@ -52,9 +52,9 @@ public class Inserter extends Operation
 
         for (int i = 0; i < session.getColumnsPerKey(); i++)
         {
-            String columnName = ("C" + Integer.toString(i));
+            columns.add(new Column(columnName(i, session.timeUUIDComparator))
-            ByteBuffer columnValue = values.get(i % values.size());
+                                .setValue(values.get(i % values.size()))
-            columns.add(new Column(ByteBufferUtil.bytes(columnName)).setValue(columnValue).setTimestamp(System.currentTimeMillis()));
+                                .setTimestamp(System.currentTimeMillis()));
         }
 
         if (session.getColumnFamilyType() == ColumnFamilyType.Super)

tools/stress/src/org/apache/cassandra/stress/util/Operation.java

@@ -27,11 +27,14 @@ import java.util.Random;
 
 import static com.google.common.base.Charsets.UTF_8;
 
+import org.apache.cassandra.db.marshal.TimeUUIDType;
 import org.apache.cassandra.stress.Session;
 import org.apache.cassandra.stress.Stress;
 import org.apache.cassandra.thrift.Cassandra;
 import org.apache.cassandra.thrift.InvalidRequestException;
+import org.apache.cassandra.utils.ByteBufferUtil;
 import org.apache.cassandra.utils.FBUtilities;
+import org.apache.cassandra.utils.UUIDGen;
 
 public abstract class Operation
 {
@@ -200,6 +203,13 @@ public abstract class Operation
         return result.toString();
     }
 
+    protected ByteBuffer columnName(int index, boolean timeUUIDComparator)
+    {
+        return timeUUIDComparator
+                ? TimeUUIDType.instance.decompose(UUIDGen.makeType1UUIDFromHost(Session.getLocalAddress()))
+                : ByteBufferUtil.bytes(String.format("C%d", index));
+    }
+
     protected String getExceptionMessage(Exception e)
     {
         String className = e.getClass().getSimpleName();