diff --git a/include/cppkafka/consumer.h b/include/cppkafka/consumer.h index 478cc5d..4241f22 100644 --- a/include/cppkafka/consumer.h +++ b/include/cppkafka/consumer.h @@ -4,6 +4,7 @@ #include #include #include +#include #include "kafka_handle_base.h" #include "message.h" @@ -14,14 +15,22 @@ class TopicConfiguration; class Consumer : public KafkaHandleBase { public: - Consumer(const Configuration& config); + using AssignmentCallback = std::function; + using RevocationCallback = std::function; + using RebalanceErrorCallback = std::function; + + Consumer(Configuration config); void set_timeout(const std::chrono::milliseconds timeout); + void set_assignment_callback(AssignmentCallback callback); + void set_revocation_callback(RevocationCallback callback); + void set_rebalance_error_callback(RebalanceErrorCallback callback); void subscribe(const std::vector& topics); void unsubscribe(); void assign(const TopicPartitionList& topic_partitions); + void unassign(); void close(); void commit(const Message& msg); @@ -38,10 +47,17 @@ public: private: static const std::chrono::milliseconds DEFAULT_TIMEOUT; + static void rebalance_proxy(rd_kafka_t *handle, rd_kafka_resp_err_t error, + rd_kafka_topic_partition_list_t *partitions, void *opaque); + void commit(const Message& msg, bool async); void commit(const TopicPartitionList& topic_partitions, bool async); + void handle_rebalance(rd_kafka_resp_err_t err, const TopicPartitionList& topic_partitions); std::chrono::milliseconds timeout_ms_; + AssignmentCallback assignment_callback_; + RevocationCallback revocation_callback_; + RebalanceErrorCallback rebalance_error_callback_; }; } // cppkafka diff --git a/include/cppkafka/topic_partition_list.h b/include/cppkafka/topic_partition_list.h index 22ae27d..1d58471 100644 --- a/include/cppkafka/topic_partition_list.h +++ b/include/cppkafka/topic_partition_list.h @@ -12,6 +12,8 @@ class TopicPartition; class TopicPartitionList { public: + static TopicPartitionList make_non_owning(rd_kafka_topic_partition_list_t* handle); + TopicPartitionList(); TopicPartitionList(rd_kafka_topic_partition_list_t* handle); TopicPartitionList(size_t size); @@ -40,11 +42,15 @@ public: private: static const size_t DEFAULT_CONTAINER_SIZE; + struct NonOwningTag { }; + using HandlePtr = std::unique_ptr; static HandlePtr make_handle(rd_kafka_topic_partition_list_t* ptr); + TopicPartitionList(rd_kafka_topic_partition_list_t* handle, NonOwningTag); + rd_kafka_topic_partition_t* get_topic_partition(const TopicPartition& topic_partition) const; HandlePtr handle_; diff --git a/src/consumer.cpp b/src/consumer.cpp index e5ea373..13f1cba 100644 --- a/src/consumer.cpp +++ b/src/consumer.cpp @@ -12,9 +12,17 @@ namespace cppkafka { const milliseconds Consumer::DEFAULT_TIMEOUT{1000}; -Consumer::Consumer(const Configuration& config) +void Consumer::rebalance_proxy(rd_kafka_t*, rd_kafka_resp_err_t error, + rd_kafka_topic_partition_list_t *partitions, void *opaque) { + TopicPartitionList list = TopicPartitionList::make_non_owning(partitions); + static_cast(opaque)->handle_rebalance(error, list); +} + +Consumer::Consumer(Configuration config) : timeout_ms_(DEFAULT_TIMEOUT) { char error_buffer[512]; + // Set ourselves as the opaque pointer + rd_kafka_conf_set_opaque(config.get_handle(), this); rd_kafka_t* ptr = rd_kafka_new(RD_KAFKA_CONSUMER, config.get_handle(), error_buffer, sizeof(error_buffer)); if (!ptr) { @@ -27,6 +35,18 @@ void Consumer::set_timeout(const milliseconds timeout) { timeout_ms_ = timeout; } +void Consumer::set_assignment_callback(AssignmentCallback callback) { + assignment_callback_ = move(callback); +} + +void Consumer::set_revocation_callback(RevocationCallback callback) { + revocation_callback_ = move(callback); +} + +void Consumer::set_rebalance_error_callback(RebalanceErrorCallback callback) { + rebalance_error_callback_ = move(callback); +} + void Consumer::subscribe(const vector& topics) { TopicPartitionList list(topics.begin(), topics.end()); rd_kafka_resp_err_t error = rd_kafka_subscribe(get_handle(), list.get_handle()); @@ -45,6 +65,11 @@ void Consumer::assign(const TopicPartitionList& topic_partitions) { check_error(error); } +void Consumer::unassign() { + rd_kafka_resp_err_t error = rd_kafka_assign(get_handle(), nullptr); + check_error(error); +} + void Consumer::close() { rd_kafka_resp_err_t error = rd_kafka_consumer_close(get_handle()); check_error(error); @@ -117,4 +142,26 @@ void Consumer::commit(const TopicPartitionList& topic_partitions, bool async) { check_error(error); } +void Consumer::handle_rebalance(rd_kafka_resp_err_t error, + const TopicPartitionList& topic_partitions) { + if (error == RD_KAFKA_RESP_ERR__ASSIGN_PARTITIONS) { + if (assignment_callback_) { + assignment_callback_(topic_partitions); + } + assign(topic_partitions); + } + else if (error == RD_KAFKA_RESP_ERR__REVOKE_PARTITIONS) { + if (revocation_callback_) { + revocation_callback_(topic_partitions); + } + unassign(); + } + else { + if (rebalance_error_callback_) { + rebalance_error_callback_(error); + } + unassign(); + } +} + } // cppkafka diff --git a/src/topic_partition_list.cpp b/src/topic_partition_list.cpp index 0981410..bec78e3 100644 --- a/src/topic_partition_list.cpp +++ b/src/topic_partition_list.cpp @@ -6,6 +6,15 @@ namespace cppkafka { const size_t TopicPartitionList::DEFAULT_CONTAINER_SIZE = 5; +void dummy_deleter(rd_kafka_topic_partition_list_t*) { + +} + +TopicPartitionList +TopicPartitionList::make_non_owning(rd_kafka_topic_partition_list_t* handle) { + return TopicPartitionList(handle, NonOwningTag()); +} + TopicPartitionList::TopicPartitionList() : TopicPartitionList(DEFAULT_CONTAINER_SIZE) { @@ -16,6 +25,12 @@ TopicPartitionList::TopicPartitionList(rd_kafka_topic_partition_list_t* handle) } +TopicPartitionList::TopicPartitionList(rd_kafka_topic_partition_list_t* handle, + NonOwningTag) +: handle_(handle, &dummy_deleter) { + +} + TopicPartitionList::TopicPartitionList(size_t size) : handle_(make_handle(rd_kafka_topic_partition_list_new(size))) {