schedule: add keyspace cpu limiter

dbms/src/Common/TiFlashMetrics.cpp

@@ -84,6 +84,18 @@ TiFlashMetrics::TiFlashMetrics()
                                                       .Name("tiflash_storage_sync_replica_ru")
                                                       .Help("RU for synchronous replica of keyspace")
                                                       .Register(*registry);
+    registered_keyspace_cpu_limiter_family = &prometheus::BuildGauge()
+                                                  .Name("tiflash_pipeline_keyspace_cpu_limiter")
+                                                  .Help("Keyspace CPU limiter state")
+                                                  .Register(*registry);
+    registered_keyspace_cpu_limiter_cpu_seconds_family = &prometheus::BuildCounter()
+                                                              .Name("tiflash_pipeline_keyspace_cpu_limiter_cpu_seconds")
+                                                              .Help("CPU time charged to the keyspace CPU limiter")
+                                                              .Register(*registry);
+    registered_keyspace_cpu_limiter_throttled_family = &prometheus::BuildCounter()
+                                                            .Name("tiflash_pipeline_keyspace_cpu_limiter_throttled")
+                                                            .Help("Failed keyspace CPU limiter admission attempts")
+                                                            .Register(*registry);
     registered_raft_proxy_thread_memory_usage_family
         = &prometheus::BuildGauge().Name(raft_proxy_thread_memory_usage).Help("").Register(*registry);
 
@@ -191,6 +203,36 @@ UInt64 TiFlashMetrics::debugQueryReplicaSyncRU(UInt32 keyspace_id)
     return counter->Value();
 }
 
+TiFlashMetrics::KeyspaceCpuLimiterMetrics & TiFlashMetrics::getKeyspaceCpuLimiterMetrics(KeyspaceID keyspace_id)
+{
+    std::lock_guard lock(keyspace_cpu_limiter_mtx);
+    auto [iter, inserted] = registered_keyspace_cpu_limiter_metrics.try_emplace(keyspace_id);
+    if (!inserted)
+        return iter->second;
+
+    const auto keyspace_id_str = std::to_string(keyspace_id);
+    auto labels = [&](const char * type) {
+        return prometheus::Labels{{"keyspace_id", keyspace_id_str}, {"type", type}};
+    };
+    auto reason_labels = [&](const char * reason) {
+        return prometheus::Labels{{"keyspace_id", keyspace_id_str}, {"reason", reason}};
+    };
+
+    iter->second.active_tasks = &registered_keyspace_cpu_limiter_family->Add(labels("active_tasks"));
+    iter->second.max_active_tasks = &registered_keyspace_cpu_limiter_family->Add(labels("max_active_tasks"));
+    iter->second.cpu_tokens_seconds = &registered_keyspace_cpu_limiter_family->Add(labels("cpu_tokens_seconds"));
+    iter->second.cpu_quota_seconds_per_second
+        = &registered_keyspace_cpu_limiter_family->Add(labels("cpu_quota_seconds_per_second"));
+    iter->second.throttled = &registered_keyspace_cpu_limiter_family->Add(labels("throttled"));
+    iter->second.cpu_seconds_total
+        = &registered_keyspace_cpu_limiter_cpu_seconds_family->Add({{"keyspace_id", keyspace_id_str}});
+    iter->second.cpu_quota_throttled_total
+        = &registered_keyspace_cpu_limiter_throttled_family->Add(reason_labels("cpu_quota"));
+    iter->second.active_tasks_throttled_total
+        = &registered_keyspace_cpu_limiter_throttled_family->Add(reason_labels("active_tasks"));
+    return iter->second;
+}
+
 prometheus::Counter * TiFlashMetrics::getReplicaSyncRUCounter(UInt32 keyspace_id, std::unique_lock<std::mutex> &)
 {
     auto itr = registered_keyspace_sync_replica_ru.find(keyspace_id);

dbms/src/Common/TiFlashMetrics.h

@@ -1373,6 +1373,20 @@ class TiFlashMetrics
         Count,
     };
 
+    struct KeyspaceCpuLimiterMetrics
+    {
+        prometheus::Gauge * active_tasks;
+        prometheus::Gauge * max_active_tasks;
+        prometheus::Gauge * cpu_tokens_seconds;
+        prometheus::Gauge * cpu_quota_seconds_per_second;
+        prometheus::Gauge * throttled;
+        prometheus::Counter * cpu_seconds_total;
+        prometheus::Counter * cpu_quota_throttled_total;
+        prometheus::Counter * active_tasks_throttled_total;
+    };
+
+    KeyspaceCpuLimiterMetrics & getKeyspaceCpuLimiterMetrics(KeyspaceID keyspace_id);
+
     void addReplicaSyncRU(UInt32 keyspace_id, UInt64 ru);
     UInt64 debugQueryReplicaSyncRU(UInt32 keyspace_id);
     enum class MemoryAllocType
@@ -1437,6 +1451,12 @@ class TiFlashMetrics
     std::mutex replica_sync_ru_mtx;
     std::unordered_map<KeyspaceID, prometheus::Counter *> registered_keyspace_sync_replica_ru;
 
+    prometheus::Family<prometheus::Gauge> * registered_keyspace_cpu_limiter_family;
+    prometheus::Family<prometheus::Counter> * registered_keyspace_cpu_limiter_cpu_seconds_family;
+    prometheus::Family<prometheus::Counter> * registered_keyspace_cpu_limiter_throttled_family;
+    std::mutex keyspace_cpu_limiter_mtx;
+    std::unordered_map<KeyspaceID, KeyspaceCpuLimiterMetrics> registered_keyspace_cpu_limiter_metrics;
+
     // TODO: Use CAS+HazPtr to remove proxy_thread_report_mtx, or hash some slots here.
     prometheus::Family<prometheus::Gauge> * registered_raft_proxy_thread_memory_usage_family;
     std::shared_mutex proxy_thread_report_mtx;

dbms/src/Flash/Pipeline/Schedule/TaskQueues/IOPriorityQueue.cpp

@@ -16,6 +16,8 @@
 #include <Flash/Pipeline/Schedule/Tasks/TaskHelper.h>
 #include <common/likely.h>
 
+#include <ext/scope_guard.h>
+
 namespace DB
 {
 namespace
@@ -46,44 +48,103 @@ IOPriorityQueue::~IOPriorityQueue()
 bool IOPriorityQueue::take(TaskPtr & task)
 {
     std::unique_lock lock(mu);
+    const bool limiter_enabled = keyspace_cpu_limiter && keyspace_cpu_limiter->isEnabled();
     while (true)
     {
         // Remaining tasks will be drained in destructor.
         if (unlikely(is_finished))
             return false;
 
+        // Snapshot before inspecting the queues. A release can happen without
+        // holding `mu`; taking the snapshot afterwards could miss that wakeup.
+        const auto previous_change_id = limiter_enabled ? keyspace_cpu_limiter->getChangeId() : 0;
+
         if (popTask(cancel_task_queue, task))
             return true;
 
         bool io_out_first = ratio_of_out_to_in * total_io_in_time_microsecond >= total_io_out_time_microsecond;
         auto & first_queue = io_out_first ? io_out_task_queue : io_in_task_queue;
         auto & next_queue = io_out_first ? io_in_task_queue : io_out_task_queue;
-        if (popTask(first_queue, task))
+        if (tryTakeTaskWithoutLock(first_queue, task))
             return true;
-        if (popTask(next_queue, task))
+        if (tryTakeTaskWithoutLock(next_queue, task))
             return true;
-        cv.wait(lock);
+        if (limiter_enabled)
+        {
+            const bool has_pending_tasks
+                = !io_out_task_queue.empty() || !io_in_task_queue.empty() || !cancel_task_queue.empty();
+            if (has_pending_tasks)
+            {
+                lock.unlock();
+                if (keyspace_cpu_limiter->isCPUQuotaEnabled())
+                    keyspace_cpu_limiter->waitForChange(
+                        previous_change_id,
+                        keyspace_cpu_limiter->getRefillWaitDuration());
+                else
+                    keyspace_cpu_limiter->waitForChange(previous_change_id);
+                lock.lock();
+            }
+            else
+            {
+                cv.wait(lock);
+            }
+        }
+        else
+        {
+            cv.wait(lock);
+        }
+    }
+}
+
+bool IOPriorityQueue::tryTakeTaskWithoutLock(std::list<TaskPtr> & task_queue, TaskPtr & task)
+{
+    if (!keyspace_cpu_limiter || !keyspace_cpu_limiter->isEnabled())
+        return popTask(task_queue, task);
+
+    for (auto it = task_queue.begin(); it != task_queue.end(); ++it)
+    {
+        const auto keyspace_id = (*it)->getKeyspaceID();
+        if (!keyspace_cpu_limiter->tryAcquire(keyspace_id))
+            continue;
+
+        bool owner_bound = false;
+        SCOPE_EXIT({
+            if (!owner_bound)
+                keyspace_cpu_limiter->release(keyspace_id);
+        });
+        task = std::move(*it);
+        task_queue.erase(it);
+        keyspace_cpu_limiter->bindOwner(keyspace_id, task.get());
+        owner_bound = true;
+        return true;
     }
+    return false;
 }
 
 void IOPriorityQueue::drainTaskQueueWithoutLock()
 {
     TaskPtr task;
     while (popTask(cancel_task_queue, task))
     {
+        if (keyspace_cpu_limiter)
+            keyspace_cpu_limiter->release(task.get());
         FINALIZE_TASK(task);
     }
     while (popTask(io_out_task_queue, task))
     {
+        if (keyspace_cpu_limiter)
+            keyspace_cpu_limiter->release(task.get());
         FINALIZE_TASK(task);
     }
     while (popTask(io_in_task_queue, task))
     {
+        if (keyspace_cpu_limiter)
+            keyspace_cpu_limiter->release(task.get());
         FINALIZE_TASK(task);
     }
 }
 
-void IOPriorityQueue::updateStatistics(const TaskPtr &, ExecTaskStatus exec_task_status, UInt64 inc_ns)
+void IOPriorityQueue::updateStatistics(const TaskPtr & task, ExecTaskStatus exec_task_status, UInt64 inc_ns)
 {
     switch (exec_task_status)
     {
@@ -95,6 +156,14 @@ void IOPriorityQueue::updateStatistics(const TaskPtr &, ExecTaskStatus exec_task
         break;
     default:; // ignore not io status.
     }
+
+    if (keyspace_cpu_limiter && keyspace_cpu_limiter->isEnabled())
+    {
+        // A cancelled task can be returned from cancel_task_queue without ever
+        // acquiring a reservation; release() is owner-aware and therefore a no-op.
+        keyspace_cpu_limiter->release(task.get());
+        notifyWaiters();
+    }
 }
 
 bool IOPriorityQueue::empty() const
@@ -109,7 +178,7 @@ void IOPriorityQueue::finish()
         std::lock_guard lock(mu);
         is_finished = true;
     }
-    cv.notify_all();
+    notifyWaiters();
 }
 
 void IOPriorityQueue::submitTaskWithoutLock(TaskPtr && task)
@@ -148,7 +217,7 @@ void IOPriorityQueue::submit(TaskPtr && task)
         std::lock_guard lock(mu);
         submitTaskWithoutLock(std::move(task));
     }
-    cv.notify_one();
+    notifyWaiters();
 }
 
 void IOPriorityQueue::submit(std::vector<TaskPtr> & tasks)
@@ -162,30 +231,37 @@ void IOPriorityQueue::submit(std::vector<TaskPtr> & tasks)
         return;
     }
 
-    std::lock_guard lock(mu);
-    for (auto & task : tasks)
     {
-        submitTaskWithoutLock(std::move(task));
-        cv.notify_one();
+        std::lock_guard lock(mu);
+        for (auto & task : tasks)
+            submitTaskWithoutLock(std::move(task));
     }
+    notifyWaiters();
 }
 
 void IOPriorityQueue::cancel(const TaskCancelInfo & cancel_info)
 {
     if unlikely (cancel_info.query_id.empty())
         return;
 
-    std::lock_guard lock(mu);
-    if (cancel_query_id_cache.add(cancel_info.query_id))
     {
-        collectCancelledTasks(cancel_task_queue, cancel_info.query_id);
-        cv.notify_all();
+        std::lock_guard lock(mu);
+        if (cancel_query_id_cache.add(cancel_info.query_id))
+            collectCancelledTasks(cancel_task_queue, cancel_info.query_id);
     }
+    notifyWaiters();
 }
 
 void IOPriorityQueue::collectCancelledTasks(std::deque<TaskPtr> & cancel_queue, const String & query_id)
 {
     moveCancelledTasks(io_in_task_queue, cancel_queue, query_id);
     moveCancelledTasks(io_out_task_queue, cancel_queue, query_id);
 }
+
+void IOPriorityQueue::notifyWaiters()
+{
+    cv.notify_all();
+    if (keyspace_cpu_limiter)
+        keyspace_cpu_limiter->notifyAll();
+}
 } // namespace DB

dbms/src/Flash/Pipeline/Schedule/TaskQueues/IOPriorityQueue.h

@@ -15,10 +15,12 @@
 #pragma once
 
 #include <Flash/Pipeline/Schedule/TaskQueues/FIFOQueryIdCache.h>
+#include <Flash/Pipeline/Schedule/TaskQueues/KeyspaceCpuLimiter.h>
 #include <Flash/Pipeline/Schedule/TaskQueues/TaskQueue.h>
 
 #include <deque>
 #include <mutex>
+#include <utility>
 
 namespace DB
 {
@@ -33,6 +35,10 @@ class IOPriorityQueue : public TaskQueue
     // The ratio of total execution time between io_out and io_in is 3:1.
     static constexpr size_t ratio_of_out_to_in = 3;
 
+    explicit IOPriorityQueue(KeyspaceCpuLimiterPtr keyspace_cpu_limiter_ = nullptr)
+        : keyspace_cpu_limiter(std::move(keyspace_cpu_limiter_))
+    {}
+
     ~IOPriorityQueue() override;
 
     void submit(TaskPtr && task) override;
@@ -52,9 +58,11 @@ class IOPriorityQueue : public TaskQueue
     void collectCancelledTasks(std::deque<TaskPtr> & cancel_queue, const String & query_id);
 
 private:
+    bool tryTakeTaskWithoutLock(std::list<TaskPtr> & task_queue, TaskPtr & task);
     void submitTaskWithoutLock(TaskPtr && task);
 
     void drainTaskQueueWithoutLock();
+    void notifyWaiters();
 
 private:
     mutable std::mutex mu;
@@ -69,5 +77,7 @@ class IOPriorityQueue : public TaskQueue
 
     FIFOQueryIdCache cancel_query_id_cache;
     std::deque<TaskPtr> cancel_task_queue;
+
+    KeyspaceCpuLimiterPtr keyspace_cpu_limiter;
 };
 } // namespace DB