package org.robotninjas.riemann.load;
import com.aphyr.riemann.Proto;
import com.google.common.base.Supplier;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.google.inject.*;
import com.yammer.metrics.core.*;
import com.yammer.metrics.reporting.ConsoleReporter;
import com.yammer.metrics.reporting.CsvReporter;
import org.robotninjas.riemann.load.annotations.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.File;
import java.util.concurrent.*;
public class LoadTestModule extends PrivateModule {
public static final MetricName SEND_RATE_METRIC_NAME = new MetricName(ClientWorker.class, "sends");
public static final MetricName ACK_RATE_METRIC_NAME = new MetricName(ClientWorker.class, "acks");
public static final MetricName EVENT_ACK_RATE_METRIC_NAME = new MetricName(ClientWorker.class, "event-acks");
public static final MetricName LATENCY_TIMER_NAME = new MetricName(ClientWorker.class, "rtt");
private final Logger logger = LoggerFactory.getLogger(getClass());
private final int workers;
private final int batchSize;
private final Supplier<Proto.Event> eventSupplier;
private final File reportdir;
public LoadTestModule(int workers, int batchSize, Supplier<Proto.Event> eventSupplier, File reportdir) {
this.workers = workers;
this.batchSize = batchSize;
this.eventSupplier = eventSupplier;
this.reportdir = reportdir;
}
@Override
protected void configure() {
bind(Integer.class).annotatedWith(WorkerCount.class).toInstance(workers);
bind(Integer.class).annotatedWith(BatchSize.class).toInstance(batchSize);
bind(new TypeLiteral<Supplier<Proto.Event>>() {
}).annotatedWith(EventSupplier.class).toInstance(eventSupplier);
bind(ClientWorker.class);
bind(LoadTestService.class);
expose(LoadTestService.class);
}
@Provides
@Exposed
@Singleton
@WorkExecutor
public Executor getExecutor() {
final ThreadFactoryBuilder builder = new ThreadFactoryBuilder();
builder.setDaemon(false);
builder.setNameFormat("Test Worker %d");
final ExecutorService executor = Executors.newCachedThreadPool(builder.build());
Runtime.getRuntime().addShutdownHook(new Thread() {
@Override
public void run() {
executor.shutdown();
}
});
return executor;
}
@Provides
@Exposed
@Singleton
@AckExecutor
public Executor getAckExecutor() {
//return new ThreadPoolExecutor(1, 1, 60, TimeUnit.SECONDS, new SynchronousQueue<Runnable>(), new ThreadPoolExecutor.CallerRunsPolicy());
return MoreExecutors.sameThreadExecutor();
}
@Provides
public Supplier<ClientWorker> getClientWorkerSupplier(final Provider<ClientWorker> provider) {
return new Supplier<ClientWorker>() {
@Override
public ClientWorker get() {
return provider.get();
}
};
}
@Provides
@LatencyTimer
public Timer getLatencyTimer(MetricsRegistry registry) {
return registry.newTimer(LATENCY_TIMER_NAME, TimeUnit.SECONDS, TimeUnit.SECONDS);
}
@Provides
@AckMeter
public Meter getAckMeter(MetricsRegistry registry) {
return registry.newMeter(ACK_RATE_METRIC_NAME, "ack", TimeUnit.SECONDS);
}
@Provides
@SendMeter
public Meter getSendMeter(MetricsRegistry registry) {
return registry.newMeter(SEND_RATE_METRIC_NAME, "send", TimeUnit.SECONDS);
}
@Provides
@EventAckMeter
public Meter getEventAckMeter(MetricsRegistry registry) {
return registry.newMeter(EVENT_ACK_RATE_METRIC_NAME, "send", TimeUnit.SECONDS);
}
@Provides
@Exposed
@Singleton
public ConsoleReporter getConsoleReporter(MetricsRegistry registry) {
return new ConsoleReporter(registry, System.out, MetricPredicate.ALL);
}
@Provides
@Exposed
@Singleton
public CsvReporter getCsvReporter(MetricsRegistry registry) {
return new CsvReporter(registry, reportdir);
}
}