StatsdHelper.java

package kr.ac.kaist.resl.lilliput.utility;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.nio.channels.DatagramChannel;
import java.util.Locale;

/**
 * @author Jaewook Byun, Ph.D Student, Department of Computer Science, KAIST
 * @email bjw0829@kaist.ac.kr
 * @see   	This code is derived from sample code of a Web Page, although I won't disrespect someone's work, the Web Page is closed
 */
public class StatsdHelper {
	private final InetSocketAddress _address;
	private final DatagramChannel _channel;

	/**
	 * 
	 * @param host Hostname of Statsd  e.g. localhost, XXX.XXX.XX.XX
	 * @param port Port of Statsd e.g. 8125
	 * @throws UnknownHostException
	 * @throws IOException
	 */
	public StatsdHelper(String host, int port) throws UnknownHostException, IOException {
		this(InetAddress.getByName(host), port);
	}

	/**
	 * 
	 * @param host InetAddress of Statsd ,    Use StatsdHelper(String host, int port)
	 * @param port Port of Statsd e.g. 8125
	 * @throws IOException
	 */
	public StatsdHelper(InetAddress host, int port) throws IOException {
		_address = new InetSocketAddress(host, port);
		_channel = DatagramChannel.open();
	}

	/**
	 * 
	 * @param key     Bucket Name
	 * @param value	 Value
	 * @return
	 */
	public boolean timing(String key, int value) {
		return timing(key, value, 1.0);
	}
	
	/**
	 * 
	 * @param key				Bucket Name
	 * @param value			Value
	 * @param sampleRate		Tells StatsD that this counter is being sent sampled every {sampleRate} of the time.
	 * @return
	 */
	public boolean timing(String key, int value, double sampleRate) {
		return send(sampleRate, String.format(Locale.ENGLISH, "%s:%d|ms", key, value));
	}
	
	/**
	 * @see			Part of implementation - Counting
	 * @param key		Bucket name, it would decrease one point
	 * @return
	 */
	public boolean decrement(String key) {
		return increment(key, -1, 1.0);
	}
	
	/**
	 * @see				Part of implementation - Counting
	 * @param key			Bucket name, it would decrease magnitude
	 * @param magnitude
	 * @return
	 */
	public boolean decrement(String key, int magnitude) {
		return decrement(key, magnitude, 1.0);
	}
	
	/**
	 * @see				Part of implementation - Counting
	 * @param key			Bucket name, it would decrease magnitude AT {sampleRate} rate
	 * @param magnitude
	 * @param sampleRate
	 * @return
	 */
	public boolean decrement(String key, int magnitude, double sampleRate) {
		magnitude = magnitude < 0 ? magnitude : -magnitude;
		return increment(key, magnitude, sampleRate);
	}
	
	/**
	 * @see			Part of implementation - Counting
	 * @param keys	names of bucket would decrease one point, respectively
	 * @return
	 */
	public boolean decrement(String... keys) {
		return increment(-1, 1.0, keys);
	}
	
	/**
	 * @see				Part of implementation - Counting
	 * @param magnitude
	 * @param keys		names of bucket would decrease {magnitude} point, respectively
	 * @return
	 */
	public boolean decrement(int magnitude, String... keys) {
		magnitude = magnitude < 0 ? magnitude : -magnitude;
		return increment(magnitude, 1.0, keys);
	}
	
	/**
	 * @see					Part of implementation - Counting
	 * @param magnitude
	 * @param sampleRate
	 * @param keys			names of bucket would decrease {magnitude} point at {sampleRate} rate, respectively
	 * @return
	 */
	public boolean decrement(int magnitude, double sampleRate, String... keys) {
		magnitude = magnitude < 0 ? magnitude : -magnitude;
		return increment(magnitude, sampleRate, keys);
	}
	
	/**
	 * @see					Part of implementation - Counting
	 * @param key
	 * @return					name of bucket would increase one point
	 */
	public boolean increment(String key) {
		return increment(key, 1, 1.0);
	}
	
	/**
	 * @see					Part of implementation - Counting
	 * @param key				name of bucket would increase {magnitude} point
	 * @return					
	 */
	public boolean increment(String key, int magnitude) {
		return increment(key, magnitude, 1.0);
	}
	
	/**
	 * @see					Part of implementation - Counting
	 * @param key				name of bucket would increase {magnitude} point at {sampleRate}
	 * @param magnitude
	 * @param sampleRate
	 * @return
	 */
	public boolean increment(String key, int magnitude, double sampleRate) {
		String stat = String.format(Locale.ENGLISH, "%s:%s|c", key, magnitude);
		return send(sampleRate, stat);
	}
	
	/**
	 * @see					Part of implementation - Counting
	 * @param key				names of bucket would increase {magnitude} point at {sampleRate}
	 * @param magnitude
	 * @param sampleRate
	 * @return
	 */
	public boolean increment(int magnitude, double sampleRate, String... keys) {
		String[] stats = new String[keys.length];
		for (int i = 0; i < keys.length; i++) {
			stats[i] = String.format(Locale.ENGLISH, "%s:%s|c", keys[i], magnitude);
		}
		return send(sampleRate, stats);
	}
	
	/**
	 * @see				Gauge
	 * @param key			Bucket name
	 * @param magnitude	If the gauge is not updated at the next flush, it will send the previous value.
	 * @return
	 */
	public boolean gauge(String key, double magnitude) {
		return gauge(key, magnitude, 1.0);
	}
	public boolean gauge(String key, double magnitude, double sampleRate) {
		final String stat = String.format(Locale.ENGLISH, "%s:%s|g", key, magnitude);
		return send(sampleRate, stat);
	}
	private boolean send(double sampleRate, String... stats) {
		boolean retval = false; // didn't send anything
		if (sampleRate < 1.0) {
			for (String stat : stats) {
				stat = String.format(Locale.ENGLISH, "%s|@%f", stat, sampleRate);
				if (doSend(stat)) {
					retval = true;
				}
			}
		} else {
			for (String stat : stats) {
				if (doSend(stat)) {
					retval = true;
				}
			}
		}
		return retval;
	}
	
	/**
	 * 
	 * @param stat message to Statsd
	 * @return true (no error), false (error)
	 */
	private boolean doSend(final String stat) {
		try {
			// Prepare Message
			final byte[] data = stat.getBytes("utf-8");
			final ByteBuffer buff = ByteBuffer.wrap(data);
			// Send
			final int nbSentBytes = _channel.send(buff, _address);
			if (data.length == nbSentBytes) {
				return true;
			} else {
				System.out.println(String.format(
						"Could not send entirely stat %s to host %s:%d. Only sent %d bytes out of %d bytes", stat,
						_address.getHostName(), _address.getPort(), nbSentBytes, data.length));
				return false;
			}
		} catch (IOException e) {
			System.out.println(String.format("Could not send stat %s to host %s:%d", stat, _address.getHostName(),
					_address.getPort()));
			return false;
		}
	}
}