co.signal.loadgen.LoadGenerator Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of kafkameter Show documentation
Kafka plugin for JMeter
The newest version!
/*
 * Copyright 2014 Signal.
 *
 * 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 co.signal.loadgen;

import java.io.File;
import java.io.IOException;
import java.util.Map;

import javax.annotation.Nullable;

import com.google.common.base.Charsets;
import com.google.common.base.Throwables;
import com.google.common.io.Files;

import org.apache.commons.lang3.tuple.Pair;
import org.apache.jmeter.config.ConfigTestElement;
import org.apache.jmeter.engine.event.LoopIterationEvent;
import org.apache.jmeter.engine.event.LoopIterationListener;
import org.apache.jmeter.testbeans.TestBean;
import org.apache.jmeter.threads.JMeterContextService;
import org.apache.jmeter.threads.JMeterVariables;
import org.apache.jorphan.logging.LoggingManager;
import org.apache.log.Logger;

/**
 * Config Element which reads a Synthetic Load Description from a file, generates
 * a domain-specific message, and exports the message under a given variableName.
 *
 * @author codyaray
 * @since 6/27/14
 */
public class LoadGenerator extends ConfigTestElement implements TestBean, LoopIterationListener {

  private static final Logger log = LoggingManager.getLoggerForClass();

  private String fileName;
  private String keyVariableName;
  private String valueVariableName;
  private String className;

  private SyntheticLoadGenerator generator;

  @Override
  public void iterationStart(LoopIterationEvent loopIterationEvent) {
    if (generator == null) {
      generator = createGenerator(getClassName(), readFile(getFileName()));
    }
    JMeterVariables variables = JMeterContextService.getContext().getVariables();
    Pair generatedMessage = generator.nextMessage();
    variables.put(getKeyVariableName(), generatedMessage.getKey());
    variables.put(getValueVariableName(), generatedMessage.getValue());
  }

  private SyntheticLoadGenerator createGenerator(String className, @Nullable String config) {
    try {
      return (SyntheticLoadGenerator) Class.forName(
          className, false, Thread.currentThread().getContextClassLoader()
      ).getConstructor(String.class).newInstance(config);
    } catch (Exception e) {
      log.fatalError("Exception initializing Load Generator class: " + className, e);
      throw Throwables.propagate(e);
    }
  }

  private static @Nullable String readFile(String filename) {
    try {
      return Files.toString(new File(filename), Charsets.UTF_8);
    } catch (IOException e) {
      log.warn("Unable to retrieve config from " + filename, e);
      return null;
    }
  }

  /**
   * @return the keyVariableName
   */
  public String getKeyVariableName() {
    return keyVariableName;
  }

  /**
   * @param keyVariableName the keyVariableName to set
   */
  public void setKeyVariableName(String keyVariableName) {
    this.keyVariableName = keyVariableName;
  }

  /**
   * @return the valueVariableName
   */
  public String getValueVariableName() {
    return valueVariableName;
  }

  /**
   * @param valueVariableName the valueVariableName to set
   */
  public void setValueVariableName(String valueVariableName) {
    this.valueVariableName = valueVariableName;
  }

  /**
   * @return the fileName
   */
  public String getFileName() {
    return fileName;
  }

  /**
   * @param fileName the fileName to set
   */
  public void setFileName(String fileName) {
    this.fileName = fileName;
  }

  /**
   * @return the className
   */
  public String getClassName() {
    return className;
  }

  /**
   * @param className the className used for generating the messages
   */
  public void setClassName(String className) {
    this.className = className;
  }

  /**
   * Helper for testing outside of JMeter
   *
   * @param args command line arguments
   */
  public static void main(String[] args) {
    LoadGenerator generator = new LoadGenerator();

    // Mock out JMeter environment
    JMeterVariables variables = new JMeterVariables();
    JMeterContextService.getContext().setVariables(variables);
    generator.setFileName("config.json");
    generator.setKeyVariableName("kafka_key");
    generator.setValueVariableName("kafka_message");

    generator.iterationStart(null);

    for (Map.Entry entry : variables.entrySet()) {
      System.out.println(entry.getKey() + " : " + entry.getValue());
    }
  }
}