Account usage from ColumnShards (#20245)

Author: abyss7

ydb/core/base/appdata_fwd.h

@@ -26,6 +26,9 @@ namespace NKikimr {
     namespace NJaegerTracing {
         class TSamplingThrottlingConfigurator;
     }
+    namespace NKqp::NScheduler {
+        class TComputeScheduler;
+    }
 }
 
 namespace NKikimrCms {
@@ -309,6 +312,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

@@ -300,7 +300,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/kqp_schedulable_actor.cpp

@@ -59,6 +59,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::IncreaseBurstThrottle(const TDuration& burstThrottle) {
     for (TTreeElementBase* parent = Query.get(); parent; parent = parent->Parent) {
         parent->BurstThrottle += burstThrottle.MicroSeconds();

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

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

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

@@ -51,6 +51,8 @@ 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;
         NMonitoring::THistogramPtr                 Delay;
     };
 

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

@@ -76,6 +76,9 @@ TPool::TPool(const TString& id, const TIntrusivePtr<TKqpCounters>& counters, con
     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);
+
     Counters.Delay     = group->GetHistogram("Delay",
         NMonitoring::ExplicitHistogram({10, 10e2, 10e3, 10e4, 10e5, 10e6, 10e7}), true); // TODO: make from MinDelay to MaxDelay.
 }
@@ -90,6 +93,9 @@ NSnapshot::TPool* TPool::TakeSnapshot() const {
     Counters.Usage->Set(BurstUsage);
     Counters.Throttle->Set(BurstThrottle);
 
+    Counters.InFlightExtra->Set(UsageExtra * 1'000'000);
+    Counters.UsageExtra->Set(BurstUsageExtra);
+
     for (const auto& child : Children) {
         newPool->AddQuery(std::shared_ptr<NSnapshot::TQuery>(std::dynamic_pointer_cast<TQuery>(child)->TakeSnapshot()));
     }

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/engines/reader/abstract/read_context.cpp

@@ -13,9 +13,10 @@ TReadContext::TReadContext(const std::shared_ptr<IStoragesManager>& storagesMana
     const std::shared_ptr<NDataAccessorControl::IDataAccessorsManager>& dataAccessorsManager,
     const NColumnShard::TConcreteScanCounters& counters, const TReadMetadataBase::TConstPtr& readMetadata, const TActorId& scanActorId,
     const TActorId& resourceSubscribeActorId, const TActorId& readCoordinatorActorId, const TComputeShardingPolicy& computeShardingPolicy,
-    const ui64 scanId, const NConveyorComposite::TCPULimitsConfig& cpuLimits)
+    const ui64 scanId, const NConveyorComposite::TCPULimitsConfig& cpuLimits, NKqp::NScheduler::TSchedulableTaskPtr schedulableTask)
     : StoragesManager(storagesManager)
     , DataAccessorsManager(dataAccessorsManager)
+    , SchedulableTask(std::move(schedulableTask))
     , Counters(counters)
     , ReadMetadata(readMetadata)
     , ResourcesTaskContext("CS::SCAN_READ", counters.ResourcesSubscriberCounters)

ydb/core/tx/columnshard/engines/reader/abstract/read_context.h

@@ -1,6 +1,7 @@
 #pragma once
 #include "read_metadata.h"
 
+#include <ydb/core/kqp/runtime/scheduler/new/kqp_schedulable_actor.h>
 #include <ydb/core/protos/tx_datashard.pb.h>
 #include <ydb/core/tx/columnshard/blobs_action/abstract/storages_manager.h>
 #include <ydb/core/tx/columnshard/columnshard_private_events.h>
@@ -50,6 +51,7 @@ class TReadContext {
 private:
     YDB_READONLY_DEF(std::shared_ptr<IStoragesManager>, StoragesManager);
     YDB_READONLY_DEF(std::shared_ptr<NDataAccessorControl::IDataAccessorsManager>, DataAccessorsManager);
+    YDB_READONLY_DEF(NKqp::NScheduler::TSchedulableTaskPtr, SchedulableTask);
     const NColumnShard::TConcreteScanCounters Counters;
     TReadMetadataBase::TConstPtr ReadMetadata;
     NResourceBroker::NSubscribe::TTaskContext ResourcesTaskContext;
@@ -151,7 +153,7 @@ class TReadContext {
         const std::shared_ptr<NDataAccessorControl::IDataAccessorsManager>& dataAccessorsManager,
         const NColumnShard::TConcreteScanCounters& counters, const TReadMetadataBase::TConstPtr& readMetadata, const TActorId& scanActorId,
         const TActorId& resourceSubscribeActorId, const TActorId& readCoordinatorActorId, const TComputeShardingPolicy& computeShardingPolicy,
-        const ui64 scanId, const NConveyorComposite::TCPULimitsConfig& cpuLimits);
+        const ui64 scanId, const NConveyorComposite::TCPULimitsConfig& cpuLimits, NKqp::NScheduler::TSchedulableTaskPtr schedulableTask);
 };
 
 class IDataReader {

ydb/core/tx/columnshard/engines/reader/actor/actor.cpp

@@ -21,7 +21,8 @@ TColumnShardScan::TColumnShardScan(const TActorId& columnShardActorId, const TAc
     const std::shared_ptr<NDataAccessorControl::IDataAccessorsManager>& dataAccessorsManager,
     const TComputeShardingPolicy& computeShardingPolicy, ui32 scanId, ui64 txId, ui32 scanGen, ui64 requestCookie, ui64 tabletId,
     TDuration timeout, const TReadMetadataBase::TConstPtr& readMetadataRange, NKikimrDataEvents::EDataFormat dataFormat,
-    const NColumnShard::TScanCounters& scanCountersPool, const NConveyorComposite::TCPULimitsConfig& cpuLimits)
+    const NColumnShard::TScanCounters& scanCountersPool, const NConveyorComposite::TCPULimitsConfig& cpuLimits,
+    NKqp::NScheduler::TSchedulableTaskPtr schedulableTask)
     : StoragesManager(storagesManager)
     , DataAccessorsManager(dataAccessorsManager)
     , ColumnShardActorId(columnShardActorId)
@@ -34,6 +35,7 @@ TColumnShardScan::TColumnShardScan(const TActorId& columnShardActorId, const TAc
     , DataFormat(dataFormat)
     , TabletId(tabletId)
     , CPULimits(cpuLimits)
+    , SchedulableTask(std::move(schedulableTask))
     , ReadMetadataRange(readMetadataRange)
     , Timeout(timeout ? timeout : COMPUTE_HARD_TIMEOUT)
     , ScanCountersPool(scanCountersPool, TValidator::CheckNotNull(ReadMetadataRange)->GetProgram().GetGraphOptional())
@@ -54,7 +56,8 @@ void TColumnShardScan::Bootstrap(const TActorContext& ctx) {
     ReadCoordinatorActorId = ctx.Register(new NBlobOperations::NRead::TReadCoordinatorActor(TabletId, SelfId()));
 
     std::shared_ptr<TReadContext> context = std::make_shared<TReadContext>(StoragesManager, DataAccessorsManager, ScanCountersPool,
-        ReadMetadataRange, SelfId(), ResourceSubscribeActorId, ReadCoordinatorActorId, ComputeShardingPolicy, ScanId, CPULimits);
+        ReadMetadataRange, SelfId(), ResourceSubscribeActorId, ReadCoordinatorActorId, ComputeShardingPolicy, ScanId, CPULimits,
+        std::move(SchedulableTask));
     ScanIterator = ReadMetadataRange->StartScan(context);
     auto startResult = ScanIterator->Start();
     StartInstant = TMonotonic::Now();

ydb/core/tx/columnshard/engines/reader/actor/actor.h

@@ -35,7 +35,8 @@ class TColumnShardScan: public TActorBootstrapped<TColumnShardScan>,
         const std::shared_ptr<NDataAccessorControl::IDataAccessorsManager>& dataAccessorsManager,
         const TComputeShardingPolicy& computeShardingPolicy, ui32 scanId, ui64 txId, ui32 scanGen, ui64 requestCookie, ui64 tabletId,
         TDuration timeout, const TReadMetadataBase::TConstPtr& readMetadataRange, NKikimrDataEvents::EDataFormat dataFormat,
-        const NColumnShard::TScanCounters& scanCountersPool, const NConveyorComposite::TCPULimitsConfig& cpuLimits);
+        const NColumnShard::TScanCounters& scanCountersPool, const NConveyorComposite::TCPULimitsConfig& cpuLimits,
+        NKqp::NScheduler::TSchedulableTaskPtr schedulableTask);
 
     void Bootstrap(const TActorContext& ctx);
 
@@ -131,6 +132,7 @@ class TColumnShardScan: public TActorBootstrapped<TColumnShardScan>,
     const NKikimrDataEvents::EDataFormat DataFormat;
     const ui64 TabletId;
     const NConveyorComposite::TCPULimitsConfig CPULimits;
+    NKqp::NScheduler::TSchedulableTaskPtr SchedulableTask;
 
     TReadMetadataBase::TConstPtr ReadMetadataRange;
     std::unique_ptr<TScanIteratorBase> ScanIterator;

ydb/core/tx/columnshard/engines/reader/common_reader/iterator/fetching.cpp

@@ -5,6 +5,7 @@
 #include "source.h"
 #include "sub_columns_fetching.h"
 
+#include <ydb/core/kqp/runtime/scheduler/new/kqp_schedulable_actor.h>
 #include <ydb/core/tx/columnshard/blobs_reader/actor.h>
 
 #include <util/string/builder.h>
@@ -65,9 +66,20 @@ TConclusion<bool> TFetchingScriptCursor::Execute(const std::shared_ptr<IDataSour
             IS_DEBUG_LOG_ENABLED(NKikimrServices::TX_COLUMNSHARD_SCAN_MEMORY));
         AFL_DEBUG(NKikimrServices::TX_COLUMNSHARD_SCAN)("scan_step", step->DebugString())("scan_step_idx", CurrentStepIdx)("source_id", source->GetSourceId());
 
+        auto& schedulableTask = source->GetContext()->GetCommonContext()->GetSchedulableTask();
+        if (schedulableTask) {
+            schedulableTask->IncreaseExtraUsage();
+        }
+
         const TMonotonic startInstant = TMonotonic::Now();
         const TConclusion<bool> resultStep = step->ExecuteInplace(source, *this);
-        FlushDuration(TMonotonic::Now() - startInstant);
+        const auto executionTime = TMonotonic::Now() - startInstant;
+
+        if (schedulableTask) {
+            schedulableTask->DecreaseExtraUsage(executionTime);
+        }
+
+        FlushDuration(executionTime);
         if (!resultStep) {
             return resultStep;
         }

ydb/core/tx/columnshard/engines/reader/transaction/tx_internal_scan.cpp

@@ -87,7 +87,7 @@ void TTxInternalScan::Complete(const TActorContext& ctx) {
     auto scanActorId = ctx.Register(new TColumnShardScan(Self->SelfId(), scanComputeActor, Self->GetStoragesManager(),
         Self->DataAccessorsManager.GetObjectPtrVerified(),
         TComputeShardingPolicy(), ScanId, LockId.value_or(0), ScanGen, requestCookie, Self->TabletID(), TDuration::Max(), readMetadataRange,
-        NKikimrDataEvents::FORMAT_ARROW, Self->Counters.GetScanCounters(), {}));
+        NKikimrDataEvents::FORMAT_ARROW, Self->Counters.GetScanCounters(), {}, {}));
 
     Self->InFlightReadsTracker.AddScanActorId(requestCookie, scanActorId);
     AFL_DEBUG(NKikimrServices::TX_COLUMNSHARD_SCAN)("event", "TTxInternalScan started")("actor_id", scanActorId)("trace_detailed", detailedInfo);

ydb/core/tx/columnshard/engines/reader/transaction/tx_scan.cpp

@@ -1,6 +1,7 @@
 #include "tx_scan.h"
 
 #include <ydb/core/formats/arrow/arrow_batch_builder.h>
+#include <ydb/core/kqp/runtime/scheduler/new/kqp_compute_scheduler_service.h>
 #include <ydb/core/sys_view/common/schema.h>
 #include <ydb/core/tx/columnshard/engines/reader/actor/actor.h>
 #include <ydb/core/tx/columnshard/engines/reader/plain_reader/constructor/constructor.h>
@@ -161,9 +162,12 @@ void TTxScan::Complete(const TActorContext& ctx) {
     TComputeShardingPolicy shardingPolicy;
     AFL_VERIFY(shardingPolicy.DeserializeFromProto(request.GetComputeShardingPolicy()));
 
+    const auto& scheduler = AppData(ctx)->ComputeScheduler;
+    auto schedulableTask = scheduler ? scheduler->CreateSchedulableTaskFactory()(txId) : nullptr;
+
     auto scanActorId = ctx.Register(new TColumnShardScan(Self->SelfId(), scanComputeActor, Self->GetStoragesManager(),
         Self->DataAccessorsManager.GetObjectPtrVerified(), shardingPolicy, scanId, txId, scanGen, requestCookie, Self->TabletID(), timeout,
-        readMetadataRange, dataFormat, Self->Counters.GetScanCounters(), cpuLimits));
+        readMetadataRange, dataFormat, Self->Counters.GetScanCounters(), cpuLimits, std::move(schedulableTask)));
     Self->InFlightReadsTracker.AddScanActorId(requestCookie, scanActorId);
 
     AFL_DEBUG(NKikimrServices::TX_COLUMNSHARD_SCAN)("event", "TTxScan started")("actor_id", scanActorId)("trace_detailed", detailedInfo);