fix groups allocation cleaning (ydb-platform#10274) (ydb-platform#10325)

Co-authored-by: ivanmorozov333 <ivanmorozov@ydb.tech>
zverevgeny · Oct 15, 2024 · 7e147f0 · 7e147f0
1 parent 5004469
commit 7e147f0
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 1 deletion.
diff --git a/ydb/core/tx/limiter/grouped_memory/service/ids.cpp b/ydb/core/tx/limiter/grouped_memory/service/ids.cpp
@@ -12,6 +12,17 @@ ui64 TIdsControl::ExtractInternalIdVerified(const ui64 externalId) {
     return result;
 }
 
+std::optional<ui64> TIdsControl::ExtractInternalIdOptional(const ui64 externalId) {
+    auto it = ExternalIdIntoInternalId.find(externalId);
+    if (it == ExternalIdIntoInternalId.end()) {
+        return std::nullopt;
+    }
+    const ui64 result = it->second;
+    InternalIdIntoExternalId.erase(result);
+    ExternalIdIntoInternalId.erase(it);
+    return result;
+}
+
 std::optional<ui64> TIdsControl::GetInternalIdOptional(const ui64 externalId) const {
     auto it = ExternalIdIntoInternalId.find(externalId);
     if (it != ExternalIdIntoInternalId.end()) {

diff --git a/ydb/core/tx/limiter/grouped_memory/service/ids.h b/ydb/core/tx/limiter/grouped_memory/service/ids.h
@@ -29,6 +29,7 @@ class TIdsControl {
     }
 
     [[nodiscard]] ui64 ExtractInternalIdVerified(const ui64 externalId);
+    [[nodiscard]] std::optional<ui64> ExtractInternalIdOptional(const ui64 externalId);
 
     ui64 GetMinInternalIdVerified() const;
     ui64 GetExternalIdVerified(const ui64 internalId) const;

diff --git a/ydb/core/tx/limiter/grouped_memory/service/process.h b/ydb/core/tx/limiter/grouped_memory/service/process.h
@@ -135,7 +135,7 @@ class TProcessMemoryScope: public NColumnShard::TMonitoringObjectsCounter<TProce
     }
 
     void UnregisterGroup(const bool isPriorityProcess, const ui64 externalGroupId) {
-        if (auto internalGroupId = GroupIds.GetInternalIdOptional(externalGroupId)) {
+        if (auto internalGroupId = GroupIds.ExtractInternalIdOptional(externalGroupId)) {
             AFL_INFO(NKikimrServices::GROUPED_MEMORY_LIMITER)("event", "remove_group")("external_group_id", externalGroupId)(
                 "internal_group_id", internalGroupId);
             UnregisterGroupImpl(*internalGroupId);