| /** |
| * 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 kafka.producer |
| |
| import async._ |
| import java.util.Properties |
| import kafka.serializer.Encoder |
| import java.util.concurrent.{ConcurrentMap, ConcurrentHashMap} |
| import kafka.cluster.{Partition, Broker} |
| import kafka.api.ProducerRequest |
| import kafka.common.{UnavailableProducerException, InvalidConfigException} |
| import kafka.utils.{Utils, Logging} |
| import kafka.message.{NoCompressionCodec, ByteBufferMessageSet} |
| |
| class ProducerPool[V](private val config: ProducerConfig, |
| private val serializer: Encoder[V], |
| private val syncProducers: ConcurrentMap[Int, SyncProducer], |
| private val asyncProducers: ConcurrentMap[Int, AsyncProducer[V]], |
| private val inputEventHandler: EventHandler[V] = null, |
| private val cbkHandler: CallbackHandler[V] = null) extends Logging { |
| |
| private var eventHandler = inputEventHandler |
| if(eventHandler == null) |
| eventHandler = new DefaultEventHandler(config, cbkHandler) |
| |
| if(serializer == null) |
| throw new InvalidConfigException("serializer passed in is null!") |
| |
| private var sync: Boolean = true |
| config.producerType match { |
| case "sync" => |
| case "async" => sync = false |
| case _ => throw new InvalidConfigException("Valid values for producer.type are sync/async") |
| } |
| |
| def this(config: ProducerConfig, serializer: Encoder[V], |
| eventHandler: EventHandler[V], cbkHandler: CallbackHandler[V]) = |
| this(config, serializer, |
| new ConcurrentHashMap[Int, SyncProducer](), |
| new ConcurrentHashMap[Int, AsyncProducer[V]](), |
| eventHandler, cbkHandler) |
| |
| def this(config: ProducerConfig, serializer: Encoder[V]) = this(config, serializer, |
| new ConcurrentHashMap[Int, SyncProducer](), |
| new ConcurrentHashMap[Int, AsyncProducer[V]](), |
| Utils.getObject(config.eventHandler), |
| Utils.getObject(config.cbkHandler)) |
| /** |
| * add a new producer, either synchronous or asynchronous, connecting |
| * to the specified broker |
| * @param bid the id of the broker |
| * @param host the hostname of the broker |
| * @param port the port of the broker |
| */ |
| def addProducer(broker: Broker) { |
| val props = new Properties() |
| props.put("host", broker.host) |
| props.put("port", broker.port.toString) |
| props.putAll(config.props) |
| if(sync) { |
| val producer = new SyncProducer(new SyncProducerConfig(props)) |
| info("Creating sync producer for broker id = " + broker.id + " at " + broker.host + ":" + broker.port) |
| syncProducers.put(broker.id, producer) |
| } else { |
| val producer = new AsyncProducer[V](new AsyncProducerConfig(props), |
| new SyncProducer(new SyncProducerConfig(props)), |
| serializer, |
| eventHandler, config.eventHandlerProps, |
| cbkHandler, config.cbkHandlerProps) |
| producer.start |
| 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 poolData the producer pool request object |
| */ |
| def send(poolData: ProducerPoolData[V]*) { |
| val distinctBrokers = poolData.map(pd => pd.getBidPid.brokerId).distinct |
| var remainingRequests = poolData.toSeq |
| distinctBrokers.foreach { bid => |
| val requestsForThisBid = remainingRequests partition (_.getBidPid.brokerId == bid) |
| remainingRequests = requestsForThisBid._2 |
| |
| if(sync) { |
| val producerRequests = requestsForThisBid._1.map(req => new ProducerRequest(req.getTopic, req.getBidPid.partId, |
| new ByteBufferMessageSet(compressionCodec = config.compressionCodec, |
| messages = req.getData.map(d => serializer.toMessage(d)): _*))) |
| debug("Fetching sync producer for broker id: " + bid) |
| val producer = syncProducers.get(bid) |
| if(producer != null) { |
| if(producerRequests.size > 1) |
| producer.multiSend(producerRequests.toArray) |
| else |
| producer.send(topic = producerRequests(0).topic, |
| partition = producerRequests(0).partition, |
| messages = producerRequests(0).messages) |
| config.compressionCodec match { |
| case NoCompressionCodec => debug("Sending message to broker " + bid) |
| case _ => debug("Sending compressed messages to broker " + bid) |
| } |
| }else |
| throw new UnavailableProducerException("Producer pool has not been initialized correctly. " + |
| "Sync Producer for broker " + bid + " does not exist in the pool") |
| }else { |
| debug("Fetching async producer for broker id: " + bid) |
| val producer = asyncProducers.get(bid) |
| if(producer != null) { |
| requestsForThisBid._1.foreach { req => |
| req.getData.foreach(d => producer.send(req.getTopic, d, req.getBidPid.partId)) |
| } |
| if(logger.isDebugEnabled) |
| config.compressionCodec match { |
| case NoCompressionCodec => debug("Sending message") |
| case _ => debug("Sending compressed messages") |
| } |
| } |
| else |
| throw new UnavailableProducerException("Producer pool has not been initialized correctly. " + |
| "Async Producer for broker " + bid + " does not exist in the pool") |
| } |
| } |
| } |
| |
| /** |
| * Closes all the producers in the pool |
| */ |
| def close() = { |
| config.producerType match { |
| case "sync" => |
| info("Closing all sync producers") |
| val iter = syncProducers.values.iterator |
| while(iter.hasNext) |
| iter.next.close |
| case "async" => |
| info("Closing all async producers") |
| val iter = asyncProducers.values.iterator |
| while(iter.hasNext) |
| iter.next.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 |
| */ |
| def getProducerPoolData(topic: String, bidPid: Partition, data: Seq[V]): ProducerPoolData[V] = { |
| new ProducerPoolData[V](topic, bidPid, data) |
| } |
| |
| class ProducerPoolData[V](topic: String, |
| bidPid: Partition, |
| data: Seq[V]) { |
| def getTopic: String = topic |
| def getBidPid: Partition = bidPid |
| def getData: Seq[V] = data |
| } |
| } |
