Add memory limit to prevent converter to AST get OOM (#8351)

Author: MrLolthe1st

ydb/library/yql/ast/yql_expr.cpp

@@ -2709,7 +2709,7 @@ TAstParseResult ConvertToAst(const TExprNode& root, TExprContext& exprContext, c
     ctx.RefAtoms = settings.RefAtoms;
     ctx.AllowFreeArgs = settings.AllowFreeArgs;
     ctx.NormalizeAtomFlags = settings.NormalizeAtomFlags;
-    ctx.Pool = std::make_unique<TMemoryPool>(4096);
+    ctx.Pool = std::make_unique<TMemoryPool>(4096, TMemoryPool::TExpGrow::Instance(), settings.Allocator);
     ctx.Frames.push_back(TFrameContext());
     ctx.CurrentFrame = &ctx.Frames.front();
     VisitNode(root, 0ULL, ctx);

ydb/library/yql/ast/yql_expr.h

@@ -2846,6 +2846,7 @@ struct TConvertToAstSettings {
     bool PrintArguments = false;
     bool AllowFreeArgs = false;
     bool NormalizeAtomFlags = false;
+    IAllocator* Allocator = TDefaultAllocator::Instance();
 };
 
 TAstParseResult ConvertToAst(const TExprNode& root, TExprContext& ctx, const TConvertToAstSettings& settings);

ydb/library/yql/core/facade/yql_facade.cpp

@@ -12,6 +12,7 @@
 #include <ydb/library/yql/core/services/yql_eval_params.h>
 #include <ydb/library/yql/utils/log/context.h>
 #include <ydb/library/yql/utils/log/profile.h>
+#include <ydb/library/yql/utils/limiting_allocator.h>
 #include <ydb/library/yql/core/services/yql_out_transformers.h>
 #include <ydb/library/yql/core/extract_predicate/extract_predicate_dbg.h>
 #include <ydb/library/yql/providers/common/provider/yql_provider_names.h>
@@ -1388,7 +1389,7 @@ TFuture<IGraphTransformer::TStatus> TProgram::AsyncTransformWithFallback(bool ap
     });
 }
 
-TMaybe<TString> TProgram::GetQueryAst() {
+TMaybe<TString> TProgram::GetQueryAst(TMaybe<size_t> memoryLimit) {
     if (ExternalQueryAst_) {
         return ExternalQueryAst_;
     }
@@ -1397,7 +1398,17 @@ TMaybe<TString> TProgram::GetQueryAst() {
     astStream.Reserve(DEFAULT_AST_BUF_SIZE);
 
     if (ExprRoot_) {
-        auto ast = ConvertToAst(*ExprRoot_, *ExprCtx_, TExprAnnotationFlags::None, true);
+        std::unique_ptr<IAllocator> limitingAllocator;
+        TConvertToAstSettings settings;
+        settings.AnnotationFlags = TExprAnnotationFlags::None;
+        settings.RefAtoms = true;
+        settings.Allocator = TDefaultAllocator::Instance();
+        if (memoryLimit) {
+            limitingAllocator = MakeLimitingAllocator(*memoryLimit, TDefaultAllocator::Instance());
+            settings.Allocator = limitingAllocator.get();
+        }
+
+        auto ast = ConvertToAst(*ExprRoot_, *ExprCtx_, settings);
         ast.Root->PrettyPrintTo(astStream, TAstPrintFlags::ShortQuote | TAstPrintFlags::PerLine);
         return astStream.Str();
     } else if (AstRoot_) {

ydb/library/yql/core/facade/yql_facade.h

@@ -223,7 +223,7 @@ class TProgram: public TThrRefBase, private TNonCopyable
         return ResultProviderConfig_->CommittedResults;
     }
 
-    TMaybe<TString> GetQueryAst();
+    TMaybe<TString> GetQueryAst(TMaybe<size_t> memoryLimit = {});
     TMaybe<TString> GetQueryPlan(const TPlanSettings& settings = {});
 
     void SetDiagnosticFormat(NYson::EYsonFormat format) {

ydb/library/yql/utils/limiting_allocator.cpp

@@ -0,0 +1,35 @@
+#include "limiting_allocator.h"
+
+#include <util/memory/pool.h>
+#include <util/generic/yexception.h>
+
+namespace {
+class TLimitingAllocator : public IAllocator {
+public:
+    TLimitingAllocator(size_t limit, IAllocator* allocator) : Alloc_(allocator), Limit_(limit) {};
+    TBlock Allocate(size_t len) override final {
+        if (Allocated_ + len > Limit_) {
+            throw std::runtime_error("Out of memory");
+        }
+        Allocated_ += len;
+        return Alloc_->Allocate(len);
+    }
+
+    void Release(const TBlock& block) override final {
+        Y_ENSURE(Allocated_ >= block.Len);
+        Allocated_ -= block.Len;
+        Alloc_->Release(block);
+    }
+
+private:
+    IAllocator* Alloc_;
+    size_t Allocated_ = 0;
+    const size_t Limit_;
+};
+}
+
+namespace NYql {
+std::unique_ptr<IAllocator> MakeLimitingAllocator(size_t limit, IAllocator* underlying) {
+    return std::make_unique<TLimitingAllocator>(limit, underlying);
+}
+}

ydb/library/yql/utils/limiting_allocator.h

@@ -0,0 +1,8 @@
+#pragma once
+
+#include <util/memory/pool.h>
+#include <memory>
+
+namespace NYql {
+std::unique_ptr<IAllocator> MakeLimitingAllocator(size_t limit, IAllocator* underlying);
+}

ydb/library/yql/utils/ya.make

@@ -10,6 +10,7 @@ SRCS(
     future_action.h
     hash.cpp
     hash.h
+    limiting_allocator.cpp
     md5_stream.cpp
     md5_stream.h
     method_index.cpp