Per query cpu limits (#8520)

Author: ssmike

ydb/core/kqp/counters/kqp_counters.cpp

@@ -830,6 +830,8 @@ TKqpCounters::TKqpCounters(const ::NMonitoring::TDynamicCounterPtr& counters, co
     FullScansExecuted = KqpGroup->GetCounter("FullScans", true);
 
     SchedulerThrottled = KqpGroup->GetCounter("NodeScheduler/ThrottledUs", true);
+    SchedulerGroupsCount = KqpGroup->GetCounter("NodeScheduler/GroupsCount", false);
+    SchedulerValuesCount = KqpGroup->GetCounter("NodeScheduler/ValuesCount", false);
     SchedulerCapacity = KqpGroup->GetCounter("NodeScheduler/Capacity");
     ComputeActorExecutions = KqpGroup->GetHistogram("NodeScheduler/BatchUs", NMonitoring::ExponentialHistogram(20, 2, 1));
     ComputeActorDelays = KqpGroup->GetHistogram("NodeScheduler/Delays", NMonitoring::ExponentialHistogram(20, 2, 1));

ydb/core/kqp/counters/kqp_counters.h

@@ -416,6 +416,8 @@ class TKqpCounters : public TKqpCountersBase, public NYql::NDq::TSpillingCounter
     NMonitoring::THistogramPtr ComputeActorDelays;
     ::NMonitoring::TDynamicCounters::TCounterPtr ThrottledActorsSpuriousActivations;
     NMonitoring::THistogramPtr SchedulerDelays;
+    NMonitoring::TDynamicCounters::TCounterPtr SchedulerGroupsCount;
+    NMonitoring::TDynamicCounters::TCounterPtr SchedulerValuesCount;
 
     // Sequences counters
     ::NMonitoring::TDynamicCounters::TCounterPtr SequencerActorsCount;

ydb/core/kqp/executer_actor/kqp_planner.cpp

@@ -237,7 +237,10 @@ std::unique_ptr<TEvKqpNode::TEvStartKqpTasksRequest> TKqpPlanner::SerializeReque
     request.SetDatabase(Database);
     if (UserRequestContext->PoolConfig.has_value()) {
         request.SetMemoryPoolPercent(UserRequestContext->PoolConfig->QueryMemoryLimitPercentPerNode);
-        request.SetMaxCpuShare(UserRequestContext->PoolConfig->TotalCpuLimitPercentPerNode / 100.0);
+        request.SetPoolMaxCpuShare(UserRequestContext->PoolConfig->TotalCpuLimitPercentPerNode / 100.0);
+        if (UserRequestContext->PoolConfig->QueryCpuLimitPercentPerNode >= 0) {
+            request.SetQueryCpuShare(UserRequestContext->PoolConfig->QueryCpuLimitPercentPerNode / 100.0);
+        }
     }
 
     return result;

ydb/core/kqp/node_service/kqp_node_service.cpp

@@ -195,6 +195,26 @@ class TKqpNodeService : public TActorBootstrapped<TKqpNodeService> {
 
         auto schedulerNow = TlsActivationContext->Monotonic();
 
+        TString schedulerGroup = msg.GetSchedulerGroup();
+
+        if (SchedulerOptions.Scheduler->Disabled(schedulerGroup)) {
+            auto share = msg.GetPoolMaxCpuShare();
+            if (share <= 0 && msg.HasQueryCpuShare()) {
+                share = 1.0;
+            }
+            if (share > 0) {
+                Scheduler->UpdateGroupShare(schedulerGroup, share, schedulerNow);
+                Send(SchedulerActorId, new TEvSchedulerNewPool(msg.GetDatabase(), schedulerGroup));
+            } else {
+                schedulerGroup = "";
+            }
+        } 
+
+        std::optional<ui64> querySchedulerGroup;
+        if (msg.HasQueryCpuShare() && schedulerGroup) {
+            querySchedulerGroup = Scheduler->MakePerQueryGroup(schedulerNow, msg.GetQueryCpuShare(), schedulerGroup);
+        }
+
         // start compute actors
         TMaybe<NYql::NDqProto::TRlPath> rlPath = Nothing();
         if (msgRtSettings.HasRlPath()) {
@@ -208,30 +228,21 @@ class TKqpNodeService : public TActorBootstrapped<TKqpNodeService> {
 
         const ui32 tasksCount = msg.GetTasks().size();
         for (auto& dqTask: *msg.MutableTasks()) {
-            TString group = msg.GetSchedulerGroup();
-
             TComputeActorSchedulingOptions schedulingTaskOptions {
                 .Now = schedulerNow,
                 .SchedulerActorId = SchedulerActorId,
                 .Scheduler = Scheduler.get(),
-                .Group = group,
+                .Group = schedulerGroup,
                 .Weight = 1,
-                .NoThrottle = false,
+                .NoThrottle = schedulerGroup.empty(),
                 .Counters = Counters
             };
 
-            if (SchedulerOptions.Scheduler->Disabled(group)) {
-                auto share = msg.GetMaxCpuShare();
-                if (share > 0) {
-                    Scheduler->UpdateMaxShare(group, share, schedulerNow);
-                    Send(SchedulerActorId, new TEvSchedulerNewPool(msg.GetDatabase(), group, share));
-                } else {
-                    schedulingTaskOptions.NoThrottle = true;
-                }
-            } 
-
             if (!schedulingTaskOptions.NoThrottle) {
                 schedulingTaskOptions.Handle = SchedulerOptions.Scheduler->Enroll(schedulingTaskOptions.Group, schedulingTaskOptions.Weight, schedulingTaskOptions.Now);
+                if (querySchedulerGroup) {
+                    Scheduler->AddToGroup(schedulerNow, *querySchedulerGroup, schedulingTaskOptions.Handle);
+                }
             }
 
             auto result = CaFactory_->CreateKqpComputeActor({