From 4efb6dc8c13e56b4d93ac7799a6edaa7dfc07e6f Mon Sep 17 00:00:00 2001
From: Benedict Elliott Smith <benedict@apache.org>
Date: Tue, 16 Sep 2014 06:28:01 +0100
Subject: [PATCH] ninja fix cassandra-stress pre-2.1.1 compatibility

---
 .../cassandra/stress/StressMetrics.java       |  9 +++++--
 .../stress/settings/SettingsNode.java         | 25 ++++++++++++++++++-
 .../cassandra/stress/util/JmxCollector.java   | 12 ++++++---
 .../stress/util/SmartThriftClient.java        |  3 +--
 4 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/tools/stress/src/org/apache/cassandra/stress/StressMetrics.java b/tools/stress/src/org/apache/cassandra/stress/StressMetrics.java
index 6d5f387d36..00f479e49a 100644
--- a/tools/stress/src/org/apache/cassandra/stress/StressMetrics.java
+++ b/tools/stress/src/org/apache/cassandra/stress/StressMetrics.java
@@ -32,6 +32,7 @@ import java.util.concurrent.TimeUnit;
 import org.apache.commons.lang3.time.DurationFormatUtils;
 
 import org.apache.cassandra.concurrent.NamedThreadFactory;
+import org.apache.cassandra.stress.settings.SettingsLog;
 import org.apache.cassandra.stress.settings.StressSettings;
 import org.apache.cassandra.stress.util.JmxCollector;
 import org.apache.cassandra.stress.util.Timing;
@@ -60,11 +61,15 @@ public class StressMetrics
         totalGcStats = new JmxCollector.GcStats(0);
         try
         {
-            gcStatsCollector = new JmxCollector(settings.node.nodes, settings.port.jmxPort);
+            gcStatsCollector = new JmxCollector(settings.node.resolveAllPermitted(settings), settings.port.jmxPort);
         }
         catch (Throwable t)
         {
-            t.printStackTrace();
+            switch (settings.log.level)
+            {
+                case VERBOSE:
+                    t.printStackTrace();
+            }
             System.err.println("Failed to connect over JMX; not collecting these stats");
             gcStatsCollector = new Callable<JmxCollector.GcStats>()
             {
diff --git a/tools/stress/src/org/apache/cassandra/stress/settings/SettingsNode.java b/tools/stress/src/org/apache/cassandra/stress/settings/SettingsNode.java
index 30fe908290..5297da8026 100644
--- a/tools/stress/src/org/apache/cassandra/stress/settings/SettingsNode.java
+++ b/tools/stress/src/org/apache/cassandra/stress/settings/SettingsNode.java
@@ -32,6 +32,8 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
+import com.datastax.driver.core.Host;
+
 public class SettingsNode implements Serializable
 {
     public final List<String> nodes;
@@ -71,7 +73,28 @@ public class SettingsNode implements Serializable
         isWhiteList = options.whitelist.setByUser();
     }
 
-    public Set<InetAddress> resolveAll()
+    public Set<String> resolveAllPermitted(StressSettings settings)
+    {
+        Set<String> r = new HashSet<>();
+        switch (settings.mode.api)
+        {
+            case THRIFT_SMART:
+            case JAVA_DRIVER_NATIVE:
+                if (!isWhiteList)
+                {
+                    for (Host host : settings.getJavaDriverClient().getCluster().getMetadata().getAllHosts())
+                        r.add(host.getAddress().getHostName());
+                    break;
+                }
+            case THRIFT:
+            case SIMPLE_NATIVE:
+                for (InetAddress address : resolveAllSpecified())
+                    r.add(address.getHostName());
+        }
+        return r;
+    }
+
+    public Set<InetAddress> resolveAllSpecified()
     {
         Set<InetAddress> r = new HashSet<>();
         for (String node : nodes)
diff --git a/tools/stress/src/org/apache/cassandra/stress/util/JmxCollector.java b/tools/stress/src/org/apache/cassandra/stress/util/JmxCollector.java
index 9611b2a1cf..8cfbebbd65 100644
--- a/tools/stress/src/org/apache/cassandra/stress/util/JmxCollector.java
+++ b/tools/stress/src/org/apache/cassandra/stress/util/JmxCollector.java
@@ -20,6 +20,7 @@ package org.apache.cassandra.stress.util;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutorService;
@@ -75,11 +76,16 @@ public class JmxCollector implements Callable<JmxCollector.GcStats>
     final NodeProbe[] probes;
 
     // TODO: should expand to whole cluster
-    public JmxCollector(List<String> hosts, int port)
+    public JmxCollector(Collection<String> hosts, int port)
     {
         probes = new NodeProbe[hosts.size()];
-        for (int i = 0 ; i < hosts.size() ; i++)
-            probes[i] = connect(hosts.get(i), port);
+        int i = 0;
+        for (String host : hosts)
+        {
+            probes[i] = connect(host, port);
+            probes[i].getAndResetGCStats();
+            i++;
+        }
     }
 
     private static NodeProbe connect(String host, int port)
diff --git a/tools/stress/src/org/apache/cassandra/stress/util/SmartThriftClient.java b/tools/stress/src/org/apache/cassandra/stress/util/SmartThriftClient.java
index b880283c8e..6367936bd2 100644
--- a/tools/stress/src/org/apache/cassandra/stress/util/SmartThriftClient.java
+++ b/tools/stress/src/org/apache/cassandra/stress/util/SmartThriftClient.java
@@ -22,7 +22,6 @@ package org.apache.cassandra.stress.util;
 
 
 import java.net.InetAddress;
-import java.net.UnknownHostException;
 import java.nio.ByteBuffer;
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
@@ -64,7 +63,7 @@ public class SmartThriftClient implements ThriftClient
         }
         else
         {
-            whiteset = settings.node.resolveAll();
+            whiteset = settings.node.resolveAllSpecified();
             whitelist = Arrays.asList(whiteset.toArray(new InetAddress[0]));
         }
     }