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

@@ -34,6 +34,7 @@
 #include <iosfwd>
 #include <algorithm>
 #include <vector>
+#include <set>
 #include <librdkafka/rdkafka.h>
 #include "macros.h"
 
@@ -54,6 +55,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 find_matches(const TopicPartitionList& partitions,
+                                             const std::set<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 find_matches(const TopicPartitionList& partitions,
+                                             const std::set<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,15 @@
  */
 
 #include <iostream>
+#include <string>
 #include "topic_partition_list.h"
 #include "topic_partition.h"
 #include "exceptions.h"
 
 using std::vector;
+using std::set;
 using std::ostream;
+using std::string;
 
 namespace cppkafka {
 
@@ -67,6 +70,40 @@ TopicPartitionsListPtr make_handle(rd_kafka_topic_partition_list_t* handle) {
     return TopicPartitionsListPtr(handle, &rd_kafka_topic_partition_list_destroy);
 }
 
+TopicPartitionList find_matches(const TopicPartitionList& partitions,
+                                const set<string>& topics) {
+    TopicPartitionList subset;
+    for (const auto& partition : partitions) {
+        for (const auto& topic : topics) {
+            if (topic.size() == partition.get_topic().size()) {
+                // compare both strings
+                bool match = equal(topic.begin(), topic.end(), partition.get_topic().begin(),
+                                   [](char c1, char c2)->bool {
+                    return toupper(c1) == toupper(c2);
+                });
+                if (match) {
+                    subset.emplace_back(partition);
+                }
+            }
+        }
+    }
+    return subset;
+}
+
+TopicPartitionList find_matches(const TopicPartitionList& partitions,
+                                const set<int>& ids) {
+    TopicPartitionList subset;
+    for (const auto& partition : partitions) {
+        for (const auto& id : ids) {
+            // 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) {

tests/topic_partition_list_test.cpp

@@ -4,6 +4,8 @@
 #include "cppkafka/topic_partition.h"
 
 using std::ostringstream;
+using std::set;
+using std::string;
 
 using namespace cppkafka;
 
@@ -42,3 +44,38 @@ TEST_CASE("topic partition list to string", "[topic_partition]") {
     output << list;
     CHECK(output.str() == "[ foo[-1:#], bar[2:#], foobar[3:4] ]");
 }
+
+TEST_CASE("find matches by topic", "[topic_partition]") {
+    TopicPartitionList list;
+    list.push_back({ "foo", 0 });
+    list.push_back({ "bar", 3 });
+    list.push_back({ "fb",  1 });
+    list.push_back({ "foo", 1 });
+    list.push_back({ "fb",  2 });
+    list.push_back({ "other", 1 });
+    list.push_back({ "a",   1 });
+
+    TopicPartitionList subset = find_matches(list, set<string>{"foo", "fb"});
+    REQUIRE(subset.size() == 4);
+    CHECK((subset[0].get_topic() == "foo" && subset[0].get_partition() == 0));
+    CHECK((subset[1].get_topic() == "fb" && subset[1].get_partition() == 1));
+    CHECK((subset[2].get_topic() == "foo" && subset[2].get_partition() == 1));
+    CHECK((subset[3].get_topic() == "fb" && subset[3].get_partition() == 2));
+}
+
+TEST_CASE("find matches by id", "[topic_partition]") {
+    TopicPartitionList list;
+    list.push_back({ "foo", 2 });
+    list.push_back({ "foo", 3 });
+    list.push_back({ "foo", 4 });
+    list.push_back({ "foo", 5 });
+    list.push_back({ "foo", 6 });
+    list.push_back({ "foo", 7 });
+    list.push_back({ "foo", 8 });
+
+    TopicPartitionList subset = find_matches(list, set<int>{2,5,8});
+    REQUIRE(subset.size() == 3);
+    CHECK(subset[0].get_partition() == 2);
+    CHECK(subset[1].get_partition() == 5);
+    CHECK(subset[2].get_partition() == 8);
+}