A utility that times creating splits and balancing

SplitTimeTest.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 org.apache.accumulo.test;

import java.util.HashMap;

public class SplitTimeTest {

    private static String TABLE_NAME = "splitTimeTest_Ic0CRwijCjjkonZIoo7l3UnyvtcoPxhh";

    static class MyOpts extends ClientOpts {
	@Parameter(names = { "--preSplit" }, description = "optionally create a small "
		+ "number of tablets and let that balance, before splitting all")
	int preSplit = 0;

	@Parameter(names = { "--numTablets" }, description = "number of tablets to create")
	int numTablets = 100;

	@Parameter(names = { "--offline" }, description = "offline table after splitting")
	boolean offline = false;
    }

    public static void main(String[] args) throws Exception {
	MyOpts opts = new MyOpts();
	opts.parseArgs(ListTables.class.getName(), args);

	Connector conn = opts.getConnector();

	try {
	    conn.tableOperations().delete(TABLE_NAME);
	} catch (TableNotFoundException tnfe) {
	}

	conn.tableOperations().create(TABLE_NAME);

	TreeSet<Text> splits = genSplits(opts.numTablets);

	long total = 0;

	if (opts.preSplit > 1) {
	    TreeSet<Text> preSplits = getPreSplits(opts.preSplit, splits);
	    System.out.print("pre ");
	    total = splitAndBalance(conn, preSplits, opts.offline);
	}

	total += splitAndBalance(conn, splits, opts.offline);

	System.out.printf("total time: %.2f secs\n", total / 1000.0);

    }

    private static TreeSet<Text> getPreSplits(int preSplit, TreeSet<Text> splits) {

	TreeSet<Text> preSplits = new TreeSet<Text>();
	int stride = splits.size() / preSplit;
	int count = 0;

	for (Text split : splits) {
	    count++;
	    if (count == stride) {
		preSplits.add(split);
		count = 0;
	    }
	}

	return preSplits;
    }

    private static long splitAndBalance(Connector conn, TreeSet<Text> splits,
	    boolean offline) throws TableNotFoundException, AccumuloException,
	    AccumuloSecurityException, Exception {
	long t1 = System.currentTimeMillis();

	conn.tableOperations().addSplits(TABLE_NAME, splits);

	long t2 = System.currentTimeMillis();

	if (offline) {
	    conn.tableOperations().offline(TABLE_NAME, true);
	    conn.tableOperations().online(TABLE_NAME, true);
	}
	waitForBalance(conn);

	long t3 = System.currentTimeMillis();

	System.out.printf("split time: %.2f secs  balance time: %.2f secs\n",
		(t2 - t1) / 1000.0, (t3 - t2) / 1000.0);

	return t3 - t1;
    }

    private static TreeSet<Text> genSplits(int numTablets) {
	TreeSet<Text> splits = new TreeSet<Text>();

	int numSplits = numTablets - 1;
	long distance = (Long.MAX_VALUE / numTablets) + 1;
	long split = distance;
	for (int i = 0; i < numSplits; i++) {

	    String s = String.format("%016x", split);

	    while (s.charAt(s.length() - 1) == '0') {
		s = s.substring(0, s.length() - 1);
	    }

	    splits.add(new Text(s));
	    split += distance;
	}
	return splits;
    }

    private static void waitForBalance(Connector conn) throws Exception {

	long diff = getBalanceDifferential(conn);

	while (diff > 3) {
	    Thread.sleep(Math.max(Math.min(10000, diff), 500));
	    diff = getBalanceDifferential(conn);
	}

    }

    static long getBalanceDifferential(Connector conn)
	    throws TableNotFoundException {
	String tableId = conn.tableOperations().tableIdMap().get(TABLE_NAME);
	Map<String, Long> counts = new HashMap<String, Long>();
	Scanner scanner = conn.createScanner(MetadataTable.NAME,
		Authorizations.EMPTY);
	scanner.setRange(TabletsSection.getRange(tableId));
	scanner.fetchColumnFamily(TabletsSection.CurrentLocationColumnFamily.NAME);
	TabletsSection.TabletColumnFamily.PREV_ROW_COLUMN.fetch(scanner);

	for (String tserver : conn.instanceOperations().getTabletServers()) {
	    counts.put(tserver, 0l);
	}

	counts.put("__none__", 0l);

	RowIterator rowIter = new RowIterator(scanner);

	while (rowIter.hasNext()) {
	    Iterator<Entry<Key, Value>> row = rowIter.next();
	    String server = "__none__";

	    while (row.hasNext()) {
		Entry<Key, Value> entry = row.next();
		if (entry
			.getKey()
			.getColumnFamily()
			.equals(TabletsSection.CurrentLocationColumnFamily.NAME)) {
		    server = entry.getValue().toString();
		}
	    }

	    Long count = counts.get(server);
	    counts.put(server, count + 1);
	}

	if (counts.get("__none__") == 0)
	    counts.remove("__none__");

        // TODO if number in __none__ is equal to the rest of the tserver, it 
        // would look balanced even though there are unassigned tablets

	// TODO code does not handle case where #tablets is less than #tservers

	long min = Long.MAX_VALUE;
	long max = Long.MIN_VALUE;

	for (Long l : counts.values()) {
	    if (l < min)
		min = l;
	    if (l > max)
		max = l;
	}

	long diff = max - min;
	return diff;
    }

}