Create consumers for Kafka tables on fly (but keep them for 1min since last used)

azat · azat · commit e7592c140e26 · 2023-12-27T09:49:07.000+01:00
Pool of consumers created a problem for librdkafka internal statistics,
you need to read from the queue always, while in ClickHouse consumers
created regardless are there any readers or not (attached materialized
views or direct SELECTs).

Otherwise, this statistics messages got queued and never released,
which:
- creates live memory leak
- and also makes destroy very slow, due to librdkafka internals (it
  moves entries from this queue into another linked list, but in a
  with sorting, which is incredibly slow for linked lists)

So the idea is simple, let's create a pool of consumers only when they
are required, and destroy them after some timeout (right now it is 60
seconds) if nobody uses them, that way this problem should gone.

This should also reduce number of internal librdkafka threads, when
nobody reads from Kafka tables.

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
@@ -604,6 +604,7 @@ KafkaConsumer::Stat KafkaConsumer::getStat() const
         .exceptions_buffer = [&](){std::lock_guard<std::mutex> lock(exception_mutex);
             return exceptions_buffer;}(),
         .in_use = in_use.load(),
+        .last_used_usec = last_used_usec.load(),
         .rdkafka_stat = [&](){std::lock_guard<std::mutex> lock(rdkafka_stat_mutex);
             return rdkafka_stat;}(),
     };
diff --git a/src/Storages/Kafka/KafkaConsumer.h b/src/Storages/Kafka/KafkaConsumer.h
@@ -57,6 +57,7 @@ class KafkaConsumer
         UInt64 num_rebalance_revocations;
         KafkaConsumer::ExceptionsBuffer exceptions_buffer;
         bool in_use;
+        UInt64 last_used_usec;
         std::string rdkafka_stat;
     };
 
@@ -113,11 +114,20 @@ class KafkaConsumer
         rdkafka_stat = stat_json_string;
     }
     void inUse() { in_use = true; }
-    void notInUse() { in_use = false; }
+    void notInUse()
+    {
+        in_use = false;
+        last_used_usec = std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
+    }
 
     // For system.kafka_consumers
     Stat getStat() const;
 
+    bool isInUse() const { return in_use; }
+    UInt64 getLastUsedUsec() const { return last_used_usec; }
+
+    std::string getMemberId() const;
+
 private:
     using Messages = std::vector<cppkafka::Message>;
     CurrentMetrics::Increment metric_increment{CurrentMetrics::KafkaConsumers};
@@ -168,6 +178,8 @@ class KafkaConsumer
     std::atomic<UInt64> num_rebalance_assignments = 0;
     std::atomic<UInt64> num_rebalance_revocations = 0;
     std::atomic<bool> in_use = 0;
+    /// Last used time (for TTL)
+    std::atomic<UInt64> last_used_usec = 0;
 
     mutable std::mutex rdkafka_stat_mutex;
     std::string rdkafka_stat;
@@ -178,8 +190,6 @@ class KafkaConsumer
     /// Return number of messages with an error.
     size_t filterMessageErrors();
     ReadBufferPtr getNextMessage();
-
-    std::string getMemberId() const;
 };
 
 }
diff --git a/src/Storages/Kafka/StorageKafka.cpp b/src/Storages/Kafka/StorageKafka.cpp
@@ -27,6 +27,7 @@
 #include <Storages/StorageMaterializedView.h>
 #include <Storages/NamedCollectionsHelpers.h>
 #include <base/getFQDNOrHostName.h>
+#include <Common/Stopwatch.h>
 #include <Common/logger_useful.h>
 #include <boost/algorithm/string/replace.hpp>
 #include <boost/algorithm/string/split.hpp>
@@ -76,6 +77,7 @@ namespace ErrorCodes
     extern const int BAD_ARGUMENTS;
     extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
     extern const int QUERY_NOT_ALLOWED;
+    extern const int ABORTED;
 }
 
 struct StorageKafkaInterceptors
@@ -262,7 +264,6 @@ StorageKafka::StorageKafka(
     , schema_name(getContext()->getMacros()->expand(kafka_settings->kafka_schema.value, macros_info))
     , num_consumers(kafka_settings->kafka_num_consumers.value)
     , log(&Poco::Logger::get("StorageKafka (" + table_id_.table_name + ")"))
-    , semaphore(0, static_cast<int>(num_consumers))
     , intermediate_commit(kafka_settings->kafka_commit_every_batch.value)
     , settings_adjustments(createSettingsAdjustments())
     , thread_per_consumer(kafka_settings->kafka_thread_per_consumer.value)
@@ -343,8 +344,8 @@ Pipe StorageKafka::read(
     size_t /* max_block_size */,
     size_t /* num_streams */)
 {
-    if (all_consumers.empty())
-        return {};
+    if (shutdown_called)
+        throw Exception(ErrorCodes::ABORTED, "Table is detached");
 
     if (!local_context->getSettingsRef().stream_like_engine_allow_direct_select)
         throw Exception(ErrorCodes::QUERY_NOT_ALLOWED,
@@ -357,12 +358,12 @@ Pipe StorageKafka::read(
 
     /// Always use all consumers at once, otherwise SELECT may not read messages from all partitions.
     Pipes pipes;
-    pipes.reserve(all_consumers.size());
+    pipes.reserve(num_consumers);
     auto modified_context = Context::createCopy(local_context);
     modified_context->applySettingsChanges(settings_adjustments);
 
     // Claim as many consumers as requested, but don't block
-    for (size_t i = 0; i < all_consumers.size(); ++i)
+    for (size_t i = 0; i < num_consumers; ++i)
     {
         /// Use block size of 1, otherwise LIMIT won't work properly as it will buffer excess messages in the last block
         /// TODO: probably that leads to awful performance.
@@ -412,19 +413,7 @@ SinkToStoragePtr StorageKafka::write(const ASTPtr &, const StorageMetadataPtr &
 
 void StorageKafka::startup()
 {
-    for (size_t i = 0; i < num_consumers; ++i)
-    {
-        try
-        {
-            auto consumer = createConsumer(i);
-            pushConsumer(consumer);
-            all_consumers.push_back(consumer);
-        }
-        catch (const cppkafka::Exception &)
-        {
-            tryLogCurrentException(log);
-        }
-    }
+    all_consumers.resize(num_consumers);
 
     // Start the reader thread
     for (auto & task : tasks)
@@ -438,21 +427,34 @@ void StorageKafka::shutdown(bool)
 {
     shutdown_called = true;
 
-    for (auto & task : tasks)
     {
-        // Interrupt streaming thread
-        task->stream_cancelled = true;
+        LOG_TRACE(log, "Waiting for streaming jobs");
+        Stopwatch watch;
+        for (auto & task : tasks)
+        {
+            // Interrupt streaming thread
+            task->stream_cancelled = true;
 
-        LOG_TRACE(log, "Waiting for cleanup");
-        task->holder->deactivate();
+            LOG_TEST(log, "Waiting for cleanup of a task");
+            task->holder->deactivate();
+        }
+        LOG_TRACE(log, "Streaming jobs finished in {} ms.", watch.elapsedMilliseconds());
     }
 
-    LOG_TRACE(log, "Closing consumers");
-    for (size_t i = 0; i < all_consumers.size(); ++i)
-        auto consumer = popConsumer();
-    LOG_TRACE(log, "Consumers closed");
+    {
+        std::lock_guard lock(mutex);
+        LOG_TRACE(log, "Closing {} consumers", consumers.size());
+        Stopwatch watch;
+        consumers.clear();
+        LOG_TRACE(log, "Consumers closed. Took {} ms.", watch.elapsedMilliseconds());
+    }
 
-    rd_kafka_wait_destroyed(CLEANUP_TIMEOUT_MS);
+    {
+        LOG_TRACE(log, "Waiting for final cleanup");
+        Stopwatch watch;
+        rd_kafka_wait_destroyed(CLEANUP_TIMEOUT_MS);
+        LOG_TRACE(log, "Final cleanup finished in {} ms (timeout {} ms).", watch.elapsedMilliseconds(), CLEANUP_TIMEOUT_MS);
+    }
 }
 
 
@@ -461,7 +463,7 @@ void StorageKafka::pushConsumer(KafkaConsumerPtr consumer)
     std::lock_guard lock(mutex);
     consumer->notInUse();
     consumers.push_back(consumer);
-    semaphore.set();
+    cv.notify_one();
     CurrentMetrics::sub(CurrentMetrics::KafkaConsumersInUse, 1);
 }
 
@@ -474,22 +476,48 @@ KafkaConsumerPtr StorageKafka::popConsumer()
 
 KafkaConsumerPtr StorageKafka::popConsumer(std::chrono::milliseconds timeout)
 {
-    // Wait for the first free buffer
-    if (timeout == std::chrono::milliseconds::zero())
-        semaphore.wait();
+    std::unique_lock lock(mutex);
+
+    KafkaConsumerPtr consumer_ptr;
+
+    /// 1. There is consumer available. Return one of them.
+    if (!consumers.empty())
+    {
+        consumer_ptr = consumers.back();
+        consumers.pop_back();
+    }
     else
     {
-        if (!semaphore.tryWait(timeout.count()))
-            return nullptr;
+        auto expired_consumer = std::find_if(all_consumers.begin(), all_consumers.end(), [](const auto & consumer_weak_ptr)
+        {
+            return consumer_weak_ptr.expired();
+        });
+
+        /// 2. There is no consumer, but we can create a new one.
+        if (expired_consumer != all_consumers.end())
+        {
+            size_t consumer_number = std::distance(all_consumers.begin(), expired_consumer);
+            /// It should be OK to create consumer under lock, since it should be fast (without subscribing).
+            consumer_ptr = createConsumer(consumer_number);
+            *expired_consumer = consumer_ptr;
+        }
+        /// 3. There is no consumer and num_consumers already created, waiting @timeout.
+        else
+        {
+            if (cv.wait_for(lock, timeout, [&]() { return !consumers.empty(); }))
+            {
+                consumer_ptr = consumers.back();
+                consumers.pop_back();
+            }
+        }
     }
 
-    // Take the first available buffer from the list
-    std::lock_guard lock(mutex);
-    auto consumer = consumers.back();
-    consumers.pop_back();
-    CurrentMetrics::add(CurrentMetrics::KafkaConsumersInUse, 1);
-    consumer->inUse();
-    return consumer;
+    if (consumer_ptr)
+    {
+        CurrentMetrics::add(CurrentMetrics::KafkaConsumersInUse, 1);
+        consumer_ptr->inUse();
+    }
+    return consumer_ptr;
 }
 
 
@@ -545,10 +573,59 @@ KafkaConsumerPtr StorageKafka::createConsumer(size_t consumer_number)
     {
         kafka_consumer_ptr = std::make_shared<KafkaConsumer>(consumer_impl, log, getPollMaxBatchSize(), getPollTimeoutMillisecond(), intermediate_commit, tasks.back()->stream_cancelled, topics);
     }
+    LOG_TRACE(log, "Created #{} consumer", consumer_number);
+
     *consumer_weak_ptr_ptr = kafka_consumer_ptr;
     return kafka_consumer_ptr;
 }
 
+void StorageKafka::cleanConsumers()
+{
+    static const UInt64 CONSUMER_TTL_USEC = 60'000'000;
+    UInt64 now_usec = std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
+
+    /// Copy consumers for closing to a new vector to close them without a lock
+    std::vector<KafkaConsumerPtr> consumers_to_close;
+
+    {
+        std::lock_guard lock(mutex);
+
+        for (auto it = consumers.begin(); it != consumers.end();)
+        {
+            auto & consumer_ptr = *it;
+
+            UInt64 consumer_last_used_usec = consumer_ptr->getLastUsedUsec();
+            chassert(consumer_last_used_usec <= now_usec);
+            if (now_usec - consumer_last_used_usec > CONSUMER_TTL_USEC)
+            {
+                /// We need this only to get the consumer number.
+                auto weak_it = std::find_if(all_consumers.begin(), all_consumers.end(), [&](const auto & consume_weak_ptr)
+                {
+                    return consumer_ptr == consume_weak_ptr.lock();
+                });
+                chassert(weak_it != all_consumers.end());
+                size_t consumer_number = std::distance(all_consumers.begin(), weak_it);
+
+                LOG_TRACE(log, "Closing #{} consumer (id: {})", consumer_number, consumer_ptr->getMemberId());
+
+                consumers_to_close.push_back(std::move(consumer_ptr));
+                it = consumers.erase(it);
+            }
+            else
+                ++it;
+        }
+    }
+
+    if (!consumers_to_close.empty())
+    {
+        Stopwatch watch;
+        size_t closed = consumers_to_close.size();
+        consumers_to_close.clear();
+        LOG_TRACE(log, "{} consumers had been closed (due to {} usec timeout). Took {} ms.",
+            closed, CONSUMER_TTL_USEC, watch.elapsedMilliseconds());
+    }
+}
+
 size_t StorageKafka::getMaxBlockSize() const
 {
     return kafka_settings->kafka_max_block_size.changed
@@ -806,6 +883,8 @@ void StorageKafka::threadFunc(size_t idx)
 
     mv_attached.store(false);
 
+    cleanConsumers();
+
     // Wait for attached views
     if (!task->stream_cancelled)
         task->holder->scheduleAfter(RESCHEDULE_MS);
diff --git a/src/Storages/Kafka/StorageKafka.h b/src/Storages/Kafka/StorageKafka.h
@@ -9,6 +9,7 @@
 
 #include <Poco/Semaphore.h>
 
+#include <condition_variable>
 #include <mutex>
 #include <list>
 #include <atomic>
@@ -102,7 +103,6 @@ class StorageKafka final : public IStorage, WithContext
     const String schema_name;
     const size_t num_consumers; /// total number of consumers
     Poco::Logger * log;
-    Poco::Semaphore semaphore;
     const bool intermediate_commit;
     const SettingsChanges settings_adjustments;
 
@@ -112,6 +112,7 @@ class StorageKafka final : public IStorage, WithContext
     std::vector<KafkaConsumerWeakPtr> all_consumers; /// busy (belong to a KafkaSource) and vacant consumers
 
     std::mutex mutex;
+    std::condition_variable cv;
 
     // Stream thread
     struct TaskContext
@@ -157,6 +158,7 @@ class StorageKafka final : public IStorage, WithContext
     bool streamToViews();
     bool checkDependencies(const StorageID & table_id);
 
+    void cleanConsumers();
 };
 
 }
diff --git a/src/Storages/System/StorageSystemKafkaConsumers.cpp b/src/Storages/System/StorageSystemKafkaConsumers.cpp
@@ -41,6 +41,7 @@ NamesAndTypesList StorageSystemKafkaConsumers::getNamesAndTypes()
         {"num_rebalance_revocations", std::make_shared<DataTypeUInt64>()},
         {"num_rebalance_assignments", std::make_shared<DataTypeUInt64>()},
         {"is_currently_used", std::make_shared<DataTypeUInt8>()},
+        {"last_used", std::make_shared<DataTypeDateTime64>(6)},
         {"rdkafka_stat", std::make_shared<DataTypeString>()},
     };
     return names_and_types;
@@ -78,6 +79,7 @@ void StorageSystemKafkaConsumers::fillData(MutableColumns & res_columns, Context
     auto & num_rebalance_revocations = assert_cast<ColumnUInt64 &>(*res_columns[index++]);
     auto & num_rebalance_assigments = assert_cast<ColumnUInt64 &>(*res_columns[index++]);
     auto & is_currently_used = assert_cast<ColumnUInt8 &>(*res_columns[index++]);
+    auto & last_used = assert_cast<ColumnDateTime64 &>(*res_columns[index++]);
     auto & rdkafka_stat = assert_cast<ColumnString &>(*res_columns[index++]);
 
     const auto access = context->getAccess();
@@ -144,6 +146,7 @@ void StorageSystemKafkaConsumers::fillData(MutableColumns & res_columns, Context
                 num_rebalance_assigments.insert(consumer_stat.num_rebalance_assignments);
 
                 is_currently_used.insert(consumer_stat.in_use);
+                last_used.insert(consumer_stat.last_used_usec);
 
                 rdkafka_stat.insertData(consumer_stat.rdkafka_stat.data(), consumer_stat.rdkafka_stat.size());
             }
