com.spotify.cassandra.opstools.DynamicSnitchDumper.java Source code

Introduction

Here is the source code for com.spotify.cassandra.opstools.DynamicSnitchDumper.java

Source

/*
 *  Licensed 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 com.spotify.cassandra.opstools;

import com.google.common.base.Functions;
import com.google.common.collect.ImmutableSortedMap;
import com.google.common.collect.Ordering;
import org.apache.cassandra.locator.DynamicEndpointSnitchMBean;

import javax.management.*;
import javax.management.remote.JMXConnector;
import javax.management.remote.JMXConnectorFactory;
import javax.management.remote.JMXServiceURL;
import java.io.IOException;
import java.net.InetAddress;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

public class DynamicSnitchDumper {
    private static final String fmtUrl = "service:jmx:rmi:///jndi/rmi://%s:%d/jmxrmi";
    private static final String dsnitchObjName = "org.apache.cassandra.db:type=DynamicEndpointSnitch";

    public static void main(String... args) throws IOException, MalformedObjectNameException {

        String host = "localhost";
        if (args.length > 0)
            host = args[0];
        int port = 7199;
        if (args.length > 1)
            port = Integer.parseInt(args[1]);
        boolean munin = false;
        if (args.length > 2 && args[2].equals("munin"))
            munin = true;
        boolean muninConfig = false;
        if (args.length > 3 && args[3].equals("config"))
            muninConfig = true;

        JMXServiceURL jmxUrl = new JMXServiceURL(String.format(fmtUrl, host, port));
        Map<String, Object> env = new HashMap<String, Object>();
        DynamicEndpointSnitchMBean dsnitch = getDSnitchMbean(jmxUrl, env);

        Map<InetAddress, Double> sorted = sortMap(dsnitch.getScores());
        if (munin) {
            if (muninConfig)
                printMuninConfig(sorted);
            else
                printForMunin(sorted);
        } else {
            for (Map.Entry<InetAddress, Double> score : sorted.entrySet()) {
                System.out.println(score.getKey().getHostName() + " : " + score.getValue());
            }
        }
    }

    private static String getMuninHost(String hostname) {
        return hostname.replace(".spotify.net", "").replace(".", "_");
    }

    private static void printForMunin(Map<InetAddress, Double> sorted) {
        for (Map.Entry<InetAddress, Double> score : sorted.entrySet()) {
            String hostName = getMuninHost(score.getKey().getHostName());
            System.out.println(hostName + ".value " + score.getValue());
        }
    }

    private static void printMuninConfig(Map<InetAddress, Double> sorted) {
        System.out.println("graph_category Cassandra");
        System.out.println("graph_title DynamicSnitch scores");
        System.out.println("graph_vlabel score");
        for (Map.Entry<InetAddress, Double> score : sorted.entrySet()) {
            String hostName = getMuninHost(score.getKey().getHostName());
            System.out.println(hostName + ".graph yes");
            System.out.println(hostName + ".label " + hostName);
        }
    }

    private static DynamicEndpointSnitchMBean getDSnitchMbean(JMXServiceURL jmxUrl, Map<String, Object> env)
            throws IOException, MalformedObjectNameException {
        JMXConnector jmxc = JMXConnectorFactory.connect(jmxUrl, env);
        MBeanServerConnection mbeanServerConn = jmxc.getMBeanServerConnection();
        Set<ObjectName> objs = mbeanServerConn.queryNames(null, null);
        ObjectName realName = null;

        // in cassandra 1.1, the mbean has a "random" instance number, listing mbeans and finding the real one.
        for (ObjectName ob : objs) {
            if (ob.getCanonicalName().contains("DynamicEndpointSnitch"))
                realName = ob;
        }
        if (realName != null)
            return JMX.newMBeanProxy(mbeanServerConn, realName, DynamicEndpointSnitchMBean.class);
        else
            throw new RuntimeException("Could not find the DynamicEndpointSnitch mbean!");
    }

    private static Map<InetAddress, Double> sortMap(Map<InetAddress, Double> scores) {
        return ImmutableSortedMap.copyOf(scores,
                Ordering.natural().onResultOf(Functions.forMap(scores)).compound(new Comparator<InetAddress>() {
                    @Override
                    public int compare(InetAddress o1, InetAddress o2) {
                        return o1.toString().compareTo(o2.toString());
                    }

                    @Override
                    public boolean equals(Object obj) {
                        return false;
                    }
                }));
    }
}