YQL-18419 create callback that checks RSS and limits memory usage

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();

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_;