content: /*
   * 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.kafka.trogdor.workload;
  
  import org.apache.kafka.trogdor.task.TaskController;
  import org.apache.kafka.trogdor.task.TaskSpec;
  import org.apache.kafka.trogdor.task.TaskWorker;
  
  import com.fasterxml.jackson.annotation.JsonCreator;
  import com.fasterxml.jackson.annotation.JsonProperty;
  
  import java.util.Collections;
  import java.util.Map;
  import java.util.Optional;
  
  /**
   * This is the spec to pass in to be able to run the `ConfigurableProducerWorker` workload.  This allows for customized
   * and even variable configurations in terms of messages per second, message size, batch size, key size, and even the
   * ability to target a specific partition out of a topic.
   *
   * This has several notable differences from the ProduceBench classes, namely the ability to dynamically control
   * flushing and throughput through configurable classes, but also the ability to run against specific partitions within
   * a topic directly.  This workload can only run against one topic at a time, unlike the ProduceBench workload.
   *
   * The parameters that differ from ProduceBenchSpec:
   *
   *     `flushGenerator` -      Used to instruct the KafkaProducer when to issue flushes.  This allows us to simulate
   *                             variable batching since batch flushing is not currently exposed within the KafkaProducer
   *                             class.  See the `FlushGenerator` interface for more information.
   *
   *     `throughputGenerator` - Used to throttle the ConfigurableProducerWorker based on a calculated number of messages
   *                             within a window.  See the `ThroughputGenerator` interface for more information.
   *
   *     `activeTopic`         - This class only supports execution against a single topic at a time.  If more than one
   *                             topic is specified, the ConfigurableProducerWorker will throw an error.
   *
   *     `activePartition`     - Specify a specific partition number within the activeTopic to run load against, or
   *                             specify `-1` to allow use of all partitions.
   *
   * Here is an example spec:
   *
   * {
   *     "startMs": 1606949497662,
   *     "durationMs": 3600000,
   *     "producerNode": "trogdor-agent-0",
   *     "bootstrapServers": "some.example.kafka.server:9091",
   *     "flushGenerator": {
   *         "type": "gaussian",
   *         "messagesPerFlushAverage": 16,
   *         "messagesPerFlushDeviation": 4
   *     },
   *     "throughputGenerator": {
   *         "type": "gaussian",
   *         "messagesPerSecondAverage": 500,
   *         "messagesPerSecondDeviation": 50,
   *         "windowsUntilRateChange": 100,
   *         "windowSizeMs": 100
   *     },
   *     "keyGenerator": {
   *         "type": "constant",
   *         "size": 8
   *     },
   *     "valueGenerator": {
   *         "type": "gaussianTimestampRandom",
   *         "messageSizeAverage": 512,
   *         "messageSizeDeviation": 100,
   *         "messagesUntilSizeChange": 100
   *     },
   *     "producerConf": {
   *         "acks": "all"
   *     },
   *     "commonClientConf": {},
   *     "adminClientConf": {},
   *     "activeTopic": {
   *         "topic0": {
   *             "numPartitions": 100,
   *             "replicationFactor": 3,
   *             "configs": {
   *                 "retention.ms": "1800000"
   *             }
   *         }
   *     },
   *     "activePartition": 5
   * }
   *
   * This example spec performed the following:
   *
   *   * Ran on `trogdor-agent-0` for 1 hour starting at 2020-12-02 22:51:37.662 GMT
   *   * Produced with acks=all to Partition 5 of `topic0` on kafka server `some.example.kafka.server:9091`.
   *   * The average batch had 16 messages, with a standard deviation of 4 messages.
   *   * The messages had 8-byte constant keys with an average size of 512 bytes and a standard deviation of 100 bytes.
   *   * The messages had millisecond timestamps embedded in the first several bytes of the value.
   *   * The average throughput was 500 messages/second, with a window of 100ms and a deviation of 50 messages/second.
   */
  
  public final class ConfigurableProducerSpec extends TaskSpec {
      private final String producerNode;
      private final String bootstrapServers;
      private final Optional<FlushGenerator> flushGenerator;
      private final ThroughputGenerator throughputGenerator;
      private final PayloadGenerator keyGenerator;
      private final PayloadGenerator valueGenerator;
      private final Map<String, String> producerConf;
      private final Map<String, String> adminClientConf;
      private final Map<String, String> commonClientConf;
      private final TopicsSpec activeTopic;
      private final int activePartition;
  
      @JsonCreator
      public ConfigurableProducerSpec(@JsonProperty("startMs") long startMs,
                                      @JsonProperty("durationMs") long durationMs,
                                      @JsonProperty("producerNode") String producerNode,
                                      @JsonProperty("bootstrapServers") String bootstrapServers,
                                      @JsonProperty("flushGenerator") Optional<FlushGenerator> flushGenerator,
                                      @JsonProperty("throughputGenerator") ThroughputGenerator throughputGenerator,
                                      @JsonProperty("keyGenerator") PayloadGenerator keyGenerator,
                                      @JsonProperty("valueGenerator") PayloadGenerator valueGenerator,
                                      @JsonProperty("producerConf") Map<String, String> producerConf,
                                      @JsonProperty("commonClientConf") Map<String, String> commonClientConf,
                                      @JsonProperty("adminClientConf") Map<String, String> adminClientConf,
                                      @JsonProperty("activeTopic") TopicsSpec activeTopic,
                                      @JsonProperty("activePartition") int activePartition) {
          super(startMs, durationMs);
          this.producerNode = (producerNode == null) ? "" : producerNode;
          this.bootstrapServers = (bootstrapServers == null) ? "" : bootstrapServers;
          this.flushGenerator = flushGenerator;
          this.keyGenerator = keyGenerator;
          this.valueGenerator = valueGenerator;
          this.throughputGenerator = throughputGenerator;
          this.producerConf = configOrEmptyMap(producerConf);
          this.commonClientConf = configOrEmptyMap(commonClientConf);
          this.adminClientConf = configOrEmptyMap(adminClientConf);
          this.activeTopic = activeTopic.immutableCopy();
          this.activePartition = activePartition;
      }
  
      @JsonProperty
      public String producerNode() {
          return producerNode;
      }
  
      @JsonProperty
      public String bootstrapServers() {
          return bootstrapServers;
      }
  
      @JsonProperty
      public Optional<FlushGenerator> flushGenerator() {
          return flushGenerator;
      }
  
      @JsonProperty
      public PayloadGenerator keyGenerator() {
          return keyGenerator;
      }
  
      @JsonProperty
      public PayloadGenerator valueGenerator() {
          return valueGenerator;
      }
  
      @JsonProperty
      public ThroughputGenerator throughputGenerator() {
          return throughputGenerator;
      }
  
      @JsonProperty
      public Map<String, String> producerConf() {
          return producerConf;
      }
  
      @JsonProperty
      public Map<String, String> commonClientConf() {
          return commonClientConf;
      }
  
      @JsonProperty
      public Map<String, String> adminClientConf() {
          return adminClientConf;
      }
  
      @JsonProperty
      public TopicsSpec activeTopic() {
          return activeTopic;
      }
  
      @JsonProperty
      public int activePartition() {
          return activePartition;
      }
  
      @Override
      public TaskController newController(String id) {
          return topology -> Collections.singleton(producerNode);
      }
  
      @Override
      public TaskWorker newTaskWorker(String id) {
          return new ConfigurableProducerWorker(id, this);
      }
  }
  