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 io.jafka.producer;
  
  
  import io.jafka.api.ProducerRequest;
  import io.jafka.cluster.Broker;
  import io.jafka.cluster.Partition;
  import io.jafka.common.InvalidPartitionException;
  import io.jafka.common.NoBrokersForPartitionException;
  import io.jafka.common.annotations.ClientSide;
  import io.jafka.producer.async.CallbackHandler;
  import io.jafka.producer.async.EventHandler;
  import io.jafka.producer.serializer.Encoder;
  import io.jafka.utils.Utils;
  import io.jafka.utils.ZKConfig;
  import io.jafka.utils.Closer;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.List;
  import java.util.Map;
  import java.util.Properties;
  import java.util.Random;
  import java.util.SortedSet;
  import java.util.concurrent.atomic.AtomicBoolean;
  
  import static io.jafka.utils.Closer.closeQuietly;
  
  /**
   * Message producer
   *
   * @author adyliu (imxylz@gmail.com)
   * @since 1.0
   */
  @ClientSide
  public class Producer<K, V> implements BrokerPartitionInfo.Callback, IProducer<K, V> {
  
      ProducerConfig config;
  
      private Partitioner<K> partitioner;
  
      ProducerPool<V> producerPool;
  
      boolean populateProducerPool;
  
      BrokerPartitionInfo brokerPartitionInfo;
  
      private final Logger logger = LoggerFactory.getLogger(Producer.class);
  
      /////////////////////////////////////////////////////////////////////////
      private final AtomicBoolean hasShutdown = new AtomicBoolean(false);
  
      private final Random random = new Random();
  
      private final boolean zkEnabled;
      private Encoder<V> encoder;
  
      public Producer(ProducerConfig config, Partitioner<K> partitioner, ProducerPool<V> producerPool, boolean populateProducerPool,
                      BrokerPartitionInfo brokerPartitionInfo) {
          super();
          this.config = config;
          this.partitioner = partitioner;
          if (producerPool == null) {
              producerPool = new ProducerPool<V>(config, getEncoder());
          }
          this.producerPool = producerPool;
          this.populateProducerPool = populateProducerPool;
          this.brokerPartitionInfo = brokerPartitionInfo;
          //
          this.zkEnabled = config.getZkConnect() != null;
          if (this.brokerPartitionInfo == null) {
              if (this.zkEnabled) {
                  Properties zkProps = new Properties();
                  zkProps.put("zk.connect", config.getZkConnect());
                  zkProps.put("zk.sessiontimeout.ms", "" + config.getZkSessionTimeoutMs());
                  zkProps.put("zk.connectiontimeout.ms", "" + config.getZkConnectionTimeoutMs());
                  zkProps.put("zk.synctime.ms", "" + config.getZkSyncTimeMs());
                  this.brokerPartitionInfo = new ZKBrokerPartitionInfo(new ZKConfig(zkProps), this);
              } else {
                  this.brokerPartitionInfo = new ConfigBrokerPartitionInfo(config);
              }
          }
          //
          // pool of producers, one per broker
          if (this.populateProducerPool) {
              for (Map.Entry<Integer, Broker> e : this.brokerPartitionInfo.getAllBrokerInfo().entrySet()) {
                  Broker b = e.getValue();
                  producerPool.addProducer(new Broker(e.getKey(), b.host, b.host, b.port,b.autocreated));
              }
          }
      }
  
      /**
       * This constructor can be used when all config parameters will be
       * specified through the ProducerConfig object
       *
       * @param config Producer Configuration object
       */
      public Producer(ProducerConfig config) {
          this(config, //
                  null,//
                  null, //
                  true, //
                  null);
      }
  
      /**
       * This constructor can be used to provide pre-instantiated objects for
       * all config parameters that would otherwise be instantiated via
       * reflection. i.e. encoder, partitioner, event handler and callback
       * handler. If you use this constructor, encoder, eventHandler,
       * callback handler and partitioner will not be picked up from the
       * config.
       *
       * @param config       Producer Configuration object
       * @param encoder      Encoder used to convert an object of type V to a
       *                     jafka.message.Message. If this is null it throws an
       *                     InvalidConfigException
       * @param eventHandler the class that implements
       *                     jafka.producer.async.IEventHandler[T] used to dispatch a
       *                     batch of produce requests, using an instance of
       *                     jafka.producer.SyncProducer. If this is null, it uses the
       *                     DefaultEventHandler
       * @param cbkHandler   the class that implements
       *                     jafka.producer.async.CallbackHandler[T] used to inject
       *                     callbacks at various stages of the
       *                     jafka.producer.AsyncProducer pipeline. If this is null, the
       *                     producer does not use the callback handler and hence does not
       *                     invoke any callbacks
       * @param partitioner  class that implements the
       *                     jafka.producer.Partitioner[K], used to supply a custom
       *                     partitioning strategy on the message key (of type K) that is
       *                     specified through the ProducerData[K, T] object in the send
       *                     API. If this is null, producer uses DefaultPartitioner
       */
      public Producer(ProducerConfig config, Encoder<V> encoder, EventHandler<V> eventHandler, CallbackHandler<V> cbkHandler, Partitioner<K> partitioner) {
          this(config, //
                  partitioner,//
                  new ProducerPool<V>(config, encoder, eventHandler, cbkHandler), //
                  true, //
                  null);
      }
  
      @SuppressWarnings("unchecked")
      @Override
      public Encoder<V> getEncoder() {
          return encoder == null ? (Encoder<V>) Utils.getObject(config.getSerializerClass()) : encoder;
      }
  
      public void send(ProducerData<K, V> data) throws NoBrokersForPartitionException, InvalidPartitionException {
          if (data == null) return;
          if (zkEnabled) {
              zkSend(data);
          } else {
              configSend(data);
          }
      }
  
  
      private void configSend(ProducerData<K, V> data) {
          producerPool.send(create(data));
      }
  
  
      private void zkSend(ProducerData<K, V> data) {
          int numRetries = 0;
          Broker brokerInfoOpt = null;
          Partition brokerIdPartition = null;
          while (numRetries <= config.getZkReadRetries() && brokerInfoOpt == null) {
              if (numRetries > 0) {
                  logger.info("Try #" + numRetries + " ZK producer cache is stale. Refreshing it by reading from ZK again");
                  brokerPartitionInfo.updateInfo();
              }
              List<Partition> partitions = new ArrayList<Partition>(getPartitionListForTopic(data));
              brokerIdPartition = partitions.get(getPartition(data.getKey(), partitions.size()));
              if (brokerIdPartition != null) {
                  brokerInfoOpt = brokerPartitionInfo.getBrokerInfo(brokerIdPartition.brokerId);
              }
              numRetries++;
          }
          if (brokerInfoOpt == null) {
              throw new NoBrokersForPartitionException("Invalid Zookeeper state. Failed to get partition for topic: " + data.getTopic() + " and key: "
                      + data.getKey());
          }
          //
          ProducerPoolData<V> ppd = producerPool.getProducerPoolData(data.getTopic(),//
                  new Partition(brokerIdPartition.brokerId, brokerIdPartition.partId),//
                  data.getData());
          producerPool.send(ppd);
      }
  
      private int getPartition(K key, int numPartitions) {
          if (numPartitions <= 0) {
              throw new InvalidPartitionException("Invalid number of partitions: " + numPartitions + "\n Valid values are > 0");
          }
          int partition = key == null ? random.nextInt(numPartitions) : getPartitioner().partition(key, numPartitions);
          if (partition < 0 || partition >= numPartitions) {
              throw new InvalidPartitionException("Invalid partition id : " + partition + "\n Valid values are in the range inclusive [0, " + (numPartitions - 1)
                      + "]");
          }
          return partition;
      }
  
      public void producerCbk(int bid, String host, int port,boolean autocreated) {
          if (populateProducerPool) {
              producerPool.addProducer(new Broker(bid, host, host, port,autocreated));
          } else {
              logger.debug("Skipping the callback since populateProducerPool = false");
          }
      }
  
      private ProducerPoolData<V> create(ProducerData<K, V> pd) {
          Collection<Partition> topicPartitionsList = getPartitionListForTopic(pd);
          //FIXME: random Broker???
          int randomBrokerId = random.nextInt(topicPartitionsList.size());
          final Partition brokerIdPartition = new ArrayList<Partition>(topicPartitionsList).get(randomBrokerId);
          return this.producerPool.getProducerPoolData(pd.getTopic(),//
                  new Partition(brokerIdPartition.brokerId, ProducerRequest.RandomPartition), pd.getData());
      }
  
      private Collection<Partition> getPartitionListForTopic(ProducerData<K, V> pd) {
          SortedSet<Partition> topicPartitionsList = brokerPartitionInfo.getBrokerPartitionInfo(pd.getTopic());
          if (topicPartitionsList.size() == 0) {
              throw new NoBrokersForPartitionException("Partition= " + pd.getTopic());
          }
          return topicPartitionsList;
      }
  
      public void close() {
          if (hasShutdown.compareAndSet(false, true)) {
              Closer.closeQuietly(producerPool);
              Closer.closeQuietly(brokerPartitionInfo);
          }
      }
  
      @SuppressWarnings("unchecked")
      @Override
      public Partitioner<K> getPartitioner() {
          if (partitioner == null) {
              partitioner = (Partitioner<K>) Utils.getObject(config.getPartitionerClass());
          }
          return partitioner;
      }
  
      public void setPartitioner(Partitioner<K> partitioner) {
          this.partitioner = partitioner;
      }
  }
  