Backport ClickHouse#88217 to 25.8: max_cpu_share alone should determine the hard cap in a workload setting, even if max_cpus is unset

Author: robot-clickhouse

src/Common/Scheduler/IResourceManager.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <Common/Scheduler/ResourceLink.h>
+#include <Common/Scheduler/WorkloadSettings.h>
 
 #include <Poco/Util/AbstractConfiguration.h>
 
@@ -37,6 +38,8 @@ class IClassifier : private boost::noncopyable
     /// Returns ResourceLink that should be used to access resource.
     /// Returned link is valid until classifier destruction.
     virtual ResourceLink get(const String & resource_name) = 0;
+    /// Returns settings that should be used to limit workload on given resource.
+    virtual WorkloadSettings getWorkloadSettings(const String & resource_name) const = 0;
 };
 
 using ClassifierPtr = std::shared_ptr<IClassifier>;

src/Common/Scheduler/Nodes/CustomResourceManager.cpp

@@ -177,6 +177,12 @@ ResourceLink CustomResourceManager::Classifier::get(const String & resource_name
         return ResourceLink{}; // unlimited access
 }
 
+WorkloadSettings CustomResourceManager::Classifier::getWorkloadSettings(const String & resource_name) const
+{
+    UNUSED(resource_name);
+    return {};
+}
+
 CustomResourceManager::CustomResourceManager()
     : state(new State())
 {

src/Common/Scheduler/Nodes/CustomResourceManager.h

@@ -84,6 +84,7 @@ class CustomResourceManager : public IResourceManager
         Classifier(const ClassifierSettings & settings_, const StatePtr & state_, const String & classifier_name);
         bool has(const String & resource_name) override;
         ResourceLink get(const String & resource_name) override;
+        WorkloadSettings getWorkloadSettings(const String & resource_name) const override;
     private:
         const ClassifierSettings settings;
         std::unordered_map<String, ResourceLink> resources; // accessible resources by names

src/Common/Scheduler/Nodes/WorkloadResourceManager.cpp

@@ -423,11 +423,32 @@ ResourceLink WorkloadResourceManager::Classifier::get(const String & resource_na
     }
 }
 
-void WorkloadResourceManager::Classifier::attach(const ResourcePtr & resource, const VersionPtr & version, ResourceLink link)
+WorkloadSettings WorkloadResourceManager::Classifier::getWorkloadSettings(const String & resource_name) const
+{
+    std::unique_lock lock{mutex};
+    auto iter = attachments.find(resource_name);
+    if (iter != attachments.end())
+    {
+        // Extract settings from the attached resource
+        return iter->second.settings;
+    }
+    return {};
+}
+
+void WorkloadResourceManager::Classifier::attach(const ResourcePtr & resource, const VersionPtr & version, UnifiedSchedulerNode * node)
 {
     std::unique_lock lock{mutex};
     chassert(!attachments.contains(resource->getName()));
-    attachments[resource->getName()] = Attachment{.resource = resource, .version = version, .link = link};
+    ResourceLink link;
+    WorkloadSettings wl_settings{};
+    if (node)
+    {
+        auto queue = node->getQueue();
+        if (queue)
+            link = ResourceLink{.queue = queue.get()};
+        wl_settings = node->getSettings();
+    }
+    attachments[resource->getName()] = Attachment{.resource = resource, .version = version, .link = link, .settings = wl_settings};
 }
 
 void WorkloadResourceManager::Resource::updateResource(const ASTPtr & new_resource_entity)
@@ -447,11 +468,12 @@ std::future<void> WorkloadResourceManager::Resource::attachClassifier(Classifier
         {
             if (auto iter = node_for_workload.find(workload_name); iter != node_for_workload.end())
             {
+                auto nodePtr = iter->second;
                 auto queue = iter->second->getQueue();
                 if (!queue)
                     throw Exception(ErrorCodes::INVALID_SCHEDULER_NODE, "Unable to use workload '{}' that have children for resource '{}'",
                         workload_name, resource_name);
-                classifier.attach(shared_from_this(), current_version, ResourceLink{.queue = queue.get()});
+                classifier.attach(shared_from_this(), current_version, nodePtr.get());
             }
             else
             {

src/Common/Scheduler/Nodes/WorkloadResourceManager.h

@@ -241,21 +241,23 @@ class WorkloadResourceManager : public IResourceManager
         /// NOTE: It is called from query threads (possibly multiple)
         bool has(const String & resource_name) override;
         ResourceLink get(const String & resource_name) override;
+        WorkloadSettings getWorkloadSettings(const String & resource_name) const override;
 
         /// Attaches/detaches a specific resource
         /// NOTE: It is called from scheduler threads (possibly multiple)
-        void attach(const ResourcePtr & resource, const VersionPtr & version, ResourceLink link);
+        void attach(const ResourcePtr & resource, const VersionPtr & version, UnifiedSchedulerNode * node);
         void detach(const ResourcePtr & resource);
 
     private:
         const ClassifierSettings settings;
         WorkloadResourceManager * resource_manager;
-        std::mutex mutex;
+        mutable std::mutex mutex;
         struct Attachment
         {
             ResourcePtr resource;
             VersionPtr version;
             ResourceLink link;
+            WorkloadSettings settings;
         };
         std::unordered_map<String, Attachment> attachments; // TSA_GUARDED_BY(mutex);
     };

src/Common/Scheduler/Nodes/tests/gtest_workload_resource_manager.cpp

@@ -18,6 +18,7 @@
 #include <Common/Scheduler/Nodes/tests/ResourceTest.h>
 #include <Common/Scheduler/Workload/WorkloadEntityStorageBase.h>
 #include <Common/Scheduler/Nodes/WorkloadResourceManager.h>
+#include <Common/getNumberOfCPUCoresToUse.h>
 
 #include <base/scope_guard.h>
 
@@ -1268,6 +1269,23 @@ TEST(SchedulerWorkloadResourceManager, CPUSchedulingIndependentPools)
     t.wait();
 }
 
+TEST(SchedulerWorkloadResourceManager, MaxCPUsDerivedFromShare)
+{
+    ResourceTest t;
+
+    t.query("CREATE RESOURCE cpu (MASTER THREAD, WORKER THREAD)");
+    // Only max_cpu_share is set, max_cpus is unset
+    t.query("CREATE WORKLOAD all SETTINGS max_cpu_share = 0.5");
+    ClassifierPtr c = t.manager->acquire("all");
+
+    // The expected hard cap is max_cpu_share * getNumberOfCPUCoresToUse()
+    WorkloadSettings settings = c->getWorkloadSettings("cpu");
+    double expected_cap = 0.5 * getNumberOfCPUCoresToUse();
+    double actual_cap = settings.max_cpus;
+
+    EXPECT_DOUBLE_EQ(actual_cap, expected_cap);
+}
+
 auto getAcquired()
 {
     return CurrentMetrics::get(CurrentMetrics::ConcurrencyControlAcquired);

src/Common/Scheduler/WorkloadSettings.cpp

@@ -232,7 +232,7 @@ void WorkloadSettings::initFromChanges(CostUnit unit_, const ASTCreateWorkloadQu
         if (share_limit > 0)
         {
             Float64 value = share_limit * getNumberOfCPUCoresToUse();
-            if (value > 0 && value < limit)
+            if (value > 0 && (limit == 0 || value < limit))
                 limit = value;
         }
         max_cpus = limit;

src/Common/Scheduler/createResourceManager.cpp

@@ -55,6 +55,16 @@ class ResourceManagerDispatcher : public IResourceManager
                 return ResourceLink{};
         }
 
+        WorkloadSettings getWorkloadSettings(const String & resource_name) const override
+        {
+            for (const auto & classifier : classifiers)
+            {
+                if (classifier->has(resource_name))
+                    return classifier->getWorkloadSettings(resource_name);
+            }
+            return {};
+        }
+
     private:
         const ClassifierSettings settings;
         std::vector<ClassifierPtr> classifiers; // should be constant after initialization to avoid races