UltimateThreadGroup.java

package kg.apc.jmeter.threads;

import kg.apc.jmeter.JMeterPluginsUtils;
import org.apache.jmeter.gui.util.PowerTableModel;
import org.apache.jmeter.testelement.TestStateListener;
import org.apache.jmeter.testelement.property.CollectionProperty;
import org.apache.jmeter.testelement.property.JMeterProperty;
import org.apache.jmeter.testelement.property.NullProperty;
import org.apache.jmeter.testelement.property.PropertyIterator;
import org.apache.jmeter.threads.JMeterThread;
import org.apache.jmeter.util.JMeterUtils;
import org.apache.jorphan.logging.LoggingManager;
import org.apache.log.Logger;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;

public class UltimateThreadGroup
        extends AbstractSimpleThreadGroup
        implements Serializable, TestStateListener {

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

    public static final String DATA_PROPERTY = "ultimatethreadgroupdata";
    public static final String EXTERNAL_DATA_PROPERTY = "threads_schedule";

    public static final int START_THREADS_CNT_FIELD_NO = 0;
    public static final int INIT_DELAY_FIELD_NO = 1;
    public static final int STARTUP_TIME_FIELD_NO = 2;
    public static final int HOLD_LOAD_FOR_FIELD_NO = 3;
    public static final int SHUTDOWN_TIME_FIELD_NO = 4;

    private PropertyIterator scheduleIT;
    private int threadsToSchedule;
    private CollectionProperty currentRecord;

    public UltimateThreadGroup() {
        super();
    }

    @Override
    protected void scheduleThread(JMeterThread thread, long tgStartTime) {
        log.debug("Scheduling thread: " + thread.getThreadName());
        if (threadsToSchedule < 1) {
            if (!scheduleIT.hasNext()) {
                throw new RuntimeException("Not enough schedule records for thread #" + thread.getThreadName());
            }

            currentRecord = (CollectionProperty) scheduleIT.next();
            threadsToSchedule = currentRecord.get(0).getIntValue();
        }

        int numThreads = currentRecord.get(START_THREADS_CNT_FIELD_NO).getIntValue();
        int initialDelay = currentRecord.get(INIT_DELAY_FIELD_NO).getIntValue();
        int startRampUp = currentRecord.get(STARTUP_TIME_FIELD_NO).getIntValue();
        int flightTime = currentRecord.get(HOLD_LOAD_FOR_FIELD_NO).getIntValue();
        int endRampUp = currentRecord.get(SHUTDOWN_TIME_FIELD_NO).getIntValue();

        long ascentPoint = tgStartTime + 1000 * initialDelay;
        final int rampUpDelayForThread = (int) Math.floor(1000 * startRampUp * (double) threadsToSchedule / numThreads);
        long startTime = ascentPoint + rampUpDelayForThread;
        long descentPoint = startTime + 1000 * flightTime + 1000 * startRampUp - rampUpDelayForThread;

        thread.setStartTime(startTime);
        thread.setEndTime(descentPoint + (int) Math.floor(1000 * endRampUp * (double) threadsToSchedule / numThreads));

        thread.setScheduled(true);
        threadsToSchedule--;
    }

    public JMeterProperty getData() {
        JMeterProperty brokenProp = getProperty(EXTERNAL_DATA_PROPERTY);
        JMeterProperty usualProp = getProperty(DATA_PROPERTY);

        if (brokenProp instanceof CollectionProperty) {
            if (usualProp == null || usualProp instanceof NullProperty) {
                log.warn("Copying '" + EXTERNAL_DATA_PROPERTY + "' into '" + DATA_PROPERTY + "'");
                JMeterProperty newProp = brokenProp.clone();
                newProp.setName(DATA_PROPERTY);
                setProperty(newProp);
            }
            log.warn("Removing property '" + EXTERNAL_DATA_PROPERTY + "' as invalid");
            removeProperty(EXTERNAL_DATA_PROPERTY);
        }

        //log.info("getData: "+getProperty(DATA_PROPERTY));
        CollectionProperty overrideProp = getLoadFromExternalProperty();
        if (overrideProp != null) {
            return overrideProp;
        }

        return getProperty(DATA_PROPERTY);
    }

    public void setData(CollectionProperty rows) {
        //log.info("setData");
        setProperty(rows);
    }


    private CollectionProperty getLoadFromExternalProperty() {
        String loadProp = JMeterUtils.getProperty(EXTERNAL_DATA_PROPERTY);
        log.debug("Profile prop: " + loadProp);
        if (loadProp != null && loadProp.length() > 0) {
            //expected format : threads_schedule="spawn(1,1s,1s,1s,1s) spawn(2,1s,3s,1s,2s)"
            log.info("GUI threads profile will be ignored");
            PowerTableModel dataModel = new PowerTableModel(UltimateThreadGroupGui.columnIdentifiers, UltimateThreadGroupGui.columnClasses);
            String[] chunks = loadProp.split("\\)");

            for (String chunk : chunks) {
                try {
                    parseChunk(chunk, dataModel);
                } catch (RuntimeException e) {
                    log.warn("Wrong  chunk ignored: " + chunk, e);
                }
            }

            log.info("Setting threads profile from property " + EXTERNAL_DATA_PROPERTY + ": " + loadProp);
            return JMeterPluginsUtils.tableModelRowsToCollectionProperty(dataModel, UltimateThreadGroup.DATA_PROPERTY);
        }
        return null;
    }

    private static void parseChunk(String chunk, PowerTableModel model) {
        log.debug("Parsing chunk: " + chunk);
        String[] parts = chunk.split("[(,]");
        String loadVar = parts[0].trim();

        if (loadVar.equalsIgnoreCase("spawn")) {
            Integer[] row = new Integer[5];
            row[START_THREADS_CNT_FIELD_NO] = Integer.parseInt(parts[1].trim());
            row[INIT_DELAY_FIELD_NO] = JMeterPluginsUtils.getSecondsForShortString(parts[2]);
            row[STARTUP_TIME_FIELD_NO] = JMeterPluginsUtils.getSecondsForShortString(parts[3]);
            row[HOLD_LOAD_FOR_FIELD_NO] = JMeterPluginsUtils.getSecondsForShortString(parts[4]);
            row[SHUTDOWN_TIME_FIELD_NO] = JMeterPluginsUtils.getSecondsForShortString(parts[5]);
            model.addRow(row);
        } else {
            throw new RuntimeException("Unknown load type: " + parts[0]);
        }
    }

    @Override
    public int getNumThreads() {
        int result = 0;

        JMeterProperty threadValues = getData();
        if (!(threadValues instanceof NullProperty)) {
            CollectionProperty columns = (CollectionProperty) threadValues;
            List<?> rows = (List<?>) columns.getObjectValue();
            for (Object row1 : rows) {
                CollectionProperty prop = (CollectionProperty) row1;
                ArrayList<JMeterProperty> row = (ArrayList<JMeterProperty>) prop.getObjectValue();
                //log.info(prop.getStringValue());
                result += row.get(0).getIntValue();
            }
        }

        return result;
    }

    @Override
    public void testStarted() {
        JMeterProperty data = getData();
        if (!(data instanceof NullProperty)) {
            scheduleIT = ((CollectionProperty) data).iterator();
        }
        threadsToSchedule = 0;
    }

    @Override
    public void testStarted(String host) {
        testStarted();
    }

    @Override
    public void testEnded() {
    }

    @Override
    public void testEnded(String host) {
        testEnded();
    }


}