Merge 4d48f87 into 946c37e

Author: ssmike

ydb/core/kqp/runtime/kqp_compute_scheduler.cpp

@@ -126,6 +126,8 @@ class TSchedulerEntity {
         std::atomic<i64> DelayedSumBatches = 0;
         std::atomic<i64> DelayedCount = 0;
 
+        double Share;
+
         TMultithreadPublisher<TGroupMutableStats> MutableStats;
     };
 
@@ -215,13 +217,16 @@ struct TComputeScheduler::TImpl {
 
     TDuration MaxDelay = TDuration::Seconds(10);
 
-    void AssignWeights() { }
+    void AssignWeights() {
+        for (auto& record : Records) {
+            record->MutableStats.Next()->Weight = SumCores * record->Share;
+        }
+    }
 
-    void CreateGroup(TString groupName, double maxShare, NMonotonic::TMonotonic now) {
+    void CreateGroup(TString groupName, double maxShare) {
         PoolId[groupName] = Records.size();
         auto group = std::make_unique<TSchedulerEntity::TGroupRecord>();
-        group->MutableStats.Next()->LastNowRecalc = now;
-        group->MutableStats.Next()->Weight = maxShare;
+        group->Share = maxShare;
         Records.push_back(std::move(group));
     }
 };
@@ -369,14 +374,18 @@ bool TComputeScheduler::Disable(TString group, TMonotonic now) {
 void TComputeScheduler::UpdateMaxShare(TString group, double share, TMonotonic now) {
     auto ptr = Impl->PoolId.FindPtr(group);
     if (!ptr) {
-        Impl->CreateGroup(group, share, now);
+        Impl->CreateGroup(group, share);
     } else {
         auto& record = Impl->Records[*ptr];
-        record->MutableStats.Next()->Weight = share;
+        record->Share = share;
     }
+    Impl->AssignWeights();
     AdvanceTime(now);
 }
 
+void TComputeScheduler::SetCapacity(ui64 cores) {
+    Impl->SumCores = cores;
+}
 
 ::NMonitoring::TDynamicCounters::TCounterPtr TComputeScheduler::GetGroupUsageCounter(TString group) const {
     return Impl->Counters
@@ -407,6 +416,7 @@ class TSchedulerActor : public TActorBootstrapped<TSchedulerActor> {
         }
         Opts.Scheduler->SetForgetInterval(Opts.ForgetOverflowTimeout);
         Opts.Scheduler->ReportCounters(Opts.Counters);
+        SetCapacity();
     }
 
     void Bootstrap() {
@@ -420,6 +430,15 @@ class TSchedulerActor : public TActorBootstrapped<TSchedulerActor> {
         Become(&TSchedulerActor::State);
     }
 
+    void SetCapacity() {
+        NActors::TExecutorPoolStats poolStats;
+        TVector<NActors::TExecutorThreadStats> threadsStats;
+        TlsActivationContext->ActorSystem()->GetPoolStats(SelfId().PoolID(), poolStats, threadsStats);
+        Y_ENSURE(poolStats.MaxThreadCount > 0);
+        Opts.Counters->SchedulerCapacity->Set(poolStats.MaxThreadCount);
+        Opts.Scheduler->SetCapacity(poolStats.MaxThreadCount);
+    }
+
     STATEFN(State) {
         switch (ev->GetTypeRewrite()) {
             hFunc(NConsole::TEvConfigsDispatcher::TEvSetConfigSubscriptionResponse, Handle);
@@ -467,6 +486,7 @@ class TSchedulerActor : public TActorBootstrapped<TSchedulerActor> {
     }
 
     void Handle(TEvents::TEvWakeup::TPtr&) {
+        SetCapacity();
         Opts.Scheduler->AdvanceTime(TlsActivationContext->Monotonic());
         Schedule(Opts.AdvanceTimeInterval, new TEvents::TEvWakeup());
     }

ydb/core/kqp/runtime/kqp_compute_scheduler.h

@@ -62,6 +62,8 @@ class TComputeScheduler {
     void ReportCounters(TIntrusivePtr<TKqpCounters>);
 
     void UpdateMaxShare(TString, double, TMonotonic now);
+    
+    void SetCapacity(ui64 cores);
 
     void SetMaxDeviation(TDuration);
     void SetForgetInterval(TDuration);