Refactor memory tracking (ydb-platform#4510)

gridnevvvit · MrLolthe1st · commit e8d75a4bdf27 · 2024-05-28T12:31:53.000Z
diff --git a/ydb/core/kqp/node_service/kqp_node_service.cpp b/ydb/core/kqp/node_service/kqp_node_service.cpp
@@ -285,12 +285,11 @@ class TKqpNodeService : public TActorBootstrapped<TKqpNodeService> {
         auto requester = ev->Sender;
 
         ui64 txId = msg.GetTxId();
-        bool isScan = msg.HasSnapshot();
         const ui64 outputChunkMaxSize = msg.GetOutputChunkMaxSize();
 
         YQL_ENSURE(msg.GetStartAllOrFail()); // todo: support partial start
 
-        LOG_D("TxId: " << txId << ", new " << (isScan ? "scan " : "") << "compute tasks request from " << requester
+        LOG_D("TxId: " << txId << ", new compute tasks request from " << requester
             << " with " << msg.GetTasks().size() << " tasks: " << TasksIdsStr(msg.GetTasks()));
 
         NKqpNode::TTasksRequest request;
@@ -339,22 +338,10 @@ class TKqpNodeService : public TActorBootstrapped<TKqpNodeService> {
             LOG_D("TxId: " << txId << ", task: " << taskCtx.TaskId << ", requested memory: " << taskCtx.Memory);
 
             requestChannels += estimation.ChannelBuffersCount;
-            request.TotalMemory += taskCtx.Memory;
         }
 
         LOG_D("TxId: " << txId << ", channels: " << requestChannels
-            << ", computeActors: " << msg.GetTasks().size() << ", memory: " << request.TotalMemory);
-
-        auto txMemory = bucket.GetTxMemory(txId, memoryPool) + request.TotalMemory;
-        if (txMemory > Config.GetQueryMemoryLimit()) {
-            LOG_N("TxId: " << txId << ", requested too many memory: " << request.TotalMemory
-                << "(" << txMemory << " for this Tx), limit: " << Config.GetQueryMemoryLimit());
-
-            Counters->RmNotEnoughMemory->Inc();
-
-            return ReplyError(txId, request.Executer, msg, NKikimrKqp::TEvStartKqpTasksResponse::QUERY_MEMORY_LIMIT_EXCEEDED,
-                TStringBuilder() << "Required: " << txMemory << ", limit: " << Config.GetQueryMemoryLimit());
-        }
+            << ", computeActors: " << msg.GetTasks().size() << ", memory: " << request.CalculateTotalMemory());
 
         TVector<ui64> allocatedTasks;
         allocatedTasks.reserve(msg.GetTasks().size());
@@ -368,29 +355,13 @@ class TKqpNodeService : public TActorBootstrapped<TKqpNodeService> {
 
             NRm::TKqpNotEnoughResources resourcesResponse;
             if (!ResourceManager()->AllocateResources(txId, task.first, resourcesRequest, &resourcesResponse)) {
-                NKikimrKqp::TEvStartKqpTasksResponse::ENotStartedTaskReason failReason = NKikimrKqp::TEvStartKqpTasksResponse::INTERNAL_ERROR;
-                TStringBuilder error;
-                if (resourcesResponse.ScanQueryMemory()) {
-                    error << "TxId: " << txId << ", NodeId: " << SelfId().NodeId() << ", not enough memory, requested " << task.second.Memory;
-                    LOG_N(error);
-
-                    failReason = NKikimrKqp::TEvStartKqpTasksResponse::NOT_ENOUGH_MEMORY;
-                }
-
-                if (resourcesResponse.QueryMemoryLimit()) {
-                    error << "TxId: " << txId << ", NodeId: " << SelfId().NodeId() << ", memory limit exceeded, requested " << task.second.Memory;
-                    LOG_N(error);
-
-                    failReason = NKikimrKqp::TEvStartKqpTasksResponse::QUERY_MEMORY_LIMIT_EXCEEDED;
-                }
-
                 for (ui64 taskId : allocatedTasks) {
                     ResourceManager()->FreeResources(txId, taskId);
                 }
 
                 ResourceManager()->FreeExecutionUnits(executionUnits);
 
-                ReplyError(txId, request.Executer, msg, failReason, error);
+                ReplyError(txId, request.Executer, msg, resourcesResponse.GetStatus(), resourcesResponse.GetFailReason());
                 return;
             }
 
@@ -553,19 +524,13 @@ class TKqpNodeService : public TActorBootstrapped<TKqpNodeService> {
 
     void TerminateTx(ui64 txId, const TString& reason) {
         auto& bucket = GetStateBucketByTx(Buckets, txId);
-        auto tasksToAbort = bucket.RemoveTx(txId);
+        auto tasksToAbort = bucket.GetTasksByTxId(txId);
 
         if (!tasksToAbort.empty()) {
-            LOG_D("TxId: " << txId << ", cancel granted resources");
-            ResourceManager()->FreeResources(txId);
-
-            for (const auto& tasksRequest: tasksToAbort) {
-                ResourceManager()->FreeExecutionUnits(tasksRequest.InFlyTasks.size());
-                for (const auto& [taskId, task] : tasksRequest.InFlyTasks) {
-                    auto abortEv = MakeHolder<TEvKqp::TEvAbortExecution>(NYql::NDqProto::StatusIds::UNSPECIFIED,
-                        reason);
-                    Send(task.ComputeActorId, abortEv.Release());
-                }
+            for (const auto& [taskId, computeActorId]: tasksToAbort) {
+                auto abortEv = MakeHolder<TEvKqp::TEvAbortExecution>(NYql::NDqProto::StatusIds::UNSPECIFIED,
+                    reason);
+                Send(computeActorId, abortEv.Release());
             }
         }
     }
diff --git a/ydb/core/kqp/node_service/kqp_node_state.h b/ydb/core/kqp/node_service/kqp_node_state.h
@@ -30,25 +30,25 @@ struct TTasksRequest {
     // when task is finished it will be removed from this map
     THashMap<ui64, TTaskContext> InFlyTasks;
     TInstant Deadline;
-    ui64 TotalMemory = 0;
     TActorId Executer;
     TActorId TimeoutTimer;
     bool ExecutionCancelled = false;
+
+    ui64 CalculateTotalMemory() const {
+        ui64 result = 0;
+        for(auto [id, task]: InFlyTasks) {
+            result += task.Memory;
+        }
+        return result;
+    }
 };
 
 struct TTxMeta {
-    ui64 TotalMemory = 0;
     NRm::EKqpMemoryPool MemoryPool = NRm::EKqpMemoryPool::Unspecified;
     ui32 TotalComputeActors = 0;
     TInstant StartTime;
 };
 
-struct TRemoveTaskContext {
-    ui64 TotalMemory = 0;
-    ui64 ComputeActorsNumber = 0;
-    bool FinixTx = false;
-};
-
 class TState {
 public:
     struct TRequestId {
@@ -57,7 +57,6 @@ class TState {
     };
 
     struct TRemoveTaskContext {
-        ui64 TotalMemory = 0;
         ui64 ComputeActorsNumber = 0;
         bool FinixTx = false;
         TActorId Requester;
@@ -73,18 +72,9 @@ class TState {
         return Requests.contains(std::make_pair(txId, requester));
     }
 
-    ui64 GetTxMemory(ui64 txId, NRm::EKqpMemoryPool memoryPool) const {
-        TReadGuard guard(RWLock);
-        if (auto* meta = Meta.FindPtr(txId)) {
-            return meta->MemoryPool == memoryPool ? meta->TotalMemory : 0;
-        }
-        return 0;
-    }
-
     void NewRequest(ui64 txId, const TActorId& requester, TTasksRequest&& request, NRm::EKqpMemoryPool memoryPool) {
         TWriteGuard guard(RWLock);
         auto& meta = Meta[txId];
-        meta.TotalMemory += request.TotalMemory;
         meta.TotalComputeActors += request.InFlyTasks.size();
         if (!meta.StartTime) {
             meta.StartTime = TAppData::TimeProvider->Now();
@@ -109,21 +99,15 @@ class TState {
 
             auto taskIt = requestIt->second.InFlyTasks.find(taskId);
             if (taskIt != requestIt->second.InFlyTasks.end()) {
-                auto task = std::move(taskIt->second);
                 requestIt->second.InFlyTasks.erase(taskIt);
-
-                Y_DEBUG_ABORT_UNLESS(requestIt->second.TotalMemory >= task.Memory);
-                requestIt->second.TotalMemory -= task.Memory;
                 requestIt->second.ExecutionCancelled |= !success;
 
                 auto& meta = Meta[txId];
-                Y_DEBUG_ABORT_UNLESS(meta.TotalMemory >= task.Memory);
                 Y_DEBUG_ABORT_UNLESS(meta.TotalComputeActors >= 1);
-                meta.TotalMemory -= task.Memory;
                 meta.TotalComputeActors--;
 
                 auto ret = TRemoveTaskContext{
-                    requestIt->second.TotalMemory, requestIt->second.InFlyTasks.size(), meta.TotalComputeActors == 0, senderIt->second
+                    requestIt->second.InFlyTasks.size(), meta.TotalComputeActors == 0, senderIt->second
                 };
 
                 if (requestIt->second.InFlyTasks.empty()) {
@@ -157,24 +141,21 @@ class TState {
         return RemoveRequestImpl(txId, requester);
     }
 
-    std::vector<TTasksRequest> RemoveTx(ui64 txId) {
+    // return the vector of pairs where the first element is a taskId
+    // and the second one is the compute actor id associated with this task.
+    std::vector<std::pair<ui64, TActorId>> GetTasksByTxId(ui64 txId) {
         TWriteGuard guard(RWLock);
-        Meta.erase(txId);
-
         YQL_ENSURE(Requests.size() == SenderIdsByTxId.size());
         const auto senders = SenderIdsByTxId.equal_range(txId);
-        std::vector<TTasksRequest> ret;
+        std::vector<std::pair<ui64, TActorId>> ret;
         for (auto senderIt = senders.first; senderIt != senders.second; ++senderIt) {
             auto requestIt = Requests.find(*senderIt);
             YQL_ENSURE(requestIt != Requests.end());
-
-            ret.push_back(std::move(requestIt->second));
-            Requests.erase(requestIt);
+            for(const auto& [taskId, task] : requestIt->second.InFlyTasks) {
+                ret.push_back({taskId, task.ComputeActorId});
+            }
         }
 
-        SenderIdsByTxId.erase(txId);
-        YQL_ENSURE(Requests.size() == SenderIdsByTxId.size());
-
         return ret;
     }
 
@@ -208,15 +189,13 @@ class TState {
         for (auto& [txId, requests] : byTx) {
             auto& meta = Meta[txId];
             str << "  TxId: " << txId << Endl;
-            str << "    Memory: " << meta.TotalMemory << Endl;
             str << "    MemoryPool: " << (ui32) meta.MemoryPool << Endl;
             str << "    Compute actors: " << meta.TotalComputeActors << Endl;
             str << "    Start time: " << meta.StartTime << Endl;
             str << "    Requests:" << Endl;
             for (auto& [requester, request] : requests) {
                 str << "      Requester: " << requester << Endl;
                 str << "        Deadline: " << request->Deadline << Endl;
-                str << "        Memory: " << request->TotalMemory << Endl;
                 str << "        In-fly tasks:" << Endl;
                 for (auto& [taskId, task] : request->InFlyTasks) {
                     str << "          Task: " << taskId << Endl;
@@ -250,9 +229,7 @@ class TState {
         YQL_ENSURE(Requests.size() == SenderIdsByTxId.size());
 
         auto& meta = Meta[txId];
-        Y_DEBUG_ABORT_UNLESS(meta.TotalMemory >= ret->TotalMemory);
         Y_DEBUG_ABORT_UNLESS(meta.TotalComputeActors >= 1);
-        meta.TotalMemory -= ret->TotalMemory;
         meta.TotalComputeActors -= ret->InFlyTasks.size();
 
         if (meta.TotalComputeActors == 0) {
diff --git a/ydb/core/kqp/node_service/kqp_node_ut.cpp b/ydb/core/kqp/node_service/kqp_node_ut.cpp
@@ -482,8 +482,8 @@ void KqpNode::NotEnoughMemory() {
         UNIT_ASSERT_VALUES_EQUAL(1, record.GetTxId());
         UNIT_ASSERT_VALUES_EQUAL(1, record.GetNotStartedTasks().size());
         auto& task = record.GetNotStartedTasks()[0];
-        UNIT_ASSERT_EQUAL(NKikimrKqp::TEvStartKqpTasksResponse::QUERY_MEMORY_LIMIT_EXCEEDED, task.GetReason());
-        UNIT_ASSERT_STRINGS_EQUAL("Required: 201000, limit: 30000", task.GetMessage());
+        UNIT_ASSERT_EQUAL(NKikimrKqp::TEvStartKqpTasksResponse::NOT_ENOUGH_MEMORY, task.GetReason());
+        UNIT_ASSERT_STRING_CONTAINS(task.GetMessage(), "Not enough memory for query, requested: 201000");
     }
 
     AssertResourceBrokerSensors(0, 0, 0, 0, 0);
@@ -670,6 +670,13 @@ void KqpNode::ExecuterLost() {
         UNIT_ASSERT_VALUES_EQUAL("executer lost", abortEvent->Get()->Record.GetLegacyMessage());
     }
 
+    {
+        NKikimr::TActorSystemStub stub;
+        TMap<ui64, TMockComputeActor> Task2ActorEmpty;
+        CompFactory->Task2Actor.swap(Task2ActorEmpty);
+        Task2ActorEmpty.clear();
+    }
+
     size_t iterations = 30;
     while (KqpCounters->RmComputeActors->Val() != 0 && iterations > 0) {
         Sleep(TDuration::MilliSeconds(300));
@@ -731,6 +738,13 @@ void KqpNode::TerminateTx() {
             auto abortEvent = Runtime->GrabEdgeEvent<TEvKqp::TEvAbortExecution>(computeActor.ActorId);
             UNIT_ASSERT_VALUES_EQUAL("terminate", abortEvent->Get()->Record.GetLegacyMessage());
         }
+
+        {
+            NKikimr::TActorSystemStub stub;
+            TMap<ui64, TMockComputeActor> Task2ActorEmpty;
+            CompFactory->Task2Actor.swap(Task2ActorEmpty);
+            Task2ActorEmpty.clear();
+        }
     }
 
     UNIT_ASSERT_VALUES_EQUAL(KqpCounters->RmComputeActors->Val(), 0);
diff --git a/ydb/core/kqp/rm_service/kqp_rm_service.cpp b/ydb/core/kqp/rm_service/kqp_rm_service.cpp
@@ -212,10 +212,12 @@ class TKqpResourceManager : public IKqpResourceManager {
 
         with_lock (Lock) {
             if (Y_UNLIKELY(!ResourceBroker)) {
-                LOG_AS_E("AllocateResources: not ready yet. TxId: " << txId << ", taskId: " << taskId);
+                TStringBuilder reason;
+                reason << "AllocateResources: not ready yet. TxId: " << txId << ", taskId: " << taskId;
                 if (details) {
-                    details->SetNotReady();
+                    details->FailReason = reason;
                 }
+
                 return false;
             }
 
@@ -228,9 +230,11 @@ class TKqpResourceManager : public IKqpResourceManager {
 
         if (!hasScanQueryMemory) {
             Counters->RmNotEnoughMemory->Inc();
-            LOG_AS_N("TxId: " << txId << ", taskId: " << taskId << ". Not enough ScanQueryMemory, requested: " << resources.Memory);
+            TStringBuilder reason;
+            reason << "TxId: " << txId << ", taskId: " << taskId << ". Not enough memory for query, requested: " << resources.Memory;
             if (details) {
-                details->SetScanQueryMemory();
+                details->FailReason = reason;
+                details->Status = NKikimrKqp::TEvStartKqpTasksResponse::NOT_ENOUGH_MEMORY;
             }
             return false;
         }
@@ -242,18 +246,21 @@ class TKqpResourceManager : public IKqpResourceManager {
         auto& txBucket = TxBucket(txId);
         with_lock (txBucket.Lock) {
             if (auto it = txBucket.Txs.find(txId); it != txBucket.Txs.end()) {
-                if (it->second.TxScanQueryMemory + resources.Memory > queryMemoryLimit) {
+                ui64 txTotalRequestedMemory = it->second.TxScanQueryMemory + resources.Memory;
+                if (txTotalRequestedMemory > queryMemoryLimit) {
                     auto unguard = ::Unguard(txBucket.Lock);
 
                     with_lock (Lock) {
                         ScanQueryMemoryResource.Release(resources.Memory);
                     } // with_lock (Lock)
 
                     Counters->RmNotEnoughMemory->Inc();
-                    LOG_AS_N("TxId: " << txId << ", taskId: " << taskId << ". Query memory limit exceeded: "
-                        << "requested " << (it->second.TxScanQueryMemory + resources.Memory));
+                    TStringBuilder reason;
+                    reason << "TxId: " << txId << ", taskId: " << taskId << ". Query memory limit exceeded: "
+                        << "requested " << txTotalRequestedMemory;
                     if (details) {
-                        details->SetQueryMemoryLimit();
+                        details->FailReason = reason;
+                        details->Status = NKikimrKqp::TEvStartKqpTasksResponse::QUERY_MEMORY_LIMIT_EXCEEDED;
                     }
                     return false;
                 }
@@ -271,10 +278,13 @@ class TKqpResourceManager : public IKqpResourceManager {
                 } // with_lock (Lock)
 
                 Counters->RmNotEnoughMemory->Inc();
-                LOG_AS_N("TxId: " << txId << ", taskId: " << taskId << ". Not enough ScanQueryMemory: "
-                    << "requested " << resources.Memory);
+                TStringBuilder reason;
+                reason << "TxId: " << txId << ", taskId: " << taskId << ". Not enough ScanQueryMemory: "
+                    << "requested " << resources.Memory;
+                LOG_AS_N(reason);
                 if (details) {
-                    details->SetScanQueryMemory();
+                    details->Status = NKikimrKqp::TEvStartKqpTasksResponse::NOT_ENOUGH_MEMORY;
+                    details->FailReason = reason;
                 }
                 return false;
             }
diff --git a/ydb/core/kqp/rm_service/kqp_rm_service.h b/ydb/core/kqp/rm_service/kqp_rm_service.h
@@ -45,17 +45,16 @@ struct TKqpResourcesRequest {
 
 /// detailed information on allocation failure
 struct TKqpNotEnoughResources {
-    std::bitset<32> State;
+    NKikimrKqp::TEvStartKqpTasksResponse::ENotStartedTaskReason Status = NKikimrKqp::TEvStartKqpTasksResponse::INTERNAL_ERROR;
+    TString FailReason;
 
-    bool NotReady() const         { return State.test(0); }
-    bool QueryMemoryLimit() const { return State.test(1); }
-    bool ScanQueryMemory() const  { return State.test(2); }
-    bool DataQueryMemory() const  { return State.test(3); }
+    NKikimrKqp::TEvStartKqpTasksResponse::ENotStartedTaskReason GetStatus() const {
+        return Status;
+    }
 
-    void SetNotReady()         { State.set(0); }
-    void SetQueryMemoryLimit() { State.set(1); }
-    void SetScanQueryMemory()  { State.set(2); }
-    void SetDataQueryMemory()  { State.set(3); }
+    TString GetFailReason() const {
+        return FailReason;
+    }
 };
 
 /// local resources snapshot
