ProducerPool.java

/**
 * 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 java.io.Closeable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;


import io.jafka.api.ProducerRequest;
import io.jafka.cluster.Broker;
import io.jafka.cluster.Partition;
import io.jafka.common.InvalidConfigException;
import io.jafka.common.UnavailableProducerException;
import io.jafka.common.annotations.ClientSide;
import io.jafka.message.ByteBufferMessageSet;
import io.jafka.message.Message;
import io.jafka.producer.async.AsyncProducer;
import io.jafka.producer.async.AsyncProducerConfig;
import io.jafka.producer.async.CallbackHandler;
import io.jafka.producer.async.DefaultEventHandler;
import io.jafka.producer.async.EventHandler;
import io.jafka.producer.serializer.Encoder;
import io.jafka.utils.Utils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author adyliu (imxylz@gmail.com)
 * @since 1.0
 */
@ClientSide
public class ProducerPool<V> implements Closeable {

    private final ProducerConfig config;

    private final Encoder<V> serializer;

    private final ConcurrentMap<Integer, SyncProducer> syncProducers;

    private final ConcurrentMap<Integer, AsyncProducer<V>> asyncProducers;

    private final EventHandler<V> eventHandler;

    private final CallbackHandler<V> callbackHandler;

    private boolean sync = true;

    private final Logger logger = LoggerFactory.getLogger(ProducerPool.class);

    public ProducerPool(ProducerConfig config,//
            Encoder<V> serializer, //
            ConcurrentMap<Integer, SyncProducer> syncProducers,//
            ConcurrentMap<Integer, AsyncProducer<V>> asyncProducers, //
            EventHandler<V> inputEventHandler, //
            CallbackHandler<V> callbackHandler) {
        super();
        this.config = config;
        this.serializer = serializer;
        this.syncProducers = syncProducers;
        this.asyncProducers = asyncProducers;
        this.eventHandler = inputEventHandler != null ? inputEventHandler : new DefaultEventHandler<V>(config, callbackHandler);
        this.callbackHandler = callbackHandler;
        if (serializer == null) {
            throw new InvalidConfigException("serializer passed in is null!");
        }
        this.sync = !"async".equalsIgnoreCase(config.getProducerType());
    }

    public ProducerPool(ProducerConfig config, Encoder<V> serializer,//
            EventHandler<V> eventHandler,//
            CallbackHandler<V> callbackHandler) {
        this(config,//
                serializer,//
                new ConcurrentHashMap<Integer, SyncProducer>(),//
                new ConcurrentHashMap<Integer, AsyncProducer<V>>(),//
                eventHandler,//
                callbackHandler);
    }

    @SuppressWarnings("unchecked")
    public ProducerPool(ProducerConfig config, Encoder<V> serializer) {
        this(config,//
                serializer,//
                new ConcurrentHashMap<Integer, SyncProducer>(),//
                new ConcurrentHashMap<Integer, AsyncProducer<V>>(),//
                (EventHandler<V>) Utils.getObject(config.getEventHandler()),//
                (CallbackHandler<V>) Utils.getObject(config.getCbkHandler()));
    }

    /**
     * add a new producer, either synchronous or asynchronous, connecting
     * to the specified broker
     * 
     * @param broker broker to producer
     */
    public void addProducer(Broker broker) {
        Properties props = new Properties();
        props.put("host", broker.host);
        props.put("port", "" + broker.port);
        props.putAll(config.getProperties());
        if (sync) {
            SyncProducer producer = new SyncProducer(new SyncProducerConfig(props));
            logger.info("Creating sync producer for broker id = " + broker.id + " at " + broker.host + ":" + broker.port);
            syncProducers.put(broker.id, producer);
        } else {
            AsyncProducer<V> producer = new AsyncProducer<V>(new AsyncProducerConfig(props),//
                    new SyncProducer(new SyncProducerConfig(props)),//
                    serializer,//
                    eventHandler,//
                    config.getEventHandlerProperties(),//
                    this.callbackHandler, //
                    config.getCbkHandlerProperties());
            producer.start();
            logger.info("Creating async producer for broker id = " + broker.id + " at " + broker.host + ":" + broker.port);
            asyncProducers.put(broker.id, producer);
        }
    }

    /**
     * selects either a synchronous or an asynchronous producer, for the
     * specified broker id and calls the send API on the selected producer
     * to publish the data to the specified broker partition
     * 
     * @param ppd the producer pool request object
     */
    public void send(ProducerPoolData<V> ppd) {
        if (logger.isDebugEnabled()) {
            logger.debug("send message: " + ppd);
        }
        if (sync) {
            Message[] messages = new Message[ppd.data.size()];
            int index = 0;
            for (V v : ppd.data) {
                messages[index] = serializer.toMessage(v);
                index++;
            }
            ByteBufferMessageSet bbms = new ByteBufferMessageSet(config.getCompressionCodec(), messages);
            ProducerRequest request = new ProducerRequest(ppd.topic, ppd.partition.partId, bbms);
            SyncProducer producer = syncProducers.get(ppd.partition.brokerId);
            if (producer == null) {
                throw new UnavailableProducerException("Producer pool has not been initialized correctly. " + "Sync Producer for broker "
                        + ppd.partition.brokerId + " does not exist in the pool");
            }
            producer.send(request.topic, request.partition, request.messages);
        } else {
            AsyncProducer<V> asyncProducer = asyncProducers.get(ppd.partition.brokerId);
            for (V v : ppd.data) {
                asyncProducer.send(ppd.topic, v, ppd.partition.partId);
            }
        }
    }

    public void send(List<ProducerPoolData<V>> poolData) {
        if (sync) {
            syncSend(poolData);
        } else {
            asyncSend(poolData);
        }
    }

    private void asyncSend(List<ProducerPoolData<V>> poolData) {
        for (ProducerPoolData<V> ppd : poolData) {
            AsyncProducer<V> asyncProducer = asyncProducers.get(ppd.partition.brokerId);
            for (V v : ppd.data) {
                asyncProducer.send(ppd.topic, v, ppd.partition.partId);
            }
        }
    }

    private void syncSend(List<ProducerPoolData<V>> poolData) {
        final Map<Integer, List<ProducerRequest>> topicBrokerIdData = new HashMap<Integer, List<ProducerRequest>>();
        for (ProducerPoolData<V> ppd : poolData) {
            List<ProducerRequest> messageSets = topicBrokerIdData.get(ppd.partition.brokerId);
            if (messageSets == null) {
                messageSets = new ArrayList<ProducerRequest>();
                topicBrokerIdData.put(ppd.partition.brokerId, messageSets);
            }
            Message[] messages = new Message[ppd.data.size()];
            int index = 0;
            for (V v : ppd.data) {
                messages[index] = serializer.toMessage(v);
                index++;
            }
            ByteBufferMessageSet bbms = new ByteBufferMessageSet(config.getCompressionCodec(), messages);
            messageSets.add(new ProducerRequest(ppd.topic, ppd.partition.partId, bbms));
        }
        for (Map.Entry<Integer, List<ProducerRequest>> e : topicBrokerIdData.entrySet()) {
            SyncProducer producer = syncProducers.get(e.getKey());
            if (producer == null) {
                throw new UnavailableProducerException("Producer pool has not been initialized correctly. " + "Sync Producer for broker " + e.getKey()
                        + " does not exist in the pool");
            }
            if (e.getValue().size() == 1) {
                ProducerRequest request = e.getValue().get(0);
                producer.send(request.topic, request.partition, request.messages);
            } else {
                producer.multiSend(e.getValue());
            }
        }
    }

    /**
     * Closes all the producers in the pool
     */
    public void close() {
        logger.info("Closing all sync producers");
        if (sync) {
            for (SyncProducer p : syncProducers.values()) {
                p.close();
            }
        } else {
            for (AsyncProducer<V> p : asyncProducers.values()) {
                p.close();
            }
        }

    }

    /**
     * This constructs and returns the request object for the producer pool
     * 
     * @param topic the topic to which the data should be published
     * @param bidPid the broker id and partition id
     * @param data the data to be published
     * @return producer data of builder
     */
    public ProducerPoolData<V> getProducerPoolData(String topic, Partition bidPid, List<V> data) {
        return new ProducerPoolData<V>(topic, bidPid, data);
    }
}