From 7e147f035f22dce102616c2c13878f94f17c338c Mon Sep 17 00:00:00 2001
From: Semyon <yentsovsemyon@ydb.tech>
Date: Fri, 11 Oct 2024 14:29:16 +0300
Subject: [PATCH] fix groups allocation cleaning (#10274) (#10325)

Co-authored-by: ivanmorozov333 <ivanmorozov@ydb.tech>
---
 ydb/core/tx/limiter/grouped_memory/service/ids.cpp   | 11 +++++++++++
 ydb/core/tx/limiter/grouped_memory/service/ids.h     |  1 +
 ydb/core/tx/limiter/grouped_memory/service/process.h |  2 +-
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/ydb/core/tx/limiter/grouped_memory/service/ids.cpp b/ydb/core/tx/limiter/grouped_memory/service/ids.cpp
index 163c90efcf12..d550e6374ef3 100644
--- 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
index acaa700411ab..93a0eaf120f3 100644
--- 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
index 4695bb2a444e..3a53ff542750 100644
--- 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);