Merge f92aa20 into 2ed4009

Author: lll-phill-lll

ydb/library/yql/dq/actors/compute/dq_compute_memory_quota.h

@@ -1,6 +1,7 @@
 
 #pragma once
 
+#include <util/system/mem_info.h>
 #include <ydb/library/services/services.pb.h>
 #include <ydb/library/yql/dq/actors/compute/dq_compute_actor.h>
 #include <ydb/library/yql/minikql/mkql_alloc.h>
@@ -55,6 +56,23 @@ namespace NYql::NDq {
             }
         }
 
+        // This callback is created for testing purposes and will be enabled only with spilling.
+        // Most likely this callback will be removed after KIKIMR-21481.
+        void TrySetIncreaseMemoryLimitCallbackWithRSSControl(NKikimr::NMiniKQL::TScopedAlloc* alloc) {
+            if (!CanAllocateExtraMemory) return;
+            const ui64 limitRSS = std::numeric_limits<ui64>::max();
+            const ui64 criticalRSSValue = limitRSS / 100 * 80;
+
+            alloc->Ref().SetIncreaseMemoryLimitCallback([this, alloc](ui64 limit, ui64 required) {
+                RequestExtraMemory(required - limit, alloc);
+                
+                ui64 currentRSS = NMemInfo::GetMemInfo().RSS;
+                if (currentRSS > criticalRSSValue) {
+                    alloc->SetMaximumLimitValueReached(true);
+                }
+            });
+        }
+
         void TryShrinkMemory(NKikimr::NMiniKQL::TScopedAlloc* alloc) {
             if (alloc->GetAllocated() - alloc->GetUsed() > MemoryLimits.MinMemFreeSize) {
                 alloc->ReleaseFreePages();
@@ -113,6 +131,7 @@ namespace NYql::NDq {
 
     private:
         void RequestExtraMemory(ui64 memory, NKikimr::NMiniKQL::TScopedAlloc* alloc) {
+            alloc->GetAllocated();
             memory = std::max(AlignMemorySizeToMbBoundary(memory), MemoryLimits.MinMemAllocSize);
 
             if (MemoryLimits.MkqlProgramHardMemoryLimit && MkqlMemoryLimit + memory > MemoryLimits.MkqlProgramHardMemoryLimit) {

ydb/library/yql/dq/actors/task_runner/task_runner_actor_local.cpp

@@ -423,7 +423,11 @@ class TLocalTaskRunnerActor
 
         auto guard = TaskRunner->BindAllocator(MemoryQuota ? TMaybe<ui64>(MemoryQuota->GetMkqlMemoryLimit()) : Nothing());
         if (MemoryQuota) {
-            MemoryQuota->TrySetIncreaseMemoryLimitCallback(guard.GetMutex());
+            if (settings.GetEnableSpilling()) {
+                MemoryQuota->TrySetIncreaseMemoryLimitCallbackWithRSSControl(guard.GetMutex());
+            } else {
+                MemoryQuota->TrySetIncreaseMemoryLimitCallback(guard.GetMutex());
+            }   
         }
 
         TaskRunner->Prepare(settings, ev->Get()->MemoryLimits, *ev->Get()->ExecCtx);

ydb/library/yql/minikql/mkql_alloc.h

@@ -213,6 +213,10 @@ class TScopedAlloc {
         Release();
     }
 
+    void SetMaximumLimitValueReached(bool IsReached) {
+        MyState_.SetMaximumLimitValueReached(IsReached);
+    }
+
 private:
     const bool InitiallyAcquired_;
     TAllocState MyState_;