org.apache.metron.performance.load.LoadGenerator.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.metron.performance.load.LoadGenerator.java

Source

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.metron.performance.load;

import com.google.common.base.Joiner;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.PosixParser;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.apache.metron.common.utils.KafkaUtils;
import org.apache.metron.performance.load.monitor.AbstractMonitor;
import org.apache.metron.performance.load.monitor.EPSGeneratedMonitor;
import org.apache.metron.performance.load.monitor.EPSThroughputWrittenMonitor;
import org.apache.metron.performance.load.monitor.MonitorTask;
import org.apache.metron.performance.load.monitor.writers.CSVWriter;
import org.apache.metron.performance.load.monitor.writers.ConsoleWriter;
import org.apache.metron.performance.load.monitor.writers.Writable;
import org.apache.metron.performance.load.monitor.writers.Writer;
import org.apache.metron.performance.sampler.BiasedSampler;
import org.apache.metron.performance.sampler.Sampler;
import org.apache.metron.performance.sampler.UnbiasedSampler;

import java.io.File;
import java.util.ArrayList;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;

public class LoadGenerator {
    public static String CONSUMER_GROUP = "metron.load.group";
    public static long SEND_PERIOD_MS = 100;
    public static long MONITOR_PERIOD_MS = 1000 * 10;
    private static ExecutorService pool;
    private static ThreadLocal<KafkaProducer<String, String>> kafkaProducer;
    public static AtomicLong numSent = new AtomicLong(0);

    public static void main(String[] args) throws Exception {
        CommandLine cli = LoadOptions.parse(new PosixParser(), args);
        EnumMap<LoadOptions, Optional<Object>> evaluatedArgs = LoadOptions.createConfig(cli);
        Map<String, Object> kafkaConfig = new HashMap<>();
        kafkaConfig.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
        kafkaConfig.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
        kafkaConfig.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        kafkaConfig.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        if (LoadOptions.ZK.has(cli)) {
            String zkQuorum = (String) evaluatedArgs.get(LoadOptions.ZK).get();
            kafkaConfig.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG,
                    Joiner.on(",").join(KafkaUtils.INSTANCE.getBrokersFromZookeeper(zkQuorum)));
        }
        String groupId = evaluatedArgs.get(LoadOptions.CONSUMER_GROUP).get().toString();
        System.out.println("Consumer Group: " + groupId);
        kafkaConfig.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
        if (LoadOptions.KAFKA_CONFIG.has(cli)) {
            kafkaConfig.putAll((Map<String, Object>) evaluatedArgs.get(LoadOptions.KAFKA_CONFIG).get());
        }
        kafkaProducer = ThreadLocal.withInitial(() -> new KafkaProducer<>(kafkaConfig));
        int numThreads = (int) evaluatedArgs.get(LoadOptions.NUM_THREADS).get();
        System.out.println("Thread pool size: " + numThreads);
        pool = Executors.newFixedThreadPool(numThreads);
        Optional<Object> eps = evaluatedArgs.get(LoadOptions.EPS);

        Optional<Object> outputTopic = evaluatedArgs.get(LoadOptions.OUTPUT_TOPIC);
        Optional<Object> monitorTopic = evaluatedArgs.get(LoadOptions.MONITOR_TOPIC);
        long sendDelta = (long) evaluatedArgs.get(LoadOptions.SEND_DELTA).get();
        long monitorDelta = (long) evaluatedArgs.get(LoadOptions.MONITOR_DELTA).get();
        if ((eps.isPresent() && outputTopic.isPresent()) || monitorTopic.isPresent()) {
            Timer timer = new Timer(false);
            long startTimeMs = System.currentTimeMillis();
            if (outputTopic.isPresent() && eps.isPresent()) {
                List<String> templates = (List<String>) evaluatedArgs.get(LoadOptions.TEMPLATE).get();
                if (templates.isEmpty()) {
                    System.out.println("Empty templates, so nothing to do.");
                    return;
                }
                Optional<Object> biases = evaluatedArgs.get(LoadOptions.BIASED_SAMPLE);
                Sampler sampler = new UnbiasedSampler();
                if (biases.isPresent()) {
                    sampler = new BiasedSampler((List<Map.Entry<Integer, Integer>>) biases.get(), templates.size());
                }
                MessageGenerator generator = new MessageGenerator(templates, sampler);
                Long targetLoad = (Long) eps.get();
                int periodsPerSecond = (int) (1000 / sendDelta);
                long messagesPerPeriod = targetLoad / periodsPerSecond;
                String outputTopicStr = (String) outputTopic.get();
                System.out.println(
                        "Generating data to " + outputTopicStr + " at " + targetLoad + " events per second");
                System.out.println("Sending " + messagesPerPeriod + " messages to " + outputTopicStr + " every "
                        + sendDelta + "ms");
                timer.scheduleAtFixedRate(new SendToKafka(outputTopicStr, messagesPerPeriod, numThreads, generator,
                        pool, numSent, kafkaProducer), 0, sendDelta);
            }
            List<AbstractMonitor> monitors = new ArrayList<>();
            if (outputTopic.isPresent() && monitorTopic.isPresent()) {
                System.out.println("Monitoring " + monitorTopic.get() + " every " + monitorDelta + " ms");
                monitors.add(new EPSGeneratedMonitor(outputTopic, numSent));
                monitors.add(new EPSThroughputWrittenMonitor(monitorTopic, kafkaConfig));
            } else if (outputTopic.isPresent() && !monitorTopic.isPresent()) {
                System.out.println("Monitoring " + outputTopic.get() + " every " + monitorDelta + " ms");
                monitors.add(new EPSGeneratedMonitor(outputTopic, numSent));
                monitors.add(new EPSThroughputWrittenMonitor(outputTopic, kafkaConfig));
            } else if (!outputTopic.isPresent() && monitorTopic.isPresent()) {
                System.out.println("Monitoring " + monitorTopic.get() + " every " + monitorDelta + " ms");
                monitors.add(new EPSThroughputWrittenMonitor(monitorTopic, kafkaConfig));
            } else if (!outputTopic.isPresent() && !monitorTopic.isPresent()) {
                System.out.println(
                        "You have not specified an output topic or a monitoring topic, so I have nothing to do here.");
            }
            int lookback = (int) evaluatedArgs.get(LoadOptions.SUMMARY_LOOKBACK).get();
            if (lookback > 0) {
                System.out.println("Summarizing over the last " + lookback + " monitoring periods ("
                        + lookback * monitorDelta + "ms)");
            } else {
                System.out.println("Turning off summarization.");
            }
            final CSVWriter csvWriter = new CSVWriter((File) evaluatedArgs.get(LoadOptions.CSV).orElse(null));
            Writer writer = new Writer(monitors, lookback, new ArrayList<Consumer<Writable>>() {
                {
                    add(new ConsoleWriter());
                    add(csvWriter);
                }
            });
            timer.scheduleAtFixedRate(new MonitorTask(writer), 0, monitorDelta);
            Optional<Object> timeLimit = evaluatedArgs.get(LoadOptions.TIME_LIMIT);
            if (timeLimit.isPresent()) {
                System.out.println("Ending in " + timeLimit.get() + " ms.");
                timer.schedule(new TimerTask() {
                    @Override
                    public void run() {
                        timer.cancel();
                        long durationS = (System.currentTimeMillis() - startTimeMs) / 1000;
                        System.out.println("\nGenerated " + numSent.get() + " in " + durationS + " seconds.");
                        csvWriter.close();
                        System.exit(0);
                    }
                }

                        , (Long) timeLimit.get());
            }
        }
    }
}