Merge 21318b9 into 2f95ed9

Author: gridnevvvit

ydb/core/kqp/compute_actor/kqp_compute_actor_factory.cpp

@@ -14,15 +14,13 @@ struct TMemoryQuotaManager : public NYql::NDq::TGuaranteeQuotaManager {
         , std::shared_ptr<IKqpNodeState> state
         , TIntrusivePtr<NRm::TTxState> tx
         , TIntrusivePtr<NRm::TTaskState> task
-        , ui64 limit
-        , ui64 reasonableSpillingTreshold)
+        , ui64 limit)
     : NYql::NDq::TGuaranteeQuotaManager(limit, limit)
     , ResourceManager(std::move(resourceManager))
     , MemoryPool(memoryPool)
     , State(std::move(state))
     , Tx(std::move(tx))
     , Task(std::move(task))
-    , ReasonableSpillingTreshold(reasonableSpillingTreshold)
     {
     }
 
@@ -57,7 +55,7 @@ struct TMemoryQuotaManager : public NYql::NDq::TGuaranteeQuotaManager {
     }
 
     bool IsReasonableToUseSpilling() const override {
-        return Tx->GetExtraMemoryAllocatedSize() >= ReasonableSpillingTreshold;
+        return Task->IsReasonableToStartSpilling();
     }
 
     TString MemoryConsumptionDetails() const override {
@@ -88,7 +86,6 @@ class TKqpCaFactory : public IKqpNodeComputeActorFactory {
     std::atomic<ui64> MkqlLightProgramMemoryLimit = 0;
     std::atomic<ui64> MkqlHeavyProgramMemoryLimit = 0;
     std::atomic<ui64> MinChannelBufferSize = 0;
-    std::atomic<ui64> ReasonableSpillingTreshold = 0;
 
 public:
     TKqpCaFactory(const NKikimrConfig::TTableServiceConfig::TResourceManager& config,
@@ -107,7 +104,6 @@ class TKqpCaFactory : public IKqpNodeComputeActorFactory {
         MkqlLightProgramMemoryLimit.store(config.GetMkqlLightProgramMemoryLimit());
         MkqlHeavyProgramMemoryLimit.store(config.GetMkqlHeavyProgramMemoryLimit());
         MinChannelBufferSize.store(config.GetMinChannelBufferSize());
-        ReasonableSpillingTreshold.store(config.GetReasonableSpillingTreshold());
     }
 
     TActorStartResult CreateKqpComputeActor(TCreateArgs&& args) override {
@@ -158,8 +154,7 @@ class TKqpCaFactory : public IKqpNodeComputeActorFactory {
             std::move(args.State),
             std::move(args.TxInfo),
             std::move(task),
-            limit,
-            ReasonableSpillingTreshold.load());
+            limit);
 
         auto runtimeSettings = args.RuntimeSettings;
         runtimeSettings.ExtraMemoryAllocationPool = args.MemoryPool;

ydb/core/kqp/compute_actor/kqp_compute_actor_factory.h

@@ -122,6 +122,8 @@ struct IKqpNodeComputeActorFactory {
         const TInstant& Deadline;
         const bool ShareMailbox;
         const TMaybe<NYql::NDqProto::TRlPath>& RlPath;
+        const TString& PoolId;
+
         TComputeStagesWithScan* ComputesByStages = nullptr;
         std::shared_ptr<IKqpNodeState> State = nullptr;
         TComputeActorSchedulingOptions SchedulingOptions = {};

ydb/core/kqp/executer_actor/kqp_planner.cpp

@@ -223,6 +223,7 @@ std::unique_ptr<TEvKqpNode::TEvStartKqpTasksRequest> TKqpPlanner::SerializeReque
     }
 
     request.SetSchedulerGroup(UserRequestContext->PoolId);
+    request.SetMemoryPoolPercent(UserRequestContext->PoolConfig->QueryMemoryLimitPercentPerNode);
 
     return result;
 }
@@ -351,7 +352,8 @@ TString TKqpPlanner::ExecuteDataComputeTask(ui64 taskId, ui32 computeTasksSize)
     NYql::NDq::TComputeRuntimeSettings settings;
     if (!TxInfo) {
         TxInfo = MakeIntrusive<NRm::TTxState>(
-            TxId, TInstant::Now(), ResourceManager_->GetCounters());
+            TxId, TInstant::Now(), ResourceManager_->GetCounters(),
+            UserRequestContext->PoolId, UserRequestContext->PoolConfig->QueryMemoryLimitPercentPerNode);
     }
 
     auto startResult = CaFactory_->CreateKqpComputeActor({
@@ -370,7 +372,8 @@ TString TKqpPlanner::ExecuteDataComputeTask(ui64 taskId, ui32 computeTasksSize)
         .StatsMode = GetDqStatsMode(StatsMode),
         .Deadline = Deadline,
         .ShareMailbox = (computeTasksSize <= 1),
-        .RlPath = Nothing()
+        .RlPath = Nothing(),
+        .PoolId = UserRequestContext->PoolId
     });
 
     if (const auto* rmResult = std::get_if<NRm::TKqpRMAllocateResult>(&startResult)) {

ydb/core/kqp/node_service/kqp_node_service.cpp

@@ -205,7 +205,8 @@ class TKqpNodeService : public TActorBootstrapped<TKqpNodeService> {
         }
 
         TIntrusivePtr<NRm::TTxState> txInfo = MakeIntrusive<NRm::TTxState>(
-            txId, TInstant::Now(), ResourceManager_->GetCounters());
+            txId, TInstant::Now(), ResourceManager_->GetCounters(),
+            msg.GetSchedulerGroup(), msg.GetMemoryPoolPercent());
 
         const ui32 tasksCount = msg.GetTasks().size();
         for (auto& dqTask: *msg.MutableTasks()) {
@@ -244,9 +245,10 @@ class TKqpNodeService : public TActorBootstrapped<TKqpNodeService> {
                 .Deadline = TInstant(),
                 .ShareMailbox = false,
                 .RlPath = rlPath,
+                .PoolId = schedulingOptions.Group,
                 .ComputesByStages = &computesByStage,
                 .State = State_,
-                .SchedulingOptions = std::move(schedulingOptions)
+                .SchedulingOptions = std::move(schedulingOptions),
             });
 
             if (const auto* rmResult = std::get_if<NRm::TKqpRMAllocateResult>(&result)) {

ydb/core/kqp/rm_service/kqp_rm_service.cpp

@@ -20,6 +20,8 @@
 
 #include <ydb/library/yql/utils/yql_panic.h>
 
+#include <library/cpp/containers/absl_flat_hash/flat_hash_map.h>
+
 namespace NKikimr {
 namespace NKqp {
 namespace NRm {
@@ -44,11 +46,14 @@ using namespace NResourceBroker;
 namespace {
 
 template <typename T>
-class TLimitedResource {
+class TLimitedResource : public TAtomicRefCount<TLimitedResource<T>> {
 public:
-    explicit TLimitedResource(T limit)
+    explicit TLimitedResource(T limit, ui64 overPercent)
         : Limit(limit)
-        , Used(0) {}
+        , Used(0)
+        , OverPercent(overPercent)
+        , ResourceIsOverCookie(MakeIntrusive<TMemoryResourceCookie>())
+    {}
 
     T Available() const {
         return Limit > Used ? Limit - Used : 0;
@@ -58,26 +63,38 @@ class TLimitedResource {
         return Available() >= amount;
     }
 
-    bool Acquire(T value) {
+    bool AcquireIfAvailable(T value) {
         if (Available() >= value) {
             Used += value;
+            UpdateCookie();
             return true;
         }
         return false;
     }
 
+    void UpdateCookie() {
+        ui64 availableTreshold = GetLimit() / 100 * (100 - OverPercent);
+        ResourceIsOverCookie->SpillingPercentReached.store(Available() < availableTreshold);
+    }
+
     void Release(T value) {
         if (Used > value) {
             Used -= value;
         } else {
             Used = 0;
         }
+
+        UpdateCookie();
     }
 
     void SetNewLimit(T limit) {
         Limit = limit;
     }
 
+    void SetOverPercent(ui64 overPercent) {
+        OverPercent = overPercent;
+    }
+
     T GetLimit() const {
         return Limit;
     }
@@ -89,6 +106,11 @@ class TLimitedResource {
 private:
     T Limit;
     T Used;
+
+    ui64 OverPercent = 100;
+
+public:
+    TIntrusivePtr<TMemoryResourceCookie> ResourceIsOverCookie;
 };
 
 struct TEvPrivate {
@@ -112,7 +134,8 @@ class TKqpResourceManager : public IKqpResourceManager {
         : Counters(counters)
         , ExecutionUnitsResource(config.GetComputeActorsCount())
         , ExecutionUnitsLimit(config.GetComputeActorsCount())
-        , ScanQueryMemoryResource(config.GetQueryMemoryLimit())
+        , SpillingPercent(config.GetSpillingPrecent())
+        , TotalMemoryResource(MakeIntrusive<TLimitedResource<ui64>>(config.GetQueryMemoryLimit(), config.GetSpillingPrecent()))
     {
         SetConfigValues(config);
     }
@@ -188,11 +211,7 @@ class TKqpResourceManager : public IKqpResourceManager {
         }
 
         bool hasScanQueryMemory = true;
-        ui64 queryMemoryLimit = 0;
 
-        // NOTE(gvit): the first memory request always satisfied.
-        // all other requests are not guaranteed to be satisfied.
-        // In the nearest future we need to implement several layers of memory requests.
         bool isFirstAllocationRequest = (resources.ExecutionUnits > 0 && resources.MemoryPool == EKqpMemoryPool::DataQuery);
         if (isFirstAllocationRequest) {
             TKqpResourcesRequest newRequest = resources;
@@ -210,17 +229,41 @@ class TKqpResourceManager : public IKqpResourceManager {
                 return result;
             }
 
-            hasScanQueryMemory = ScanQueryMemoryResource.Has(resources.Memory);
-            if (hasScanQueryMemory) {
-                ScanQueryMemoryResource.Acquire(resources.Memory);
-                queryMemoryLimit = QueryMemoryLimit.load();
+            hasScanQueryMemory = TotalMemoryResource->AcquireIfAvailable(resources.Memory);
+            task->TotalMemoryCookie = TotalMemoryResource->ResourceIsOverCookie;
+
+            if (hasScanQueryMemory && !tx->PoolId.empty() && tx->MemoryPoolPercent > 0) {
+                ui64 poolLimit = 1_MB;
+
+                {
+                    double poolLimitDouble = static_cast<double>(TotalMemoryResource->GetLimit()) * tx->MemoryPoolPercent;
+                    poolLimit = std::max(poolLimit, static_cast<ui64>(poolLimitDouble));
+                }
+
+                auto [it, success] = MemoryNamedPools.emplace(tx->PoolId, nullptr);
+
+                if (success) {
+                    it->second = MakeIntrusive<TLimitedResource<ui64>>(poolLimit, SpillingPercent.load());
+                } else if (it->second->GetLimit() != poolLimit) {
+                    it->second->SetNewLimit(poolLimit);
+                    it->second->SetOverPercent(SpillingPercent.load());
+                }
+
+                auto& poolMemory = it->second;
+                if (!poolMemory->AcquireIfAvailable(resources.Memory)) {
+                    hasScanQueryMemory = false;
+                    TotalMemoryResource->Release(resources.Memory);
+                }
+
+                task->PoolMemoryCookie = poolMemory->ResourceIsOverCookie;
             }
-        } // with_lock (Lock)
+        }
 
         if (!hasScanQueryMemory) {
             Counters->RmNotEnoughMemory->Inc();
             TStringBuilder reason;
-            reason << "TxId: " << txId << ", taskId: " << taskId << ". Not enough memory for query, requested: " << resources.Memory;
+            reason << "TxId: " << txId << ", taskId: " << taskId << ". Not enough memory for query, requested: " << resources.Memory
+                << ". " << tx->ToString();
             result.SetError(NKikimrKqp::TEvStartKqpTasksResponse::NOT_ENOUGH_MEMORY, reason);
             return result;
         }
@@ -232,28 +275,23 @@ class TKqpResourceManager : public IKqpResourceManager {
             if (!result) {
                 Counters->RmNotEnoughMemory->Inc();
                 with_lock (Lock) {
-                    ScanQueryMemoryResource.Release(resources.Memory);
-                } // with_lock (Lock)
+                    TotalMemoryResource->Release(resources.Memory);
+                    auto it = MemoryNamedPools.find(tx->PoolId);
+                    if (it != MemoryNamedPools.end()) {
+                        it->second->Release(resources.Memory);
+                    }
+                }
             }
         };
 
-        ui64 txTotalRequestedMemory = tx->GetExtraMemoryAllocatedSize() + resources.Memory;
-        if (txTotalRequestedMemory > queryMemoryLimit) {
-            TStringBuilder reason;
-            reason << "TxId: " << txId << ", taskId: " << taskId << ". Query memory limit exceeded: "
-                << "requested " << txTotalRequestedMemory;
-            result.SetError(NKikimrKqp::TEvStartKqpTasksResponse::QUERY_MEMORY_LIMIT_EXCEEDED, reason);
-            return result;
-        }
-
         bool allocated = ResourceBroker->SubmitTaskInstant(
             TEvResourceBroker::TEvSubmitTask(rbTaskId, rbTaskName, {0, resources.Memory}, "kqp_query", 0, {}),
             SelfId);
 
         if (!allocated) {
             TStringBuilder reason;
-            reason << "TxId: " << txId << ", taskId: " << taskId << ". Not enough ScanQueryMemory: "
-                << "requested " << resources.Memory;
+            reason << "TxId: " << txId << ", taskId: " << taskId << ". Not enough memory for query, requested: " << resources.Memory
+                << ". " << tx->ToString();
             LOG_AS_N(reason);
             result.SetError(NKikimrKqp::TEvStartKqpTasksResponse::NOT_ENOUGH_MEMORY, reason);
             return result;
@@ -302,14 +340,19 @@ class TKqpResourceManager : public IKqpResourceManager {
 
         if (resources.Memory > 0) {
             with_lock (Lock) {
-                ScanQueryMemoryResource.Release(resources.Memory);
-            } // with_lock (Lock)
+                TotalMemoryResource->Release(resources.Memory);
+                auto it = MemoryNamedPools.find(tx->PoolId);
+                if (it != MemoryNamedPools.end()) {
+                    it->second->Release(resources.Memory);
+                }
+            }
         }
 
-        LOG_AS_D("TxId: " << tx->TxId << ", taskId: " << task->TaskId << ". Released resources, "
-            << "ScanQueryMemory: " << resources.Memory << ", "
-            << "ExternalDataQueryMemory " << resources.ExternalMemory << ", "
-            << "ExecutionUnits " << resources.ExecutionUnits << ".");
+        LOG_AS_D("TxId: " << tx->TxId << ", taskId: " << task->TaskId
+            << ". Released resources, "
+            << "Memory: " << resources.Memory << ", "
+            << "Free Tier: " << resources.ExternalMemory << ", "
+            << "ExecutionUnits: " << resources.ExecutionUnits << ".");
 
         FireResourcesPublishing();
     }
@@ -346,7 +389,7 @@ class TKqpResourceManager : public IKqpResourceManager {
 
         with_lock (Lock) {
             result.ExecutionUnits = ExecutionUnitsResource.load();
-            result.Memory[EKqpMemoryPool::ScanQuery] = ScanQueryMemoryResource.Available();
+            result.Memory[EKqpMemoryPool::ScanQuery] = TotalMemoryResource->Available();
         }
 
         return result;
@@ -435,7 +478,8 @@ class TKqpResourceManager : public IKqpResourceManager {
     // limits (guarded by Lock)
     std::atomic<i32> ExecutionUnitsResource;
     std::atomic<i32> ExecutionUnitsLimit;
-    TLimitedResource<ui64> ScanQueryMemoryResource;
+    std::atomic<ui64> SpillingPercent;
+    TIntrusivePtr<TLimitedResource<ui64>> TotalMemoryResource;
     std::atomic<i64> ExternalDataQueryMemory = 0;
 
     // current state
@@ -448,6 +492,8 @@ class TKqpResourceManager : public IKqpResourceManager {
     // state for resource info exchanger
     std::shared_ptr<TResourceSnapshotState> ResourceSnapshotState;
     TActorId ResourceInfoExchanger = TActorId();
+
+    absl::flat_hash_map<TString, TIntrusivePtr<TLimitedResource<ui64>>> MemoryNamedPools;
 };
 
 struct TResourceManagers {
@@ -594,7 +640,7 @@ class TKqpResourceManagerActor : public TActorBootstrapped<TKqpResourceManagerAc
 
         if (queueConfig.GetLimit().GetMemory() > 0) {
             with_lock (ResourceManager->Lock) {
-                ResourceManager->ScanQueryMemoryResource.SetNewLimit(queueConfig.GetLimit().GetMemory());
+                ResourceManager->TotalMemoryResource->SetNewLimit(queueConfig.GetLimit().GetMemory());
             }
             LOG_I("Total node memory for scan queries: " << queueConfig.GetLimit().GetMemory() << " bytes");
         }
@@ -711,7 +757,7 @@ class TKqpResourceManagerActor : public TActorBootstrapped<TKqpResourceManagerAc
             PRE() {
                 str << "State storage key: " << WbState.Tenant << Endl;
                 with_lock (ResourceManager->Lock) {
-                    str << "ScanQuery memory resource: " << ResourceManager->ScanQueryMemoryResource.ToString() << Endl;
+                    str << "ScanQuery memory resource: " << ResourceManager->TotalMemoryResource->ToString() << Endl;
                     str << "External DataQuery memory: " << ResourceManager->ExternalDataQueryMemory.load() << Endl;
                     str << "ExecutionUnits resource: " << ResourceManager->ExecutionUnitsResource.load() << Endl;
                 }
@@ -806,13 +852,13 @@ class TKqpResourceManagerActor : public TActorBootstrapped<TKqpResourceManagerAc
         ActorIdToProto(MakeKqpResourceManagerServiceID(SelfId().NodeId()), payload.MutableResourceManagerActorId()); // legacy
         with_lock (ResourceManager->Lock) {
             payload.SetAvailableComputeActors(ResourceManager->ExecutionUnitsResource.load()); // legacy
-            payload.SetTotalMemory(ResourceManager->ScanQueryMemoryResource.GetLimit()); // legacy
-            payload.SetUsedMemory(ResourceManager->ScanQueryMemoryResource.GetLimit() - ResourceManager->ScanQueryMemoryResource.Available()); // legacy
+            payload.SetTotalMemory(ResourceManager->TotalMemoryResource->GetLimit()); // legacy
+            payload.SetUsedMemory(ResourceManager->TotalMemoryResource->GetLimit() - ResourceManager->TotalMemoryResource->Available()); // legacy
 
             payload.SetExecutionUnits(ResourceManager->ExecutionUnitsResource.load());
             auto* pool = payload.MutableMemory()->Add();
             pool->SetPool(EKqpMemoryPool::ScanQuery);
-            pool->SetAvailable(ResourceManager->ScanQueryMemoryResource.Available());
+            pool->SetAvailable(ResourceManager->TotalMemoryResource->Available());
         }
 
         LOG_I("Send to publish resource usage for "