Java tutorial
/* * 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.autobalance; import org.apache.cassandra.dht.Murmur3Partitioner; import org.apache.cassandra.dht.RandomPartitioner; import org.apache.cassandra.tools.NodeProbe; import org.apache.commons.cli.BasicParser; import org.apache.commons.cli.CommandLine; import org.apache.commons.cli.CommandLineParser; import org.apache.commons.cli.Options; import org.apache.commons.cli.ParseException; import java.io.IOException; import java.math.BigInteger; import java.net.InetAddress; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; public class Main { private void run(CommandLine cmd) throws IOException, InterruptedException { boolean dryrun = cmd.hasOption("d"); boolean force = cmd.hasOption("f"); boolean noresolve = cmd.hasOption("r"); int port = cmd.hasOption("p") ? Integer.parseInt(cmd.getOptionValue("p")) : 7199; String nodehost = cmd.hasOption("h") ? cmd.getOptionValue("h") : "localhost"; System.out.println("Collecting information about the cluster..."); NodeProbe nodeProbe = new NodeProbe(nodehost, port); if (nodeProbe.getTokens().size() != 1) { System.err.println("Cluster is using vnodes and should already be automatically balanced!"); System.exit(1); } boolean hasData = false; if (!dryrun) { Map<String, String> loadMap = nodeProbe.getLoadMap(); for (String s : loadMap.values()) { if (s.contains("KB")) continue; if (s.contains("MB") || s.contains("GB") || s.contains("TB")) { hasData = true; continue; } throw new RuntimeException("Unknown suffix in load map; don't dare to continue"); } } String partitioner = nodeProbe.getPartitioner(); BigInteger minToken, maxToken; if (partitioner.equals(RandomPartitioner.class.getName())) { minToken = RandomPartitioner.ZERO; maxToken = RandomPartitioner.MAXIMUM; } else if (partitioner.equals(Murmur3Partitioner.class.getName())) { minToken = BigInteger.valueOf(Murmur3Partitioner.MINIMUM.token); maxToken = BigInteger.valueOf(Murmur3Partitioner.MAXIMUM); } else { throw new RuntimeException("Unsupported partitioner: " + partitioner); } // Get current mapping of all live nodes List<String> liveNodes = nodeProbe.getLiveNodes(); Map<String, BigInteger> hostTokenMap = new HashMap<String, BigInteger>(); Map<String, String> hostDcMap = new HashMap<String, String>(); for (String host : liveNodes) { String dc = nodeProbe.getEndpointSnitchInfoProxy().getDatacenter(host); String decoratedHost = host; if (!noresolve) { // Prefix host with canonical host name. // This makes things prettier and also causes tokens to be assigned in logical order. decoratedHost = InetAddress.getByName(host).getCanonicalHostName() + "/" + host; } else { decoratedHost = "/" + host; } hostDcMap.put(decoratedHost, dc); List<String> tokens = nodeProbe.getTokens(host); if (tokens.size() > 1) { throw new RuntimeException("vnodes not supported"); } if (tokens.size() == 0) { throw new RuntimeException("No token for " + host + "; aborting"); } hostTokenMap.put(decoratedHost, new BigInteger(tokens.get(0))); } Balancer balancer = new Balancer(hostTokenMap, hostDcMap, minToken, maxToken); Map<String, BigInteger> newMap = balancer.balance(); List<Operation> operations = new ArrayList<Operation>(); boolean movesNeeded = false; for (Map.Entry<String, BigInteger> entry : hostTokenMap.entrySet()) { String host = entry.getKey(); BigInteger oldToken = entry.getValue(); BigInteger newToken = newMap.get(host); if (!oldToken.equals(newToken)) { movesNeeded = true; } operations.add(new Operation(host, hostDcMap.get(host), oldToken, newToken)); } if (movesNeeded && hasData && !dryrun && !force) { dryrun = true; System.out.println( "The cluster is unbalanced but has data, so no operations will actually be carried out. Use --force if you want the cluster to balance anyway."); } Collections.sort(operations); boolean unbalanced = false, moved = false; for (Operation op : operations) { if (op.oldToken.equals(op.newToken)) { System.out.println(op.host + ": Stays on token " + op.oldToken); } else { System.out.println(op.host + ": Moving from token " + op.oldToken + " to token " + op.newToken); if (!dryrun) { String ip = op.host.substring(op.host.lastIndexOf("/") + 1); NodeProbe np = new NodeProbe(ip, 7199); np.move(op.newToken.toString()); moved = true; } else { unbalanced = true; } } } if (!unbalanced && moved) { System.out.println("The cluster is now balanced!"); } } private static class Operation implements Comparable<Operation> { public String host; public String dataCenter; public BigInteger oldToken; public BigInteger newToken; private Operation(String host, String dataCenter, BigInteger oldToken, BigInteger newToken) { this.host = host; this.dataCenter = dataCenter; this.oldToken = oldToken; this.newToken = newToken; } @Override public int compareTo(Operation o) { if (!dataCenter.equals(o.dataCenter)) { return dataCenter.compareTo(o.dataCenter); } return newToken.compareTo(o.newToken); } } public static void main(String[] args) throws IOException, InterruptedException, ParseException { final Options options = new Options(); options.addOption("f", "force", false, "Force auto balance"); options.addOption("d", "dryrun", false, "Dry run"); options.addOption("r", "noresolve", false, "Don't resolve host names"); options.addOption("h", "host", true, "Host to connect to (default: localhost)"); options.addOption("p", "port", true, "Port to connect to (default: 7199)"); CommandLineParser parser = new BasicParser(); CommandLine cmd = parser.parse(options, args); new Main().run(cmd); } }