diff --git a/ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp b/ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp index 4c2b2beedfcd..46abf76845a6 100644 --- a/ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp +++ b/ydb/library/yql/dq/actors/compute/dq_async_compute_actor.cpp @@ -461,18 +461,6 @@ class TDqAsyncComputeActor : public TDqComputeActorBaseFreeSpace; } - TGuard BindAllocator() override { - return TypeEnv->BindAllocator(); - } - - std::optional> MaybeBindAllocator() override { - std::optional> guard; - if (TypeEnv) { - guard.emplace(TypeEnv->BindAllocator()); - } - return guard; - } - void OnTaskRunnerCreated(NTaskRunnerActor::TEvTaskRunnerCreateFinished::TPtr& ev) { const auto& secureParams = ev->Get()->SecureParams; const auto& taskParams = ev->Get()->TaskParams; diff --git a/ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h b/ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h index 5b4a3fa59793..e0d3d3fa297f 100644 --- a/ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h +++ b/ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h @@ -477,7 +477,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped } { - auto guard = MaybeBindAllocator(); // Source/Sink could destroy mkql values inside PassAway, which requires allocator to be bound + auto guard = BindAllocator(); // Source/Sink could destroy mkql values inside PassAway, which requires allocator to be bound for (auto& [_, source] : SourcesMap) { if (source.Actor) { @@ -1087,12 +1087,8 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped TerminateSources(TIssues({TIssue(message)}), success); } - virtual TGuard BindAllocator() { - return TaskRunner->BindAllocator(); - } - - virtual std::optional> MaybeBindAllocator() { - return TaskRunner->BindAllocator(); + TGuard BindAllocator() { + return Guard(GetAllocator()); } virtual bool SayHelloOnBootstrap() { @@ -1503,7 +1499,7 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped .TypeEnv = typeEnv, .HolderFactory = holderFactory, .TaskCounters = TaskCounters, - .Alloc = TaskRunner ? Alloc : nullptr, + .Alloc = Alloc, .MemoryQuotaManager = MemoryLimits.MemoryQuotaManager, .SourceSettings = (!settings.empty() ? settings.at(inputIndex) : nullptr), .Arena = Task.GetArena(), diff --git a/ydb/library/yql/dq/actors/compute/dq_sync_compute_actor_base.h b/ydb/library/yql/dq/actors/compute/dq_sync_compute_actor_base.h index 8d2ffb00924b..eb69e10eea82 100644 --- a/ydb/library/yql/dq/actors/compute/dq_sync_compute_actor_base.h +++ b/ydb/library/yql/dq/actors/compute/dq_sync_compute_actor_base.h @@ -51,8 +51,9 @@ class TDqSyncComputeActorBase: public TDqComputeActorBaseTaskRunner); - auto guard = this->TaskRunner->BindAllocator(this->MemoryQuota->GetMkqlMemoryLimit()); + auto guard = TBase::BindAllocator(); auto* alloc = guard.GetMutex(); + alloc->SetLimit(this->MemoryQuota->GetMkqlMemoryLimit()); this->MemoryQuota->TrySetIncreaseMemoryLimitCallback(alloc);