YQL-18314 use shared alloc in all task runners

zverevgeny · zverevgeny · commit d80571ecc879 · 2024-05-02T12:58:49.000Z
diff --git a/ydb/core/kqp/compute_actor/kqp_pure_compute_actor.cpp b/ydb/core/kqp/compute_actor/kqp_pure_compute_actor.cpp
@@ -69,7 +69,7 @@ void TKqpComputeActor::DoBootstrap() {
         settings.ReadRanges.push_back(readRange);
     }
 
-    auto taskRunner = MakeDqTaskRunner(TBase::GetAllocator(), execCtx, settings, logger);
+    auto taskRunner = MakeDqTaskRunner(TBase::GetAllocatorPtr(), execCtx, settings, logger);
     SetTaskRunner(taskRunner);
 
     auto wakeup = [this]{ ContinueExecute(); };
diff --git a/ydb/core/kqp/compute_actor/kqp_scan_compute_actor.cpp b/ydb/core/kqp/compute_actor/kqp_scan_compute_actor.cpp
@@ -216,7 +216,7 @@ void TKqpScanComputeActor::DoBootstrap() {
         };
     }
 
-    auto taskRunner = MakeDqTaskRunner(GetAllocator(), execCtx, settings, logger);
+    auto taskRunner = MakeDqTaskRunner(GetAllocatorPtr(), execCtx, settings, logger);
     TBase::SetTaskRunner(taskRunner);
 
     auto wakeup = [this] { ContinueExecute(); };
diff --git a/ydb/core/kqp/executer_actor/kqp_executer_impl.cpp b/ydb/core/kqp/executer_actor/kqp_executer_impl.cpp
@@ -49,7 +49,7 @@ void TEvKqpExecuter::TEvTxResponse::TakeResult(ui32 idx, NDq::TDqSerializedBatch
 
 TEvKqpExecuter::TEvTxResponse::~TEvTxResponse() {
     if (!TxResults.empty() && Y_LIKELY(AllocState)) {
-        with_lock(AllocState->Alloc) {
+        with_lock(*AllocState->Alloc) {
             TxResults.crop(0);
         }
     }
diff --git a/ydb/core/kqp/executer_actor/kqp_executer_impl.h b/ydb/core/kqp/executer_actor/kqp_executer_impl.h
@@ -166,7 +166,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
     TString BuildMemoryLimitExceptionMessage() const {
         if (Request.TxAlloc) {
             return TStringBuilder() << "Memory limit exception at " << CurrentStateFuncName()
-                << ", current limit is " << Request.TxAlloc->Alloc.GetLimit() << " bytes.";
+                << ", current limit is " << Request.TxAlloc->Alloc->GetLimit() << " bytes.";
         }
         return TStringBuilder() << "Memory limit exception at " << CurrentStateFuncName();
     }
diff --git a/ydb/core/kqp/executer_actor/kqp_literal_executer.cpp b/ydb/core/kqp/executer_actor/kqp_literal_executer.cpp
@@ -166,7 +166,7 @@ class TKqpLiteralExecuter {
         UpdateCounters();
     }
 
-    void RunTask(NMiniKQL::TScopedAlloc& alloc, TTask& task, const TDqTaskRunnerContext& context, const TDqTaskRunnerSettings& settings) {
+    void RunTask(std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, TTask& task, const TDqTaskRunnerContext& context, const TDqTaskRunnerSettings& settings) {
         auto& stageInfo = TasksGraph.GetStageInfo(task.StageId);
         auto& stage = stageInfo.Meta.GetStage(stageInfo.Id);
 
@@ -211,7 +211,7 @@ class TKqpLiteralExecuter {
         auto status = taskRunner->Run();
         YQL_ENSURE(status == ERunStatus::Finished);
 
-        with_lock (alloc) { // allocator is used only by outputChannel->PopAll()
+        with_lock (*alloc) { // allocator is used only by outputChannel->PopAll()
             for (auto& taskOutput : task.Outputs) {
                 for (ui64 outputChannelId : taskOutput.Channels) {
                     auto outputChannel = taskRunner->GetOutputChannel(outputChannelId);
@@ -276,7 +276,7 @@ class TKqpLiteralExecuter {
 
 private:
     void CleanupCtx() {
-        with_lock(Request.TxAlloc->Alloc) {
+        with_lock(*Request.TxAlloc->Alloc) {
             TaskRunners.erase(TaskRunners.begin(), TaskRunners.end());
             Request.Transactions.erase(Request.Transactions.begin(), Request.Transactions.end());
             ComputeCtx.reset();
diff --git a/ydb/core/kqp/provider/yql_kikimr_exec.cpp b/ydb/core/kqp/provider/yql_kikimr_exec.cpp
@@ -577,7 +577,7 @@ class TKiSourceCallableExecutionTransformer : public TAsyncCallbackTransformer<T
             return peepHoleStatus;
         }
 
-        auto guard = Guard(SessionCtx->Query().QueryData->GetAllocState()->Alloc);
+        auto guard = Guard(*SessionCtx->Query().QueryData->GetAllocState()->Alloc);
 
         auto input = Build<TDqPhyStage>(ctx, pos)
             .Inputs()
diff --git a/ydb/core/kqp/query_data/kqp_query_data.cpp b/ydb/core/kqp/query_data/kqp_query_data.cpp
@@ -117,19 +117,17 @@ void TKqpExecuterTxResult::FillYdb(Ydb::ResultSet* ydbResult, TMaybe<ui64> rowsL
 
 TTxAllocatorState::TTxAllocatorState(const IFunctionRegistry* functionRegistry,
     TIntrusivePtr<ITimeProvider> timeProvider, TIntrusivePtr<IRandomProvider> randomProvider)
-    : Alloc(__LOCATION__, NKikimr::TAlignedPagePoolCounters(), functionRegistry->SupportsSizedAllocators())
-    , TypeEnv(Alloc)
+    : Alloc(std::make_shared<NKikimr::NMiniKQL::TScopedAlloc>(__LOCATION__, NKikimr::TAlignedPagePoolCounters(), functionRegistry->SupportsSizedAllocators(), false))
+    , TypeEnv(*Alloc)
     , MemInfo("TQueryData")
-    , HolderFactory(Alloc.Ref(), MemInfo, functionRegistry)
+    , HolderFactory(Alloc->Ref(), MemInfo, functionRegistry)
 {
-    Alloc.Release();
     TimeProvider = timeProvider;
     RandomProvider = randomProvider;
 }
 
 TTxAllocatorState::~TTxAllocatorState()
 {
-    Alloc.Acquire();
 }
 
 std::pair<NKikimr::NMiniKQL::TType*, NUdf::TUnboxedValue> TTxAllocatorState::GetInternalBindingValue(
diff --git a/ydb/core/kqp/query_data/kqp_query_data.h b/ydb/core/kqp/query_data/kqp_query_data.h
@@ -166,7 +166,7 @@ struct TTimeAndRandomProvider {
 
 class TTxAllocatorState: public TTimeAndRandomProvider {
 public:
-    NKikimr::NMiniKQL::TScopedAlloc Alloc;
+    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> Alloc;
     NKikimr::NMiniKQL::TTypeEnvironment TypeEnv;
     NKikimr::NMiniKQL::TMemoryUsageInfo MemInfo;
     NKikimr::NMiniKQL::THolderFactory HolderFactory;
diff --git a/ydb/core/kqp/runtime/kqp_tasks_runner.cpp b/ydb/core/kqp/runtime/kqp_tasks_runner.cpp
@@ -70,7 +70,7 @@ IDqOutputConsumer::TPtr KqpBuildOutputConsumer(const NDqProto::TTaskOutput& outp
 
 
 TKqpTasksRunner::TKqpTasksRunner(google::protobuf::RepeatedPtrField<NDqProto::TDqTask>&& tasks,
-    NKikimr::NMiniKQL::TScopedAlloc& alloc,
+    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc,
     const TDqTaskRunnerContext& execCtx, const TDqTaskRunnerSettings& settings, const TLogFunc& logFunc)
     : LogFunc(logFunc)
     , Alloc(alloc)
@@ -230,13 +230,13 @@ const NYql::NDq::TDqTaskSettings& TKqpTasksRunner::GetTask(ui64 taskId) const {
 
 TGuard<NMiniKQL::TScopedAlloc> TKqpTasksRunner::BindAllocator(TMaybe<ui64> memoryLimit) {
     if (memoryLimit) {
-        Alloc.SetLimit(*memoryLimit);
+        Alloc->SetLimit(*memoryLimit);
     }
-    return TGuard(Alloc);
+    return TGuard(*Alloc);
 }
 
 TIntrusivePtr<TKqpTasksRunner> CreateKqpTasksRunner(google::protobuf::RepeatedPtrField<NDqProto::TDqTask>&& tasks,
-    NKikimr::NMiniKQL::TScopedAlloc& alloc,
+    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc,
     const TDqTaskRunnerContext& execCtx, const TDqTaskRunnerSettings& settings, const TLogFunc& logFunc)
 {
     return new TKqpTasksRunner(std::move(tasks), alloc, execCtx, settings, logFunc);
diff --git a/ydb/core/kqp/runtime/kqp_tasks_runner.h b/ydb/core/kqp/runtime/kqp_tasks_runner.h
@@ -16,7 +16,7 @@ NYql::NDq::IDqOutputConsumer::TPtr KqpBuildOutputConsumer(const NYql::NDqProto::
 class TKqpTasksRunner : public TSimpleRefCount<TKqpTasksRunner>, private TNonCopyable {
 public:
     TKqpTasksRunner(google::protobuf::RepeatedPtrField<NYql::NDqProto::TDqTask>&& tasks,
-                    NKikimr::NMiniKQL::TScopedAlloc& alloc,
+                    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc,
                     const NYql::NDq::TDqTaskRunnerContext& execCtx, const NYql::NDq::TDqTaskRunnerSettings& settings,
                     const NYql::NDq::TLogFunc& logFunc);
 
@@ -51,15 +51,15 @@ class TKqpTasksRunner : public TSimpleRefCount<TKqpTasksRunner>, private TNonCop
     // otherwise use particular memory limit
     TGuard<NMiniKQL::TScopedAlloc> BindAllocator(TMaybe<ui64> memoryLimit = Nothing());
 
-    ui64 GetAllocatedMemory() const { return Alloc.GetAllocated(); }
+    ui64 GetAllocatedMemory() const { return Alloc->GetAllocated(); }
 
     const TMap<ui64, const NYql::NDq::TDqTaskRunnerStats*> GetTasksStats() const { return Stats; }
 private:
     TMap<ui64, TIntrusivePtr<NYql::NDq::IDqTaskRunner>> TaskRunners;
     TMap<ui64, NYql::NDq::TDqTaskSettings> Tasks;
     TMap<ui64, const NYql::NDq::TDqTaskRunnerStats*> Stats;
     NYql::NDq::TLogFunc LogFunc;
-    NMiniKQL::TScopedAlloc& Alloc;
+   std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> Alloc;
     NMiniKQL::TKqpComputeContextBase* ComputeCtx;
     NMiniKQL::TKqpDatashardApplyContext* ApplyCtx;
 
diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
@@ -692,11 +692,11 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
         ui64 mkqlMaxLimit = phaseLimitsProto.GetComputeNodeMemoryLimitBytes();
         mkqlMaxLimit = mkqlMaxLimit ? mkqlMaxLimit : ui64(Settings.MkqlMaxMemoryLimit);
 
-        alloc->Alloc.SetLimit(mkqlInitialLimit);
-        alloc->Alloc.Ref().SetIncreaseMemoryLimitCallback([this, &alloc, mkqlMaxLimit](ui64 currentLimit, ui64 required) {
+        alloc->Alloc->SetLimit(mkqlInitialLimit);
+        alloc->Alloc->Ref().SetIncreaseMemoryLimitCallback([this, &alloc, mkqlMaxLimit](ui64 currentLimit, ui64 required) {
             if (required < mkqlMaxLimit) {
                 LOG_D("Increase memory limit from " << currentLimit << " to " << required);
-                alloc->Alloc.SetLimit(required);
+                alloc->Alloc->SetLimit(required);
             }
         });
 
@@ -2322,7 +2322,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
     TString BuildMemoryLimitExceptionMessage() const {
         if (QueryState && QueryState->TxCtx) {
             return TStringBuilder() << "Memory limit exception at " << CurrentStateFuncName()
-                << ", current limit is " << QueryState->TxCtx->TxAlloc->Alloc.GetLimit() << " bytes.";
+                << ", current limit is " << QueryState->TxCtx->TxAlloc->Alloc->GetLimit() << " bytes.";
         } else {
             return TStringBuilder() << "Memory limit exception at " << CurrentStateFuncName();
         }
diff --git a/ydb/library/yql/dq/actors/compute/dq_compute_actor.cpp b/ydb/library/yql/dq/actors/compute/dq_compute_actor.cpp
@@ -55,7 +55,7 @@ class TDqComputeActor : public TDqSyncComputeActorBase<TDqComputeActor> {
             };
         }
 
-        auto taskRunner = TaskRunnerFactory(GetAllocator(), Task, RuntimeSettings.StatsMode, logger);
+        auto taskRunner = TaskRunnerFactory(GetAllocatorPtr(), Task, RuntimeSettings.StatsMode, logger);
         SetTaskRunner(taskRunner);
         auto wakeup = [this]{ ContinueExecute(EResumeSource::CABootstrapWakeup); };
         TDqTaskRunnerExecutionContext execCtx(TxId, std::move(wakeup));
diff --git a/ydb/library/yql/dq/actors/compute/dq_compute_actor.h b/ydb/library/yql/dq/actors/compute/dq_compute_actor.h
@@ -367,7 +367,7 @@ struct TComputeMemoryLimits {
 };
 
 using TTaskRunnerFactory = std::function<
-    TIntrusivePtr<IDqTaskRunner>(NKikimr::NMiniKQL::TScopedAlloc& alloc, const TDqTaskSettings& task, NDqProto::EDqStatsMode statsMode, const TLogFunc& logFunc)
+    TIntrusivePtr<IDqTaskRunner>(std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const TDqTaskSettings& task, NDqProto::EDqStatsMode statsMode, const TLogFunc& logFunc)
 >;
 
 void FillAsyncStats(NDqProto::TDqAsyncBufferStats& proto, TDqAsyncStats stats);
diff --git a/ydb/library/yql/dq/actors/task_runner/task_runner_actor_local.cpp b/ydb/library/yql/dq/actors/task_runner/task_runner_actor_local.cpp
@@ -405,7 +405,7 @@ class TLocalTaskRunnerActor
     void OnDqTask(TEvTaskRunnerCreate::TPtr& ev) {
         ParentId = ev->Sender;
         auto settings = NDq::TDqTaskSettings(&ev->Get()->Task);
-        TaskRunner = Factory(*Alloc.get(), settings, ev->Get()->StatsMode, [this](const TString& message) {
+        TaskRunner = Factory(Alloc, settings, ev->Get()->StatsMode, [this](const TString& message) {
             LOG_D(message);
         });
 
diff --git a/ydb/library/yql/dq/runtime/dq_tasks_runner.cpp b/ydb/library/yql/dq/runtime/dq_tasks_runner.cpp
@@ -1011,10 +1011,10 @@ class TDqTaskRunner : public IDqTaskRunner {
     }
 };
 
-TIntrusivePtr<IDqTaskRunner> MakeDqTaskRunner(NKikimr::NMiniKQL::TScopedAlloc& alloc, const TDqTaskRunnerContext& ctx, const TDqTaskRunnerSettings& settings,
+TIntrusivePtr<IDqTaskRunner> MakeDqTaskRunner(std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const TDqTaskRunnerContext& ctx, const TDqTaskRunnerSettings& settings,
     const TLogFunc& logFunc)
 {
-    return new TDqTaskRunner(alloc, ctx, settings, logFunc);
+    return new TDqTaskRunner(*alloc, ctx, settings, logFunc);
 }
 
 } // namespace NYql::NDq
diff --git a/ydb/library/yql/dq/runtime/dq_tasks_runner.h b/ydb/library/yql/dq/runtime/dq_tasks_runner.h
@@ -405,7 +405,7 @@ class IDqTaskRunner : public TSimpleRefCount<IDqTaskRunner>, private TNonCopyabl
 };
 
 TIntrusivePtr<IDqTaskRunner> MakeDqTaskRunner(
-    NKikimr::NMiniKQL::TScopedAlloc& alloc, 
+    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, 
     const TDqTaskRunnerContext& ctx, 
     const TDqTaskRunnerSettings& settings,
     const TLogFunc& logFunc
diff --git a/ydb/library/yql/providers/dq/actors/compute_actor.cpp b/ydb/library/yql/providers/dq/actors/compute_actor.cpp
@@ -56,7 +56,7 @@ IActor* CreateComputeActor(
         }
     }
 
-    auto taskRunnerFactory = [factory = options.Factory](NKikimr::NMiniKQL::TScopedAlloc& alloc, const NDq::TDqTaskSettings& task, NDqProto::EDqStatsMode statsMode, const NDq::TLogFunc& logger) {
+    auto taskRunnerFactory = [factory = options.Factory](std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const NDq::TDqTaskSettings& task, NDqProto::EDqStatsMode statsMode, const NDq::TLogFunc& logger) {
         Y_UNUSED(logger);
         return factory->Get(alloc, task, statsMode, {});
     };
diff --git a/ydb/library/yql/providers/dq/local_gateway/yql_dq_gateway_local.cpp b/ydb/library/yql/providers/dq/local_gateway/yql_dq_gateway_local.cpp
@@ -67,7 +67,7 @@ class TLocalServiceHolder {
         lwmOptions.FunctionRegistry = functionRegistry;
         lwmOptions.TaskRunnerInvokerFactory = new NDqs::TTaskRunnerInvokerFactory();
         lwmOptions.TaskRunnerActorFactory = NDq::NTaskRunnerActor::CreateLocalTaskRunnerActorFactory(
-            [factory=lwmOptions.Factory](NKikimr::NMiniKQL::TScopedAlloc& alloc, const NDq::TDqTaskSettings& task, NDqProto::EDqStatsMode statsMode, const NDq::TLogFunc& )
+            [factory=lwmOptions.Factory](std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const NDq::TDqTaskSettings& task, NDqProto::EDqStatsMode statsMode, const NDq::TLogFunc& )
                 {
                     return factory->Get(alloc, task, statsMode);
                 });
diff --git a/ydb/library/yql/providers/dq/provider/exec/yql_dq_exectransformer.cpp b/ydb/library/yql/providers/dq/provider/exec/yql_dq_exectransformer.cpp
@@ -124,7 +124,7 @@ class TLocalExecutor: public TCounters
             ? CreateDeterministicRandomProvider(1)
             : State->RandomProvider;
 
-        TScopedAlloc alloc(
+        auto alloc = std::make_shared<NKikimr::NMiniKQL::TScopedAlloc>(
             __LOCATION__, 
             NKikimr::TAlignedPagePoolCounters(), 
             State->FunctionRegistry->SupportsSizedAllocators(),
diff --git a/ydb/library/yql/providers/dq/runtime/task_command_executor.cpp b/ydb/library/yql/providers/dq/runtime/task_command_executor.cpp
@@ -724,14 +724,14 @@ class TTaskCommandExecutor {
 
             Y_ABORT_UNLESS(!Alloc);
             Y_ABORT_UNLESS(FunctionRegistry);
-            Alloc = std::make_unique<NKikimr::NMiniKQL::TScopedAlloc>(
+            Alloc = std::make_shared<NKikimr::NMiniKQL::TScopedAlloc>(
                 __LOCATION__,
                 NKikimr::TAlignedPagePoolCounters(),
                 FunctionRegistry->SupportsSizedAllocators(),
                 false
             );
 
-            Runner = MakeDqTaskRunner(*Alloc.get(), Ctx, settings, nullptr);
+            Runner = MakeDqTaskRunner(Alloc, Ctx, settings, nullptr);
         });
 
         auto guard = Runner->BindAllocator(DqConfiguration->MemoryLimit.Get().GetOrElse(0));
@@ -761,7 +761,7 @@ class TTaskCommandExecutor {
         result.Save(&output);
     }
 private:
-    std::unique_ptr<NKikimr::NMiniKQL::TScopedAlloc> Alloc;
+    std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> Alloc;
     NKikimr::NMiniKQL::TComputationNodeFactory ComputationFactory;
     TTaskTransformFactory TaskTransformFactory;
     NKikimr::NMiniKQL::IStatsRegistry* JobStats;
diff --git a/ydb/library/yql/providers/dq/task_runner/tasks_runner_local.cpp b/ydb/library/yql/providers/dq/task_runner/tasks_runner_local.cpp
@@ -220,11 +220,11 @@ class TLocalFactory: public IProxyFactory {
         ExecutionContext.PatternCache = patternCache;
     }
 
-    ITaskRunner::TPtr GetOld(NKikimr::NMiniKQL::TScopedAlloc& alloc, const TDqTaskSettings& task, const TString& traceId) override {
+    ITaskRunner::TPtr GetOld(std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const TDqTaskSettings& task, const TString& traceId) override {
         return new TLocalTaskRunner(task, Get(alloc, task, NDqProto::DQ_STATS_MODE_BASIC, traceId));
     }
 
-    TIntrusivePtr<NDq::IDqTaskRunner> Get(NKikimr::NMiniKQL::TScopedAlloc& alloc, const TDqTaskSettings& task, NDqProto::EDqStatsMode statsMode, const TString& traceId) override {
+    TIntrusivePtr<NDq::IDqTaskRunner> Get(std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const TDqTaskSettings& task, NDqProto::EDqStatsMode statsMode, const TString& traceId) override {
         Y_UNUSED(traceId);
         NDq::TDqTaskRunnerSettings settings;
         settings.TerminateOnError = TerminateOnError;
diff --git a/ydb/library/yql/providers/dq/task_runner/tasks_runner_pipe.cpp b/ydb/library/yql/providers/dq/task_runner/tasks_runner_pipe.cpp
@@ -1291,6 +1291,7 @@ class TDqSink : public IDqAsyncOutputBuffer {
 class TTaskRunner: public IPipeTaskRunner {
 public:
     TTaskRunner(
+        std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc,
         const NDqProto::TDqTask& task,
         TFilesHolder::TPtr&& filesHolder,
         THolder<TChildProcess>&& command,
@@ -1299,8 +1300,7 @@ class TTaskRunner: public IPipeTaskRunner {
         : TraceId(traceId)
         , Task(task)
         , FilesHolder(std::move(filesHolder))
-        , Alloc(new NKikimr::NMiniKQL::TScopedAlloc(__LOCATION__, NKikimr::TAlignedPagePoolCounters(), true),
-            [](NKikimr::NMiniKQL::TScopedAlloc* ptr) { ptr->Acquire(); delete ptr; })
+        , Alloc(alloc)
         , AllocatedHolder(std::make_optional<TAllocatedHolder>(*Alloc, "TDqTaskRunnerProxy"))
         , Running(true)
         , Command(std::move(command))
@@ -1310,16 +1310,14 @@ class TTaskRunner: public IPipeTaskRunner {
         , TaskId(Task.GetId())
         , StageId(stageId)
     {
-        Alloc->Release();
         StderrReader->Start();
         InitTaskMeta();
         InitChannels();
     }
 
     ~TTaskRunner() {
-        Alloc->Acquire();
+        auto guard = Guard(*Alloc);
         AllocatedHolder.reset();
-        Alloc->Release();
         Command->Kill();
         Command->Wait(TDuration::Seconds(0));
     }
@@ -1629,12 +1627,13 @@ class TTaskRunner: public IPipeTaskRunner {
 class TDqTaskRunner: public NDq::IDqTaskRunner {
 public:
     TDqTaskRunner(
+        std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc,
         const NDqProto::TDqTask& task,
         TFilesHolder::TPtr&& filesHolder,
         THolder<TChildProcess>&& command,
         ui64 stageId,
         const TString& traceId)
-        : Delegate(new TTaskRunner(task, std::move(filesHolder), std::move(command), stageId, traceId))
+        : Delegate(new TTaskRunner(alloc, task, std::move(filesHolder), std::move(command), stageId, traceId))
         , Task(task)
     { }
 
@@ -1914,17 +1913,17 @@ class TPipeFactory: public IProxyFactory {
         TaskScheduler.Start();
     }
 
-    ITaskRunner::TPtr GetOld(NKikimr::NMiniKQL::TScopedAlloc& alloc, const NDq::TDqTaskSettings& tmp, const TString& traceId) override {
+    ITaskRunner::TPtr GetOld(std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const NDq::TDqTaskSettings& tmp, const TString& traceId) override {
         Y_UNUSED(alloc);
         Yql::DqsProto::TTaskMeta taskMeta;
         tmp.GetMeta().UnpackTo(&taskMeta);
         ui64 stageId = taskMeta.GetStageId();
         auto result = GetExecutorForTask(taskMeta.GetFiles(), taskMeta.GetSettings());
         auto [task, filesHolder] = PrepareTask(tmp, result.Get());
-        return new TTaskRunner(task, std::move(filesHolder), std::move(result), stageId, traceId);
+        return new TTaskRunner(alloc, task, std::move(filesHolder), std::move(result), stageId, traceId);
     }
 
-    TIntrusivePtr<NDq::IDqTaskRunner> Get(NKikimr::NMiniKQL::TScopedAlloc& alloc, const NDq::TDqTaskSettings& tmp, NDqProto::EDqStatsMode statsMode, const TString& traceId) override
+    TIntrusivePtr<NDq::IDqTaskRunner> Get(std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const NDq::TDqTaskSettings& tmp, NDqProto::EDqStatsMode statsMode, const TString& traceId) override
     {
         Y_UNUSED(statsMode);
         Y_UNUSED(alloc);
@@ -1933,7 +1932,7 @@ class TPipeFactory: public IProxyFactory {
 
         auto result = GetExecutorForTask(taskMeta.GetFiles(), taskMeta.GetSettings());
         auto [task, filesHolder] = PrepareTask(tmp, result.Get());
-        return new TDqTaskRunner(task, std::move(filesHolder), std::move(result), taskMeta.GetStageId(), traceId);
+        return new TDqTaskRunner(alloc, task, std::move(filesHolder), std::move(result), taskMeta.GetStageId(), traceId);
     }
 
 private:
diff --git a/ydb/library/yql/providers/dq/task_runner/tasks_runner_proxy.h b/ydb/library/yql/providers/dq/task_runner/tasks_runner_proxy.h
@@ -83,9 +83,9 @@ class IProxyFactory: public TThrRefBase, private TNonCopyable {
 public:
     using TPtr = TIntrusivePtr<IProxyFactory>;
 
-    virtual ITaskRunner::TPtr GetOld(NKikimr::NMiniKQL::TScopedAlloc& alloc, const NDq::TDqTaskSettings& task, const TString& traceId = "") = 0;
+    virtual ITaskRunner::TPtr GetOld(std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const NDq::TDqTaskSettings& task, const TString& traceId = "") = 0;
 
-    virtual TIntrusivePtr<NDq::IDqTaskRunner> Get(NKikimr::NMiniKQL::TScopedAlloc& alloc, const NDq::TDqTaskSettings& task, NDqProto::EDqStatsMode statsMode, const TString& traceId = "TODO") = 0;
+    virtual TIntrusivePtr<NDq::IDqTaskRunner> Get(std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> alloc, const NDq::TDqTaskSettings& task, NDqProto::EDqStatsMode statsMode, const TString& traceId = "TODO") = 0;
 };
 
 } // namespace NYql::NTaskRunnerProxy
diff --git a/ydb/library/yql/providers/dq/task_runner_actor/task_runner_actor.cpp b/ydb/library/yql/providers/dq/task_runner_actor/task_runner_actor.cpp
diff --git a/ydb/library/yql/tools/dq/worker_node/main.cpp b/ydb/library/yql/tools/dq/worker_node/main.cpp

Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ void TKqpComputeActor::DoBootstrap() {`
`69`	`69`	`settings.ReadRanges.push_back(readRange);`
`70`	`70`	`}`
`71`	`71`
`72`		`- auto taskRunner = MakeDqTaskRunner(TBase::GetAllocator(), execCtx, settings, logger);`
	`72`	`+ auto taskRunner = MakeDqTaskRunner(TBase::GetAllocatorPtr(), execCtx, settings, logger);`
`73`	`73`	`SetTaskRunner(taskRunner);`
`74`	`74`
`75`	`75`	`auto wakeup = [this]{ ContinueExecute(); };`
Original file line number	Diff line number	Diff line change
`@@ -216,7 +216,7 @@ void TKqpScanComputeActor::DoBootstrap() {`
`216`	`216`	`};`
`217`	`217`	`}`
`218`	`218`
`219`		`- auto taskRunner = MakeDqTaskRunner(GetAllocator(), execCtx, settings, logger);`
	`219`	`+ auto taskRunner = MakeDqTaskRunner(GetAllocatorPtr(), execCtx, settings, logger);`
`220`	`220`	`TBase::SetTaskRunner(taskRunner);`
`221`	`221`
`222`	`222`	`auto wakeup = [this] { ContinueExecute(); };`
Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,7 @@ void TEvKqpExecuter::TEvTxResponse::TakeResult(ui32 idx, NDq::TDqSerializedBatch`
`49`	`49`
`50`	`50`	`TEvKqpExecuter::TEvTxResponse::~TEvTxResponse() {`
`51`	`51`	`if (!TxResults.empty() && Y_LIKELY(AllocState)) {`
`52`		`- with_lock(AllocState->Alloc) {`
	`52`	`+ with_lock(*AllocState->Alloc) {`
`53`	`53`	`TxResults.crop(0);`
`54`	`54`	`}`
`55`	`55`	`}`
Original file line number	Diff line number	Diff line change
`@@ -166,7 +166,7 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {`
`166`	`166`	`TString BuildMemoryLimitExceptionMessage() const {`
`167`	`167`	`if (Request.TxAlloc) {`
`168`	`168`	`return TStringBuilder() << "Memory limit exception at " << CurrentStateFuncName()`
`169`		`- << ", current limit is " << Request.TxAlloc->Alloc.GetLimit() << " bytes.";`
	`169`	`+ << ", current limit is " << Request.TxAlloc->Alloc->GetLimit() << " bytes.";`
`170`	`170`	`}`
`171`	`171`	`return TStringBuilder() << "Memory limit exception at " << CurrentStateFuncName();`
`172`	`172`	`}`
Original file line number	Diff line number	Diff line change
`@@ -577,7 +577,7 @@ class TKiSourceCallableExecutionTransformer : public TAsyncCallbackTransformer<T`
`577`	`577`	`return peepHoleStatus;`
`578`	`578`	`}`
`579`	`579`
`580`		`- auto guard = Guard(SessionCtx->Query().QueryData->GetAllocState()->Alloc);`
	`580`	`+ auto guard = Guard(*SessionCtx->Query().QueryData->GetAllocState()->Alloc);`
`581`	`581`
`582`	`582`	`auto input = Build<TDqPhyStage>(ctx, pos)`
`583`	`583`	`.Inputs()`
Original file line number	Diff line number	Diff line change
`@@ -117,19 +117,17 @@ void TKqpExecuterTxResult::FillYdb(Ydb::ResultSet* ydbResult, TMaybe<ui64> rowsL`
`117`	`117`
`118`	`118`	`TTxAllocatorState::TTxAllocatorState(const IFunctionRegistry* functionRegistry,`
`119`	`119`	`TIntrusivePtr<ITimeProvider> timeProvider, TIntrusivePtr<IRandomProvider> randomProvider)`
`120`		`- : Alloc(__LOCATION__, NKikimr::TAlignedPagePoolCounters(), functionRegistry->SupportsSizedAllocators())`
`121`		`- , TypeEnv(Alloc)`
	`120`	`+ : Alloc(std::make_shared<NKikimr::NMiniKQL::TScopedAlloc>(__LOCATION__, NKikimr::TAlignedPagePoolCounters(), functionRegistry->SupportsSizedAllocators(), false))`
	`121`	`+ , TypeEnv(*Alloc)`
`122`	`122`	`, MemInfo("TQueryData")`
`123`		`- , HolderFactory(Alloc.Ref(), MemInfo, functionRegistry)`
	`123`	`+ , HolderFactory(Alloc->Ref(), MemInfo, functionRegistry)`
`124`	`124`	`{`
`125`		`- Alloc.Release();`
`126`	`125`	`TimeProvider = timeProvider;`
`127`	`126`	`RandomProvider = randomProvider;`
`128`	`127`	`}`
`129`	`128`
`130`	`129`	`TTxAllocatorState::~TTxAllocatorState()`
`131`	`130`	`{`
`132`		`- Alloc.Acquire();`
`133`	`131`	`}`
`134`	`132`
`135`	`133`	`std::pair<NKikimr::NMiniKQL::TType*, NUdf::TUnboxedValue> TTxAllocatorState::GetInternalBindingValue(`