Account usage from ColumnShard scans in KQP scheduler

Author: abyss7

ydb/core/base/appdata_fwd.h

@@ -25,6 +25,9 @@ namespace NKikimr {
     namespace NJaegerTracing {
         class TSamplingThrottlingConfigurator;
     }
+    namespace NKqp::NScheduler {
+        class TComputeScheduler;
+    }
 }
 
 namespace NKikimrCms {
@@ -303,6 +306,8 @@ struct TAppData {
     // Tracing configurator (look for tracing config in ydb/core/jaeger_tracing/actors_tracing_control)
     TIntrusivePtr<NKikimr::NJaegerTracing::TSamplingThrottlingConfigurator> TracingConfigurator;
 
+    std::shared_ptr<NKqp::NScheduler::TComputeScheduler> ComputeScheduler;
+
     TAppData(
             ui32 sysPoolId, ui32 userPoolId, ui32 ioPoolId, ui32 batchPoolId,
             TMap<TString, ui32> servicePools,

ydb/core/kqp/proxy_service/kqp_proxy_service.cpp

@@ -297,7 +297,7 @@ class TKqpProxyService : public TActorBootstrapped<TKqpProxyService> {
                 MakeKqpCompileComputationPatternServiceID(SelfId().NodeId()), CompileComputationPatternService);
         }
 
-        auto scheduler = std::make_shared<NScheduler::TComputeScheduler>(Counters);
+        auto scheduler = AppData()->ComputeScheduler = std::make_shared<NScheduler::TComputeScheduler>(Counters);
 
         ResourceManager_ = GetKqpResourceManager();
         CaFactory_ = NComputeActor::MakeKqpCaFactory(

ydb/core/kqp/runtime/scheduler/new/fwd.h

@@ -43,7 +43,7 @@ namespace NKikimr::NKqp::NScheduler {
     }
 
     struct TSchedulableTask;
-    using TSchedulableTaskPtr = THolder<TSchedulableTask>;
+    using TSchedulableTaskPtr = std::shared_ptr<TSchedulableTask>;
     using TSchedulableTaskFactory = std::function<TSchedulableTaskPtr(const NHdrf::TQueryId&)>;
 
 } // namespace NKikimr::NKqp::NScheduler

ydb/core/kqp/runtime/scheduler/new/kqp_compute_scheduler_service.cpp

@@ -194,7 +194,7 @@ TComputeScheduler::TComputeScheduler(TIntrusivePtr<TKqpCounters> counters)
 
 TSchedulableTaskFactory TComputeScheduler::CreateSchedulableTaskFactory() {
     return [ptr = this->shared_from_this()](const NHdrf::TQueryId& queryId) {
-        return MakeHolder<TSchedulableTask>(ptr->GetQuery(queryId));
+        return std::make_shared<TSchedulableTask>(ptr->GetQuery(queryId));
     };
 }
 

ydb/core/kqp/runtime/scheduler/new/kqp_schedulable_actor.cpp

@@ -57,6 +57,19 @@ void TSchedulableTask::DecreaseUsage(const TDuration& burstUsage) {
     }
 }
 
+void TSchedulableTask::IncreaseExtraUsage() {
+    for (TTreeElementBase* parent = Query.get(); parent; parent = parent->Parent) {
+        ++parent->UsageExtra;
+    }
+}
+
+void TSchedulableTask::DecreaseExtraUsage(const TDuration& burstUsageExtra) {
+    for (TTreeElementBase* parent = Query.get(); parent; parent = parent->Parent) {
+        --parent->UsageExtra;
+        parent->BurstUsageExtra += burstUsageExtra.MicroSeconds();
+    }
+}
+
 void TSchedulableTask::IncreaseThrottle() {
     for (TTreeElementBase* parent = Query.get(); parent; parent = parent->Parent) {
         ++parent->Throttle;

ydb/core/kqp/runtime/scheduler/new/kqp_schedulable_actor.h

@@ -18,6 +18,10 @@ struct TSchedulableTask {
     void IncreaseUsage(const TDuration& burstThrottle);
     void DecreaseUsage(const TDuration& burstUsage);
 
+    // Account extra usage which doesn't affect scheduling
+    void IncreaseExtraUsage();
+    void DecreaseExtraUsage(const TDuration& burstUsage);
+
     void IncreaseThrottle();
     void DecreaseThrottle();
 

ydb/core/kqp/runtime/scheduler/new/tree/common.h

@@ -51,6 +51,9 @@ namespace NKikimr::NKqp::NScheduler::NHdrf {
         NMonitoring::TDynamicCounters::TCounterPtr FairShare;
         NMonitoring::TDynamicCounters::TCounterPtr InFlight;
         NMonitoring::TDynamicCounters::TCounterPtr Waiting;
+
+        NMonitoring::TDynamicCounters::TCounterPtr InFlightExtra;
+        NMonitoring::TDynamicCounters::TCounterPtr UsageExtra;
     };
 
 } // namespace NKikimr::NKqp::NScheduler::NHdrf

ydb/core/kqp/runtime/scheduler/new/tree/dynamic.cpp

@@ -75,6 +75,9 @@ TPool::TPool(const TString& id, const TIntrusivePtr<TKqpCounters>& counters, con
     Counters.Usage     = group->GetCounter("Usage",     true);
     Counters.Throttle  = group->GetCounter("Throttle",  true);
     Counters.FairShare = group->GetCounter("FairShare", true);  // snapshot
+
+    Counters.InFlightExtra = group->GetCounter("InFlightExtra",  false);
+    Counters.UsageExtra    = group->GetCounter("UsageExtra",     true);
 }
 
 NSnapshot::TPool* TPool::TakeSnapshot() const {

ydb/core/kqp/runtime/scheduler/new/tree/dynamic.h

@@ -34,10 +34,12 @@ namespace NKikimr::NKqp::NScheduler::NHdrf::NDynamic {
 
     struct TTreeElementBase : public TStaticAttributes {
         std::atomic<ui64> Usage = 0;
+        std::atomic<ui64> UsageExtra = 0;
         std::atomic<ui64> Demand = 0;
         std::atomic<ui64> Throttle = 0;
 
         std::atomic<ui64> BurstUsage = 0;
+        std::atomic<ui64> BurstUsageExtra = 0;
         std::atomic<ui64> BurstThrottle = 0;
 
         TTreeElementBase* Parent = nullptr;

ydb/core/tx/columnshard/columnshard_impl.cpp

@@ -1213,11 +1213,13 @@ class TTxAskPortionChunks: public TTransactionBase<TColumnShard> {
 
         AFL_TRACE(NKikimrServices::TX_COLUMNSHARD)("stage", "finished");
         NConveyorComposite::TScanServiceOperator::SendTaskToExecute(
-            std::make_shared<TAccessorsParsingTask>(FetchCallback, std::move(FetchedAccessors)), 0);
+            std::make_shared<TAccessorsParsingTask>(FetchCallback, std::move(FetchedAccessors)), 0, {}); // TODO(scheduler)
         return true;
     }
+
     void Complete(const TActorContext& /*ctx*/) override {
     }
+
     TTxType GetTxType() const override {
         return TXTYPE_ASK_PORTION_METADATA;
     }