Merge d47fd48 into dddbb74

Author: ivanmorozov333

ydb/core/protos/config.proto

@@ -634,6 +634,7 @@ message TLimiterConfig {
 message TGroupedMemoryLimiterConfig {
     optional bool Enabled = 1 [default = true];
     optional uint64 MemoryLimit = 2;
+    optional uint64 HardMemoryLimit = 3;
 }
 
 message TExternalIndexConfig {

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

@@ -52,6 +52,7 @@ class TReadContext {
     const TActorId ResourceSubscribeActorId;
     const TActorId ReadCoordinatorActorId;
     const TComputeShardingPolicy ComputeShardingPolicy;
+    TAtomic AbortFlag = 0;
 
 public:
     template <class T>
@@ -61,6 +62,13 @@ class TReadContext {
         return result;
     }
 
+    void AbortWithError(const TString& errorMessage) {
+        if (AtomicCas(&AbortFlag, 1, 0)) {
+            NActors::TActivationContext::Send(
+                ScanActorId, std::make_unique<NColumnShard::TEvPrivate::TEvTaskProcessedResult>(TConclusionStatus::Fail(errorMessage)));
+        }
+    }
+
     bool IsReverse() const {
         return ReadMetadata->IsDescSorted();
     }

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

@@ -189,6 +189,14 @@ TAllocateMemoryStep::TFetchingStepAllocation::TFetchingStepAllocation(
     , TasksGuard(source->GetContext()->GetCommonContext()->GetCounters().GetResourcesAllocationTasksGuard()) {
 }
 
+void TAllocateMemoryStep::TFetchingStepAllocation::DoOnAllocationImpossible(const TString& errorMessage) {
+    auto sourcePtr = Source.lock();
+    if (sourcePtr) {
+        sourcePtr->GetContext()->GetCommonContext()->AbortWithError(
+            "cannot allocate memory for step " + Step.GetName() + ": '" + errorMessage + "'");
+    }
+}
+
 TConclusion<bool> TAllocateMemoryStep::DoExecuteInplace(
     const std::shared_ptr<IDataSource>& source, const TFetchingScriptCursor& step) const {
 

ydb/core/tx/columnshard/engines/reader/plain_reader/iterator/fetching.h

@@ -237,7 +237,7 @@ class TAllocateMemoryStep: public IFetchingStep {
         NColumnShard::TCounterGuard TasksGuard;
         virtual bool DoOnAllocated(std::shared_ptr<NGroupedMemoryManager::TAllocationGuard>&& guard,
             const std::shared_ptr<NGroupedMemoryManager::IAllocation>& allocation) override;
-
+        virtual void DoOnAllocationImpossible(const TString& errorMessage) override;
     public:
         TFetchingStepAllocation(const std::shared_ptr<IDataSource>& source, const ui64 mem, const TFetchingScriptCursor& step);
     };

ydb/core/tx/columnshard/engines/reader/plain_reader/iterator/merge.h

@@ -77,6 +77,9 @@ class TBaseMergeTask: public IDataTasksProcessor::ITask, public NGroupedMemoryMa
     virtual bool DoApply(IDataReader& indexedDataRead) const override;
     virtual bool DoOnAllocated(std::shared_ptr<NGroupedMemoryManager::TAllocationGuard>&& guard,
         const std::shared_ptr<NGroupedMemoryManager::IAllocation>& allocation) override;
+    virtual void DoOnAllocationImpossible(const TString& errorMessage) override {
+        Context->GetCommonContext()->AbortWithError("cannot allocate memory for merge task: '" + errorMessage + "'");
+    }
 
 public:
     TBaseMergeTask(const std::shared_ptr<TMergingContext>& mergingContext, const std::shared_ptr<TSpecialReadContext>& readContext)

ydb/core/tx/columnshard/engines/reader/sys_view/chunks/chunks.cpp

@@ -195,7 +195,12 @@ TStatsIterator::TFetchingAccessorAllocation::TFetchingAccessorAllocation(
     , AccessorsManager(context->GetDataAccessorsManager())
     , Request(request)
     , WaitingCountersGuard(context->GetCounters().GetFetcherAcessorsGuard())
-    , OwnerId(context->GetScanActorId()) {
+    , OwnerId(context->GetScanActorId())
+    , Context(context) {
+}
+
+void TStatsIterator::TFetchingAccessorAllocation::DoOnAllocationImpossible(const TString& errorMessage) {
+    Context->AbortWithError("cannot allocate memory for take accessors info: " + errorMessage);
 }
 
 }   // namespace NKikimr::NOlap::NReader::NSysView::NChunks

ydb/core/tx/columnshard/engines/reader/sys_view/chunks/chunks.h

@@ -109,12 +109,15 @@ class TStatsIterator: public NAbstract::TStatsIterator<NKikimr::NSysView::Schema
         std::shared_ptr<TDataAccessorsRequest> Request;
         NColumnShard::TCounterGuard WaitingCountersGuard;
         const NActors::TActorId OwnerId;
+        const std::shared_ptr<NReader::TReadContext> Context;
+
         virtual bool DoOnAllocated(std::shared_ptr<NGroupedMemoryManager::TAllocationGuard>&& guard,
             const std::shared_ptr<NGroupedMemoryManager::IAllocation>& /*selfPtr*/) override {
             Guard = std::move(guard);
             AccessorsManager->AskData(std::move(Request));
             return true;
         }
+        virtual void DoOnAllocationImpossible(const TString& errorMessage) override;
 
         virtual void DoOnRequestsFinished(TDataAccessorsResult&& result) override {
             if (result.HasErrors()) {

ydb/core/tx/limiter/grouped_memory/service/allocation.h

@@ -48,12 +48,16 @@ class TAllocationInfo: public NColumnShard::TMonitoringObjectsCounter<TAllocatio
         AFL_TRACE(NKikimrServices::GROUPED_MEMORY_LIMITER)("event", "allocated")("allocation_id", Identifier)("stage", Stage->GetName());
         AFL_VERIFY(Allocation)("status", GetAllocationStatus())("volume", AllocatedVolume)("id", Identifier)("stage", Stage->GetName())(
             "allocation_internal_group_id", AllocationInternalGroupId);
+        auto allocationResult = Stage->Allocate(AllocatedVolume);
+        if (allocationResult.IsFail()) {
+            AllocationFailed = true;
+            Allocation->OnAllocationImpossible(allocationResult.GetErrorMessage());
+            return false;
+        }
         const bool result = Allocation->OnAllocated(
             std::make_shared<TAllocationGuard>(ProcessId, ScopeId, Allocation->GetIdentifier(), ownerId, Allocation->GetMemory()), Allocation);
-        if (result) {
-            Stage->Allocate(AllocatedVolume);
-        } else {
-            Stage->Free(AllocatedVolume, false);
+        if (!result) {
+            Stage->Free(AllocatedVolume, true);
             AllocationFailed = true;
         }
         Allocation = nullptr;

ydb/core/tx/limiter/grouped_memory/usage/abstract.h

@@ -5,6 +5,7 @@
 #include <ydb/library/actors/core/actor.h>
 #include <ydb/library/actors/core/actorid.h>
 #include <ydb/library/actors/core/log.h>
+#include <ydb/library/conclusion/status.h>
 
 namespace NKikimr::NOlap::NGroupedMemoryManager {
 
@@ -95,6 +96,7 @@ class TStageFeatures {
 private:
     YDB_READONLY_DEF(TString, Name);
     YDB_READONLY(ui64, Limit, 0);
+    YDB_READONLY(ui64, HardLimit, 0);
     YDB_ACCESSOR_DEF(TPositiveControlInteger, Usage);
     YDB_ACCESSOR_DEF(TPositiveControlInteger, Waiting);
     std::shared_ptr<TStageFeatures> Owner;
@@ -114,24 +116,33 @@ class TStageFeatures {
         return Usage.Val() + Waiting.Val();
     }
 
-    TStageFeatures(
-        const TString& name, const ui64 limit, const std::shared_ptr<TStageFeatures>& owner, const std::shared_ptr<TStageCounters>& counters)
+    TStageFeatures(const TString& name, const ui64 limit, const ui64 hardLimit, const std::shared_ptr<TStageFeatures>& owner,
+        const std::shared_ptr<TStageCounters>& counters)
         : Name(name)
         , Limit(limit)
+        , HardLimit(hardLimit)
         , Owner(owner)
         , Counters(counters) {
     }
 
-    void Allocate(const ui64 volume) {
+    [[nodiscard]] TConclusionStatus Allocate(const ui64 volume) {
+        if (HardLimit < Usage.Val() + volume) {
+            return TConclusionStatus::Fail(TStringBuilder() << "limit:" << HardLimit << ";val:" << Usage.Val() << ";delta=" << volume << ";");
+        }
         Waiting.Sub(volume);
         Usage.Add(volume);
         if (Counters) {
             Counters->Add(volume, true);
             Counters->Sub(volume, false);
         }
         if (Owner) {
-            Owner->Allocate(volume);
+            const auto ownerResult = Owner->Allocate(volume);
+            if (ownerResult.IsFail()) {
+                Free(volume, true);
+                return ownerResult;
+            }
         }
+        return TConclusionStatus::Success();
     }
 
     void Free(const ui64 volume, const bool allocated) {
@@ -199,6 +210,7 @@ class IAllocation {
     YDB_READONLY(ui64, Identifier, Counter.Inc());
     YDB_READONLY(ui64, Memory, 0);
     bool Allocated = false;
+    virtual void DoOnAllocationImpossible(const TString& errorMessage) = 0;
     virtual bool DoOnAllocated(
         std::shared_ptr<TAllocationGuard>&& guard, const std::shared_ptr<NGroupedMemoryManager::IAllocation>& allocation) = 0;
 
@@ -216,6 +228,10 @@ class IAllocation {
         return Allocated;
     }
 
+    void OnAllocationImpossible(const TString& errorMessage) {
+        DoOnAllocationImpossible(errorMessage);
+    }
+
     [[nodiscard]] bool OnAllocated(
         std::shared_ptr<TAllocationGuard>&& guard, const std::shared_ptr<NGroupedMemoryManager::IAllocation>& allocation);
 };

ydb/core/tx/limiter/grouped_memory/usage/config.cpp

@@ -7,13 +7,16 @@ bool TConfig::DeserializeFromProto(const NKikimrConfig::TGroupedMemoryLimiterCon
     if (config.HasMemoryLimit()) {
         MemoryLimit = config.GetMemoryLimit();
     }
+    if (config.HasHardMemoryLimit()) {
+        HardMemoryLimit = config.GetHardMemoryLimit();
+    }
     Enabled = config.GetEnabled();
     return true;
 }
 
 TString TConfig::DebugString() const {
     TStringBuilder sb;
-    sb << "MemoryLimit=" << MemoryLimit << ";Enabled=" << Enabled << ";";
+    sb << "MemoryLimit=" << MemoryLimit << ";HardMemoryLimit=" << HardMemoryLimit << ";Enabled=" << Enabled << ";";
     return sb;
 }