fixed potential data race

Author: gridnevvvit

ydb/core/kqp/rm_service/kqp_rm_service.cpp

@@ -107,18 +107,29 @@ struct TTxState {
     ui32 TxExecutionUnits = 0;
     TInstant CreatedAt;
 
-    bool IsDataQuery = false;
-
-    TTaskState& Allocated(ui64 taskId, TInstant now, const TKqpResourcesRequest& resources) {
+    TTaskState& Allocated(ui64 taskId, TInstant now, const TKqpResourcesRequest& resources, bool memoryAsExternal = false) {
+        ui64 externalMemory = resources.ExternalMemory;
+        ui64 resourceBrokerMemory = 0;
+        if (memoryAsExternal) {
+            externalMemory += resources.Memory;
+        } else {
+            resourceBrokerMemory = resources.Memory;
+        }
 
-        TxScanQueryMemory += resources.Memory;
+        TxExternalDataQueryMemory += externalMemory;
+        TxScanQueryMemory += resourceBrokerMemory;
         if (!CreatedAt) {
             CreatedAt = now;
         }
 
+        if (resources.ExecutionUnits) {
+            Y_ABORT_UNLESS(!Tasks.contains(taskId));
+        }
+
         auto& taskState = Tasks[taskId];
         taskState.ExecutionUnits += resources.ExecutionUnits;
-        taskState.ScanQueryMemory += resources.Memory;
+        taskState.ScanQueryMemory += resourceBrokerMemory;
+        taskState.ExternalDataQueryMemory += externalMemory;
         if (!taskState.CreatedAt) {
             taskState.CreatedAt = now;
         }
@@ -233,19 +244,28 @@ class TKqpResourceManager : public IKqpResourceManager {
             }
         };
 
-        if (resources.MemoryPool == EKqpMemoryPool::DataQuery) {
-            NotifyExternalResourcesAllocated(txId, taskId, resources);
-            return result;
-        }
-
-        Y_ABORT_UNLESS(resources.MemoryPool == EKqpMemoryPool::ScanQuery);
         if (Y_UNLIKELY(resources.Memory == 0)) {
             return result;
         }
 
         auto now = ActorSystem->Timestamp();
         bool hasScanQueryMemory = true;
         ui64 queryMemoryLimit = 0;
+        // NOTE(gvit): the first memory request from the data query pool always satisfied.
+        // all other requests are not guaranteed to be satisfied.
+        // In the nearest future we need to implement several layers of memory requests.
+        bool isFirstAllocationRequest = (resources.ExecutionUnits > 0 && resources.MemoryPool == EKqpMemoryPool::DataQuery);
+        if (isFirstAllocationRequest) {
+            auto& txBucket = TxBucket(txId);
+            with_lock(txBucket.Lock) {
+                auto& tx = txBucket.Txs[txId];
+                tx.Allocated(taskId, now, resources, /*memoryAsExternal=*/true);
+                ExternalDataQueryMemory.fetch_add(resources.Memory + resources.ExternalMemory);
+                Counters->RmExternalMemory->Add(resources.Memory + resources.ExternalMemory);
+            }
+
+            return result;
+        }
 
         with_lock (Lock) {
             if (Y_UNLIKELY(!ResourceBroker)) {
@@ -286,15 +306,14 @@ class TKqpResourceManager : public IKqpResourceManager {
                 }
             };
 
-            if (auto it = txBucket.Txs.find(txId); it != txBucket.Txs.end()) {
-                ui64 txTotalRequestedMemory = it->second.TxScanQueryMemory + resources.Memory;
-                if (txTotalRequestedMemory > queryMemoryLimit) {
-                    TStringBuilder reason;
-                    reason << "TxId: " << txId << ", taskId: " << taskId << ". Query memory limit exceeded: "
-                        << "requested " << txTotalRequestedMemory;
-                    result.SetError(NKikimrKqp::TEvStartKqpTasksResponse::QUERY_MEMORY_LIMIT_EXCEEDED, reason);
-                    return result;
-                }
+            auto& tx = txBucket.Txs[txId];
+            ui64 txTotalRequestedMemory = tx.TxScanQueryMemory + resources.Memory;
+            if (txTotalRequestedMemory > queryMemoryLimit) {
+                TStringBuilder reason;
+                reason << "TxId: " << txId << ", taskId: " << taskId << ". Query memory limit exceeded: "
+                    << "requested " << txTotalRequestedMemory;
+                result.SetError(NKikimrKqp::TEvStartKqpTasksResponse::QUERY_MEMORY_LIMIT_EXCEEDED, reason);
+                return result;
             }
 
             bool allocated = ResourceBroker->SubmitTaskInstant(
@@ -310,7 +329,7 @@ class TKqpResourceManager : public IKqpResourceManager {
                 return result;
             }
 
-            auto& taskState = txBucket.Txs[txId].Allocated(Counters, taskId, now, resources);
+            auto& taskState = tx.Allocated(taskId, now, resources);
             if (!taskState.ResourceBrokerTaskId) {
                 taskState.ResourceBrokerTaskId = rbTaskId;
             } else {
@@ -334,7 +353,6 @@ class TKqpResourceManager : public IKqpResourceManager {
     void FreeResources(ui64 txId, ui64 taskId, const TKqpResourcesRequest& resources) override {
 
         if (resources.MemoryPool == EKqpMemoryPool::DataQuery) {
-            NotifyExternalResourcesFreed(txId, taskId, resources);
             return;
         }
 
@@ -374,6 +392,7 @@ class TKqpResourceManager : public IKqpResourceManager {
 
     void FreeResources(ui64 txId, ui64 taskId) override {
         ui64 releaseScanQueryMemory = 0;
+        ui64 releaseExternalDataQueryMemory = 0;
         ui32 remainsTasks = 0;
 
         auto& txBucket = TxBucket(txId);
@@ -387,34 +406,39 @@ class TKqpResourceManager : public IKqpResourceManager {
                 return;
             }
 
-            auto taskIt = txIt->second.Tasks.find(taskId);
-            if (taskIt == txIt->second.Tasks.end()) {
+            auto& tx = txIt->second;
+            auto taskIt = tx.Tasks.find(taskId);
+            if (taskIt == tx.Tasks.end()) {
                 return;
             }
 
-            if (taskIt->second.ExecutionUnits) {
-                FreeExecutionUnits(std::exchange(taskIt->second.ExecutionUnits, 0));
-            }
-
-            if (txIt->second.IsDataQuery) {
-                guard.Clear();
-                return NotifyExternalResourcesFreed(txId, taskId);
+            auto& task = taskIt->second;
+            if (task.ExecutionUnits) {
+                FreeExecutionUnits(task.ExecutionUnits);
             }
 
-            releaseScanQueryMemory = taskIt->second.ScanQueryMemory;
+            releaseExternalDataQueryMemory = task.ExternalDataQueryMemory;
+            releaseScanQueryMemory = task.ScanQueryMemory;
 
-            bool finished = ResourceBroker->FinishTaskInstant(
-                TEvResourceBroker::TEvFinishTask(taskIt->second.ResourceBrokerTaskId), SelfId);
-            Y_DEBUG_ABORT_UNLESS(finished);
+            if (task.ResourceBrokerTaskId) {
+                bool finished = ResourceBroker->FinishTaskInstant(
+                    TEvResourceBroker::TEvFinishTask(task.ResourceBrokerTaskId), SelfId);
+                Y_DEBUG_ABORT_UNLESS(finished);
+            }
 
-            remainsTasks = txIt->second.Tasks.size() - 1;
+            remainsTasks = tx.Tasks.size() - 1;
 
             if (remainsTasks == 0) {
                 txBucket.Txs.erase(txIt);
             } else {
-                txIt->second.Tasks.erase(taskIt);
-                txIt->second.TxScanQueryMemory -= releaseScanQueryMemory;
+                tx.Tasks.erase(taskIt);
+                tx.TxScanQueryMemory -= releaseScanQueryMemory;
+                tx.TxExternalDataQueryMemory -= releaseExternalDataQueryMemory;
             }
+
+            i64 prev = ExternalDataQueryMemory.fetch_sub(releaseExternalDataQueryMemory);
+            Counters->RmExternalMemory->Sub(releaseExternalDataQueryMemory);
+            Y_DEBUG_ABORT_UNLESS(prev >= 0);
         } // with_lock (txBucket.Lock)
 
         with_lock (Lock) {
@@ -442,101 +466,11 @@ class TKqpResourceManager : public IKqpResourceManager {
 
         auto& txBucket = TxBucket(txId);
         with_lock (txBucket.Lock) {
-            auto& tx = txBucket.Txs[txId];
-            tx.IsDataQuery = true;
-            auto& task = tx.Tasks[taskId];
-
-            task.ExternalDataQueryMemory = resources.Memory;
-            tx.TxExternalDataQueryMemory += resources.Memory;
+            txBucket.Txs[txId].Allocated(taskId, TInstant(), resources);
+            ExternalDataQueryMemory.fetch_add(resources.ExternalMemory);
+            Counters->RmExternalMemory->Add(resources.ExternalMemory);
         } // with_lock (txBucket.Lock)
 
-        with_lock (Lock) {
-            ExternalDataQueryMemory += resources.Memory;
-        } // with_lock (Lock)
-
-        Counters->RmExternalMemory->Add(resources.Memory);
-
-        FireResourcesPublishing();
-    }
-
-    void NotifyExternalResourcesFreed(ui64 txId, ui64 taskId, const TKqpResourcesRequest& resources) override {
-        LOG_AS_D("TxId: " << txId << ", taskId: " << taskId << ". External free: " << resources.ToString());
-
-        YQL_ENSURE(resources.MemoryPool == EKqpMemoryPool::DataQuery);
-
-        ui64 releaseMemory = 0;
-
-        auto& txBucket = TxBucket(txId);
-        with_lock (txBucket.Lock) {
-            auto txIt = txBucket.Txs.find(txId);
-            if (txIt == txBucket.Txs.end()) {
-                return;
-            }
-
-            auto taskIt = txIt->second.Tasks.find(taskId);
-            if (taskIt == txIt->second.Tasks.end()) {
-                return;
-            }
-
-            if (taskIt->second.ExternalDataQueryMemory <= resources.Memory) {
-                releaseMemory = taskIt->second.ExternalDataQueryMemory;
-                if (txIt->second.Tasks.size() == 1) {
-                    txBucket.Txs.erase(txId);
-                } else {
-                    txIt->second.Tasks.erase(taskIt);
-                    txIt->second.TxExternalDataQueryMemory -= releaseMemory;
-                }
-            } else {
-                releaseMemory = resources.Memory;
-                taskIt->second.ExternalDataQueryMemory -= resources.Memory;
-            }
-        } // with_lock (txBucket.Lock)
-
-        with_lock (Lock) {
-            Y_DEBUG_ABORT_UNLESS(ExternalDataQueryMemory >= releaseMemory);
-            ExternalDataQueryMemory -= releaseMemory;
-        } // with_lock (Lock)
-
-        Counters->RmExternalMemory->Sub(releaseMemory);
-        Y_DEBUG_ABORT_UNLESS(Counters->RmExternalMemory->Val() >= 0);
-
-        FireResourcesPublishing();
-    }
-
-    void NotifyExternalResourcesFreed(ui64 txId, ui64 taskId) override {
-        LOG_AS_D("TxId: " << txId << ", taskId: " << taskId << ". External free.");
-
-        ui64 releaseMemory = 0;
-
-        auto& txBucket = TxBucket(txId);
-        with_lock (txBucket.Lock) {
-            auto txIt = txBucket.Txs.find(txId);
-            if (txIt == txBucket.Txs.end()) {
-                return;
-            }
-
-            auto taskIt = txIt->second.Tasks.find(taskId);
-            if (taskIt == txIt->second.Tasks.end()) {
-                return;
-            }
-
-            releaseMemory = taskIt->second.ExternalDataQueryMemory;
-
-            if (txIt->second.Tasks.size() == 1) {
-                txBucket.Txs.erase(txId);
-            } else {
-                txIt->second.Tasks.erase(taskIt);
-                txIt->second.TxExternalDataQueryMemory -= releaseMemory;
-            }
-        } // with_lock (txBucket.Lock)
-
-        with_lock (Lock) {
-            Y_DEBUG_ABORT_UNLESS(ExternalDataQueryMemory >= releaseMemory);
-            ExternalDataQueryMemory -= releaseMemory;
-        } // with_lock (Lock)
-
-        Counters->RmExternalMemory->Sub(releaseMemory);
-        Y_DEBUG_ABORT_UNLESS(Counters->RmExternalMemory->Val() >= 0);
 
         FireResourcesPublishing();
     }
@@ -629,7 +563,7 @@ class TKqpResourceManager : public IKqpResourceManager {
     std::atomic<i32> ExecutionUnitsResource;
     std::atomic<i32> ExecutionUnitsLimit;
     TLimitedResource<ui64> ScanQueryMemoryResource;
-    ui64 ExternalDataQueryMemory = 0;
+    std::atomic<i64> ExternalDataQueryMemory = 0;
 
     // current state
     std::array<TTxStatesBucket, BucketsCount> Buckets;
@@ -957,7 +891,7 @@ class TKqpResourceManagerActor : public TActorBootstrapped<TKqpResourceManagerAc
                 str << "State storage key: " << WbState.Tenant << Endl;
                 with_lock (ResourceManager->Lock) {
                     str << "ScanQuery memory resource: " << ResourceManager->ScanQueryMemoryResource.ToString() << Endl;
-                    str << "External DataQuery memory: " << ResourceManager->ExternalDataQueryMemory << Endl;
+                    str << "External DataQuery memory: " << ResourceManager->ExternalDataQueryMemory.load() << Endl;
                     str << "ExecutionUnits resource: " << ResourceManager->ExecutionUnitsResource.load() << Endl;
                 }
                 str << "Last resource broker task id: " << ResourceManager->LastResourceBrokerTaskId.load() << Endl;

ydb/core/kqp/rm_service/kqp_rm_service.h

@@ -37,10 +37,11 @@ struct TKqpResourcesRequest {
     ui64 ExecutionUnits = 0;
     EKqpMemoryPool MemoryPool = EKqpMemoryPool::Unspecified;
     ui64 Memory = 0;
+    ui64 ExternalMemory = 0;
 
     TString ToString() const {
         return TStringBuilder() << "TKqpResourcesRequest{ MemoryPool: " << (ui32) MemoryPool << ", Memory: " << Memory
-           << " }";
+            << "ExternalMemory: " << ExternalMemory << " }";
     }
 };
 
@@ -83,14 +84,11 @@ class IKqpResourceManager : private TNonCopyable {
     virtual TKqpRMAllocateResult AllocateResources(ui64 txId, ui64 taskId, const TKqpResourcesRequest& resources) = 0;
 
     using TResourcesAllocatedCallback = std::function<void(NActors::TActorSystem* as)>;
-    using TNotEnoughtResourcesCallback = std::function<void(NActors::TActorSystem* as, const TString& reason, bool byTimeout)>;
 
     virtual void FreeResources(ui64 txId, ui64 taskId, const TKqpResourcesRequest& resources) = 0;
     virtual void FreeResources(ui64 txId, ui64 taskId) = 0;
 
     virtual void NotifyExternalResourcesAllocated(ui64 txId, ui64 taskId, const TKqpResourcesRequest& resources) = 0;
-    virtual void NotifyExternalResourcesFreed(ui64 txId, ui64 taskId, const TKqpResourcesRequest& resources) = 0;
-    virtual void NotifyExternalResourcesFreed(ui64 txId, ui64 taskId) = 0;
 
     virtual void RequestClusterResourcesInfo(TOnResourcesSnapshotCallback&& callback) = 0;
 

ydb/core/tx/datashard/build_kqp_data_tx_out_rs_unit.cpp

@@ -85,13 +85,13 @@ EExecutionStatus TBuildKqpDataTxOutRSUnit::Execute(TOperation::TPtr op, TTransac
 
         NKqp::NRm::TKqpResourcesRequest req;
         req.MemoryPool = NKqp::NRm::EKqpMemoryPool::DataQuery;
-        req.Memory = txc.GetMemoryLimit();
+        req.ExternalMemory = txc.GetMemoryLimit();
         ui64 taskId = dataTx->GetFirstKqpTaskId();
 
         NKqp::GetKqpResourceManager()->NotifyExternalResourcesAllocated(tx->GetTxId(), taskId, req);
 
         Y_DEFER {
-            NKqp::GetKqpResourceManager()->NotifyExternalResourcesFreed(tx->GetTxId(), taskId);
+            NKqp::GetKqpResourceManager()->FreeResources(tx->GetTxId(), taskId);
         };
 
         LOG_T("Operation " << *op << " (build_kqp_data_tx_out_rs) at " << tabletId

ydb/core/tx/datashard/execute_kqp_data_tx_unit.cpp

@@ -189,7 +189,7 @@ EExecutionStatus TExecuteKqpDataTxUnit::Execute(TOperation::TPtr op, TTransactio
         };
 
         auto [validated, brokenLocks] = op->HasVolatilePrepareFlag()
-            ? KqpValidateVolatileTx(tabletId, sysLocks, kqpLocks, useGenericReadSets, 
+            ? KqpValidateVolatileTx(tabletId, sysLocks, kqpLocks, useGenericReadSets,
                 txId, tx->DelayedInReadSets(), awaitingDecisions, outReadSets)
             : KqpValidateLocks(tabletId, sysLocks, kqpLocks, useGenericReadSets, inReadSets);
 
@@ -219,12 +219,12 @@ EExecutionStatus TExecuteKqpDataTxUnit::Execute(TOperation::TPtr op, TTransactio
 
         NKqp::NRm::TKqpResourcesRequest req;
         req.MemoryPool = NKqp::NRm::EKqpMemoryPool::DataQuery;
-        req.Memory = txc.GetMemoryLimit();
+        req.ExternalMemory = txc.GetMemoryLimit();
         ui64 taskId = dataTx->GetFirstKqpTaskId();
         NKqp::GetKqpResourceManager()->NotifyExternalResourcesAllocated(txId, taskId, req);
 
         Y_DEFER {
-            NKqp::GetKqpResourceManager()->NotifyExternalResourcesFreed(txId, taskId);
+            NKqp::GetKqpResourceManager()->FreeResources(txId, taskId);
         };
 
         LOG_T("Operation " << *op << " (execute_kqp_data_tx) at " << tabletId