From 297ab0deba6897fabc2e26e42ce807d9d65deb04 Mon Sep 17 00:00:00 2001
From: Radovan Zvoncek <zvo@spotify.com>
Date: Fri, 28 Nov 2014 14:55:00 +0100
Subject: [PATCH] Added JMX client who can fetch tokens from the ring

Change-Id: I3cad0fae9b0fa90f30dcb426ac847c06ebb75e92
---
 pom.xml                                       |  6 ++
 .../com/spotify/reaper/ReaperException.java   | 16 +++++
 .../spotify/reaper/cassandra/ClusterInfo.java | 30 +++++++++
 .../reaper/cassandra/IClusterInfo.java        |  9 +++
 .../spotify/reaper/cassandra/JMXProxy.java    | 67 +++++++++++++++++++
 5 files changed, 128 insertions(+)
 create mode 100644 src/main/java/com/spotify/reaper/ReaperException.java
 create mode 100644 src/main/java/com/spotify/reaper/cassandra/ClusterInfo.java
 create mode 100644 src/main/java/com/spotify/reaper/cassandra/IClusterInfo.java
 create mode 100644 src/main/java/com/spotify/reaper/cassandra/JMXProxy.java

diff --git a/pom.xml b/pom.xml
index cff4de076..a6fd39099 100644
--- a/pom.xml
+++ b/pom.xml
@@ -13,6 +13,7 @@
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
         <dropwizard.version>0.7.0</dropwizard.version>
+        <cassandra.version>2.0.11</cassandra.version>
     </properties>
 
     <dependencies>
@@ -21,6 +22,11 @@
             <artifactId>dropwizard-core</artifactId>
             <version>${dropwizard.version}</version>
         </dependency>
+        <dependency>
+            <groupId>org.apache.cassandra</groupId>
+            <artifactId>cassandra-all</artifactId>
+            <version>${cassandra.version}</version>
+        </dependency>
         <dependency>
             <groupId>junit</groupId>
             <artifactId>junit</artifactId>
diff --git a/src/main/java/com/spotify/reaper/ReaperException.java b/src/main/java/com/spotify/reaper/ReaperException.java
new file mode 100644
index 000000000..e6c7d8fe6
--- /dev/null
+++ b/src/main/java/com/spotify/reaper/ReaperException.java
@@ -0,0 +1,16 @@
+package com.spotify.reaper;
+
+public class ReaperException extends Exception {
+
+  public ReaperException(String s) {
+    super(s);
+  }
+
+  public ReaperException(Exception e) {
+    super(e);
+  }
+
+  public ReaperException(String s, Exception e) {
+    super(s,e);
+  }
+}
diff --git a/src/main/java/com/spotify/reaper/cassandra/ClusterInfo.java b/src/main/java/com/spotify/reaper/cassandra/ClusterInfo.java
new file mode 100644
index 000000000..6737e3bca
--- /dev/null
+++ b/src/main/java/com/spotify/reaper/cassandra/ClusterInfo.java
@@ -0,0 +1,30 @@
+package com.spotify.reaper.cassandra;
+
+import com.spotify.reaper.ReaperException;
+
+import java.util.List;
+
+public class ClusterInfo implements IClusterInfo {
+
+  private String seedHost;
+  private int seedPort = 0;
+
+  private List<String> tokens;
+
+  public ClusterInfo(String seedHost, int seedPort) {
+    this.seedHost = seedHost;
+    this.seedPort = seedPort;
+  }
+
+  public void init() throws ReaperException {
+    JMXProxy jmx = seedPort == 0 ? JMXProxy.connect(seedHost) : JMXProxy.connect(seedHost, seedPort);
+    tokens = jmx.getTokens();
+    jmx.close();
+  }
+
+  @Override
+  public List<String> getTokens() {
+    return tokens;
+  }
+
+}
diff --git a/src/main/java/com/spotify/reaper/cassandra/IClusterInfo.java b/src/main/java/com/spotify/reaper/cassandra/IClusterInfo.java
new file mode 100644
index 000000000..192b53c48
--- /dev/null
+++ b/src/main/java/com/spotify/reaper/cassandra/IClusterInfo.java
@@ -0,0 +1,9 @@
+package com.spotify.reaper.cassandra;
+
+import java.util.List;
+
+public interface IClusterInfo {
+
+  public List<String> getTokens();
+
+}
diff --git a/src/main/java/com/spotify/reaper/cassandra/JMXProxy.java b/src/main/java/com/spotify/reaper/cassandra/JMXProxy.java
new file mode 100644
index 000000000..259b0674e
--- /dev/null
+++ b/src/main/java/com/spotify/reaper/cassandra/JMXProxy.java
@@ -0,0 +1,67 @@
+package com.spotify.reaper.cassandra;
+
+import com.google.common.collect.Lists;
+import com.spotify.reaper.ReaperException;
+import org.apache.cassandra.service.StorageServiceMBean;
+
+import java.io.IOException;
+import java.net.MalformedURLException;
+import java.util.List;
+
+import javax.management.JMX;
+import javax.management.MBeanServerConnection;
+import javax.management.MalformedObjectNameException;
+import javax.management.ObjectName;
+import javax.management.remote.JMXConnector;
+import javax.management.remote.JMXConnectorFactory;
+import javax.management.remote.JMXServiceURL;
+
+public class JMXProxy {
+
+  private static final int DEFAULT_JMX_PORT = 7199;
+
+  private JMXConnector jmxc = null;
+  private StorageServiceMBean ssProxy;
+
+  public JMXProxy(JMXConnector jmxc, StorageServiceMBean ssProxy) {
+    this.jmxc = jmxc;
+    this.ssProxy = ssProxy;
+  }
+
+  public static JMXProxy connect(String host) throws ReaperException {
+    return connect(host, DEFAULT_JMX_PORT);
+  }
+
+  public static JMXProxy connect(String host, int port) throws ReaperException {
+    JMXServiceURL jmxUrl;
+    ObjectName name;
+    try {
+      jmxUrl = new JMXServiceURL(String.format("service:jmx:rmi:///jndi/rmi://%s:%d/jmxrmi",
+          host, port));
+      name = new ObjectName("org.apache.cassandra.db:type=StorageService");
+    } catch (MalformedURLException | MalformedObjectNameException e) {
+      throw new ReaperException("Failure during preparations for JMX connection", e);
+    }
+    try {
+      JMXConnector jmxc = JMXConnectorFactory.connect(jmxUrl);
+      MBeanServerConnection mbeanServerConn = jmxc.getMBeanServerConnection();
+      StorageServiceMBean
+          ssProxy = JMX.newMBeanProxy(mbeanServerConn, name, StorageServiceMBean.class);
+      return new JMXProxy(jmxc, ssProxy);
+    } catch (IOException e) {
+      throw new ReaperException("Failure when establishing JMX connection", e);
+    }
+  }
+
+  public List<String> getTokens() {
+    return Lists.newArrayList(ssProxy.getTokenToEndpointMap().keySet());
+  }
+
+  public void close() throws ReaperException {
+    try {
+      jmxc.close();
+    } catch (IOException e) {
+      throw new ReaperException(e);
+    }
+  }
+}