Merge a4a82ed into 7cab7c7

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,21 @@ 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) {
+                alloc->Ref().SetIncreaseMemoryLimitCallback([this, alloc](ui64 limit, ui64 required) {
+                    RequestExtraMemory(required - limit, alloc);
+                    const ui64 RSSLimit = std::numeric_limits<ui64>::max();
+                    auto currentRSS = NMemInfo::GetMemInfo().RSS;
+                    if (currentRSS > RSSLimit) {
+                        alloc->SetMaximumLimitValueReached(true);
+                    }
+                });
+            }
+        }
+
         void TryShrinkMemory(NKikimr::NMiniKQL::TScopedAlloc* alloc) {
             if (alloc->GetAllocated() - alloc->GetUsed() > MemoryLimits.MinMemFreeSize) {
                 alloc->ReleaseFreePages();
@@ -113,6 +129,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_;