simplify scheduler constraints

src/Common/Scheduler/ISchedulerConstraint.h

@@ -15,8 +15,7 @@ namespace DB
  * When constraint is again satisfied, scheduleActivation() is called from finishRequest().
  *
  * Derived class behaviour requirements:
- *  - dequeueRequest() must fill `request->constraint` iff it is nullptr;
- *  - finishRequest() must be recursive: call to `parent_constraint->finishRequest()`.
+ *  - dequeueRequest() must call `request->addConstraint()`.
  */
 class ISchedulerConstraint : public ISchedulerNode
 {
@@ -29,30 +28,8 @@ public:
     /// Should be called outside of scheduling subsystem, implementation must be thread-safe.
     virtual void finishRequest(ResourceRequest * request) = 0;
 
-    void setParent(ISchedulerNode * parent_) override
-    {
-        ISchedulerNode::setParent(parent_);
-
-        // Assign `parent_constraint` to the nearest parent derived from ISchedulerConstraint
-        for (ISchedulerNode * node = parent_; node != nullptr; node = node->parent)
-        {
-            if (auto * constraint = dynamic_cast<ISchedulerConstraint *>(node))
-            {
-                parent_constraint = constraint;
-                break;
-            }
-        }
-    }
-
     /// For introspection of current state (true = satisfied, false = violated)
     virtual bool isSatisfied() = 0;
-
-protected:
-    // Reference to nearest parent that is also derived from ISchedulerConstraint.
-    // Request can traverse through multiple constraints while being dequeue from hierarchy,
-    // while finishing request should traverse the same chain in reverse order.
-    // NOTE: it must be immutable after initialization, because it is accessed in not thread-safe way from finishRequest()
-    ISchedulerConstraint * parent_constraint = nullptr;
 };
 
 }

src/Common/Scheduler/Nodes/SemaphoreConstraint.h

@@ -69,10 +69,7 @@ public:
         if (!request)
             return {nullptr, false};
 
-        // Request has reference to the first (closest to leaf) `constraint`, which can have `parent_constraint`.
-        // The former is initialized here dynamically and the latter is initialized once during hierarchy construction.
-        if (!request->constraint)
-            request->constraint = this;
+        request->addConstraint(this);
 
         // Update state on request arrival
         std::unique_lock lock(mutex);
@@ -87,10 +84,6 @@ public:
 
     void finishRequest(ResourceRequest * request) override
     {
-        // Recursive traverse of parent flow controls in reverse order
-        if (parent_constraint)
-            parent_constraint->finishRequest(request);
-
         // Update state on request departure
         std::unique_lock lock(mutex);
         bool was_active = active();

src/Common/Scheduler/Nodes/ThrottlerConstraint.h

@@ -79,10 +79,7 @@ public:
         if (!request)
             return {nullptr, false};
 
-        // Request has reference to the first (closest to leaf) `constraint`, which can have `parent_constraint`.
-        // The former is initialized here dynamically and the latter is initialized once during hierarchy construction.
-        if (!request->constraint)
-            request->constraint = this;
+        // We don't do `request->addConstraint(this)` because `finishRequest()` is no-op
 
         updateBucket(request->cost);
 
@@ -93,12 +90,8 @@ public:
         return {request, active()};
     }
 
-    void finishRequest(ResourceRequest * request) override
+    void finishRequest(ResourceRequest *) override
     {
-        // Recursive traverse of parent flow controls in reverse order
-        if (parent_constraint)
-            parent_constraint->finishRequest(request);
-
         // NOTE: Token-bucket constraint does not require any action when consumption ends
     }
 

src/Common/Scheduler/ResourceRequest.cpp

@@ -1,13 +1,42 @@
 #include <Common/Scheduler/ResourceRequest.h>
 #include <Common/Scheduler/ISchedulerConstraint.h>
 
+#include <Common/Exception.h>
+
+#include <ranges>
+
 namespace DB
 {
 
+namespace ErrorCodes
+{
+    extern const int LOGICAL_ERROR;
+}
+
 void ResourceRequest::finish()
 {
-    if (constraint)
-        constraint->finishRequest(this);
+    // Iterate over constraints in reverse order
+    for (ISchedulerConstraint * constraint : std::ranges::reverse_view(constraints))
+    {
+        if (constraint)
+            constraint->finishRequest(this);
+    }
+}
+
+void ResourceRequest::addConstraint(ISchedulerConstraint * new_constraint)
+{
+    for (auto & constraint : constraints)
+    {
+        if (!constraint)
+        {
+            constraint = new_constraint;
+            return;
+        }
+    }
+    // TODO(serxa): is it possible to validate it during enqueue of resource request to avoid LOGICAL_ERRORs in the scheduler thread? possible but will not cover case of moving queue with requests inside to invalid position
+    throw Exception(ErrorCodes::LOGICAL_ERROR,
+        "Max number of simultaneous workload constraints exceeded ({}). Remove extra constraints before using this workload.",
+        ResourceMaxConstraints);
 }
 
 }

src/Common/Scheduler/ResourceRequest.h

@@ -2,6 +2,7 @@
 
 #include <boost/intrusive/list.hpp>
 #include <base/types.h>
+#include <array>
 #include <limits>
 
 namespace DB
@@ -15,6 +16,10 @@ class ISchedulerConstraint;
 using ResourceCost = Int64;
 constexpr ResourceCost ResourceCostMax = std::numeric_limits<int>::max();
 
+// TODO(serxa): validate hierarchy to avoid too many constrants
+/// Max number of constraints for a request to pass though (depth of constaints chain)
+constexpr size_t ResourceMaxConstraints = 8;
+
 /*
  * Request for a resource consumption. The main moving part of the scheduling subsystem.
  * Resource requests processing workflow:
@@ -49,9 +54,10 @@ public:
     /// NOTE: If cost is not known in advance, ResourceBudget should be used (note that every ISchedulerQueue has it)
     ResourceCost cost;
 
-    /// Scheduler node to be notified on consumption finish
-    /// Auto-filled during request enqueue/dequeue
-    ISchedulerConstraint * constraint;
+    /// Scheduler nodes to be notified on consumption finish
+    /// Auto-filled during request dequeue
+    /// Vector is not used to avoid allocations in the scheduler thread
+    std::array<ISchedulerConstraint *, ResourceMaxConstraints> constraints;
 
     explicit ResourceRequest(ResourceCost cost_ = 1)
     {
@@ -62,7 +68,8 @@ public:
     void reset(ResourceCost cost_)
     {
         cost = cost_;
-        constraint = nullptr;
+        for (auto & constraint : constraints)
+            constraint = nullptr;
         // Note that list_base_hook should be reset independently (by intrusive list)
     }
 
@@ -79,6 +86,9 @@ public:
     /// ResourceRequest should not be destructed or reset before calling to `finish()`.
     /// WARNING: this function MUST not be called if request was canceled.
     void finish();
+
+    /// Is called from the scheduler thread to fill `constraints` chain
+    void addConstraint(ISchedulerConstraint * new_constraint);
 };
 
 }