LoadTest.java

package org.robotninjas.riemann.load;

import com.aphyr.riemann.Proto;
import com.google.common.base.Supplier;
import com.google.common.net.HostAndPort;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.palominolabs.metrics.guice.InstrumentationModule;
import com.yammer.metrics.reporting.ConsoleReporter;
import com.yammer.metrics.reporting.CsvReporter;
import org.apache.commons.cli.*;
import org.apache.commons.pool.impl.GenericObjectPool;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.util.List;
import java.util.concurrent.TimeUnit;

import static java.lang.Integer.parseInt;

public class LoadTest {

  private static final int NUM_CLIENT_WORKERS = 1;
  private static final int BATCH_SIZE = 1;
  private static final int NUM_CONNECTIONS = 1;
  private static final int NUM_NETTY_WORKERS = 1;
  private static final String DEFAULT_REPORTDIR = ".";

  private final Logger logger = LoggerFactory.getLogger(getClass());
  private final String address;
  private final int port;
  private final int clientWorkers;
  private final int batchSize;
  private final int numConnections;
  private final int numNettyWorkers;
  private final Supplier<Proto.Event> eventSupplier;
  private final File reportdir;

  public LoadTest(String address, int port) {
    this(address, port, NUM_CLIENT_WORKERS, BATCH_SIZE, NUM_CONNECTIONS,
      NUM_NETTY_WORKERS, new DefaultEventSupplier(), new File(DEFAULT_REPORTDIR));
  }

  public LoadTest(String address, int port, int clientWorkers) {
    this(address, port, clientWorkers, BATCH_SIZE, NUM_CONNECTIONS,
      NUM_NETTY_WORKERS, new DefaultEventSupplier(), new File(DEFAULT_REPORTDIR));
  }

  public LoadTest(String address, int port, int clientWorkers, int batchSize, int numConnections,
                  int numNettyWorkers, Supplier<Proto.Event> eventSupplier, File reportdir) {

    this.address = address;
    this.port = port;
    this.clientWorkers = clientWorkers;
    this.batchSize = batchSize;
    this.numConnections = numConnections;
    this.numNettyWorkers = numNettyWorkers;
    this.eventSupplier = eventSupplier;
    this.reportdir = reportdir;
  }

  public void start() {

    final GenericObjectPool.Config poolConfig = new GenericObjectPool.Config();
    poolConfig.maxActive = numConnections;

    final long filename = System.nanoTime();
    final File thisReportdir = new File(reportdir, Long.toString(filename));
    thisReportdir.mkdir();

      logger.info("Reports will be placed at {}", thisReportdir.toString());

    final Injector injector = Guice.createInjector(
      new InstrumentationModule(),
      new InstrumentedClientModule(address, port, numNettyWorkers, poolConfig),
      new LoadTestModule(clientWorkers, batchSize, eventSupplier, thisReportdir));

    final ConsoleReporter consoleReporter = injector.getInstance(ConsoleReporter.class);
    Runtime.getRuntime().addShutdownHook(new Thread() {
      @Override
      public void run() {
        consoleReporter.shutdown();
      }
    });
    consoleReporter.start(1, TimeUnit.SECONDS);

    final CsvReporter csvReporter = injector.getInstance(CsvReporter.class);
    Runtime.getRuntime().addShutdownHook(new Thread() {
      @Override
      public void run() {
        csvReporter.shutdown();
      }
    });
    csvReporter.start(1, TimeUnit.SECONDS);

    final LoadTestService loadTestService = injector.getInstance(LoadTestService.class);
    Runtime.getRuntime().addShutdownHook(new Thread() {
      @Override
      public void run() {
        loadTestService.stopAndWait();
      }
    });
    loadTestService.startAndWait();

  }

  public static void main(String[] args) {

    final Options opts = new Options();
    opts.addOption("w", "client-workers", true, "number of client workers, default is 1");
    opts.addOption("b", "batch-size", true, "number of Events to send in each Msg, default is 1");
    opts.addOption("c", "connections", true, "number of concurrent connections, if not specified then equal to -w");
    opts.addOption("n", "netty-workers", true, "number of netty worker threads, if not specified then equal to -w");
    opts.addOption("s", "buffer-size", true, "netty pipeline buffer size in bytes, default is 1");
    opts.addOption("r", "reports-dir", true, "base directory for generated reports, default is ./<timestamp>");

    try {

      final Parser parser = new PosixParser();
      final CommandLine line = parser.parse(opts, args);

      final int clientWorkers = line.hasOption('w') ? parseInt(line.getOptionValue('w')) : NUM_CLIENT_WORKERS;
      final int batchSize = line.hasOption('b') ? parseInt(line.getOptionValue('b')) : BATCH_SIZE;
      final int numConnections = line.hasOption('c') ? parseInt(line.getOptionValue('c')) : clientWorkers;
      final int numNettyWorkers = line.hasOption('n') ? parseInt(line.getOptionValue('n')) : clientWorkers;
      final String reportdir = line.getOptionValue("r", DEFAULT_REPORTDIR);

      final List<String> otherArgs = line.getArgList();
      HostAndPort riemannHostAndPort = HostAndPort.fromParts("localhost", 5555);
      if (otherArgs.size() >= 1) {
        riemannHostAndPort = HostAndPort.fromString(otherArgs.get(0));
      }

      final LoadTest loadTest =
        new LoadTest(riemannHostAndPort.getHostText(), riemannHostAndPort.getPort(), clientWorkers,
          batchSize, numConnections, numNettyWorkers, new DefaultEventSupplier(), new File(reportdir));

      loadTest.start();

    } catch (ParseException e) {
      final HelpFormatter formatter = new HelpFormatter();
      formatter.printHelp("load-test", opts);
    }


  }

}