package com.github.projectflink.streaming;

import com.dataartisans.flink.example.eventpattern.Event;
import com.dataartisans.flink.example.eventpattern.EventsGenerator;
import com.dataartisans.flink.example.eventpattern.StateMachineMapper;
import com.dataartisans.flink.example.eventpattern.kafka.EventDeSerializer;
import com.github.projectflink.streaming.utils.PimpedKafkaSink;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.utils.ParameterTool;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.flink.streaming.connectors.kafka.api.config.PartitionerWrapper;
import org.apache.flink.streaming.connectors.kafka.partitioner.SerializableKafkaPartitioner;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Properties;


public class KafkaGenerator {

	private static final Logger LOG = LoggerFactory.getLogger(KafkaGenerator.class);

	public static void main(String[] args) throws Exception {
		StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
		final ParameterTool pt = ParameterTool.fromArgs(args);
		see.getConfig().setGlobalJobParameters(pt);
		if(pt.has("p")) {
			see.setParallelism(pt.getInt("p"));
		}

		DataStreamSource<Event> src = see.addSource(new RichParallelSourceFunction<Event>() {

			int min;
			int max;
			EventsGenerator eg;
			boolean running = true;

			@Override
			public void open(Configuration parameters) throws Exception {
				super.open(parameters);

				int range = Integer.MAX_VALUE / getRuntimeContext().getNumberOfParallelSubtasks();
				this.min = range * getRuntimeContext().getIndexOfThisSubtask();
				this.max = min + range;
				eg = new EventsGenerator();
				LOG.info("Creating new EventsGenerator from {} to {}", min, max);
			}


			@Override
			public void run(SourceContext<Event> sourceContext) throws Exception {
				long generated = 0;
				long logfreq = pt.getInt("logfreq");
				long lastLog = -1;
				long lastElements = 0;

				while (running) {
					Event gen = eg.next(min, max);
					sourceContext.collect(gen);
					//	LOG.info("Generated event {}", gen);

					generated++;
					if (generated % logfreq == 0) {
						// throughput over entire time
						long now = System.currentTimeMillis();

						// throughput for the last "logfreq" elements
						if (lastLog == -1) {
							// init (the first)
							lastLog = now;
							lastElements = generated;
						} else {
							long timeDiff = now - lastLog;
							long elementDiff = generated - lastElements;
							double ex = (1000 / (double) timeDiff);
							LOG.info("During the last {} ms, we generated {} elements. That's {} elements/second/core", timeDiff, elementDiff, elementDiff * ex);
							// reinit
							lastLog = now;
							lastElements = generated;
						}
					}

				}
			}

			@Override
			public void cancel() {
				running = false;
			}
		});

		// run the state machine here as well to see illegal transistions
		src.partitionByHash(new KeySelector<Event, Integer>() {
			@Override
			public Integer getKey(Event event) throws Exception {
				return event.sourceAddress();
			}
		}).flatMap(new StateMachineMapper(pt));


		// write stuff into Kafka

		String zkServer = pt.get("zookeeper");
		Properties props = pt.getProperties();

		if(pt.has("localPartitioner")) {
			SerializableKafkaPartitioner part = new PimpedKafkaSink.LocalKafkaPartitioner(zkServer, pt.getRequired("topic"));
			props.put("partitioner.class", PartitionerWrapper.class.getCanonicalName());
			// java serialization will do the rest.
			props.put(PartitionerWrapper.SERIALIZED_WRAPPER_NAME, part);
		}

		src.addSink(new PimpedKafkaSink<Event>(pt.getRequired("brokerList"), pt.getRequired("topic"), props, new EventDeSerializer()));


		see.execute();
	}
}

@@ -2232,25 +2232,273 @@
 %09%09%09%09
-while (running) %7B
+int delay = pt.getInt(%22delay%22);%0A%09%09%09%09int sleepFreq = pt.getInt(%22sleepFreq%22);%0A%0A%09%09%09%09while (running) %7B%0A%09%09%09%09%09if(delay %3E 0) %7B%0A%09%09%09%09%09%09if(generated %25 sleepFreq == 0) %7B%0A%09%09%09%09%09%09%09try %7B Thread.sleep(delay); %7D catch (InterruptedException e) %7B e.printStackTrace();%7D%0A%09%09%09%09%09%09%7D%0A%09%09%09%09%09%7D%0A
 %0A%09%09%09
@@ -3434,17 +3434,16 @@
 l transi
-s
 tions%0A%09%09

Make kafka generator throttleable