Allocators (#118)

* Added allocator support for consumers and buffered producer

* Changed MessageList back to std::vector<Message> for consistency with the allocator API

include/cppkafka/consumer.h

@@ -379,10 +379,14 @@ public:
      * This can return one or more messages
      *
      * \param max_batch_size The maximum amount of messages expected
+     * \param alloc The optionally supplied allocator for allocating messages
      *
      * \return A list of messages
      */
-    MessageList poll_batch(size_t max_batch_size);
+    template <typename Allocator>
+    std::vector<Message, Allocator> poll_batch(size_t max_batch_size,
+                                               const Allocator& alloc);
+    std::vector<Message> poll_batch(size_t max_batch_size);
 
     /**
      * \brief Polls for a batch of messages
@@ -391,10 +395,16 @@ public:
      *
      * \param max_batch_size The maximum amount of messages expected
      * \param timeout The timeout for this operation
+     * \param alloc The optionally supplied allocator for allocating messages
      *
      * \return A list of messages
      */
-    MessageList poll_batch(size_t max_batch_size, std::chrono::milliseconds timeout);
+    template <typename Allocator>
+    std::vector<Message, Allocator> poll_batch(size_t max_batch_size,
+                                               std::chrono::milliseconds timeout,
+                                               const Allocator& alloc);
+    std::vector<Message> poll_batch(size_t max_batch_size,
+                                    std::chrono::milliseconds timeout);
     
     /**
      * \brief Get the global event queue servicing this consumer corresponding to
@@ -430,6 +440,7 @@ public:
 private:
     static void rebalance_proxy(rd_kafka_t *handle, rd_kafka_resp_err_t error,
                                 rd_kafka_topic_partition_list_t *partitions, void *opaque);
+    static Queue get_queue(rd_kafka_queue_t* handle);
     void close();
     void commit(const Message& msg, bool async);
     void commit(const TopicPartitionList* topic_partitions, bool async);
@@ -440,6 +451,30 @@ private:
     RebalanceErrorCallback rebalance_error_callback_;
 };
 
+// Implementations
+template <typename Allocator>
+std::vector<Message, Allocator> Consumer::poll_batch(size_t max_batch_size,
+                                                     const Allocator& alloc) {
+    return poll_batch(max_batch_size, get_timeout(), alloc);
+}
+
+template <typename Allocator>
+std::vector<Message, Allocator> Consumer::poll_batch(size_t max_batch_size,
+                                                     std::chrono::milliseconds timeout,
+                                                     const Allocator& alloc) {
+    std::vector<rd_kafka_message_t*> raw_messages(max_batch_size);
+    // Note that this will leak the queue when using rdkafka < 0.11.5 (see get_queue comment)
+    Queue queue(get_queue(rd_kafka_queue_get_consumer(get_handle())));
+    ssize_t result = rd_kafka_consume_batch_queue(queue.get_handle() , timeout.count(), raw_messages.data(),
+                                                  raw_messages.size());
+    if (result == -1) {
+        check_error(rd_kafka_last_error());
+        // on the off-chance that check_error() does not throw an error
+        return std::vector<Message, Allocator>(alloc);
+    }
+    return std::vector<Message, Allocator>(raw_messages.begin(), raw_messages.begin() + result, alloc);
+}
+
 } // cppkafka
 
 #endif // CPP_KAFKA_CONSUMER_H