Pause/resume a consumer by topic

include/cppkafka/consumer.h

@@ -193,6 +193,16 @@ class CPPKAFKA_API Consumer : public KafkaHandleBase {
      */ 
     void unassign();
 
+    /**
+     * \brief Pauses all consumption
+     */
+    void pause();
+
+    /**
+     * \brief Resumes all consumption
+     */
+     void resume();
+
     /**
      * \brief Commits the current partition assignment
      *
@@ -364,7 +374,6 @@ class CPPKAFKA_API Consumer : public KafkaHandleBase {
 private:
     static void rebalance_proxy(rd_kafka_t *handle, rd_kafka_resp_err_t error,
                                 rd_kafka_topic_partition_list_t *partitions, void *opaque);
-
     void close();
     void commit(const Message& msg, bool async);
     void commit(const TopicPartitionList* topic_partitions, bool async);

include/cppkafka/topic_partition_list.h

@@ -54,6 +54,16 @@ CPPKAFKA_API TopicPartitionList convert(const TopicPartitionsListPtr& topic_part
 CPPKAFKA_API TopicPartitionList convert(rd_kafka_topic_partition_list_t* topic_partitions);
 CPPKAFKA_API TopicPartitionsListPtr make_handle(rd_kafka_topic_partition_list_t* handle);
 
+// Extracts a partition list subset belonging to the provided topics (case-insensitive)
+CPPKAFKA_API TopicPartitionList make_subset(const TopicPartitionList& partitions,
+                                            const std::vector<std::string>& topics);
+
+// Extracts a partition list subset belonging to the provided partition ids
+// Note: this assumes that all topic partitions in the original list belong to the same topic
+//       otherwise the partition ids may not be unique
+CPPKAFKA_API TopicPartitionList make_subset(const TopicPartitionList& partitions,
+                                            const std::vector<int>& ids);
+
 CPPKAFKA_API std::ostream& operator<<(std::ostream& output, const TopicPartitionList& rhs);
 
 } // cppkafka

src/consumer.cpp

@@ -27,6 +27,8 @@
  *
  */
 #include <sstream>
+#include <algorithm>
+#include <cctype>
 #include "consumer.h"
 #include "exceptions.h"
 #include "logging.h"
@@ -39,6 +41,8 @@ using std::move;
 using std::make_tuple;
 using std::ostringstream;
 using std::chrono::milliseconds;
+using std::toupper;
+using std::equal;
 
 namespace cppkafka {
 
@@ -125,6 +129,14 @@ void Consumer::unassign() {
     check_error(error);
 }
 
+void Consumer::pause() {
+    pause_partitions(get_assignment());
+}
+
+void Consumer::resume() {
+    resume_partitions(get_assignment());
+}
+
 void Consumer::commit() {
     commit(nullptr, false);
 }

src/topic_partition_list.cpp

@@ -28,12 +28,14 @@
  */
 
 #include <iostream>
+#include <string>
 #include "topic_partition_list.h"
 #include "topic_partition.h"
 #include "exceptions.h"
 
 using std::vector;
 using std::ostream;
+using std::string;
 
 namespace cppkafka {
 
@@ -67,6 +69,42 @@ TopicPartitionsListPtr make_handle(rd_kafka_topic_partition_list_t* handle) {
     return TopicPartitionsListPtr(handle, &rd_kafka_topic_partition_list_destroy);
 }
 
+TopicPartitionList make_subset(const TopicPartitionList& partitions,
+                               const vector<string>& topics) {
+    vector<bool> skip(partitions.size(), false);
+    TopicPartitionList subset;
+    for (const auto& topic : topics) {
+        for (size_t i = 0; i < partitions.size(); ++i) {
+            if (!skip[i] && (topic.size() == partitions[i].get_topic().size())) {
+                // compare both strings
+                bool match = equal(topic.begin(), topic.end(), partitions[i].get_topic().begin(),
+                                   [](char c1, char c2)->bool {
+                    return toupper(c1) == toupper(c2);
+                });
+                if (match) {
+                    skip[i] = true;
+                    subset.emplace_back(partitions[i]);
+                }
+            }
+        }
+    }
+    return subset;
+}
+
+TopicPartitionList make_subset(const TopicPartitionList& partitions,
+                               const vector<int>& ids) {
+    TopicPartitionList subset;
+    for (const auto& id : ids) {
+        for (const auto& partition : partitions) {
+            // compare both partition ids
+            if (id == partition.get_partition()) {
+                subset.emplace_back(partition);
+            }
+        }
+    }
+    return subset;
+}
+
 ostream& operator<<(ostream& output, const TopicPartitionList& rhs) {
     output << "[ ";
     for (auto iter = rhs.begin(); iter != rhs.end(); ++iter) {