Move StorageKafka::createConsumer() into KafkaConsumer

azat · azat · commit ebad1bf4f3c4 · 2023-12-27T09:49:07.000+01:00
Signed-off-by: Azat Khuzhin &lt;a.khuzhin@semrush.com&gt;
diff --git a/src/Storages/Kafka/KafkaConsumer.cpp b/src/Storages/Kafka/KafkaConsumer.cpp
@@ -13,6 +13,7 @@
 
 #include <Common/CurrentMetrics.h>
 #include <Common/ProfileEvents.h>
+#include <base/defines.h>
 
 namespace CurrentMetrics
 {
@@ -63,9 +64,22 @@ KafkaConsumer::KafkaConsumer(
 {
 }
 
-void KafkaConsumer::setConsumer(const ConsumerPtr & consumer_)
+void KafkaConsumer::createConsumer(cppkafka::Configuration consumer_config)
 {
-    consumer = consumer_;
+    chassert(!consumer.get());
+
+    /// Using this should be safe, since cppkafka::Consumer can poll messages
+    /// (including statistics, which will trigger the callback below) only via
+    /// KafkaConsumer.
+    if (consumer_config.get("statistics.interval.ms") != "0")
+    {
+        consumer_config.set_stats_callback([this](cppkafka::KafkaHandleBase &, const std::string & stat_json)
+        {
+            setRDKafkaStat(stat_json);
+        });
+    }
+    consumer = std::make_shared<cppkafka::Consumer>(consumer_config);
+    consumer->set_destroy_flags(RD_KAFKA_DESTROY_F_NO_CONSUMER_CLOSE);
 
     // called (synchronously, during poll) when we enter the consumer group
     consumer->set_assignment_callback([this](const cppkafka::TopicPartitionList & topic_partitions)
diff --git a/src/Storages/Kafka/KafkaConsumer.h b/src/Storages/Kafka/KafkaConsumer.h
@@ -72,7 +72,7 @@ class KafkaConsumer
 
     ~KafkaConsumer();
 
-    void setConsumer(const ConsumerPtr & consumer);
+    void createConsumer(cppkafka::Configuration consumer_config);
     bool hasConsumer() const { return consumer.get() != nullptr; }
     ConsumerPtr && moveConsumer() { return std::move(consumer); }
 
diff --git a/src/Storages/Kafka/StorageKafka.cpp b/src/Storages/Kafka/StorageKafka.cpp
@@ -530,8 +530,11 @@ KafkaConsumerPtr StorageKafka::popConsumer(std::chrono::milliseconds timeout)
     else if (!ret_consumer_ptr && closed_consumer_index.has_value())
     {
         ret_consumer_ptr = consumers[*closed_consumer_index];
+
+        cppkafka::Configuration consumer_config = getConsumerConfiguration(*closed_consumer_index);
         /// It should be OK to create consumer under lock, since it should be fast (without subscribing).
-        ret_consumer_ptr->setConsumer(createConsumer(*ret_consumer_ptr, *closed_consumer_index));
+        ret_consumer_ptr->createConsumer(consumer_config);
+        LOG_TRACE(log, "Created #{} consumer", *closed_consumer_index);
     }
     /// 3. There is no free consumer and num_consumers already created, waiting @timeout.
     else
@@ -576,29 +579,6 @@ KafkaConsumerPtr StorageKafka::createKafkaConsumer(size_t consumer_number)
     return kafka_consumer_ptr;
 }
 
-ConsumerPtr StorageKafka::createConsumer(KafkaConsumer & kafka_consumer, size_t consumer_number)
-{
-    cppkafka::Configuration consumer_config = getConsumerConfiguration(consumer_number);
-
-    /// Using KafkaConsumer by reference should be safe, since
-    /// cppkafka::Consumer can poll messages (including statistics, which will
-    /// trigger the callback below) only via KafkaConsumer.
-    if (consumer_config.get("statistics.interval.ms") != "0")
-    {
-        consumer_config.set_stats_callback([&kafka_consumer](cppkafka::KafkaHandleBase &, const std::string & stat_json)
-        {
-            kafka_consumer.setRDKafkaStat(stat_json);
-        });
-    }
-
-    auto consumer_ptr = std::make_shared<cppkafka::Consumer>(consumer_config);
-    consumer_ptr->set_destroy_flags(RD_KAFKA_DESTROY_F_NO_CONSUMER_CLOSE);
-
-    LOG_TRACE(log, "Created #{} consumer", consumer_number);
-
-    return consumer_ptr;
-}
-
 cppkafka::Configuration StorageKafka::getConsumerConfiguration(size_t consumer_number)
 {
     cppkafka::Configuration conf;
diff --git a/src/Storages/Kafka/StorageKafka.h b/src/Storages/Kafka/StorageKafka.h
@@ -130,8 +130,6 @@ class StorageKafka final : public IStorage, WithContext
     SettingsChanges createSettingsAdjustments();
     /// Creates KafkaConsumer object without real consumer (cppkafka::Consumer)
     KafkaConsumerPtr createKafkaConsumer(size_t consumer_number);
-    /// Creates real cppkafka::Consumer object
-    ConsumerPtr createConsumer(KafkaConsumer & kafka_consumer, size_t consumer_number);
     /// Returns consumer configuration with all changes that had been overwritten in config
     cppkafka::Configuration getConsumerConfiguration(size_t consumer_number);
 
