Handle (un)assign on consumer mock

mfontanini · mfontanini · commit 8480cf3b1e74 · 2017-10-22T15:31:36.000-07:00
diff --git a/mocking/include/cppkafka/mocking/consumer_mock.h b/mocking/include/cppkafka/mocking/consumer_mock.h
@@ -28,6 +28,8 @@ class ConsumerMock : public HandleMock {
     ~ConsumerMock();
 
     void subscribe(const std::vector<std::string>& topics);
+    void assign(const std::vector<TopicPartitionMock>& topic_partitions);
+    void unassign();
     void set_opaque(void* opaque);
 private:
     static uint64_t make_consumer_id();
@@ -49,14 +51,14 @@ class ConsumerMock : public HandleMock {
     using TopicPartitionId = std::tuple<std::string, int>;
 
     static TopicPartitionId make_id(const TopicPartitionMock& topic_partition);
-    void on_assignment(std::vector<TopicPartitionMock>& topic_partitions);
-    void on_revocation(const std::vector<TopicPartitionMock>& topic_partitions);
+    void on_assignment(const std::vector<TopicPartitionMock>& topic_partitions);
+    void on_revocation();
     void on_message(const std::string& topic_name, unsigned partition, uint64_t offset,
                     const KafkaMessageMock* message);
-    template <typename List>
-    void handle_rebalance(rd_kafka_resp_err_t type, List& topic_partitions);
+    void handle_rebalance(rd_kafka_resp_err_t type,
+                          const std::vector<TopicPartitionMock>& topic_partitions);
     void handle_assign(const TopicPartitionMock& topic_partition);
-    void handle_unassign(const TopicPartitionMock& topic_partition);
+    void handle_unassign();
     void fetch_existing_messages(unsigned partition, uint64_t next_offset,
                                  KafkaTopicMock& topic);
 
diff --git a/mocking/include/cppkafka/mocking/kafka_cluster.h b/mocking/include/cppkafka/mocking/kafka_cluster.h
@@ -16,8 +16,8 @@ namespace mocking {
 
 class KafkaCluster {
 public:
-    using AssignmentCallback = std::function<void(std::vector<TopicPartitionMock>&)>;
-    using RevocationCallback = std::function<void(const std::vector<TopicPartitionMock>&)>;
+    using AssignmentCallback = std::function<void(const std::vector<TopicPartitionMock>&)>;
+    using RevocationCallback = std::function<void()>;
     using MessageCallback = std::function<void(std::string topic, unsigned partition,
                                                uint64_t offset, const KafkaMessageMock*)>;
 
@@ -32,8 +32,6 @@ class KafkaCluster {
     void create_topic(const std::string& name, unsigned partitions);
     bool topic_exists(const std::string& name) const;
     void produce(const std::string& topic, unsigned partition, KafkaMessageMock message);
-    template <typename Functor>
-    void acquire_topic(const std::string& topic, const Functor& functor);
     KafkaTopicMock& get_topic(const std::string& name);
     const KafkaTopicMock& get_topic(const std::string& name) const;
     void subscribe(const std::string& group_id, uint64_t consumer_id,
@@ -72,19 +70,6 @@ class KafkaCluster {
     mutable std::mutex consumer_data_mutex_;
 };
 
-template <typename Functor>
-void KafkaCluster::acquire_topic(const std::string& topic, const Functor& functor) {
-    std::unique_lock<std::mutex> lock(topics_mutex_);
-    auto iter = topics_.find(topic);
-    if (iter == topics_.end()) {
-        throw std::runtime_error("Topic " + topic + " doesn't exist");
-    }
-    // Unlock and execute callback. We won't remove topics so this is thread safe on a
-    // cluster level
-    lock.unlock();
-    functor(iter->second);
-}
-
 } // mocking
 } // cppkafka
 
diff --git a/mocking/src/consumer_mock.cpp b/mocking/src/consumer_mock.cpp
@@ -13,6 +13,7 @@ using std::bind;
 using std::runtime_error;
 using std::make_tuple;
 using std::tie;
+using std::get;
 using std::lock_guard;
 using std::unique_lock;
 using std::mutex;
@@ -47,11 +48,22 @@ void ConsumerMock::subscribe(const vector<string>& topics) {
         consumer_id_,
         topics,
         bind(&ConsumerMock::on_assignment, this, _1),
-        bind(&ConsumerMock::on_revocation, this, _1),
+        bind(&ConsumerMock::on_revocation, this),
         bind(&ConsumerMock::on_message, this, _1, _2, _3, _4)
     );
 }
 
+void ConsumerMock::assign(const vector<TopicPartitionMock>& topic_partitions) {
+    for (const TopicPartitionMock& topic_partition : topic_partitions) {
+        handle_assign(topic_partition);
+    }
+}
+
+void ConsumerMock::unassign() {
+    lock_guard<mutex> _(assigned_partitions_mutex_);
+    assigned_partitions_.clear();
+}
+
 void ConsumerMock::set_opaque(void* opaque) {
     opaque_ = opaque;
 }
@@ -60,18 +72,23 @@ ConsumerMock::TopicPartitionId ConsumerMock::make_id(const TopicPartitionMock& t
     return make_tuple(topic_partition.get_topic(), topic_partition.get_partition());
 }
 
-void ConsumerMock::on_assignment(vector<TopicPartitionMock>& topic_partitions) {
+void ConsumerMock::on_assignment(const vector<TopicPartitionMock>& topic_partitions) {
     handle_rebalance(RD_KAFKA_RESP_ERR__ASSIGN_PARTITIONS, topic_partitions);
-    for (const TopicPartitionMock& topic_partition : topic_partitions) {
-        handle_assign(topic_partition);
-    }
 }
 
-void ConsumerMock::on_revocation(const vector<TopicPartitionMock>& topic_partitions) {
+void ConsumerMock::on_revocation() {
+    // Fetch and reset all assigned topic partitions
+    vector<TopicPartitionMock> topic_partitions = [&]() {
+        lock_guard<mutex> _(assigned_partitions_mutex_);
+        vector<TopicPartitionMock> output;
+        for (const auto& topic_partition_pair : assigned_partitions_) {
+            const TopicPartitionId& id = topic_partition_pair.first;
+            output.emplace_back(get<0>(id), get<1>(id));
+        }
+        assigned_partitions_.clear();
+        return output;
+    }();
     handle_rebalance(RD_KAFKA_RESP_ERR__REVOKE_PARTITIONS, topic_partitions);
-    for (const TopicPartitionMock& topic_partition : topic_partitions) {
-        handle_unassign(topic_partition);
-    }
 }
 
 void ConsumerMock::on_message(const string& topic_name, unsigned partition, uint64_t offset,
@@ -91,8 +108,8 @@ void ConsumerMock::on_message(const string& topic_name, unsigned partition, uint
     message_queue_condition_.notify_one();
 }
 
-template <typename List>
-void ConsumerMock::handle_rebalance(rd_kafka_resp_err_t type, List& topic_partitions) {
+void ConsumerMock::handle_rebalance(rd_kafka_resp_err_t type,
+                                    const vector<TopicPartitionMock>& topic_partitions) {
     auto rebalance_callback = config_.get_rebalance_callback();
     if (rebalance_callback) {
         auto handle = to_rdkafka_handle(topic_partitions);
@@ -115,14 +132,8 @@ void ConsumerMock::handle_assign(const TopicPartitionMock& topic_partition) {
 
     // Fetch any existing messages and push them to the available message queue
     auto& cluster = get_cluster();
-    cluster.acquire_topic(topic_partition.get_topic(), [&](KafkaTopicMock& topic) {
-        fetch_existing_messages(topic_partition.get_partition(), next_offset, topic);
-    });
-}
-
-void ConsumerMock::handle_unassign(const TopicPartitionMock& topic_partition) {
-    lock_guard<mutex> _(assigned_partitions_mutex_);
-    assigned_partitions_.erase(make_id(topic_partition));
+    KafkaTopicMock& topic = cluster.get_topic(topic_partition.get_topic());
+    fetch_existing_messages(topic_partition.get_partition(), next_offset, topic);
 }
 
 void ConsumerMock::fetch_existing_messages(unsigned partition_index, uint64_t next_offset,
diff --git a/mocking/src/kafka_cluster.cpp b/mocking/src/kafka_cluster.cpp
@@ -163,7 +163,7 @@ void KafkaCluster::generate_revocations(const TopicConsumersMap& topic_consumers
             ConsumerMetadata& consumer = consumer_data_[consumer_id];
             // Execute revocation callback and unsubscribe from the partition object
             if (!consumer.partitions_assigned.empty()) {
-                consumer.revocation_callback(consumer.partitions_assigned);
+                consumer.revocation_callback();
                 for (const auto& topic_subscription : consumer.subscriptions) {
                     KafkaTopicMock& topic = *topic_subscription.first;
                     for (const auto& partition_subscription : topic_subscription.second) {
